arXiv:2409.04275cs.LG2024-09

通过分布式优化视角引入注意力共识差异,提升模型性能。

AttentionX: Exploiting Consensus Discrepancy In Attention from A Distributed Optimization Perspective

  • 借鉴拉格朗日乘子法,将注意力中的共识差异融入更新过程。
  • 在ViT和nanoGPT上验证,显著改善了模型表现。
  • 适合对自注意力机制改进感兴趣的开发者与研究者。

本文从分布式优化视角出发,扩展了标准Transformer中的注意力机制,提出AttentionX。该方法受对偶乘子法(PDMM)启发,其核心思想是通过拉格朗日乘子捕捉线性边约束残差中的历史共识差异,从而促进节点间逐步达成一致。在注意力机制中,$KQ$-softmax加权求和相当于信息收集,而前馈网络(FFN)则对应局部信息融合。本工作将共识差异定义为加权求和的$V$表示与原始$V$表示之间的差值,并将其引入注意力输出更新表达式。实验结果表明,在ViT和nanoGPT模型上,AttentionX均表现出色,验证了其有效性。

原文摘要 · Abstract (English)

In this paper, we extend the standard Attention in transformer by exploiting the consensus discrepancy from a distributed optimization perspective, referred to as AttentionX. It is noted that the primal-dual method of multipliers (PDMM) \cite{Zhang16PDMM} is designed to iteratively solve a broad class of distributed optimization problems over a pear-to-pear (P2P) network, where neighbouring nodes gradually reach consensus as specified by predefined linear edge-constraints in the optimization process. In particular, at each iteration of PDMM, each node in a network first performs information-gathering from neighbours and then performs local information-fusion. From a high-level point of view, the $KQ$-softmax-based weighted summation of $V$-representations in Attention corresponds information-gathering from neighbours while the feature-processing via the feed-forward network (FFN) in transformer corresponds to local information fusion. PDMM exploits the Lagrangian multipliers to capture the historical consensus discrepancy in the form of residual errors of the linear edge-constraints, which plays a crucial role for the algorithm to converge. Inspired by PDMM, we propose AttentionX to incorporate the consensus discrepancy in the output update-expression of the standard Attention. The consensus discrepancy in AttentionX refers to the difference between the weighted summation of $V$-representations and scaled $V$-representions themselves. Experiments on ViT and nanoGPT show promising performance.

注意力机制分布式优化Transformer

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