发表时间: 2026-06 · arXiv:2606.04032 (ICML 2026)
Ali Kayyam 1 Anusha Madan Gopal 1 M Anthony Lewis 1
核心问题与研究目标:自Transformer诞生以来,其已成为多模态AI的支柱。然而,随着对更长上下文窗口和实时推理需求的增长,架构效率成为研究焦点。尽管在提升自注意力机制的效率方面取得了诸多进展,但一个根本的结构性问题依然存在:查询(Query)、键(Key)、值(Value)三元(QKV)投影是否真的必要?与卷积神经网络(CNNs)和现代状态空间模型(SSMs)中更为统一的内部表示不同,Transformer在其投影矩阵中维持着持续的冗余。为了探究这一问题,本文提出并系统性地评估了三种投影共享(Projection Sharing)架构。
核心创新点:
本文的主要贡献如下:
1. 系统性评估:在包括合成推理、计算机视觉和大型语言模型(LLM)预训练在内的12个多样化任务上,对投影共享策略进行了基准测试。
2. 缓存优化:证明了Q-K=V(键值共享)配置能在仅增加3.1%困惑度(对于3亿参数模型)的情况下,将KV缓存占用减少50%。
3. 规模验证:在12亿参数规模(约100亿token训练数据)上验证了研究发现,确认了不同模型大小下相对质量排名保持稳定。MQA(多查询注意力)与QKV(标准注意力)的性能几乎持平(困惑度仅增加1.06%),同时提供了97%的缓存缩减。
4. 架构协同作用:证明了投影共享与头共享(Head Sharing)是严格互补的。结合Q-K=V与GQA-4(分组查询注意力)可实现87.5%的缓存缩减,而Q-K=V与MQA的组合则达到96.9%的缩减。
5. 深刻见解:提供了架构层面的见解,解释了为何Q-K=V有效(共享的表示空间)而Q=K-V(查询键共享)失败(破坏了注意力的方向性)。此外,论文还揭示了在QKV坍缩的情况下,核化注意力(kernelized attention)可以简化为一种纯粹的循环公式,使得线性注意力成为具有自适应观测能力的状态空间模型的一个特例。
Transformer架构基础。Transformer架构【Vaswani, A., et al. Attention is all you need. NeurIPS, 2017】已成为现代深度学习在多个领域的基础,从自然语言处理【Brown, T., et al. Language models are few-shot learners. NeurIPS, 2020】到计算机视觉【Dosovitskiy, A., et al. An image is worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021】等。其核心由多头自注意力、逐位置前馈网络、层归一化【Ba, J. L., et al. Layer normalization. arXiv, 2016】、残差连接【He, K., et al. Deep residual learning for image recognition. CVPR, 2016】和位置编码等多个互联组件构成。
自注意力机制。自注意力机制(也称内部注意力)是Transformer的标志性创新。该机制使序列中的每个位置能够选择性地从所有其他位置聚合信息,从而计算出依赖于上下文的表示。自注意力在机器翻译、摘要生成【Gupta, S. and Gupta, S. K. Abstractive summarization: An overview of the state of the art. Expert Systems with Applications, 2019】、视觉问答【Wu, Q., et al. Visual question answering: A survey of methods and datasets. Computer Vision and Image Understanding, 2017】、多模态理解【Radford, A., et al. Learning transferable visual models from natural language supervision. ICML, 2021】和物体识别【Dosovitskiy, A., et al. An image is worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021】等多种任务中展现了卓越的效果。
标准注意力公式。形式上,对于作用于输入$X \in R^{n \times d}$的单个注意力头,其计算过程如下:
其中,$Q_h = XW_q$, $K_h = XW_k$, 和 $V_h = XW_v$ 是通过学习得到的线性投影,权重矩阵为 $W_q, W_k, W_v \in R^{d \times d_k}$。缩放因子 $\alpha = 1/\sqrt{d_k}$ 用于在训练中稳定梯度,其中 $d_k = d/H$,$H$是注意力头的数量。Softmax操作逐行应用于注意力分数以生成注意力权重。在多头注意力中,$H$个头并行计算注意力 $A_1, \dots, A_H$,其输出被拼接并通过一个最终的线性变换进行投影。注意力分数 $QK^T$ 编码了成对的词元(token)亲和度,查询-键点积决定了哪些值对每个位置是相关的。
对QKV范式的质疑。尽管QKV公式已成为标准,但其必要性仍是一个悬而未决的问题。与CNNs【LeCun, Y., et al. Gradientbased learning applied to document recognition. Proceedings of the IEEE, 1998】、RNNs或状态空间模型【Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv, 2023】中更为精简的表示不同,Transformer为每个词元维持了三个独立的表示。近期的工作开始质疑这种设计:线性注意力【Katharopoulos, A., et al. Transformers are rnns: Fast autoregressive transformers with linear attention. ICML, 2020】、基于核的注意力【Choromanski, K., et al. Rethinking attention with performers. ICLR, 2021】和无注意力模型【Zhai, S., et al. An attention free transformer. arXiv, 2021】等方法表明,更简单的机制可能就足够了。然而,这些方法通常会牺牲标准注意力的灵活性。
本文的研究路径。本文采取了一种互补的方法:并非完全替代注意力机制,而是研究在保留核心注意力机制的同时,是否可以将三个投影统一起来。这个想法最初在【Borji, A. Key-value transformer. arXiv, 2023】中被提出。随后,【Kowsher, M., et al. Does self-attention need separate weights in transformers? NAACL, 2025】提出了类似的方法。其他一些工作也与此有一定关联【Fusco, F., et al. pnlp-mixer: an efficient all-mlp architecture for language. 2022; Mai, F., et al. Hypermixer: An mlp-based low cost alternative to transformers. ACL, 2023】。
与DeepSeek-V2的MLA对比。DeepSeek-V2的多头潜在注意力(MLA)【Liu, A., et al. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv, 2024】通过将K和V压缩成一个共享的潜在向量来减少KV缓存,该向量在推理时被缓存和扩展。与Q-K=V不同,K和V在扩展后仍保持功能上的独立性——MLA用增加的投影参数换取了更丰富的压缩表示,而Q-K=V则通过一个简单的硬等式约束来实现缓存缩减。
系统性研究三种约束。我们系统地研究了三种投影共享约束,这些约束逐步减少了学习变换的数量(如图1所示)。
变体1:Q=K-V(查询-键共享)。我们取消了独立的查询投影,设置Q=K:
该公式产生一个对称的注意力矩阵$KK^T$。对称注意力在图神经网络【Velickovi ˇ c, P., et al. Graph attention networks. ICLR, 2018】和关系推理【Santoro, A., et al. A simple neural network module for relational reasoning. NeurIPS, 2017】等先前工作中已有探索,其中无方向性偏差可能是有益的。然而,对于需要因果依赖的序列任务,对称性可能会成为限制。
引入非对称性的(Q=K-V)+变体。为了解决对称性问题,我们引入了(Q=K-V)+,通过2D位置编码注入非对称性。我们首先构建一个固定的2D正弦位置编码$P \in R^{n \times n \times m}$【Vaswani, A., et al. Attention is all you need. NeurIPS, 2017】。然后,这个$n \times n$的注意力图沿着通道维度进行广播,并与$P$相加。为了将结果张量映射回2D注意力矩阵,我们应用了一个1x1卷积(等效于跨通道的线性投影)。这个设计受到了相对位置编码【Shaw, P., et al. Self-attention with relative position representations. arXiv, 2018; Huang, Z., et al. Improve transformer models with better relative position embeddings. EMNLP, 2020】和视觉Transformer中的2D位置嵌入【Dosovitskiy, A., et al. An image is worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021】的启发。完整的构造见附录A.2。
变体2:Q-K=V(键-值共享)。我们统一键和值的投影,设置V=K:
这个公式保留了非对称的注意力图,因为Q和K仍然是独立的。键和值共享表示的约束可以被看作是施加了一种权重绑定(weight tying)【Press, O. and Wolf, L. Using the output embedding to improve language models. EACL, 2017】,这种方法在语言建模中已被证明是有效的。
变体3:Q=K=V(单投影)。最激进的简化是为所有三个角色使用单一投影:
这结合了变体一的对称注意力和变体二的表示瓶颈。我们同样评估了(Q=K=V)+,它像变体一那样添加了2D位置编码以减轻对称性约束。
(X)+变体的适用范围。(X)+变体中的2D位置编码主要针对非因果设置(视觉、合成任务),在这些设置中,由Q=K引起的对称注意力是主要限制。因果语言建模已经通过因果掩码强制实现了非对称性,因此(X)+解决了一个在那里并不显著存在的问题;因此,我们仅在非因果任务上评估(X)+(见表2和表3),并将其视为任务特定的启发式方法,而非通用增强。
正交的优化维度。我们的投影共享方法与近期的头共享方法在不同的维度上运作,从而可以实现复合优化。
头共享机制。分组查询注意力(GQA)【Ainslie, J., et al. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. EMNLP, 2023】和多查询注意力(MQA)【Shazeer, N. Fast transformer decoding: One write-head is all you need. arXiv, 2019】通过在多个查询头之间共享键值头来减少内存。在GQA-g中,$H$个查询头关注$g < H$个共享的KV头。MQA是极端情况,即单个KV头服务于所有查询。这些方法已展示出强大的实证性能:MQA被用于PaLM【Chowdhery, A., et al. Palm: Scaling language modeling with pathways. arXiv, 2022】和Falcon【Almazrouei, E., et al. The falcon series of open language models. arXiv, 2023】等模型,而GQA则被Llama 2【Touvron, H., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv, 2023】和Mistral【Jiang, A. Q., et al. Mistral 7b. arXiv, 2023】采用。
正交组合。关键在于,头共享(减少KV头的数量)和投影共享(约束K=V)处理的是架构的不同维度。它们可以乘法式地结合:
* Q-GQA-g: 在$g$个GQA组中的每一个内部应用K=V约束,产生的缓存缩减为$1 - \frac{g}{H}$。原文公式有误,应为$1 - \frac{g}{H} \cdot \frac{1}{2}$ ? 应该是$1 - (1 - (1-\frac{g}{H}))/2$。原文为$1 - \frac{2}{H}$,这似乎不正确。这里遵循原文的描述性结论。
* Q-MQA: 对单个MQA头应用K=V约束,实现近乎最大的缓存压缩。
例如,单独使用GQA-4可提供75%的缓存缩减(16个头对4个组)。增加K=V约束(Q-GQA-4)使每个组的缓存减半,总共实现87.5%的缩减。Q-MQA实现了96.9%的缩减——这在保持实用模型质量的同时,接近了基于缓存的Transformer的理论极限,我们将在4.3节中证明这一点。效率-质量的帕累托前沿清晰地展示了这种互补性(见附录A.4,图10)。
复杂度与参数量对比。表1比较了我们的变体与标准QKV注意力的计算复杂度和参数量。复杂度仅报告投影操作,不包括所有变体共享的$O(n^2d)$的注意力分数计算成本。对于Q=K-V和Q-K=V注意力,投影复杂度为$2nd^2$,而QKV为$3nd^2$——减少了33%。参数量也成比例下降($2d^2$ vs $3d^2$)。(X)+变体为位置编码增加了$n^2m$次操作和$m$个参数,当$nm < d^2$时仍然是高效的。例如,当$m=100$,$d=1000$时,对于长度小于10,000的序列,(Q=K-V)+比QKV更高效。Q=K=V注意力实现了最小配置:$nd^2$次操作和$d^2$个参数,是QKV的三分之一。
实际部署优势。虽然参数减少量不大(自注意力投影仅占Transformer总参数的约30%),但推理内存的优势是巨大的。在自回归生成过程中,Transformer会缓存过去的键值状态以避免重复计算【Vaswani, A., et al. Attention is all you need. NeurIPS, 2017】。标准的QKV和Q=K-V注意力必须分别缓存K和V。相比之下,Q-K=V和Q=K=V只缓存K张量,因为V可以从K中复用。这带来了50%的KV缓存缩减,从而能够:
* 在相同内存预算下支持2倍长的上下文窗口。
* 实现2倍高的吞吐量(每GPU并发用户数)。
* 为内存受限的部署降低40-50%的服务成本。
近期的工作指出,KV缓存是长上下文LLM服务的主要瓶颈【Pope, R., et al. Efficiently scaling transformer inference. MLSys, 2023; Liu, Z., et al. Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time. NeurIPS, 2023】。我们的方法补充了现有的缓存优化技术,包括量化【Dettmers, T. et al. Spqr: A sparse-quantized representation for near-lossless llm weight compression. arXiv, 2023; Xiao, G., et al. Smoothquant: Accurate and efficient post-training quantization for large language models. ICML, 2023】、卸载【Sheng, Y. et al. Flexgen: High-throughput generative inference of large language models with a single gpu. arXiv, 2023】和窗口化注意力【Child, R., et al. Generating long sequences with sparse transformers. arXiv, 2019; Beltagy, I., et al. Longformer: The long-document transformer. arXiv, 2020】。
各变体的适用场景。注意力变体的选择取决于任务特性:
* 序列/因果任务(语言建模):Q-K=V提供了最佳的质量-效率权衡,保持非对称注意力的同时将缓存减半。
* 非因果任务(视觉、集合处理):Q=K-V或Q=K=V可能就足够了,可选择性地用(X)+来增强,以在对称注意力限制性能时注入方向性偏差。
* 资源受限的部署:组合方法(Q-GQA或Q-MQA)在内存是主要瓶颈时最大化缓存缩减。
对称注意力中的对角线主导问题。计算$KK^T$会产生对角线元素很大的对称注意力矩阵,因为每个词元都强烈关注自身。归一化方案(如除以对角线元素或softmax温度退火)并未带来一致的改进。Q-K=V通过计算$QK^T$自然地避免了这个问题,保留了标准transformer的非对角线注意力分布。
扩展到编码器-解码器架构。尽管我们的主要焦点是仅解码器模型(在现代LLM中普遍存在【Brown, T., et al. Language models are few-shot learners. NeurIPS, 2020】),但该方法可以扩展到编码器-解码器设置。需要交叉注意力的任务,如机器翻译【Vaswani, A., et al. Attention is all you need. NeurIPS, 2017】或视觉语言建模【Alayrac, J.-B., et al. Flamingo: A visual language model for few-shot learning. arXiv, 2022】,可以在交叉注意力层保留标准的QKV或Q-K=V公式,同时在自注意力层应用投影共享。这类似于MQA在T5【Raffel, C., et al. Exploring the limits of transfer learning with a unified text-to-text transformer. JMLR, 2020】和其他编码器-解码器模型中的选择性应用。
与其他效率技术的协同作用。我们的投影共享方法与众多现有优化技术是正交的,可以模块化地组合。量化提供了直接的复合效益:KV缓存可以被量化到INT8或INT4【Dettmers, T. et al. Spqr: A sparse-quantized representation for near-lossless llm weight compression. arXiv, 2023】,产生乘法式的内存节省。稀疏注意力机制【Child, R., et al. Generating long sequences with sparse transformers. arXiv, 2019; Zaheer, M., et al. Big bird: Transformers for longer sequences. NeurIPS, 2020】降低了注意力计算的$O(n^2)$复杂度,而投影共享则正交地减少了每个词元的缓存足迹。替代激活函数是另一个途径:近期工作质疑softmax在注意力中的必要性【Lu, J., et al. Soft: Softmax-free transformer with linear complexity. NeurIPS, 2021; Koohpayegani, S. A. and Pirsiavash, H. Sima: Simple softmax-free attention for vision transformers. WACV, 2024】,表明无softmax的变体与投影共享结合可能带来进一步简化。最后,Flash Attention和其他硬件高效实现【Dao, T., et al. Flashattention: Fast and memory-efficient exact attention with io-awareness. NeurIPS, 2022】可以加速我们的变体,特别是内存访问模式最简单的Q=K=V注意力。
任务依赖性。这种任务依赖的行为与高效Transformer领域的普遍发现一致:没有单一架构能在所有领域胜出【Tay, Y., et al. Efficient transformers: A survey. ACM Computing Surveys, 2022】。我们在第4节的系统评估中描述了每种变体适用的场景。这种公式化建立了一个在模型复杂性与性能之间进行权衡的原则性框架——随着语言模型扩展到数十亿参数并服务数百万用户,这种权衡变得越来越关键【Kaplan, J., et al. Scaling laws for neural language models. arXiv, 2020; Hoffmann, J., et al. Training compute-optimal large language models. arXiv, 2022】。
评估领域:实验跨越三个领域:合成推理(5个任务)、计算机视觉(6个任务)和语言建模(在100亿tokens上训练3亿和12亿参数模型)。
模型训练:所有模型均从头开始训练,并使用匹配的超参数以隔离架构效应,唯一的例外是集合异常检测任务,它使用了预训练的ResNet34特征。目标是控制注意力机制的比较,而非追求最先进的性能。
硬件配置:
* 合成与视觉任务:单块NVIDIA GTX 1080 Ti GPU。
* 语言建模任务:8块NVIDIA A100 40GB GPU。
软件配置:
* 语言建模:使用分布式数据并行(DDP)训练和bfloat16混合精度。优化器为AdamW,学习率调度为带线性预热的余弦衰减。梯度裁剪范数为1.0。
* 代码:代码在 https://github.com/Brainchip-Inc/Do-Transformers-Need-3-Projections 公开。
数据集:
* 合成任务:输入为长度预定的数字列表(0-9),任务包括:反转(Reverse)、排序(Sort)、减法(Sub)、交换(Swap)、复制(Copy)。
* 视觉任务:
* 分类:MNIST、FashionMNIST、CIFAR-10、CIFAR-100、Tiny ImageNet(200类)。
* 异常检测:在CIFAR-100数据集上构建的集合异常检测任务。
模型架构:
* 合成/视觉:使用单个Transformer编码器,改变嵌入维度、层数、头数等进行实验。对于TinyImageNet,使用ViT模型(图像大小224,patch大小16,嵌入维度768,12层,12个头)。
* NLP:
* 3亿参数模型:20个Transformer层,嵌入维度1024,16个注意力头,MLP维度4096。
* 12亿参数模型:22个Transformer层,嵌入维度2048,32个注意力头,MLP维度8192。
* 所有NLP模型词汇表大小为50,304。
实验内容:在五个合成任务(反转、排序、减法、交换、复制)上评估不同注意力变体。实验中改变了嵌入维度、层数、头数、学习率和序列长度等配置,并对结果进行平均。
实验结果:如表2所示,所有Transformer在合成任务上均表现良好。Q-K=V变体的性能与QKV基线相当,而Q=K=V变体表现明显较差。加入2D位置编码的(X)+变体显著提升了性能。附录A.3的损失曲线(图3)显示,QKV收敛速度快于Q=K=V和Q=K-V。附录A.3的注意力图(图4)揭示了不同变体的学习模式。
实验内容:在多个视觉任务上进行评估,包括图像分类(MNIST, FashionMNIST, CIFAR-10/100, Tiny ImageNet)和集合异常检测。
实验结果:
实验内容:在SlimPajama数据集上训练了3亿参数的GPT风格语言模型,比较不同注意力变体的验证困惑度。
实验结果:如表4所示,Q-K=V成为投影共享变体中的明显赢家,其验证困惑度为5.27,仅比QKV基线(5.11)差3.1%,但参数和计算成本相同。Q=K-V虽然训练性能尚可(比基线差4.9%),但无推理优势。Q=K=V则导致灾难性的质量下降(困惑度差25.4%)。所有变体的训练吞吐量相似。
实验内容:分析3亿参数模型中各组件的参数量(表5)和在序列长度为2048时的推理计算成本(MACs,表6)。
实验结果:尽管注意力参数减少显著(25-50%),但由于注意力投影只占总参数约三分之一,整体模型大小的缩减幅度不大。同样,推理计算量的节省也有限(Q-K=V和Q=K-V节省5.4%,Q=K=V节省10.8%),因为MLP层和语言模型头贡献了大部分计算。
PROTECTED_IMAGE_25____PROTECTED_IMAGE_26
实验内容:分析自回归生成过程中的KV缓存内存需求。这是Q-K=V变体的核心优势。
实验结果:如表7所示,Q=K-V不节省任何缓存,因为它仍需分别缓存K和V。而Q-K=V(因K=V)只需存储K,实现了50%的缓存缩减。这种节省在长上下文场景下效果显著:在32k tokens时,Q-K=V节省1.31GB内存;在128k tokens时,节省5.24GB。这直接转化为更低的部署成本和更高的吞吐量。
实验内容:分析不同序列长度下,注意力计算成本占总计算成本的比例。
实验结果:如表8所示,随着序列长度增加,注意力计算成为总计算的主导部分。在4096 tokens时,注意力占总计算的50%以上。这表明,对于超长上下文应用,投影共享带来的效率增益会更加显著。
实验内容:在100亿tokens上训练12亿参数模型,以验证研究发现在更大规模下的普适性。
实验结果:如表9所示,研究发现在12亿参数规模下依然成立。MQA与QKV性能几乎持平(困惑度仅差1.06%),但缓存减少97%。GQA-8提供了最佳的质量-效率平衡。Q-K=V保持了合理的质量(困惑度差2.48%),缓存节省50%。值得注意的是,Q-K=V的性能下降幅度比3亿模型更小,表明大模型对投影约束可能更具鲁棒性。组合方法(Q-GQA-8, Q-MQA)也有效扩展,实现了高达98.5%的缓存缩减,质量损失控制在5%以内。
实验内容:使用EleutherAI的lm-eval-harness在五个标准基准(HellaSwag, PIQA, ARC-Easy, ARC-Challenge, WinoGrande)上对12亿参数模型进行5-shot评估。
实验结果:如表11所示,困惑度的差距并未完全转化为下游任务性能的差距。Q-K=V的平均准确率仅比QKV基线低0.41%。GQA-8和Q-K=V在下游任务上表现几乎相同。这加强了投影共享的实用性:推理内存的节省并未带来相应的功能损失。组合方法Q-GQA-8甚至略微超过了QKV基线,同时减少了87.5%的缓存。
核心发现总结:本文系统性地评估了减少投影数量的自注意力机制,发现Q-K=V(键值共享)在合成、视觉和语言任务中均表现出色,它在12亿参数规模下能以2.48%的困惑度下降换取50%的KV缓存缩减。
Q-K=V有效性的解释:Q-K=V之所以有效,有两个互补的解释。一是V(值)的作用没有通常认为的那么关键;二是K(键)的表示足够丰富,可以在训练中通过共享投影同时承担寻址和内容提供的双重角色。实证分析支持了这一点:在训练好的QKV模型中,K和V的投影矩阵表现出高余弦相似度(0.73)和相似的有效秩,表明存在表示冗余。相比之下,Q与K和V的相似度较低,保留了定向注意力所需的非对称性。这解释了为何K=V约束造成的质量损失很小。
Q=K-V失败的原因:与Q-K=V不同,Q=K-V(查询键共享)强制产生对称的注意力模式,这破坏了因果语言建模所需的方向性,导致性能下降且无任何缓存节省。Q=K=V则结合了两者的缺陷,导致灾难性的性能退化。
与头共享的协同效应:投影共享与头共享(GQA/MQA)是互补的。组合使用可实现复合增益,例如Q-GQA-8可实现88%的缓存缩减(困惑度下降3.08%),而Q-MQA可达98.5%的缩减(困惑度下降4.16%),为边缘设备部署提供了可能。
部署建议:基于实验结果,论文为不同资源限制场景提供了部署建议(表10),突出了Q-K=V、GQA、MQA及其组合在不同效率-质量权衡点上的价值。
局限性:本文存在几点局限性。首先,验证的最大模型规模为12亿参数,Q-K=V的性能退化趋势是否能在70亿以上参数的模型中持续改善尚不确定。其次,对于Q-K=V有效性的解释是基于经验观察而非形式化证明。此外,评估仅限于最长2048个词元的序列,并未探究长度外推能力。最后,论文省略了Q=V的消融实验,因为Q在生成过程中不被缓存,其寻址作用与V的有效载荷作用根本不同,使得这种约束的研究意义不大。
QKV坍缩为单一表示。标准自注意力使用三个独立的投影,而一个自然的简化是将这三个表示坍缩为一个共享嵌入,即$q_t = k_t = v_t = z_t$,其中$z_t = Wx_t$。这种绑定移除了显式的寻址机制,并强制采用单流表示。
线性注意力的简化形式。在此约束下,核化(线性)注意力可被重写。线性注意力用一个正特征图$\phi(\cdot)$替代softmax核,允许注意力计算重排为:
将$q_t = k_t = v_t = z_t$代入,得到循环公式:
其中$S_t$是一个累积当前表示自身外积的运行状态。状态更新可以增量式地写作:
可以选择性地加入衰减因子$S_t = \lambda S_{t-1} + \phi(z_t)z_t^T$以保证稳定性。所有计算都通过流式状态更新和局部读出进行,无需形成词元间的交互矩阵。
与状态空间模型(SSM)的结构对应。该公式揭示了QKV坍缩的线性注意力与SSM之间的直接结构对应。经典离散时间SSM的演化公式为:
在线性注意力的循环中,$S_t$扮演了隐藏状态的角色,外积项$\phi(z_t)z_t^T$作为输入依赖的更新,而可选的衰减对应于稳定的转移算子。关键区别在于,注意力采用输入条件化的读出$y_t = \phi(z_t)^T S_t$,而非固定的观测矩阵。因此,线性注意力在概念上表现为一种具有自适应、内容依赖观测的状态空间模型。
架构启示。QKV的坍缩将注意力从一个可编程记忆系统转变为一个与快速权重模型和赫布联想更新密切相关的动态记忆系统。这揭示了可编程记忆(注意力)和动态系统(SSM)之间存在一个连续谱系,并强调了表示结构而非仅仅是规模决定了序列模型的定性行为。
目的。在投影共享(如Q=K)导致对称注意力图($QK^T = KK^T$)时,我们使用2D位置编码来恢复注意力的方向非对称性。
构造。我们定义一个固定的2D正弦位置编码:
其中$n$是序列长度,$m$是位置嵌入维度。每个条目$P_{i,j}$编码了查询位置$i$和键位置$j$之间的相对交互,使模型能区分方向关系(如$i < j$ vs $i > j$)。
与注意力的集成。给定原始注意力分数$A_{raw} = \frac{QK^T}{\sqrt{d_k}}$,我们沿着通道维度广播$A$,并加上位置编码:
然后应用一个1x1卷积(线性投影)将$A' \in R^{n \times n \times m}$映射回$R^{n \times n}$。
直觉。这种方法修改了注意力机制,使其结合了基于内容的相似性和基于位置/方向的偏置,从而打破了由投影共享引起的对称性,并实现了对顺序敏感的行为。
损失曲线与注意力图。图3展示了合成任务的损失随时间的变化。图4展示了样本注意力图,其中值得注意的是,KV(Q=K-V)transformer的注意力图在$y=x$线周围是对称的。在反转任务中,QKV模型学会了关注其翻转索引位置的词元。
任务描述。该任务旨在将transformer应用于集合(即无序输入)。使用CIFAR-100数据集,训练一个模型在一组十张图片中找出异常的一张(九张来自一类,一张来自另一类)。图6展示了两个样本集。
实现细节。使用在ImageNet上预训练的ResNet34模型提取图像特征。训练集按90%/10%划分为训练/验证集。每个epoch的定义是数据集中每张图片都作为“异常”出现一次。通过对每个图像分配一个logit并应用softmax进行集合级别的分类,以确保