💻 computer science

AMUSE: Anytime Muon with Stable Gradient Evaluation

本論文は、Muon の迅速なバッチ処理の進展と Schedule-Free 平均化を組み合わせ、学習率スケジューリングの必要性を排除しつつトレーニングを安定化させ、視覚および言語タスクにおいて優れた性能を達成する新しいオプティマイザ AMUSE を紹介する。

原著者： Jueun Kim, Baekrok Shin, Jihun Yun, Beomhan Baek, Minhak Song, Chulhee Yun

公開日 2026-05-22

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Jueun Kim, Baekrok Shin, Jihun Yun, Beomhan Baek, Minhak Song, Chulhee Yun

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

以下は、論文「AMUSE: Anytime Muon with Stable Gradient Evaluation」を平易な言葉と創造的な比喩を用いて解説したものです。

全体像：AI の学習は山登りに似ている

あなたがコンピュータ（ニューラルネットワーク）に、猫を認識したり物語を書いたりするといったスキルを教えようとしていると想像してください。そのためには、コンピュータは「誤差」という巨大で複雑な山を下り、最も低い地点（完璧な解）を見つける必要があります。

この山の形状は厄介です。単なる滑らかな斜面ではありません。

川: 広々とした平坦で穏やかな谷底があり、ここでは素早く歩き、大きな進歩を遂げることができます。ここが真の学習が行われる場所です。
谷の壁: この谷の両側には、信じられないほど急峻で岩だらけの地面があります。壁の方へ一歩でも踏み出しすぎると、激しく跳ね返され、エネルギーを浪費して何の進歩も得られなくなります。

既存手法の問題点

長らく、この山を登る（最適化する）標準的な方法は、AdamWと呼ばれる厳格な地図を必要とする登山者のようなものでした。この地図は、どのくらいの速さで歩き、いつ減速すべきかを正確に指示します。登山者がこの地図を無視すれば、道に迷ってしまいます。

最近、2 つの新しい登山戦略が登場しました。

Schedule-Free (SF): これは地図を必要としない登山者のようなものです。彼らは歩き続け、谷底の平坦な場所に留まるために、時折自分がどこを歩いたか振り返ります。非常に安定していますが、スタートが遅い傾向があります。
Muon: これは超高速な新しい登山者です。特別なトリックを持っています。急峻な壁に引っかからないように、歩幅を「まっすぐ」にするのです。彼らは川を信じられないほど速く駆け下ります。しかし、彼らはあまりに速く攻撃的であるため、時折谷の壁に跳ね返され、揺れ動いて不安定になることがあります。

解決策：AMUSE

この論文の著者たちは、Muon は速いが不安定で、Schedule-Free は安定しているが時折遅いことに気づきました。彼らは、両者の長所を組み合わせることを目指しました。

彼らが生み出したのがAMUSE（Anytime Muon with Stable gradient Evaluation）です。

比喩：「スマートな補間」を行う登山者

あなたが川（谷）を走る曲がりくねった道を車で運転していると想像してください。

Muonは、時速 100 マイルでスポーツカーを運転するようなものです。早く着きますが、段差（急峻な壁）に当たると制御を失う可能性があります。
Schedule-Freeは、時速 40 マイルで重いトラックを運転するようなものです。非常に安定していますが、どこへ行くにも時間がかかります。

AMUSEは、時間帯に応じて運転スタイルを変えるスマートなドライバーのようなものです。

旅の初期（スタート）: ドライバーはスポーツカーモードです。素早く動き出し、距離を稼ぐために速く運転します（Muon のように）。スピードアップのためにいくつかのリスクを取ることをいといません。
旅の後半（ゴール）: 目的地に近づくにつれ、ドライバーは徐々に「トラックモード」へ移行します。これまでの平均的な経路をより重視し、旋回を滑らかにします。これにより壁に跳ね返されるのを防ぎ、完璧に平坦な川の底にとどまり続けます。

AMUSE の仕組み（技術的な魔法）

この論文は、「時間変数係数」（時間とともに変化するダイヤルという言い換え）を用いてこれを説明しています。

ダイヤル（ $\beta_t$ ）: 開始時には、ダイヤルは「速い」経路に焦点を合わせるように設定されています。学習が進むにつれ、ダイヤルはゆっくりと「安定した」経路に焦点を合わせるように回っていきます。
結果: AMUSE は、開始時にMuonの速度を、終了時にSchedule-Freeの安定性を獲得します。いつ減速すべきかを指示する事前の地図（学習率スケジューリング）は不要で、自らそれを判断します。

論文が明らかにした結果

著者たちは、この新しい手法を多くの異なる「山」（タスク）でテストしました。

画像認識: コンピュータに画像（猫、犬、手書きの数字など）を識別させる学習。
大規模言語モデル: チャットボットの背後にあるモデルのように、テキストを書き、理解する AI の学習。

結果:

高速: AMUSE は他の手法と同じ高品質な結果に達しましたが、ステップ数は少なくて済みました。例えば、大規模言語モデルのタスクでは、次点の手法よりも1.5 倍速くゴールに到達しました。
安定: Muon が時折見せるような揺れやクラッシュは起こりませんでした。
Anytime（いつでも）: 学習をいつでも停止でき、モデルは良好な状態にあります。「学習終了」の特定の瞬間を待つことなく、良い結果を得ることができます。

まとめ

この論文は、AI 学習のための新しいツールであるAMUSEを紹介しています。これは、Muonオプティマイザの不安定さ（速いが揺れ動く）を、Schedule-Free最適化からの安定化テクニックを借用することで修正します。

それは、いつスプリントし、いつ着実に歩くべきかを知っている登山者のようなものです。これにより、山頂から転落することなく、山の底に素早く到達することが保証されます。この論文は、画像やテキストのタスクにおいて、複雑なスケジューリングを必要とせずに、既存の標準的な手法よりも優れていると主張しています。

技術的サマリー：AMUSE（Anytime Muon with Stable Gradient Evaluation）

問題定義

現代の深層学習最適化は、通常、AdamW オプティマイザと指定された学習率スケジュール（例：コサイン減衰）の組み合わせに依存しています。最近の進展は、これらの構成要素を個別に挑戦してきました：Schedule-Free (SF) 最適化は、現在の反復点と平均化された反復点の間の補間点で勾配を評価することで、明示的なスケジュールを排除し、Muon は、モーメントを直交化することで行列値パラメータの更新幾何学を改善します。

Muon が大規模言語モデル（LLM）の事前学習において経験的に成功を収めているにもかかわらず、その背後にあるメカニズムは完全には理解されていません。本論文は、Muon における特定の不安定性を特定しています。すなわち、その直交化は損失地形の「平坦な」方向（「川」）に沿った進捗を加速しますが、意図せずして「急峻な」高曲率方向（「谷の壁」）におけるノイズを増幅させ、振動的な軌道を引き起こします。逆に、Schedule-Free 手法は軌道を川に近い位置に保つことで軌道を安定化させますが、Muon の直交化が提供する迅速なバルク進捗を欠く可能性があります。核心的な問題は、手動の学習率スケジュールを必要とすることなく、Muon の迅速な収束と Schedule-Free 平均化の安定性を組み合わせたオプティマイザを設計することです。

手法：AMUSE

著者は、時間変化する補間係数を用いて、Muon の直交化された行列更新を Schedule-Free フレームワークに統合する AMUSE（Anytime Muon with Stable gradient Evaluation）を提案します。

理論的基盤：川 - 谷の損失地形

分析は、川 - 谷の損失地形という視点に基づいています。

川：有用な学習進捗が発生する、高次元かつ低曲率のバルク部分空間。
谷の壁： 振動を引き起こす、低次元かつ高曲率の支配的部分空間。
Muon のメカニズム： モーメント行列を直交化することで、更新の相対的なバルク成分を増加させ、川に沿った進捗を加速します。しかし、支配的な成分の減衰を防ぐため、支配的な方向における小さなノイズが谷の壁を横断する振動に増幅されてしまいます。
Schedule-Free のメカニズム： 平均化された軌道（川に近い）で勾配を評価することで、更新ステップ前に高曲率の揺らぎをフィルタリングします。

アルゴリズム設計

AMUSE は、高速なベース系列 $Z_t$ （Muon 更新）と安定した平均化系列 $X_t$ の間の勾配評価点 $Y_t$ を決定するために、時間変化する補間係数 $\beta_t$ を採用します。
$Y_t = (1 - \beta_t)Z_t + \beta_t X_t$

$\beta_t$ のスケジュールは、初期適応と後期の安定性のバランスを取るよう設計されています。

ウォームアップ段階 ( $t \le T_0$ )： $\beta_t$ は小さな値 $\beta_1$ に固定されます。これにより、オプティマイザは高速な Muon 系列 ( $Z_t$ ) の近くで勾配を評価し、迅速な初期適応を可能にします。
ウォームアップ後 ( $t > T_0$ )： $\beta_t$ は以下に従って 1 に向かって徐々に増加します。
$\beta_t = 1 - \left( \frac{T_0 - 1}{t - 1} \right)^\rho (1 - \beta_1)$
ここで、 $\rho \in [0, 1]$ は移行の速度を制御します。 $\beta_t$ が増加するにつれて、勾配評価点は平均化系列 $X_t$ へと移行します。これにより、Muon の直交化を受ける前に勾配の高曲率（支配的）成分が抑制され、軌道が安定化します。

重要なのは、このスケジュールが総学習期間に依存しない点であり、これによりAnytime 学習（任意の時点で学習を停止しても性能が低下しない学習）が可能になります。AMUSE は学習率減衰スケジュールを必要としません。

主要な貢献

Muon の幾何学的分析： 本論文は、Muon の効率性と不安定性に対する幾何学的な説明を提供します。直交化がバルク指向の更新を促進し（川の進捗を加速）、支配的方向のノイズを増幅させる（谷の壁での振動を引き起こす）ことを実証しています。
補間による安定化： 著者は、平均化された軌道に近い位置（Schedule-Free 経由）で勾配を評価することで、直交化前に支配的成分を削減し、振動を軽減することを示しています。
AMUSE アルゴリズム： 高速な Muon 軌道から安定した平均化軌道へ勾配評価点を動的にシフトさせる新しいオプティマイザです。これにより、Muon の迅速な収束と Schedule-Free 手法の安定性が組み合わされます。
Anytime 学習のサポート： 総ステップ数に依存しない時間変化するスケジュールを使用し、学習率減衰を回避することで、AMUSE は柔軟な学習期間をサポートします。

実験結果

AMUSE は、画像分類（CIFAR-10/100、SVHN、ImageNet-1k）、画像セグメンテーション（ISIC 2018）、および LLM 事前学習（FineWeb 上の 124M、720M、1B パラメータの Llama 型モデル）を含む多様なベンチマークで評価されました。

性能 - 反回数パレートフロンティア： AMUSE は、すべてのタスクにおいて (Schedule-Free) AdamW および Muon を一貫して上回ります。
効率性の向上：
- 720M LLM 事前学習において、AMUSE は Muon の最終性能に 1.51 $\times$ 少ないステップで到達します。
- ResNet-50 を用いた ImageNet において、同じ性能を 1.12 $\times$ 少ないステップで達成します。
- ViT のファインチューニングを用いた ImageNet において、3.08 $\times$ 少ないステップで済みます。
安定性： 事後の指数重み平均（EWA）や学習率減衰から大幅に恩恵を受けるバニラ Muンとは異なり、AMUSE は学習全体を通じて安定した軌道を維持し、これらの事後手法を適用してもほとんど改善が見られません。
スケーラビリティ： AMUSE は、他の Schedule-Free 変種が時に苦労する大規模バッチサイズ（最大 2048）および大規模モデル規模（1B パラメータ）まで効果的にスケーリングします。

意義と主張

本論文は、AMUSE が学習ダイナミクスと損失地形の幾何学に対する体系的な理解に由来する、オプティマイザ設計の原理的なアプローチであると主張しています。迅速なバルク進捗と谷の壁の安定性というトレードオフを明示的に扱うことで、AMUSE は手動の学習率スケジュールを必要とせずに優れた性能を達成します。

著者は、AMUSE が既存のオプティマイザを上回る一方で、バニラ Muンと比較してわずかなメモリオーバーヘッド（ベース系列 $Z_t$ のための追加の状態コピー 1 つ）を伴うが、AdamW や Schedule-Free AdamW とは同等であると控えめに指摘しています。彼らは、将来の研究が、この追加のメモリコストなしに同様の安定性と性能を達成することに焦点を当てうることを示唆しています。論文は、「川 - 谷」という視点が、効率的かつ安定したオプティマイザを設計するための堅牢な枠組みを提供することを強調しています。