Nexusformer: Nonlinear Attention Expansion for Stable and Inheritable Transformer Scaling

📄 arXiv: 2604.19147v1 📥 PDF

作者: Weijie Zhao, Mingquan Liu, Bolun Wang, Simo Wu, Nuobei Xie, Rui-Jie Zhu, Peng Zhou

分类: cs.LG, cs.AI

发布日期: 2026-04-21


💡 一句话要点

Nexusformer:通过非线性注意力扩展实现Transformer的稳定和可继承缩放

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: Transformer 模型缩放 非线性注意力 增量学习 零初始化 语言建模 推理 Nexus-Rank层

📋 核心要点

  1. 现有Transformer模型扩展受限于注意力机制的线性投影,导致表达能力和增量学习能力不足。
  2. Nexusformer通过引入Nexus-Rank层,利用非线性映射克服线性约束,实现无损结构增长。
  3. 实验表明,Nexusformer在渐进缩放时,能以更少的计算资源匹配Tokenformer的性能,并具有稳定的收敛轨迹。

📝 摘要(中文)

Transformer的扩展通常需要从头开始训练更大的模型,因为标准架构难以在不丢弃已学习表示的情况下进行扩展。本文确定了注意力机制中线性投影的主要瓶颈,它严格地将特征提取限制在固定维度的子空间中,从而限制了表达能力和增量容量。为了解决这个问题,本文提出了Nexusformer,它用Nexus-Rank层替换了线性Q/K/V投影,Nexus-Rank层是由在逐渐更高维度空间中的双重激活驱动的三阶段非线性映射。这种设计克服了线性约束,并实现了无损结构增长:可以通过零初始化的块沿两个轴注入新的容量,从而保留预训练的知识。在语言建模和推理基准上的实验表明,Nexusformer在使用渐进缩放(240M到440M)期间,使用比Tokenformer少41.5%的训练计算量即可匹配其困惑度。此外,对增长动态的分析表明,零初始化会诱导稳定的收敛轨迹,从而可以推导出几何缩放定律,该定律可以准确地预测跨扩展尺度的性能。

🔬 方法详解

问题定义:Transformer模型在扩展时,通常需要从头开始训练,无法有效利用已有的预训练知识。这是因为标准Transformer的注意力机制中的线性投影将特征提取限制在固定维度的子空间中,限制了模型的表达能力和增量学习能力。因此,如何实现Transformer模型的稳定和可继承的缩放是一个关键问题。

核心思路:本文的核心思路是通过引入非线性映射来克服注意力机制中线性投影的限制。具体来说,作者提出了Nexus-Rank层,它使用三阶段的非线性映射,在逐渐更高维度的空间中进行特征提取,从而打破了线性约束,提高了模型的表达能力。此外,通过零初始化新增的Nexus-Rank层,可以保证在扩展模型时,保留原有的预训练知识。

技术框架:Nexusformer的核心是使用Nexus-Rank层替换了标准Transformer中的线性Q/K/V投影。Nexus-Rank层包含三个阶段的非线性映射,每个阶段都使用双重激活函数。整体架构与标准的Transformer类似,可以堆叠多个Nexusformer层来构建更深的模型。在模型扩展时,可以通过零初始化新增的Nexus-Rank层来实现无损的结构增长。

关键创新:Nexusformer的关键创新在于Nexus-Rank层,它通过非线性映射克服了线性投影的限制,提高了模型的表达能力和增量学习能力。此外,通过零初始化新增的Nexus-Rank层,可以保证在扩展模型时,保留原有的预训练知识,从而实现了Transformer模型的稳定和可继承的缩放。

关键设计:Nexus-Rank层使用三阶段的非线性映射,每个阶段都使用双重激活函数,例如ReLU和Tanh。零初始化用于初始化新增的Nexus-Rank层,以保证在扩展模型时,保留原有的预训练知识。作者还发现,零初始化可以诱导稳定的收敛轨迹,从而可以推导出几何缩放定律,该定律可以准确地预测跨扩展尺度的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,Nexusformer在语言建模和推理基准上表现出色。在渐进缩放(240M到440M)期间,Nexusformer使用比Tokenformer少41.5%的训练计算量即可匹配其困惑度。此外,作者还通过实验验证了零初始化可以诱导稳定的收敛轨迹,并推导出了几何缩放定律,该定律可以准确地预测跨扩展尺度的性能。

🎯 应用场景

Nexusformer的潜在应用领域包括自然语言处理、计算机视觉等。它可以用于构建更大规模、更强大的Transformer模型,从而提高各种任务的性能。此外,Nexusformer的可继承缩放特性使得它可以更容易地将模型部署到资源受限的设备上,例如移动设备和嵌入式系统。未来,该技术有望推动人工智能在各个领域的应用。

📄 摘要(原文)

Scaling Transformers typically necessitates training larger models from scratch, as standard architectures struggle to expand without discarding learned representations. We identify the primary bottleneck in the attention mechanism's linear projections, which strictly confine feature extraction to fixed-dimensional subspaces, limiting both expressivity and incremental capacity. To address this, we introduce Nexusformer, which replaces linear $Q/K/V$ projections with a Nexus-Rank layer, a three-stage nonlinear mapping driven by dual activations in progressively higher dimensional spaces. This design overcomes the linearity constraint and enables lossless structured growth: new capacity can be injected along two axes via zero-initialized blocks that preserve pretrained knowledge. Experiments on language modeling and reasoning benchmarks demonstrate that Nexusformer matches Tokenformer's perplexity using up to 41.5\% less training compute during progressive scaling (240M to 440M). Furthermore, our analysis of growth dynamics reveals that zero initialization induces a stable convergence trajectory, allowing us to derive a geometric scaling law that accurately predicts performance across expansion scales.