A Unified Adaptive Feature Composition Framework for Multi-Task Generalization in Wireless Foundation Models

問題提起

無線基盤モデル（WFM）は、普遍的なチャネル表現を学習する可能性を示しているが、現在のパラダイムでは、多様なダウンストリーム・タスクへの適応において大きなボトルネックに直面している。既存のアプローチは一般的に以下の2つのカテゴリに分類されるが、いずれも限界がある：

1. フル・ファインチューニング（全層微調整）： 柔軟ではあるが、膨大な計算およびストレージのオーバーヘッドを導入するため、ヘテロジニアスなタスクにおけるデプロイメントやメンテナンスを複雑にする。
2. 凍結された特徴抽出（Frozen Feature Extraction）： このパラダイムは通常、単一の固定された表現（通常は最終層の出力）に依存する。しかし、これは性能のボトルネックを生み出す。なぜなら、異なる無線タスクには異なるレベルの抽象化が必要だからである。例えば、チャネル推定にはきめ細かな時間・周波数およびマルチパス構造（低次レベルの特徴）が必要であるが、ビーム管理や位置推定には、グローバルな空間的・幾何学的パターン（高次レベルの特徴）が必要となる。

本論文は、最終層のみに依存することは、深い表現が特定のタスクに必要な微細な詳細を失ってしまう「特徴の過度な平滑化（feature over-smoothing）」を引き起こすと主張している。さらに、自然言語処理（NLP）やコンピュータビジョンからの既存のマルチレイヤー融合技術は、パラメータが多すぎるか、あるいは単一のタスクに密結合しており、WFMのための統一された軽量なインターフェースを提供できていない。

手法

これらの問題に対処するため、著者らは、**特徴構成のためのルーティング・アダプター（RAFC: Routing Adapter for Feature Composition）**と呼ばれる軽量なモジュールを中心とした、統一適応型特徴構成フレームワークを提案している。

1. コア概念：適応型マルチレベル集約
最終層の出力のみを抽出するのではなく、本フレームワークは、異なるTransformerの深さからの隠れ状態を、再利用可能なマルチレベル特徴のプールとして扱う。RAFCモジュールは、ダウンストリーム・タスクの特定の要件に基づいて、これらの階層的な表現を動的に集約する。

2. RAFCアーキテクチャ
RAFCモジュールは、以下の3つの連続するステップを持つ、モデルに依存しないプラグインとして機能する：

• 層のコンパクトな記述（Layer Compact Description）： フル・トークン・シーケンスを処理することによる二次的な計算コストを避けるため、選択された層からのトークン・シーケンスは、まずトークンごとの平均プーリングによって圧縮される。これにより、コンパクトなグローバル層記述が生成される。
• 学習可能な重みの割り当て（Learnable Weights Allocation）： これらの記述は結合され、軽量なボトルネック型多層パーセプトロン（MLP）に投入される。このMLPは、特徴を投影・縮小し、ReLU活性化関数を適用した後、再び拡張して一連のルーティング重み（$\alpha$）を生成する。ソフトマックス正規化により、これらの重みの合計は1となる。
• 加重和（Weighted Summation）： 元の階層的な特徴テンソルは、生成された重みを用いて、サンプル適応型の加重和によって融合される。これにより、局所的な構造とグローバルな意味論のバランスをとる複合表現 $z^{(q)}$ が生成される。

3. 層の選択戦略
冗長性とオーバーヘッドを避けるため、フレームワークはすべての層を使用するのではなく、代表的な層のセット $S = \{2, \lfloor N/2 \rfloor + 1, N\}$ を選択し、低次、中次、および高次の特徴を捉える。

• なぜ第2層なのか？ 著者らは、第1層は観測ノイズに対して非常に敏感である（高い摂動比を持つ）ことを経験的に示している。一方、第2層は低次のノイズを効果的に抑制しつつ、微細な物理的変動を保持するため、最適な低次インターフェースとなる。
• なぜ全層ではないのか？ 一様なサンプリングは、表現の多様性とデプロイメント効率の間の実用的なトレードオフを提供する。

4. 学習パラダイム
学習済みのWFMバックボーンは厳格に凍結されたままとなる。最適化されるのは、タスク固有のRAFCモジュールとダウンストリーム・タスク・ヘッドのパラメータのみである。これにより、フレームワークは無視できる程度の追加パラメータ（約0.045M）のみを導入し、基盤モデルの低遅延という利点を維持することを保証する。

主な貢献

• 統一フレームワーク： 単一層の抽出から動的なマルチレイア集約へと適応パラダイムを転換し、深いWFMにおける過度な平滑化の問題を軽減する、モデルに依存しない特徴構成フレームワーク。
• 軽量設計： RAFCモジュールは、トークンごとの圧縮層と小さなMLPのみで構成されており、無視できる程度の計算およびストレージのオーバーヘッド（追加パラメータ50K未満）しか導入しない。
• 解釈可能性： 学習されたルーティング重みは物理的な解釈可能性を提供し、異なるタスクが本質的に異なる物理的表現レベル（例：再構成のための微細な詳細 vs 空間知覚のためのグローバルな抽象化）を優先していることを明らかにしている。

実験結果

本フレームワークは、学習済みのTransformerベースのWFM（8層、6.62Mパラメータ）を用い、4つの代表的なダウンストリーム・タスク（チャネル推定、チャネル予測、ビーム予測、都市部位置推定）で評価された。

• 性能向上：
  • チャネル推定： 最終層のみのベースラインと比較して、正規化平均二乗誤差（NMSE）を**38%**削減。
  • チャネル予測： NMSEを**80%**削減。
  • ビーム予測： Top-1精度において**3%**の向上を達成。
  • 都市部位置推定： 絶対誤差を**34%**削減。
• 効率性： 本手法は総パラメータサイズの約0.7%しか追加せず、遅延の増加も極めて小さい（0.0002 msから0.0011 ms）。
• 転移性： 本フレームワークは別のバックボーンであるLWMv1.1にも適用可能であり、一貫した性能向上を示した。これは、特定の基盤アーキテクチャから切り離されていることを証明している。
• アブレーション研究：
  • 特定の層（2, 5, 8）を選択することは、性能とオーバーヘッドのトレードオフの観点から、浅い層のみ、深い層のみ、または全層を使用する場合よりも優れた結果となった。
  • 平均プーリングは、アテンションベースまたはフル・トークン・ルーティングと比較して、最小限のFLOPsで最大に近い性能を提供し、最適な集約戦略であることが判明した。

意義と主張

本論文は、RAFCが、WFMを多様なダウンストリーム・シナリオに適応させるための、低複雑度、スケーラブル、かつ説明可能なインターフェースを提供すると主張している。

• タスク固有の優先順位： 学習された重みの分析により、信号再構成タスク（推定・予測）は「粗から精（coarse-to-fine）」の学習軌跡を示し、最終的には微細な物理的変動を回復するために低次レベルの特徴に大きく依存することが明らかになった。対照的に、空間知覚タスク（ビーム予測・位置推定）は、グローバルな幾何学的構造を捉えるために、中次から高次の抽象化を好む。
• 動的な適応性： ルーティング・メカニメントは静的ではなく、入力条件に適応する。例えば、非見通し（NLoS）シナリオでは、ルーターは抽象的なマクロ構造を利用するために深い層へと重みをシフトさせ、端末速度が高い場合は、瞬時的な信号変動を捉えるために浅い層へとシフトする。
• 実用的なデプロイメント： フル・ファインチューニングや重いクロスアテンション・メカニズムを必要としないことで、本フレームワークは無線システムにおける厳格な軽量性と低遅延の要件に対処し、ヘテロジニアスなタスク間で学習済みWFMを効率的に再利用することを可能にする。