技術要約: WeightCLIP

問題提起

重み空間学習（Weight Space Learning: WSL）は、ニューラルネットワーク（NN）のパラメータを表現学習のためのデータモダリティとして扱い、モデルの特性予測や新しいモデルの合成といったタスクを可能にする。既存のアプローチは有望な成果を示しているものの、決定的な限界に直面している。それは、学習された重み空間の表現に明示的な意味論的参照フレームが欠けていることである。その結果、純粋に再構成や特性予測のために学習された潜在空間はナビゲートが困難であり、潜在空間における近接性が必ずしもデータセットレベルの変動に対応するわけではなく、「ナビゲーション」が特定のデータ分布に直感的にマッピングされることもない。これは、高次元の重み空間における高コストなサンプリングを行うことなく、ゼロショット転移、制御可能な重みの生成、あるいは効果的な洗練（refinement）を行う能力を制限している。

著者らは、学習済みモデルは、自身が学習したデータセットという明示的な「キャプション」を内包していると仮定している。解決すべき核心的な問題は、このデータセット情報をどのように活用して重み空間の表現を再形成し、テキストが視覚と言語のモデルを導くのと同様に、データセットが自然な監督信号として機能する共有潜在空間を構築するかである。

手法: WeightCLIP

WeightCLIPは、対照学習（contrastive learning）の目的関数を用いて、データセットの表現とモデルの重みの表現を整合させる手法を提案する。このフレームワークは、主に以下の3つのコンポーネントで構成される：

1. 重み空間オートエンコーダ: 先行研究（Schürholt et al., 2024）に従い、ニューラルネットワークは「重みトークン」（層のパラメータの平坦化された部分集合）のシーケンスとして表現される。エンコーダ $g_\phi$ はこれらのトークンを潜在表現へと圧縮し、デコーダ $h_\psi$ は重みを再構成する。
2. データセットエンコーダ: データセットエンコーダ $d_\theta$（DeepSetsアーキテクチャに基づく）は、データセットからの少数の画像サブセットを、データセット埋め込み $e_D$ へとマッピングする。
3. 対照的整合（Contrastive Alignment）: コアとなる革新性は、モデルの重みのトークン表現とデータセットの埋め込みを整合させるトークンレベルの対照損失である。独自の投影空間を学習する手法とは異なり、WeightCLIPはオートエンコーダの潜在空間を直接再形成する。損失関数には、再構成項（$L_{recon}$）と、モデルのトークン表現とその学習データセットの埋め込みとの間の類似性を最大化し、他のデータセットとの類似性を最小化する双方向の対照的整合項（$L_{align}$）が含まれる。

ダウンストリーム・コンポーネント

潜在空間が整合された後、WeightCLIPは以下の2つのメカニズムをダウンストリームタスクに導入する：

• データセットからモデルへのマッピング: 新しいデータセットプロンプトからモデルを生成するために、軽量なマッパーがデータセット埋め込み $e_D$ を一連のモデルトークン表現 $Z$ へと変換する。著者らは2つのマッパーを提案している：
  • リニアマッパー（Linear Mapper）: データセット埋め込みからトークンシーケンスを回帰する単純なアフィン変換。
  • メモリバンクマッパー（Memory-Bank Mapper）: 学習済みのモデルの多様体（manifold）へと予測をバイアスさせる手法。これは、ソフトな近傍分布によって重み付けされた、保存された学習済みトークン表現の凸結合を通じてトークンを合成する。
• 潜在的洗練（Latent Refinement）: 生成後の不一致（例：分類器ヘッドの違い）に対処するため、著者らは重みのファインチューニングではなく、潜在空間内で直接モデルを洗練させる手法を提案している。このプロセスは以下の手順で行われる：
  1. グローバル事前分布: 初期潜在トークンを、学習データから推定された超球面シェル上に投影し、それらが有効なモデル多様体上に存在することを保証する。
  2. 局所最適化: 小規模なターゲットデータセットに対するタスク損失を最小化するように潜在トークン $Z$ を最適化する。その際、更新が学習された多様体の近くに留まるよう制約を課す。これは標準的な重み空間のファインチューニングとは構造的に異なる。

主な貢献

本論文は、主に3つの貢献を行っている：

1. データセットに整合された潜在空間: 対照的整合を通じてデータセットを参照点として用いることで、重み空間の潜在空間が整理されることを実証した。これにより、モデルが類似したデータセットで学習された場合にクラスター化するという、直感的なナビゲーションが可能になる。
2. データセット誘導型の重み生成: データ埋め込み（データプロンプト）を直接一連のモデル埋め込みへとマッピングし、それをデコードしてニューラルネットワークの重みを生成する手法を導入した。これにより、アウトオブディストリビューション（OOD）のデータセットを含む、ターゲットデータセットに特化したモデルの生成が可能となる。
3. 潜在多様体の洗練: 超球面シェルおよび局所的な近傍に制約をかけながら、デコーダを通じてタスク損失を最適化する洗練手順を提案した。これにより、標準的な生成重みのファインチューニングよりも効果的なテスト時適応が可能となる。

実験結果

著者らは、10,000個のCNNチェックポイントを含む大規模なモデルズーと、1,000個のResNet-18チェックポイント（多様な画像分類データセットで学習）を用いてWeightCLIPを評価した。

• 潜在空間の組織化: t-SNE可視化とk-NN分類により、整合（alignment）がモデルの表現を学習データセットごとに有意にクラスター化することを示した（非整合のベースラインと比較して、分類精度が77%から98%に向上）。データセットプロンプト間の補間は、モデル性能の滑らかなトレードオフをもたらし、ナビゲーション可能性を確認した。
• 検索（Retrieval）: WeightCLIPは、データセットからモデルへの検索を大幅に改善した。インディストリビューションの設定では、Recall@1で95%を達成し、非整合の学習およびTANSのベースラインを上回った。OOD設定においても、検索されたモデルはベースラインよりも高いゼロショットおよびファインチューニング精度を示した。
• 生成: データセットプロンプトをマッピングして新しいモデルを生成する場合、WeightCLIPはスクラッチからの学習、TANS、およびハイパーネットワークのベースライン（Text2Model）を上回った。初期状態（ファインチューニング0エポック）であっても、生成されたモデルはターゲットタスクに対して強い整合性を示した。メモリバンクマッパーは一般に初期性能において最良の結果を提供したが、リニアマッパーも競争力のある性能を維持した。
• 洗練: 潜在的洗練は、同じ計算予算を用いて、標準的な重み空間のファインチューニングと同等またはそれ以上の性能を達成した。アブレーション研究により、超球面制約と初期の整合の両方がこれらの利得に必要であることが確認された。

意義と主張

本論文は、データセット情報を明示的に組み込むことで重み空間の表現を再形成することにより、既存のWSL手法における主要な限界に対処していると主張している。データセットを意味論的参照フレームとして扱うことで、以下のことが可能になる：

• 構造化されたナビゲーション: 潜在空間は、データセットの意味論に基づいてクエリ可能かつナビゲート可能になる。
• 効率的なモデル合成: 高価なサンプリングやスクラッチからの学習を行うことなく、新しいデータセットに対して強力なモデルの初期化を生成できる。
• 改善された適応: 提案された潜在的洗練は、標準的なファインチューニングとは構造的に異なり、しばしばより効果的な代替手段となる。

著者らは、本研究を、既知のアーキテクチャに対して新しいデータセットへの適用をより時間的・計算量的に効率的にするためのステップとして位置づけている。なお、現在の範囲はコンピュータビジョンのデータセットと均質なモデルアーキテクチャに限定されていること、また、本手法は重み空間の対称性を明示的に強制するのではなく、学習時の拡張（augmentation）に依存して対称性を扱うことによる制限があることを認めている。