[论文解读] Making Scalable Meta Learning Practical

Sang Keun Choe, Sanket Vaibhav Mehta+5 more|arXiv (Cornell University)|Oct 9, 2023

Domain Adaptation and Few-Shot LearningComputer Science被引用 3

一句话总结

该论文提出 SAMA，一种可扩展的元学习框架，通过避免显式计算二阶梯度并支持自适应优化器，实现了基于梯度的元学习的高效分布式训练。在多 GPU 环境下，SAMA 实现了最高 4.8 倍的吞吐量提升和 3.8 倍的内存使用降低，展示了在 BERT 和 RoBERTa 模型的数据剪枝与微调任务中的最先进性能。

ABSTRACT

Despite its flexibility to learn diverse inductive biases in machine learning programs, meta learning (i.e., learning to learn) has long been recognized to suffer from poor scalability due to its tremendous compute/memory costs, training instability, and a lack of efficient distributed training support. In this work, we focus on making scalable meta learning practical by introducing SAMA, which combines advances in both implicit differentiation algorithms and systems. Specifically, SAMA is designed to flexibly support a broad range of adaptive optimizers in the base level of meta learning programs, while reducing computational burden by avoiding explicit computation of second-order gradient information, and exploiting efficient distributed training techniques implemented for first-order gradients. Evaluated on multiple large-scale meta learning benchmarks, SAMA showcases up to 1.7/4.8x increase in throughput and 2.0/3.8x decrease in memory consumption respectively on single-/multi-GPU setups compared to other baseline meta learning algorithms. Furthermore, we show that SAMA-based data optimization leads to consistent improvements in text classification accuracy with BERT and RoBERTa large language models, and achieves state-of-the-art results in both small- and large-scale data pruning on image classification tasks, demonstrating the practical applicability of scalable meta learning across language and vision domains.

研究动机与目标

• 解决基于梯度的元学习因计算/内存成本过高且缺乏分布式训练支持而导致的可扩展性差的问题。
• 克服现有隐式微分方法在使用 Adam 等自适应优化器时的不稳定性和低效性，尤其针对大模型。
• 提出一种系统-算法协同设计的整体方法，以实现在真实世界 NLP 和视觉应用中元学习的实际部署。
• 实现在多样化任务（包括数据优化、剪枝和大模型如 BERT 与 RoBERTa 的微调）中高效且可扩展的元学习。
• 通过支持现代深度学习工作流中的自适应优化器和分布式训练，弥合理论元学习与实际部署之间的差距。

提出的方法

• 通过使用单位矩阵近似基础雅可比逆矩阵，消除昂贵的矩阵求逆操作，降低计算开销。
• 通过链式法则扩展元梯度计算，以支持 Adam 等自适应优化器，实现与现代深度学习模型的兼容性。
• 实现一种新颖的通信-计算重叠策略，以高效利用多 GPU 环境下的分布式训练资源。
• 利用 JaxOpt 框架，通过一阶梯度隐式微分实现 SAMA，避免显式计算二阶梯度。
• 设计一种系统优化的流水线，支持元梯度反向传播计算，无需自定义 autograd 扩展。
• 将原本专为一阶方法设计的高效分布式训练技术整合到元学习框架中，实现跨多 GPU 的可扩展性。

Figure 1: Top: Table showing a scalability comparison. Bottom left: Plot of throughput vs memory of different GBML algorithms on the noisy finetuning of BERT-base experiment. SAMA achieves better memory/compute efficiency overall given a fixed model, and the gap further widens by distributing comput

实验结果

研究问题

• RQ1 能否使元学习中的隐式微分在使用 Adam 等自适应优化器训练大模型时具备可扩展性和高效性？
• RQ2 如何在不牺牲训练稳定性的前提下，降低基础雅可比逆矩阵计算的计算与内存开销？
• RQ3 一阶优化的分布式训练技术在多大程度上可被适配以支持可扩展的元学习？
• RQ4 与现有 GBML 方法相比，SAMA 在大规模基准测试中能否实现吞吐量和内存效率的显著提升？
• RQ5 SAMA 是否能在实际应用（如大语言模型的数据剪枝与微调）中实现一致的性能增益？

主要发现

• 在单 GPU 环境下，SAMA 相较基线元学习算法，吞吐量最高提升 1.7 倍，内存消耗降低 2.0 倍。
• 在多 GPU 环境下，SAMA 相较基线方法，吞吐量最高提升 4.8 倍，内存消耗降低 3.8 倍。
• 基于 SAMA 的数据优化方法在多个文本分类基准上提升了 BERT 和 RoBERTa 大型语言模型的微调准确率。
• SAMA 在 ImageNet-1k 图像分类任务中，实现了小规模与大规模数据剪枝的最先进结果。
• 随着模型规模增大，SAMA 与基线方法之间的效率差距进一步扩大，SAMA 展现出最小的 GPU 内存使用增长。
• 在 4 块 GPU 上，SAMA 实现了相比 MWN（缺乏分布式训练支持）15–20 倍的加速。

Figure 2: The overall workflow of meta gradient computation with SAMA in the distributed data parallel setting. In detail, SAMA consists of three first-order backward passes performed with the underlying automatic differentiation engine, and one manual backward pass for algorithmic adaptation for th