Token Merging (ToMe)
学习 Token 合并 (ToMe) 如何优化 Transformer 和 ViT 模型。发现如何减少 FLOPs、加速实时推理并提升生成式 AI 速度。
Token Merging (ToMe) 是一项尖端技术,旨在通过减少前向传递过程中处理的标记数量来优化 Transformer 架构的性能和效率。ToMe 最初是为了加速 Vision Transformer (ViT) 模型而开发的,它通过系统地识别并合并网络中的冗余标记来工作,而无需任何额外训练。由于 self-attention 机制的计算复杂度随标记数量呈二次方扩展,合并相似标记能显著减少总浮点运算 (FLOPs),从而实现更快的 real-time inference。
Link to this section了解 Token Merging 流程#
ToMe 与 tokenization 有本质区别,后者是将图像或文本分解为单个 tokens 的初始预处理步骤。虽然 tokenization 创建了离散元素,但 Token Merging 在模型的前向执行过程中充当一种动态降采样机制。
该算法通常使用二分匹配来评估标记相似度,通常通过计算注意力层中标记的键(keys)之间的余弦相似度来实现。具有高度相似视觉或语义信息的标记会被合并——通常通过对它们的特征取平均值。这确保了在丢弃不必要的计算负载的同时保留了基本的空间或上下文信息,从而使 PyTorch 等框架能够更快地处理复杂的视觉模型。
Link to this sectionToken Merging 的实际应用#
Token Merging 已成为在计算受限环境中部署大型注意力架构的关键优化策略。
-
Generative AI and Image Synthesis:在流行的文本到图像扩散模型中,ToMe 经常被用来加速图像生成。通过合并背景或低细节标记,生成过程所需的步骤更少,节省了巨大的 GPU 资源,并降低了依赖生成模型的最终用户的延迟。你可以在 arXiv 的基础研究中了解更多关于扩散过程的信息。
-
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)现代机器学习流水线需要仔细平衡准确性和速度。无论你是使用 Token Merging 来优化自定义 ViT,还是依赖 YOLO26 的前沿效率,使用 Ultralytics Platform 都可以极大地简化这些复杂的数据工作流。该平台为自动化 data annotation、无缝云端训练以及跨不同 edge computing 硬件环境的稳健模型部署提供了直观的生态系统。正在扩展其 computer vision 计划的组织依赖这些工具,可靠且高效地将尖端模型推向生产环境。






