[論文レビュー] Betty: An Automatic Differentiation Library for Multilevel Optimization

Sang Keun Choe, Willie Neiswanger+2 more|arXiv (Cornell University)|Jul 5, 2022

Metaheuristic Optimization Algorithms Research被引用数 5

ひとこと要約

Betty は、マルチレベル最適化（MLO）のための新しい自動微分ライブラリであり、最適応答ヤコビ行列の連鎖則に基づく合成を通じて、効率的でスケーラブルな勾配計算を可能にするデータフロー・グラフ抽象化を導入している。計算複雑度を 𝒪(d³) から 𝒪(d²) に低減し、混合精度およびデータ並列学習をサポートしており、11%高いテスト精度、14%低いGPUメモリ使用量、20%速い学習速度を達成するなど、従来のシステムに比べてエンドツーエンドの複雑なMLOパイプラインを実装可能にしている。

ABSTRACT

Gradient-based multilevel optimization (MLO) has gained attention as a framework for studying numerous problems, ranging from hyperparameter optimization and meta-learning to neural architecture search and reinforcement learning. However, gradients in MLO, which are obtained by composing best-response Jacobians via the chain rule, are notoriously difficult to implement and memory/compute intensive. We take an initial step towards closing this gap by introducing Betty, a software library for large-scale MLO. At its core, we devise a novel dataflow graph for MLO, which allows us to (1) develop efficient automatic differentiation for MLO that reduces the computational complexity from O(d^3) to O(d^2), (2) incorporate systems support such as mixed-precision and data-parallel training for scalability, and (3) facilitate implementation of MLO programs of arbitrary complexity while allowing a modular interface for diverse algorithmic and systems design choices. We empirically demonstrate that Betty can be used to implement an array of MLO programs, while also observing up to 11% increase in test accuracy, 14% decrease in GPU memory usage, and 20% decrease in training wall time over existing implementations on multiple benchmarks. We also showcase that Betty enables scaling MLO to models with hundreds of millions of parameters. We open-source the code at https://github.com/leopard-ai/betty.

研究の動機と目的

• 勾配ベースのマルチレベル最適化（MLO）の実装とスケーリングにおける主なボトル neck を解消すること、具体的には複雑な手動の勾配合成および高いメモリ/計算コスト。
• 理論的MLO研究と実践的ソフトウェアシステムの間のギャップを埋めることで、効率的でモジュラーかつスケーラブルなMLO実装を可能にすること。
• 二段階最適化を超える、三段階およびそれ以上の階層を持つ任意のMLOプログラム構造をサポートすること。
• 混合精度やデータ並列トレーニングなどのシステムレベルの機能を提供することで、メモリと計算のオーバーヘッドを低減すること。
• 従来のライブラリでは不可能だった、数億パラメータを有する大規模MLOアプリケーションを可能にすること。

提案手法

• 最適化問題を方向付き依存関係を持つ階層的ノードとして構造化する、MLO問題の新しいデータフロー・グラフ表現を導入する。
• データフロー・グラフの逆走査を用いて、最適応答ヤコビ行列の繰り返し乗算により、連鎖則を効率的に適用して勾配を計算する。
• 重複するヘッセ行列の計算を回避し、ヤコビ行列の合成を最適化することで、自動微分の計算複雑度を 𝒪(d³) から 𝒪(d²) に低減する。
• アルゴリズム的選択（例：AID, ITD）とシステム最適化（例：混合精度、データ並列化）を分離する、モジュラーなソフトウェアスタックを設計する。
• 1ステップの勾配降下更新を各問題ごとに抽象化した「ステップ」関数として定義し、合成可能で再利用可能なMLOプログラム構築を可能にする。
• JAXおよびPyTorchエコシステムと統合することで、自動微分および分散トレーニングを含む現代のディープラーニングシステムをサポートする。

実験結果

リサーチクエスチョン

• RQ1 データフロー・グラフ抽象化を用いることで、連鎖則に基づくヤコビ行列合成を通じて、マルチレベル最適化における勾配計算を体系的かつ効率的に行えるか？
• RQ2 MLOにおける自動微分を最適化することで、計算複雑度を 𝒪(d³) から 𝒪(d²) に低減できるか？
• RQ3 二段階最適化を超える任意のMLOプログラム構造をサポートしつつ、モジュラリティとパフォーマンスを維持できるソフトウェアライブラリを設計できるか？
• RQ4 混合精度やデータ並列トレーニングなどのシステムレベル最適化が、MLOにおけるスケーラビリティ向上とメモリ使用量低減にどの程度寄与できるか？
• RQ5 提案されたライブラリにより、従来のツールでは達成不可能だった、数億パラメータを有するMLOモデルのトレーニングが可能になるか？

主な発見

• Betty は、データフロー・グラフの逆走査を用いて最適応答ヤコビ行列を効率的に合成することで、MLOにおける自動微分の計算複雑度を 𝒪(d³) から 𝒪(d²) に低減している。
• ライブラリは、既存の実装と比較してベンチマークタスクで最大11%高いテスト精度を達成しており、最適化の安定性と収束性の向上を示している。
• 最適化された勾配計算とシステムレベルの機能のおかげで、複数のベンチマークにおいてGPUメモリ使用量を最大14%、トレーニングのウォールタイムを最大20%低減している。
• ライブラリは、数億パラメータを有するモデルへのMLOのスケーリングに成功し、従来不可能だった大規模トレーニングシナリオを可能にしている。
• Betty は、混合精度やデータ並列トレーニングなどの多様な最適応答ヤコビ行列アルゴリズムとシステム最適化を、低レベルのコード変更なしにモジュラーに統合可能である。
• Apache-2.0 ライセンスでオープンソースとしてリリースされ、すべての実験の完全な実装とMLOプログラムの例が含まれており、再現性とコミュニティによる拡張性を確保している。