🤖 machine learning

Sticky Routing: Training MoE Models for Memory-Efficient Inference

本論文は、意味的に一貫したトークンスパン全体でエキスパートの割り当てを維持するようにMixture-of-Expertsモデルを学習させる微分可能なルーティング一貫性損失であるStickyMoEを提案しており、これにより、パープレキシティの低下を最小限に抑えつつ、エッジデバイスにおけるメモリ集約的な重みのスワッピングを大幅に削減する。

原著者： Ali Kayyam

公開日 2026-07-13

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

CC BY 4.0

原著者： Ali Kayyam

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

想像してみてください。あなたには64人の異なる専門家である司書たちがいる巨大な図書室がありますが、あなたの小さなエッジデバイス（スマートフォンやスマートスピーカーなど）には、一度にデスクの上に置いておけるのはわずか2人の司書だけです。残りの司書は、遅くて埃っぽい地下室に保管されています。

標準的な「Mixture-of-Experts（混合エキスパート）」モデルでは、AIが新しい単語を読むたびに、どの司書に頼むかを決めるためにコインを投げます。問題は、そのコイン投げを単語ごとに行っていることです。つまり、「apple」については司書Aに聞き、「pie」については司書Bに聞き、「crust」については司書Cに聞く、という具合です。切り替えるたびに、AIは地下室まで走り、司書Bを追い出し、司書Cを連れて上がってこなければなりません。この「あちこちへの往復」があまりに遅いため、AIは物語を読むことよりも、司書を呼び出すことに多くの時間を費やしてしまいます。

この論文は、モデルがこれほど激しくコイン投げをしないようにするための新しい方法であるStickyMoEを紹介しています。モデルに「粘着性（sticky）」を持たせるように学習させるのです。もしモデルが「apple」に対して司書Aに頼んだなら、そのトピックが意味を成す限り、「pie」や「crust」についても司書Aに頼み続けるよう促されます。

大きな発見：正しく訓練せよ、後から直そうとするな

著者たちは、この「あちこちへの往復」問題を解決するためにいくつかの方法を試みました。そして以下のことを発見しました。

「事後的（Post-Hoc）」な修正は失敗した： 彼らは、すでに学習済みのモデルを取り出し、コイン投げの仕組み（ルーター）を、よりランダムではなくなるように「微調整（ファインチューニング）」することを試みました。彼らはこれをReMoEと呼びました。結果は、うまくいきませんでした。 切り替え率はほとんど変化しませんでした（絶対変化量は0.5%未満）。著者らは、司書たち（エキスパート）がすでに特定の孤立した単語のスペシャリストとして訓練されてしまっているからだと主張しています。もし司書自身の脳が長い連続したトピックを扱う準備ができていないのであれば、コイン投げ器に対して「もっと粘着性を持て」と指示しても意味がありません。それは、短距離走者に、スタートの構え方だけを変えることでマラソンランナーになろうと教えるようなものです。筋肉（エキスパートの重み）は、長距離走のためではなく、短距離の爆発力のために作られているのです。
「ハード（強制的）」な修正は硬直的すぎた： 彼らはまた、一定の単語数にわたって一つの司書に固執することを厳格なルールで強制する手法（「ハード・ウィンドウ」）も試しました。これにより切り替えは減少しましたが、AIの回答の質が大幅に低下しました（中規模モデルにおいてパープレキシティが最大**6.8%**増加）。それは、たとえジャンルが完全に変わったとしても、司書に小説を一冊丸ごと読ませるよう強制するようなものでした。その結果、物語の質が損なわれてしまったのです。
「スティッキー（粘着性のある）」解決策が成功した： 勝者となった手法であるStickyMoEは、トレーニングの極めて初期の段階で、緩やかな「一貫性ペナルティ（consistency penalty）」を加えます。これはモデルに対し、「もし隣り合う二つの単語の間でエキスパートを切り替えたら、それは少しペナルティになるよ」と伝えるものです。これにより、モデルは言語を学びながら、同時に「粘着性」も学ぶようになります。

結果：粘着性がある方が良い

彼らがWikiText-2データセット（Wikipediaの記事のコレクション）を用いて小規模および中規模のモデルでテストしたところ、結果は明白でした。

切り替えの減少： たった一つのつまみ（λと呼ばれる数値）を調整することで、エキスパートを入れ替える回数を最大59%削減できました。例えば、中規模モデルでは、切り替え率は0.71（ほぼ毎回切り替わっている状態）から0.29へと減少しました。
優れた速度（シミュレーション）： モデルの切り替えが減ったため、2人のエキスパートを保持できるキャッシュ（デスクの棚）をシミュレートしました。その結果、「キャッシュヒット率」（適切な司書がすでにデスク上にいる確率）は0.54から0.88へと跳ね上がりました。これは、モデルが以前よりも地下室へ走りに行く回数が3.92回少なくて済むことを意味します。
より高い品質： 驚くべきことに、中規模モデルにおいて、適度な「粘着性」（λ = 0.05）を使用すると、AIはむしろ賢くなり、パープレキシティ（混乱の指標）が**4.1%**低下しました。モデルは、トピックに固執することが文脈の理解に役立つことを学んだのです。
崩壊（Collapse）なし： モデルが「怠けて」しまい、一つの司書を何にでも使ってしまうのではないかという大きな懸念がありました。しかし、「利用エントロピー（Utilisation Entropy）」は高い水準（最大値2.0のうち1.92ビット以上）を維持しており、4つのエキスパートが依然として公平に使われていることを示しました。モデルは怠けておらず、ただ整理整頓されただけなのです。

何がうまくいかなかったのか（とその理由）

この論文は、モデルが構築された後にこの問題を修正できるという考えを明確に否定しています。ReMoEのアプローチ（既存のモデルを微調整する手法）は、エキスパートの表現（司書の脳）が、トピックに固執することを考慮しないプロセスによってすでに形成されてしまっているため、効果的ではないことが示されました。局所性を後付けすることはできません。最初から組み込んでおく必要があるのです。

また、彼らはモデルの最初の層（レイヤー0）が、最も「粘着性」を持たせることが難しいことも指摘しています。最適な設定を用いても、そこでの切り替え率は0.49を上回り続けました。これは、最初の層が文脈と混ざり合う前の「生の単語」を見ているため、「apple」と「pie」がこの段階では純粋に全く別物に見える可能性があるからです。

結論

この論文は、デバイスのメモリが限られている環境でAIを高速に動作させたいのであれば、単に巧妙なキャッシュルールを考案したり、後からモデルを修正しようとしたりすべきではないと示唆しています。代わりに、最初の日からモデルに「粘着性」を持たせるように訓練すべきです。隣接する単語間でエキスパートを頻繁に切り替えることに対して単純なペナルティを加えることで、自然とトピックに集中し、遅い地下室への往復を減らし、実際に物語をより良く理解できるモデルが得られます。

著者らは、これがシミュレーションに基づいていること、および小・中規模モデル（最大22Mパラメータ）を用いたWikiText-2データセットでの訓練であることを認めています。彼らは、モデルが大きくなるにつれて、この「粘着性」はさらに効果的になる可能性があると考えていますが、数十億パラメータを持つ大規模なモデルについてはまだテストを行っていません。また、彼らの手法は現在のところ文章の終わりを認識できないため、段落の区切りを越えて「粘着性が強すぎる」状態になる可能性があり、これは今後の課題であるとも述べています。

技術要約：メモリ効率の高いMoE推論のためのスティッキー・ルーティング（Sticky Routing）

1. 問題提起：エキスパート・スワッピングのボトルネック

Mixture-of-Experts (MoE) モデルは、トークンごとに活性化されるエキスパートのサブセットを限定することで、大規模言語モデル（LLM）をスケールアップさせるよう設計されています。理論的には、これにより膨大なパラメータ数を持ちながら、1トークンあたりの計算量を低く抑えることができ、高速メモリ（VRAM/SRAM）に一度に少数のエキスパートのみを配置すればよいエッジ展開において非常に魅力的です。

しかし、本論文は実用上の決定的な失敗モードを特定しています。それは**時間的なルーティングの不整合（temporal routing inconsistency）**です。標準的なMoEルーターは、各トークンの隠れ状態のみに基づいて独立した決定を下します。その結果、連続するトークンが全く異なるエキスパートを活性化してしまうことが頻繁に起こります。メモリ制約のあるハードウェアでは、各切り替えごとにキャッシュされたエキスパートの重みを破棄し、新しい重気を低速なストレージ（CPU RAM/Flash）からロードする必要があります。PCIeやNVMeの帯域幅はGPUメモリの帯域幅よりも数桁低いため、これらの頻繁な「エキスパート・スワップ」が推論レイテンシを支配し、結果としてMoEモデルを同規模のデンス（dense）モデルよりも遅くしてしまうことがよくあります。

既存の解決策は不十分であるとみなされています：

システムレベルのアプローチ（例：キャッシング・ヒューリスティック、プリフェッチャー）：これらは存在する局所性を利用するものの、事後的（reactive）なものであり、本質的に不整合なルーター自体を修正することはできません。
ポストホック（事後）アプローチ（例：ルーターのみのファインチューニング）：エキスパートの表現が、局所性を促すような学習プロセスを経て形成されていないため、これらは失敗します。ルーターは、設計段階で意図されていない表現空間において、決定を再調整することが容易ではありません。
アーキテクチャの再設計（例：Oracle-MoE）：これらはモデルや学習パイプラインに大幅な構造変更を必要とします。

2. 手法：StickyMoE

著者らは、ルーティングの時間的局所性を明示的な目的関数として扱う学習時アプローチであるStickyMoEを提案しています。その中核となるイノベーションは、隣接するトークン間でのエキスパート割り当ての急激な変化を罰する、**微分可能なルーティング一貫性損失（differentiable routing consistency loss, $L_{cons}$ ）**です。

ルーティング一貫性損失

この損失関数は、連続するトークンのソフトマックス・ゲート確率ベクトル（ $g_t$ ）の $\ell_2$ 距離を最小化します：
$L_{cons} = \frac{1}{T-1} \sum_{t=2}^{T} \| g_t - g_{t-1} \|_2^2$
ここで、 $g_t \in \mathbb{R}^N$ はトークン $t$ に対するルーターの出力です。

この損失は、標準的なクロスエントロピー損失（ $L_{CE}$ ）およびロードバランシング損失（ $L_{bal}$ ）とともに、総学習目的関数に統合されます：
$L = L_{CE} + \lambda L_{cons} + \mu L_{bal}$

主要なメカニズム

共適応（Co-adaptation）：ポストホックな手法とは異 달리、StickyMoEはルーターとエキスパート表現を同時に学習します。 $L_{cons}$ を最小化することで、ルーターの重み（ $W_r$ ）を更新し、隠れ状態（ $h_t$ ）を通じてバックプロパゲーションを行い、意味的に一貫したスパンにおいて自然に一貫性を持つ表現を学習するように促します。
ソフト・ハード・バリアント（Soft-Hard Variant）：「長距離ドリフト（小さなステップごとの変化が蓄積して遠くのスイッチにつながること）」に対処するため、著者らはアンカー制約を導入しています。シーケンスはウィンドウに分割され、ウィンドウの最初のトークンがアンカーを設定し、後続のトークンはアンカーからの逸脱に対してペナルティを受けます。このペナルティは、ウィンドウの開始点からの距離に応じて線形に増加します。
互換性：この手法はアーキテクチャの変更を必要とせず、計算オーバーヘッドも無視できる程度です。標準的なtop- $k$ MoEおよび既存のロードバランシング技術と互換性があります。

3. 主な貢献

ルーティング一貫性損失：隣接するトークン間のエキスパートの切り替えを直接罰する、単純で微分可能な学習目的関数であり、単一のハイパーパラメータ（ $\lambda$ ）のみを必要とします。
ロードバランシングとの相補性：実験的証拠により、一貫性損失とロードバランシング損失は相補的であることが示されています。一貫性損失は、エキスパートの崩壊（利用エントロピーが理論的最大値付近の高い状態を維持すること）を引き起こすことなく、エキスパートが「いつ」使用されるかを形成します。
ソフト・ハード・バリアント：ステップごとのソフトなペナルティと、セグメントレベルのアンカー制約を組み合わせることで、追加の品質コストなしに強力な局所性を保証します。
包括的な評価：WikiText-2を用い、小規模（~~8.8M params）および中規模（~~22M params）のMoEモデルに対して、Vanilla MoE、ポストホック・ルーター・ファインチューニング（ReMoE）、およびアーキテクチャ・ベースライン（Oracle-MoE）と比較する制御実験を実施しました。

4. 実験結果

実験は、WikiText-2上で小規模（~~8.8M params）および中規模（~~22M params）のMoEモデルを用いて行われました。

エキスパート・スイッチ率（SR）：StickyMoEは、エキスパートのスイッチ率を最大**59%**削減しました（例：中規模モデルにおいて0.71から0.29へ減少）。
キャッシュヒット率（CHR）：スイッチングの減少は、2エキスパート・スロットのキャッシュを用いた場合、中規模モデルにおいて3.92倍のキャッシュミス減少（CHRが~~0.54から~~0.88へ上昇）に直結しました。
パープレキシティ（品質）：
- 適度な $\lambda$ の値（例：0.05）において、StickyMoEはベースラインと比較してパープレキシティを改善しました（例：中規模モデルで-4.1%）。これは、一貫性損失が有益な正則化として機能していることを示唆しています。
- 過度な $\lambda$ の値はトレードオフを生じさせますが、StickyMoEのパレート最適解は他の手法を厳密に上回っています。
ポストホック・ファインチューニングの失敗：ReMoEベースライン（学習済みモデルに対してルーターのみをファインチューニングする手法）は、スイッチ率を意味のあるレベルで減少させることに失敗しました（変化は<0.5%）。これは、表現が固定された後に局所性を後付けすることができないことを裏付けています。
層別分析：一貫性損失は、深い層（L1–L3）で最も効果的でした。入力層（L0）は局所性の向上が困難であり、これはトークン埋め込みにコンテキストの混合が欠けているためと考えられます。

5. 意義と主張

本論文は、ルーティングの時間的局所性は、事後的なシステムレベルの検討事項としてではなく、事前学習中に植え付けるべき正当な学習目的であると主張しています。

パラダイムシフト：量子化を考慮した学習（Quantization-aware training）が標準となったのと同様に、著者らは「局所性を考慮した学習（Locality-aware training）」が、メモリ制約のある展開を目指すMoEモデルのデフォルトになるべきだと提唱しています。
根本原因の解決：StickyMoEは、症状（キャッシュミス）ではなく、根本的な原因（一貫性のないルーティング決定）に対処しています。表現とルーティング決定を共同で学習することで、ポストホックな手法を悩ませる表現のミスマッチを回避しています。
実用性：この手法はアーキテクチャに依存せず、推論パイプラインに変更を必要としません。学習されたルーターは自然に一貫したゲート分布を生成するため、標準的なキャッシング戦略（LRU、LFUなど）が追加のエンジニアリングなしに大幅に効率的に機能するようになります。

著者らは、システムレベルの最適化や事後の調整も有用な補完要素ではあるものの、最初からキャッシュフレンドリーになるように訓練されていないモデルの代わりにはなり得ないと結論付けています。StickyMoEは、エッジデバイスにおいてMoEアーキテクチャの理論的なメモリ効率を実現するための道筋を示しています。