2022年のGoogleのAI研究の成果と今後の展望～MLとコンピュータシステム編～(3/3)

１．2022年のGoogleのAI研究の成果と今後の展望～MLとコンピュータシステム編～(3/3)まとめ

・モジュール式なコンパイラLLVMにML技術を体系的に統合した産業レベルのフレームワークであるMLGOをリリース
・MLの実践者はモデル学習時の二酸化炭素換算排出量を「4M」を考慮する事で桁違いに削減できる事を示した
・2023年にはトレーニングのマシン、機械化、場所を慎重に選択する事でCO2eを最大20倍まで削減する方法を実証予定

２．AIと持続可能性

以下、ai.googleblog.comより「Google Research, 2022 & beyond: Responsible AI」の意訳です。元記事の投稿は2023年2月2日、Mangpo PhothilimthanaさんとAdam Paszkeさんによる投稿です。

モデルをトレーニングする際の二酸化炭素の排出はそんなに大したレベルではなさそうに思えるかもしれませんが、Googleが結構気にしている理由の１つは過去に「自動車を製造する際のCO2排出量を100倍過大評価し、その数字を自動車を運転する際のCO2排出量として使用するような試算」で、批判を受けた事があるからなんですね。この過去の経緯を知るとやや気候変動に関しては懐疑的な気持ちになるかもしれませんが、それでも「気候変動に対してAIは何が出来るのか？」を一読しておく事をお勧めします。

アイキャッチ画像はstable diffusionのカスタムモデルによる生成

制約付き探索空間をナビゲートする機械学習

ハードウェアとワークロードを効率的に変更する以外に、ワークロードを複数のデバイスにマッピングするパーティショナーや、ワークロードをハードウェアが理解できる低レベルの表現に変換するコンパイラーなどのミドルレイヤーも最適化することができます。昨年、私たちは、より良いデバイス配置とコンパイラの決定を見つけるために、どのように機械学習(ML:Machine Learning)を適用することができるかを示しました。過去一年以内に、私たちはこの方向性をさらに追求し、多くの最適化探索空間が大きな制約を受けており、有効な解が非常にまばらであることを発見しました。

この課題に対処するため、私たちは、学習済みモデルが制約のある探索空間を効果的に案内できるようにするためのいくつかの技術を開発しました。Telamallocは、複数の選択肢がある場合に、MLモデルと経験則を組み合わせて意思決定を行い、さらに依存する意思決定を推論するために制約解決法(constraint solver)を活用します。Telamallocは、Edge TPUコンパイラのメモリ割り当てパスを整数線形プログラミング手法よりも高速化し、他の方法ではサポートされない重要な実世界のモデルを可能にします。

「A Transferable Approach for Partitioning Machine Learning Models on Multi-Chip-Modules」は、少し異なるアプローチを提案しています。これは、強化学習(RL:Reinforcement Learning)を適用して、1つのステップで意思決定を提案し、制約解決法が提案した解決策が有効であるように調整するよう依頼するものです。Edge TPUベースのマルチチップメッシュ上のBERTモデルについて、このアプローチは、非学習型探索戦略と比較してはるかに少ない時間予算で、デバイス間でより良い分布のモデルを発見します。

大規模生産システム向けML

私たちは、実稼働している様々な大規模システムの効率を向上させるために、MLを導入しています。

最近、モジュール式で再利用可能なコンパイラとツールのコレクションであるLLVMインフラストラクチャにML技術を体系的に統合する最初の工業使用に耐えうる汎用的なフレームワークであるMLGOをリリースしました。

MLGOは、LLVMの経験則をRLポリシーに置き換えて、最適化の判断を行うことができます。社内の大規模アプリケーションのセットでテストしたところ、学習されたポリシーは、インライン化の決定を最適化する際にバイナリサイズを3%～7%削減でき、レジスタ割り当ての意思決定を最適化する際にスループットを0.3%～1.5%改善できることがわかりました。

数年前に発表された学習済みコストモデルであるXLAが、最近、Google製品のMLコンパイラに導入され、トップレベルのMLワークロードのTPUカーネルの最適タイルサイズの選択をガイドし、Googleのデータセンター全体のTPU計算時間の2%を削減しました。

また、YouTube のキャッシュ交換アルゴリズムにおける既存の経験則を、単純な経験則と学習済みモデルを組み合わせた新しいハイブリッドアルゴリズムに置き換え、ピーク時のバイトミス率を9%向上させました。

関数インライン化におけるMLGOの説明図
“#bbs”, “#users”, “callsite height” は、caller-callee ペアの特徴量の例です。

AIと持続可能性

世界的な気候変動の危機を考えると、MLが環境に与える影響について懸念がある事は理解できます。最近の論文では、最善の方法に従うことで、MLの実践者はトレーニングによる二酸化炭素換算排出量(CO2e)を桁違いに削減できることを示しました。私たちはその実践を4M(Model, Machine, Mechanization, Map Optimization)と呼んでいます。

(1)モデル(Model)
最初のステップは、最も効率的なMLモデルアーキテクチャを選択することです。例えば、Primerは、4年前に開発された人気のTransformerと比較して、同じハードウェアで同じ品質スコアを達成しながら～4倍速く実行されます。

(2)マシン(Machine)
2つ目は、最もエネルギー効率の高いコンピュータを使用することです。例えば、2017年にTransformerモデルが初めて公開されたとき、人気のGPUはNvidia P100でした。TPU v4のようなMLトレーニングに最適化された最近のプロセッサを使用すると、ワットあたりの性能が～15倍向上します。

(3)設備の高度化(Mechanization)
学習用のコンピュータは、データセンターに収容する必要がありました。大規模なクラウドデータセンターは、一般的な小規模なオンプレミスのデータセンターと比較して、エネルギー効率が～1.4倍高いです。

(4)地域の最適化(Map Optimization)
今回の調査で最大の驚きは、最適な場所を選ぶことで、エネルギー供給のクリーン度に影響を与えるということでした。しかも、クラウドでは、4つの要素の中で立地が最も変更しやすいのです。一般的な場所とよく選ばれた場所との差は、同じ国内でも～9倍にもなるのです。

この例では、4Mを掛け合わせると、2017年のGPUを使ったオリジナルのTransformerモデルのトレーニングからで、4年間で4 × 15 × 1.4 × 9、つまり～750倍のCOe削減が可能になります。

私達はこの分野の探求を続けており、2023年には、トレーニングのマシン、機械化、場所を慎重に選択することで、現在のモデルトレーニングのCO2eを最大20倍まで削減する方法を実証する、さらなる研究を発表する予定です。