Softmax as Linear Attention in the Large-Prompt Regime: a Measure-based Perspective

📄 arXiv: 2512.11784v1 📥 PDF

作者: Etienne Boursier, Claire Boyer

分类: cs.LG, stat.ML

发布日期: 2025-12-12


💡 一句话要点

提出基于测度的统一框架,分析大提示下的Softmax注意力机制。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: Softmax注意力 Transformer模型 线性注意力 无限提示 集中不等式

📋 核心要点

  1. Softmax注意力机制的非线性特性使其理论分析面临挑战。
  2. 论文提出基于测度的统一框架,将无限提示下的Softmax注意力近似为线性算子。
  3. 通过非渐近集中界,量化了有限提示模型向无限提示模型的收敛速度。

📝 摘要(中文)

Softmax注意力是Transformer架构的核心组成部分,但其非线性结构给理论分析带来了显著挑战。本文开发了一个统一的、基于测度的框架,用于研究有限和无限提示下的单层Softmax注意力。对于独立同分布的高斯输入,我们利用Softmax算子在无限提示极限下收敛到作用于底层输入token测度的线性算子的事实。基于此,我们建立了Softmax注意力输出和梯度的非渐近集中界,量化了有限提示模型逼近其无限提示对应模型的速度,并证明了这种集中在具有亚高斯token的一般上下文学习设置中的整个训练轨迹上保持稳定。在线性上下文回归的情况下,我们使用易于处理的无限提示动态来分析有限提示长度的训练。我们的结果允许为线性注意力开发的优化分析直接转移到提示足够长的Softmax注意力,表明大提示Softmax注意力继承了其线性对应物的分析结构。这反过来又为研究大提示机制中Softmax注意力层的训练动态和统计行为提供了一个原则性的、广泛适用的工具包。

🔬 方法详解

问题定义:论文旨在解决Softmax注意力机制由于其非线性结构,在理论分析上存在的困难。现有方法难以有效分析其训练动态和统计行为,尤其是在长文本提示的情况下。

核心思路:核心思路是将Softmax注意力在无限提示长度的极限情况下,近似为一个作用于输入token测度的线性算子。通过这种近似,可以将复杂的非线性Softmax注意力问题转化为更易于分析的线性问题。

技术框架:论文构建了一个基于测度的统一框架,用于分析单层Softmax注意力。该框架首先在有限提示长度下研究Softmax注意力,然后推导出无限提示长度下的极限形式。通过比较有限提示和无限提示下的行为,可以量化有限提示模型逼近其无限提示对应模型的速度。该框架还包括对Softmax注意力输出和梯度的集中界分析。

关键创新:关键创新在于将非线性的Softmax注意力机制与线性算子联系起来,从而可以使用线性分析工具来研究Softmax注意力。此外,论文还提供了非渐近的集中界,可以量化有限提示长度下的近似误差。

关键设计:论文假设输入token是独立同分布的高斯变量或亚高斯变量。在无限提示长度下,Softmax算子收敛到一个线性算子,该算子作用于输入token的经验测度。论文使用集中不等式来分析有限提示长度下的Softmax注意力输出和梯度的行为,并证明了其在训练过程中的稳定性。

📊 实验亮点

论文建立了Softmax注意力输出和梯度的非渐近集中界,量化了有限提示模型逼近其无限提示对应模型的速度。结果表明,当提示足够长时,Softmax注意力继承了线性注意力的分析结构,从而可以将为线性注意力开发的优化分析直接应用于Softmax注意力。

🎯 应用场景

该研究成果可应用于自然语言处理领域,特别是Transformer模型的优化和理论分析。通过将Softmax注意力近似为线性算子,可以更深入地理解Transformer模型的训练动态和泛化能力,并为设计更高效的Transformer变体提供理论指导。此外,该方法还可以应用于其他使用注意力机制的领域,例如图像识别和语音识别。

📄 摘要(原文)

Softmax attention is a central component of transformer architectures, yet its nonlinear structure poses significant challenges for theoretical analysis. We develop a unified, measure-based framework for studying single-layer softmax attention under both finite and infinite prompts. For i.i.d. Gaussian inputs, we lean on the fact that the softmax operator converges in the infinite-prompt limit to a linear operator acting on the underlying input-token measure. Building on this insight, we establish non-asymptotic concentration bounds for the output and gradient of softmax attention, quantifying how rapidly the finite-prompt model approaches its infinite-prompt counterpart, and prove that this concentration remains stable along the entire training trajectory in general in-context learning settings with sub-Gaussian tokens. In the case of in-context linear regression, we use the tractable infinite-prompt dynamics to analyze training at finite prompt length. Our results allow optimization analyses developed for linear attention to transfer directly to softmax attention when prompts are sufficiently long, showing that large-prompt softmax attention inherits the analytical structure of its linear counterpart. This, in turn, provides a principled and broadly applicable toolkit for studying the training dynamics and statistical behavior of softmax attention layers in large prompt regimes.