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.实验