技術概要：フェデレーテッド・ネストド・ラーニング（FedNL）

問題定義

大規模言語モデル（LLM）におけるフェデレーテッド・ラーニング（FL）は、データ異質性（Non-IID）とロングテール分布という、二つの持続的かつ密接に関連する課題に直面しています。現実的なシナリオでは、クライアントのデータ分布は極めて偏っており（例：医療データ対コーディングデータ）、重要な知識はしばしばロングテールに存在し、グローバル集約時にヘッドクラスによって容易に霞んでしまいます。

既存のアプローチは、主に静的なモデル重みの正則化やパーソナライゼーションに焦点を当てています（例：FedProx、FedSSI）。これらは破滅的な忘却の緩和には効果的ですが、グローバルモデルを静的知識のコンテナとして扱っています。その結果、未見で高度に異質なデータを持つクライアントにデプロイされた際、モデルは勾配更新なしで現在の文脈に適応するテスト時の可塑性を欠きます。これにより、文脈に依存するニュアンスが重要なドメイン固有タスクにおいて、性能が最適化されないことになります。

手法：フェデレーテッド・ネストド・ラーニング（FedNL）

本論文は、静的知識の集約から最適化則の協調的学習へと FL を再定義するフレームワーク、**フェデレーテッド・ネストド・ラーニング（FedNL）**を提案します。**ネストド・ラーニング（NL）**理論に着想を得て、FedNL は推論を受動的な検索プロセスではなく、モデルが文脈を一時的なメモリ状態に能動的に圧縮する高頻度の「内側ループ」トレーニングプロセスとして扱います。

コアアーキテクチャ

FedNL は、FL を3 段階のネストド最適化システムとして再定式化します：

1. レベル 0（外側ループ - サーバー）：協調的汎化。サーバーはメモリ内容ではなく、メタ則（パラメータ $\theta$）を集約します。これらの則は、メモリが構築され更新される方法を支配します。
2. レベル 1（中間ループ - クライアント）：則の学習。ローカルトレーニング中、クライアントは内側ループが下流タスクに有用なメモリ状態を生成するように $\theta$（例：LoRA アダプター、ゲーティング機構）を最適化します。これには、メモリ軌跡を通じた逆伝播が含まれます。
3. レベル 2（内側ループ - クライアント）：テスト時適応。推論中、クライアントはグローバル則 $\theta$ を用いて、デルタ則を通じてプライベートで文脈を認識するメモリ状態 $S_t$ をオンラインで更新します。この状態は一時的であり、デバイスから決して送信されません。

技術的実装

• Titans ベースの線形アテンション：FedNL は、線形化されたアテンション機構を利用するTitansアーキテクチャを活用します。メモリ状態 $S_t$ はデルタ則によって更新されます：
  $$S_t = S_{t-1} + \eta(v_t - S_{t-1}k_t)k_t^\top$$
  この更新は、連想メモリ目的関数に対する勾配降下の一ステップと数学的に等価であり、推論をオンライン最適化プロセスとして効果的に扱います。
• 知識と則の分離：モデルは、一般的な言語能力を提供するフリーズされたバックボーンと、学習可能なメタパラメータ（$\theta$）で構成されます。メタパラメータには、クエリ/キー/バリュー生成用の低ランク射影（LoRA）と、静的および動的アテンションのバランスを取るメモリゲーティング機構（$\alpha$）が含まれます。
• 効率性：フルモデル重みやメモリ状態ではなく、小さなメタ則（アダプター）のみを集約することで、FedNL は通信オーバーヘッドを劇的に削減します。さらに、線形アテンション機構は、標準的なアテンションが $O(T)$ の KV キャッシュを必要とするのに対し、シーケンス長に関わらず一定の推論メモリ複雑度（$O(1)$）を維持します。

主要な貢献

1. 理論的再定義：本論文は、FL に対してネストド・ラーニングの視点を導入し、静的表現ではなく最適化則の協調的トレーニングとして問題を形式化します。これにより、テスト時適応を通じて Non-IID 課題に対処するための理論的基盤が提供されます。
2. FedNL フレームワーク：Titans ベースの線形アテンションを FL パイプラインに統合する実用的なアルゴリズムです。デルタ則をオンライン勾配降下ステップとして扱うことで、重い計算負荷なしに未見ドメインでのゼロショット・テスト時適応を可能にします。
3. 実証的優位性：実験により、FedNL は短い文脈推論において競争力のある性能を達成し、強力なフェデレーテッドベースラインと比較して、長い文脈検索およびストリーミング診断において優れた向上を示すことが実証されました。

実験結果

著者らは、5 クライアントの Non-IID MMLU分割と、7 クライアントの Non-IID **Needle In A Haystack（NIAH）**長文脈ベンチマークの 2 つのフェデレーテッド設定で FedNL を評価しました。

• Non-IID MMLU（短文脈推論）：

  • Qwen2.5-1.5B バックボーンにおいて、FedNL は最高の平均精度（58.88%）を達成し、最良のベースラインである FedSSI（58.70%）をわずかに上回りました。
  • STEM（+2.0 ポイント）および Math/CS（+3.8 ポイント）分野で顕著な向上が見られました。
  • Llama-3.2-1B バックボーンにおいて、FedNL は**42.64%**に達し、FedSSI の 42.10% を上回りました。
  • 重要なのは、FedNL が「集約による低下」（ローカルモデルからフェデレーテッドモデルへ移行する際の性能損失）を最小化し、分野全体でほぼゼロに抑えたことです。一方、静的ベースラインは顕著な劣化を示しました。

• 長文脈検索（NIAH）：

  • FedNL は、7 つの検索テンプレート全体で最高平均精度（29.7%）を達成し、FedALA（28.6%）を上回りました。
  • 向上は MK-NIAH、MV-NIAH、UUID コードタスクで最も顕著でした。
  • ストリーミング交差エントロピー（CE）診断：16K トークンのストリーミングプローブにおいて、FedAvg の正規化 CE は文脈が展開されるにつれて 3.1% 増加しました（不確実性の蓄積）。これに対し、FedNL は2.1% 減少し、反復的なメモリ状態が長いストリームにわたって有用な文脈を成功裡に吸収したことを示しました。
  • メモリ効率：FedNL は一定の推論メモリを維持し、静的アテンションベースラインに影響を与えた 16K トークンにおけるメモリ不足エラーを回避しました。

• 通信効率：

  • FedNL は、テスト設定における学習可能パラメータについて、FedAvg と比較してラウンドあたりの通信を約350 倍削減しました（0.26 MB 対 22.5 MB）。

意義と主張

本論文は、FedNL が、モデルを静的知識の貯蔵庫から継続的かつ文脈を認識する学習のパラダイムへと進化させる、フェデレーテッド・ラーニングの新たな方向性を確立すると主張しています。

• パラダイムシフト：一般的な言語能力（フリーズされたバックボーン）をメモリ構築則（学習可能なアダプター）から分離することで、FedNL は多様なクライアントが矛盾する「知識状態」ではなく、「学習能力」を共有することを可能にします。
• 異質性への頑健性：グローバルメタパラメータは、異なるドメイン（例：コード対医療）間の矛盾を符号化する必要はありません。代わりに、「文脈が X を含む場合、メモリを Y になるように更新する」という普遍的な則を符号化します。これにより、パラメータ空間における Non-IID 対立が緩和されます。
• ゼロショット適応：このフレームワークは、モデルがテスト時トレーニング（TTT）手法がしばしば直面する「コールドスタート」問題を解決するため、学習されたメモリ更新則を実行するだけで、推論の最初の数トークンの間に特定のローカル分布に適応することを可能にします。

著者らは、1B〜1.5B スケールでは結果が一貫しているものの、メモリ則集約の一般性を検証するために、FedNL をより大規模なファウンデーションモデルへ拡張することが次の必要なステップであると指摘しています。