Learning Theory of Transformers: Local-to-Global Approximation via Softmax Partition of Unity¶

会议: NeurIPS 2026
arXiv: 2605.08811
OpenReview: https://openreview.net/forum?id=R1SSvVTsSW
论文: https://arxiv.org/abs/2605.08811
领域: 优化/理论
关键词: Transformer 学习理论, 局部到全局逼近, Softmax 单位分解, 泛化误差界, 流形适应性

一句话总结¶

本文提出了一种基于 Softmax 单位分解（POU）的 Transformer 构造性逼近理论框架，首次证明了仅需 2 层编码器块、标准 Softmax 点积注意力和正弦位置编码的浅宽稠密 Transformer 即可达到 $\alpha$-Hölder 连续函数的最优逼近阶 $\mathcal{O}(\varepsilon^{-d/\alpha})$ 与近极大极小泛化率 $\mathcal{O}(n^{-\frac{2\alpha}{2\alpha+d}}\log n)$，并在黎曼流形上克服维数灾难。

研究背景与动机¶

过去十年间，基于自注意力机制的 Transformer 架构在大语言模型与视觉生成等领域取得了主导地位，但对其表征能力与统计泛化特性的理论刻画仍落后于工程实践。在深度学习逼近论的传统范式中，针对 $\alpha$-Hölder 连续函数的局部到全局（Local-to-Global）逼近主要依赖显式空间网格构建单位分解（Partition of Unity, POU）。对于经典的深层 ReLU 网络，构造紧支集 POU 面临根本性的结构瓶颈：要么采用“Spike POU”技术通过复合多元仿射极小值构造单纯形剖分，导致网络深度必须达到 $d^2+d$；要么通过一维梯形函数的张量积生成高维网格，由于 ReLU 无法原生计算连续乘法，必须依靠深度为 $\mathcal{O}(\log(1/\varepsilon))$ 的深层网络级联来逼近乘积。这种对超平面显式几何剖分的刚性依赖，不仅无法直接推广到低维嵌入的紧致弯曲流形等非欧几何空间，也将以往 Transformer 理论逼近分析牢牢束缚在超深层或人为简化的硬注意力假定上。

为了回避 ReLU 逼近乘法带来的深度惩罚，已有关于 Transformer 逼近与泛化理论的研究普遍诉诸偏离工程实际的数学妥协。例如，部分代表性工作直接将标准的连续 Softmax 替换为非连续的 Hardmax 激活或无界 ReLU 注意力以强行获得精确乘法或选择性门控；另一些工作则采用非标准的规范化位置编码，或迫使 Transformer 深度扩展至多项式对数阶 $\mathcal{O}(\text{poly}(\log(1/\varepsilon)))$。这些人为假设割裂了理论证明与现实主流 Transformer 架构（即浅层堆叠、标准 Softmax 点积注意力、正弦位置编码）之间的联系，导致自注意力和 Softmax 机制在特征表征学习中的核心数学本质依然未被揭示。

本文的切入点在于反思局部加权逼近的几何与代数结构：标准的 Softmax 激活函数不仅天然满足权重归一化的单位分解条件，更蕴含着一种优美的代数抵消机制。核心 idea：利用 Softmax 激活函数分子与分母中输入范数平方项的精确代数抵消，将无网格欧几里得距离动态加权降维为原生点积注意力的仿射变换，仅用 2 层标准编码器块即可实现局域专家逼近到全局连续函数的无缝聚合，并在低维黎曼流形上实现内在维数自适应的近极大极小最优泛化。

方法详解¶

整体框架¶

本文构建的 Transformer 回归模型接收输入向量 $\bm{x} \in \mathcal{X}$（欧几里得紧支集 $[0,1]^d$ 或嵌入在 $\mathbb{R}^{\bar{d}}$ 中的 $d$ 维紧致黎曼流形 $\mathcal{M}$），经由单层仿射预处理嵌入为初级隐藏序列 $\bm{Z}_0 \in \mathbb{R}^{D \times P}$，随后连续通过 $L=2$ 个标准 Transformer 编码器块，最后由线性读出向量提取回归标量输出。每个编码器块均严格遵循“多头自注意力（MHA）+ 单隐藏层点级前馈网络（FFN）”的标准现代结构。整个计算流向高度并行，模型容量主要集中在序列长度维度 $P$（对应空间覆盖中心数），而深度固定为常数级 $L=2$。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    IN["输入向量 x ∈ ℝᵈ"] --> PRE["预处理层：输入单点嵌入与全序列正弦编码"]
    PRE --> MHA1["第一自注意力层：局部特征并行提取与正弦编码对称保持"]
    MHA1 --> FFN1["第一前馈网络：非线性硬阈值去噪与规范恢复"]
    FFN1 --> MHA2["第二自注意力层：全局Softmax聚合与流形适应"]
    MHA2 --> FFN2["第二前馈网络：恒等传递映射与线性读出"]
    FFN2 --> OUT["标量回归输出 g^T(x)"]

关键设计¶

1. Softmax单位分解与二次项代数抵消：无网格欧氏距离向原生点积注意力的精确转化 传统逼近理论在欧氏网格上的剖分要求刚性超平面，难以迁移至流形结构。本文设计了一种全局支撑且指数局域化的平滑无网格门控机制——Softmax 单位分解。设以有限个点集 $\{\bm{c}_i\}_{i=1}^{C_g}$ 为球心、半径为 $r_g = (\varepsilon/(4C_H))^{1/\alpha}$ 的闭球族构成区域的有效覆盖。给定尺度缩放参数 $M_g > 0$，定义局部空间权重为： $$\beta_i(\bm{x}) = \frac{\exp\left(M_g(r_g^2 - \|\bm{x}-\bm{c}_i\|_2^2)\right)}{\sum_{k=1}^{C_g} \exp\left(M_g(r_g^2 - \|\bm{x}-\bm{c}_k\|_2^2)\right)} = \frac{\exp\left(2M_g \langle\bm{x}, \bm{c}_i\rangle - M_g\|\bm{c}_i\|_2^2\right)}{\sum_{k=1}^{C_g} \exp\left(2M_g \langle\bm{x}, \bm{c}_k\rangle - M_g\|\bm{c}_k\|_2^2\right)}$$ 该设计的关键数学机理在于：分子与分母中均包含非线性项 $\exp(M_g(r_g^2 - \|\bm{x}\|_2^2))$，两者在 Softmax 的指数商中发生完全代数抵消。原本需要高阶多项式拟合的非线性欧氏距离计算，被精确简化为关于输入 $\bm{x}$ 的纯仿射变换 $2M_g \langle\bm{x}, \bm{c}_i\rangle - M_g\|\bm{c}_i\|_2^2$。局部近似误差被严格分解为近邻球内的 Hölder 连续性偏差（被 $r_g^\alpha$ 控制在 $\varepsilon/2$ 以内）以及远场球外的指数级衰减尾部（通过选取足够大的 $M_g \propto \varepsilon^{-2/\alpha}\log(1/\varepsilon)$ 压制在 $\varepsilon/2$ 以内），整体逼近误差严格控制在 $\varepsilon$ 以内。

2. 第一自注意力层局部特征并行提取与正弦编码对称保持：全通道多任务解耦 第一层自注意力模块需要在浅层单步计算中，将输入 token 的信息分发至 $P$ 个序列位置以并行生成局部专家的仿射特征，同时确保附加的正弦位置编码在多头变换中不受混叠污染。本文设置 $H^1 = P+2$ 个独立的注意力头，键/查询/值维度均为 $d_k=d_v=2$： - 前 $P$ 个仿射头（$h \in [P]$）：查询矩阵引入二维平面旋转矩阵 $\bm{R}(\theta_1 - \theta_h)$，利用键向量上的缩放因子 $M$ 产生相对相位差。当且仅当查询位置 $j=h$ 时，未归一化得分在源 token（$k=1$）处取得最大值 $M$，从而使得第 $h$ 个头将注意力集中在首个 token 上，并在第 $h$ 个输出列上并行提取空间仿射特征 $\widetilde{T}_h(\bm{x}) \approx 2M_g \bm{c}_h^\top\bm{x} - M_g\|\bm{c}_h\|_2^2$ 与目标函数局部采样值 $\widetilde{g}(\bm{c}_h) \approx g(\bm{c}_h)$。 - 结构指示符头（$h=P+1$）：通过保留单位投影提取首 token 指示向量，在目标位置输出确定性峰值 $\eta(M) \approx 1$。 - 位置编码保持头（$h=P+2$）：基于正弦编码在均匀角网格 $\theta_k = 2\pi k/P$ 上的全周期对称性，利用三角函数共轭对相消特性，严格证明了注意力加权后的位置编码向量严格保持输入角度 $\theta_j$ 的几何方向不变，仅产生统一的全局标量缩放因子 $\lambda(M) = \frac{\sum_k e^{M\cos\theta_k}\cos\theta_k}{\sum_l e^{M\cos\theta_l}}$，且满足 $0 \le 1 - \lambda(M) \le 2(P-1)e^{-cM}$。

3. 第一前馈网络非线性硬阈值去噪与规范恢复：微小残差截断与离散状态净化 自注意力的 Softmax 操作在非目标列上不可避免地产生以 $\mathcal{O}(e^{-cM})$ 衰减的微小残差泄漏。若不加抑制，这类串扰误差会在跨层传播中指数级恶化。本文构造隐藏层宽度仅为 $d_{\text{ff}} = 2D = 2d+8$ 的点级前馈层 $\mathcal{F}_1$，利用激活恒等式 $x = \sigma(x) - \sigma(-x)$ 实现高精度的非线性硬阈值滤波： $$\bm{z}_j = \bm{\Lambda}\sigma\left(\widehat{\bm{z}}_j - 2e^{-cM}\bm{e}_{d+2}\right) - \bm{\Lambda}\sigma\left(-\widehat{\bm{z}}_j\right)$$ 其中对角缩放矩阵 $\bm{\Lambda}$ 的第 $d+2$ 项配置为逆补偿权重 $w_{\text{ind}} = \frac{1}{\eta(M) - 2e^{-cM}}$。对于指示符神经元，由于非目标列的泄漏幅度满足 $\widetilde{I}_j \le e^{-cM} < 2e^{-cM}$，偏置位移直接将其推入 ReLU 的零区死区，被彻底截断为精确的 $0$；而在源 token 位置（$j=1$），$\eta(M) - 2e^{-cM} > 0$ 且经逆权重放大后精确恢复为标量 $1$，从而无误差地恢复出离散克罗内克指示符 $\delta_{j,1}$。同时，对位置编码维度应用标量补偿系数 $1/\lambda(M)$，使得正弦位置编码恢复为无偏的标准单位模长，保证后续层输入特征的规范性。

4. 第二自注意力层全局Softmax聚合与流形适应：单头凸组合与本征维数自适应 第二编码器块专注于执行局域专家向全局预测的非线性空间融合。该层仅需设置单头注意力（$H^2=1, d_k=d_v=1$）。查询向量由投影矩阵提取上一层恢复的指示符 $\delta_{j,1}$，键向量提取各个位置存储的局域空间特征 $\widetilde{T}_k(\bm{x})$，值向量则提取局域采样值 $\widetilde{g}(\bm{c}_k)$。由此，在首个序列位置（$j=1$），自注意力的未归一化得分精确等于 $\widetilde{T}_k(\bm{x})$，经列向 Softmax 处理后直接生成了归一化的 Softmax POU 权重 $\widetilde{\beta}_k(\bm{x})$。输出层随之执行显式凸组合： $$\widehat{g}_T(\bm{x}) = \sum_{k=1}^P \widetilde{\beta}_k(\bm{x}) \widetilde{g}(\bm{c}_k)$$ 当输入数据分布于具有正内蕴半径（reach）$\tau > 0$ 的 $d$ 维紧致连通黎曼流形 $\mathcal{M} \subset [0,1]^{\bar{d}}$ 上时，利用测地距离与环境欧几里得距离在近邻区域内的局部等价性 $d_{\mathcal{M}}(\bm{x}, \bm{y}) \le 2\|\bm{x}-\bm{y}\|_2$，Softmax POU 仍能直接在环境欧氏度量下运作。覆盖数与模型参数规模仅以内蕴维数 $d$ 呈指数增长，即 $P \le C_P \varepsilon^{-d/\alpha}$，彻底绕开了高维环境空间 $\mathbb{R}^{\bar{d}}$ 的维数灾难。

一个完整示例¶

考虑一个二维紧致域上逼近目标 $\alpha$-Hölder 函数 $g(\bm{x})$ 的具体计算过程： 1. 预处理阶段：输入向量 $\bm{x} \in [0,1]^2$（$d=2, D=6$），序列长度设定为 $P$。预处理层将 $\bm{x}$ 写入首个 token $\bm{Z}_{0,:,1} = [\bm{x}^\top, 1, 0, \sin\theta_1, \cos\theta_1]^\top$，其余列 $\bm{Z}_{0,:,j}$ 前 3 项置零，后 2 项填入均匀角正弦位置编码 $[\sin\theta_j, \cos\theta_j]^\top$。 2. 第一层 MHA 提取：配置 $P+2$ 个注意力头。第 $j$ 个仿射头通过相位差匹配精准捕获首 token 的 $\bm{x}$，并在列 $j$ 计算出 $\widetilde{T}_j(\bm{x}) \approx 2M_g \bm{c}_j^\top\bm{x} - M_g\|\bm{c}_j\|_2^2$ 和 $\widetilde{g}(\bm{c}_j) \approx g(\bm{c}_j)$；头 $P+1$ 产生首列为 $\eta(M) \approx 0.9999$、其余列 $\le 10^{-4}$ 的指示符；头 $P+2$ 保持位置角。 3. 第一层 FFN 去噪：点级 ReLU 前馈层利用偏置 $-2\times 10^{-4}$，将所有非首列的指示符残留截断为绝对的 $0$，首列恢复为标准 $1$；同时按 $1/\lambda(M)$ 归一化位置编码。 4. 第二层 MHA 聚合：单头注意力机制以首列的指示符作为 Query，以各列的 $\widetilde{T}_k(\bm{x})$ 作为 Key，以 $\widetilde{g}(\bm{c}_k)$ 作为 Value。首列的自注意力权重自发形成 Softmax 分布 $\widetilde{\beta}_k(\bm{x}) = \frac{\exp(\widetilde{T}_k(\bm{x}))}{\sum_l \exp(\widetilde{T}_l(\bm{x}))}$，并在 Value 聚合后将最终逼近标量 $\widehat{g}_T(\bm{x})$ 精准输出到首列第一维度。 5. 读出映射：读出向量 $\bm{c}_3 = [1, 0, \dots, 0]^\top$ 直接提取该标量值完成推理。

实验关键数据¶

主实验：理论架构与逼近特性全景对比¶

本文在理论层面上将提出的 Softmax POU 构造与文献中代表性的神经网络及 Transformer 逼近理论进行了系统对照，如下表所示：

理论工作 / 模型架构	注意力机制	位置编码形式	编码器块深度 $L$	总参数量复杂度 $\mathcal{N}_{total}$	适用输入几何领域
Gurevych et al. [17]	Hardmax	无 / 人工指示	$\mathcal{O}(1)$	$\mathcal{O}(\varepsilon^{-d/\alpha})$	欧几里得网格空间
Takakura & Suzuki [41]	Softmax	正弦位置编码	$\mathcal{O}(\text{poly}(\log(1/\varepsilon)))$	$\mathcal{O}(\varepsilon^{-d/\alpha}\text{poly}(\log(1/\varepsilon)))$	欧几里得各向异性空间
Havrilla & Liao [18]	ReLU 注意力	正弦位置编码	$\mathcal{O}(\log(1/\varepsilon))$	$\mathcal{O}(\varepsilon^{-d/\alpha}\log(1/\varepsilon))$	紧致黎曼低维流形
Shi et al. [36]	Softmax	规范标准编码	$\mathcal{O}(\log(1/\varepsilon))$	$\mathcal{O}(\varepsilon^{-d/\alpha}\log(1/\varepsilon))$	层次复合欧氏空间
Jiao et al. [20]	Softmax	坐标投影编码	$\mathcal{O}(\log(1/\varepsilon))$	$\mathcal{O}(\varepsilon^{-d/\alpha}\log(1/\varepsilon))$	连续 Hölder 函数空间
本文方法 (Ours)	原生 Softmax	标准正弦编码	2 (常量级)	$\bm{\mathcal{O}(\varepsilon^{-d/\alpha})}$	欧氏空间与紧致黎曼流形

消融实验：欧氏空间与黎曼流形理论界分解与统计收敛率定量对比¶

下表展示了本文在欧几里得空间与紧致黎曼流形两套理论设定下的具体网络参数规模常数、假说集覆盖数上界以及经验风险最小化（ERM）估计量的统计泛化收敛率：

几何空间设定	空间覆盖半径 $r_g$	隐藏序列长度 $P$	隐藏维度 $D$	参数模长界 $M_{max}$	ERM 超额风险界 $\mathbb{E}[\\|\pi_B f_{\mathcal{S}} - f_\rho\\|_\rho^2]$
欧几里得紧支集 $[0,1]^d$	$(\frac{\varepsilon}{4C_H})^{1/\alpha}$	$\le d^{d/2}(8C_H)^{d/\alpha}\varepsilon^{-d/\alpha}$	$d+4$	$\le \widetilde{C}_{mag}\varepsilon^{-2d/\alpha}\log\frac{1}{\varepsilon}$	$\le C_4 n^{-\frac{2\alpha}{2\alpha+d}}\log n$
紧致黎曼流形 $\mathcal{M}\subset[0,1]^{\bar{d}}$	$(\frac{\varepsilon}{8C_H})^{1/\alpha}$	$\le C_{\mathcal{M}}(16C_H)^{d/\alpha}\varepsilon^{-d/\alpha}$	$\bar{d}+4$	$\le \widetilde{C}_{mag}\varepsilon^{-2d/\alpha}\log\frac{1}{\varepsilon}$	$\le \widetilde{C}_4 n^{-\frac{2\alpha}{2\alpha+d}}\log n$
消融对照：环境欧氏直接逼近	$(\frac{\varepsilon}{4C_H})^{1/\alpha}$	$\le \bar{d}^{\bar{d}/2}(8C_H)^{\bar{d}/\alpha}\varepsilon^{-\bar{d}/\alpha}$	$\bar{d}+4$	$\le \widetilde{C}_{mag}\varepsilon^{-2\bar{d}/\alpha}\log\frac{1}{\varepsilon}$	$\le C_4' n^{-\frac{2\alpha}{2\alpha+\bar{d}}}\log n$

关键发现¶

深度极度压缩：相较于以往所有依赖深层堆叠复合（$\mathcal{O}(\log(1/\varepsilon))$）来计算乘法与连续指示滤波的工作，本文证明了仅需 $L=2$ 的网络深度即可达到最优逼近阶，展示了注意力竞争在空间解耦中的高效性。
环境无关的统计收敛率：在数据支撑于 $\bar{d}$ 维环境空间中具有正 reach 的 $d$ 维低维流形时，泛化收敛率中的样本指数纯粹取决于内在维数 $d$，即指数项为 $-\frac{2\alpha}{2\alpha+d}$，完全摆脱了环境维数 $\bar{d}$ 的诅咒。
宽度与容量的对偶性：逼近精度 $\varepsilon$ 的提升被完全吸收到了序列长度维度 $P \propto \varepsilon^{-d/\alpha}$ 与第一层注意力头数 $H^1 = P+2$ 中，这表明 Transformer 能够通过空间序列的横向宽度扩展换取纵向深度的极大简化。

亮点与洞察¶

二次项代数抵消机制：利用 Softmax 算子对输入模长平方项 $\|\bm{x}\|_2^2$ 的自然消去，将非线性局域核计算无损转换为标准点积注意力的仿射矩阵变换，揭示了 Softmax 自注意力机制在空间局部化任务中的内在数学优势。
正弦位置编码的旋转对称保持：在全周期等距角网格上利用对称性与正交互补性证明了注意力输出对位置角的不变性，为实用的正弦位置编码提供了严格的保真度理论保证。
无网格流形几何自适应：通过欧氏度量与测地度量的局部等价性，使得环境点积注意力能够直接感知流形内蕴拓扑，省去了传统微分几何流形逼近对图册（chart）和局部切空间坐标系的复杂构造。

局限与展望¶

光滑性阶数限制：当前构造性证明严格局限于 $\alpha \in (0,1]$ 的 Hölder 连续函数，未进一步拓展至高阶 Sobolev 空间或解析光滑函数空间；对于更高阶光滑性，可能需要高阶局部多项式专家代替常数逼近。
宽网络与参数模长上界：模型为了在极浅层（$L=2$）内完成逼近，将复杂度转移到了注意力头数与序列长度上（$P \propto \varepsilon^{-d/\alpha}$），且参数模长存在 $\varepsilon^{-2d/\alpha}\log(1/\varepsilon)$ 的多项式膨胀。
上下文学习（ICL）延伸：作者指出本套 Softmax POU 框架有望为解释 Transformer 的上下文学习机制（In-Context Learning）提供全新的自适应逼近视角，值得后续深入探索。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙利用 Softmax 代数抵消构造无网格单位分解，完全摒弃了传统 ReLU 网络的几何剖分桎梏。
实验充分度: ⭐⭐⭐⭐☆ 纯理论研究，提供了详尽严密的构造性引理、参数量化界、覆盖数估计与 ERM 泛化界，推导无懈可击。
写作质量: ⭐⭐⭐⭐⭐ 结构清晰，数学符号定义严谨规范，理论洞察与几何直觉阐述透彻。
价值: ⭐⭐⭐⭐⭐ 奠定了浅宽实用 Transformer 架构的逼近与泛化理论基础，为自注意力的可解释性与流形适应性提供了权威基准。

几何空间设定	空间覆盖半径 \(r_g\)	隐藏序列长度 \(P\)	隐藏维度 \(D\)	参数模长界 \(M_{max}\)	ERM 超额风险界 \(\mathbb{E}[\\|\pi_B f_{\mathcal{S}} - f_\rho\\|_\rho^2]\)
欧几里得紧支集 \([0,1]^d\)	\((\frac{\varepsilon}{4C_H})^{1/\alpha}\)	\(\le d^{d/2}(8C_H)^{d/\alpha}\varepsilon^{-d/\alpha}\)	\(d+4\)	\(\le \widetilde{C}_{mag}\varepsilon^{-2d/\alpha}\log\frac{1}{\varepsilon}\)	\(\le C_4 n^{-\frac{2\alpha}{2\alpha+d}}\log n\)
紧致黎曼流形 \(\mathcal{M}\subset[0,1]^{\bar{d}}\)	\((\frac{\varepsilon}{8C_H})^{1/\alpha}\)	\(\le C_{\mathcal{M}}(16C_H)^{d/\alpha}\varepsilon^{-d/\alpha}\)	\(\bar{d}+4\)	\(\le \widetilde{C}_{mag}\varepsilon^{-2d/\alpha}\log\frac{1}{\varepsilon}\)	\(\le \widetilde{C}_4 n^{-\frac{2\alpha}{2\alpha+d}}\log n\)
消融对照：环境欧氏直接逼近	\((\frac{\varepsilon}{4C_H})^{1/\alpha}\)	\(\le \bar{d}^{\bar{d}/2}(8C_H)^{\bar{d}/\alpha}\varepsilon^{-\bar{d}/\alpha}\)	\(\bar{d}+4\)	\(\le \widetilde{C}_{mag}\varepsilon^{-2\bar{d}/\alpha}\log\frac{1}{\varepsilon}\)	\(\le C_4' n^{-\frac{2\alpha}{2\alpha+\bar{d}}}\log n\)