アリ・カイアム氏は6月12日(現地時間)、arXiv（アーカイヴ）にて、Mixture-of-Experts（MoE）モデルの推論時メモリ効率を大幅に高める新トレーニング手法「StickyMoE」に関する研究論文を発表した。この手法は、MoEモデルが連続するトークン間で頻繁にエキスパートを切り替えることで生じる、メモリ内のウェイトスワッピング問題の解消を目指す。StickyMoEは、エキスパート割り当ての一貫性を高める新たな損失関数を導入し、推論の効率化に貢献する。
Mixture-of-Experts（MoE）モデルは、大規模言語モデル（LLM）の効率的なスケーリングを可能にする技術として注目を集めている。個々の入力トークンに対し、多数のエキスパート（専門ネットワーク）の中から一部の疎なサブセットのみを活性化させることで、高いモデル容量と計算効率を両立させる。しかし、このアーキテクチャには課題がある。連続するトークンの処理中に活性化されるエキスパートが頻繁に変動するため、特にエッジデバイスやメモリ帯域が限られた環境では、低速ストレージと高速メモリ間で絶え間ないウェイトの入れ替え（スワッピング）が発生し、推論速度の低下を招いていた。
従来の対策として、システムレベルでのキャッシングヒューリスティクスや、モデルトレーニング後のルータのファインチューニングなどが試みられてきた。これらのアプローチは根本的な原因、すなわちモデルの事前トレーニング段階におけるエキスパートの非一貫性には対処できていなかった。この問題に対し、アリ・カイアム氏らはStickyMoEと呼ばれる微分可能なルーティング一貫性損失（routing consistency loss）を提案している。
StickyMoEの核心は、隣接するトークン間でエキスパートが急激に切り替わることに対し、明確なペナルティを課す点にある。これにより、MoEモデルのルータは、意味的に関連性の高いトークンのスパン全体で可能な限り同じエキスパート割り当てを維持するように学習を促される。この手法は既存のMoEアーキテクチャにいかなる変更も必要とせず、単一のハイパーパラメーター「lambda」を追加するだけで導入できる簡便性を持つ。また、後処理としてのファインチューニングとは異なり、最初のトレーニングステップからエキスパート表現とルーティング決定が相互に適応できる点が、その有効性を高めている。
小規模なMoE言語モデルを用いた実験では、StickyMoEがエキスパートの切り替え率を最大60%削減できることが実証された。この削減にもかかわらず、モデルのパープレキシティ（予測の不確実性）の劣化は4%未満に抑えられている。この結果は、モデル品質とルーティングの時間的局所性のトレードオフにおいて、事後ファインチューニングよりもStickyMoEが優れたバランスを提供することを示唆している。ルーティングの時間的局所性は、モデルのトレーニング段階で最も効率的に組み込まれるべき特性であると結論付けられている。
この研究成果は、特にリソース制約の厳しい環境下でのMoEモデル展開において重要な意味を持つ。本手法によって達成されるパープレキシティの僅かな劣化と推論効率向上とのトレードオフは、今後のLLM研究において詳細に評価され、より広範なMoEモデルやタスクへの適用可能性が探求されるだろう。
参考: arXiv cs.LG — 2026年7月13日 13:00 (JST)
原文ハイライト"Sticky Routing: Training MoE Models"