Laplacian Heads Improve Transformers by Smoothing Token Representations

会议: NeurIPS 2026
arXiv: 2602.09297
OpenReview: https://openreview.net/forum?id=o6wVXdc2L1
论文: https://arxiv.org/abs/2602.09297
领域: 自监督学习
关键词: Laplacian Head, Token 平滑, 序列内方差控制, 图热扩散, 神经坍塌 (Neural Collapse)

一句话总结

将 Transformer 中部分常规注意力头替换为无额外参数的拉普拉斯头（Laplacian Head），直接控制序列内表示方差并模拟图热扩散，在图像分类、语言模型和自监督学习中普遍提升表征几何质量与任务性能。

研究背景与动机

标准的 Transformer 架构通过多头注意力与残差连接不断融合并更新 token 表示。每一层的注意力矩阵通常是行随机矩阵，其更新机制实质上是取全序列 token 的加权平均并叠加到原特征上。在既有文献中，“过度平滑（Oversmoothing）”和“秩坍塌（Rank Collapse）”一直被视为注意力层堆叠过程中的核心失效模式——研究者们普遍认为，若不同 token 之间的表示越来越相似，模型就会损失表达能力，因而设计了大量复杂的正则化损失、跳跃连接缩放或正交约束来竭力抑制平滑。

然而，这种将“平滑”一律视为缺陷的传统视角忽略了一个关键机制性盲区：标准 Transformer 依靠残差连接叠加注意力加权均值，实质上只能直接调控 token 表示的均值，却缺乏显式、直接控制序列内部方差（within-sequence variance）的数学算子。序列内 token 之间的相对离散程度只能被动地受到特征变换影响，无法在层间有针对性地去噪或聚焦在紧凑低维流形上。现有注意力难以在保持类别/语义区分度的同时主动收拢样本内部的不相关扰动。

本文的切入角度是重新审视平滑在表征学习中的建设性角色。既然行随机注意力矩阵可自然视作 token 构成的有向图邻接矩阵，那么对应的随机游走归一化图拉普拉斯算子即可直接刻画每个 token 与其邻域加权均值的偏差。核心 idea：将多头注意力中的一部分注意力矩阵 \(P^{(i)}\) 替换为图拉普拉斯算子 \(I - P^{(i)}\)，无需引入任何额外参数，直接通过图热扩散步显式控制序列内表示方差，使 token 表征在任务目标下收敛为更优的几何结构。

方法详解

整体框架

在标准 Transformer 的单层多头自注意力机制中，输入序列特征矩阵为 \(X \in \mathbb{R}^{T \times d}\)。模型维护 \(H\) 个并行注意力头，其中第 \(i\) 个头计算归一化的 Softmax 注意力矩阵 \(P^{(i)} = \mathrm{softmax}(X W_{Q_i} W_{K_i}^\top / \sqrt{d_k}) \in \mathbb{R}^{T \times T}\)，并通过 Value 投影与输出投影得到特征增量更新。

本文提出的改进不改变任何投影矩阵的参数量与维度，而是将总共 \(H\) 个头拆分为两组不相交的集合：常规注意力头集合 \(\mathcal{A}\)（大小为 \(|\mathcal{A}| = H - |\mathcal{L}|\)）与拉普拉斯头集合 \(\mathcal{L}\)（大小为 \(|\mathcal{L}|\)）。对于属于 \(\mathcal{L}\) 的头，将注意力权重矩阵 \(P^{(i)}\) 直接替换为其图拉普拉斯算子 \(I - P^{(i)}\)（其中 \(I \in \mathbb{R}^{T \times T}\) 为单位矩阵）。模型在前向传递中协同计算两类头的更新量，并在残差连接处与原特征相加。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列特征 X"] --> B["Q/K/V 线性投影"]
    B --> C["注意力概率矩阵 P"]
    C --> D["双路算子分流<br/>注意力头组 vs 拉普拉斯头组"]
    D --> E["均值更新与方差控制<br/>P 调整均值 / (I - P) 调控序列内方差"]
    E --> F["残差叠加与图热扩散<br/>X + Attention + Laplacian"]
    F --> G["优化后的紧凑几何表示"]

关键设计

1. 序列内方差显式调控：均值更新与差值缩减的正交解耦 在标准自注意力残差项中，单头输出为 \(P^{(i)} X W_{V_i} W_{o_i}\)。由于 \(P^{(i)}\) 每一行和为 1，它为序列中的每一个 token 计算了全序列其他 token 的凸组合。将该项加回自身相当于直接拉动 token 表示的均值，但无法精准压缩或调节序列内部各个 token 围绕局部均值的离散程度。拉普拉斯头采用算子 \((I - P^{(i)}) X W_{V_i} W_{o_i}\)，其实质是先计算每个 token 投影向量与其注意力加权平均之间的差值（偏差向量）。当该偏差向量经过可学习的残差投影加回原特征时，模型获得了沿注意力邻域几何方向主动收拢或扩展样本内部离散度的自由度。数学更新形式定义为：

\[X \leftarrow X + \sum_{i \in \mathcal{A}} P^{(i)} X W_{V_i} W_{o_i} + \sum_{i \in \mathcal{L}} (I - P^{(i)}) X W_{V_i} W_{o_i}\]

这种设计使 Transformer 在同一注意力块内同时具备了“调整全局/局部均值”与“显式压缩/拉伸序列内方差”的双重调控能力，从根本上打破了标准注意力无法直接控制表示分散程度的结构局限。

2. 基于注意力图的热扩散机制：离散动力学的物理对齐 将一个包含 \(T\) 个 token 的序列视作图 \(\mathcal{G} = (\mathcal{V}, \mathcal{E})\)，各 token 对应节点，自注意力矩阵 \(P^{(i)}\) 的元素 \(P^{(i)}_{jk}\) 对应从节点 \(k\) 到节点 \(j\) 的有向边权重。因每行和为 1，\(P^{(i)}\) 是天然的行随机转移矩阵，而 \(I - P^{(i)}\) 即为随机游走归一化图拉普拉斯算子。若将任意特征信号作用于图节点上，经典图热传导方程由微分动力学给出：

\[\frac{d}{dt} x(t) = - (I - P^{(i)}) x(t)\]

当采用一阶欧拉离散化且步长 \(\Delta t = 1\) 时，\(x_{t+1} = x_t - (I - P^{(i)}) x_t = P^{(i)} x_t\) 恰好退化为标准注意力；而在特征残差更新中直接保留 \((I - P^{(i)})\) 算子，相当于在注意力自适应诱导的动态流形上直接执行受控的热扩散步。负拉普拉斯算子促使特征向邻居均值自然流动并平滑高频噪声，从而为 Transformer 在不丢失长程依赖的前提下带来了强大的局部降噪与结构诱导偏置。

损失函数 / 训练策略

论文完全沿用任务基线的默认损失函数与超参数，不增减任何损失项或可学习参数： - 图像分类：基于 DeiT-3 (ViT-B) 架构，在 CIFAR-10/100 上使用交叉熵损失配合 AdamW 优化器，在 ImageNet-1k 上使用二值交叉熵（BCE）损失配合 LAMB 优化器，均采用 Cosine 退火调度与 RandAugment/3-Aug 数据增强。 - 自回归语言模型：采用 561M 参数的 Decoder-only 架构（引入 QK-Norm、RoPE、Squared ReLU、Soft Logits Clipping），在 11.2B FineWeb-Edu 上预训练，继而在 SmolTalk、MMLU、GSM8K 等混合数据集上做 Mid-training 与 SFT。 - 自监督学习：基于 DINO (ViT-S/16)，结合 iBOT 掩码重构头与 KoLeo 正则损失，在 ImageNet-1k 上自监督预训练 100 个 epoch。

实验关键数据

主实验

在图像分类（DeiT-3 ViT-B）、语言模型推理（561M GPT-2）以及自监督学习（DINO ViT-S）三大主流范式下，增加拉普拉斯头均取得了超越纯注意力基线的显著性能增益。

数据集 / 评测基准	指标	基线 (0 Laplacian)	本文最佳 Laplacian 配置	性能提升
CIFAR-10 (ViT-B)	Top-1 准确率 (%)	95.35 ± 0.16	95.88 ± 0.14 (11L)	+0.53%
CIFAR-100 (ViT-B)	Top-1 准确率 (%)	75.12 ± 0.46	77.04 ± 0.40 (12L)	+1.92%
ImageNet-1k (ViT-B)	Top-1 准确率 (%)	81.45 ± 0.08	82.73 ± 0.19 (3L)	+1.28%
LM 推理综合均分 (561M)	Avg Accuracy (%)	23.93	30.09 (9L)	+6.16%
ARC-Easy (Zero-shot)	准确率 (%)	31.03 ± 4.29	47.12 ± 1.66 (5L)	+16.09%
ARC-Challenge (Zero-shot)	准确率 (%)	26.36 ± 1.29	34.22 ± 0.86 (5L)	+7.86%
GSM8K (Pass@10)	准确率 (%)	22.29 ± 0.13	24.00 ± 0.27 (3L)	+1.71%
DINO ImageNet 线性探测	Top-1 准确率 (%)	71.55	72.98 (3L / 5L)	+1.43%
DINO ADE20K 语义分割	mIoU	29.81	32.27 (5L)	+2.46

消融实验

以 ViT-B (总头数 12) 在各视觉分类数据集上的头数分配消融，以及 DINO ViT-S (总头数 6) 的拉普拉斯头数变化为例：

模型配置	CIFAR-100 Top-1 (%)	ImageNet-1k Top-1 (%)	DINO ADE20K mIoU	说明
基线 (0L, 纯 Attention)	75.12 ± 0.46	81.45 ± 0.08	29.81	传统多头自注意力
1 Laplacian (1L)	-	-	30.73	引入单头拉普拉斯平滑即见涨点
3 Laplacians (3L)	75.81 ± 0.14	82.73 ± 0.19	31.70	ImageNet 最佳配置，提点 1.28%
5 Laplacians (5L)	-	-	32.27	DINO 下游密集分割最佳，mIoU 涨 2.46
6 Laplacians (6L)	76.39 ± 0.26	82.70 ± 0.02	29.64	注意力头与拉普拉斯头对半开
9 Laplacians (9L)	76.80 ± 0.39	82.47 ± 0.08	-	语言模型上取得最高综评分数 (30.09%)
11 Laplacians (11L)	76.83 ± 0.26	82.40 ± 0.12	-	CIFAR-10 上取得最高准确率 (95.88%)
12 Laplacians (全 12L)	77.04 ± 0.40	82.32 ± 0.07	-	纯拉普拉斯更新，CIFAR100 相对涨幅近 2%

关键发现

• 方差转移诱导神经坍塌（Neural Collapse）：方差分析（ANOVA）表明，拉普拉斯头显著降低了序列内方差占总方差的比例（WithinSeqVar），将原本分散的能量重定向转移至类间方差（BetweenClassVar）。在最后一层，类间表征清晰收敛为单纯形等角紧框架（Simplex ETF），在分类边界上实现了高纯度的几何聚集。
• 语言模型的语义坍塌与谱衰减：在自回归模型中，拉普拉斯头加快了奇异值谱的衰减速率。解释 99% 方差所需的有效主成分维度从基线的 547.7 骤降至 430.4（降幅达 21%），使得具有相同预测目标（Next-token target）的 token 聚集成更紧致的簇，显著加深了语言模型的“语言坍塌（Linguistic Collapse）”。
• 自监督学习下的边界保留与低维凝聚：DINO 模型的 PCA 特征可视化显示，拉普拉斯头生成的表征在前 3 个主成分上展现出清晰锐利的物体边缘与微小细节（如头盔、球衣号码），直接推动下游密集预测 ADE20K mIoU 提升 2.46。

亮点与洞察

• 颠覆“过度平滑有害”的传统教条：传统观点视特征平滑为表达力下降的根源，而本文表明适度且受控的图拉普拉斯平滑反而是消除序列内部冗余扰动、促使语义表示向内在低维流形与单纯形几何收敛的有效动力学机制。
• 零额外参数的即插即用改造：仅需在矩阵加法时用单位矩阵相减 \((I - P)\)，不增加任何浮点参数量或计算开销，便能跨模态（视觉分类、代码/推理语言生成、自监督密集感知）一致获得稳健的性能收益。
• 图扩散与深度学习几何统一：成功将连续图热传导方程、序列内信噪比（SNR）和分类末期神经坍塌（Neural Collapse）理论串联，为分析 Transformer 的表征几何提供了坚实的物理与数学解释工具。

局限与展望

• 层间分配策略过于简单：全文在所有层中均硬性采用相同数量的拉普拉斯头，这一全局静态策略显然处于次优状态。更深或更浅层可能对均值传播与方差收缩有着动态不同的需求。
• 任务依赖的最佳头数差异：在全监督小数据集上全拉普拉斯（12L）最优，而在大规模 ImageNet 或自监督预训练中 3L~5L 最佳；如何针对不同数据尺度自适应学习头数配比尚未解决。
• 因果掩码下的单向扩散理论待完善：在自回归解码场景下，因果 Mask 导致注意力矩阵变为下三角，对应的图属于有向无环图，此时非对称拉普拉斯算子的扩散动力学与无向图热传导存在理论差异，值得深入探究。

相关工作与启发

• vs. 抗过度平滑方法（DropEdge, PairNorm, Nonlocal Functionals）：传统工作视 token 相似度上升为表征退化并竭力对抗；本文反其道而行之，不仅不压制平滑，反而主动构造拉普拉斯扩散机制，证明有益的平滑能够催化类间可分性与任务对齐。
• vs. 神经坍塌理论（Neural Collapse & Linguistic Collapse）：Papyan 等人证明了分类任务末期特征会自发塌缩为单纯形 ETF；本文从架构设计切入，证明拉普拉斯头是驱动序列内方差压缩并主动诱导神经坍塌的强力结构偏置。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 颠覆性反思过度平滑问题，以极简优雅的图拉普拉斯算子赋能 Transformer 方差控制。
• 实验充分度: ⭐⭐⭐⭐⭐ 涵盖图像分类、大规模自回归语言建模、自监督预训练三大领域，并辅以严谨的 ANOVA、PCA 谱分析及神经坍塌几何验证。
• 写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻透彻，物理图像（热传导）与数学推导自洽连贯。
• 价值: ⭐⭐⭐⭐⭐ 无需额外参数、跨模态通用性强，对理解与改进 Transformer 内部表征机制极具启发性。

相关论文

• [CVPR 2025] SATA: Spatial Autocorrelation Token Analysis for Enhancing the Robustness of Vision Transformers
• [NeurIPS 2026] How Data Augmentation Shapes Neural Representations
• [ICML 2026] The Geometry of Projection Heads: Conditioning, Invariance and Collapse
• [NeurIPS 2026] RATS! Patches Talk Through Registers: Emergent Parts in Register Attention Transformers
• [ICML 2026] NITP: Next Implicit Token Prediction for LLM Pre-training