技術要約: SnapViT – 事前学習済みモデルから再学習なしで弾力的なVision Transformerを生成する

問題提起

ビジョン基盤モデルは、多様なタスクにおいて最先端の性能を達成していますが、通常、限定された事前決定されたサイズ（例：DINOv3やDeITファミリーの特定のパラメータ数）でリリースされます。この硬直性は、計算量、レイテンシ、コストに関する実世界の制約に直面している実務家に対し、不適切なデプロイメントの選択を強いることになります。既存のモデルの柔軟性に関するソリューションには、以下のような重大な限界があります：

• 知識蒸留 (Knowledge Distillation): ターゲットとなるアーキテクチャが既定である必要があり、多くの場合、公開されていない事前学習データセットに依存するため、堅牢性が損なわれます。
• 弾力的な推論 (Elastic Inference): Matryoshka表現やSlimmable Networksのような手法は、事前学習の段階で特定のネストされた構造を設計しておく必要があり、既存のオフザシェルフの事前学習済みモデルには適用できません。
• 構造化プルーニング (Structured Pruning): メモリと計算量を削減するのに効果的ですが、現在のほとんどのプルーニング技術は特定のタスクや制約に特化しており、通常は再学習を必要とします。これは、ラベル付きデータや分類ヘッドを必要としない、普遍的に適応可能なモデルとの間にギャップを残しています。

手法: SnapViT

著者らは、SnapViT（Single-shot network approximation for pruned Vision Transformers）を提案します。これは、再学習やラベル付きデータを必要とせずに、単一の事前学習済みモデルから連続的な弾力的なサブネットワークを生成する、ポスト事前学習型の構造化プルーニング手法です。本手法の核となるのは、削除するパラメータをランク付けするためのプルーナビリティ・スコア ($P$) であり、これは以下の2つのコンポーネントに基づいています。

1. ローカル・ヘッセ行列近似 (自己教師あり勾配)

分類ヘッドなしでローカルなパラメータ感度を捉えるために、本手法は自己教師あり勾配を用いてヘッセ行列の対角成分を近似します。具体的には、DINO目的関数を用いて、画像データセットに対する勾配を計算します。ローカルな曲率 $H^{(l)}$ は、勾配の二乗ノルムの平均として推定されます：
$$H^{(l)} \approx \frac{1}{N_D} \sum_{i=1}^{N_D} \|\nabla_\theta L_i\|^2$$
これにより、特定のブロック（例：フィードフォワード・ニューロンやアテンション・ヘッド）内における個々のパラメータの重要度の基準値が得られます。

2. グローバル・ヘッセ行列推定 (進化アルゴリズム)

ブロック間の依存関係（非対角ヘッセ成分）を無視するという制限に対処するため、SnapViTは指数自然進化戦略 (xNES) を用いた新しい近似を導入しています。

• xNESは、完全なヘッセ行列を明示的に計算する代わりに、ローカルな感度を再重み付けする構造単位のスケーリング係数 ($c$) を学習します。
• アルゴリズムは、これらのスケーリング係数に対する多変量ガウス分布を最適化します。適合度関数はラベルフリーであり、PCA圧縮後の元のモデルとプルーニング後のモデルの埋め込み間のコサイン類似性を測定します。
• 自然勾配更新を通じて、進化戦略の共分散行列 ($\Sigma$) は、真のブロック間ヘッセ行列を近似するように進化します（$H^{(g)} \approx \alpha \Sigma^{-1}$）。これにより、表現の質を低下させることなく、どのブロックを共同でプルーニングできるかを特定することが可能になります。

3. 統一されたプルーナビリティ・スコアと弾力的な生成

最終的なプルーナビリティ・スコアは、ローカルな感度とグローバルなスケーリング係数の要素ごとの積です：
$$P = \text{diag}(H^{(l)}) \odot (M c)$$
ここで、$M$ はブロックレベルの係数を個々のパラメータへと拡張するメンバーシップ行列です。

• シングルショット生成: 一度 $P$ が計算されれば、すべてのパラメータがグローバルにランク付けされます。この単一のランキングにより、さらなる最適化を行うことなく、任意の望ましいスパース性レベル ($S \in [0, 1]$) のサブネットワークを抽出できます。
• 範囲: 本手法は、フィードフォワード・ブロック内の行・列の組み合わせと、アテンション・ヘッド全体の両方のプルーニングを行います。

主な貢献

1. 弾力的なプルーニング戦略: 再学習なしで、あらゆる計算制約に適応可能な弾力的なモデルを生み出す、事前学習済みVision Transformerのための効率的かつ高速なプルーニング戦略。
2. 進化論的ヘッセ行列近似: 進化アルゴリズム (xNES) を用いて、ネットワーク構造の非対角成分を近似する新しい戦略。これにより、従来の対角近似が見落としがちなブロック間の相関関係を捉えることができます。
3. ラベルフリーかつヘッドに依存しない性能: ラベル付きデータや分類ヘッドを必要とせずに、様々なスパース性レベルにおいて最先端の性能を達成し、教師ありおよび自己教師ありの基盤モデルの両方に汎用可能です。

実験結果

著者らは、k-最近傍 (k-NN) 分類、線形プロービング、および線形セマンティック・セグメンテーションを用いて、8つのデータセット（ImageNet-1k、Pascal VOC、および各種ファイングレインド・ベンチマークを含む）でSnapViTを評価しました。

• SOTAとの比較性能: SnapViTは、様々なスパース性レベルにおいて、最先端の手法（LAMP, LLM Surgeon, FPTP, SparseGPT, SNIP Magnitude, NViT）と同等、あるいはそれ以上の性能を一貫して示しました。
  • DINOv1 ViT-B/16 において、本手法は40%のスパース性までプルーニングしても精度低下を5%未満に抑え、1.58倍の推論高速化を達成しました。
  • 教師ありAugRegモデルにおける50%のスパース性では、線形分類においてSNIPを7%、FPTPを12.3%上回りました。
  • SigLIPv2 および DINOv3（大規模基盤モデル）において、本手法はベースラインを大幅に上回り、線形分類において2番目に優れた手法に対して最大34.3%の改善を示しました。
• 効率性: モデルファミリー（例：DeIT-III）全体のプルーニングプロセスは、単一のA100 GPU上で5分未満で完了します。
• 後処理: 本手法は、ポストプルーニングによる重みの補正（SparseGPTスタイル）およびフルファインチューニングと互換性があります。単一の重み補正ステップにより、極端なスパース性（例：SigLIPv2における50%）においても性能がほぼ維持されます。
• スケーラビリティ: 大規模な事前学習（例：DINOv3, SigLIPv2）は、分散された表現知識のためにプルーニングをより困難にしますが、SnapViTは、重み補正やより長い最適化ホライゾン（500回のxNES反復）を組み合わせることで、依然として効果的です。

意義と主張

本論文は、SnapViTが現在の基盤モデルのデプロイメントにおける硬直性と、柔軟でリソース適応型の推論へのニーズとの間の溝を埋めるものであると主張しています。単一の事前学習済みモデルから連続的なサブネットワークをシングルショットで抽出することを可能にすることで、ユーザーは自身の特定の計算予算に合わせて最先端のモデルを精密にカスタマイズできます。

著者らは、本アプローチが以下の点を強調しています：

• 再学習の必要性を排除し、計算リソースと時間を大幅に節約します。
• ラベル付きデータへの依存を排除し、分類ヘッドを持たない自己教師あり基盤モデルにも適用可能にします。
• 進化論的な近似を通じてグローバルな構造的相関を捕捉し、ローカルのみのプルーニング指標が抱える限界を克服します。
• Vision Transformerの計算要件と電力使用量を低減することで、環境負荷を軽減し、推論によるCO2排出量の削減に貢献する可能性があります。

本研究は、高い性能を維持しながら、制約のある環境下で大規模なビジョンモデルをデプロイするための実用的なソリューションとして提示されています。