Token Merging (ToMe)

トークンマージ (ToMe) は、フォワードパス中に処理するトークン数を削減することで、Transformerアーキテクチャのパフォーマンスと効率を最適化するために設計された最先端の技術です。もともとビジョントランスフォーマー (ViT)モデルを高速化するために開発されたToMeは、追加のトレーニングを必要とせず、ネットワーク内の冗長なトークンを体系的に特定して結合します。セルフアテンションメカニズムの計算量はトークン数に対して二次関数的に増加するため、類似したトークンをマージすると浮動小数点演算 (FLOPs) の総数が大幅に減少し、リアルタイム推論を大幅に高速化できます。

トークンマージプロセスの概要#

ToMeは、画像やテキストを個々のトークンに分割する初期前処理ステップであるトークン化とは根本的に異なります。トークン化が離散的な要素を作成するのに対し、トークンマージはモデルのフォワード実行中に動的なダウンサンプリングメカニズムとして機能します。

このアルゴリズムでは通常、二部グラフマッチングを使用してトークンの類似度を評価し、アテンションレイヤー内のトークンのキー間のコサイン類似度を計算します。視覚的または意味的に非常に類似した情報を共有するトークンは、多くの場合、それらの特徴量を平均することで融合されます。これにより、重要な空間情報やコンテキスト情報を保持しながら、不要な計算負荷を削減できるため、PyTorchなどのフレームワークで複雑なビジョンモデルをはるかに高速に処理できます。

トークンマージの実際の応用#

トークンマージは、計算リソースに制約のある環境で、大規模なアテンションベースのアーキテクチャをデプロイするための重要な最適化戦略になっています。

1. 生成AIと画像合成: 人気の高いテキストから画像を生成する拡散モデルでは、画像生成を高速化するためにToMeが頻繁に使用されます。背景やディテールの少ないトークンをマージすることで、生成プロセスに必要なステップ数が減少し、GPUリソースを大幅に節約するとともに、生成モデルを利用するエンドユーザーのレイテンシを短縮できます。拡散プロセスについては、arXivの基礎研究で詳しく学習できます。

2. エッジAIのデプロイ: セグメント・エニシング・モデル (SAM)のような大規模モデルをモバイルデバイスにデプロイすることは、メモリ制約のため非常に困難です。ToMeはメモリフットプリントを動的に縮小し、複雑な画像セグメンテーションタスクをエッジハードウェア上で実行できるようにします。純粋な速度が重要なシナリオでは、エンジニアは高速なエンドツーエンドのエッジ推論を実現するため、Ultralytics YOLO26のようなネイティブに最適化されたアテンションフリーアーキテクチャへ移行することがよくあります。

Pythonの例: トークン類似度の計算#

ToMeを完全なアーキテクチャに統合するにはアテンションブロックの変更が必要ですが、核となる概念は類似したトークンを見つけることにあります。次のPyTorchスニペットでは、トークンの候補を特定するために、トークンの集合間のコサイン類似度を計算する方法を示します。

import torch
import torch.nn.functional as F

# Simulate a batch of 4 image patches (tokens) with 64-dimensional features
tokens = torch.randn(1, 4, 64)

# Normalize the tokens to easily compute cosine similarity via dot product
normalized_tokens = F.normalize(tokens, p=2, dim=-1)

# Compute the similarity matrix between all tokens (1 x 4 x 4)
similarity_matrix = torch.matmul(normalized_tokens, normalized_tokens.transpose(1, 2))

# Tokens with high similarity scores (close to 1.0) off the diagonal
# are prime candidates for Token Merging.
print("Similarity Matrix:", similarity_matrix)

最新の機械学習パイプラインでは、精度と速度の慎重なバランス調整が必要です。カスタムViTの最適化にトークンマージを使用する場合でも、YOLO26の最先端の効率性を活用する場合でも、こうした複雑なデータワークフローの管理はUltralytics Platformによって大幅に簡素化されます。このPlatformは、自動化されたデータアノテーション、シームレスなクラウドトレーニング、多様なエッジコンピューティングハードウェア環境に対応した堅牢なモデルデプロイのための直感的なエコシステムを提供します。コンピュータービジョンの取り組みを拡大する組織は、最先端モデルを確実かつ効率的に本番環境へ導入するために、これらのツールを活用しています。