4ビットAdamW、丸め方の再設計で精度差縮小

研究の概要

Hanyang Li氏らの研究チームは、大規模言語モデルの学習で広く使われる最適化手法AdamWについて、内部に保持する状態量を4ビットへ圧縮する際の「丸め方」を根本から設計し直した。AdamWは学習の進み方を調整するため、勾配の平均にあたる一次モーメントと、勾配の二乗の平均にあたる二次モーメントを、パラメータごとに記憶する。この記憶領域は、モデル本体と同程度かそれ以上のメモリを占めるため、圧縮の価値が大きい。

従来の4ビット化では、圧縮で生じる誤差が毎ステップの再帰計算を通じて蓄積し、後続の更新量を乱す問題があった。研究チームは、二次モーメントのゼロ付近の量子化区間を局所的に解析し、状態そのものの平均誤差が小さくても、次の更新で使われる「プレコンディショナー」の平均誤差は小さいとは限らないことを示した。プレコンディショナーとは、二次モーメントの平方根の逆数として更新量を調整する係数である。

提案は二つある。​ZIP-SR​は、二次モーメントの符号表にゼロを残し、確率的丸めの確率をプレコンディショナー空間で計算する。​ZE-EDEN​は、ゼロを除いた符号表を使い、量子化ブロックを再スケールして正の下限値による歪みを補正する。いずれも一次モーメントにはNF4を用い、学習終盤10%ではLMヘッドの一次モーメントに確率的丸めを適用する。

130Mから2.7Bパラメータの GPT系・Llama系の事前学習実験で、両手法はTorchAOの4ビットAdamWと32ビットAdamWとの検証損失の平均差を全モデルサイズで縮め、最大の縮小幅は​70%​に達した。全パラメータの教師あり微調整でも、TorchAOより低い検証損失を示し、下流タスクでは32ビット版に近い水準を保った。

しくみをやさしく

量子化とは、細かな数値を少数の「目盛り」に置き換える作業である。4ビットでは目盛りが16段階しかなく、元の値は隣り合う二つの目盛りのどちらかに丸められる。確率的丸めは、近い側へ機械的に寄せるのではなく、距離に応じた確率でどちらかを選ぶ方式であり、平均すると偏りが出にくい。

本研究の核心は、「どの座標で」確率を計算するかにある。更新量は二次モーメントそのものではなく、その平方根の逆数に比例する。二次モーメントがゼロ近傍にあると、わずかな丸め誤差が逆数を通じて極端に増幅される。そこで、更新量に直結するプレコンディショナーの側で丸めの確率を決めれば、実際に効く量の平均が保たれる。これがZIP-SRの発想である。

身近な例では、売上そのものではなく、最終的に意思決定に使う利益率の側で端数処理を揃えるようなものである。ZE-EDENは別の経路で、ゼロを目盛りから外す代わりに、ブロック単位で全体の大きさを調整して歪みを打ち消す。

ビジネスへの示唆

最も恩恵を受けるのは、自社データでLLMを継続事前学習・微調整している組織である。具体的には、金融機関のAI推進部門が社内文書で行うモデル調整、製造業のDX部門が設備保全記録で行う専用モデル構築、医療・法務・コールセンター向けの業界特化モデル開発が該当する。

読者に問いたいのは、自部署の学習・微調整業務で「GPUメモリ不足を理由に、モデルを小さくした、バッチを減らした、外部クラウドを借りた」場面がないか、という点である。メモリの余裕が生まれれば、同じGPU台数でより大きなモデルを扱える、あるいは同じモデルをより少ない台数で回せる可能性がある。

影響を受けやすいKPIは次のとおりである。

• 学習1回あたりのGPUコスト、および実験の試行回数
• 微調整済みモデルの業務精度（回答の正確性、分類の適合率など）
• 開発リードタイムと、MLエンジニアのインフラ調整工数

ただし、本論文の評価は事前学習と微調整の損失・下流タスクが中心であり、自社業務の精度が同様に維持される保証はない。また圧縮対象は最適化状態であり、推論コストが直接下がるわけではない点にも留意が必要である。

現場での導入イメージ

仮に導入支援に入るなら、次の流れが現実的である。

1. ​棚卸し​：社内で実施している学習・微調整ジョブを洗い出し、メモリ不足で見送った案件、モデル縮小やバッチ縮小の妥協をした案件を可視化する。到達点は、最適化状態の圧縮が効く案件の優先順位表を持つことである。

2. ​差し込み先の設計​：新しい学習基盤を増やさず、すでに使っているPyTorch系の学習スクリプトで、オプティマイザの差し替えとして試せる箇所を特定する。比較基準はTorchAOの4ビットAdamWと32ビットAdamWとする。到達点は、既存の学習フローを変えずに切り替えられる構成の確定である。

3. ​小さく試す（PoC）​：一部署の小規模モデルで、数週間の範囲で32ビット版と並走させる。検証損失に加え、担当者が業務で使う評価セットの出来を見て、実感として差が許容範囲かを確かめる。到達点は、自社データで精度差が小さいという確認である。

4. ​運用と横展開​：採用可否の基準、学習終盤での丸め設定などの運用ルールを文書化し、他部署の案件へ広げる。到達点は、メモリ制約を理由にした妥協が減り、実験の回数を増やせる状態である。

全体として目指すのは、精度を大きく損なわずに、自社専用モデルの開発を小さく・安く回し続けられる体制である。

今後の展望

今回の検証は最大2.7Bパラメータまでであり、より大規模なモデルや長期学習での挙動は今後の確認課題である。一方、丸めの座標を使途に合わせて選ぶという考え方は、他の圧縮手法にも応用できる可能性がある。学習コストが競争力を左右する局面で、地味ながら効く基盤技術として注目される。