AI・機械学習
上級
ToMe(Token Merging)(トークンマージング)
Vision Transformerの中間層で類似したトークンを統合し、計算量を削減するViT推論高速化手法。Meta AI発表のToken Mergingは精度劣化を最小化しながら処理速度を2倍以上に向上させる。
0 回閲覧
0 いいね
2026/6/7 更新
関連タグ
LLM推論
Token Merging
Vision Transformer
推論高速化
Meta AI
ToMe(Token Merging)— 類似トークン統合による推論高速化
ToMe(Token Merging)は2022年にMeta AIのBolya et al.が発表したVision Transformer(ViT)の推論高速化手法だ。Transformerの各レイヤーで類似した視覚トークンを統合(マージ)することで、後続レイヤーの計算量を削減し、精度をほぼ損なわずに処理速度を向上させる。
基本アイデア
Vision Transformerは画像を16x16ピクセルのパッチに分割し、各パッチを「トークン」として処理する。例えば224x224の画像は196個のトークンになる。しかし実際の画像には:
- 空: 均一な青で隣接トークンがほぼ同一
- 背景: 類似テクスチャが繰り返される
- 余白: 情報が少ない領域
など、意味的に冗長なトークンが多数存在する。ToMeはこれを活用する。
アルゴリズム詳細
マージの基準: 各Transformerレイヤーでアテンションのキー(K)行列を使い、トークン間のコサイン類似度を計算。上位r件の類似ペアを選択し、それらを1つのトークンに統合(平均化)する。
2段階処理:
- Bipartite Soft Matching: トークンを「source」と「destination」の2グループに交互分割し、source内の各トークンに最も似ているdestinationトークンを見つけてマージ
- Unmerge(必要時): 出力を元のトークン数に戻す場合は、マージ先の情報を分配
実装コスト: 既存のViT実装へのパッチは約30行、追加の学習は不要(Training-free)。
効果
ViT-H(Vision Transformer Huge)での実験:
| 設定 | スループット向上 | ImageNet Top-1精度低下 |
|---|---|---|
| r=8 per layer | 1.5倍 | -0.1% |
| r=16 per layer | 2.0倍 | -0.3% |
| r=32 per layer | 3.1倍 | -1.2% |
LLMへの拡張
ToMeはViT向けに設計されたが、テキストのLLMへの適用(LLM-ToMe)も研究されている。テキストでも意味的に冗長なトークン(助詞・冗長な接続詞等)が存在するが、言語の離散性から画像ほど単純ではない。後続研究のH2O・PyramidKV等がLLMのKVキャッシュ圧縮という形で同様の思想を発展させた。