选择性注意力:通过有原则的上下文控制增强Transformer

Transformer架构中的注意力机制使模型能够根据token与查询的相关性对其进行加权和组合。尽管自注意力取得了巨大成功,但它显著地以相同方式处理所有查询$q$,即应用映射$V^\top\text{softmax}(Kq)$,其中$V,K$分别是值嵌入和键嵌入。在这项工作中,我们认为这种均匀处理阻碍了控制上下文稀疏性和相关性的能力。作为解决方案,我们引入了$\textit{Selective Self-Attention}$ (SSA)层,通过有原则的温度缩放策略增强了softmax非线性。通过控制温度,SSA使注意力图的上下文稀疏性适应查询嵌入及其在上下文窗口中的位置。通过理论和实验,我们证明这缓解了注意力稀释,有助于优化过程,并增强了模型控制单个查询softmax尖锐度的能力。我们还为值嵌入引入了温度缩放,并表明这提高了模型抑制无关/噪声token的能力。值得注意的是,SSA是一种轻量级方法,通过权重共享策略引入了不到0.5%的新参数,并且可以在现有的LLM上进行微调。广泛的实证评估表明,配备SSA的模型在语言建模基准上取得了显著且一致的准确率提升。