[论文解读] EcoFormer: Energy-Saving Attention with Linear Complexity

Jing Liu, Zizheng Pan+3 more|arXiv (Cornell University)|Sep 19, 2022

Advanced Neural Network Applications被引用 16

一句话总结

EcoFormer 提出了一种线性复杂度、能效更高的注意力机制，通过核化哈希将查询和键映射为紧凑的二进制编码，保留成对相似性的同时，用低能耗累加操作替代昂贵的乘加运算。与标准注意力相比，其在 ImageNet-1K 上仅损失 0.33% 准确率，但片上能耗降低了 73%。

ABSTRACT

Transformer is a transformative framework that models sequential data and has achieved remarkable performance on a wide range of tasks, but with high computational and energy cost. To improve its efficiency, a popular choice is to compress the models via binarization which constrains the floating-point values into binary ones to save resource consumption owing to cheap bitwise operations significantly. However, existing binarization methods only aim at minimizing the information loss for the input distribution statistically, while ignoring the pairwise similarity modeling at the core of the attention. To this end, we propose a new binarization paradigm customized to high-dimensional softmax attention via kernelized hashing, called EcoFormer, to map the original queries and keys into low-dimensional binary codes in Hamming space. The kernelized hash functions are learned to match the ground-truth similarity relations extracted from the attention map in a self-supervised way. Based on the equivalence between the inner product of binary codes and the Hamming distance as well as the associative property of matrix multiplication, we can approximate the attention in linear complexity by expressing it as a dot-product of binary codes. Moreover, the compact binary representations of queries and keys enable us to replace most of the expensive multiply-accumulate operations in attention with simple accumulations to save considerable on-chip energy footprint on edge devices. Extensive experiments on both vision and language tasks show that EcoFormer consistently achieves comparable performance with standard attentions while consuming much fewer resources. For example, based on PVTv2-B0 and ImageNet-1K, Ecoformer achieves a 73% on-chip energy footprint reduction with only a 0.33% performance drop compared to the standard attention. Code is available at https://github.com/ziplab/EcoFormer.

研究动机与目标

• 解决标准 Transformer 在边缘设备上能耗过高的问题。
• 克服传统二值化方法在注意力机制中忽略成对相似性的局限性。
• 开发一种专为 Softmax 注意力机制设计的二值化策略，以保留 token 级别的相似性关系。
• 在实现注意力计算线性复杂度的同时，将性能下降最小化。
• 通过能效高且硬件友好的操作，实现资源受限边缘设备上 Transformer 模型的可扩展部署。

提出的方法

• 使用 RBF 核的核化哈希，将高维查询和键映射到汉明空间中的低维二进制编码。
• 以自监督方式训练哈希函数，使其汉明亲和度（二进制编码的点积）与真实注意力得分对齐。
• 利用矩阵乘法的结合性质，将注意力表示为二进制编码的线性复杂度点积。
• 将大部分乘加运算替换为简单累加，以减少片上能耗。
• 使用标准注意力的相似性矩阵作为哈希函数学习的监督信号。
• 引入支持样本策略，通过超参数 $ m $ 提高训练过程中哈希函数的泛化能力。

实验结果

研究问题

• RQ1 是否可以对自注意力中的查询和键进行定制化二值化，以保留成对相似性关系，而非最小化统计量化误差？
• RQ2 使用 RBF 核的核化哈希能否在保持竞争力性能的同时实现线性复杂度注意力？
• RQ3 通过将 FP32/FP16 乘加运算替换为二进制累加，Transformer 中的能耗可降低多少？
• RQ4 支持样本数量 $ m $ 对所提方法的性能和能效有何影响？
• RQ5 EcoFormer 是否能在视觉和语言基准上实现与标准注意力相当的准确率，同时显著降低计算和能耗成本？

主要发现

• 与标准注意力相比，EcoFormer 在 PVTv2-B0 上实现了 73% 的片上能耗降低，且在 ImageNet-1K 上仅损失 0.33% 准确率。
• 在 CIFAR-100 上，使用 $ m=25 $ 个支持样本的 EcoFormer 达到 71.23% 的 Top-1 准确率，优于支持样本更少的模型。
• 该方法将计算复杂度从 $ frac{1}{2}N^2 $ 降低至 $ frac{1}{2}N $，实现了线性时间注意力计算。
• 在多个基准测试中验证，EcoFormer 的延迟和片上能耗显著低于其他高效注意力方法。
• 在 ImageNet-1K 上从零开始训练表明，EcoFormer 在 PVTv2-B0 上减少了 72% 的乘法运算和 72% 的能耗。
• 二值化导致的性能下降主要源于不可微符号函数中的梯度近似，该问题可通过先进优化技术缓解。