Flow-based Spectral Kernel Learning for Nonstationary Attention¶

会议: NeurIPS 2026
arXiv: 2610.30465
代码: https://github.com/ziyang18/FSKA
领域: 学习理论
关键词: 非平稳核、谱核学习、正规化流、线性注意力、无投影注意力

一句话总结¶

提出了基于正规化流的谱核注意力（FSKA），将非平稳核学习转化为双变量谱密度估计并结合正随机特征实现线性计算复杂度，同时利用核表达能力吸收 QK 投影并固定正交 V 投影，削减约 40%–47% 投影参数。

研究背景与动机¶

标准自注意力机制本质上是一种基于相似度的平滑聚合机器，其软最大化（softmax）操作可以严格解释为一种特定的核函数（Kernel）。为了克服全自注意力随序列长度二次方膨胀的计算与显存瓶颈，大量高效注意力方法（如 Performer、Random Feature Attention 等）基于 Bochner 定理，利用随机傅里叶特征（RFF）逼近核函数，成功将注意力计算复杂度降至线性。然而，这些经典方法绝大多数依赖于预先设定且固定的平移不变（平稳）核函数，核的几何相似度仅依赖于 token 间的差值向量，无法根据输入数据的复杂上下文自适应地调整交互拓扑，极大地限制了模型捕捉富上下文特征的能力。

为了打破固定核函数的僵化约束，后续研究尝试通过高斯混合模型（GMM-RKS）或隐式神经网络（IKAN）来学习核的谱分布。但这两类方案各自走向了极端：高斯混合模型在拟合复杂多峰谱结构时需要大量混合分量，导致参数规模和优化难度急剧上升；而隐式神经网络则完全丧失了显式概率密度的解析形式，无法引入具有先验约束的谱正则化项，极易导致谱塌陷（Spectral Collapse）或剧烈过拟合，使得估计出的注意力亲和力矩阵发生畸变。此外，传统注意力架构无一例外地保留了复杂的查询（Query）、键（Key）和值（Value）三路可训练线性投影，这在长序列和大模型下引入了沉重的参数负担。

针对上述挑战，本文选择从泛函分析与广义 Bochner 定理切入，将非平稳核函数的参数化转化为双变量谱密度建模，并借助可逆正规化流（Normalizing Flows）实现显式且高度灵活的谱空间表征。核心 idea：利用正规化流显式建模非平稳双变量谱密度并通过结构化谱先验正则化，结合正随机特征实现数值稳定的线性注意力；并从泛函逼近理论证明高表达力核可直接作用于原始输入吸收 Query-Key 投影、仅保留固定的正交 Value 投影，实现参数精简与线性缩放的双赢。

方法详解¶

整体框架¶

FSKA 的整体流水线围绕“显式非平稳谱密度建模 $\to$ 稳定正随机特征线性化 $\to$ 投影解耦与参数极简”展开。首先，模型不再对输入进行可训练的 Query 和 Key 投影变换，而是直接将输入的原始词元表征输入到核注意力计算中；接着，利用 RealNVP 架构的正规化流将标准高斯分布变换为双变量谱密度分布，采样成对的随机频率向量；随后，将这些频率代入基于广义 Bochner 定理构建的双变量随机特征映射中，并通过偏移指数线性单元（$\text{ELU}+1$）转化为非负正随机特征，避免分母相消崩溃；最后，利用结合律将特征内积与固定正交矩阵变换后的 Value 向量重新组合，以线性时间与线性显存完成序列聚合。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列 X"] --> B["原始表征直入<br/>完全免除 Query 与 Key 投影"]
    A --> C["固定正交 Value 投影<br/>多头相互正交且不参与训练"]
    D["先验噪声 z ~ N(0, I)"] --> E["正规化流双变量谱密度建模<br/>RealNVP 可逆变换生成频率对"]
    E --> F["非平稳正随机特征映射<br/>Bochner 余弦正弦分解与 ELU+1"]
    B --> F
    F --> G["线性核注意力矩阵分解聚合<br/>交换求和次序实现 O(L) 复杂度"]
    C --> G
    G --> H["多头拼接与多层输出"]

关键设计¶

1. 基于正规化流的双变量谱密度建模：解决灵活建模与显式正则化的矛盾

传统平稳核仅依赖于差值向量，无法建模输入依赖的非平稳相似性。本文基于通用 Bochner 定理，任何连续复值有界核函数 $k(\mathbf{x}_1, \mathbf{x}_2)$ 均可唯一表示为双变量谱测度 $u(\mathrm{d}\bm{\omega}_1, \mathrm{d}\bm{\omega}_2)$ 上的积分： $$k(\mathbf{x}_1, \mathbf{x}_2) = \frac{1}{4}\int \Big(\exp\left(i(\mathbf{x}_1\bm{\omega}_1^\top - \mathbf{x}_2\bm{\omega}_2^\top)\right) + \exp\left(i(\mathbf{x}_1\bm{\omega}_2^\top - \mathbf{x}_2\bm{\omega}_1^\top)\right) + \exp\left(i(\mathbf{x}_1\bm{\omega}_1^\top - \mathbf{x}_2\bm{\omega}_1^\top)\right) + \exp\left(i(\mathbf{x}_1\bm{\omega}_2^\top - \mathbf{x}_2\bm{\omega}_2^\top)\right)\Big) u(\mathrm{d}\bm{\omega}_1, \mathrm{d}\bm{\omega}_2)$$ 为了兼顾高表达力与显式可计算的概率密度，FSKA 引入 $N_f$ 层 RealNVP 可逆变换 $F_\theta$。从低维隐变量 $\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ 出发，生成成对谱向量 $(\bm{\omega}_1, \bm{\omega}_2) = F_\theta(\mathbf{z})$。根据变量代换公式，其精确似然为 $p_\theta(\bm{\omega}_1, \bm{\omega}_2) = p_0(\mathbf{z}) |\det(\partial F_\theta(\mathbf{z}) / \partial \mathbf{z})|^{-1}$。为防止神经网络参数化流发生谱塌陷，FSKA 构建了基于非平稳高斯谱混合（NG-SM）作为参考先验的 KL 散度正则项： $$\mathcal{L}_{\text{spec}} = \mathrm{KL}\left(p_\theta(\bm{\omega}_1, \bm{\omega}_2) \,\|\, p_{\text{NGSM}}(\bm{\omega}_1, \bm{\omega}_2)\right)$$ 通过蒙特卡洛采样软约束流的几何形态，使得谱密度在保持自适应拓扑的同时兼具平滑与稳定性。

2. 保持非负性的双变量正随机特征映射：消除线性注意力分母抵消

在将上述双变量积分转化为蒙特卡洛随机特征估计时，根据定理 3.2，可构造 $M$ 维随机特征向量 $\bm{\phi}(\mathbf{x})$： $$\bm{\phi}(\mathbf{x}) = \sqrt{\frac{1}{2M}}\begin{bmatrix}\cos\left(\mathbf{x}\bm{\omega}_{1}^{(1:M/2)\top}\right) + \cos\left(\mathbf{x}\bm{\omega}_{2}^{(1:M/2)\top}\right)\\ \sin\left(\mathbf{x}\bm{\omega}_{1}^{(1:M/2)\top}\right) + \sin\left(\mathbf{x}\bm{\omega}_{2}^{(1:M/2)\top}\right)\end{bmatrix} \in \mathbb{R}^M$$ 然而，由于三角函数包含负值，在有限样本维度 $M$ 下，线性注意力归一化分母 $\sum_j \bm{\phi}(\mathbf{q}_i)^\top \bm{\phi}(\mathbf{k}_j)$ 极易因正负对消而趋于接近零甚至出现致命的负数，引发梯度爆炸与数值发散。为此，FSKA 引入单调非负变换： $$\widetilde{\bm{\phi}}(\mathbf{x}) = g(\bm{\phi}(\mathbf{x})), \quad g(u) = \text{ELU}(u) + 1 = \begin{cases} u + 1, & u > 0 \\ \exp(u), & u \leq 0 \end{cases}$$ 该变换强制使得特征映射的所有坐标均严格大于零，使得任意序列的归一化分母天然半正定，彻底根除了相消问题，并诱导了始终数值稳定的半正定亲和度矩阵 $\widetilde{k}_\theta(\mathbf{x}, \mathbf{y}) = \widetilde{\bm{\phi}}(\mathbf{x})^\top \widetilde{\bm{\phi}}(\mathbf{y})$。

3. 核表达力吸收 QK 投影与固定正交 V 投影：重塑极简参数结构

标准自注意力机制之所以需要可训练的 $\mathbf{W}_Q, \mathbf{W}_K$，是因为其底层核函数（如 Softmax 或高斯径向基）是固定的，必须依靠两路投影矩阵将输入映射到适应特定核几何的潜在空间中。根据本文定理 3.3（投影吸收定理），只要非平稳核函数族 $\mathcal{K}$ 在连续函数空间 $C(\mathcal{X} \times \mathcal{X})$ 中具备一致稠密性，对于任意紧集上的连续相似度度量与任意投影矩阵 $\mathbf{W}_Q, \mathbf{W}_K$，均存在某个核 $\tilde{k}_\epsilon \in \mathcal{K}$ 能够以任意精度 $\epsilon$ 一致逼近复合相似度 $f(\mathbf{x}_1, \mathbf{x}_2) = k(\mathbf{x}_1\mathbf{W}_Q, \mathbf{x}_2\mathbf{W}_K)$。

这意味着可学习的非平稳核完全有能力直接在原始输入空间建立 token 间的相似度关系，显式的 Query 与 Key 投影完全是参数冗余。与此同时，Value 投影承担的是信息聚合的表示子空间分配功能，而非相似度调节。为了保证多头注意力在无需训练参数的情况下依旧具备丰富的子空间表征多样性，FSKA 移除了所有可学习的 Value 权重，代之以预先初始化、相互严格正交的固定投影矩阵： $$(\mathbf{W}_V^{(h)})^\top \mathbf{W}_V^{(h')} = \mathbf{0}, \quad \forall h \neq h'$$ 这一设计直接将每层注意力模块中与投影相关的可训练参数削减了约 40%–47%，大幅压缩了模型足迹。

损失函数 / 训练策略¶

模型端到端训练的总损失函数由任务监督损失（如文本分类的交叉熵损失 $\mathcal{L}_{\text{task}}$）与双变量谱密度的 KL 散度正则化项复合而成： $$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \lambda_{\text{spec}} \mathcal{L}_{\text{spec}}$$ 其中默认设置谱正则化权重 $\lambda_{\text{spec}} = 10^{-3}$，基准 NGSM 先验采用包含 $Q=10$ 个分量的混合模型。优化器选用 Adam，权重衰减为 $10^{-4}$，学习率采用余弦退火策略调度。

实验关键数据¶

主实验¶

论文在长文档分类（SCOTUS、BBC News、20 Newsgroups 以及 arXiv）和句子级文本分类基准上对 FSKA 进行了全面验证。文档分类严格输入完整全文，不采用任何截断、分块或滑动窗口。

模型	SCOTUS Acc (%)	SCOTUS F1 (%)	BBC News Acc (%)	BBC News F1 (%)	arXiv Acc (%)	arXiv F1 (%)
Transformer	55.12 ± 1.24	35.76 ± 0.91	88.43 ± 1.52	88.01 ± 1.57	75.62 ± 1.61	74.90 ± 1.66
FlashAttention	55.68 ± 0.79	36.21 ± 0.87	88.91 ± 0.49	88.47 ± 0.54	76.18 ± 0.58	75.41 ± 0.63
Performer	52.38 ± 0.77	31.05 ± 0.88	83.74 ± 1.26	82.96 ± 1.21	69.81 ± 1.04	68.93 ± 1.09
Mamba	59.84 ± 0.68	37.91 ± 0.78	89.34 ± 0.46	88.96 ± 0.51	76.92 ± 0.54	76.24 ± 0.59
RKA	53.06 ± 0.99	31.72 ± 1.26	85.18 ± 0.48	84.71 ± 0.55	70.15 ± 1.58	69.42 ± 1.62
GMM-RKS	39.44 ± 0.93	12.31 ± 0.82	66.92 ± 0.75	65.81 ± 0.79	73.08 ± 0.51	72.47 ± 0.28
KPCA-Scaled	56.37 ± 0.71	37.42 ± 0.83	89.56 ± 0.43	89.18 ± 0.49	76.84 ± 0.55	76.09 ± 0.60
MIKAN	40.28 ± 0.86	17.05 ± 0.74	84.31 ± 0.59	83.92 ± 0.63	64.92 ± 0.70	64.15 ± 0.73
FSKA (本文)	58.91 ± 1.12	39.68 ± 0.71	90.76 ± 0.39	90.31 ± 0.44	78.43 ± 0.48	77.81 ± 0.53

在短文本分类基准（SST-2, SST-5, RT-Polarity, TREC）上，FSKA 同样表现亮眼：在 SST-2 上 Macro-F1 达到 78.74%（标准 Transformer 仅 76.07%），在 SST-5 上提升至 38.16%（标准 Transformer 仅 35.11%）。

消融实验¶

在 MNIST 图像序列分类任务上，论文对各项架构设计组件进行了控制变量消融分析：

模型配置	QK 投影	可学习 V	灵活核学习	正随机特征	测试准确率 Acc (↑)
标准 Transformer	✓	✓	✗	–	0.8556
Transformer 无 QKV 投影	✗	✗	✗	–	0.5267
隐式核注意力 (IKAN)	✓	✓	✓	✗	0.8656
FSKA (本文默认)	✗	✗	✓	✓	0.8722
FSKA + 直接使用 RFF	✗	✗	✓	✗	0.3956
FSKA + 启用可学习 V	✗	✓	✓	✓	0.8725

在长文档任务（BBC News 和 20NG）上对比“学习谱密度”与“固定双变量高斯谱”：学习谱密度将 BBC News 的 Macro-F1 从 63.50% 暴升至 90.31%，20NG 从 20.40% 提升至 28.86%，验证了流学习复杂谱拓扑的不可替代性。

关键发现¶

去 QK 投影的本质依赖于核表达力：在固定 softmax 的标准 Transformer 中粗暴剔除 QKV 投影会导致准确率崩塌（从 0.8556 跌至 0.5267）；而在 FSKA 中，由于非平稳谱核具备全局连续逼近能力，去掉 QK 投影反而以更少参数获得更优泛化性能（0.8722）。
正随机特征是线性稳定性的生命线：若不采用 $\text{ELU}+1$ 变换而直接使用经典 RFF，在有限维度 $M=96$ 时准确率断崖式下跌至 0.3956，其根本原因是有限抽样下的正负对消导致分母不稳定。
可学习 Value 冗余验证：为 FSKA 重新加入可训练的 Value 投影仅带来微不足道的增益（0.8722 变至 0.8725），但引入了大量参数，证明固定正交基在特征聚合层面完全足够。
显存与参数扩展性卓越：在序列长度扩展至 4096 的长序列测试中，标准二次自注意力显存迅速攀升到 80k–120k MB，而 FSKA 始终稳定在 10k MB 以下，与典型线性注意力保持完全一致的平缓线性增长斜率，同时参数减少 40%–47%。

亮点与洞察¶

泛函理论与工程架构的高度协同：论文不仅提出了用正规化流采样双变量谱密度的数学实现，更通过定理证明了核的强表达力可以“向下兼容”线性投影，从而为 Transformer 精简 QK 投影提供了扎实的理论支点。
正规化流在谱学习中的绝妙落地：不同于常规生成模型用流去拟合像素或文本，本文用流来建模核的傅里叶频域空间，既绕开了 GMM 分量过多的参数膨胀，又解决了隐式网络无法进行密度估计与 KL 先验正则化的死穴。
正交子空间分配代替学习聚合：用各头固定的正交投影矩阵代替可学习参数来划分特征子空间，这一设计思维极具启发性，有望迁移到 MoE 路由投影或轻量级微调架构中。

局限与展望¶

ELU+1 引入的内核偏差：严格理论上的 Bochner 无偏性仅存在于变换前的有符号特征映射；加入非负平移变换后，核相似度虽然彻底杜绝了分母崩溃，但在数学上诱导了基线漂移，且压缩了极端相似度的对比度（离线测试绝对差距保持在 ~0.8）。
大规模自回归预训练验证缺失：论文实验主要集中在自编码分类（双向上下文）和长文档分类上，尚未在百亿级自回归因果语言模型预训练（Causal LM）中验证因果掩码下的训练稳定性。
未来方向：探索无偏且天然非负的双变量特征分解方案；以及将 FSKA 的固定正交投影思路扩展至更深层级的 Transformer 模块架构剪枝与参数高效微调。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将正规化流引入双变量谱核学习，并给出吸收 QK 投影的理论证明。
实验充分度: ⭐⭐⭐⭐ 覆盖长文本、短文本、Toy 回归及系统性能剖析，消融极为详尽；若能补全大语言模型因果预训练则更佳。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备，从定理支撑到架构精简的叙事逻辑极其丝滑。
价值: ⭐⭐⭐⭐ 提出了一种高效、精简且兼具高表达力的非平稳注意力新范式，对轻量化长序列建模启发重大。