Muonを超えて: MUD（MomentUm Decorrelation）による高速なトランスフォーマー学習

要旨: Muon のような直交化モーメンタム最適化器は、短い極分解イテレーションを介して、行列値モーメンタム更新をほぼホワイトニング/直交化することにより、トランスフォーマーの訓練を改善します。しかし、極因子近似は通常、複数の大規模な行列乗算を必要とし、その結果のオーバーヘッドはかなり大きく、ハードウェア依存性が高くなることがあります。我々は MUD（MomentUm Decorrelation）を導入します。これは Muon の極更新を、古典的なグラム-シュミットおよびガウス-ザイデルの考え方に触発された三角行列（コレスキー風）ホワイトニング代理手法に置換する補完的なホワイトニング手法です。行直交行列が MUD 写像の不動点であることを示し、内部ステップをグラム矩陣の対称ガウス-ザイデル前処理に関連づけ、不動点の近傍で二次的局所収束を証明します。パープレキシティまでの実時間の観点では、MUD は調整済み AdamW および Muon に対して、10-50% の実時間改善を一貫して提供します。通常、Muon よりは各ステップの収束がわずかに遅くなるものの、オプティマイザのオーバーヘッドは大幅に低くなります。Muon と比較して、ほとんどの設定でピークトークン/秒を約 1.3-2.6 倍向上させ、A100 上の GPT-2 Large ではほぼ 3 倍に達します。さらに、ESM-2 150M タンパク質言語モデルの訓練を示し、MUD が Muon レベルの検証パープレキシティと同等の値を、はるかに短い実時間で達成することを示します。