[論文レビュー] 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

ひとこと要約

この論文では、明示的な2階微分計算を回避し、適応的最適化手法をサポートすることで、勾配ベースのメタラーニングの効率的で分散可能な学習を可能にするスケーラブルなメタラーニングフレームワークSAMAを紹介する。SAMAは、マルチGPU環境下で最大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にまたがる分散学習リソースを効率的に活用するための新規な通信-計算オーバーラップ戦略を実装する。
• SAMAを最初の順勾配に基づく暗黙的微分により実装するために、JaxOptフレームワークを活用し、明示的な2階微分計算を回避する。
• カスタム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 SAMAは、大規模ベンチマークにおいて、既存のGBML手法と比較してスループットとメモリ効率に顕著な向上を達成できるか。
• RQ5 SAMAは、大規模言語モデルにおけるデータプリーニングやファインチューニングといった実用的応用において一貫したパフォーマンス向上を実現できるか。

主な発見

• 単一GPU環境下では、SAMAはベースラインのメタラーニングアルゴリズムと比較して最大1.7倍のスループット向上と2.0倍のメモリ消費量削減を達成した。
• マルチGPU環境下では、SAMAはベースライン手法と比較して最大4.8倍のスループット向上と3.8倍のメモリ消費量削減を達成した。
• SAMAを用いたデータ最適化は、複数のテキスト分類ベンチマークでBERTおよびRoBERTaの大規模言語モデルのファインチューニング精度を向上させた。
• SAMAは、ImageNet-1kにおける画像分類タスクの小規模および大規模なデータプリーニングにおいて、最先端の結果を達成した。
• モデルサイズが増加するにつれて、SAMAとベースライン手法の間の効率格差が拡大し、SAMAはGPUメモリ使用量の増加が最小限に抑えられた。
• 4GPUでの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