Learnable Low-Rank Polynomial Sketch for Effective Linear Attention¶

会议: NeurIPS 2026
arXiv: 2610.14219
领域: LLM效率
关键词: 线性注意力、多项式核、低秩张量分解、CP分解、长序列建模

一句话总结¶

本文提出可学习低秩多项式草图（LLoPS），通过理论证明最优多项式草图并采用 CP 低秩分解将其解耦为轻量投影与逐元素乘积，以 \(O(ndmtp)\) 的线性复杂度实现了严格优于现有方法的逼近误差与性能。

研究背景与动机¶

标准 Transformer 中的 Softmax 自注意力机制在序列长度上具有二次方时间与显存复杂度 \(\mathcal{O}(n^2)\)，这一计算瓶颈在处理长文档理解、代码大模型建模和检索增强生成（RAG）等超长上下文场景时尤为严重。线性注意力方法通过将核化点积与矩阵乘法结合律相结合，将状态更新简化为随时间累积的循环记忆状态，从而使计算与显存复杂度降低至随序列长度呈线性 \(\mathcal{O}(n)\) 增长。然而，以往许多线性特征映射（如基于随机特征或简单激活函数）在实际大规模任务中相较标准注意力表现出明显的性能退化。

近期研究表明，高阶多项式核是逼近 Softmax 指数核极具潜力的候选方案，原因在于高次幂能够模拟 Softmax 注意力分布中的“稀疏尖锐（sparse and spiky）”特性——强烈放大高相关性的查询-键点积，同时抑制无关的背景干扰。然而，直接计算 \(p\) 阶多项式特征映射需要显式展开阶数高达 \(\mathcal{O}(d^p)\) 的高维单项式外积，在高维向量表征下在算力和显存上完全不可行。为了规避此瓶颈，先前工作主要依赖启发式或受限的近似手段：PolySketchFormer 采用固定、数据无关的随机草图矩阵进行投影；BASED 则先通过低维线性投影将向量压缩至 \(r \ll d\) 的极小子空间再计算低阶（2阶）精确多项式特征。

这些既有方案在多项式核逼近的理论视角下仍是次优的。数据无关的随机投影未结合表征分布，无法在紧凑维度下达到理论最优；而预先降维再展开的做法等价于施加了极其苛刻的低秩结构约束。核心 idea：本文证明了固定输出特征维度下存在最小化多项式核截断误差的最优线性投影映射，并基于 CP 张量秩分解将其高效结构化为多组可学习投影矩阵的列式 Kronecker 积，使模型能够在不显式构建 \(\mathcal{O}(d^p)\) 外积的情况下以轻量矩阵运算精确逼近最优高阶多项式注意力。

方法详解¶

整体框架¶

LLoPS 的目标是在固定的输出特征维度 \(m\) 下，为 Query 矩阵 \(\bm{Q} \in \mathbb{R}^{n \times d}\) 与 Key 矩阵 \(\bm{K} \in \mathbb{R}^{n \times d}\) 分别学习高效的特征映射 \(\phi_Q(\bm{Q}), \phi_K(\bm{K}): \mathbb{R}^{n \times d} \rightarrow \mathbb{R}^{n \times m}\)，使得其内积逼近 \(p\) 阶多项式核 \((\bm{Q}\bm{K}^\top)^{*p}\)，并以线性注意力形式完成对 Value 矩阵 \(\bm{V}\) 的聚合：

\[\mathrm{Attn}(\bm{Q}, \bm{K}, \bm{V}) = \phi_Q(\bm{Q}) \left( \phi_K(\bm{K})^\top \bm{V} \right)\]

整体流程首先由理论分析指出最优投影矩阵 \(\bm{W}^* \in \mathbb{R}^{d^p \times m}\) 的存在性；接着将投影矩阵的每一列重塑为 \(p\) 阶张量并进行秩为 \(t\) 的 CP（CANDECOMP/PARAFAC）张量分解；最后通过列式 Kronecker 积（Column-wise Kronecker Product）的数学恒等式，将对高维外积 \(\psi_p(\bm{x})=\mathrm{vec}(\bm{x}^{\otimes p})\) 的高维投影转化为对原始输入向量分别进行 \(p\) 次轻量线性投影并在通道维度上做逐元素乘积与求和。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    InQ["输入 Query / Key 向量<br/>Q, K in R^{n x d}"] --> P1["多阶投影矩阵变换<br/>计算 Q Theta_{j,r} 与 K Theta_{j,r}"]
    P1 --> P2["CP 分解逐元素 Hadamard 积<br/>各阶特征沿通道做 element-wise 乘积"]
    P2 --> P3["秩分量聚合与非负约束<br/>跨 t 个 rank 分量求和 (可选平方保证非负)"]
    P3 --> Out["线性注意力循环聚合<br/>phi_Q(Q) (phi_K(K)^T V)"]

关键设计¶

1. 最优多项式草图存在性定理：确立最小截断误差的理论边界

以往线性注意力特征映射多基于启发式激活函数或数据无关的随机哈希，缺乏关于多项式逼近精度的明确保证。针对这一问题，本文首先从低秩逼近理论出发探讨最佳压缩。令 \(\psi_p(\bm{x}) = \mathrm{vec}(\bm{x}^{\otimes p}) \in \mathbb{R}^{d^p}\)，目标多项式核矩阵为 \(\bm{A} = (\bm{Q}\bm{K}^\top)^{*p} = \psi_p(\bm{Q}) \psi_p(\bm{K})^\top\)。设 \(\bm{A}\) 的奇异值分解（SVD）为 \(\bm{U} \bm{\Sigma} \bm{V}^\top\)（其中奇异值降序排列）。根据 Eckart-Young-Mirsky 定理，对于任意秩不超过 \(m\) 的矩阵分解 \(\phi_Q(\bm{Q})\phi_K(\bm{K})^\top\)，其在 Frobenius 范数下的最小逼近误差严格由截断奇异值确定：

\[\min_{\phi_Q, \phi_K} \|\phi_Q(\bm{Q})\phi_K(\bm{K})^\top - (\bm{Q}\bm{K}^\top)^{*p}\|_F^2 = \sum_{i=m+1}^{n} \sigma_i^2\]

本文证明了（Theorem 3.1）：无需引入复杂的非线性神经网络映射，仅需存在两个线性投影矩阵 \(\bm{W}_Q^*, \bm{W}_K^* \in \mathbb{R}^{d^p \times m}\)，使得 \(\phi_Q^*(\bm{Q}) = \psi_p(\bm{Q})\bm{W}_Q^*\) 且 \(\phi_K^*(\bm{K}) = \psi_p(\bm{K})\bm{W}_K^*\)，即可精确达到该理论下界 \(\bm{U}_m \bm{\Sigma}_m \bm{V}_m^\top\)。这一理论结论揭示了高维外积特征经由适当的线性映射足以实现最优压缩，为后续架构设计提供了坚实的理论落脚点。

2. CP 张量分解与列式 Kronecker 积解耦：规避指数维度的快速运算

尽管最优投影矩阵 \(\bm{W}^*\) 理论存在，但其行数高达 \(d^p\)，显式存储和相乘的时间与空间复杂度均为 \(\mathcal{O}(n m d^p)\)，在实际大模型中依然不可计算。为化解这一指数爆炸难题，本文提出将投影矩阵 \(\bm{W}^*\) 的第 \(i\) 列向量 \(\bm{W}_i \in \mathbb{R}^{d^p}\) 重塑为一个 \(p\) 阶张量 \(\mathcal{W}_i \in \mathbb{R}^{d \times \dots \times d}\)，并利用张量秩分解（CP 分解）将其展开为 \(t\) 个秩-1 张量的线性组合：

\[\mathcal{W}_i \approx \sum_{r=1}^{t} \lambda_r^{(i)} \bm{w}_{1,r}^{(i)} \otimes \bm{w}_{2,r}^{(i)} \otimes \dots \otimes \bm{w}_{p,r}^{(i)}\]

基于外积内积等于各分量内积连乘的代数性质，向量与该张量化投影的内积可转换为 \(\langle \bm{W}_i, \mathrm{vec}(\bm{x}^{\otimes p}) \rangle = \sum_{r=1}^{t} \lambda_r^{(i)} \prod_{j=1}^{p} \langle \bm{w}_{j,r}^{(i)}, \bm{x} \rangle\)。通过定义参数矩阵 \(\bm{\Theta}_{j,r} \in \mathbb{R}^{d \times m}\)（第 \(i\) 列为 \(\lambda_r^{(i)} \bm{w}_{j,r}^{(i)}\)），整个投影矩阵在代数上严格等价于这组矩阵的列式 Kronecker 积（\(\odot_{col}\)）：

\[\hat{\bm{W}} = \sum_{r=1}^{t} \left( \bm{\Theta}_{1,r} \odot_{col} \bm{\Theta}_{2,r} \odot_{col} \dots \odot_{col} \bm{\Theta}_{p,r} \right)\]

由此，计算特征映射时完全绕过了对高维向量 \(\psi_p(\bm{Q})\) 的实例化，直接转化为先计算 \(p\) 次低维矩阵乘法 \(\bm{Q}\bm{\Theta}_{j,r}\)，再执行逐元素哈达玛积与秩累加：\(\psi_p(\bm{Q})\hat{\bm{W}} = \sum_{r=1}^{t} \prod_{j=1}^{p} \bm{Q}\bm{\Theta}_{j,r}\)。计算复杂度从 \(\mathcal{O}(nmd^p)\) 骤降至 \(\mathcal{O}(ndmtp)\)。在实际工程实现中，设定 \(t=1\) 即可达到优异的拟合精度并控制参数量，此时映射形式简化为 \(p\) 个投影矩阵变换后的逐元素连乘，大幅提升了在现代 GPU 上的矩阵乘法并行利用率。

3. 统一多项式核视角的泛化与误差保证

现有主流近似多项式核方法在数学本质上均是 LLoPS 框架的受限特例。具体而言：(1) 元素级多项式核（如 Flatten Transformer）：形式为 \(\phi(\bm{x}) = (\bm{\Theta}\bm{x})^{*p}\)，等价于 LLoPS 强制所有 \(p\) 个投影矩阵完全共享参数（即 \(\bm{\Theta}_1 = \dots = \bm{\Theta}_p\)）；(2) BASED：通过 \(\bm{W} \in \mathbb{R}^{m' \times d}\) 降维后再展开为二次项 \(\mathrm{vec}(\bm{W}\bm{x} \otimes \bm{W}\bm{x}) = (\bm{W} \otimes_{\mathrm{kron}} \bm{W}) \mathrm{vec}(\bm{x}\bm{x}^\top)\)，等价于令 \(\bm{\Theta}_1^\top = \bm{W} \odot_{col} \bm{1}\) 与 \(\bm{\Theta}_2^\top = \bm{1} \odot_{col} \bm{W}\) 的特殊参数约束；(3) PolySketchFormer：使用满足 AMM 性质的固定稀疏随机矩阵进行张量草图构建，等价于将 \(\bm{\Theta}_j\) 固定为未经数据驱动优化的随机稀疏矩阵。

此外，本文定理（Theorem 3.4）给出了 LLoPS 特征映射逼近最优投影矩阵 \(\bm{W}^*\) 的误差解析：

\[\min_{\bm{\Theta}_1, \dots, \bm{\Theta}_p} \|\bm{W}^* - \bm{\Theta}_1 \odot \dots \odot \bm{\Theta}_p\|_F^2 = \|\bm{W}^*\|_F^2 - \sum_{i=1}^{m} \tau_i^2\]

其中 \(\tau_i\) 为重塑张量 \(\mathcal{W}_i\) 的张量谱范数（最大张量奇异值）。由于 LLoPS 放开了参数共享与固定随机约束，并随端到端训练数据自适应优化，其表示空间严格覆盖上述基线，在相同输出特征维度下必然取得严格更小或相等的逼近误差。

损失函数 / 训练策略¶

LLoPS 与标准自回归语言模型完全兼容，在预训练与微调阶段直接采用端到端的因果语言建模交叉熵损失进行优化。在具体的实验部署中： - 权重继承与继续训练：由于 BASED 的注意力核可精确嵌入为 LLoPS 的特例，为节约从头预训练的高昂开销，模型直接加载公开的 BASED 预训练权重作为初始点，并在 SlimPajama-6B 数据集上进行相同步数的持续训练（Continual Training），从而在严格控制训练步数与数据分布的条件下独立评估核函数改进的净收益。 - 阶数设置与非负约束：在 360M 模型与 LRA 任务中，多项式阶数默认设为 \(p=4\)；在 1B 模型中设为 \(p=2\)。若针对特定长序列训练需要保证注意力权重的严格非负性以增强数值稳定性，可在输出前应用平方操作，即 \(\phi(\bm{Q}) = (\prod_{j=1}^{p} \bm{Q}\bm{\Theta}_j)^{*2}\)。

实验关键数据¶

主实验¶

评估覆盖了两大核心基准：(1) 基于 LM-Evaluation-Harness 的 9 个零样本 NLP 基准（涵盖常识推理与信息抽取任务）；(2) 面向长序列上下文建模的 Long Range Arena (LRA) 基准（输入长度 1k-4k）。

在零样本 NLP 评估中，在 360M 与 1B 规模下对比了 Softmax 注意力、Mamba、BASED 以及持续训练对照组（BASED+CT）：

模型规模	方法	Wino	PIQA	LAM	Hella	ARC-E	ARC-C	FDA	SWDE	SQuAD	平均准确率 (%)
360M	Softmax Attn	51.78	63.71	29.32	33.60	42.85	24.32	57.89	56.08	27.88	43.05
360M	Mamba	50.36	63.76	27.89	33.88	42.42	24.66	5.90	17.37	24.83	32.34
360M	BASED	51.30	63.55	26.64	33.23	41.88	24.23	17.42	27.18	27.41	34.76
360M	BASED+CT	50.91	65.40	27.73	36.85	43.10	25.51	8.35	30.87	26.27	35.00
360M	BASED+LLoPS (本文)	53.04	66.70	29.23	37.07	43.94	25.51	7.53	31.41	31.94	36.26
1B	Softmax Attn	52.49	67.30	38.07	39.12	51.05	22.10	68.78	68.68	35.79	49.26
1B	Mamba	51.30	66.81	38.87	39.34	46.80	26.11	11.07	28.17	29.46	37.55
1B	BASED	51.62	66.70	37.86	39.10	45.54	25.60	24.05	42.84	34.58	40.88
1B	BASED+CT	51.22	68.44	36.43	41.59	47.01	27.30	24.41	46.80	36.23	42.16
1B	BASED+LLoPS (本文)	51.62	68.72	38.11	42.69	47.31	27.13	26.32	47.97	36.36	42.91

在 Long Range Arena (LRA) 长上下文基准上，LLoPS (\(p=4\)) 显著超越了包括 PolySketchFormer 和 BASED 在内的全部线性注意力基准：

方法	Image	Text	Retrieval	ListOps	Pathfinder	平均准确率 (%)
Softmax Attn	64.36	82.96	83.23	48.49	70.45	69.90
Nyström	58.66	66.19	83.90	45.46	69.71	64.78
Skyformer	58.54	65.25	79.55	46.77	72.73	64.57
Hedgehog	52.68	66.59	82.24	41.03	74.16	63.34
PolaFormer	47.74	73.06	80.50	38.86	70.53	62.14
Flatten Transformer	51.48	66.13	78.72	38.51	70.27	61.02
PolySketchFormer	51.40	64.60	81.10	43.00	66.20	61.26
BASED	51.13	65.73	80.76	43.55	72.98	62.83
LLoPS (本文)	57.54	73.94	82.17	48.29	73.84	67.16

消融实验¶

论文针对多项式核逼近的相对平均绝对误差（RMAE）、多项式阶数 \(p\) 以及特征维度 \(m\) 开展了充分的分析与消融研究。

1. 多项式核逼近相对平均绝对误差（RMAE）对比

在训练于合成联想召回（Associative Recall）任务的 Transformer 模型中，提取不同层（Layer 0, Layer 8, Layer 15）的查询与键向量，评估不同方法对 \(p=2, 3\) 真实多项式核 \((\bm{q}^\top \bm{k})^p\) 的逼近误差（RMAE 越低越好）：

阶数 \(p\)	方法类型	映射方案	Layer 0	Layer 8	Layer 15	平均 RMAE
\(p=2\)	数据无关草图	PolySketchFormer (PSF)	0.186	0.247	0.445	0.293
\(p=2\)	神经网络映射	2层 MLP (NN)	0.035	0.070	0.058	0.054
\(p=2\)	结构化低秩草图	LLoPS (本文)	0.008	0.032	0.025	0.022
\(p=3\)	数据无关草图	PolySketchFormer (PSF)	0.740	0.738	0.745	0.741
\(p=3\)	神经网络映射	2层 MLP (NN)	0.059	0.114	0.143	0.105
\(p=3\)	结构化低秩草图	LLoPS (本文)	0.018	0.040	0.058	0.038

2. 合成联想召回（AR）任务上阶数 \(p\) 与草图维度 \(m\) 的影响

配置参数	设定值	AR 准确率 (%)	说明
多项式阶数 \(p\) (\(m=256, d=128\))	\(p=2\)	88.7	低阶多项式，尖锐度稍欠缺
多项式阶数 \(p\) (\(m=256, d=128\))	\(p=4\)	93.2	表达能力与近似难度的最佳平衡点
多项式阶数 \(p\) (\(m=256, d=128\))	\(p=6\)	86.6	阶数过高导致在固定 \(m\) 下逼近误差与过拟合增加
隐藏维度 \(d=128\) (BASED对比)	BASED (\(r=16\))	91.9	二次项展开
隐藏维度 \(d=128\) (LLoPS)	LLoPS (\(p=4, m=256\))	93.2	相对提升 +1.3%
隐藏维度 \(d=192\) (BASED对比)	BASED (\(r=16\))	98.1	二次项展开
隐藏维度 \(d=192\) (LLoPS)	LLoPS (\(p=4, m=256\))	99.0	相对提升 +0.9%

关键发现¶

拟合精度与收敛稳定性的双重超越：LLoPS 不仅在逼近误差上以数量级优势超越基于随机草图的 PSF（\(p=2\) 时平均 RMAE 0.022 vs 0.293），同时也优于深层非线性 MLP 拟合（0.054）。这一现象直接印证了定理 3.1 的预测：线性映射在外积空间中已具完备逼近能力，复杂的非线性结构不仅没有理论增益，反而引入了额外的非凸优化困难。
阶数 \(p\) 的折中权衡机制：消融实验揭示多项式阶数并非越高越好。从 \(p=2\) 增加至 \(p=4\) 时，由于更好地复现了 Softmax 的稀疏尖锐选择性，AR 任务准确率显著提升至 93.2%；但当进一步提升至 \(p=6\) 时，受限于固定草图维度 \(m\) 下逼近难度激增以及参数增多带来的过拟合倾向，准确率回落至 86.6%。
实际推理吞吐效率更佳：在端到端耗时基准评测（序列长度 1,000 至 10,000）中，LLoPS 的前向耗时不仅远低于二次复杂度的 Softmax 注意力，也优于同维度的 BASED。这是因为 BASED 依赖显式的低维外积展开，而 LLoPS 完全由标准矩阵乘法与逐元素积构成，在现代硬件张量核心上的计算密度与缓存局部性更为理想。

亮点与洞察¶

理论与架构的无缝贯通：将多项式核逼近抽象为张量外积的最优线性草图问题，巧妙运用 CP 张量分解的代数变换将其拆解为轻量投影的逐元素积，兼顾了数学优雅性与硬件加速友好性。
广义包容现有方案：揭示了包括 BASED、Flatten Transformer 与 PolySketchFormer 在内的多项式线性注意力方法皆为本文框架在特定约束（全参数共享、固定随机矩阵、特殊分块约束）下的特例，提供了一个统一的理论分析谱系。
端到端即插即用性：无需复杂重构状态更新规则，仅通过改造查询与键向量的特征映射函数即可实现。这种纯粹的特征映射设计可无缝叠加到 Gated DeltaNet、Kimi Linear 等前沿的线性循环状态更新规则中。

局限与展望¶

原预训练数据缺失导致的控制变量局限：由于 BASED 最初采用的 800GB The Pile 数据集因版权下架，作者无法在完全相同的环境下从头预训练，而是在 SlimPajama-6B 上继续微调。尽管在严格配对对照（BASED+CT）下验证了优越性，但缺少超大规模（数十亿至数万亿 Token）从头预训练的完整 Scaling Law 检验。
高阶逼近面临的维度维数灾难：虽然 \(t=1\) 的设计大幅压缩了计算量，但在阶数极高时（\(p \ge 6\)），单秩张量拟合对于极高阶交互项的捕获能力仍然受限；如何在不引入过多参数的前提下动态扩展张量秩 \(t\) 仍是值得探索的课题。
未来方向：探索将可学习多项式草图与数据驱动的状态遗忘衰减门控（如 Mamba2 的半可分离矩阵、Delta Rule）深度融合，构建下一代高性能长文本基础模型。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从张量秩分解与草图理论严格证明了多项式线性注意力的最优线性逼近并给出了可计算结构。
实验充分度: ⭐⭐⭐⭐☆ 覆盖零样本 NLP、LRA、理论逼近误差与合成 AR 任务，受限于 Pile 数据集缺失未做从头大规模预训练。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备，从问题定义、最优定理、CP分解到实际算法流程一气呵成。
价值: ⭐⭐⭐⭐⭐ 为长上下文大语言模型中的线性注意力核设计提供了坚实的理论基准与实用高性能方案。