一种可理解的线性transformer

找到一种更简洁的形式，如下：

message passing公式。其含义是，每个token产生一个消息$Q_i$。然后消息通过权重$m_i$加权合并，最后通过权重$s_i$消息分发。极限情况下(即softmax中有一个token的值为1，其余token值为0)每一次传递，会将某个token的消息传给另一个token。问题在于极限情况下，每一层只能设定固定的传递次数。在全连接图的情况下，传递次数应为$O(n^2)$。在softmax情况下，这种情况可能有所改善。权重平均的情况下，相当于每个token加入了一个相同常数。令$m_i=w_1{x}_i,s_i=w_2{x}_i$，有
$$Q=\sum _i\frac{exp(m_i)}{{\sum }_{j}exp(m_j)}{Q}_i$$

$$x_n^{\prime }=x_n+\frac{exp(s_i)}{{\sum }_{j}exp(s_j)}Q$$
如果进行$k$次合并，和$k$次分发，总体公式如下：

$$x^{\prime }=x+softmax(Sx,dim=l)softmax(Mx,dim=l{)}^{T}Qx$$

$$，其中Qx:[batch,len,c],softmax(Mx{)}^T:[b,k,l],exp(Sx):[b,l,k]$$

从公式上来看，更像一个低秩方法。需要看一下与其它的低秩方法的区别。

与之对比的self_attention公式如下，可以发现其实就是拆解了self attention的softmax，比较巧合。

$$x^{\prime }=x+softmax((Qx{)}^T(Kx),dim=l)\ast (Vx)$$

其中

$$Qx:[b,l,k],Kx:[b,k,l],Vx:[b,l,c]$$

参考了一下linear transformer的论文，主要是https://blog.csdn.net/hymn1993/article/details/125254897。发现确实是低秩方法的一种，只不过用softmax作为核进行映射。不过比较巧妙的是，这套方法有明确的可解释意义。

transformer$O(n^2)$复杂度的关键点在于，对每个token都查询了一次。因此降低复杂度的一个行之有效的方法是降低查询的次数。因此提出竞争查询的方法。公式如下:

$$Q=\sum _i\frac{exp(z_i)}{{\sum }_{j}exp(z_j)}{Q}_i$$

$$x_n^{\prime }=x_n+\frac{exp(z_n)}{{\sum }_{j}exp(z_j)}\sum _m(\frac{exp(Q\ast K_m)}{{\sum }_{k}exp(Q\ast K_k)}{V}_m)$$

$Z$向量为竞争向量，通过softmax归一化得到分布在tokens上的权重，根据$Z$的权重对所有的query向量$Q_i$进行求和，得到竞争成功的$Q$向量。可以理解为这一步将所有要查询的东西编码到同一个向量中。然后正常按照transformer的办法用Q向量与每个token的key向量$K_m$和$V_m$得到更新向量。然后按照$Q$向量的比例，依次按照比例把更新向量赋值给所有的token。可以看出当$Q_i$的比例为极限情况(0,0,…,1,…,0,0)时，相当于只对比例为1的token做查询。

另外一种公式是:

$$z_i=w{x}_i$$

$$Q=\sum _i\frac{exp(z_i)}{{\sum }_{j}exp(z_j)}{Q}_i$$

$$V=\sum _i\frac{exp(z_i)}{{\sum }_{j}exp(z_j)}{V}_i$$

$$x_n^{\prime }=x_n+relu(Q{K}_n)V$$

也是具备非常明确的意义。

第二种存在GPT形式(casual mask attention):

要依次根据mask，生成每个token的$Q^n,V^n$。每个token的更新如下:

$$Q^n=\sum _{i=0}^n\frac{exp(z_i)}{{\sum }_{j=0}^{n}exp(z_j)}{Q}_i$$

$$V^n=\sum _{i=0}^n\frac{exp(z_i)}{{\sum }_{j=0}^{n}exp(z_j)}{V}_i$$

$$x_n^{\prime }=x_n+relu(Q^n{K}_n)V^n$$

注意$Q_n$的计算是存在递归公式的，因此其复杂度为seq len线性相关，缺点是无法并行。

第二个公式是有明确意义的。每一层都预定了一个w，用来判断query的重要程度。因此w是一个先验知识，用来决定应该先查询什么，然后再查询什么。但是真正的查询向量Q又是和序列有关的，不同的序列有不同的查询向量。