Dense2MoE: Pushing the Pareto Frontier of On-Device LLMs via Unified Pruning and Upcycling

技術概要：Dense2MoE

問題定義
リソース制約のあるエッジデバイス（自動運転、エッジ AI など）への大規模言語モデル（LLM）の展開は、「推論メモリ壁」によって深刻に阻害されている。ハードウェアの Roofline 理論によれば、密な LLM は自己回帰的デコーディング中にメモリに束縛されており、理論的な FLOPs を単純に削減するだけではレイテンシの改善には不十分である。既存の解決策は、根本的なトレードオフに直面している：

1. 構造化プルーニング：UIDL などの手法は、冗長な層を破棄することでレイテンシを低減するが、LLM の主要な知識リポジトリである Multi-Layer Perceptron（MLP）を除去することで、やむを得ずモデル精度を犠牲にする。
2. MoE アップサイクリング：密なモデルを Mixture-of-Experts（MoE）アーキテクチャに変換することで、動的計算を線形的に増加させることなく容量を増加させる。しかし、標準的なアップサイクリングは、同質なエキスパートの初期化（ネイティブ MLP の複製）に起因する「表現の崩壊」に悩まされ、専門化を誘発するために大規模な継続的事前学習を必要とする。さらに、最近のスパースな再利用手法は、層内冗長性に焦点を当てることが多く、冗長なアテンションおよび KV キャッシュの読み込みオーバーヘッドを引き起こすマクロレベルの層間構造的重複に対処できていない。

手法：Dense2MoE
著者は、Layer-Fusion Upcycling（LF-UC） と呼ばれる新規技術を通じて、構造化された層プルーニングと MoE アップサイクリングを統合する統一フレームワーク Dense2MoE を提案する。この手法は 3 つの段階で動作する：

1. 層間類似性分析：フレームワークは、デコーダー出力と MLP 入力双方における隣接層の余弦類似度を分析することで、構造的冗長性を特定する。重要なのは、冗長な MLP の再利用が活性化の異常やスケールシフトを導入しないように、厳格な特徴量ノルム一致制約を課すことである。
2. Layer-Fusion Upcycling（LF-UC）：冗長な層を完全に破棄するのではなく、Dense2MoE はこれらの層のアテンションモジュール（MHA）のみを物理的にプルーニングし、メモリ帯域幅を多く消費する操作を排除する。これらのプルーニングされた層からの事前学習済み MLP は保持され、保持されたベース層内の異質なエキスパートとして再利用される。
   • これにより、「ベースエキスパート」が複製されたネイティブ MLP であり、「補完的エキスパート」がプルーニングされた層から保持された MLP であるというエキスパートプールが作成される。
   • 動的なトークンレベルのルーターは、これらのエキスパートへトークンを選択的に配信するか、残差パスを介してバイパスする。
   • この設計は、同質なアップサイクリングで一般的に見られる「勝者総取り」のルーティング崩壊を緩和するために、層間構造的異質性を活用する。
3. 継続的事前学習：修正されたアーキテクチャは、ルーティング機構の最適化と融合エキスパートの整合化を行うため、軽量な継続的事前学習フェーズ（2250 億トークン）を経る。この計算量は、ゼロから MoE モデルを学習するために必要な計算量の 1% 未満である。

最適なプルーニング深度を決定するために、著者は MnasNet に着想を得たハードウェア対応の報酬メトリックを採用し、特定のターゲットハードウェア（例：NVIDIA Jetson Thor-U）に対するレイテンシ対精度のパレートフロンティアにプルーニング制約をマッピングする。

主要な貢献

• 統一フレームワーク：Dense2MoE は、構造化された層プルーニングと MoE アップサイクリングをシームレスに統合し、推論効率とモデル精度の間のトレードオフを解決する。
• 新規アーキテクチャ（LF-UC）：アテンションのボトルネックをプルーニングしつつ、MLP を異質なエキスパートとしてリサイクルすることで、表現能力を破壊することなくメモリ帯域幅の消費を削減する。
• 最先端のパレート最適性：このフレームワークは、デバイス内推論のレイテンシ対モデル精度のパレートフロンティアを前進させ、密なベースライン、最先端の圧縮手法、および標準的なアップサイクリング手法を上回る。
• 堅牢なスケーラビリティ：この手法は、さまざまなモデル規模（0.5B から 7B）および異なるアーキテクチャファミリー（Qwen2.5 および LLaMA2）において検証されている。

実験結果
Qwen2.5 および LLaMA2 モデルを用いて、MMLU、ARC、GSM8K、HumanEval などを含む 9 つの多様なベンチマークで評価された：

• 性能対レイテンシ：最適な Dense2MoE 構成（24 層の Qwen2.5-0.5B シードから 19 層を保持）は、シードモデルの 307.5 ms から 271.0 ms（Thor-U SoC 上）に物理的推論レイテンシを低減しつつ、平均ベンチマークスコア 47.31 を達成し、密なベースライン（37.53）を大幅に上回った。
• ベースラインとの比較：推論およびコーディングタスクで深刻な容量崩壊に悩まされた従来のプルーニング手法（UIDL、LLM-Streamline など）やスパース再利用アプローチ（ToMoE、LLaMA-MoE など）とは異なり、Dense2MoE は堅牢な表現保持を維持した。平均スコアにおいて密なベースラインを 9.78 ポイント以上上回りながら、0.42B のパラメータのみを活性化させた。
• スケーラビリティ：より大規模なモデル（Qwen2.5-1.5B および LLaMA2-7B）に適用した場合、Dense2MoE は同一の 2250 億トークンデータセットで微調整した後、それぞれ密な対応モデルを 3.04 ポイントおよび 3.54 ポイント上回った。

意義
本論文は、Dense2MoE が、オープンソースの密な LLM を高効率なデバイス内アーキテクチャへ再利用するための実用的かつ産業標準的なレシピを確立すると主張している。層間構造的冗長性を体系的に活用することで、推論レイテンシとモデル精度の同時最適化を可能にする。この研究は、ゼロから MoE モデルを学習する際の莫大なコストや、過激なプルーニングに伴う精度低下を伴わずに、推論メモリ壁を打破し、複雑な推論能力を保持することが可能であることを実証している。これは、レイテンシがクリティカルなエッジ展開に最適なモデルの柔軟な連続体を提供する。