arXiv:2411.07176cs.CLcs.AI2024-11被引 13

让注意力权重可为负,提升模型表达能力与鲁棒性

More Expressive Attention with Negative Weights

  • 引入可学习负权重的注意力机制,动态通过查询键内积符号实现删减或复制操作
  • 在语言建模和图像生成任务中均优于传统softmax注意力,性能更优
  • 适合关注注意力机制创新、模型表达力提升的研究者

我们提出一种新型注意力机制——Cog Attention,允许注意力权重为负以增强表达能力。其核心在于两点:(1) 提升参数灵活性,例如,不同于传统softmax注意力头使用固定输出值(OV)矩阵来删除或复制输入,Cog Attention能自然地利用动态查询-键(QK)内积的符号表示这些操作,使单个注意力头可同时执行多种操作,同时使OV矩阵更专注于精细化调整;(2) 通过防止早期令牌被过度压缩到后期位置,增强模型对表征坍缩的鲁棒性。我们构建了基于Cog Attention的Transformer类模型,涵盖用于语言建模的解码器单向模型及不同规模的U-ViT扩散模型。实验表明,采用Cog Attention的模型在各项任务上均优于传统softmax注意力模型。该方法为突破传统softmax注意力中非负权重等固有约束提供了新方向。

原文摘要 · Abstract (English)

We propose a novel attention mechanism, named Cog Attention, that enables attention weights to be negative for enhanced expressiveness, which stems from two key factors: (1) Cog Attention enhances parameter flexibility. For example, unlike traditional softmax attention heads that use a static output-value (OV) matrix to delete or copy inputs that the heads attend to, Cog Attention naturally learns to use the sign of dynamic query-key (QK) inner products to represent these operations. This enables Cog Attention to perform multiple operations simultaneously within a single head. Meanwhile, Cog Attention's OV matrix can focus more on refinement or modification. (2) Cog Attention enhances the model's robustness against representational collapse by preventing the ``over-squashing'' of earlier tokens into later positions. We develop Transformer-like models which use Cog Attention as attention modules, including decoder-only models at various scales for language modeling and U-ViT diffusion models for image generation. Experiments show that models using Cog Attention exhibit superior performance compared to those employing traditional softmax attention modules. Our approach suggests a promising research direction for rethinking and breaking the entrenched constraints of traditional softmax attention, such as the requirement for non-negative weights.

注意力机制表达能力负权重Transformer

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。