技術的概要：LoRA-Mixer

問題定義

大規模言語モデル（LLM）をマルチタスクドメインに適応させるための最近の取り組みでは、低ランク適応（LoRA）とエキスパート混合（MoE）アーキテクチャを組み合わせることが多い。しかし、既存のアプローチはパラメータ効率とタスク特化において重大な限界を抱えている。現在の手法は通常、2 つのパターンに従う：(1) 全体の注意層またはフィードフォワードネットワーク（FFN）層をスイッチ型エキスパートに置換するか、(2) 出力をメインパスに融合させる並列エキスパート分岐を追加する。これらの設計は、しばしばパラメータ効率を損ない、すべてのエキスパートの共同学習を必要とする（事前学習済みで入手可能な LoRA モジュールの再利用を妨げる）、また、コアの注意メカニズムや状態遷移メカニズムを十分に活用できない浅い出力融合をもたらす。さらに、一般的な補助的ルーティング損失は均一な負荷分散を優先するため、入力に応じた特化を抑制し、入力の意味に関係なくエキスパートを均等に使用することを強制する。

手法：LoRA-Mixer

著者は、複数の事前学習済み LoRA モジュールを動的でプラグ可能なメモリセルとして扱うように設計されたモジュール型 MoE フレームワークであるLoRA-Mixerを提案する。

1. アーキテクチャ統合

以前の研究が全体の層を置換または並列化するのとは異なり、LoRA-Mixer は LoRA エキスパートを注意モジュール（および SSM の場合は状態空間モデル）の線形射影層（入力/出力線形層）に直接統合する。

• メカニズム: このフレームワークは、$E$ 個の低ランクエキスパートのセットを用いて、事前学習済みの射影行列 $W$ を強化する。各エキスパートは $\Delta W^{(e)} = A^{(e)}B^{(e)}$ としてパラメータ化される。
• ルーティング: 軽量なルーター $\alpha(x)$ が、入力意味に基づいてこれらのエキスパートを動的に融合する。出力は以下のように計算される：
  $$y = Wx + \text{F}_{\text{route}}\left( \left\{ \alpha_e(x) \cdot \Delta W^{(e)}x \right\}_{e=1}^E \right)$$
• 互換性: 線形射影層は Transformer と状態空間モデル（SSM）の両方のアーキテクチャにおいて普遍的であるため、LoRA-Mixer はアーキテクチャに依存せず、両方をサポートする。

2. ルーティング特化バランス損失（RSL）

従来の補助的損失が引き起こす「過度の平均化」の問題に対処するため、著者は**ルーティング特化バランス損失（RSL）**を導入する。

• 動機: 標準的な補助的損失（例：負荷分散項）は均一なエキスパート使用を促すため、入力に応じた特化の必要性と矛盾する。
• 定式化: RSL は、補助的損失にエントロピー正則化項を追加することで、グローバルな負荷分散とローカルな選択性を統合する：
  $$L_{RSL} = \alpha \cdot \sum_{i=1}^K \bar{p}_i \cdot \bar{f}_i - \lambda \cdot \mathbb{E}_{x \sim D}[H(p(x))]$$
  ここで、$\bar{p}_i$ は平均ルーティングスコア、$\bar{f}_i$ は正規化された使用スコア、$H(p(x))$ はルーティング分布のエントロピーである。
• 効果: エントロピー項はルーティング単体における曲率提供者として機能し、強い凸性を促進して、ルーターが均一分布に収束するのを防ぐ。これは、グローバルな負荷分散を維持しつつ、入力意味と整合する高分散で鋭い分布を促進する。

3. 学習レジーム

LoRA-Mixer は、2 つの明確な学習モードをサポートする：

1. 共同最適化: 異なるハード - ソフト top-k ルーティング方式を用いて、アダプターとルーターを共同最適化する。
2. プラグアンドプレイ: このフレームワークは、パブリックリポジトリ（例：LoRAHub）から供給される凍結済みの事前学習済み LoRA モジュールに対して、ルーターのみを学習するための最小限の追加データ（例：2k の混合データポイント）だけでルーティングを行うことができる。

主要な貢献

1. モジュール型 MoE フレームワーク: タスク固有の LoRA エキスパートを注意/SSM モジュールのコア射影行列にルーティングする新規設計により、基盤アーキテクチャを乱すことなく微細なトークンレベルの特化を可能にする。
2. ルーティング特化バランス損失（RSL）: 負荷分散と入力応答型特化の間の矛盾を解決する最適化目的関数。これは効果的なルーティング学習に必要なデータを削減し、従来の補助的損失の「過度の平均化」バイアスを防止する。
3. データおよびパラメータ効率: この手法は、既存の LoRA-MoE ベースラインの学習可能パラメータのわずか**48%**のみを使用して最先端のパフォーマンスを達成する。標準的な補助的損失を使用する手法と比較して、ルーターの学習に必要な学習データが大幅に少ない高いデータ効率を示す。
4. クロスモデル転送性: このフレームワークは、異なるモデル間での LoRA モジュールの再利用（例：Mistral-7B で学習したルーターを LLaMA3-8B に転送）をサポートし、多様なソースからの異種 LoRA モジュールの構成を可能にする。

実験結果

著者は、Medical QA、常識推論、NLP、数学、コーディングの 5 つのドメインにまたがる15 のベンチマークで LoRA-Mixer を評価した。

• パフォーマンス: RSL 最適化された LoRA-Mixer は、GSM8K で**+3.79%、CoLA で+2.90%、ARC-C で+3.95%**という顕著な改善を示し、MoLE、MixLoRA、LoRAHub などの最先端のルーティングおよび LoRA-MoE ベースラインの多くで、ほとんどのタスクにおいて上回った。
  • 効率性: これらの結果は、同等の手法の学習可能パラメータの 48% しか使用せずに達成された。
• データ効率: 低データシナリオ（例：2k の学習サンプル）において、RSL 最適化された LoRA-Mixer は、標準的な補助的損失を使用するベースラインを大幅に上回った（2k データで平均 +1.97% の改善）。
• プラグアンドプレイ機能: インターネットから入手した LoRA（LoRAHub）に適用し、ルーター学習のために 2k の追加データポイントのみを使用した場合、GLUE タスク（SST-2、CoLA、MRPC、RTE、QQP）で優れたパフォーマンスを達成した。
• 汎化能力: この手法は、ドメイン間タスク（例：数学 - コーディング、医療 - 数学）および分布外（OOD）ベンチマークにおいて強力な汎化性能を示した。
• 転送性: ルーターの微調整なしに Mistral-7B から LLaMA3-8B へルーターを転送する実験は、学習されたルーティングメカニズムの堅牢性とアーキテクチャ非依存性を確認した。

意義と主張

本論文は、LoRA-Mixer が以下の点でパラメータ効率的なマルチタスク適応において重要な前進を代表すると主張している：

• 再利用の最大化: 広範な再適応やすべてのエキスパートの共同学習を必要とせず、独立して学習された LoRA モジュールのシームレスな構成を可能にする。
• 特化と負荷のバランス: 歴史的に低データ領域における MoE パフォーマンスを阻害してきた負荷分散とエキスパート特化のトレードオフに対する、理論的根拠のある解決策（RSL）を提供する。
• 汎用性: Transformer と SSM 両方のアーキテクチャに対応する「ドロップイン」互換ソリューションを提供し、タスク固有のメモリと強力な転送性を持つ大規模モジュール型言語モデルの構築を容易にする。

著者は、現在の固定された top-K ルーティングが曖昧な入力への適応性を制限する可能性があるものの、このフレームワークは、最小限のデータと計算オーバーヘッドで LoRA エキスパートの効率的なモジュール構成が実現可能であることを成功裏に実証していると結論づけている。