跳转至

Ordinary Least Squares as an Attention Mechanism

会议: NeurIPS 2026
arXiv: 2504.09663
领域: 时间序列 / Time Series
关键词: 普通最小二乘法 (OLS), 注意力机制 (Attention Mechanism), 样本相似度检索, 回归块 (Regression Block), 表格数据 Transformer

一句话总结

论文证明了普通最小二乘法 (OLS) 样本外预测与受限线性注意力机制严格等价,将相关系数估计重塑为正交特征空间中的样本内积检索,并由此提出基于多头非线性注意力回归与基于多项式特征回归的轻量级 Transformer 模块(Regression Block)。

研究背景与动机

自注意力机制(Self-Attention)是现代大语言模型和序列建模的核心支柱,但对于受过传统应用统计学与计量经济学训练的研究者而言,源自信息检索领域的查询-键-值(Query-Key-Value)范式往往带有黑盒色彩与概念隔阂。长久以来,统计学界将线性回归视为求解参数系数矢量的工具,而深度学习界则将注意力视为数据驱动的加权检索系统。已有文献虽然注意到了注意力机制与核平滑(Kernel Smoothers)或上下文学习(In-Context Learning)中收敛至最小二乘更新的经验关联,但大多是从复杂网络训练后“涌现”出统计行为的单向视角切入,未曾揭示传统回归本身所深藏的检索本质。

这种理解割裂的核心矛盾在于:我们通常认为最小二乘法完全是“全局相关性驱动”的,而注意力机制是“样本间局部接近度(Proximity)驱动”的。事实上,当考查样本外预测公式时,预测值可被完全重组为训练集标签的加权平均,其权重恰好对应于经样本协方差矩阵白化正交后的测试样本与训练样本的内积。这意味着即使是最朴素的未正则化 OLS,在其代数底层也过着双重生活——它天然就是一个预先计算好最优嵌入(即样本逆协方差矩阵 \((\boldsymbol{X}'\boldsymbol{X})^{-1}\))的注意力检索模块。

本文的切入角度是直接建立从经典最小二乘到注意力模块的双向代数桥梁,不仅消除注意力机制在经典统计视角下的神秘感,更把非参数回归与协方差估计的先验知识反哺回神经网络设计中。核心 idea:将 OLS 样本外预测形式化表述为恒等激活下的线性交叉注意力机制,指出其最优几何嵌入唯一收敛于样本精度矩阵 \((\boldsymbol{X}'\boldsymbol{X})^{-1}\);并在反方向上用显式多项式特征回归替代 Transformer 中的多头注意力子层,以约 30% 参数量在表格任务上匹敌甚至超越标准架构。

方法详解

整体框架

论文从线性代数的谱分解出发,建立预测公式与检索权重的精确映射,进而将线性注意力向降维(主成分回归 PCR)、收缩正则化(Ridge)、非线性 Softmax(注意力回归 Attention Regression)以及深度网络结构扩展。整个推导与衍生流程由线性代数恒等式贯穿至网络模块构建:首先将特征矩阵白化为正交因子,以正交因子的内积作为注意力得分;随后引入投影矩阵降维和指数归一化非线性;最终在工程实践中反转视角,用显式二次多项式回归块(Regression Block)替换 FT-Transformer 中的注意力子层。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入设计矩阵 X_train, y, X_test"] --> Decomp["正交化嵌入映射<br/>W = U Λ^{-1/2} 投影至正交因子空间"]
    Decomp --> LinearAttn["等价受限注意力表述<br/>Q=X_test W, K=X_train W, V=y"]
    LinearAttn --> Ext["扩展与演变路径"]
    Ext --> ExtPCR["降维与正则化约束<br/>截断特征向量 PCR 与 Ridge 缩减"]
    Ext --> ExtSoftmax["非线性注意力回归<br/>引入 Softmax 约束与多头 Cholesky 参数化"]
    Ext --> RegBlock["反向工程:回归块架构<br/>多项式交叉项 + 冻结 PCA + 岭回归初始化"]
    ExtPCR --> Out["样本外预测输出 y_hat"]
    ExtSoftmax --> Out
    RegBlock --> Out

关键设计

1. 正交化嵌入映射:将最小二乘重塑为内积相似度检索 传统 OLS 估计系数为 \(\hat{\boldsymbol{\beta}} = (\boldsymbol{X}_{\text{train}}'\boldsymbol{X}_{\text{train}})^{-1}\boldsymbol{X}_{\text{train}}'\boldsymbol{y}\),样本外测试集 \(\boldsymbol{X}_{\text{test}} \in \mathbb{R}^{J \times P}\) 的预测值为 \(\hat{\boldsymbol{y}}_{\text{test}} = \boldsymbol{X}_{\text{test}}(\boldsymbol{X}_{\text{train}}'\boldsymbol{X}_{\text{train}})^{-1}\boldsymbol{X}_{\text{train}}'\boldsymbol{y}\)。论文指出,利用样本自相关矩阵的谱分解 \((\boldsymbol{X}_{\text{train}}'\boldsymbol{X}_{\text{train}})^{-1} = \boldsymbol{U}\Lambda^{-1}\boldsymbol{U}'\),可将逆特征值矩阵拆分为 \(\Lambda^{-\frac{1}{2}}\Lambda^{-\frac{1}{2}}\),从而将预测值严格重写为因子得分矩阵的乘积: $\(\hat{\boldsymbol{y}}_{\text{test}} = \underbrace{\boldsymbol{X}_{\text{test}}\boldsymbol{U}\Lambda^{-\frac{1}{2}}}_{\boldsymbol{F}_{\text{test}}} \underbrace{\Lambda^{-\frac{1}{2}}\boldsymbol{U}'\boldsymbol{X}_{\text{train}}'}_{\boldsymbol{F}_{\text{train}}'} \boldsymbol{y}_{\text{train}} = [\boldsymbol{F}_{\text{test}}\boldsymbol{F}_{\text{train}}']\boldsymbol{y}_{\text{train}}\)$ 其中 \(\boldsymbol{F}_{\text{train}} \in \mathbb{R}^{N \times P}\) 满足 \(\boldsymbol{F}_{\text{train}}'\boldsymbol{F}_{\text{train}} = \boldsymbol{I}_P\),是彼此完全不相关的标准正交因子。单条测试样本的预测值因此完全分解为训练集响应值的相似度加权平均:\(\hat{y}_j = \sum_{i=1}^N \langle F_j, F_i \rangle y_i\)。该设计揭示了 OLS 与 \(k\) 近邻等基于接近度估计器的本质差异:OLS 允许相似度为负值(符号相反的极端点提供负权重),且在正交化基底下去除了特征共线性导致的冗余加权。

2. 等价受限注意力表述:QKV 映射与最优嵌入的唯一性 根据标准自注意力公式 \(\text{Attention}(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = g(\frac{\boldsymbol{Q}\boldsymbol{K}'}{\sqrt{P}})\boldsymbol{V}\),若将激活函数 \(g(\cdot)\) 设为恒等映射 \(I(\cdot)\),令查询投影 \(\boldsymbol{W}_Q = \boldsymbol{U}\Lambda^{-\frac{1}{2}}\)、键投影 \(\boldsymbol{W}_K = \boldsymbol{U}\Lambda^{-\frac{1}{2}}\)、值投影 \(\boldsymbol{w}_v = \mathbf{1}\),则有: $\(\boldsymbol{Q} = \boldsymbol{X}_{\text{test}}\boldsymbol{W}_Q, \quad \boldsymbol{K} = \boldsymbol{X}_{\text{train}}\boldsymbol{W}_K, \quad \boldsymbol{V} = \boldsymbol{y}\boldsymbol{w}_v \implies \boldsymbol{Q}\boldsymbol{K}' = \boldsymbol{X}_{\text{test}}(\boldsymbol{X}_{\text{train}}'\boldsymbol{X}_{\text{train}})^{-1}\boldsymbol{X}_{\text{train}}'\)$ 此时线性注意力完全等同于 OLS 预测。进一步,论文证明若将求解系数 \(\boldsymbol{\beta}\) 转化为求解几何嵌入矩阵 \(\boldsymbol{\Omega} \equiv \boldsymbol{W}_Q\boldsymbol{W}_K'\): $\(\hat{\boldsymbol{\Omega}} \in \arg\min_{\boldsymbol{\Omega} \in \mathbb{R}^{P \times P}} \|\boldsymbol{y} - \boldsymbol{X}_{\text{train}}\boldsymbol{\Omega}\boldsymbol{X}_{\text{train}}'\boldsymbol{y}\|^2\)$ 对于任意可能的目标响应向量 \(\boldsymbol{y}\),只要设计矩阵满秩,该最小二乘嵌入矩阵的解是唯一的,且严格等于逆样本协方差矩阵 \((\boldsymbol{X}_{\text{train}}'\boldsymbol{X}_{\text{train}})^{-1}\)。这表明如果用梯度下降训练一个线性注意力层来拟合残差平方和,其权重收敛极限就是 OLS。

3. 降维与正则化约束:主成分回归与岭回归在注意力中的统一 当特征维度 \(P\) 较大或存在多重共线性时,直接求逆会导致方差膨胀。将映射空间维度压缩至 \(L < P\),即 \(\boldsymbol{W}_Q = \boldsymbol{W}_K = \boldsymbol{U}_L\Lambda_L^{-\frac{1}{2}}\),此时模型直接退化为主成分回归(PCR),其预测值等价于在低维正交子空间上施加线性注意力。同理,施加 \(L_2\) 正则化的岭回归仅需将对角矩阵调整为 \((\Lambda + \lambda\boldsymbol{I})^{-1}\)。由此,OLS、Ridge 与 PCR 统一为同一个注意力模块内部不同的奇异值缩放规则: $\(\boldsymbol{D}_{\text{OLS}} = \Lambda^{-1}, \quad \boldsymbol{D}_{\text{ridge}} = (\Lambda + \lambda\boldsymbol{I})^{-1}, \quad \boldsymbol{D}_{\text{PCR}} = \text{diag}(\lambda_1^{-1}, \dots, \lambda_L^{-1}, 0, \dots, 0)\)$ 这种统一性为神经网络注意力层提供了全新的协方差正则化视角,例如通过 Ledoit-Wolf 收缩或稀疏阈值化来增强注意力权重矩阵的条件数。

4. 非线性注意力回归:Softmax 单纯形约束下的多头参数化 标准 Transformer 中引入了逐行 Softmax 归一化,迫使权重落在非负概率单纯形上(\(\sum_i \omega_{ji} = 1, \omega_{ji} \ge 0\))。这彻底剥夺了 OLS 自由利用负权重的机制,将估计器转变为类似随机森林的凸组合平滑器。为了在连续空间中优化该非线性目标,论文提出了多头注意力回归(Attention Regression): $\(\hat{\boldsymbol{y}} = \sum_{m=1}^M \alpha_m \left[ \text{softmax}\left(\boldsymbol{X}^*(\boldsymbol{L}_m\boldsymbol{L}_m')\boldsymbol{X}_{\text{train}}'\right)\boldsymbol{y} \right]\)$ 其中每个头的相似度度量通过下三角 Cholesky 分解 \(\boldsymbol{\Omega}_m = \boldsymbol{L}_m\boldsymbol{L}_m'\) 参数化以保证半正定性(即内积空间合法性),且每一头均以岭回归精度矩阵的 Cholesky 分子结合高斯噪声进行热启动,利用 L-BFGS 进行端到端惩罚平方误差最小化。

5. 反向工程:回归块架构(Regression Block)对注意力子层的替代 受“注意力本质即特征回归”这一发现的启发,论文反向设计了专为表格数据设计的“回归块”(Regression Block)。在 FT-Transformer 的残差与 LayerNorm 框架中,论文彻底剔除了多头自注意力机制和 Softmax,替换为显式的二次多项式交叉项扩展: $\(\text{LN}\left(\boldsymbol{h} + \boldsymbol{W}\,\text{PCA}(\text{polycross}(\boldsymbol{h}))\right)\)$ 输入 Token 经包含平方项与两两乘积项的二次扩展后,使用在训练集上预先拟合并冻结的 PCA 压缩至 200 个主成分,再通过线性投影层 \(\boldsymbol{W}\) 与残差流相加。该投影层的第一通道在训练前使用岭回归(\(\lambda=1\))进行热启动。该模型仅需单层即可工作,无需多层堆叠,参数量仅为基线 FT-Transformer 的约 30%。

损失函数 / 训练策略

注意力回归模型(Attention Regression)的目标损失函数采用加入 Frobenius 范数惩罚项的均方误差: $\(\mathcal{L} = \frac{1}{N}\left\|\boldsymbol{y} - \sum_{m=1}^M \alpha_m \text{softmax}\left(\boldsymbol{X}_{\text{train}}(\boldsymbol{L}_m\boldsymbol{L}_m')\boldsymbol{X}_{\text{train}}'\right)\boldsymbol{y}\right\|^2 + \lambda\sum_{m=1}^M \|\boldsymbol{L}_m\|_F^2\)$ 优化采用带有早停机制的 L-BFGS 算法;当验证集拟合发生异常(如测试集 \(R^2 < 0.1\))时触发自动重拟合机制。而在反向替代的 Regression Block 实验中,采用 Adam 优化器,恒定学习率 \(10^{-3}\),权重衰减 \(10^{-4}\),批大小 256,Dropout 0.1,梯度裁剪阈值为 1,并在验证集上施加耐心度为 10 轮的早停策略。

实验关键数据

主实验

在 8 个标准 UCI/OpenML 表格回归基准数据集上进行 5 折 80/20 划分测试(样本量上限截断为 5000)。对比基线涵盖经典线性模型(OLS)、强非线性基线(Random Forest、MLP)、前沿表格 Transformer(FT-Transformer)以及论文提出的 Attention Regression(5头)与单层 Regression Block。

数据集 样本量 \(N\) 特征数 \(P\) OLS Random Forest MLP FT-Transformer Att. Reg (本文) Reg. Block (本文)
California 5000 8 0.597 0.777 0.763 0.764 0.738 0.766
Yacht 308 6 0.562 0.979 0.970 0.988 0.958 0.989
Energy 768 8 0.913 0.996 0.994 0.994 0.995 0.995
Concrete 1030 8 0.624 0.892 0.895 0.898 0.837 0.907
Airfoil 1503 5 0.497 0.910 0.917 0.924 0.757 0.927
Abalone 4177 7 0.260 0.322 0.323 0.324 0.312 0.329
Kin8nm 5000 8 0.430 0.664 0.912 0.914 0.851 0.920
Protein 5000 9 0.285 0.521 0.486 0.426 0.361 0.455

消融实验

论文在附录 E 中对 Regression Block(RB)开展了系统的 10 种配置消融研究,探讨多项式扩展、PCA 降维、Softmax 注意力机制、模型深度及参数规模的作用机制:

配置代号 架构变体说明 均值 \(R^2\) 参数量 (\(P=8\)) 说明与关键发现
A0 标准 FT-Transformer 基线(3层,64维) 0.779 101,697 传统自注意力表格基线
A1 A0 + 输入端加入二次多项式特征 0.773 106,305 标准 Transformer 难以有效吸收显式多项式,反而微弱掉点
A2 Regression Block 移除多项式扩展 0.572 21,697 性能暴跌 0.215,证明显式非线性交叉是替代注意力的关键
A3 Regression Block 移除 PCA 降维压缩 0.773 173,249 参数激增 5.7 倍但准确率微降 0.013,证明 PCA 主要贡献参数压缩
A4 单层骨干网换回 Softmax 注意力 0.669 34,177 单层 Softmax 严重退化(-0.117),显式回归头远优于浅层注意力
A5 FT-Transformer 窄化至 RB 同等参数量 0.772 33,013 压缩容量后三层 Softmax 依然稳定,说明并非单纯容量差异
A6 Regression Block 堆叠至 3 层深层结构 0.772 89,921 堆叠多层反而下降 0.015,符合线性算子幂等与局部化特征
A7 Regression Block 单层展宽至 FT-T 参数量 0.787 101,249 展宽后增益不显著 (+0.001),单层紧凑设计已近饱和
RB 本文完整 Regression Block 0.786 30,401 以 30% 参数量在所有 8 个数据集上超越 FT-Transformer
T 极限精简版(无 FFN,50 个 PCA 分量) 0.774 7,489 参数量降至 FT-T 的 7.4%,均值 \(R^2\) 仅相差 0.005

关键发现

  • 显式特征交叉是核心源泉:消融表表明,从 Regression Block 中剔除多项式特征(A2)会导致均值 \(R^2\) 骤降 0.215(\(p < 0.001\))。且数据集中 OLS 与随机森林的差距越大(代表真实非线性程度越强),剔除多项式带来的损失越大(等级相关系数达 0.74)。
  • 回归头与浅层优势显著:将 Softmax 注意力放回 Regression Block 的单层骨架中(A4),均值 \(R^2\) 仅为 0.669。这证明在单层网络中,显式回归头配合岭回归热启动具备远优于 Softmax 权重的拟合效率与凸几何鲁棒性。
  • 深度红利的不对称性:Softmax 注意力高度依赖深层堆叠(A0 达到 0.779,A4 仅 0.669);而 Regression Block 单层表现最优(RB 为 0.786),堆叠到 3 层(A6)反而下降至 0.772,呼应了线性投影算子具备幂等性(Idempotence)的理论分析。
  • 配对统计显著性:在所有 8 个基准数据集上,Regression Block 的表现均优于 FT-Transformer(符号检验 \(p = 0.008\)),在 40 组数据集/划分的配对双侧 \(t\) 检验中取得平均 +0.0070 的优势(\(p = 0.009\))。

亮点与洞察

  • 将 OLS 重新解读为样本内积检索系统:打破了“最小二乘法只是解系数方程”的传统教条,证明了 OLS 的样本外预测本质上是数据在正交精度空间里的内积相似度加权平均,让统计学与自注意力机制底层的 QKV 计算在代数层面合流。
  • 阐明负注意力的本质作用:揭示了 Softmax 强制非负单纯形约束的代价——它排除了统计回归天然具备的“反向抵消与负向加权”能力。在表格数据上舍弃 Softmax 转而采用显式回归,使得浅层网络能够直接学到跨特征的全局交互。
  • 可复用的高参数效益模块(Regression Block):提供了一个极简高效的工程范式——将输入 Token 展开二次多项式交互并经无监督 PCA 压缩,随后直接用岭回归初始化映射层。该设计在参数量锐减 70% 的前提下实现了更优的表格拟合能力。

局限与展望

  • 特征维度扩展时的计算瓶颈:Regression Block 生成的两两交互多项式特征数随特征数呈二次方增长(\(O(P^2)\))。在特征维度达到数十甚至上百维时(如附录 F 中测试的 40-48 维数据),生成多项式与进行全量 PCA 的内存与计算开销急剧放大。
  • 深度注意力等价性的局部性局限:严格的 OLS 等价性仅在单层受限注意力(恒等激活、无前馈非线性激活)下精确成立。对于包含多层激活、LayerNorm 和复杂 MLP 的深层大语言模型,该等价性仅能在“以当前层输入表征为条件”的局部意义下理解,尚不能直接推演深层 Transformer 的全局闭式解。
  • 时间序列伪样本外泄漏问题:论文在附录 A.4 指出,虽然自注意力机制可通过因果下三角掩码避免直接使用未来时间步,但只要相似度度量依赖于基于全样本计算的协方差逆矩阵 \(\boldsymbol{\Omega}\),就依然存在隐式信息泄漏,无法完全替代计量经济学中耗时的递推重估计。

相关工作与启发

  • vs 核方法与表示定理(Kernel Methods & Representer Theorem):经典理论通常从对偶问题将预测表示为训练样本与测试样本的核函数线性组合,但对偶参数系数的具体几何角色较为晦涩。本文则指出即便不经任何非线性核变换,原始 OLS 的内积空间本身就由样本精度矩阵唯一确定,直接对应于注意力机制中的权重投影。
  • vs 上下文学习(In-Context Learning)回归研究(Garg et al. / von Oswald et al.):已有工作通过大量预训练任务证明 Transformer 在上下文中能够“涌现”出类似梯度下降或岭回归的行为;本文则是反向思考,证明回归算法本身就已经内生了注意力机制的全部结构,从而解释了为何 Transformer 天然容易学会回归。
  • vs 表格 Transformer(FT-Transformer):FT-Transformer 依赖于多层多头自注意力机制去隐式捕捉特征之间的特征交叉;Regression Block 则直接将特征两两交叉显式展开,用冻结 PCA 加单层线性回归替代多头 Softmax,在大幅减少参数量与深度的同时获得了更具鲁棒性的表现。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 颠覆性地将统计学基石 OLS 与现代深度学习核心注意力机制直接统一,理论视角令人耳目一新。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 6 个合成 DGP 的 96 组蒙特卡洛实验、8 个公开表格基准及 10 组细致的消融实验,并给出了严格的配对显著性检验。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严密流畅,从经典线性代数平滑过渡到现代 Transformer 架构,直觉解释透彻明晰。
  • 价值: ⭐⭐⭐⭐⭐ 不仅在理论上弥合了计量统计与深度学习的认知鸿沟,更提出了具有实用价值的轻量级表格回归模块架构。