启用 FlashAttention 的令牌压缩 “表象移动 “的工作原理和效果

摘要

本文提出了一种新方法来应对变压器模型计算成本增加的挑战。

近年来，变换器已被广泛应用于自然语言处理和图像/视频理解领域，但随着规模的扩大，处理效率已成为一个严重问题，因为自我注意机制的计算复杂度与输入词块数量的平方成正比增长。

传统上，人们试图从两个方向解决这一问题。

一种是以 FlashAttention 为代表的内存效率方法，另一种是基于标记压缩的计算量减少方法。

然而，标记压缩通常与 FlashAttention 等不建立注意力图谱的机制不兼容，因为它使用注意力图谱来估计标记的重要性。

因此，作者提出了一种新的指标–表征偏移。这可以衡量每个标记在通过各层时的表征发生了多大变化，也是定义重要性的一种方法。

该指标无需训练，与模型无关，可与 FlashAttention 结合使用。

实验结果表明，这种方法在效率和准确性方面都优于传统方法，推理速度提高了 5.5 倍。

建议的方法

所提出的方法 “表征偏移”（Representation Shift）可以测量每个标记在层的输入和输出中的嵌入表征之间的差异，从而量化标记在模型中被增强了多少信息。

具体来说，计算向量通过 MLP 层和注意力层前后的距离，并将此值作为重要性得分。L2 准则在测量距离方面表现最为稳定。

传统方法依赖于注意力图谱，而这种方法可以独立于注意力机制来估计标记的重要性，因此可以自然地与不构建注意力图谱的计算方法（如 FlashAttention）相结合。

该框架还具有足够的通用性，不仅适用于 Transformer，还适用于 CNN 和状态空间模型 (SSM)。

作者还详细研究了设计选择，如在哪一层测量表征偏移，以及基于哪种操作（注意力或 MLP）。
结果表明，使用 MLP 层的变化量最为有效。

这种设计可以将信息损失降到最低，同时消除标记冗余。

实验

作者对图像分类和视频理解任务进行了广泛的实验，以测试所提方法的有效性。

首先，在视频任务中，使用无掩码教师（UMT）对视频文本检索和视频质量保证进行了评估，每层标记减少了 20%。

结果表明，与现有的基于注意力分数的方法相比，表征转移与 FlashAttention 的结合更快、更准确，吞吐量最多提高了 5.5 倍。

与单纯的模型微型化相比，它还能更好地权衡速度/准确性。

接下来，使用 DeiT 序列在 ImageNet 上对图像分类进行了验证，与 FlashAttention 相结合，推理速度提高了 1.2 倍，同时准确率也高于传统的基于注意力的方法。

它还被应用于 CNN/SSM，如 ResNet 和 Vision Mamba，其有效性在这些非基于变换器的架构中得到了证实。

特别是，在 ResNet-50 上逐行标记剪枝的速度提高了 18% 以上，并在很大程度上保持了准确性。

这些实验表明，表征转换是一种通用而强大的标记压缩标准。