ToaSt: Token Channel Selection and Structured Pruning for Efficient ViT

ViT の高計算コストを解決するため、マルチヘッド自己注意機構への結合型構造化プルーニングと、フィードフォワードネットワーク向けにグローバル伝播を回避するトークンチャネル選択（TCS）を組み合わせたフレームワーク「ToaSt」を提案し、精度と効率性の両立を実現した。

要約

この論文「ToaSt」は、画像認識 AI（Vision Transformer）を**「もっと軽く、もっと速く、そしてもっと賢く」**するための新しい方法を紹介しています。

AI は頭が良すぎて、スマホや小さな機械で動かすのが大変なほど重くなってしまっています。これを解決するために、研究者たちは「不要な部分を削ぎ落とす（圧縮する）」技術を開発しました。しかし、これまでの方法は「削るのに時間がかかりすぎる」か「削りすぎて頭が悪くなってしまう」という問題がありました。

ToaSt は、この問題を**「2 つの異なる道具」**を使って、まるで料理のレシピを最適化するように解決しました。

---

🍳 料理の例えで理解する ToaSt

AI を「巨大な料理屋」と想像してください。この料理屋には、2 つの主要な作業場（モジュール）があります。

1. 注目する場所を決める作業場（MHSA）：「この料理のどこが重要か？」を全体的に見渡して決める場所。
2. 味付けと調理をする作業場（FFN）：実際に材料を混ぜて味を決める場所。実は、ここが全体の作業量の60% 以上を占めています。

これまでの圧縮技術は、この 2 つの作業場を**「同じ方法」**で無理やり縮めようとしていました。でも、それは「作業場全体を壊して作り直す」ようなもので、とても時間がかかり、味（精度）が落ちてしまうことがありました。

ToaSt は、**「それぞれの作業場に合った、専用の道具」**を使います。

1. 作業場 1（MHSA）へのアプローチ：「チームの人数調整」

• これまでの方法：チーム全員をバラバラに減らすと、連携が取れなくなって失敗します。
• ToaSt の方法：「チームの連携」を重視します。
  • 料理人のチーム（ヘッド）が 8 人いるとします。ToaSt は、「質問をする人（Q）」と「答えを探す人（K）」はセットで減らし、「材料を渡す人（V）」と「結果を渡す人（P）」もセットで減らすというルールを作りました。
  • これにより、チームの人数は減っても、連携は完璧に保たれます。
  • 結果：作業場を壊さずに、必要な人数だけ減らして、「再トレーニング（味付けのやり直し）」の時間を大幅に短縮できました。

2. 作業場 2（FFN）へのアプローチ：「不要なスパイスの選別（Token Channel Selection）」

• これまでの問題：味付け作業場には、実は**「ほとんど使われていないスパイス（ノイズ）」**が大量に入っていました。でも、それを削ると味が変わってしまうと恐れ、削りきれませんでした。
• ToaSt の発見：
  • 深い層（後半の工程）に行くと、スパイスの**「90% 以上が実は使われていない（無駄）」**ことがわかりました。
  • さらに、**「一部のスパイスを使えば、残りのスパイスの味をほぼ再現できる」**という性質（線形依存性）があることも発見しました。
• ToaSt の方法：
  • **「トレーニング不要（Training-free）」**な選別機を導入しました。
  • 全スパイスを味見する必要はなく、**「少しだけサンプルを味見する」**だけで、「どのスパイスが本当に必要か」を瞬時に判断します。
  • 不要なスパイス（ノイズ）を大胆に捨てます。
  • 結果：作業場のスペースが空いて、**「余計なノイズがなくなるので、逆に味が（精度が）良くなる」**という奇跡が起きました。

---

🚀 ToaSt がもたらすすごい成果

この「2 つの道具」を組み合わせることで、以下のような驚くべき結果が出ました。

• 超高速化：AI の計算量が約 40% 減りました。つまり、スマホでもサクサク動くようになります。
• 精度向上：削ったはずなのに、精度が逆に上がりました（例：ViT-MAE-Huge で +1.64%）。
  • なぜ？ → 不要なノイズ（無駄なスパイス）を捨てたおかげで、AI が本当に重要な情報に集中できるようになったからです。
• 再トレーニングが楽：これまでの方法では、削った後に元の味を取り戻すために何ヶ月もかかることがありました。ToaSt は、最大モデルでも 15 回程度の短い練習（微調整）だけで、元の味（以上）を取り戻せます。

🎯 まとめ：なぜこれが画期的なのか？

これまでの AI 圧縮は、「重たい荷物を下ろすために、荷物を壊して作り直す」ようなものでした。

ToaSt は、**「荷物の中身を確認して、本当に不要なゴミだけを捨て、残った荷物を整理整頓する」**というスマートな方法です。

• ゴミ（ノイズ）を捨てる → 精度が上がる。
• 整理整頓 → 計算が速くなる。
• 作り直し不要 → すぐに使える。

この技術を使えば、高性能な AI が、もっと手軽なデバイス（スマホやドローンなど）で、より賢く、より速く動くようになる未来が近づきます。

技術概要

ToaSt: Token Channel Selection and Structured Pruning for Efficient ViT の技術的サマリー

本論文は、Vision Transformer (ViT) の高い計算コストという課題に対し、ToaSt（Token Channel Selection and Structured Pruning）と呼ばれる新しい圧縮フレームワークを提案しています。ToaSt は、重みの構造化プルーニングとトークンチャネルの選択を組み合わせることで、再学習（リトレーニング）のオーバーヘッドを最小化しつつ、精度を維持・向上させながら大幅な効率化を実現します。

以下に、問題定義、手法、主要な貢献、実験結果、そして意義について詳細に解説します。

---

1. 問題定義 (Problem)

Vision Transformer (ViT) は画像分類や物体検出など多様なタスクで優れた性能を発揮していますが、実用化には以下の重大な計算コストの課題が存在します。

1. 計算量の膨大さ:
   • 自己注意機構 (Self-Attention): トークン数 $N$ に対して二次的な複雑さ $O(N^2)$ を持ちます。
   • フィードフォワードネットワーク (FFN): 隠れ次元 $D$ と $D_{mlp}$ に比例する線形変換が行われます。ViT の総 FLOPs の約 61% は FFN 層で消費されており、これがボトルネックとなっています。
2. 既存手法の限界:
   • 構造化プルーニング: 重みのチャネルやヘッドを削除する方法ですが、精度回復のために元のトレーニング時間と同等の長い再学習（リトレーニング）が必要であり、大規模モデルでは非現実的です。また、多くの手法は Attention 機構に焦点を当て、FFN の冗長性を十分に扱えていません。
   • トークン圧縮 (Token Compression): トークン数を減らすことで Attention の計算量を削減しますが、FFN 内のチャネル次元 ($D^2$) の冗長性は解決できず、また層間依存性により最適化が複雑になるという問題があります。

2. 手法 (Methodology)

ToaSt は、ViT の異なるコンポーネントに対して特化した戦略を適用する非結合（Decoupled）フレームワークです。再学習のオーバーヘッドを排除し、層間の依存性を避ける設計が特徴です。

2.1. MHSA に対する構造化結合重みプルーニング (Structured Coupled Weight Pruning)

Multi-Head Self-Attention (MHSA) モジュールに対して、ヘッドごとの内部次元 $d_k$ を削減します。

• 結合された重み行列のプルーニング:
  • Attention 計算において、Query ($W_Q$) と Key ($W_K$)、および Value ($W_V$) と Projection ($W_{proj}$) は数学的に強く結合しています。
  • ToaSt は、$W_Q$ と $W_K$ の対応する列、および $W_V$ と $W_{proj}$ の対応する行を同期して削除します。これにより、アテンション計算の数学的整合性が保たれ、残りの層への入力次元を変更せずに済みます。
• 重要度指標:
  • 事前学習済み重みの幾何中央値 (Geometric Median, GM) を用いて重要度を評価します。分布の中心に近い次元（他の次元で近似可能な冗長な次元）を特定し、優先的に削除します。
• 戦略:
  • 最初の層（パッチ埋め込みと特徴のインターフェース）は保護し、それ以降の層で aggressive なプルーニング（最大 90%）を適用します。
  • ヘッドごとのプルーニング比率を統一し、ハードウェアでの効率的な行列乗算を可能にします。

2.2. FFN に対するトークンチャネル選択 (Token Channel Selection: TCS)

FFN モジュール（$D \to 4D \to D$）の冗長性を、再学習なしで削減します。

• FFN の冗長性分析:
  • 事前学習済みモデルの分析から、深い層において以下の 3 つの現象が確認されました。
    1. 高い線形再構成忠実度 ($R^2$): 特定のチャネルの活性化が、他のチャネルの線形結合で高精度に再構成可能（$R^2 \approx 1.0$）。
    2. 有効ランクの低下 (Collapsing Effective Rank): 深い層では情報の本質的な次元が低次元に収束している。
    3. スパース性の増加: 深い層では多くのニューロンがほぼゼロの活性化を示す（「死んだニューロン」）。
• トレーニングフリーな選択プロセス:
  • 統計的サンプリング: 全トークンの活性化を計算する代わりに、ランダムにサンプリングした少量のトークン（2%〜20%）のみを用いてチャネル重要度を推定します。高い線形依存性により、この少量のサンプルでグローバルな重要度分布を正確に推定できます。
  • 注意ガイドされた重要度スコア: CLS トークンの活性化と、パッチ間の注意重み ($A_{cls, i}$) を組み合わせたスコアを計算し、チャネルを評価します。
  • 層適応型プルーニング:
    • FC1 (拡張層): 初期層では保守的に、深い層では中程度にプルーニング。
    • FC2 (縮小層): 深い層では冗長性が高いため、最大 90% まで aggressive にプルーニング。
  • 構造化削減: 特定のチャネル全体（FC1 の列と FC2 の行）を削除することで、疎行列ではなく密行列構造を維持し、標準 GPU での高速化を実現します。

3. 主要な貢献 (Key Contributions)

1. MHSA 用の構造化結合プルーニング:
   • 重み行列間の制約（Q-K, V-Proj の同期）を厳密に守ることで、アテンション機構の整合性を保ちつつ、レイヤー非依存の圧縮を実現しました。
2. FFN 用のトレーニングフリーなトークンチャネル選択 (TCS):
   • 深い層における冗長性（高い $R^2$、低い有効ランク、高いスパース性）を分析し、これらを利用した層適応型のチャネル選択手法を提案しました。これにより、従来の重みプルーニングに必要な高コストな再学習を不要にしました。
3. 広範なモデルとタスクでの優れた性能:
   • DeiT, ViT-MAE, Swin Transformer などの 9 種類のモデルで検証され、精度と効率のトレードオフにおいて既存の最良手法を上回りました。
   • 大規模モデルほど微調整（Fine-tuning）の必要 epochs が少なくなるという逆相関を発見しました。

4. 実験結果 (Results)

ImageNet-1K 分類タスクおよび COCO 物体検出タスクでの評価結果は以下の通りです。

• ImageNet-1K 分類:
  • ViT-MAE-Huge: 39.4% の FLOPs 削減に対し、Top-1 精度が 88.52%（ベースラインより +1.64% 向上）を達成。微調整はわずか 15 エポックで完了しました。
  • DeiT-Small: 45.7% の FLOPs 削減で、精度が 83.40%（+3.58% 向上）となり、H100 GPU 上でのスループットは 2.07 倍 向上しました。
  • 既存のトークン圧縮手法（ToMe, DiffRate など）と比較して、同等の FLOPs 予算において 1〜4% 高い精度を達成しました。
• COCO 物体検出:
  • Swin-Base を用いた Cascade Mask R-CNN において、圧縮モデルは 52.2 mAP を達成し、圧縮前のベースライン (51.9 mAP) を上回りました。
  • これは、TCS がノイズとなる冗長な特徴を除去し、タスクに不可欠な識別特徴を保持していることを示しています。
• アブレーション研究:
  • MHSA プルーニング単独では精度が低下しますが、TCS を組み合わせることで精度が回復し、さらにベースラインを上回る結果となりました。これにより、両手法が相補的であることが証明されました。

5. 意義と結論 (Significance & Conclusion)

ToaSt は、ViT の実用化における「計算コスト」と「精度」のジレンマを解決する画期的なアプローチです。

• 再学習コストの劇的削減: 従来の構造化プルーニングが抱えていた「数ヶ月単位の再学習」という障壁を、トレーニングフリーな TCS と効率的な微調整により克服しました。
• 大規模モデルへの適応性: 大規模なファウンデーションモデルほど内部冗長性が高く、ToaSt の圧縮による恩恵が大きいことを実証しました。
• 汎用性と実用性: 単一のアーキテクチャに限定されず、DeiT、MAE、Swin など多様な構造に対応し、下流タスク（物体検出）でも有効性が確認されました。

将来的には、層ごとのプルーニング比率の自動学習や、ビジョン・ランゲージモデルへの拡張、量子化との組み合わせなどが期待されています。本手法は、リソース制約のある環境（モバイル、エッジデバイス）での高性能 ViT の展開を可能にする重要なステップです。