跳转至

Sparse Attention as Compact Kernel Regression

会议: NeurIPS 2026
arXiv: 2601.22766
OpenReview: https://openreview.net/forum?id=v1sbV2aAVF
代码: https://github.com/deep-spin/sparse_kernel_regression
领域: LLM 效率
关键词: 稀疏注意力、紧支撑核、核回归、α-entmax、ReLUmax

一句话总结

本文建立了稀疏注意力与紧支撑核回归之间的形式化理论映射，证明 normalized ReLU、sparsemax 及 \(\alpha\)-entmax 分别对应固定与自适应带宽下的 Epanechnikov 与整流多项式核，并提出消除分母退化且具备动态裕度稀疏特性的 \(r\)-ReLUmax 注意力机制。

研究背景与动机

标准 Transformer 中的自注意力机制在测试期常被视为非参数统计中的 Nadaraya-Watson 核回归估计器，其中 softmax 变换在数学上等价于键值缓存（key-value cache）上的高斯核（Gaussian kernel）加权回归。然而，高斯核具备无限支撑集（infinite support），导致所有上下文位置均被分配非零概率质量。在长上下文序列建模中，这种全局密集分配会引发严重的注意力分散（attention dispersion）与信息过度挤压（information over-squashing），使得模型在大量不相关词元微弱噪声的干扰下，逐渐丧失对关键局部表征的辨别能力与记忆检索精度。

从认知科学与神经生物学视角来看，人类海马体联想记忆索引理论表明，大脑对过去情景记忆的检索依赖于少量高度分离、极致稀疏的神经元群落活动，而非在漫长生命周期积累的所有记忆痕迹上进行微弱且全量的密集积分。为了在神经网络中复现这种局部性与抗干扰能力，学术界此前探索了启发式 top-\(k\) 截断以及基于 Tsallis 熵的 sparsemax 和 \(\alpha\)-entmax 等可微稀疏激活函数。尽管这些稀疏机制在经验上表现出色，但相比于 softmax 与高斯核之间成熟的统计学联系，稀疏注意力在测试期核回归框架下究竟对应何种核函数、其稀疏性如何从第一性原理中自然导出，长期缺乏严格的理论刻画。

针对这一理论空白，本文的目标是建立一套涵盖密集与稀疏注意力的统一非参数核回归分析框架，阐明稀疏性与统计学中紧支撑核（compact support kernels）之间的本质对偶关系。核心 idea：将稀疏注意力机制形式化统一为紧支撑核函数的 Nadaraya-Watson 核回归，证明 sparsemax 与 \(\alpha\)-entmax 分别对应自适应带宽自归一化的 Epanechnikov 与高阶整流多项式核，并由此导出以最大相似度为锚点、自适应动态活跃集且彻底消除零分母退化的 \(r\)-ReLUmax 注意力。

方法详解

整体框架

在 Nadaraya-Watson 核回归视角下，给定历史观测键值对 \(\{(\bm{k}_i, \bm{v}_i)\}_{i=1}^{n-1}\) 与当前查询 \(\bm{q} \equiv \bm{k}_n\)，注意力机制的目标是通过核函数 \(K_h\) 计算局部加权平均来预测下一个目标值向量 \(\hat{\bm{v}}_n(\bm{q})\)。在键与查询均进行 \(\ell_2\) 归一化（\(\|\bm{q}\| = \|\bm{k}_i\| = 1\)）的前提下，欧氏距离平方与内积满足恒等式 \(\frac{1}{2}\|\bm{k}_i - \bm{q}\|^2 = 1 - \bm{k}_i^\top \bm{q}\)。替换经典高斯核为紧支撑核，能够直接迫使超出核带宽半径的数据点获得严格为零的权重。

整个计算体系主要由两条紧支撑核回归路径构成：第一条是自适应带宽自归一化路径，通过让核带宽 \(h\) 依赖于输入数据，使紧支撑整流响应的求和天然为 1，直接推导出连续可微的 sparsemax 与 \(\alpha\)-entmax；第二条是锚定紧支撑核路径，通过硬截断（kNN 均匀核与截断高斯核）或软锚定（以最高得分键为基准的 \(r\)-ReLUmax）构建自适应邻域。作者将上述机制实装于显式实现 Nadaraya-Watson 估计器的 Memory Mosaics 架构中进行端到端验证。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入查询与键值对<br/>Query q, Keys K, Values V"] --> B["度量空间对齐<br/>单位球投影: 1/2||k_i - q||^2 = 1 - k_i^T q"]
    B --> C["紧支撑核回归路径选择"]
    C -->|自归一化约束| D["自适应带宽自归一化<br/>Epanechnikov 与 α-entmax 核"]
    C -->|最大相似度锚定| E["锚定紧支撑核与 ReLUmax<br/>最大值差分裕度与动态活跃集"]
    D --> F["非参数局部加权聚合<br/>Nadaraya-Watson 估计预测目标值"]
    E --> F
    F --> G["输出表征与下游任务<br/>语言建模 / 上下文推理 / 长度外推"]

关键设计

1. 自适应带宽自归一化：从 Epanechnikov 核到 α-entmax 的统一表示

针对固定带宽核回归在整流截断时容易遭遇分母全零（所有预激活值均为负数导致 \(0/0\) 无法定义）的退化痛点，本文提出了核尺度的自归一化准则。对于统计密度估计中最优的 Epanechnikov 紧支撑核 \(K(\bm{u}) \propto [1 - \|\bm{u}\|^2]_+\)，固定带宽下的 Nadaraya-Watson 归一化退化为 normalized ReLU 形式。而自归一化则打破固定带宽假设，要求选择数据自适应带宽 \(h\)，使得整流响应自身之和恰好等于 1：

\[\sum_{i=1}^{n-1} \left[ \frac{\bm{k}_i^\top \bm{q}}{\gamma} - \left( \frac{1}{\gamma} - \frac{h^2}{2\gamma} \right) \right]_+ = 1 \implies \hat{\bm{v}}_n(\bm{q}) = \sum_{i=1}^{n-1} \text{sparsemax}_i \left( \frac{\bm{K}\bm{q}}{\gamma} \right) \bm{v}_i\]

其中阈值 \(\tau = \frac{1}{\gamma} - \frac{h^2}{2\gamma}\) 恰好对应 sparsemax 的截断常数，有效自适应带宽为 \(h = \sqrt{2 - 2\gamma\tau}\)。更进一步，作者在 Proposition 3.1 中证明了这一结论的高阶泛化：对于任意阶数 \(r \ge 1\) 的整流多项式核 \(K_h(\bm{u}) \propto [1 - \|\bm{u}\|^2 / h^2]_+^r\)，在自归一化条件下的 Nadaraya-Watson 回归严格等价于参数为 \(\alpha = 1 + 1/r\) 的 \(\alpha\)-entmax 注意力，自适应带宽为 \(h = \sqrt{2 - 2r\gamma\tau}\)。这一理论成果将 Epanechnikov 核（\(r=1 \leftrightarrow \alpha=2\) 即 sparsemax）、双权重核（biweight, \(r=2 \leftrightarrow \alpha=1.5\)）及三权重核（triweight, \(r=3 \leftrightarrow \alpha=4/3\)）完全统一在非参数核回归谱系中，而高斯核正是 \(r \to \infty\)（\(\alpha \to 1\)）时的极限特例。

2. 锚定紧支撑核与 ReLUmax：基于最大值裕度消除退化与动态稀疏

为了提供一种不需要迭代求解隐式归一化阈值 \(\tau\)、但又能彻底杜绝 normalized ReLU 零分母崩溃的替代方案，本文提出了 \(r\)-ReLUmax 变换。该机制的核心是将核的支撑集锚定在与查询相似度最高的键向量上。定义当前序列中的最大内积得分为 \(m := \max_j \bm{k}_j^\top \bm{q}\)，并引入正超参数 \(b > 0\)，\(r\)-ReLUmax 核及其回归估计定义为：

\[K_{r\text{-relumax}}(\bm{k}_i - \bm{q}) \propto \left[ b + \frac{\bm{k}_i^\top \bm{q} - m}{h^2} \right]_+^r \implies \hat{\bm{v}}_n(\bm{q}) = \sum_{i=1}^{n-1} \frac{\left[ b + \frac{\bm{k}_i^\top \bm{q} - m}{h^2} \right]_+^r}{\sum_{k=1}^{n-1} \left[ b + \frac{\bm{k}_k^\top \bm{q} - m}{h^2} \right]_+^r} \bm{v}_i\]

由于对于获得最高分的键必定满足 \(\bm{k}_i^\top \bm{q} - m = 0\)，其括号内取值为 \(b > 0\)，因此至少存在一个正响应项，保证分母严格大于零。更为关键的是，该核激活的充要条件为 \(\bm{k}_i^\top \bm{q} > m - bh^2\)。这意味着 \(r\)-ReLUmax 在机制上构成了核回归视角下的连续版核采样（nucleus / top-\(p\) selection）：与硬性固定邻居个数的 top-\(k\) 截断不同，\(r\)-ReLUmax 根据得分差距动态调整活跃集规模——当候选键得分高度密集时自动扩充保留数量，而在最佳候选显著胜出时自动收缩为极稀疏分布。

损失函数 / 训练策略

在 Memory Mosaics 架构中，Transformer 块被解耦为情境记忆（contextual memory）与持久记忆（persistent memory）。情境记忆严格遵循因果 Nadaraya-Watson 估计，计算键的滑动平均 \(\bm{k}_n = \text{Norm}(\tilde{\bm{k}}_n + \lambda_\varphi \bar{\bm{k}}_{n-1})\)，值向量则前瞻一步构建 \(\bm{v}_n = \text{Norm}(\tilde{\bm{v}}_n + \lambda_\psi \tilde{\bm{v}}_{n+1})\) 以弥补排除自注意力主对角线的信息损失。

模型端到端采用标准交叉熵损失进行自回归训练，优化器为 Adam（\(\beta_1=0.9, \beta_2=0.95\)），配置余弦退火学习率调度与线性预热（warmup 2k 步，峰值学习率 \(5 \times 10^{-3}\)，12 层模型降至 \(1 \times 10^{-3}\)）。在 GPU 算子实现上，整流多项式核与 AdaSplash 稀疏内核无缝兼容，避免了传统 top-\(k\) 机制中耗时的离散排序操作，在长序列稀疏状态下展现出媲美 FlashAttention 的计算吞吐。

实验关键数据

主实验

主实验在探测正则语言推理能力的 RegBench 上下文学习基准上进行，评估不同样本量下模型预测最终词元的准确率（Acc \(\uparrow\)）与目标分布的全变差距离（TVD \(\downarrow\)）。

模型 / 核配置 1k Acc \(\uparrow\) 1k TVD \(\downarrow\) 5k Acc \(\uparrow\) 5k TVD \(\downarrow\) 10k Acc \(\uparrow\) 10k TVD \(\downarrow\) 40k Acc \(\uparrow\) 40k TVD \(\downarrow\)
RetNet 51.3 67.4 81.5 33.3 88.4 25.7 92.6 18.9
H3 51.5 65.1 70.9 47.4 76.9 41.4 87.7 25.4
Mamba 69.4 51.0 88.7 30.6 90.5 28.4 91.5 27.0
Standard Transformer 54.5 62.6 92.7 18.9 94.8 15.7 95.5 13.9
MM: Gaussian (Softmax) 93.9 34.1 94.6 17.7 95.0 15.9 95.4 14.2
MM: Gaussian kNN (\(k=16\)) 93.4 33.3 94.4 17.2 94.8 16.0 95.0 14.9
MM: Uniform kNN (\(k=16\)) 84.6 41.1 86.1 27.2 86.2 26.4 86.4 26.4
MM: Epanechnikov (normReLU) 84.3 43.6 90.1 24.6 91.6 21.7 93.3 18.6
MM: Epanechnikov (sparsemax) 94.2 32.5 94.4 17.1 94.9 15.9 95.4 14.0
MM: Epanechnikov (1-ReLUmax) 93.8 36.8 94.3 17.8 94.9 15.7 95.4 14.1
MM: Biweight (1.5-entmax) 93.8 33.0 94.7 16.6 95.2 15.7 95.5 13.8
MM: Triweight (4/3-entmax) 93.9 32.9 94.8 16.2 94.9 16.1 95.3 14.4

消融实验

消融实验聚焦于长上下文序列长度外推评测，考察多查询多词元联想召回（MQMTAR, \(L=4\)）以及序列排序（Sort, \(L=2\)）在上下文长度从 \(1\times\)（训练长度 64）外推至 \(64\times\)（测试长度 4096）时的精确匹配率（Exact Match Accuracy）。

模型 / 核机制 MQMTAR \(1\times\) MQMTAR \(8\times\) MQMTAR \(16\times\) MQMTAR \(32\times\) MQMTAR \(64\times\) Sort \(1\times\) Sort \(2\times\) Sort \(4\times\)
Standard Transformer 1.00 0.09 0.01 0.00 0.00 1.00 0.00 0.00
MM: Gaussian 0.99 0.85 0.64 0.14 0.00 1.00 0.06 0.00
MM: Gaussian kNN (\(k=32\)) 0.99 0.84 0.68 0.27 0.05 1.00 0.80 0.01
MM: Uniform kNN (\(k=32\)) 0.04 0.00 0.00 0.00 0.00 1.00 0.00 0.00
MM: Epanechnikov (normReLU) 1.00 0.66 0.44 0.33 0.31 1.00 0.05 0.00
MM: Epanechnikov (1-ReLUmax) 1.00 0.92 0.74 0.42 0.12 1.00 0.80 0.30
MM: Epanechnikov (sparsemax) 1.00 0.98 0.86 0.63 0.31 1.00 0.92 0.29
MM: Biweight (1.5-entmax) 1.00 0.95 0.75 0.51 0.23 1.00 0.39 0.26
MM: Triweight (4/3-entmax) 1.00 0.97 0.89 0.61 0.25 1.00 0.98 0.07

关键发现

  • 在数据稀缺的情境学习任务中（RegBench 1k-5k 样本），以 \(\alpha\)-entmax 为代表的自适应带宽整流多项式核（Triweight、Biweight、sparsemax）全面超越高斯核，获得更低的 TVD 与更高的准确率；而缺乏自适应带宽机制的 normalized ReLU 与离散均匀 kNN 则表现出显著性能坍塌。
  • 在长度外推场景下，高斯核与标准 Transformer 遭遇严重注意力分散，在 \(32\times\) 序列长度下准确率跌破 15%；而紧支撑核凭借精确局部截断展现出极强的外推稳定性，sparsemax 在 \(32\times\) 长度下依然维持 0.63 准确率。
  • 位置偏差消融显示，为 Epanechnikov 核追加 ALiBi 显式位置先验反而导致 \(64\times\) 长度下的 MQMTAR 准确率从 0.31 骤降至 0.14，证明基于数据驱动的紧支撑核稀疏检索本身即可完美解耦局部性与位置依赖，无需脆弱的启发式距离衰减。

亮点与洞察

  • 理论优美且完备：首次证明了非参数核回归中的经典紧支撑多项式核（Epanechnikov、biweight、triweight）与 Tsallis 熵诱导的 \(\alpha\)-entmax 稀疏注意力族具有严格的一一对应关系，且高斯/softmax 正好为其阶数趋于无穷的特例。
  • 机制创新兼顾实用：提出的 \(r\)-ReLUmax 通过将截断阈值绑定到序列内最大键得分，用极其轻量的代数操作彻底解决了传统 ReLU 归一化的除以零隐患，同时实现了连续可微的类核采样动态稀疏度。
  • 硬件友好的稀疏加速：相较于需要离散排序的 top-\(k\) 机制，紧支撑核与 AdaSplash 框架天然兼容，为兼顾统计学严谨性与现代 GPU 硬件执行效率的长文本注意力架构开辟了新路径。

局限与展望

  • 实验验证主要建立在 Memory Mosaics 架构与百兆参数级（<1B）小型模型及受控合成任务（BabiStories、RegBench、MQMTAR）上，尚未在数十亿参数的大规模 LLaMA 级主流因果语言模型全量预训练中完成端到端硬件加速验证。
  • \(r\)-ReLUmax 中的超参数 \(b\) 目前采用固定常数设置，未来可探索随序列长度或层深度动态自适应调整 \(b\) 与有效带宽 \(h\) 的元学习策略。

相关工作与启发

  • vs Softmax 高斯核回归 (Wang et al., 2025; Vaswani et al., 2017): 前者基于无限支撑的高斯核，导致长序列下注意力被背景噪声稀释；本文通过紧支撑核实现严格零权重的局部稀疏，彻底根除注意力分散。
  • vs 启发式 Top-\(k\) 截断 (Gupta et al., 2021): Top-\(k\) 机制需要离散排序使得梯度不连续，且固定的 \(k\) 无法感知局部相似度分布密度；本文的自适应带宽核与 \(r\)-ReLUmax 保证了连续可微性，并能自适应动态调整活跃词元数量。
  • vs 现代 Hopfield 网络 (Hu et al., 2023; Santos et al., 2025): 既往工作主要从能量函数与 Fenchel-Young 对偶角度推导稀疏联想记忆；本文开辟了非参数核密度估计与测试期 Nadaraya-Watson 回归的新视角，完成了两种学术流派的理论闭环。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次建立稀疏注意力与紧支撑核回归的形式化精确对应，理论体系严谨完备。
  • 实验充分度: ⭐⭐⭐⭐ 涵盖语言建模、正则语言上下文推理、超长序列外推及位置偏差消融，唯缺超大模型规模。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导清晰优雅，从动机分析到理论构建再到实证检验逻辑一气呵成。
  • 价值: ⭐⭐⭐⭐⭐ 为长上下文稀疏注意力设计提供了坚实的第一性原理理论支撑与实用的算子替代方案。

相关论文