提出高效令牌合并方法ToMA，让扩散模型生成更快更省显存。
ToMA: Token Merge with Attention for Diffusion Models
- 将令牌合并重构成子模优化问题，选多样令牌避免信息丢失。
- 用类注意力矩阵运算实现合并/还原，适配GPU加速，降低延迟。
- 利用潜在空间局部性和重复模式，显著减少计算开销，适合部署。
扩散模型在高质量图像生成上表现优异，但受限于Transformer的二次注意力复杂度，难以扩展。现有即插即用的令牌压缩方法（如ToMeSD、ToFu）虽能减少计算量，却依赖效率低下的操作（如排序、分散写入），在搭配FlashAttention等优化注意力实现时，反而因额外开销抵消理论加速效果。为此，我们提出离线可用的令牌合并与注意力方法（ToMA），有三大贡献：1）将令牌合并重构为子模优化问题，以选择多样性令牌；2）通过友好的GPU矩阵运算实现合并/还原，类似注意力的线性变换；3）利用潜在空间局部性和序列冗余（模式复用）最小化运行开销。ToMA使SDXL和Flux的生成延迟分别降低24%和23%（DINO差异小于0.07），超越已有方法。本工作弥合了扩散模型中Transformer理论与实际效率之间的差距。代码已开源：https://github.com/WenboLuu/ToMA。
原文摘要 · Abstract (English)
Diffusion models excel in high-fidelity image generation but face scalability limits due to transformers' quadratic attention complexity. Plug-and-play token reduction methods like ToMeSD and ToFu reduce FLOPs by merging redundant tokens in generated images but rely on GPU-inefficient operations (e.g., sorting, scattered writes), introducing overheads that negate theoretical speedups when paired with optimized attention implementations (e.g., FlashAttention). To bridge this gap, we propose Token Merge with Attention (ToMA), an off-the-shelf method that redesigns token reduction for GPU-aligned efficiency, with three key contributions: 1) a reformulation of token merge as a submodular optimization problem to select diverse tokens; 2) merge/unmerge as an attention-like linear transformation via GPU-friendly matrix operations; and 3) exploiting latent locality and sequential redundancy (pattern reuse) to minimize overhead. ToMA reduces SDXL/Flux generation latency by 24%/23%, respectively (with DINO $Δ< 0.07$), outperforming prior methods. This work bridges the gap between theoretical and practical efficiency for transformers in diffusion. Code available at https://github.com/WenboLuu/ToMA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。