通过温度调控提升注意力机制的上下文控制能力,增强模型表现。
Selective Attention: Enhancing Transformer through Principled Context Control

- 引入选择性自注意力层,用温度调节控制查询的上下文稀疏性。
- 在语言建模任务中实现稳定且显著的准确率提升。
- 轻量级设计,新增参数不足0.5%,可直接微调现有大模型。
Transformer架构中的注意力机制使模型能够根据查询与标记的相关性进行加权和融合。尽管自注意力取得了巨大成功,但其对所有查询均采用统一的映射方式 $V^ op ext{softmax}(Kq)$,其中 $V$、$K$ 分别为值和键嵌入。本文认为这种均匀处理限制了对上下文稀疏性和相关性的控制。为此,我们提出选择性自注意力(SSA)层,通过原理化的温度缩放策略增强softmax非线性。通过调节温度,SSA能根据查询嵌入及其在上下文窗口中的位置动态调整注意力图的稀疏性。理论与实验表明,该方法缓解了注意力稀释问题,促进优化过程,并增强对单个查询softmax尖锐度的控制。我们还对值嵌入引入温度缩放,进一步提升模型抑制无关或噪声标记的能力。值得注意的是,SSA为轻量级方法,通过参数共享策略引入的新增参数少于0.5%,可在现有大语言模型上进行微调。大量实证评估显示,配备SSA的模型在语言建模基准上实现明显且一致的性能提升。
原文摘要 · Abstract (English)
The attention mechanism within the transformer architecture enables the model to weigh and combine tokens based on their relevance to the query. While self-attention has enjoyed major success, it notably treats all queries $q$ in the same way by applying the mapping $V^\top\text{softmax}(Kq)$, where $V,K$ are the value and key embeddings respectively. In this work, we argue that this uniform treatment hinders the ability to control contextual sparsity and relevance. As a solution, we introduce the $\textit{Selective Self-Attention}$ (SSA) layer that augments the softmax nonlinearity with a principled temperature scaling strategy. By controlling temperature, SSA adapts the contextual sparsity of the attention map to the query embedding and its position in the context window. Through theory and experiments, we demonstrate that this alleviates attention dilution, aids the optimization process, and enhances the model's ability to control softmax spikiness of individual queries. We also incorporate temperature scaling for value embeddings and show that it boosts the model's ability to suppress irrelevant/noisy tokens. Notably, SSA is a lightweight method which introduces less than 0.5% new parameters through a weight-sharing strategy and can be fine-tuned on existing LLMs. Extensive empirical evaluations demonstrate that SSA-equipped models achieve a noticeable and consistent accuracy improvement on language modeling benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。