Selective attention improves transformer详细解读

Selective attention improves transformer Google 2024.10.3
一句话：简单且无需额外参数的选择性注意力机制，通过选择性忽略不相关信息并进行上下文剪枝，在不增加计算复杂度的情况下显著提升了Transformer模型的语言建模性能和推理效率。
论文链接：https://arxiv.org/pdf/2410.02703v1
1.为什么引入selective attention
Transformer的记忆负担
Transformer，核心self-attention关注输入序列中的所有部分，不仅仅是局部信息（也是RNN、LSTM改进可以关注更长字符串，但是也引入了更高计算量），所有信息都保存在上下文缓冲区，计算所有上下文信息的相关性
Selective attention高效信息筛选器
自动删除不再有用的信息，从上下文缓冲区中移除不必要的元素，提高模型性能，减少计算和内存需求

标记token b无法影响标记c从标记a中读取信息的程度，标记b确定标记a对于后续标记c是不相关甚至是误导性的
Selective attention允许一个标记决定另一个标记不再被需要，从而减少后续标记对该标记的关注度
2.selective attention可视化剔除token过程
（1）变量赋值
y=7; x=1; x=3; z=5; x=? 则：x=3，即不管前面x=1赋值是多少，都与之无关

红色线代表对前面token的掩蔽程度，在变量赋值中，掩蔽程度非0即1（一般为[0,1]）
变量赋值中，绿色箭头处，当出现第二次 ‘Z=’ token时，前面Z=、177直接掩蔽掉
（2）自然语言模型
序列：Bar,##ack, Obama

##ack直接掩蔽了bar，这里红色线有深浅，代表掩蔽程度不同，比如day对a的掩蔽程度比较浅，说明保留了部分a的信息
3.选择函数
selection matrix SNxN，Sij表示标记xi对标记xj的掩蔽程度

S矩阵限制条件：1.通过ReLU限制负值为0，只降低注意力，不增强注意力 2. Begin of Sentence标记，初始句首不屏蔽，本身不关注不屏蔽







第一行：QK/sqrt(dk)
第二行：mask引入设置
第三行：选定head 0
第四行-第六行：S的三个约束条件
第七行：右移，且右移后对角线为0？
第八行：S累加得到F，为什么累加？
第九行：从标准attention中减去F
第十行：归一化权重

文心一言代码解读

4.context pruning 上下文剪枝
上下文缓冲区修剪元素来减少注意力模块的内存和计算需求。每层的稀疏性在样本之间是稳定的（本文实验有验证），为每一层设定各自固定的内存预算。
上下文剪枝步骤：
1.初始化K = K1, . … , KL= N，为每层内存预算，其中N的上下文缓冲区大小
2.前Kl个token保持，后续每个token和前面对比，丢弃最高F值对应token
3.贪婪迭代方法分配总的内存预算，迭代直到模型性能达到预定义阈值，即标准attention模型性能
5.loss

每层内存之和/层数token数，我们希望M（内存）越小越好，M越小，L越小，相关性一致，同时Ln≠pad即同aqrt(dk)限定范围一样，将分子大小限定一定范围内

内存计算， τ= 1限定F矩阵范围不超过1
Lppl： standard log-perplexity loss 标准对角困惑度损失函数
ϵ is a small weight factor: ϵ = 0.1， τ= 1，固定数值
L表示层数，n≠pad表示非填充标记的数量（字符串输入固定，缺失填充padding，对应token来说即非填充token，实际有效信息token）
6.selective attention改进及其效果
（1）简单且无需额外参数
（2）减小注意力机制的上下文缓冲区大小，推理过程中显著减少内存和计算需求
（3）标准Attention模型拥有约两倍多的头数和参数与selective attention效果相当
（4）上下文大小为512、1024和2048时，内存分别比未采用选择性注意力的相同验证困惑度的模型减少16倍、25倍和47倍
7.待改进
（1）Decoder-only
（2）上下文减少提高推理效率，但并不能提高训练效率，探索在训练过程中迭代减少上下文缓冲区的大小
（3）移除元素后，没有对模型进行进一步的训练，在上下文减少后进行一些额外的训练可能会实现进一步的改进
（4）仅对具有选择性注意力的预训练模型进行了实验，微调步骤中将其应用于现有模型
8.实验