Nexus高阶注意力机制

Nexus是一种新型的高阶注意力架构，通过递归框架增强Transformer的表示能力。该方法发表于ICLR 2025，能够在单层注意力内建模多跳关系，有效打破标准注意力的低秩瓶颈。¹

1. 低秩瓶颈问题

标准自注意力的局限

标准Transformer的自注意力机制定义为：

$$\text{Attention}(X)=\text{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$$

其中 $Q=W_{q}X$、$K=W_{k}X$、$V=W_{v}X$ 分别为查询、键、值矩阵。

核心问题：注意力权重矩阵 $A=\text{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)\in {\mathbb{R}}^{n\times n}$ 仅建模token对之间的直接成对交互，无法捕获高阶依赖关系。

多跳推理的困境

考虑三个token $x_i,x_j,x_k$ 的交互，标准注意力需要：

$$A_{ij}=\text{softmax}\left(\frac{Q_i{K}_j^{\top }}{\sqrt{d_k}}\right),A_{ik}=\text{softmax}\left(\frac{Q_i{K}_k^{\top }}{\sqrt{d_k}}\right),A_{jk}=\text{softmax}\left(\frac{Q_j{K}_k^{\top }}{\sqrt{d_k}}\right)$$

建模三元关系需要堆叠多层注意力或进行多步推理，导致：

• 计算负担增加
• 信息损失和梯度消失
• 训练/推理效率低下

理论分析：线性瓶颈

Bhojanapalli等人（2020）指出当 $d_k<n$ 时，标准注意力无法表达任意注意力权重 $A$。²

定理3.1（线性瓶颈）：

1. 给定 $N$ 个输入 $X_m\in {\mathbb{R}}^{n\times d}$ 和对应的目标行随机矩阵 $A_m\in {\mathbb{R}}^{n\times n}$，只要 $\text{rank}(\log (A_m))\le d_k$，总存在映射 $Q,K:{\mathbb{R}}^{n\times d}\to {\mathbb{R}}^{n\times d_k}$ 使得：

   $$\text{softmax}\left(\frac{Q(X_m)K(X_m{)}^{\top }}{\sqrt{d_k}}\right)=A_m$$

2. 若 $d<n-1$，存在 $\text{rank}(\log (A_m))=1$ 的 $A_m$，但对于所有线性变换 $Q(X)=X{W}_q$、$K(X)=X{W}_k$，式(14) 仍不成立。

推论：标准注意力机制甚至无法表达秩为1的对数注意力权重矩阵，这深刻揭示了其表达能力不足的本质。

---

2. Nexus架构

核心思想

与标准方法使用静态线性投影不同，Nexus通过嵌套自注意力机制动态精炼Query和Key表示：

$$\text{H-Attention}(X)=\text{Attention}\left({\text{Attention}}_q(X),{\text{Attention}}_k(X),V\right)$$

高阶注意力定义

首先，通过自注意力分别精炼查询和键：

$${\text{Attention}}_q(X)=\text{softmax}\left(\frac{Q{Q}^{\top }}{\sqrt{d_k}}\right)Q$$ $${\text{Attention}}_k(X)=\text{softmax}\left(\frac{K{K}^{\top }}{\sqrt{d_k}}\right)K$$

完整的高阶注意力为：

$$\text{H-Attention}(X)=\text{softmax}\left(\frac{{\text{Attention}}_q(X)\cdot {\text{Attention}}_k(X{)}^{\top }}{\sqrt{d_k}}\right)V$$

递归扩展

构建递归高阶注意力以捕获更深层的依赖：

$$H^2\text{-Attention}(X)=\text{Attention}\left({\text{H-Attention}}_q(X),{\text{H-Attention}}_k(X),V\right)$$

推广到 $m$ 阶：

$$H^m\text{-Attention}(X)=\text{Attention}\left(H^{m-1}{\text{-Attention}}_q(X),H^{m-1}{\text{-Attention}}_k(X),V\right)$$

┌─────────────────────────────────────────────────────────────┐
│                    Nexus Layer                               │
│  ┌─────────┐    ┌─────────┐    ┌─────────┐                 │
│  │   Q     │───→│Inner Attn│───→│Refined Q│──┐             │
│  │(Linear) │    └─────────┘    └─────────┘  │             │
│  └─────────┘                            ┌───┴───┐         │
│  ┌─────────┐    ┌─────────┐    ┌─────────┐    │         │
│  │   K     │───→│Inner Attn│───→│Refined K│───→│Attn→Out │
│  │(Linear) │    └─────────┘    └─────────┘    │         │
│  └─────────┘                            ┌───┴───┘         │
│  ┌─────────┐                             │                 │
│  │   V     │────────────────────────────→│                 │
│  └─────────┘                             └─────────────────┘
└─────────────────────────────────────────────────────────────┘

---

3. 动态Query/Key精炼

语义对齐机制

Nexus的核心洞见是：Query和Key是内层注意力循环的输出。这意味着：

1. 上下文感知投影：标准Transformer中 $K_i$ 仅依赖 $x_i$，而Nexus中 $K_i$ 通过内层注意力聚合了全局信息

2. 语义高亮：内层Key注意力呈现明显的垂直条纹，表明特定位置被许多后续token关注——这相当于在主注意力计算前识别并聚合全局相关语义

3. 预推理能力：这种设计在最终注意力计算前完成了”预推理”步骤，简化了主注意力的任务

注意力模式可视化

实验可视化表明：

注意力类型	特点	功能
Baseline	对角线强聚焦 + 左侧垂直带	建模序列顺序
Nexus Outer	与baseline类似的因果结构	保持基础能力
Nexus Inner Q	内层查询聚合	动态调整查询表示
Nexus Inner K	垂直条纹模式	语义高亮与全局聚合

---

4. 高阶相关性建模

多Token依赖

高阶注意力机制允许模型直接捕获多token依赖，无需堆叠多层：

$$\text{H-Attention}(X)=\text{softmax}\left(\frac{\bar{Q}\cdot {\bar{K}}^{\top }}{\sqrt{d_k}}\right)V$$

其中 $\bar{Q}={\text{Attention}}_q(X)$、$\bar{K}={\text{Attention}}_k(X)$。

三元关系建模

对于三个token $x_i,x_j,x_k$：

• 标准注意力：需要顺序成对计算 $A_{ij},A_{ik},A_{jk}$
• 高阶注意力：在单次注意力层中同时考虑其组合影响

与CoT推理的联系

递归高阶注意力类似于在注意力层内执行链式思维（Chain-of-Thought）推理：

$$H^m\text{-Attention}\approx \underset{m\text{ 次}}{\underbrace{\text{Reason}\circ \text{Reason}\circ \cdots \circ \text{Reason}}}$$

---

5. 参数共享策略

问题与解决方案

高阶机制的朴素实现需要为每个递归步骤使用独立的投影矩阵，导致参数量随 $m$ 线性增长。

权重共享策略

基于假设：将向量投影到Query或Key空间的语义变换在不同递归层次上基本相似。

修改递归步骤，使内层注意力复用外层权重 $\theta =\{W_q,W_k,W_v\}$：

$${\text{H-Attention}}_q(X;\theta )=\text{Attention}(X,X,X;\theta )\cdot W_q$$ $${\text{H-Attention}}_k(X;\theta )=\text{Attention}(X,X,X;\theta )\cdot W_k$$

参数量分析

配置	参数量	说明
标准注意力	$O(1)$	基准
m阶朴素实现	$O(m)$	随阶数线性增长
m阶权重共享	$O(1)$	与标准注意力相同

关键优势：Nexus的参数量与标准Transformer完全相同，仅增加计算密度（推理FLOPs），而不增加模型存储。

---

6. 复杂度分析

时间复杂度

标准自注意力的计算复杂度为 $O(n^2{d}_k)$，其中 $n$ 为序列长度。

对于 $m$ 阶高阶注意力：

$$T(m)=2\cdot T(m-1)+O(n^2{d}_k)$$

其中 $T(0)=O(n^2{d}_k)$，递归展开得：

$$T(m)=O(2^m\cdot n^2{d}_k)$$

实践权衡

• $m=2$：约2倍计算量，但显著性能提升
• $m=3$：约4倍计算量，进一步提升
• 实验表明 $m=2$ 是效率与性能的最佳平衡点

内存效率

方面	标准Transformer	Nexus
模型参数量	$P$	$P$（相同）
推理FLOPs	$F$	$\approx 2F$（$m=2$）
存储需求	$S$	$S$（相同）

---

7. 打破线性瓶颈的理论证明

核心洞察

Nexus通过非线性映射 $Q(X)={\text{Attention}}_q(X)$ 和 $K(X)={\text{Attention}}_k(X)$ 替代线性投影，从而增强表达能力。

高阶注意力的表达能力

定义非线性映射：

$$Q(X)={\text{Attention}}_q(X),K(X)={\text{Attention}}_k(X)$$

则 $Q(X)$ 和 $K(X)$ 是上下文感知的聚合表示，而非简单的线性变换。

注意力模式对比

可视化实验表明，Nexus的注意力矩阵展现出：

• 更复杂的互联模式
• 更高的连接度和多样性
• 更强的多面关系建模能力

这验证了Nexus能够捕获标准注意力和其它高阶注意力无法表达的依赖关系。

---

8. 实验结果

Pythia模型对比

在Pythia（70M-1B）上的零样本准确率对比：

规模	模型	ARC-C	ARC-E	Hellaswag	LogiQA	PiQA	SciQ	平均
70M	Pythia	0.208	0.359	0.356	0.276	0.569	0.615	0.397
70M	Nexus	0.204	0.382	0.358	0.287	0.586	0.685	0.417
160M	Pythia	0.200	0.385	0.380	0.260	0.600	0.686	0.419
160M	Nexus	0.211	0.405	0.385	0.285	0.605	0.713	0.434
410M	Pythia	0.225	0.394	0.375	0.285	0.601	0.708	0.431
410M	Nexus	0.226	0.415	0.384	0.294	0.608	0.733	0.443
1B	Pythia	0.232	0.440	0.395	0.296	0.625	0.758	0.458
1B	Nexus	0.230	0.455	0.399	0.290	0.636	0.777	0.465

关键发现：

• Nexus在所有规模上平均准确率均优于基线
• SciQ任务提升最大（+6%@70M），该任务需要多步推理
• 在需要复杂逻辑依赖的任务上优势明显

消融实验

模型配置	投影	共享	阶数	平均准确率
Baseline	-	-	-	0.397
Nexus-Q	Q	否	2	0.400
Nexus-QK	Q, K	否	2	0.409
Nexus-QKV	Q, K, V	否	2	0.409
Nexus-QK-Shared	Q, K	是	2	0.406
Nexus-Recursive	Q, K	是	3	0.415

结论：

1. Q和K的高阶处理是核心，V无需高阶
2. 权重共享仅轻微性能下降（0.409→0.406）
3. 增加递归深度（$m=3$）可进一步提升

Qwen2.5模型升级实验

在数学推理基准上的表现：

基础模型	方法	MATH-500	AIME24	GPQA-Diamond	平均
Qwen2.5-1.5B	Standard SFT	0.786	0.194	0.276	0.419
Qwen2.5-1.5B	Nexus-SFT	0.801	0.194	0.280	0.425
Qwen2.5-7B	Standard SFT	0.921	0.452	0.401	0.591
Qwen2.5-7B	Nexus-SFT	0.921	0.475	0.407	0.601

关键发现：Nexus可作为现有LLM的”架构升级套件”，在微调阶段即可带来推理能力提升。

---

9. 与相关工作的比较

vs. 高效注意力

方法	目标	权衡
Linformer	$O(n)$复杂度	低秩近似损失表达能力
Performer	$O(n)$复杂度	核近似可能不准确
Reformer	$O(n\log n)$	依赖LSH近似
Nexus	增强表达能力	参数不增加，计算略增

vs. 高阶方法

方法	应用领域	特点
Attention on Attention	图像描述	增加注意力层
Deformable Attention	视觉	可变形卷积思想
高阶关系Transformer	多模态	多级结构
Nexus	NLP + 多模态	递归+权重共享

vs. 标准Transformer

Nexus保留了标准Transformer的核心特性：

• 因果结构（对角线聚焦）
• 全局上下文建模
• 可扩展的深度和宽度

同时增强了：

• 多token依赖捕获
• 层级关系建模
• 推理密度

---

10. 实际应用建议

何时使用Nexus

适合场景：

• 多跳推理任务（LogiQA、ARC-C）
• 数学问题求解（MATH、AIME）
• 需要建模层级依赖的任务
• 希望在不增加参数量的情况下提升模型能力

可能不适合：

• 计算资源极其受限的场景
• 对延迟要求极高的在线推理

实现建议

class NexusAttention(nn.Module):
    def __init__(self, d_model, num_heads, order=2, share_weights=True):
        super().__init__()
        self.order = order
        self.share_weights = share_weights
        
        # 共享或独立的投影矩阵
        if share_weights:
            self.W_q = self.W_k = self.W_v = nn.Linear(d_model, d_model)
        else:
            self.W_q = nn.Linear(d_model, d_model)
            self.W_k = nn.Linear(d_model, d_model)
            self.W_v = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        # 标准投影
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)
        
        # 递归高阶注意力
        for _ in range(self.order):
            # 内层注意力精炼Q和K
            Q = self._inner_attention(Q, Q, Q) @ self.W_q.weight
            K = self._inner_attention(K, K, K) @ self.W_k.weight
        
        # 外层主注意力
        return self._outer_attention(Q, K, V)

---

11. 总结

Nexus通过递归高阶注意力机制解决了标准Transformer的核心瓶颈：

方面	标准注意力	Nexus
投影方式	线性静态	非线性递归
依赖建模	成对交互	多token依赖
参数量	$O(1)$	$O(1)$（共享后）
表达能力	低秩瓶颈	打破瓶颈
推理密度	1x	~2x（m=2）

Nexus不仅是预训练的全新架构，更可作为现有LLM的”升级套件”，通过微调即可释放更强的推理能力。

---

参考

---

相关词条

• Transformer与注意力机制 — 标准Transformer架构基础
• 稀疏注意力与长度外推 — 另一类注意力优化方向

Footnotes

1. Chen et al., Nexus: Higher-Order Attention Mechanisms in Transformers, ICLR 2025 ↩

2. Bhojanapalli et al., Low-rank bottleneck in multi-head attention models, ICML 2020 ↩