Token Merging (ToMe)

トークンマージ（ToMe）は、ネットワーク内で類似するトークンを結合して各ブロックの処理数を減らし、Transformer、特にビジョントランスフォーマー（ViT）を高速化する手法です。画像には、ほとんど同じに見える空や道路のパッチなど、多くの冗長性があります。このようなトークンを平均してまとめることで、情報を保ちながら計算量を減らします。ToMeは再学習なしで学習済みモデルに適用でき、学習中にも適用できます。Georgia TechとMeta AIのDaniel Bolyaらは、ICLR 2023の口頭発表論文「Token Merging: Your ViT But Faster」でこの手法を発表しました。そのまま適用した場合、512ピクセルのViT-Lと518ピクセルのViT-Hでスループットが2倍になり、精度低下はわずか0.2～0.3%でした。また、ViT-Lの動画モデルを同程度の精度低下で2.2倍高速化しました。

トークンマージの仕組み#

ViTは画像をパッチに分割し、各パッチを1つのトークンにします。各ブロックではセルフアテンションを実行し、その計算量はトークン数の2乗に比例して増加します。続くMLPの計算量はトークン数に比例します。ToMeは各ブロックのアテンションとMLPの間にマージ処理を挿入し、毎回一定数rのトークンを削除します。そのため、r = 8の24ブロックモデルでは、最終的に開始時より192個少ないトークンになります。マージには、処理が高速で反復クラスタリングが不要な二部グラフのソフトマッチングを使います：

1. 分割：トークンを交互に2つの集合A、Bに分けます。
2. マッチング：A内の各トークンを、B内で最も類似するトークンに接続します。類似度は、ヘッド間で平均したアテンションキーのコサイン類似度で測定します。
3. 最良のエッジを選択：類似度が最も高いr組を残します。
4. マージ：選択した各ペアを1つのトークンに平均統合します。重みには各トークンがすでに表す元のパッチ数を使い、最後に2つの集合を再度連結します。

マージ後のトークンは複数のパッチを表すため、ToMeは各トークンのサイズの対数をアテンションスコアに加算します。比例アテンションと呼ばれるこの処理により、複数のパッチから構成されるトークンが相応のアテンションを受け取ります。全体の処理は微分可能なので、ToMeを使ってモデルを学習し、さらに高速化できます。

このPyTorchの例では、448 x 448のバス画像を使い、Hugging Face TransformersのMeta製DINOv2-small ViTに二部グラフのソフトマッチングを適用します。12個の各ブロックの後にr個のトークンをマージし、得られた画像埋め込みを、マージを行わないモデルの埋め込みと比較します：

from urllib.request import urlopen

import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoImageProcessor, AutoModel

processor = AutoImageProcessor.from_pretrained("facebook/dinov2-small")
model = AutoModel.from_pretrained("facebook/dinov2-small")
image = Image.open(urlopen("https://ultralytics.com/images/bus.jpg")).resize((448, 448))
pixels = processor(images=image, do_resize=False, do_center_crop=False, return_tensors="pt").pixel_values


def merge(x, r):
    """Bipartite soft matching: merge the r most similar pairs between alternating token sets A and B."""
    cls, a, b = x[:, :1], x[:, 1::2], x[:, 2::2]
    similarity = F.normalize(a[0], dim=-1) @ F.normalize(b[0], dim=-1).T
    best, match = similarity.max(dim=-1)  # each A token's most similar B token
    order = best.argsort(descending=True)
    merged, kept = order[:r], order[r:]
    index = match[merged][None, :, None].expand(1, r, x.shape[-1])
    b = b.scatter_reduce(1, index, a[:, merged], reduce="mean")  # average merged A tokens into their B matches
    return torch.cat([cls, a[:, kept], b], dim=1)


with torch.no_grad():
    for r in (0, 32, 64):
        x, counts = model.embeddings(pixels), []
        for layer in model.encoder.layer:
            counts.append(x.shape[1])
            x = merge(layer(x), r)  # merge r tokens after every transformer block
        embedding = model.layernorm(x)[:, 0]  # CLS token: the image embedding
        reference = embedding if r == 0 else reference
        similarity = F.cosine_similarity(embedding, reference).item()
        print(f"r={r}: mean tokens per block {sum(counts) / 12:.0f}, CLS similarity to unmerged {similarity:.3f}")

マージなしのモデルは、各ブロックで1025個のトークンを処理します。これは1,024個のパッチとCLSトークンです。ブロックごとに32個のトークンをマージすると、平均トークン数は849に減りますが、画像埋め込みは元の埋め込みと0.983のコサイン類似度を保ちます。ブロックごとに64個をマージすると、類似度は0.887でトークン数が673まで減り、rによって計算量と忠実度がどのようにトレードオフになるかが分かります。

トークンマージの用途#

高速Stable Diffusion向けトークンマージは、Stable DiffusionのデノイズネットワークにあるTransformerブロックに同じ手法を適用します。追加学習なしでトークンを最大60%削減でき、画像生成を最大2倍高速化し、メモリ使用量を最大5.6分の1にできます。画像以外では、ToMe論文で、マージを使って学習したオーディオViT-Bのスループットが2倍になり、mAPの低下は0.4%でした。また、動画のマスク付きオートエンコーダーのファインチューニングを最大2倍高速化しました。研究者は視覚言語モデルの視覚トークンにも同様のマージや枝刈りを適用しています。たとえばLLaVA-1.5では、画像ごとに言語モデルのコンテキストを576トークン消費します。

Ultralyticsパッケージでは、MetaのSegment Anything Model（SAM）が各画像を1024 x 1024ピクセルでエンコードし、4,096個のパッチトークンを生成します。MobileSAMの比較では、GPU上でSAMが画像1枚あたりに必要とする456 msのうち、ViT-H画像エンコーダーが452 msを占めます。これは、トークン削減手法が対象とする計算量の割合です。

トークンマージと関連概念#

• トークン枝刈り：DynamicViTやEViTなどの手法は、重要でないと判断したトークンを破棄します。その結果、情報が失われ、通常は学習も必要です。ToMeはトークンをマージするため、情報を平均して取り込み、ToMeを使わずに学習したモデルにも適用できます。
• トークナイズ：ネットワークの実行前に、パッチやテキストから初期トークンを作成します。ToMeはネットワーク内でトークン数を減らします。
• Flash attentionとスパースアテンション：トークン数を維持したまま、アテンションを低コスト化または高速化します。ToMeはトークン数を減らし、MLP層も高速化するため、これらの手法と組み合わせられます。
• Mixture of Depths（MoD）：言語モデルで、学習済みルーターがどのトークンをブロックからスキップさせるか決定します。ToMeは固定の類似度ルールを使って、冗長なトークンを恒久的に統合します。
• 入力解像度の低下：画像全体のトークン数を一様に減らします。ToMeは背景などの冗長な領域を主にマージし、細部の多い領域は元の解像度で保持します。