Rethinking Factor Sharing in Federated LoRA: A Rank-Aware Adaptive Approach

技術要約：Federated LoRAにおけるファクター共有の再考

問題提起
連邦学習（Federated Learning, FL）において、低ランク適応（LoRA）を用いた大規模言語モデル（LLM）の微調整は、「集約の不一致（aggregation mismatch）」と呼ばれる重大な課題に直面しています。標準的なLoRAは、モデルの更新を2つの低ランク・ファクターの積 $\Delta W_i = B_i A_i$ として表現します。$N$ 個のクライアントが存在する連邦設定では、単にファクターを独立して平均化すること（$\frac{1}{N}\sum B_i \cdot \frac{1}{N}\sum A_i$）は、更新の平均（$\frac{1}{N}\sum B_i A_i$）とは一致しません。

これらを緩和するための既存のアプローチは、多くの場合、「ハードコードされた」ファクター共有戦略を採用しています。例えば、全クライアントで共通のファクター $A$ を固定し、$B$ をローカルに保持する（Share-A/Local-B）、あるいはその逆（Share-B/Local-A）といった戦略です。しかし、本論文は、どちらの戦略も普遍的に最適であるわけではないことを示しています。どのファクターを共有するかという選択は、特定のデータ分布、LoRAのランク、およびシステム設定に依存します。固定された割り当ては、異なるクライアント間で必要とされる入力側と出力側の表現空間における構造的な非対称性を考慮できないため、サブオプティマル（劣最適）な性能を招く可能性があります。

手法
著者らは、トレーニング開始前に共有戦略を適応的に選択するフレームワークである FedAS-LoRA (Federated Adaptive Factor Sharing Low-Rank Adaptation) を提案しています。

1. 理論的洞察（最小二乗近似サロゲート）:
   本論文では、最小二乗近似サロゲートを用いてLoRAファクターの構造的な非対称性を分析しています。これにより、以下のことが明らかになります：

   • Share-A/Local-B は、全クライアントに共通のランク-$r$ の入力側空間を課します。その性能は、望ましい更新をこの共有された入力空間へ投影した際の投影残差によって制限されます。
   • Share-B/Local-A は、全クライアントに共通のランク-$r$ の出力側空間を課します。その性能は、共有された出力空間への投影残差によって制限されます。
     好ましい戦略は、クライアント間の集約的な投影残差を最小化する方となります。

2. ランクを考慮した共有部分空間充足度（RSS）指標:
   最適な共有側をトレーニングなしで決定するために、著者らはRSS指標を設計しました。この指標は、凍結されたLLMバックボーンから抽出された表現を用いて、グローバルなランク-$r$ の入力部分空間がローカルなデータ分布を捉えるのに十分であるかどうかを評価します。

   • これは、グローバルなランク-$r$ の部分空間によって捉えられる「表現エネルギー（第2次変動によって捕捉されるもの）」を、クライアント固有のランク-$r$ の部分空間によって保持されるエネルギーと比較します。
   • 閾値 $\tau(r)$ は、推定誤差を考慮するために、ランダムなサンプル再割り当てとブートストラップ・リサンプリングを用いて校正されます。
   • 決定ルール: グローバルな入力部分空間を使用することによるエネルギー損失（RSS欠損）が閾値を超えた場合、システムは Share-B/Local-A を選択します。そうでなければ、Share-A/Local-B を選択します。

3. 収束解析:
   本論文では、任意のクライアント参加（部分的参加）条件下において、FedAS-LoRAが定常的な近傍に収束するという理論的な証明を提供しています。この解析は両方の共有戦略に対して成立し、反復回数、ローカルステップ、およびクライアントサンプリングのドリフトに依存する収束界を満たすことを示しています。

主な貢献

• 投影ベースの比較: 本論文は、Share-A/Local-BとShare-B/Local-Aの最初の形式的な比較を提供し、それらがそれぞれ入力側および出力側の投影残差に対応することを明らかにしました。これにより、なぜ固定された戦略が多様な連邦設定において失敗するのかを説明しています。
• 適応型フレームワーク (FedAS-LoRA): ポリシーをハードコードするのではなく、データ特性に基づいて共有ファクターを動的に選択する新しいフレームワークです。
• トレーニングフリーの指標 (RSS): 凍結されたバックボーンの表現のみを使用して、事前のトレーニングを行うことなく最適な共有側を効率的に決定するRSS指標の設計。
• 収束保証: 従来の多くの研究で一般的であった全参加の仮定を超え、任意のクライアント参加条件下での提案手法の収束に関する理論的証明。

実験結果
自然言語理解（GLUEベンチマーク：MNLI, SST-2, QNLI, QQP, RTE）および自然言語生成（LLaMA3-8Bを用いたGSM8K）のタスクにおいて実験が行われました。

• 性能: FedAS-LoRAは、様々なタスクにおいて、標準的なLoRA、FFA-LoRA、FedDPA-LoRA、およびFedSA-LoRAを含むベースラインを一貫して上回りました。GLUEベンチマークにおいて、FedAS-LoRAは平均精度 90.77% を達成し、FedSA-LoRAを0.93ポイント上回りました。
• 堅牢性: 本手法は、以下を含む様々な条件下で優れた性能を示しました：
  • クライアントサンプリング: 一様および非一様サンプリングの設定。
  • データの不均一性: IIDおよびnon-IID（ディリクレ分布）の分布。
  • LoRAランク: ランク $r \in \{2, 4, 8, 16\}$ において競争力のある性能を維持し、単一の共有ポリシーが全てのランクに対して最適ではないことを検証しました。
• 汎用性: 本アプローチは、異なるLoRAのバリアント（rsLoRA, VeRA）およびモデルサイズ（RoBERTa-large, LLaMA3-8B）において効果的でした。

意義
本研究の意義は、単一のファクター共有戦略が全ての連邦LoRAシナリオに十分であるという仮定に異を唱える点にあると著者らは主張しています。最適な戦略がデータ分布とLoRAランクの相互作用に依存することを実証することで、著者らは原理に基づいた適応的なソリューションを提供しています。RSS指標は、リソース制約のある連邦環境において微調整の性能を高めるための、実用的かつトレーニング不要のメカニズムを提供し、共有されるファクターがクライアントデータの基礎となる構造的要求と一致することを保証します。