論文「REPRESENTATIONAL ALIGNMENT ACROSS MODEL LAYERS AND BRAIN REGIONS WITH MULTI-LEVEL OPTIMAL TRANSPORT」の技術的サマリー

この論文は、異なる深さやアーキテクチャを持つニューラルネットワーク（人工知能モデル）間、あるいは人工モデルと生物学的脳（fMRI データ）間の表現を比較・整合させるための新しい枠組み**「Multi-Level Optimal Transport (MOT)」**を提案しています。従来の層ごとの貪欲なマッチング手法の限界を克服し、グローバルに整合性のある柔軟な対応付けを実現する手法です。

以下に、問題設定、手法、主要な貢献、結果、および意義について詳細にまとめます。

1. 背景と問題設定 (Problem)

ニューラルネットワークの内部表現を比較する際、従来の標準的なアプローチは**「層ごとのペアマッチング（Layer-wise Matching）」**です。これは、ソースモデルの各層を、ターゲットモデルの中で最も類似する単一の層と 1 対 1 で対応付ける方法です（RSA, CKA, Procrustes 距離などを使用）。

しかし、この手法には以下の根本的な限界があります：

• 硬直的な 1 対 1 対応の強制: ネットワークの深さ（層数）が異なる場合、またはある層の表現がターゲットモデルの複数の層に分散している場合に、不適切な対応付けを強要してしまいます。
• 非対称性とグローバルスコアの欠如: 比較の方向によって結果が変わり（非対称）、ネットワーク全体としての統一的なアライメントスコアを提供できません。
• グローバル構造の無視: 各層を独立して最適化するため、ネットワーク全体の活性化構造を考慮せず、ノイズに過剰適合するリスクがあります。
• 階層構造の再現失敗: 貪欲な手法では、浅い層と深い層の対応関係が明確に現れず、特定の層だけが支配的になるなど、生物学的な階層性（例：V1→V2→V3）が再現されません。

2. 提案手法：Multi-Level Optimal Transport (MOT)

著者らは、Multi-Level Optimal Transport (MOT) という統合フレームワークを提案しました。これは、最適輸送（Optimal Transport: OT）理論をネットワークの「層」と「ニューロン」の両レベルで同時に適用するものです。

2.1 基本的な考え方

MOT は、ソース層の「表現質量（mass）」をターゲットの複数の層に柔軟に分配（ソフトマッチング）することを許可します。これにより、深さの不一致を自然に処理し、グローバルに整合性のある対応付けを実現します。

2.2 手法の詳細

MOT は 2 つのレベルで最適化を行います：

1. 内層レベル（Inner Level: ニューロン間輸送）:

   • 各層のペア $(\ell, m)$ に対して、ニューロン間のチューニングプロファイルに基づいた dissimilarity 行列を計算します。
   • 最適輸送問題（ソフトマッチング）を解き、層 $\ell$ の各ニューロンが層 $m$ のどのニューロンに対応するかを示す「ソフトな結合行列 $Q_{\ell m}$」を推定します。
   • 層の幅が異なる場合でも、ニューロンが複数のターゲットニューロンの混合に対応することを許容します。

2. 外層レベル（Outer Level: 層間輸送）:

   • 内層レベルで得られたコスト行列 $C$ を用いて、層間のグローバルな結合行列 $P$ を推定する 2 番目の最適輸送問題を解きます。
   • 制約条件:
     • 各ソース層は、その質量の 100% をターゲット層全体に分配する（情報損失なし）。
     • 各ターゲット層は、すべてのソース層から受け取る質量の合計がバランスよく配分される。
   • これにより、$L \neq M$（層数が異なる）場合でも、1 つのソース層が複数のターゲット層にまたがって対応付けられる「ソフトな多対多マッピング」が自然に生成されます。

2.3 回転不変性の拡張 (MOT + R)

多くの表現類似性指標（CKA, RSA など）は回転不変ですが、OT ベースの手法は通常、ニューロンの順序や基底の回転に敏感です。これを克服するため、各層ペアに対して直交変換行列 $R_{\ell m}$ を導入し、再構成誤差を最小化する「MOT + R」を提案しています。これにより、基底が回転していても幾何学的な等価性を捉えることが可能になります。

2.4 評価指標

学習した輸送計画（$P$ と $Q$）を用いて、ソース層の表現をターゲット層から再構成し、保持データ（held-out data）上での再構成相関係数を「アライメントスコア」として評価します。

3. 主要な貢献 (Key Contributions)

1. 原理的なグローバルアライメント枠組み:
   • 貪欲なペアマッチングとは異なり、すべての層の対応関係を同時に最適化し、対称的でグローバルに整合したアライメントと単一のスコアを提供します。
2. 深さの不一致への自然な対応:
   • 1 対 1 対応を強制せず、ソフトな多対多マッピングを通じて、異なる深さのネットワーク間でも適切な対応付けを可能にします。
3. 回転不変性の拡張:
   • 表現空間の回転を考慮した拡張版（MOT + R）を提案し、基底の違いに頑健な高品質なアライメントを実現します。
4. 階層的構造の自然な発現:
   • 明示的な順序制約を課さなくても、MOT は「浅い層は浅い層へ、深い層は深い層へ」という階層的な対応関係を自然に回復します。これは貪欲な手法では見られない特徴です。
5. 3 レベル MOT によるトレーニング軌道の整合:
   • 2 つのモデルのトレーニング軌道（チェックポイント間）を整合させる「3 レベル MOT」を実証し、単なる層間比較を超えた動的な対応付けの可能性を示しました。

4. 実験結果 (Results)

MOT は、ビジョンモデル、大規模言語モデル（LLM）、人間の視覚皮質（fMRI）の 3 つの異なるドメインで評価されました。

• LLM 間（LLaMA, Qwen）:
  • MOT はベースライン（ペアワイズ OT, 線形予測など）を上回る再構成スコアを達成しました。
  • 輸送計画は明確な対角構造（階層性）を示し、浅いモデルの 1 層が深いモデルの複数の連続する層に分散して対応していることが確認されました。
• 視覚皮質間（fMRI 被験者間）:
  • 異なる被験者間の脳領域（V1-V4 など）の対応付けにおいて、MOT はペアワイズ手法よりも構造化された領域間対応を回復しました。
  • 再構成スコアはベースラインと同等かそれ以上であり、生物学的に意味のある対応付けを提供します。
• ビジョンモデル間（ViT, DINOv2）:
  • 回転不変性を考慮しない標準 MOT はペアワイズ手法と同等程度でしたが、MOT + R（回転拡張版） は劇的に高い再構成スコアを達成し、明確な階層的対応を回復しました。
  • これは、Transformer の残差ストリームが回転不変であるという性質を正しく捉えていることを示しています。
• トレーニング軌道の整合:
  • Pythia モデルのトレーニングチェックポイント間での 3 レベル MOT 適用により、ペアワイズ手法では見逃される、トレーニング軌道全体の一貫した対応関係が明らかになりました。

5. 意義と結論 (Significance)

この研究は、異なるアーキテクチャや深さを持つシステム間（AI モデル同士、AI と脳）の表現を比較する際、**「グローバルな構造を考慮した柔軟な対応付け」**の重要性を浮き彫りにしました。

• 理論的意義: 貪欲な局所最適化ではなく、大域的最適化によって得られるアライメントが、より解釈可能で生物学的・構造的に意味のある対応関係（階層性、深さの不一致の解決）を自然に発現することを示しました。
• 実用的意義: 深さやアーキテクチャが異なるモデル間、あるいは脳と AI の比較において、単なるスコアだけでなく、表現がどのように分散・再構成されているかという「メカニズム」の洞察を提供します。
• 将来展望: MOT は、モデルのトレーニングダイナミクスや、異なる学習条件下での普遍的原理（Universality）の解明に向けた強力な測定基盤となります。

総じて、MOT は、複雑化するニューラルネットワークと生物学的システムの比較において、より豊かで解釈可能な対話を可能にする画期的な手法です。