连续化 Softmax Attention

Softmax Attention想必大家已经不陌生，它是定义在有限向量集合下的一个加权混合算子。近日，论文《Continuum Attention》提出了一个很有意思的想法——将注意力连续化。这让笔者想起，之前在《为什么DeltaNet要加L2 Normalize？》介绍的EFLA也是将离散RNN连续化，然后又重新离散化，得到一些有意思的结果，类似的还有早期的SSM系列工作。

这样看来，“连续化—离散化”似乎是一条屡试不爽的探索路径。那么，Softmax Attention连续化又能得到什么有意思的结果呢？让我们一起来品鉴一番。

问题描述[#](https://spaces.ac.cn/archives/11980#%E9%97%AE%E9%A2%98%E6%8F%8F%E8%BF%B0)

标准的Softmax Attention定义为

o t=∑t i=1 e q t⋅k i v i∑t i=1 e q t⋅k i

o t=∑t i=1 e q t⋅k i v i∑t i=1 e q t⋅k i(1)

所谓连续化，是想将它改成

o t=∫t 0 e q t⋅k s v s d s∫t 0 e q t⋅k s d s o t=∫t 0 e q t⋅k s v s d s∫t 0 e q t⋅k s d s(2)

其中非整数点的 k,v k,v，定义为线性插值，即对于任意 s∈[i−1,i],i∈N+s∈[i−1,i],i∈N+，定义

k s=(i−s)k i−1+(s−i+1)k i,v s=(i−s)v i−1+(s−i+1)v i k s=(i−s)k i−1+(s−i+1)k i,v s=(i−s)v i−1+(s−i+1)v i(3)

该定义需要额外引入 k 0,v 0 k 0,v 0 两个向量（用于 s∈[0,1]s∈[0,1]），其取值我们最后再讨论。 注意，式(2)(2)的定义跟原论文略有不同，原论文出于对称考虑，将积分区间改为了[1/2,t+1/2][1/2,t+1/2]，但笔者感觉这样调整虽然更对称了，但对Causal场景反而没那么自然，所以本文改用积分区间[0,t][0,t]。当然，接下来的推导其实跟区间的选择关系不大，技术上都是一样的。

准备工作[#](https://spaces.ac.cn/archives/11980#%E5%87%86%E5%A4%87%E5%B7%A5%E4%BD%9C)

式(2)(2)只是一个形式化定义，想要实际使用它，就必须把积分显式地计算出来。根据积分的可加性，我们只需要将每一段∫i+1 i∫i+1 i 都分别算出来。记 a t,i=q t⋅k i,δ t,i=a t,i−a t,i−1 a t,i=q t⋅k i,δ t,i=a t,i−a t,i−1，不难得到

∫i i−1 e q t⋅k s v s d s=∫1 0 e(1−s)a t,i−1+s a t,i((1−s)v i−1+s v i)d s=v i∫1 0 e a t,i−1+s δ t,i d s+(v i−1−v i)∫1 0 e a t,i−1+s δ t,i(1−s)d s=e a t,i−1 φ 1(δ t,i)v i+e a t,i−1 φ 2(δ t,i)(v i−1−v i)∫i i−1 e q t⋅k s d s=∫1 0 e(1−s)a t,i−1+s a t,i d s=∫1 0 e a t,i−1+s δ t,i d s=e a t,i−1 φ 1(δ t,i)

∫i i−1 e q t⋅k s v s d s===∫1 0 e(1−s)a t,i−1+s a t,i((1−s)v i−1+s v i)d s v i∫1 0 e a t,i−1+s δ t,i d s+(v i−1−v i)∫1 0 e a t,i−1+s δ t,i(1−s)d s e a t,i−1 φ 1(δ t,i)v i+e a t,i−1 φ 2(δ t,i)(v i−1−v i)∫i i−1 e q t⋅k s d s=∫1 0 e(1−s)a t,i−1+s a t,i d s=∫1 0 e a t,i−1+s δ t,i d s=e a t,i−1 φ 1(δ t,i)(4)(5)

这里

φ 1(x)≜∫1 0 e s x d s=e x−1 x,φ 2(x)≜∫1 0 e s x(1−s)d s=e x−1−x x 2 φ 1(x)≜∫1 0 e s x d s=e x−1 x,φ 2(x)≜∫1 0 e s x(1−s)d s=e x−1−x x 2(6)

其中 x=0 x=0 时按照极限来处理，分别有 φ 1(0)=1,φ 2(0)=1/2 φ 1(0)=1,φ 2(0)=1/2。它们还成立如下恒等关系

φ 1(x)−φ 2(x)=x e x−e x+1 x 2=e x⋅e−x−1−(−x)(−x)2=e x φ 2(−x) φ 1(x)−φ 2(x)=x e x−e x+1 x 2=e x⋅e−x−1−(−x)(−x)2=e x φ 2(−x)(7)

这个恒等式在下一节化简时会用到。另外，从积分定义可以看出，φ 1(x)φ 1(x)和 φ 2(x)φ 2(x)都是恒正的，这样后面的取对数运算就不会有未定义的问题。

组装化简[#](https://spaces.ac.cn/archives/11980#%E7%BB%84%E8%A3%85%E5%8C%96%E7%AE%80)

将上述积分代入到式(2)(2)，得到

o t=∑t i=1 e a t,i−1 φ 1(δ t,i)v i+e a t,i−1 φ 2(δ t,i)(v i−1−v i)∑t i=1 e a t,i−1 φ 1(δ t,i)

o t=∑t i=1 e a t,i−1 φ 1(δ t,i)v i+e a t,i−1 φ 2(δ t,i)(v i−1−v i)∑t i=1 e a t,i−1 φ 1(δ t,i)(8)

先化简分子

t∑i=1 e a t,i−1 φ 1(δ t,i)v i+e a t,i−1 φ 2(δ t,i)(v i−1−v i)=t∑i=1 e a t,i−1(φ 1(δ t,i)−φ 2(δ t,i))v i+e a t,i−1 φ 2(δ t,i)v i−1=t∑i=1 e a t,i−1+δ t,i φ 2(−δ t,i)v i+e a t,i−1 φ 2(δ t,i)v i−1=t∑i=1 e a t,i φ 2(−δ t,i)v i+t−1∑i=0 e a t,i φ 2(δ t,i+1)v i=e a t,0 φ 2(δ t,1)v 0+e a t,t φ 2(−δ t,t)v t+t−1∑i=1 e a t,i(φ 2(−δ t,i)+φ 2(δ t,i+1))v i ====∑i=1 t e a t,i−1 φ 1(δ t,i)v i+e a t,i−1 φ 2(δ t,i)(v i−1−v i)∑i=1 t e a t,i−1(φ 1(δ t,i)−φ 2(δ t,i))v i+e a t,i−1 φ 2(δ t,i)v i−1∑i=1 t e a t,i−1+δ t,i φ 2(−δ t,i)v i+e a t,i−1 φ 2(δ t,i)v i−1∑i=1 t e a t,i φ 2(−δ t,i)v i+∑i=0 t−1 e a t,i φ 2(δ t,i+1)v i e a t,0 φ 2(δ t,1)v 0+e a t,t φ 2(−δ t,t)v t+∑i=1 t−1 e a t,i(φ 2(−δ t,i)+φ 2(δ t,i+1))v i(9)

由此可见分子具有∑t i=0 e a t,i+b t,i v i∑t i=0 e a t,i+b t,i v i 的形式，其中

a t,i=q t⋅k i,e b t,i={φ 2(a t,1−a t,0)i=0 φ 2(a t,i−1−a t,i)+φ 2(a t,i+1−a t,i)1≤i≤t−1 φ 2(a t,t−1−a t,t)i=t a t,i=q t⋅k i,e b t,i=⎧⎩⎨⎪⎪φ 2(a t,1−a t,0)φ 2(a t,i−1−a t,i)+φ 2(a t,i+1−a t,i)φ 2(a t,t−1−a t,t)i=0 1≤i≤t−1 i=t(10)

注意分母正好是全体 v i v i 换成 1 1 的结果，所以直接可得分母为∑t i=0 e a t,i+b t,i∑t i=0 e a t,i+b t,i，于是

o t=∑t i=0 e a t,i+b t,i v i∑t i=0 e a t,i+b t,i o t=∑t i=0 e a t,i+b t,i v i∑t i=0 e a t,i+b t,i(11)

结果分析[#](https://spaces.ac.cn/archives/11980#%E7%BB%93%E6%9E%9C%E5%88%86%E6%9E%90)

相比标准的Softmax Attention

o t=∑t i=1 e a t,i v i∑t i=1 e a t,i

o t=∑t i=1 e a t,i v i∑t i=1 e a t,i(12)

连续化带来的变化，就是引入了额外的 k 0,v 0 k 0,v 0，以及给每个 a t,i a t,i 加上了修正项 b t,i b t,i。多出来的 k 0,v 0 k 0,v 0 其实不算很本质的改动，我们等会再简单讨论一下，现在主要看 b t,i b t,i 对于一般的 1≤i≤t−1 1≤i≤t−1，b t,i b t,i 同时依赖于 a t,i a t,i 以及左右的 a t,i−1,a t,i+1 a t,i−1,a t,i+1，这在形式上很像是给Attention矩阵加上了Short Conv，尤其是对 b t,i b t,i 做泰勒展开后，这一点更为明显

b t,i=log(φ 2(a t,i−1−a t,i)+φ 2(a t,i+1−a t,i))=a t,i−1+a t,i+1−2 a t,i 6+⋯

b t,i=log(φ 2(a t,i−1−a t,i)+φ 2(a t,i+1−a t,i))=a t,i−1+a t,i+1−2 a t,i 6+⋯(13)

其中 a t,i−1+a t,i+1−2 a t,i a t,i−1+a t,i+1−2 a t,i 正是 a t,i a t,i 的二阶差分，属于一种简单的Short Conv运算。所以，连续化Softmax Attention确实与“给Attention矩阵加Short Conv”有颇多异曲同工之处。 至于 k 0,v 0 k 0,v 0，其实就相当于我们为解决Attention Sink问题引入的Sink token，对于连续化来说，主要是为了补足 t t 个区间，是一个单纯的数学技巧，我们一般有三种选择：全零、k 0=k 1,v 0=v 1 k 0=k 1,v 0=v 1 以及可训练，笔者在下面的实验中选择了全零。

简单实验[#](https://spaces.ac.cn/archives/11980#%E7%AE%80%E5%8D%95%E5%AE%9E%E9%AA%8C)

b t,i b t,i 同时依赖于 a t,i a t,i 及其相邻值，也意味着它无法复用现有的Flash Attention，但实现原理上是通用的，并且理论FLOPs也基本一致。笔者让K3简单写了一个Kernel，速度只有Softmax Attention的1/4左右，主要是里边细碎的运算太多了，难以加速，不过总归可以先跑一下。

第一个实验结果是，连续化虽然引入了局域先验，但它无法完全取代位置编码，在全Full Attention的架构下，连续化Softmax Attention如果不加RoPE，会比标准Softmax Attention加RoPE要明显差，相反如果加上RoPE，则能够追回这个差距。

至于训练动态，在大家都加RoPE的情况下，连续化Softmax Attention前期会比标准Softmax Attention明显好，但差距逐渐缩窄，等到训练结束时，两者基本重合，优势消失。这里猜测核心原因是，自然语言并不是真正的时间序列，连续化并不能带来实质好处，也许对于真正的时间序列，才有机会测出优势。

笔者还做了从标准Softmax Attention中途切换连续化版本的实验，发现起跳Loss并不算高，最后也是慢慢收敛到跟标准版基本一致的效果。

文章小结[#](https://spaces.ac.cn/archives/11980#%E6%96%87%E7%AB%A0%E5%B0%8F%E7%BB%93)

本文简单推导了一下连续化的Softmax Attention，其结果类似于对Attention矩阵加Short Conv操作。虽然在自然语言任务上实测并没有发现这种连续化带来的优势，但它仍不失作为一个有趣的结果来赏析。

_转载到请包括本文地址：[https://spaces.ac.cn/archives/11980](https://spaces.ac.cn/archives/11980 "连续化 Softmax Attention")_

_更详细的转载事宜请参考：_[《科学空间FAQ》](https://spaces.ac.cn/archives/6508#%E6%96%87%E7%AB%A0%E5%A6%82%E4%BD%95%E8%BD%AC%E8%BD%BD/%E5%BC%95%E7%94%A8 "《科学空间FAQ》")