Linear Attention

type

status

date

slug

summary

传统自注意力机制回顾

给定输入序列，Transformer 中通常会先通过线性变换得到查询（Query，）、键（Key，）与值（Value，）：

其中是输入序列的表示，为可学习的参数矩阵。

接着，传统自注意力的输出可以写为：

这里给出了序列间相似性矩阵

操作会按行进行归一化，使得每个时间步的注意力权重之和为 1

最终再将得到的注意力权重与相乘，得到输出

一般化的形式：

💡

其实很容易看出，制约self-attention的实际上就是softmax，如果去除softmax后，可以很直觉的使用结合律先计算KV，来大幅度的降低计算复杂度。

Transformer：

计算：

计算复杂度。

Softmax 归一化：

计算 Softmax 需要遍历矩阵，每一行计算指数和归一化，复杂度。

计算：

计算复杂度。

最终计算复杂度：

想象中的如果先计算，再计算

先计算 而不是 ：

计算复杂度。

计算：

计算复杂度。

最终计算复杂度：

Linear Attention

💡

显然Linear Attention的实现最重要的就是对 Softmax 的去处，当去处了 Softmax 后能够大幅度的降低计算Self-Attention的计算复杂度

1. 使用 kernel function 来替代 Softmax

其中

是一个核函数，如 ReLU、ELU，它将 Query 和 Key 投影到一个新空间，使得点积仍能有效计算注意力。

💡

需要了解Softmax起什么作用，才能知道Kernel function替代 softmax的合理性

softmax函数的作用是：

归一化注意力权重：

保证所有注意力权重的总和为1。

避免注意力的数值爆炸。

突出主要的注意点：

通过指数函数，softmax能明显凸显query与最相关的key之间的差异（如本例中key1权重明显高于其他key）。

即使差异不大，softmax也会放大微小的差别。

Kernel attention如何近似实现softmax类似作用

故

2.Causal Mask

在标准 Transformer 中，Causal Mask 通过 Softmax 归一化自然生效（隐式实现）。

而在自回归任务中，去除 Softmax 后显然需要面对Causal Mask问题，即可能看到未来的数据

在Linear Attention中，通过拆分注意力计算来实现这一点

Linear Attention 的发展

1、linear kernel func 逼近 softmax

2、最近工作发现不需要核函数也挺好，通过使用normalization 来保持概率特性

3、分母做normalization会造成数值不稳定，变成output做一个normalization

从数据分布的角度来思考self -attention

Self-Attention

在特征空间中的主要步骤：

线性投影到查询、键、值空间

原始输入特征经过三个投影矩阵变换：

其中，是学习的参数，将输入从原始特征空间投影到不同的查询、键和值空间。

相似性计算与注意力分数

计算实际上是查询和键在投影空间中的相似度计算，其中体现了不同 token 之间的关系：

这相当于在投影后的特征空间中进行点积相似度计算，衡量不同 token 之间的相关性。

归一化注意力分数

通过 softmax 归一化，使得注意力分数具有概率分布的性质：

💡

这样，Self-Attention 机制本质上是在一个投影特征空间中学习一个加权矩阵 A ，其中的值表示某个 token 对另一个 token 的依赖程度。

加权投影

最终通过计算输出：

这一步本质上是将 V 空间的特征进行加权聚合，实现信息的全局交互。

Linear attention

在 Linear Attention 里，计算方式变成：

💡

是一个 “KV 聚合特征空间”

它不是原始的 Key 或 Value，而是Key 经过核映射后，与 Value 进行加权累积的全局特征表征。

与 Softmax Attention 不同，Linear Attention 不是在计算注意力矩阵后再聚合，而是 先聚合 Key-Value 关系，然后 Query 再进行变换。

比较 Softmax attention和Linear attention，可以打一个比方

Softmax 就像是查询书本的目录，然后找到对应的内容

Linear attention 实际上是将书本变成一张巨大的A4纸，然后将内容压缩显示。所以只需要直接查询这个A4纸就可以了

容量限制（Capacity Limitation）

给定查询向量为，我们希望检索对应的value，计算方法为：

这个过程的含义：

当查询向量等于时，我们是沿着第一个key的方向进行记忆检索。

因为和完全正交，第二个记忆对（）根本不会干扰查询。

只返回对应的第一个记忆对（）的值。

假设此时加入一个新记忆：

存入新记忆后，矩阵为：

此时，虽然每个元素叠加了新值，但实际上旧的记忆仍然存在：

若用原本的检索：

结果为：；

此时由于新key 与旧key不正交，导致记忆混叠。

因此问题根本不在于元素被覆盖，而是key不再正交。