混合エキスパート(MoE)モデルは、各トークンに対してエキスパートの疎なサブセットのみを活性化することで効率を向上させますが、エッジデバイス上では連続するトークンが異なるエキスパートを活性化するため、モデルの重みを低速ストレージ(フラッシュメモリなど)と高速メモリの間で頻繁に交換する必要があります。この重みの交換は推論遅延を増加させ、メモリボトルネックを引き起こし、スマートフォンやIoTデバイスなどのリソース制約のあるプラットフォームへのMoEモデルの展開を制限しています。
既存の対策は、システムレベルのキャッシュヒューリスティックと事後的なルーターファインチューニングの2つに大別されますが、いずれもトレーニング中のルーティングの不安定性という根本原因に対処していません。キャッシュ戦略は症状を緩和するだけであり、事後ファインチューニングは局所性を改善できるものの、モデル品質を損なう可能性があり、追加のトレーニングステップが必要です。
arXivで発表された新しい研究では、StickyMoEを提案しています。これは微分可能なルーティング一貫性損失関数であり、トレーニング中に隣接トークン間の急激なエキスパート切り替えをペナルティし、意味的に一貫したスパン内で同じエキスパート割り当てを維持するようにルーターを促します。この手法はアーキテクチャの変更を必要とせず、単一のハイパーパラメータλを追加するだけで済みます。また、事後手法とは異なり、エキスパート表現とルーティング決定が最初のトレーニングステップから共適応できるため、モデル品質を維持しながら時間的局所性を持つルーティング戦略を自然に学習できます。
小規模MoE言語モデルでの実験では、StickyMoEはエキスパート切り替え率を最大60%削減し、パープレキシティの低下は4%未満に抑えられました。品質-局所性フロンティアにおいて、StickyMoEは事後ファインチューニングをパレート改善で上回りました。この研究は、ルーティングの時間的局所性をトレーニング時に注入することが最も効率的であることを示しており、エッジデバイス向けMoEモデルの展開に新たな道を開くものです。エッジAIの急速な発展に伴い、StickyMoEはMoEモデルの実用化を促進する重要な技術となるでしょう。