今回の論文
今回取り上げるのは、Daniel Bolya、Cheng-Yang Fu、Xiaoliang Dai、Peizhao Zhang、Christoph Feichtenhofer、Judy Hoffman による論文「Token Merging: Your ViT But Faster」です。初出は 2022 年 10 月の arXiv です。研究分野はコンピュータビジョン、Transformer 高速化、推論最適化です。URL は https://doi.org/10.48550/arXiv.2210.09461 です。
この論文を選んだ理由は、既存の Vision Transformer を作り直さずに高速化する、かなり実務的な発想だからです。大きいモデルをより安く回したい場面は、画像分類だけでなく動画理解、音声処理、マルチモーダル、生成モデル周辺でも増えています。そのとき「重要でないトークンを消す」のではなく「似たトークンをまとめる」という考え方は、推論最適化の引き出しとしてかなり応用範囲があります。
どんな技術か
Token Merging は、Vision Transformer に入ってくる多数のパッチトークンのうち、互いに似たものを途中の層で少しずつ統合していく技術です。論文中では ToMe と略されています。
Vision Transformer は画像を小さなパッチに分け、それぞれをトークンとして self-attention に流します。しかし画像の中には、空、壁、芝生、服の模様のように、近い情報を何度も持つ領域があります。ToMe は、そうした冗長なトークンを毎層少しずつまとめ、トークン数を減らしながら意味をできるだけ保つようにします。
重要なのは、単にトークンを捨てるわけではないことです。似たトークンを平均的に統合するので、情報を完全に失うのではなく、情報密度を上げながら計算量を下げる設計になっています。これが「精度低下を抑えながら速くする」ポイントです。
課題
この技術が解決しようとしているのは、Transformer 系の視覚モデルが高性能である一方、トークン数に応じて計算コストが膨らみやすいという課題です。
何が難しいのかというと、Vision Transformer では self-attention の計算がトークン間の相互作用に依存するため、入力パッチ数が多いほど計算量とメモリ使用量が重くなりやすいからです。高解像度画像、長い動画、音声スペクトログラムのように入力トークンが増えるタスクでは、この負担がそのまま推論コストになります。
既存の方法では、軽量アーキテクチャに作り替える、注意機構そのものを近似する、不要そうなトークンを pruning で削る、といった方向が中心でした。ただし pruning は、削ったトークンの情報をそのまま失いやすく、学習し直しが必要だったり、入力ごとに削減数が変わってバッチ推論しづらかったりします。
なぜこの課題を解く必要があるのかというと、実際の AI システムでは、1 回の推論性能よりも「大量リクエストをどれだけ安く回せるか」が重要になるからです。画像検索、動画解析、監視、製造検査、マルチモーダル検索、生成パイプラインの前処理などでは、数パーセントの精度差より、レイテンシと GPU コストのほうが導入可否に直結することも多いです。
つまり ToMe は、Transformer の表現力を大きく壊さずに、冗長なトークンだけを圧縮することで、実運用可能な速度帯へ近づけようとする技術です。
用語解説
- Vision Transformer（ViT）
- 画像をパッチ列として扱う Transformer です。ToMe はこのパッチトークンの数を途中で減らす技術なので、CNN の特徴マップ圧縮ではなく、トークン列の長さ最適化だと理解することが重要です。
- トークン冗長性
- 入力トークンの一部が、ほかのトークンとかなり近い情報を持っている状態です。ToMe はこの冗長性を利用し、意味的に近いトークンを統合しても性能を大きく落とさない、という前提に立っています。
- Token Pruning
- 不要そうなトークンを途中で削除する高速化手法です。ToMe の比較対象として重要で、削除ではなく統合にすることで情報損失を抑えるのがこの論文の狙いです。
- Self-Attention の Key
- Transformer の attention で類似度計算に使う表現です。ToMe はトークン同士の近さを生の埋め込みではなく key ベクトルで測るため、何を似ているとみなすかの中核になります。
- MAE（Masked Autoencoder）
- 画像や動画でよく使われる自己教師あり事前学習手法です。論文では MAE 事前学習済み ViT に ToMe を適用しており、大きい事前学習済みモデルを後から速くする実務上の価値につながっています。
技術の仕組み
ToMe の基本アイデアは、各 Transformer ブロックの途中で「似ているトークン同士だけを少数ペアでマッチングし、平均統合する」というものです。クラスタリングのように大きくまとめるのではなく、毎層少しずつ圧縮するのが特徴です。
基本アイデア
多くの画像では、隣接パッチや同じ物体の一部がかなり近い特徴を持ちます。ToMe は、こうした重複情報を持つトークンをそのまま全部 attention に残すのは無駄だと考えます。
そこで、各層で r 個ずつトークンを減らすように、似たトークンのペアを選びます。重要なのは、どの画像でも削減数を一定に保つことです。これにより、入力ごとにトークン数がバラバラになる pruning より、バッチ処理しやすい設計になります。
どこで統合するのか
論文では、トークン統合を Transformer ブロックの attention と MLP の間に入れています。ブロックの最初ではなく途中に置くことで、その層の attention が一度計算した情報を使って、どのトークンをまとめるべきか決められます。
これは地味ですが重要な工夫です。先に消してしまうと、そのトークンが持っていた情報を attention に反映できません。ToMe は一度情報伝播したあとに統合するので、単純な前処理より精度低下が小さくなります。
似ているトークンをどう見つけるか
ToMe は、トークンの類似度計算に attention の key 表現を使います。論文の考え方では、key は「このトークンが何に注目されるべきか」を要約した表現なので、単純な特徴ベクトル距離より、統合判断に向いています。
類似度計算そのものは、key 同士の内積ベースで行います。ここで全トークンの完全な組み合わせを毎回比較すると重いので、論文では高速な近似マッチングを使います。
Bipartite Soft Matching
ToMe の中核は bipartite soft matching です。流れはかなり明快です。
- トークン列を交互に 2 グループへ分けます。
- 片方の各トークンについて、もう片方で最も似ている相手を 1 つ探します。
- 類似度が高いペアだけを上位
r個選びます。 - そのペアを平均的に統合し、残りはそのまま残します。
この方法の良い点は、反復的なクラスタリングよりずっと軽いことです。しかも、一気に巨大クラスタを作らず、局所的で保守的な統合に留めるので、モデルの振る舞いを壊しにくいです。
統合後の重み付け
2 つ以上のパッチを表すトークンは、元の 1 パッチトークンより情報量が大きくなります。そのため論文では、統合されたトークンの「サイズ」を追跡し、attention の softmax に反映させる proportional attention を導入しています。
直感的には、まとめたトークンを 1 個の普通のトークンとして扱うと、存在感が過小評価されるので、その分だけ重みを補正する処理です。ここを入れないと、統合後の attention 分布が不自然になりやすく、精度低下につながります。
学習方法と推論方法
ToMe は 2 通りの使い方ができます。1 つは、既存の学習済み ViT にそのまま差し込む「学習なし適用」です。もう 1 つは、ToMe を入れた状態で fine-tuning や学習を進める方法です。
学習なし適用の価値は、巨大な事前学習済みモデルを後から速くできることです。一方、ToMe ありで学習すると、統合によるわずかな精度低下をさらに小さくできます。論文では、この方式で訓練時間の短縮も確認しています。
実験と結果
論文では、ToMe が画像だけでなく、動画や音声でも有効かを検証しています。単一タスク専用の小技ではなく、「トークン冗長性があるモダリティ全般」に効くかを見る構成です。
何を検証したのか
主な検証ポイントは 3 つです。1 つ目は、学習済み ViT に後付けしても速度改善が出るか。2 つ目は、精度低下をどこまで小さく抑えられるか。3 つ目は、画像以外の動画・音声にも同じ考え方が通用するかです。
加えて、token pruning 系の既存手法と比べて、速度・実装の単純さ・学習時の扱いやすさで優位があるかも検証しています。
画像での結果
画像分類では ImageNet-1k を使い、AugReg、MAE、SWAG、DeiT など異なる学習済み ViT で評価しています。本文の比較では、ViT-L 系で 85.7% 精度・93 images/s のモデルに対し、ToMe 適用版が 85.1%・183 images/s を達成しています。ViT-H 系でも 86.9%・35 images/s に対し、ToMe 適用版は 86.1%・81 images/s です。つまり、精度を 0.6〜0.8 ポイント程度の低下に抑えつつ、実効スループットをほぼ 2 倍にできています。
さらに大きいモデルでは、ViT-L や ViT-H を「1 つ下の速度帯」に押し下げられることも示されています。つまり、小さい別モデルへ置き換えなくても、大きいモデルの表現力をある程度保ったまま現実的な速度で動かせるわけです。
pruning 手法との比較
DeiT-S を使った比較では、ToMe は 79.4% 精度、2.7 GFLOPs、1552 images/s で、DynamicViT の 79.3%、1505 images/s と同程度以上の速度精度バランスを示しています。
ここで重要なのは、ToMe が複雑な補助モジュールや特殊な学習 tricks なしにこの性能を出している点です。論文では、pruning 系は学習時に padding やマスク処理が必要になりやすく、訓練高速化にはつながりにくい一方、ToMe は DeiT 学習で約 1.5 倍の訓練高速化も確認しています。
動画での結果
動画では Kinetics-400 を用い、MAE 事前学習した ViT-L ベースで比較しています。論文要旨では、ToMe により ViT-L のスループットを 2.2 倍にしつつ、精度低下を 0.2〜0.3 ポイント程度に抑えたと報告しています。
動画はフレーム方向にトークンが増えるぶん、画像以上に計算負荷が重くなりやすいです。そのため、ToMe のような系列長圧縮の効果が特に出やすく、学習時の高速化にもつながりやすいことが示唆されます。
音声での結果
音声では Audio MAE ベースの ViT-B を AudioSet-2M で評価しています。論文では「ほぼ 2 倍のスループットで、mAP 低下は 0.4 ポイント程度」とまとめています。
この結果が示すのは、ToMe が画像専用の空間的発想ではなく、スペクトログラムのようなトークン化された時系列表現にも効くということです。冗長なパッチがあるなら、モダリティが変わっても成立しやすいわけです。
何に使える？
ToMe の価値は、モデルを作り替えずに推論コストを圧縮できることです。そのため、大きい ViT 系モデルを使いたいが、そのままでは重すぎる場面で幅広く使えます。
画像分類・検索・検査システムの高速化
製造検査、棚監視、画像タグ付け、類似画像検索のように大量画像をさばく処理では、1 枚あたりのコスト差がそのまま運用費に響きます。ToMe は学習済み ViT を後付け高速化できるので、既存の推論基盤を大きく壊さずにスループット改善を狙えます。
動画解析や監視のバッチ処理
動画はフレーム数のぶんだけトークンが増え、推論費用が急激に上がります。ToMe は動画でも有効性が確認されているため、行動認識、スポーツ解析、監視映像のイベント抽出などで、同じ GPU 予算でより多くの映像を処理したい場面に向いています。
音声スペクトログラム処理
音声イベント検出や環境音分類のように、スペクトログラムを ViT に入れる系ではそのまま応用しやすいです。会議音声、工場音、設備監視などで大量バッチ推論を回す場合、ToMe 的な圧縮は現実的な選択肢になります。
マルチモーダル基盤の前段圧縮
画像エンコーダや音声エンコーダに ViT 系を使うマルチモーダルモデルでは、エンコーダ側のコストが全体のボトルネックになることがあります。ToMe を前段に入れると、後続のクロスアテンションや融合処理へ渡すトークン数自体を減らせるので、全体最適に効きやすいです。
生成AI周辺の高速化
この論文自体は主に分類系評価ですが、ViT バックボーンや画像トークン処理を含む生成パイプラインでも発想は有効です。たとえば画像理解付きエージェントや画像条件付けモデルの前処理で、視覚トークンの冗長性を圧縮する設計は十分考えられます。ここは本文の実験から一歩進んだ応用推測ですが、技術的には自然です。
開発や事業へのヒント
ToMe から得られるヒントは、「モデルを軽くしたいなら、新しい軽量モデルへ乗り換える前に、トークン列そのものを最適化できないかを見る」ということです。
推論最適化はアーキテクチャ刷新だけではない
多くの現場では、精度を維持したまま速くしたいとき、より小さいモデルへ置き換えるか、量子化や蒸留を検討します。もちろん有効ですが、ToMe はその前に「入力の冗長性を消す」という別のレイヤーで最適化できることを示しています。既存モデル資産を活かしたまま改善できる点が強みです。
後付け導入できる設計は事業上強い
再学習なしで導入できる手法は、PoC から本番移行しやすいです。大規模再学習は時間も GPU 予算も重いですが、後付け最適化なら AB テストしやすく、費用対効果の判断も早くできます。既存の画像分類 API や動画解析 SaaS の利益率改善にも直結しやすいです。
小規模チームでも試しやすい
ToMe の中核は、巨大な新モデル開発ではなく、トークン統合モジュールの差し込みです。すでに ViT 系モデルを使っているなら、まずは精度許容ラインを決めて r を調整し、どこまでスループットを上げられるかを見る、という実験がしやすいです。
今後注目すべき方向性
今後は token pruning、token merging、windowing、KV 圧縮のような「系列長そのものを扱う最適化」がさらに重要になります。モデル本体の巨大化が進むほど、各層で何トークンを残すかという問題がコストを左右するからです。LLM でも画像モデルでも、長い系列をどう圧縮するかは共通テーマだと見ておくとよいです。
限界
ToMe にも注意点はあります。まず、どれだけトークンを統合するかで精度と速度のトレードオフが変わるため、万能な設定はありません。重要領域までまとめすぎると、細かい境界や局所特徴が効くタスクでは性能が落ちやすいです。
また、論文では主に分類系で強い結果を示していますが、検出、セグメンテーション、キーポイント推定のように位置精度が重要なタスクでは、どの程度そのまま効くかを個別検証する必要があります。トークン統合は空間分解能を下げる側面があるからです。
実装面では、Transformer の中間表現に手を入れるため、利用中のフレームワークや推論基盤によっては差し込みが簡単でないこともあります。特に最適化済み推論エンジンでは、カスタム処理の導入コストを見積もる必要があります。
さらに、ToMe は冗長性があることを前提に効く手法です。入力が最初から非常に情報密度の高い場合や、細部の差が重要な場合は、統合のメリットが小さいか、悪影響が大きい可能性があります。
よくある質問
Q. Token pruning と Token Merging は何が違うのですか？
A. pruning は不要そうなトークンを捨てますが、ToMe は似たトークンをまとめます。つまり、情報をゼロにするのではなく、圧縮して残す発想です。そのため、同じ削減率でも精度低下を抑えやすいのが利点です。
Q. 学習し直さなくても本当に使えますか？
A. 論文の重要なポイントの 1 つがそこです。既存の学習済み ViT に後付けしても大きな速度改善が出ています。ただし、実運用では許容できる精度低下がタスクごとに違うので、r の調整とベンチマークは必要です。
Q. どんなタスクで特に効きやすいですか？
A. トークン数が多く、しかも入力に冗長な領域が多いタスクで効きやすいです。高解像度画像、動画、スペクトログラム入力などは典型です。逆に、細部の局所差が極めて重要なタスクでは慎重な検証が必要です。
Q. LLM にも同じ考え方は使えますか？
A. そのまま同一実装を入れるのは簡単ではありませんが、「系列中の冗長な要素を統合して計算量を下げる」という考え方自体は共通です。長文コンテキスト圧縮や KV 圧縮、要約トークン導入など、近い発想の最適化は LLM 側でも重要です。
Q. まず実務で試すなら何から始めるべきですか？
A. 既存の ViT 系推論に対し、精度許容範囲を先に決めたうえで、小さい r から段階的に増やすのが安全です。速度だけでなく、難例サンプルや境界ケースで何が壊れるかを見ながら決めるのが実務向きです。
今日の学び
この論文は、Vision Transformer が高性能な反面、トークン数の多さで推論コストが重くなる課題を扱いました。これに対して、似たトークンを各層で少しずつまとめる Token Merging によって、情報をできるだけ保ちながら計算量を減らす方法を示しました。
ここから得られるヒントは、推論最適化は量子化や小型化だけではないということです。入力系列の冗長性をどう圧縮するかを設計できれば、既存モデル資産を活かしたまま、開発コストと運用コストの両方を下げられます。