Vision Transformerで“軽くて速い”モデルを作るには？

画像分析で人気のViT（Vision Transformer）。でも「重い・遅い・高コスト」が悩みどころ。今日は、パッチサイズの最適化と位置埋め込みの工夫で、ViTを軽くする実践Tipsを紹介します。

なぜパッチサイズが重要？

ViTは画像をパッチ（例：16×16）に分割し、各パッチを“単語”のように処理します。

• 小さいパッチ（8×8）：細かい情報が取れるが、トークン数が増えて計算コストが爆増
• 大きいパッチ（32×32）：情報は粗くなるが、処理は軽くて速い

ポイントは、解像度 × パッチサイズ = トークン数で、Attentionの計算量はトークン数²に比例すること。
たとえば：224x224 の画像なら

• 16×16パッチ → 14×14 = 196トークン
• 8×8パッチ → 28×28 = 784トークン（約16倍の計算量）

位置情報はどうする？

CNNと違い、Transformerはピクセルの“位置”を知らないため、**位置埋め込み（Positional Embedding）**が必須です。

よくある工夫2つ：

1. Learnable Embedding：各パッチ位置ごとに学習する。デフォルト方式。
2. 相対位置（Relative Position）：距離や方向だけで位置関係を学習。画面サイズが変わっても対応しやすいのが利点。

技術Tips（PyTorch）

パッチ数 N が大きくなると、ViTのSelf-Attentionが重くなるのは、次の式から明らか：

attn = (Q @ K.transpose(-1, -2)) / (dk ** 0.5)
weights = attn.softmax(dim=-1)
out = weights @ V

• Q, K, V の形：(B, H, N, d_k)
• K.transpose(-1, -2)：キーの次元を転置して、QK^T を計算
• dk ** 0.5：スケーリング（softmaxが安定）
• dim=-1：各トークンが他のトークンとの関係を見て重みを学習

計算量は O(N²)

実務での選択

• 異常検知など高精度が求められる場面→ 小パッチ（8×8〜14×14）
• 分類・物体検出など全体が見えればOK→ 中〜大パッチ（16×16〜32×32）
• Edge AI/組込機器では、パッチサイズと位置埋め込みでモデルを半分以下の計算量に圧縮できることも

まとめ

以下が結論として、軽くて速いViTモデルを作る際の具体的なチェックリストです：

1. パッチサイズの調整
   • タスクに応じて 小～大パッチ（8×8～32×32） を選ぶ。
   • トークン数を減らすとAttentionコストが劇的に下がる。
2. 位置埋め込みの簡易化
   • 相対位置を使うと学習が軽量化可能。
   • 必要ない場合は Learnable Embedding を省略し、単純なサイン・コサインのみでもOK。
3. モデル規模の最適化
   • 層深（Transformerブロック数）や ヘッド数 を減らす。
   • 小型モデル（DeiT-Tiny / DeiT-Small）をベースに選択。
4. 畜学習と蒸留
   • 大きなモデルで学習した後、知識蒸留で軽量モデルへ性能を転移。
5. 量子化・プルーニング
   • 8bit 混合精度（AMP）や 量子化、不要な重みのプルーニングで計算量を圧縮。
6. ハードウェア最適化
   • TensorRT / ONNX Runtime でモデルを最適化。
   • レイヤフュージョンや バッチサイズ調整も効果的。

明日は「自己教師あり学習のSimCLRを異常検知に応用する例」を紹介します！

ViTの性能は、**「何をどこまで見たいか」**でパッチサイズ・位置表現を調整するのがコツ。軽くしたいなら、パッチ数を減らす／位置埋め込みを簡易化するだけでも大きな効果があります。