Mixture-of-Experts（MoE）アーキテクチャは，パラメトリックルータにより各トークンを少数のエキスパートへ疎に割当て，計算コストを抑えたままモデル容量を拡張する．一方で，ルータのパラメータは学習後に凍結されることが多く，分布シフト下ではエキスパート割当が不安定になりやすい．本稿では、検索拡張型エキスパート割当手法であるkNN-MoEを提案する．本手法では，参照集合上で正解トークンの尤度を最大化するように推定した「最適エキスパート割当」をメモリに保存し，推論時にkNN検索を用いて割当を再利用する．近傍の類似度から混合係数を算出し，類似事例が乏しい場合にはルータへ自動的にフォールバックする．3つのMoEモデルで評価した結果，kNN-MoEはZero-shotを一貫して上回り，パラメータ更新なしにSFTに匹敵する性能を達成した．

大規模言語モデル（LLM）のスケーリングは，Transformer[1]とMixture-of-Experts（MoE）アーキテクチャ[2,3,4,5]の組み合わせにより進展してきた．MoEは疎な活性化により，計算コストを大きく増やさずにモデル容量を増やすことができる[5]．ルーティングネットワーク（以下ルータ）は，各トークンに対して少数のエキスパートを動的に有効化する[5,6]．しかし，MoEモデルの性能はこのルータの判断の質に強く依存する．標準的なルータは，隠れ状態からエキスパート割当を予測する軽量な分類器として学習されるが，学習後は凍結され，推論時に割当を調整できない．学習分布から外れた入力（例えば，パープレキシティが高い事例）では学習した判断規則を外挿するしかルータ入力類似度に基づく検索メモリルータ(過去の最適割当)信頼度信頼度に基づく適応混合予測割当検索割当最終的なエキスパート割当Expert4Expert3Expert2Expert1タの予測とメモリから検索した近傍の割当を取得する．「信頼度に基づく適応混合」は近傍類似度に基づく係数で両者を補間し，最終的なエキスパート割当を出力する．なく，非最適なエキスパートを選びうる[7,8,9]．この硬直性は，特に専門領域や分布外タスクでMoEモデルの潜在能力を制約する．本稿では，ルータ入力の𝑘近傍[10]を検索し，参照集合上で正解トークン尤度を最大化する割当を再利用してルータの判断を補正するkNN-MoEを提案する（図1）．具体的には，ルータ入力（キー）と最適エキスパート割当（バリュー）を予めメモリに保存し，推論時にkNN検索で近傍を取得する．さらに，近傍類似度に基づく信頼度係数で検索割当とルータの予測割当を補間し，類似度が低い場合はルータへフォールバックすることで検索ノイズを抑える．本アプローチは，パラメータ更新を伴わず，推論時にエキスパート割当を非パラメトリックに最適化する新たな枠組みを提示するものである．3種類のMoEモデル（OLMoE[11]，GPT-OSS[12]，Qwen3[13]）に対し，GPQA[14]，SuperGPQA[15]，MMLU[16]，USMLE[17]，MedMCQA[18]で評価しThisworkispublishedwithoutpeerreviewandislicensedbytheauthor(s)underCCBY4.0(https://creativecommons.org/licenses/by/4.0/).org/licenses/by/4.0/).める．ロジット𝑟∈ℝ𝑁から疎なエキスパート重みへの写像を𝜋(𝑟)=TopK(Softmax(𝑟))とする．トークン𝑡に対して，パラメトリックロジット𝑥𝑡𝑊𝑟の代わりに学習可能なロジットベクトル𝑟𝑡を導入し，負の対数尤度を最小化する：𝑟∗𝑡=argmin𝑟∈ℝ𝑁L𝑡(𝑟),(3)L𝑡(𝑟)=−log𝑝𝜃(𝑦𝑡|𝑥𝑡,routing=𝜋(𝑟)),(4)ここで𝜃はネットワークの残り部分の凍結パラメータを表す．式3は勾配降下法で解く．初期値として元のパラメトリックロジットを用い，𝑟(0)𝑡=𝑥𝑡𝑊𝑟とする．その後，𝑆ステップの更新を行う：𝑟(𝑠+1)𝑡=𝑟(𝑠)𝑡−𝜂∇𝑟L𝑡(𝑟(𝑠)𝑡),for𝑠=0,...,𝑆−1,(5)ここで𝜂は学習率である．この更新は全MoE層ℓ∈LMoEのルーティングロジット{𝑟(ℓ)𝑡}に対して同時に適用され，各系列ごとに独立して最適化される．最終的な最適エキスパート割当は𝑎∗(𝑥𝑡)=𝜋(𝑟(𝑆)𝑡)である．この値をメモリのバリューとして𝑣𝑡=𝑎∗(𝑥𝑡)と保存する．以上より，メモリは次のように構成される：M={(𝑥𝑡,𝑎∗(𝑥𝑡))|y∈Dref,𝑡∈[1,|y|]}.(6)この処理をLMoEの各層について繰り返す．4.2信頼度を考慮した適応混合推論時には，各MoE層でルータ入力𝑥に対し，4.1節で構築したメモリから𝐾近傍N(𝑥)⊂Mを検索する．検索近傍を類似度で重み付けして集約し，メモリに基づく割当𝑎mem(𝑥)を得る．検索された𝐾個のキー・バリューペアを{(𝑘𝑗,𝑣𝑗)}𝐾𝑗=1とすると，𝐾∑𝑠(𝑥,𝑘𝑗)∑𝐾𝑚=1𝑠(𝑥,𝑘𝑚)𝑣𝑗,(7)𝑎mem(𝑥)=𝑗=1ここで𝑠(·,·)は類似度関数である．次に，パラメトリックルータの予測割当𝑎(𝑥)（式1）と𝑎mem(𝑥)を融合するため，検索信頼度に基づく混合係数𝜆(𝑥)を導入する．本稿では，近傍の平均類似度を信頼度として用いる：𝐾∑𝜆(𝑥)=1𝑗=1𝑠(𝑥,𝑘𝑗).(8)𝐾最終的なエキスパート割当は，両者を線形補間して得る：𝑎ﬁnal(𝑥)=(1−𝜆(𝑥))𝑎(𝑥)+𝜆(𝑥)𝑎mem(𝑥).(9)最後に，𝑎ﬁnal(𝑥)を用いてMoE層の出力ℎを計算する：𝑁∑ℎ=𝑖=1𝑎ﬁnal(𝑥)𝑖·𝐸𝑖(𝑥).(10)これにより，類似度が高い場合は検索割当を重視してルータの予測割当を補正し，類似度が低い場合は予測割当へフォールバックする．5実験5.1実験設定データセット一般推論ベンチマーク（GPQA，SuperGPQA，MMLU）と医療ベンチマーク（USMLE，MedMCQA）でkNN-MoEを評価した．各ベンチマークで，互いに重ならないテスト集合Dtestと参照集合Drefを用意し，Drefは全手法で共有した．kNN-MoEはメモリ構築に，5-shot[23]はin-context例の検索プールに，SFTは学習データにDrefを用いた．各ベンチマークの分割とデータ規模は付録Aに示す．モデルOLMoE（OLMoE-1B-7B-0125-Instruct）[11]，GPT-OSS（gpt-oss-20b）[12]，Qwen3（Qwen3-30B-A3B-Instruct-2507）[13]の3種類のMoEモデルを用いた．ベースラインkNN-MoEを，計算コストの低い標準推論（Zero-shot，5-shot）と，SFTベースの手法（SFT，SFT（RouterOnly））と比較した．ここで，5-shotは参照集合Drefからテスト例の類似事例を5件検索してin-context例に含める手法，SFT（RouterOnly）はルータパラメータのみをファインチューニングする手法である．なお，各ベースラインの詳細設定は付録Bに記載する．評価方法全データセットに対し，公式解答を用いてDtest上の正解率（%）を報告する．5.2結果3種類のモデルにおける主要な結果を表1に示す．Zero-shotおよびSFT（RouterOnly）との比較全モデル・データセットで，kNN-MoEはZero-shotとSFT（RouterOnly）を一貫して上回った．この結果は，パラメータ更新なしでも検索割当がルータのみより有効であることを示す．5-shotとの比較kNN-MoEは5-shotより安定して性能を改善した．Qwen3では，5-shotとkNN-MoEの両方がZero-shotを上回った一方で，GPT-OSSとOLMoEでは5-shotがZero-shotを下回ることが多org/licenses/by/4.0/).モデル手法GPQAMMLUSuperGPQAUSMLEMedMCQAOLMoEGPT-OSSQwen3かった．これは，in-context例の連結により入力が大幅に長くなりモデルの性能が不安定になったこと，およびin-context例の検索ノイズが増えること（低類似事例の混入）が一因と考えられる．SFTとの比較kNN-MoEは一部の設定でSFTを上回ったが，相対性能は条件に依存した．参照集合が小さい設定（GPQAとUSMLE）では，SFTがZero-shotを下回る場合があった．これは，少量データによる過学習の可能性がある．一方，kNN-MoEはデータ不足条件でもZero-shotを安定して上回った．5.3考察:なぜ提案手法が有効か本稿では，検索はパラメトリックルータが不確かなときほど有効であり，その不確かさはパープレキシティ（PPL）の増加と相関する，という仮説を立てた．これを検証するため，元モデル（Zero-shot）のPPLでテスト例を高／中／低の3群に分け，Zero-shotに対するkNN-MoEの正解率差を群ごとに算出した．表2にOLMoEでの結果を示す．を，元モデル（Zero-shot）のPPLで3群に分けて示す．ベンチマーク高PPL中PPL低PPLGPQA+6.15+2.99-1.52SuperGPQA+0.17+0.65-0.06MMLU+2.07+1.15-0.13USMLE+6.32+1.41-1.13MedMCQA+6.04+1.13+0.56高PPL例で改善が最大となり（例：GPQAで+6.15，USMLEで+6.32），全ベンチマークで同傾向が見られた．中PPL例でも改善は小さいが一貫して向上しZero-shot27.2746.7713.0232.8135.575-shot21.7233.0611.0931.3125.36SFT21.7246.8913.6737.8434.78SFT（RouterOnly）24.2445.2711.6631.2234.23kNN-MoE（提案）29.8047.8113.2735.0437.01Zero-shot43.9470.2023.5267.2956.205-shot36.8763.7219.9560.0252.76SFT41.9270.2824.8365.6156.06SFT（RouterOnly）41.4169.1818.8365.0554.89kNN-MoE（提案）45.4570.2824.3568.3157.06Zero-shot41.4178.5934.8575.3062.515-shot44.4480.4835.5677.1766.58SFT39.9079.0840.2080.8066.60SFT（RouterOnly）43.9478.3633.1576.0566.03kNN-MoE（提案）44.9578.8635.1576.7066.65た．一方，低PPL例では改善がほぼゼロで，わずかに低下する場合もあった．これは，ルータ予測が十分なときに検索がまれにノイズを導入するためと考えられる．総合すると，高PPL入力ほどルータの割当予測が不確かになる一方で，kNN-MoEは類似事例の最適エキスパート割当でそれを補正できるため，PPLと改善幅が強く相関したと考える．6結論と今後の課題本稿では，ラベル付き参照集合からオフラインで推定した最適エキスパート割当をメモリ化し，推論時のkNN検索でルータの予測割当を補うエキスパート割当枠組みkNN-MoEを提案した．近傍類似度に基づく混合により，検索割当とルータの予測割当を補間してパラメータ更新なしにエキスパート割当を改善することが可能となった．3つのMoEモデルと5つのベンチマークにおいて，kNN-MoEはZero-shotおよびSFT（RouterOnly）を一貫して上回り，SFTに匹敵する性能を達成できることを示した．本手法の新規性は，静的ルータの予測割当を検索に基づく動的な割当で補間する点であり，MoEモデルの適応能力を拡張する新たな視点を提示する．今後の課題として，メモリやキー（ルータ入力）の圧縮による検索高速化が挙げられる．例えば次元削減や量子化により距離計算を軽量化しつつ，精度の維持を目指す．第二に，ラベルなし参照集合への拡張である．LLM-as-a-judge[24]で参照集合に擬似ラベルを付与し，同様の尤度最大化で最適エキスパート割当を構築することを検討する．org/licenses/by/4.0/).