技術要約: PreLort – ランクの不均一性下における連邦微調整のためのプレフィックス・ネスト型LoRA

1. 問題提起

連邦学習（Federated Learning, FL）は、生のデータを共有することなく大規模言語モデル（LLM）の協調的なトレーニングを可能にしますが、低ランク適応（LoRA）のようなパラメータ効率の良い微調整（PEFT）手法を用いて基盤モデルを適応させる際、重大な課題に直面します。主要な障害は**ランクの不均一性（rank heterogeneity）**です。現実的なクロスデバイス設定では、クライアントはそれぞれ異なる計算およびメモリの制約を持っており、それによって異なるランク（$r$）を持つLoRAアダプタを初期化することになります。

既存の不均一な連邦LoRAの解決策は、主に集約フェーズに焦点を当てており、異なるランクのアダプタを平均化する前にゼロパディング（zero-padding）を用いる手法が一般的です。しかし、著者らはこれらのアプローチにおける根本的な限界を指摘しています。それは**表現の不整合（representation misalignment）**です。集約が数学的に定義されていたとしても、異なるランクで独立して訓練されたモデルを組み合わせることは、不整合をもたらします。具体的には、低ランクのクライアントは高ランクのクライアントからの更新の恩恵を受けることができません。なぜなら、情報はランク次元間で階層的に分布されていないからです。さらに、ゼロパディングは、ゼロ値のセグメントを導入することで高ランククライアントの貢献を希釈し、共有された低ランク表現の最適な利用を妨げます。

2. 手法: PreLort

著者らは、訓練のダイナミクスと集約ルールを統合的に構造化することで、ランクの不均一性に対処するフレームワークであるPreLort（Prefix-Nested LoRA）を提案しています。その核心となる哲学は、アダプタの次元を**プレフィックス階層（prefix hierarchy）**として構成することです。これにより、低ランクの次元には全クライアントで共有されるタスク関連の情報をエンコードさせ、高ランクの次元には増分的な容量を担わせるようにします。

2.1 ネストされた低ランク構造

PreLortは、アダプタの次元に対してプレフィックス構造化された構成を課します。ランク $r_k$ を持つクライアントの場合、最初の $r_k$ 次次元のみがアクティブになります。重要な点は、ローカル訓練中、アダプタは単にそのフルランクにおいてのみ最適化されるのではなく、**複数の切り詰められた構成（プレフィックス・サブランク）**の下で最適化されることです。

• メカニズム: ランク $r_k$ を持つクライアントは、$r_k$ だけでなく、連邦内に存在するすべての低ランクレベル（$r \in \{r \mid r \le r_k\}$）においても良好に機能するように訓練されます。
• 目的: これにより、モデルがタスクに関連する信号を先頭の（低ランクの）次元に集中させるよう強制し、これらの次元が切り詰められた場合でも情報量があり有用であり続けるようにします。

2.2 ローカル訓練戦略

クライアント $k$ のローカル目的関数は、すべての有効なプレフィックス・サブランク $R_k$ について計算されたクロスエントロピー損失の平均として定義されます：
$$\mathcal{L}_k = \frac{1}{|R_k|} \sum_{r \in R_k} \mathcal{L}_{CE}(f(\cdot; \theta, r), \mathcal{D}_k)$$
実際には、各訓練ステップにおいて、フルランクでのフォワード/バックワードパスに加え、すべてのサブランクに対する追加のパスが行われます。この戦略により、プレフィックスセグメントが集約で使用されるものと同じ粒度で明示的に最適化されるようになり、低次元が最も重要な情報を捉えるという一貫した階層構造が促進されます。

2.3 セグメント単位の集約

ゼロパディングによる希釈を避けるため、PreLortは勾配（更新）空間で動作するセグメント単位の集約ルールを導入しています。

• プロセス: すべてのクライアントについてすべての次元を平均化するのではなく、サーバーは、特定のランクセグメント $[r_{l-1}, r_l)$ に対して、そのセグメントに寄与するのに十分なランク $r_k$ を持つクライアントのサブセット（すなわち、$r_k \ge r_l$）に対してのみ勾配を集約します。
• 数式: 第 $l$ セグメントについて、集約された勾配は以下の通りです：
  $$\nabla A_{global}^{r_{l-1}:r_l} = \frac{\sum_{k: r_k \ge r_l} n_k \nabla A_k^{r_{l-1}:r_l}}{\sum_{k: r_k \ge r_l} n_k}$$
• 利点: これにより、低ランクのクライアントによるゼロパディングからの信号希釈を防ぎ、高ランクの次元が実際にそれらを学習する能力を持つクライアントによってのみ更新されるようにします。一方で、低次元は全集団から恩恵を受けることができます。

3. 主な貢献

本論文は、主に3つの貢献を挙げています：

1. 表現の不整合の特定: 著者らは、ランク次元間における構造化された整合性の欠如が、現在の不均一な連邦LoRA手法における主要な制限であり、それが一貫性のない表現につながっていることを突き止めました。
2. ネストされた定式化と統一された戦略: 共通のプレフィックス構造を課すネストされた低ランク定式化を提案しています。これは、ゼロパディングなしで不均一なランク間での一貫した集約を可能にする、統一された訓練および集約戦略と組み合わされています。
3. 実証的検証: 本手法は、指示チューニング（instruction-tuning）のベンチマークにおいて、従来の不均一な連邦LoRA手法よりも一貫して高い品質と安定性を示すことが実証されました。

4. 実験結果

著者らは、2つのベースモデル（TinyLlama-1.1B および Qwen2.5-0.5B）を用い、3つのデータセット（Alpaca, Databricks-dolly-15, 20 Newsgroups）においてPreLortを評価しました。

• 性能指標: PreLortは、精度（MMLU, 20 Newsgroups）およびROUGE-Lスコアにおいて、一貫してベースライン（ZeroPad, HetLoRA, FLoRA）を上回りました。
  • TinyLlama-1.1Bにおいて、PreLortはMMLU (Dolly) で33.86%、Alpacaで30.76%の精度を達成し、次に優れたベースラインを上回りました。
  • Qwen2.5-0.5Bにおいて、MMLU (Dolly) で39.85%の精度を達成しました。
• パープレキシティ: PreLortは、両方のベースモデルにおいてDollyおよびNewsGroupで最も低いパープレキシティを達成し、優れたトークンレベルの尤度を示しました。
• ランクの堅牢性: 分析により、PreLortは異なるプレフィックス切り詰めレベルにおいても安定した性能を維持することが示されました。対照的に、「Nested Aggregation Only」のようなベースラインは、特に高ランクにおいて不安定さを示しており、集約だけではネストされた訓練目的関数なしには不十分であることを裏付けています。
• 次元ごとの重要性: 次元重要性の可視化により、PreLortはタスク関連の信号を先頭の次元に集中させている（構造化された分布）のに対し、標準的な訓練では一様な分布になることが明らかになりました。この階層的な優先順位付けが、不均一なランク間での安定した更新を容易にします。

5. 重要性と主張

本論文は、PreLortがランクの不均一性への対処におけるパラダイムシフトを提示していると主張しています。すなわち、ランクの不均一性を事後的な集約問題として扱うのではなく、訓練プロセス自体を構造化することによって解決するというものです。

• 核心的な議論: 著者らは、ランクの不均一性に対処するには、改善された集約だけでなく、ランク間の表現の構造化された整合が必要であると主張しています。
• 効率性: この手法により、低ランクのクライアントは、プレフィックス次元が一貫して学習・集約されるため、高ランクのクライアントが提供するより豊かな情報の恩恵を受けることができます。
• 限界に関する謙虚な姿勢: 著者らは、ネストされた訓練戦略により、バッチあたりの複数のフォワード/バックワードパスが生じるため、約40%の追加計算オーバーヘッドが発生することを認めています。また、非常に限定的なローカルエポック数では、モデルが十分に整合した低ランクのプレフィックス表現を完全に学習できない可能性があることも指摘しています。

結論として、PreLortは、ランクの不均一性は、アダプタがどのように訓練されるか（ネストされたプレフィックス最適化を通じて）と、それらの更新がどのように集約されるか（セグメント単位のルールを通じて）を結合して構造化することによって、最も効果的に対処できることを示しています。これにより、既存の不均一な集約手法を凌駕する、整合され、堅牢なプレフィックス表現が得られます。