论文标题：单注意力层的无限宽度极限：基于张量程序的分析 作者： Mana Sakai¹,³, Ryo Karakida²,³, Masaaki Imaizumi¹,³ 机构： ¹The University of Tokyo, ²National Institute of Advanced Industrial Science and Technology, ³RIKEN Center for Advanced Intelligence Project 发表信息： 本文发表于第39届神经信息处理系统会议（NeurIPS 2025）。

一、 学术背景与研究目标 本研究属于机器学习理论，特别是深度神经网络的理论分析领域。一个核心的分析范式是研究神经网络在宽度（即每层神经元数量）趋于无穷时的极限行为，这被称为“无限宽度极限”。在此极限下，许多架构（如多层感知机）的神经元预激活值会收敛到高斯分布，这为使用神经网络高斯过程（Neural Network Gaussian Process, NNGP）或神经正切核（Neural Tangent Kernel, NTK）等理论工具进行分析提供了便利。然而，这一基于高斯近似的理论框架在处理Transformer模型中的自注意力层时遇到了根本性困难。除了在“无限多头”或特殊的“1/𝑛缩放”等特殊设定下，现有理论无法准确刻画在标准“1/√𝑛缩放”和有限头数下注意力层的极限分布，而这两种设定正是实践中广泛使用的。

因此，本研究旨在解决一个核心理论空白：在标准的1/√𝑛缩放和有限头数的现实设定下，单层注意力机制在无限宽度极限下的精确输出分布是什么？ 研究的目标是利用张量程序（Tensor Programs）这一强大的数学框架，严格推导并描述这一极限分布，阐明其非高斯特性，并通过实验验证理论预测的有效性。

二、 研究流程与方法 本研究主要分为理论推导和数值验证两大流程。

流程一：理论框架构建与极限分布推导 本研究不涉及传统的实验对象或样本，其“研究对象”是数学定义的神经网络计算图及其在无限宽度下的随机行为。核心方法是利用Greg Yang等人发展的“张量程序”（Tensor Programs）框架，这是一个用于分析各类神经网络架构无限宽度极限行为的统一概率分析框架。

1. 模型设定： 研究首先明确定义了分析对象——一个单层多头注意力机制。输入是一个长度为𝑠的序列，每个token是𝑛维向量。对于每个注意力头𝑎，通过独立的权重矩阵𝑊^𝑄,𝑎, 𝑊^𝐾,𝑎, 𝑊^𝑉,𝑎, 𝑊^𝑂,𝑎 ∈ ℝ^(𝑛×𝑛) 计算查询（Query）、键（Key）、值（Value）和输出投影。缩放点积注意力分数定义为𝑝_(𝑖,𝑗)^(𝑎) = (1/√𝑛) (𝑊^𝑄,𝑎 𝑥_𝑖)^⊤ (𝑊^𝐾,𝑎 𝑥_𝑗)。然后对每行的分数应用softmax得到注意力权重𝐴^(𝑎)，最终多头注意力输出为各头输出的平均投影。所有权重矩阵的元素初始化为独立同分布的高斯随机变量𝑁(0, σ²/𝑛)，输入向量也来自某个多元高斯分布。这是深度学习中标准的初始化方案。

2. 核心挑战与创新定义： 标准张量程序框架能够处理由矩阵乘法（MatMul）和逐元素非线性变换（Nonlin）规则生成的变量，并证明它们在无限宽度下收敛到高斯分布。然而，注意力机制中的点积分数𝑝_(𝑖,𝑗)^(𝑎)是两个由MatMul生成的𝑛维向量的内积，它本身是一个标量随机变量，其行为无法被现有的框架直接捕获。这是分析的核心难点。

3. 引入新的变量类别与极限分布定义： 为了克服这一困难，本研究对张量程序框架进行了关键扩展。研究者在定义3.1中，除了定义标准网络变量ℎ的极限随机变量𝑍ℎ（遵循高斯或非线性变换规则）之外，**特别为点积分数𝑝(𝑖,𝑗)^(𝑎)这类标量内积定义了一个新的极限随机变量集合。论文证明了，在无限宽度极限下，这些点积分数变量共同收敛为一个零均值的高斯向量，并且这个高斯向量与网络中其他所有变量（𝑍_ℎ）的极限分布是统计独立的**。其协方差由相关输入变量的极限分布的四阶矩决定。这是理论上的一个核心创新点。

4. 主要定理的陈述与证明： 基于上述扩展，论文提出了核心定理3.1。该定理指出，对于一个包含上述点积操作的张量程序，任何由网络内部变量{𝑔_𝛼}和点积分数{𝑝}通过伪利普希茨（pseudo-Lipschitz）函数𝜑构造出的新向量ℎ，其经验分布（对𝑛个坐标取平均）在𝑛→∞时，会依分布收敛到条件期望E[𝜓(𝑍_ℎ1, …, 𝑍_ℎ𝑘) | 𝑝1, …, 𝑝𝑟]。这意味着注意力层的输出分布是一个分层（或条件）高斯分布：在给定随机注意力分数（𝑝）的条件下，输出是高斯的；但由于𝑝本身是随机的，边缘分布（即最终输出分布）是非高斯的。

5. 应用到注意力层的具体形式： 在示例3.1中，作者将定理3.1具体应用于标准的多头注意力层。推导出在无限宽度下，注意力输出𝑦_𝑖的第𝛼个分量𝑦𝑖^𝛼依分布收敛到随机变量𝑍(𝑦𝑖)。𝑍(𝑦𝑖)的表达式为(1/√𝐻) Σ(𝑎,𝑗) softmax𝑗(𝑝(𝑖,1)^(𝑎), …, 𝑝(𝑖,𝑠)^(𝑎)) * 𝑍(𝑣 ̃(𝑎,𝑗))。其中，{𝑍(𝑣 ̃(𝑎,𝑗))}是一个与{𝑝(𝑖,𝑗)^(𝑎)}独立的联合高斯向量，而{𝑝_(𝑖,𝑗)^(𝑎)}本身也是一个联合高斯向量。这清晰地展示了输出的非高斯性源于softmax函数的参数（即注意力分数）是随机的。

流程二：数值实验验证 为了验证理论推导的正确性及其在有限宽度下的近似效果，研究进行了系统的数值模拟。

1. 实验对象与设置： “实验对象”是依据理论模型生成的有限宽度（𝑛为有限值，如16, 64, 256, 1024）注意力层。输入通过一个裁剪（clipping）激活函数生成，以确保有界性。权重按理论要求初始化。研究主要观测输出向量𝑦_1的第一个分量𝑦_1^1的分布。

2. 实验方法与分析： 对于每个设定的宽度𝑛和头数𝐻，通过蒙特卡洛方法重复采样权重和输入，生成大量𝑦1^1的样本。同时，利用理论推导出的极限分布公式（通过蒙特卡洛采样𝑝和𝑍(𝑣 ̃)来模拟𝑍_(𝑦_1)）生成极限分布的样本。然后，通过核密度估计（Kernel Density Estimation, KDE）比较有限宽度经验分布与理论极限分布的形态，并计算两者之间的KL散度来衡量收敛程度。

3. 对比实验：

   • 有限宽度效应（图1）： 固定头数𝐻=2，逐渐增加宽度𝑛。结果显示，随着𝑛增大，有限宽度输出𝑦1^1的经验分布曲线逐渐逼近理论极限分布𝑍(𝑦_1)的曲线。KL散度随log₄(𝑛)增大而单调下降，证实了收敛性。
   • 缩放规则的影响（图2a）： 比较了标准1/√𝑛缩放与之前理论中使用的1/𝑛缩放。在𝑛=256时，1/𝑛缩放下的点积分数𝑝_(1,1)^(1)分布高度集中在0附近，而1/√𝑛缩放下的分布则较为分散，与理论极限分布匹配。这解释了为什么1/𝑛缩放会导致极限下注意力机制失效（所有分数趋于0，softmax输出均匀）。
   • 头数的影响（图2b）： 在𝑛=256下，比较了𝐻=1和𝐻=256的情况。有限宽度分布与各自对应的理论极限分布（𝐻=1或256）高度吻合。同时，当𝐻很大时，理论极限分布（𝑍_(𝑦_1)）接近[HBSDN20]中“无限头”极限下的高斯分布，这表明本文的理论将有限头和无限头情况统一在了一个框架下。
   • 低秩注意力（附录B.2，图3）： 模拟了实际Transformer中常见的设置，即头维度𝑛_𝐻固定（如64），总宽度𝑛增加时头数𝐻也同比增加（𝑛 = 𝐻 * 𝑛_𝐻）。结果显示，即使在这种更现实的低秩投影设置下，有限宽度分布仍收敛于调整缩放后的理论极限分布。
   • 鲁棒性测试（附录B.3）： 增加了序列长度𝑠和使用ReLU激活函数，结果均显示理论预测与模拟结果高度一致，证明了理论的鲁棒性。

三、 主要结果 1. 非高斯极限分布的严格证明： 本研究的主要理论结果是定理3.1及其推论3.2。它首次严格证明了，在标准的1/√𝑛缩放和有限头数下，单层注意力机制在无限宽度极限下的输出分布收敛于一个非高斯分布。该分布具有分层高斯结构：在给定随机注意力分数（𝑝）的条件下，输出是高斯的；而𝑝本身是一个高斯随机向量。这一结果从根本上与之前基于NNGP或NTK的高斯近似理论区分开来。 2. 与数值实验的高度一致性： 所有数值实验均强有力地支持了理论预测。图1和图3显示，有限宽度网络的输出分布随着宽度增加，以可量化的方式（KL散度下降）收敛到理论推导的极限分布。图2验证了理论对缩放规则和头数变化的准确描述能力。 3. 统一了有限头与无限头情况： 如备注3.2所述，本研究的极限分布中，注意力分数𝑝的分布与[HBSDN20]中无限头极限下的分布完全一致。同时，有限头输出分布𝑍_(𝑦)的方差理论上等于无限头高斯极限的方差。这表明无限头极限正确捕捉了输出的二阶统计量，而本文的有限头分析则进一步刻画了高阶统计量导致的非高斯形态，提供了一个更完整的图像。 4. 揭示了1/𝑛缩放的局限性： 实验结果（图2a）直观展示了为何1/𝑛缩放（[Yan19b]所用）不适用于分析标准注意力机制：它迫使点积分数在极限下退化为零，从而抹杀了注意力机制根据相似性进行加权选择的核心功能。

四、 结论与意义 本研究得出了一个明确的结论：在现实的标准参数化（1/√𝑛缩放，有限头数）下，单层注意力机制的无限宽度极限分布是非高斯的，它是一个以随机高斯注意力分数为条件的分层高斯分布。

这项研究的价值体现在多个层面： * 科学价值： 它填补了深度神经网络无限宽度理论中的一个关键空白，首次为理解注意力层在最具现实意义的设定下的渐进行为提供了严格的理论基础。它表明，对于包含注意力机制的架构，传统的高斯过程或神经正切核理论可能不再适用，需要发展新的、能够处理非高斯极限分布的理论框架。 * 理论创新价值： 本研究发展了一种新颖的分析技术，通过将注意力点积分数视为一类新的、与网络其他部分独立的极限随机变量，成功地将张量程序框架扩展到了包含内积运算的架构。这为分析其他具有类似交互结构的网络组件提供了方法论上的借鉴。 * 应用价值： 对注意力层精确极限分布的理解，有助于从理论上解释其独特的行为，例如为什么注意力机制更容易避免信号传播中的秩崩溃（rank collapse），或者其优化景观（optimization landscape）可能具有何种特性。这为设计更稳定、更高效的Transformer架构提供了理论指导。 * 未来研究方向： 本文为分析深度Transformer（堆叠的注意力层和MLP层）的无限宽度极限奠定了基础。作者推测，在深度架构中，不仅注意力层，MLP层的极限分布也可能变为非高斯。这预示着需要一个全新的、超越现有高斯基理论的统一框架来理解深度Transformer的动态特性。

五、 研究亮点 1. 重要的理论发现： 首次严格推导出标准设定下注意力层无限宽度极限的精确非高斯分布形式，解决了该领域一个长期存在的公开问题。 2. 方法上的新颖性： 提出了针对点积（内积）运算的新分析技术，扩展了强大的张量程序框架的适用范围，这是理论推导成功的关键。 3. 理论与实验的紧密结合： 不仅提供了严格证明，还通过系统、全面的数值实验验证了理论在有限宽度下的准确性和鲁棒性，涵盖了不同宽度、头数、缩放规则、甚至实际中的低秩注意力设置。 4. 澄清与统一了先前工作： 明确了先前两种简化设定（无限头和1/𝑛缩放）的局限性，并将有限头情况与无限头极限在同一个理论框架内联系起来，揭示了它们之间的关系。 5. 对后续研究的启发性： 明确指出了现有高斯基理论在分析Transformer类模型时的不足，并为其超越高斯近似、发展更一般的非高斯极限理论铺平了道路，具有重要的引领意义。