プリコンディショナー空間での丸め：4ビットAdamW最適化器状態量子化の再設計

AdamWは大規模モデル学習の標準的な最適化器だが、そのコストは過小評価されがちである。重みと勾配に加えて、各パラメータは一次モーメントと二次モーメントという二つの状態を持つ。32ビットで保存すると、この常駐メモリは重みそのものと同程度か、それ以上になる。最適化器の状態を4ビットに圧縮することは、学習時のメモリを減らす最も直接的な方法の一つである。TorchAOはすでに4ビットAdamWを提供しているが、32ビットAdamWとの間には測定可能な検証損失の差が残っている。2026年10月8日にarXiv cs.LGで公開されたHanyang Liら五名の論文は、コードブックの形やブロックサイズを調整するのではなく、より基本的な問いを立てる。量子化器が隣り合う二つの再構成レベルのどちらかを選ぶとき、どの座標で選ぶべきか、という問いである。

著者はこの座標を「丸め空間」と呼ぶ。重要なのは、量子化誤差がその場にとどまらないことだ。誤差はモーメントの再帰式を通って伝播し、以後のすべての適応的な更新を乱す。一次モーメントでは、更新量が状態にほぼ線形に依存するため、状態空間での丸めは大きな問題にならない。二次モーメントは事情が異なる。Adamは二次モーメントの平方根の逆数で割り、この逆数がプリコンディショナーである。論文はゼロに隣接する量子化セルを局所的に解析し、状態の平均誤差が小さくても、次のステップのプリコンディショナーの平均誤差が小さいとは限らないことを示した。状態がゼロから少しずれるだけで、急峻な逆数関数を通って誤差が拡大する。さらに一次元の二次関数による構成例で、状態空間での丸めとプリコンディショナー空間での丸めが、数値の違いにとどまらず、性質の異なる最適化の挙動を生むことを示している。

この観察から第一の手法ZIP-SR、すなわちゼロを含むプリコンディショナー空間の確率的丸めが生まれる。二次モーメントのコードブックにゼロを残し、確率的丸めの確率を状態空間ではなくプリコンディショナー空間で計算する。確率的丸めは期待値が不偏であり、決定的な系統誤差を平均ゼロの雑音に変える。空間を正しく選べば、その不偏性は更新幅を実際に決める量の上に置かれる。第二の手法ZE-EDENは補完的な道をとる。ゼロを含まない二次モーメントのコードブックを使うので、量子化値には正の下限があり、ゼロ付近で逆数が発散しない。ただしその下限自体がプリコンディショナーを歪めるため、EDEN方式の校正で量子化後の二次モーメントのブロックを再スケールし、歪みを打ち消す。どちらの構成も、一次モーメントには4ビットのNormalFloat、つまりNF4を使い、学習の最後の10%では言語モデルの出力ヘッドの一次モーメントに対して狙いを絞った確率的丸めを行う。

実験はGPT型とLlama型の事前学習を対象とし、モデル規模は1億3000万から27億パラメータに及ぶ。評価したすべての規模で、両手法はTorchAOの4ビットAdamWと32ビットAdamWの平均検証損失差を縮め、報告された最大の縮小幅は70%に達した。全パラメータの教師ありファインチューニングでは、両方のレシピがTorchAOより低い検証損失を達成し、下流タスクでは32ビットAdamWに近い水準を保った。一度だけの大きな勝利よりも、一貫性のほうが重要である。すべての規模で改善が現れるなら、それは偶然ではなく仕組みの効果を示唆する。

私たちの見立てでは、この論文の主な貢献は設計空間の引き直しにある。従来の議論はコードブックの設計やブロックの切り方に集中していた。本論文は第三の軸、すなわち丸めが行われる座標を加えた。この視点は、保存した状態に非線形な写像を適用する他の最適化器にも移せる可能性が高い。ただし慎重さも要る。要旨にはメモリ削減量やスループットの数値がなく、最大のモデルは27億パラメータで、より大きな規模や長い学習でも効果が続くかは独立した再現を待つ必要がある。学習メモリに悩むチームにとっては、本番のレシピを直接置き換えるのではなく、自分たちの規模でTorchAOと比べる小規模な試験の候補として扱うのが妥当だ。