MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates
本論文では、ランク1更新とシャーマン・モリソンに基づく逆行列計算を利用することで、計算量と通信量をそれぞれ3次から2次および1次へと削減し、それによって最先端の1次および2次手法と比較して大規模言語モデルの学習と収束を大幅に加速させる、モーメンタムを有効化した2次最適化手法であるMKORを提案する。
現代の人工知能という広大な風景において、最も強力なシステムは、パターンを認識し、言語を翻訳し、テキストを生成することを学習する複雑な数学的構造である、ディープニューラルネットワークの上に構築されています。これらのシステムを教えるために、研究者は巨大な最適化問題、すなわち、ネットワークが正確な予測を行うことを可能にする、内部の数値（重み）の完璧なセットを見つけ出すという問題を解決しなければなりません。数十年にわたり、標準的なアプローチは、足元の傾斜を確認しながら霧の中の山を下るハイカーのように、誤差を最も速く減少させる方向へとこれらの重みを押し進めることでした。この一次形式の手法は信頼性が高く広く使用されていますが、谷の底に到達するまでにしばしば数千ステップを要するなど、速度が遅くなることがあります。
二次最適化として知られる、より野心的な戦略は、ハイカーに地形全体の地図（地面の曲率を含む）を与えることで、より高速な下降を可能にしようとするものです。傾斜がどのように変化するかを理解することで、これらの手法は理論的には、より少ないステップで最適解に到達できます。しかし、このアプローチは歴史的に、今日使用されている大規模なモデルにとって計算コストが高すぎ、実用的ではありませんでした。大規模なネットワークに必要な曲率情報を計算するには膨大な処理能力とメモリが必要であり、その結果、この手法が置き換えを目指している単純なアプローチよりも遅くなってしまうことがよくあります。課題は、地図を運ぶという重い負担を負うことなく、地図を読むハイカーのスピードをどのように捉えるかという点にありました。
トロント大学、ラトガーズ大学、およびテキサス高度計算センターの研究チームは、このジレンマを解決するMKORと呼ばれる新しいツールを開発しました。彼らは、特定の数学的なショートカットを使用して、あらゆる詳細を計算することなく地形の形状を近似する、モーメンタム（慣性）を備えたオプティマイザを作成しました。損失関数の地形の完全で複雑な曲率を計算しようとする代わりに、MKORはランク1更新と呼ばれる手法を使用しています。このアプローチは、変化の最も重要な方向に焦点を当てることで計算を簡素化し、計算コストを実質的に立方体の関係から二乗の関係へと削減します。実用的な観点では、これはモデルが大きくなるにつれて手法のスケールアップが大幅に向上することを意味し、これまで二次形式の手法が大規模言語モデルで使用されるのを妨げてきたボトルネックを回避します。
研究者たちは、BERT-Large-Uncasedとして知られる大規模言語モデルを用いて、64基の高性能グラフィックスプロセッサで実験を行い、この新しいオプティマイザをテストしました。彼らはMKOKを、既存の最高の一次形式の手法および最新の二次形式の実装と比較しました。結果は驚くべきものでした。質問回答を伴うタスクにおいて、標準的な一次形式の手法は特定の精度に達するまでに1,536ステップを必要としました。新しいMKOR法は、わずか600ステップで同等またはそれ以上の精度に達しました。トレーニングを完了するために必要な総時間を測定したところ、標準的な手法がNearly 8時間を要したのに対し、MKORはわずか3時間強でタスクを完了しました。これは2.57倍の高速化を意味し、実質的にトレーニング時間を半分以下に短縮しています。
純粋な速度を超えて、この研究は、二次形式の手法をしばしば悩ませる決定的な安定性の問題にも対処しました。これらの手法は複雑な数学的要因に基づいて学習プロセスを調整するため、「勾配爆発」を起こしやすい、つまり調整が大きくなりすぎて学習プロセスがクラッシュしたり発散したりする現象が発生しやすいのです。従来の解決策では、安全性を維持するために学習率を大幅に遅くする必要があることが多く、それが速度の利点を打ち消していました。MKORには、これらの調整の大きさを監視する組み込みのスタビライザーが含まれています。値が大きくなりすぎると、システムは自動的にそれらを再スケーリングし、二次形式のアプローチの積極的なペースを維持しながら、学習が破綻するのを防ぎます。これにより、オプティマイザは不安定性のリスクなしに迅速に移動できます。
研究者たちはまた、二次形式の手法の利点は、モデルがデータの広範な輪郭を学習しているトレーニングの初期段階で最も顕著になることも観察しました。モデルが最適解に近づくにつれて、複雑な曲率計算の恩恵は減少します。これに対処するため、彼らはMKOR-Hと呼ばれるハイブリッド版のオプティマイザを導入しました。このバージョンは、強力な二次形式の手法で初期の学習フェーズを加速させますが、トレーニングの損失が急速に減少しなくなると、自動的により単純な一次形式の手法に切り替わります。このハイブリッドアプローチにより、システムはさらに少ないステップで収束し、標準的な手法と同じ最終精度を半分以下の時間で達成しました。
この研究では、なぜこの特定の数学的ショートカットがこれほど上手く機能するのかについても調査しました。研究者たちはニューラルネットワーク内のデータ構造を分析し、ネットワークの異なる部分がどのように相互作用するかを記述する共分散行列が、しばしば低ランクであることを発見しました。これは、データの高次元性にもかかわらず、情報が少数の主要な方向に集中していることを意味します。これらの行列をランク1更新を用いて近似することで、MKORは完全な行列計算のオーバーヘッドなしに、最適化に必要な本質的な情報を捉えます。チームは広範なテストを通じて、この近似が非常に少ない誤差しか導入しないことを検証し、簡素化が精度の犠牲の上に成り立っていないことを確認しました。
速度と安定性に加えて、この新手法は、トレーニング中に異なるプロセッサ間で転送する必要があるデータ量を大幅に削減します。数千のプロセッサが連携して動作する分散コンピューティングにおいて、通信は主要なボトルネックになり得ます。MKORはこの通信の複雑さを二乗の関係から線形の関係へと減少させます。これは、プロセッサが追加されるにつれて、データ転送の遅延によって停滞することなく、システムが効率的にスケールアップすることを意味します。この効率性により、この手法は、従来の二次形式のアプローチでは制限であった、現代の大規模言語モデルに求められる巨大なバッチサイズを扱うことが可能になります。
この研究の意義は、単に一つの特定のモデルをトレーニングすることに留まりません。大規模なシステムに対して二次最適化を実用的なものにすることで、MKORは人工知能のより効率的なトレーニングへの道を提供します。研究者たちは、彼らの手法が機械学習の実践者にとって、かなりのエネルギーと時間の節約に役立つ可能性があると指摘しました。技術自体は中立的なツールですが、その応用は医療診断から気候モデリングに至るまでの分野において科学的発見を加速させる可能性があります。しかし、著者らはまた、効率的な最適化アルゴメントは、社会に負の影響を与えるものを含むあらゆるモデルに適用可能であることを認め、そのような技術の責任ある使用は、それを用いる人々の選択に依存することを強調しています。
最終的に、この論文はディープラーニングにおける長年の問題に対する具体的な解決策を提示しています。モーメンタム、ランク1近似、およびハイブリッドな切り替えメカニズムを組み合わせることで、研究者たちは、従来のものよりも高速で、より安定し、よりスケーラブルなオプティマイザを作り上げました。実験は、二次形式の手法の力を、それらを傍観させてきた法外な計算コストなしに活用することが可能であることを示しています。人工知能の分野がより大規模で複雑なモデルへと進み続ける中で、MKORのようなツールは、これらのシステムをより効率的にトレーニングする方法を提供し、かつては理論的な利点であったものを実用的な現実へと変えています。
自分の分野の論文に埋もれていませんか？
研究キーワードに一致する最新の論文のダイジェストを毎日受け取りましょう——技術要約付き、あなたの言語で。