大規模AIモデル（Transformerや言語モデルなど）の訓練は、AI分野において不可欠な要素となっていますが、高い計算コスト、メモリ消費、エネルギー需要という課題も抱えています。例えば、OpenAIのGPT-3は1750億個のパラメータを持ち、GPUによる訓練に数週間を要します。このような莫大なリソース需要は、大規模な計算リソースを持つ組織でのみこの技術の応用を制限するだけでなく、エネルギー効率と環境への影響に対する懸念も高めています。これらの課題を解決することは、AI発展のより広範な普及と持続可能性を確保するために極めて重要です。
従来の訓練方法は非効率的で、革新的な解決策が求められています
大規模モデルの訓練が非効率的な主な原因は、密行列への依存にあります。これにより、膨大なメモリと計算能力が必要となります。現代のGPUは最適化された低精度または低ランク演算に対するサポートが限られており、これらの需要をさらに増大させています。行列分解やヒューリスティックなランク削減などの手法が提案されてきましたが、実用上は依然として制約があります。例えば、GaLoreは単一バッチ設定下で訓練できますが、非現実的な実行時オーバーヘッドが存在します。同様に、低ランクアダプターLTEを採用した場合、大規模タスクでの収束性に問題があります。メモリ使用量、計算コスト、訓練時間を同時に削減し、性能を損なわない方法が不足しているため、革新的な解決策の必要性が切迫しています。
CoMERAフレームワーク：適応的テンソル最適化による効率的な訓練
ニューヨーク州立大学オルバニー校、カリフォルニア大学サンタバーバラ校、Amazon Alexa AI、Metaの研究者らが共同で、CoMERA（Computing-and Memory-Efficient training method via Rank-Adaptive tensor optimization）と呼ばれる新しいフレームワークを発表しました。このフレームワークは、メモリ効率と計算速度を両立させ、適応的ランクテンソル圧縮技術によって実現しています。従来の方法が圧縮のみに焦点を当てているのに対し、CoMERAは多目的最適化手法を用いて、圧縮率とモデル精度をバランス良く調整します。テンソル化された埋め込みと高度なテンソルネットワーク縮約を利用してGPU使用率を最適化することで、実行時オーバーヘッドを削減しながら、高い性能を維持します。また、CUDAグラフを導入することで、GPU操作中のカーネル起動遅延（従来のテンソル圧縮手法における主要なボトルネック）を最小限に抑えています。
CoMERAの基礎は適応的テンソル表現であり、モデル層がリソース制約に応じてランクを動的に調整することを可能にします。テンソルランクを変更することで、ニューラルネットワーク操作の完全性を損なうことなく圧縮を実現します。この動的な最適化は、2段階の訓練プロセスによって実現されます。
初期段階：安定した収束に重点を置きます。
後期段階：特定の圧縮目標を満たすようにランクを微調整します。
6エンコーダTransformerモデルにおいて、CoMERAは初期段階で最大43倍の圧縮率を実現し、後期最適化段階では最大361倍の圧縮率を実現しました。さらに、GaLoreと比較して、メモリ消費量を9倍削減し、1回の訓練速度を2～3倍向上させました。
CoMERAの優れた性能を示す複数のテスト結果
MNLIデータセットで訓練されたTransformerモデルに適用した場合、CoMERAはモデルサイズを256MBからわずか3.2MBに縮小しながら、精度を維持しました。DLRMなどの大規模推薦システムでは、CoMERAはモデルを99倍圧縮し、ピークメモリ使用量を7倍削減しました。このフレームワークは、事前訓練されたCodeBERT（特定分野の大規模言語モデル）においても優れた性能を発揮し、全体で4.23倍の圧縮率を達成し、一部の訓練段階では2倍の高速化を実現しました。これらの結果は、様々なタスクやアーキテクチャに対応できる能力を示しており、様々な分野への適用可能性を広げています。
CoMERAフレームワークの主な利点のまとめ
この研究の主な結論は以下のとおりです。
CoMERAは、特定の層で最大361倍、モデル全体で最大99倍の圧縮率を実現し、ストレージとメモリの需要を大幅に削減しました。
このフレームワークは、Transformerと推薦システムの1回の訓練時間を2～3倍短縮し、計算リソースと時間を節約しました。
テンソル化表現とCUDAグラフを使用することで、CoMERAはピークメモリ消費量を7倍削減し、より小型のGPUでの訓練を可能にしました。
CoMERAの手法は、Transformerや大規模言語モデルを含む様々なアーキテクチャをサポートし、精度を維持または向上させます。
訓練に必要なエネルギーとリソースを削減することで、CoMERAはより持続可能なAIの実践に貢献し、より幅広い層が最先端モデルを使用できるようになります。
