3、本文方法
为了减少self-attention的能量消耗,可以对查询{qt}Nt=1和键{kt}Nt=1执行二进制量化。在这种情况下,我们可以用高能效的逐位运算代替大部分高能耗的乘法运算。然而,现有的二进制量化方法只关注于最小化原始全精度值与等式(5)中的二进制值之间的量化误差,无法保持注意力中不同标记之间的成对语义相似性,导致性能下降.
请注意,注意力可以看作是对成对的令牌应用内核平滑器,其中内核分数表示令牌对的相似性,如 3.2 节所述。受此启发,我们提出了一种新的二值化方法,该方法应用带有高斯 RBF 的核散列,将原始高维查询/键映射到汉明空间中的低维相似性保持二进制码。我们称之为 EcoFormer 的提议框架如图 1 (c) 所示。为了保持注意力的语义相似性,我们以自我监督的方式学习哈希函数。通过利用二进制代码之间线性点积的关联特性以及代码内积(即汉明亲和度)与汉明距离之间的等价性,我们能够以较低的能量成本在线性时间内逼近自注意力。下面,我们首先在 4.1 节介绍核化哈希注意力,然后在 4.2 节展示如何以自监督的方式学习哈希函数。
3.1、Kernelized Hashing Attention
在应用哈希函数之前,我们让查询 {qt}Nt=1 和键 {kt}Nt=1 相同。这样,我们就可以应用核化散列函数 H : RDp 7→ {1, -1}b 而无需显式应用 3.2 节中提到的变换 φ(·) 将 qi 和 kj 映射为 b 位二进制码 H(qi) 和 H(kj ),分别(见第 4.2 节)。在这种情况下,它们之间的汉明距离可以定义为
其中 Hr(·) 是二进制码的第 r 位;1{A} 是一个指示函数,如果满足 A,则返回 1,否则返回 0。使用 D (H(qi), H(kj )),H(qi) 和 H(kj ) 之间的代码内积可以是 制定为
重要的是,等式(7)显示了汉明距离和代码内积之间的等价性,因为存在一一对应关系。通过用等式(4a)中的散列查询和键替换,我们可以将自注意力近似为
注意 H(qt)> H(kj ) ∈ [−b, b]。为了避免分母为零,我们在每个内积中引入一个偏置项 2c,使得 H(qt)> H(kj ) + 2c > 0,对相似性度量没有影响。在这里,我们可以简单地将 c 设置为 d log2(b + 1)e,其中 due 返回大于或等于 u 的最小整数。利用矩阵乘法的关联属性,我们将自注意力近似为
在实践中,等式(9)中的二进制代码和全精度值之间的乘法可以用简单的加法和减法代替,这大大减少了片上能量足迹方面的计算开销。此外,与 2c 的二次幂的乘法也可以通过有效的位移操作来实现。结果,唯一的乘法来自分子和分母之间的元素除法。
3.2、Self-supervised Hash Function Learning
给定查询 Q = {q1, ... , qN } ⊂ RDp ,我们寻求学习一组哈希函数 h : RDp 7→ {1, -1}。我们没有显式应用第 3.2 节中提到的变换函数 φ(·),而是使用核函数 κ(qi, qj ) 计算散列函数:RDp × RDp 7→ R。给定 Q = [q1,..., qN ]> ∈ RN×Dp ,我们从 Q 中随机抽取 m 个查询 q(1), ..., q(m) 作为支持样本,遵循基于内核的监督散列 (KSH) [36] 并定义一个散列函数 h 为
为了指导二进制代码的学习,我们希望相似的标记对具有最小的汉明距离,而不同的标记对具有最大的距离。然而,由于方程(6)中的非凸和非光滑公式,直接优化汉明距离是困难的。利用等式(7)中代码内积和汉明距离之间的等价性,我们改为基于汉明亲和力进行优化以最小化重构误差为
其中 k·kF 是 Frobenius 范数,Y ∈ RN×N 是目标汉明亲和矩阵。为了保持查询和键之间的相似关系,我们使用注意力分数作为自监督信息来构造 Y。让 S 和 U 是相似和不相似的标记对。我们通过选择具有 Top-l 最大和最小注意力分数的标记对来获得 S 和 U。然后我们构造成对标签 Y 为
然而,问题 (12) 是 NP 难的。为了有效地解决它,我们采用离散循环坐标下降来顺序学习二进制代码。具体来说,我们只在之前的 a1,...,ar−1 优化后才求解 ar。令 ˆYr−1 = bY − P r−1 t=1 ht(Q)ht(Q)> 为残差矩阵,其中 ˆY0 = bY。然后,我们可以最小化以下目标以获得 ar
其中 C 是一个常数。注意 ^Yr−1 是一个对称矩阵。因此,问题 (14) 是一个标准的二元二次规划问题,可以通过许多现有的方法有效地解决,例如 LBFGS-B 求解器和块图切割。为了结合网络参数学习 ar,我们建议使用基于梯度的方法来解决问题 (14)。对于不可微分的符号函数,我们使用 STE 来近似梯度,使用硬 tanh,如 3.3 节所述。请注意,为每个 epoch 学习散列函数在计算上是昂贵的,但却是不必要的。我们只学习每个 τ epoch 的哈希函数。
4、实验
5、参考
[1].EcoFormer: Energy-Saving Attention with Linear Complexity.
6、推荐阅读
NeurIPS 2022 | 百度提出超快Transformer分割模型RTFormer,180FPS+81mIOU