Token Merging (ToMe)

Token merging（ToMe）是一种通过在网络内部合并相似token来加速 Transformer 的方法，尤其适用于视觉 Transformer（ViT），这样每个块处理的 token 就会减少。图像中有很多冗余内容，例如外观几乎相同的天空或道路图像块；将这些 token 求平均并合并，可以在保留信息的同时减少计算量。ToMe 可用于已经训练好的模型，无需重新训练，也可以在训练期间应用。Georgia Tech 和 Meta AI 的 Daniel Bolya 等人在 Token Merging: Your ViT But Faster 中介绍了这项技术，该论文在 ICLR 2023 上作口头报告。开箱即用时，它使 512 像素的 ViT-L 和 518 像素的 ViT-H 吞吐量翻倍，准确率仅下降 0.2% 到 0.3%；它还使 ViT-L 视频模型提速 2.2 倍，准确率下降幅度相近。

Token merging 的工作原理#

ViT 会将图像拆分成图像块，每个图像块对应一个 token；每个块先执行自注意力，其计算成本随 token 数量的平方增长，再执行 MLP，其计算成本随 token 数量线性增长。ToMe 会在每个块的注意力层与 MLP 之间插入合并步骤，每次移除固定数量 r 的 token，因此 24 个块、r = 8 的模型最终会比初始少 192 个 token。合并采用二分图软匹配，因为它速度快且不需要迭代聚类：

1. 拆分： 将 token 交替划分为两个集合 A 和 B。
2. 匹配： 将 A 中的每个 token 与 B 中最相似的 token 相连，相似度根据跨注意力头平均后的注意力键的余弦相似度计算。
3. 保留最佳边： 保留最相似的 r 对。
4. 合并： 将每个保留的 token 对按各自所代表的原始图像块数量加权求平均，合并为一个 token，然后重新拼接两个集合。

由于合并后的 token 代表多个图像块，ToMe 会将每个 token 大小的对数加到注意力分数中，这称为比例注意力；这样，由许多图像块构成的 token 就能获得相应的注意力份额。整个过程可微，因此可以使用 ToMe 训练模型以进一步提速。

这个 PyTorch 示例将二分图软匹配应用于 Meta 的 DINOv2-small ViT。该模型来自 Hugging Face Transformers，输入是一张 448 x 448 的公交车图像；示例会在 12 个块之后分别合并 r 个 token，并将得到的图像嵌入与未合并模型的结果进行比较：

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 个 token，即 1,024 个图像块加上 CLS token。每块合并 32 个 token 后，平均 token 数降至 849，而图像嵌入与原始结果的余弦相似度为 0.983；每块合并 64 个 token 后，token 数降至 673，相似度为 0.887，这展示了 r 如何在计算量与保真度之间进行权衡。

Token merging 的应用场景#

Token Merging for Fast Stable Diffusion将相同思路应用于 Stable Diffusion 去噪网络的 Transformer 块中，无需额外训练即可移除最多 60% 的 token，使图像生成速度最高提升两倍、内存占用最高减少 5.6 倍。除了图像之外，ToMe 论文还报告称，在使用合并技术训练的音频 ViT-B 上，吞吐量提升了一倍，mAP 下降 0.4%；在视频上进行掩码自编码器微调时，速度最高提升 2 倍。研究人员也会对视觉语言模型的视觉 token 采用类似的合并和剪枝方法；例如，LLaVA-1.5 每张图像会占用语言模型上下文中的 576 个 token。

在 Ultralytics 包中，Meta 的 Segment Anything Model (SAM) 会将每张 1024 x 1024 像素的图像编码为 4,096 个图像块 token；在 MobileSAM 对比中，ViT-H 图像编码器占用了 SAM 在 GPU 上处理每张图像所需 456 毫秒中的 452 毫秒，这正是 token 缩减方法要针对的计算量。

Token merging 与相关概念#

• Token 剪枝： DynamicViT 和 EViT 等方法会丢弃被判定为不重要的 token，这会损失信息，且通常需要训练。ToMe 会合并 token，而不是丢弃它们，因此信息会被平均保留下来；它也适用于未经 ToMe 训练的模型。
• Token 化： 在网络运行之前，从图像块或文本生成初始 token。ToMe 会在网络内部减少 token 数量。
• Flash attention 和稀疏注意力： 在 token 数量不变的情况下，让注意力计算更省资源或更快。ToMe 会减少 token 数量，也能加速 MLP 层，并且可以与这两种方法结合使用。
• 深度混合（MoD）： 让学习到的路由器决定语言模型中的哪些 token 跳过某个块。ToMe 则使用固定的相似度规则，永久融合冗余 token。
• 降低输入分辨率： 均匀减少整张图像中的 token 数量。ToMe 主要在背景等冗余区域合并 token，并让细节丰富的区域保持完整分辨率。