単一パスにおける自己注意の剪定と畳み込み層への変換

Vision Transformers (ViTs) は、さまざまなコンピュータビジョンタスクにおいて優れた性能を達成しています。しかし、多頭自己注意（Multi-Head Self-Attention, MSA）層を用いて全局的な相関関係をモデル化すると、計算リソースの大量消費と局所的な視覚パターンのモデリングに対する内在的な誘導バイアスの欠如という、広く認識されている2つの問題が生じます。これらの問題を解決するために、単純かつ効果的な手法である Single-Path Vision Transformer pruning (SPViT) を考案しました。この手法により、事前学習された ViTs を効率的に自動で圧縮し、適切な局所性が追加されたコンパクトなモデルに変換します。具体的には、まず MSA と畳み込み操作間の新しい重み共有スキームを提案します。これにより、すべての候補操作をエンコードする単一パス空間が提供されます。この方法により、各 MSA 層でどのパラメータの部分集合を使用するかを見つけるという操作探索問題が大幅に簡素化され、計算コストと最適化の難易度が著しく低下します。また、事前学習された MSA パラメータを使用して畳み込みカーネルを良好に初期化することができます。単一パス空間に依存して、学習可能なバイナリゲートを導入し、MSA 層での操作選択をエンコードします。同様に、学習可能なゲートを使用して FFN 層における細かい MLP 拡大比率もエンコードします。これにより、SPViT は広大で統一された探索空間から自動的に最適な解を探し出し、各個別の稠密モデルに対して MSA-FFN の圧縮比率を柔軟に調整することができます。私たちは代表的な2つの ViTs を用いて広範な実験を行い、SPViT が ImageNet-1k 上での圧縮において新たな最先端（State-of-the-Art, SOTA）の結果を得ることを示しました。例えば、SPViT は DeiT-B の FLOPs を52.0%削減しつつ、トップ1精度で0.6%の向上を達成しています。ソースコードは https://github.com/ziplab/SPViT で公開されています。