提出新方法实现无需重训练的高效注意力压缩,兼容FlashAttention。
Representation Shift: Unifying Token Compression with FlashAttention
- 基于表示变化度衡量令牌重要性,无需注意力图或训练。
- 视频图文检索与视频问答分别提速5.5%和4.4%。
- 适用于Transformer、CNN及状态空间模型,通用性强。
Transformers在视觉、语言和视频任务中表现卓越,但任务复杂度提升导致模型增大、令牌增多,自注意力计算成本呈二次增长,且显存访问开销显著。为降低自注意力计算开销,已有工作提出删除冗余或低信息量令牌的令牌压缩技术。同时,如FlashAttention等融合内核通过避免构建注意力图及其对高带宽内存(HBM)的输入输出操作,缓解了内存开销。然而,这使得大多数无需训练的令牌压缩方法无法兼容,因其依赖注意力图判断令牌重要性。本文提出Representation Shift——一种无需训练、模型无关的度量方法,通过衡量每个令牌表示的变化程度来评估其重要性。该方法可无缝集成至FlashAttention,无需注意力图或重新训练。实验表明,Representation Shift能有效实现与FlashAttention兼容的令牌压缩,在视频图文检索和视频问答任务中分别获得5.5%和4.4%的显著加速。代码已开源。
原文摘要 · Abstract (English)
Transformers have demonstrated remarkable success across vision, language, and video. Yet, increasing task complexity has led to larger models and more tokens, raising the quadratic cost of self-attention and the overhead of GPU memory access. To reduce the computation cost of self-attention, prior work has proposed token compression techniques that drop redundant or less informative tokens. Meanwhile, fused attention kernels such as FlashAttention have been developed to alleviate memory overhead by avoiding attention map construction and its associated I/O to HBM. This, however, makes it incompatible with most training-free token compression methods, which rely on attention maps to determine token importance. Here, we propose Representation Shift, a training-free, model-agnostic metric that measures the degree of change in each token's representation. This seamlessly integrates token compression with FlashAttention, without attention maps or retraining. Our method further generalizes beyond Transformers to CNNs and state space models. Extensive experiments show that Representation Shift enables effective token compression compatible with FlashAttention, yielding significant speedups of up to 5.5% and 4.4% in video-text retrieval and video QA, respectively. Code is available at https://github.com/mlvlab/Representation-Shift.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。