Token Merging (ToMe)

Token Merging (ToMe) は、フォワードパス中に処理されるトークン数を削減することで、Transformer アーキテクチャのパフォーマンスと効率を最適化するために設計された最先端の技術です。もともとは Vision Transformer (ViT) モデルを高速化するために開発されたToMeは、ネットワーク内の冗長なトークンを体系的に特定し統合することで機能し、追加のトレーニングを必要としません。self-attention メカニズムの計算複雑性はトークン数の2乗でスケールするため、類似したトークンをマージすることで総浮動小数点演算数（FLOPs）が劇的に削減され、real-time inference を大幅に高速化できます。

Link to this sectionToken Merging プロセスの理解#

ToMe is fundamentally different from tokenization, which is the initial preprocessing step of breaking down an image or text into individual tokens. While tokenization creates the discrete elements, Token Merging acts as a dynamic downsampling mechanism during the model's forward execution.

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

Link to this sectionToken Merging の実世界における応用#

Token Merging は、計算資源が制限された環境で大規模なアテンションベースのアーキテクチャをデプロイするための重要な最適化戦略となっています。

1. Generative AI and Image Synthesis: 一般的なテキストから画像を生成する拡散モデルでは、ToMeは画像生成を高速化するためによく使用されます。背景や詳細度の低いトークンをマージすることで、生成プロセスに必要なステップ数が減り、GPUリソースを大幅に節約するとともに、生成モデルを利用するエンドユーザーのレイテンシを削減できます。拡散プロセスに関する詳細は、arXiv の基礎研究を参照してください。

2. Edge AI Deployments: Segment Anything Model (SAM) のような巨大なモデルをモバイルデバイスにデプロイすることは、メモリ制限のために非常に困難です。ToMeはメモリフットプリントを動的に削減し、複雑な image segmentation タスクをエッジハードウェアで実行できるようにします。純粋な速度が重要なシナリオでは、エンジニアはより高速なエンドツーエンドのエッジ推論のために、Ultralytics YOLO26 のような、ネイティブに最適化されたアテンションフリーのアーキテクチャを採用することが一般的です。

Link to this sectionPython の例: トークン類似度の計算#

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 を最適化するために Token Merging を採用する場合でも、YOLO26 の最先端の効率に頼る場合でも、これらの複雑なデータワークフローの管理は Ultralytics Platform によって大幅に簡素化されます。このプラットフォームは、自動化された data annotation、シームレスなクラウドトレーニング、および多様な edge computing ハードウェア環境全体での堅牢なモデルデプロイのための直感的なエコシステムを提供します。computer vision イニシアチブを拡大する組織は、これらのツールを利用して、最先端のモデルを確実かつ効率的に本番環境へ移行しています。