解決策
Mixture-of-Expertsモデルは高速な推論を実現するが、ゼロから学習する場合や密なLLMをアップサイクルする場合には、膨大な計算資源とコーパスが必要となる。DeepSeek-V3とQwen2.5-Maxはそれぞれ14.8Tトークンと20Tトークンで事前学習された。これに対しL0-MoEは、事前学習済みの密なチェックポイントと厳選された30Bトークンのコーパスから、実用的なMoEを構築する。
その仕組みは次のとおりである。BGE-M3埋め込みとK-meansを用いてコーパスを意味ドメインにクラスタリングし(クラスタ混同行列によって意味的距離の順に並べる)、次に各ドメインについて、FFNの中間次元上のバイナリマスクをL0正則化を用いて学習する。この際、保持率を100%から目標値までアニーリングし、MLP以外のすべてを凍結する。これにより64のエキスパートが得られ、それぞれが元のモデルの一部を切り出した、1つのドメインに特化したものとなる。
次に動的バッチ処理によって、2つのループでルーターを学習させる。まず同一ドメインのバッチを用いてエキスパート選択を初期化し、その後、ドメインを周期的に混在させたバッチを用いることで、トークンレベルのマルチエキスパートルーティングを学習させる。
Llama-3-8B、Mistral-7B、Qwen2-7Bにおいて、L0-MoEはそれぞれ2.0倍、2.1倍、2.5倍の推論高速化を達成し、ベンチマーク性能は元の密なモデルに匹敵する。Mistralの変種では平均で約1%向上することさえある。これはGPTQ量子化、LLM shearing、RKD+CoT蒸留の各ベースラインを上回る。これらの手法は高速化するものの品質を劣化させる。アブレーション実験では、ランダム、マグニチュード、OBS、SVDによる次元選択がいずれもL0で学習したマスクに劣ることが示されている。
達成したこと
エキスパートをゼロから育てるのではなく、ドメインごとに密なモデル自身の隠れ次元のうちどれが重要かを選び、それだけを保持する
うまくいった理由
L0のハードコンクリートゲートにより次元選択が微分可能になるため、「このドメインにはどのニューロンか」を探索するのではなく学習できる。
FFNマスク以外のすべてを凍結することで、各エキスパートは元のモデルの真の一部であり続けるため、品質が維持される。
クラスタ混同行列によってドメインが順序付けられるため、ルーターは混在ドメインより先に馴染みのあるドメインに出会う。
ドメインを一致させたバッチにより、混在バッチでトークンレベルの選択を学習する前に、シーケンスレベルのルーティングをブートストラップする。
応用できること
密なモデルはすでにそのエキスパートを含んでいる。鍵となるのは、どのニューロンがどのドメインに属するかを見つけることである。ドメインごとのL0マスクとドメインカリキュラムにより、MoE学習コストのほんの一部でそれらを切り出せる。
その後
第63回ACL年次総会(2025年)の議事録に掲載された。コードと厳選データセットはgithub.com/zhangzhenyu13/L0-MOEで公開されている。著者らは、ゼロからのMoE学習が現実的でないコストに敏感な産業展開向けの手法として、このアプローチを位置づけている。
出典をたどる
出典
- Accelerating Dense LLMs via L0-regularized Mixture-of-Experts arXiv AI · arxiv.org ·