arXiv:2412.17810cs.LG2024-12被引 64

提出线性复杂度注意力机制,让Transformer更快更可解释。

Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction

  • 通过变分优化编码率减少目标,设计新注意力模块TSSA。
  • TSSA计算与内存复杂度均为线性,性能媲美标准Transformer。
  • 适合追求高效、可解释性的视觉/语言/长序列任务研究者。

注意力机制是Transformer架构的核心,但在多种任务中表现出卓越性能的同时,其计算开销随令牌数量呈二次增长。本文提出一种新型注意力机制,使计算复杂度随令牌数线性增长。基于先前工作,我们发现变压器结构自然源自白盒架构设计,即每一层实现最大编码率减少(MCR²)目标的增量优化步骤。具体地,我们推导出MCR²目标的新变分形式,并证明沿该变分目标进行梯度下降展开后,可得到一种名为令牌统计自注意力(TSSA)的新注意力模块。TSSA具有线性计算和内存复杂度,且完全摒弃了传统方法中令牌间两两相似度计算的方式。在视觉、语言及长序列任务上的实验表明,仅将标准自注意力替换为TSSA(即令牌统计Transformer,ToST),即可在保持竞争力性能的同时显著提升计算效率并增强可解释性。结果也对‘两两相似度注意力是Transformer成功的关键’这一普遍认知提出了质疑。代码将发布于https://github.com/RobinWu218/ToST。

原文摘要 · Abstract (English)

The attention operator is arguably the key distinguishing factor of transformer architectures, which have demonstrated state-of-the-art performance on a variety of tasks. However, transformer attention operators often impose a significant computational burden, with the computational complexity scaling quadratically with the number of tokens. In this work, we propose a novel transformer attention operator whose computational complexity scales linearly with the number of tokens. We derive our network architecture by extending prior work which has shown that a transformer style architecture naturally arises by "white-box" architecture design, where each layer of the network is designed to implement an incremental optimization step of a maximal coding rate reduction objective (MCR$^2$). Specifically, we derive a novel variational form of the MCR$^2$ objective and show that the architecture that results from unrolled gradient descent of this variational objective leads to a new attention module called Token Statistics Self-Attention (TSSA). TSSA has linear computational and memory complexity and radically departs from the typical attention architecture that computes pairwise similarities between tokens. Experiments on vision, language, and long sequence tasks show that simply swapping TSSA for standard self-attention, which we refer to as the Token Statistics Transformer (ToST), achieves competitive performance with conventional transformers while being significantly more computationally efficient and interpretable. Our results also somewhat call into question the conventional wisdom that pairwise similarity style attention mechanisms are critical to the success of transformer architectures. Code will be available at https://github.com/RobinWu218/ToST.

Transformer注意力机制线性复杂度可解释性

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