表示偏移：统一令牌压缩与FlashAttention

Transformer 在视觉、语言和视频等多个领域已展现出卓越的性能。然而，随着任务复杂性的增加，模型规模和序列长度不断增大，导致自注意力机制的计算开销呈二次增长，同时GPU内存访问的开销也显著上升。为降低自注意力的计算成本，先前的研究提出了多种令牌压缩技术，通过移除冗余或信息量较低的令牌来减少计算量。与此同时，诸如 FlashAttention 这类融合注意力核函数的方法被提出，通过避免显式构建注意力矩阵及其对高带宽内存（HBM）的输入/输出操作，有效缓解了内存访问的开销。然而，这类方法与大多数无需训练的令牌压缩技术不兼容，因为后者通常依赖注意力矩阵来判断令牌的重要性。本文提出一种名为 Representation Shift 的新方法，这是一种无需训练、与模型无关的度量指标，用于衡量每个令牌表示的变化程度。该方法可无缝集成至 FlashAttention 之中，无需依赖注意力矩阵，也无需重新训练模型。此外，该方法还具有良好的泛化能力，可扩展至卷积神经网络（CNNs）和状态空间模型（State Space Models）等架构。大量实验结果表明，Representation Shift 能够实现与 FlashAttention 兼容的有效令牌压缩，在视频-文本检索和视频问答任务中分别带来高达 5.5% 和 4.4% 的加速效果。代码已开源，地址为：https://github.com/mlvlab/Representation-Shift。