关于亚二次架构:从应用到底层原理 On Subquadratic Architectures: From Applications to Principles
系统比较xLSTM、Mamba-2和Gated DeltaNet在复杂依赖任务上的性能
前置知识
亚二次架构
亚二次架构是指计算复杂度低于$O(T^2)$的序列建模架构,其中$T$是序列长度。Transformer的标准自注意力机制需要对每个token与其他所有token进行交互,导致$O(T^2)$的二次复杂度。亚二次架构通过引入状态空间模型、线性注意力或递归机制,将复杂度降低到$O(T)$或$O(T\log T)$,从而能够处理更长的序列。常见方法包括Mamba系列、xLSTM、Gated DeltaNet等,它们都通过引入显式状态存储来避免全注意力计算。
论文的核心是比较不同亚二次架构的性能差异,理解这些架构的复杂度特性和内存动力学是理解论文贡献的基础
线性注意力
线性注意力是标准自注意力的一种变体,通过将softmax替换为其他核函数来避免$O(T^2)$的复杂度。标准注意力可以表示为$\text{softmax}(QK^T)V$,而线性注意力则使用$Q(K^TV)$的等价形式,将复杂度降低到$O(T^2D)$。通过引入显式矩阵状态$C_t$,可以进一步实现递归形式:$h_t = q_tC_t/n_t$,其中$C_t = C_{t-1} + k_t \otimes v_t$。这使得线性注意力在推理时具有线性复杂度,同时保持了对长程依赖的建模能力。
论文分析的xLSTM、Mamba-2和Gated DeltaNet都可以统一到线性注意力的框架下理解,这是论文架构分析的理论基础
状态跟踪与累加
状态跟踪是指模型能够根据输入序列维护和更新有限状态的内部表示,例如跟踪函数调用的返回值或变量的绑定。累加是指模型能够对序列中的元素进行计数或聚合操作,不受训练长度的限制。理论研究表明,现代序列模型很难同时具备这两种能力:Mamba类状态空间模型受到TC0复杂度类的限制,无法解决排列组合等困难的状态跟踪问题;而需要显式计数机制的任务则对累加能力有要求。xLSTM通过矩阵状态的线性注意力和非线性递归层的结合,试图同时提供这两种能力。
论文提出的核心假设是xLSTM在复杂任务上的优势源于其能够结合状态跟踪和累加两种能力,理解这个概念是把握论文技术贡献的关键
长度泛化
长度泛化是指模型在训练长度以外的序列长度上仍能保持良好性能的能力。这是一个衡量序列模型真正理解结构模式而非记忆训练数据特征的重要指标。论文中,模型在长度为128的序列上训练,然后在512(4倍)和2048(16倍)长度上评估。如果一个模型真正理解了任务结构(例如AnBn的平衡括号性质),它应该能够在更长的序列上保持性能。长度泛化失败通常表明模型过度依赖特定长度的模式或启发式规则,而不是掌握了任务的底层规律。
论文通过长度泛化实验验证了架构假设,这是证明xLSTM确实具有更强建模能力的关键证据
研究动机
Transformer架构在序列建模中占据主导地位,但其自注意力机制的$O(T^2)$二次计算成本使其在处理长序列时面临严重挑战。虽然已有多种亚二次架构被提出,包括xLSTM、Mamba-2和Gated DeltaNet等,并且它们已经在实际的混合基础模型中得到应用,例如Samba、Nemotron Nano、Kimi Linear和Olmo Hybrid,但缺乏直接的对比研究来指导架构选择。现有工作主要在标准语言建模和常识推理基准上评估这些架构,发现性能差异很小且难以区分。然而,这些基准可能不够敏感,无法揭示架构之间的本质差异,特别是在具有长程、结构化依赖的任务上。
本文的目标是本文的第一个目标是在代码和时间序列这两个复杂依赖的数据域上进行严格的头对头比较,通过从零开始的预训练和Transformer到亚二次架构的蒸馏两种设置来评估xLSTM、Mamba-2和Gated DeltaNet的性能。第二个目标是解释xLSTM在这些任务上的优势,通过将三种架构统一到同一个框架中,分析它们在内存动力学方面的差异,特别是门控机制、归一化和覆盖机制如何影响累加和有限状态跟踪这两种基础能力。第三个目标是通过控制性的合成长度泛化任务来验证从理论分析中得出的假设。
与已有工作不同的是,与现有工作相比,本文的独特之处在于将评估扩展到标准语言建模之外,选择了代码和时间序列这两个对架构归纳偏差更敏感的复杂域。代码结合了类语言的token分布和形式结构,包括语法、变量绑定和作用域;时间序列需要模型从连续值的历史记录中推断和更新复杂的动态。此外,本文首次提出了将xLSTM、Mamba-2和Gated DeltaNet统一到线性注意力框架下的方法,使得它们在如何写入、遗忘、覆盖和读取状态的层面上直接可比。最后,本文还提供了对累加和状态跟踪假设的实验验证,这在以往的对比研究中是缺失的。
核心方法
本文的方法分为三个层面:首先是在复杂依赖任务上的实证比较,通过代码预训练、代码蒸馏和时间序列基础模型预训练三个实验来建立xLSTM的性能优势模式;其次是架构层面的理论分析,将三种亚二次架构统一到线性注意力的框架中,分析它们的门控机制和内存动力学;最后是假设验证,通过控制性的合成任务测试累加和状态跟踪能力。整体思路是从实证现象到理论解释再到实验验证的完整闭环,不仅展示了xLSTM在什么任务上表现更好,还解释了为什么在这些任务上表现更好。
论文的核心创新点是提出将xLSTM、Mamba-2和Gated DeltaNet统一到线性注意力框架下,并从这个框架中识别出累加和状态跟踪两个基础能力来解释性能差异。统一框架显示,这三种架构都可以用输入门$i_t$、遗忘门$f_t$和矩阵状态$C_t$的递归更新来表示,但它们在门控机制的灵活性和覆盖行为上存在本质区别。xLSTM使用指数输入门,可以像softmax一样灵活地调整之前权重的权重,实现最灵活的内存修正;Mamba-2的输入和遗忘门是绑定的,表达受限;Gated DeltaNet则显式覆盖旧值,有利于检索任务但不利于计数。论文假设xLSTM的优势源于其矩阵状态提供的累加能力和递归层提供的有限状态跟踪能力的结合。
方法步骤详情
代码预训练实验使用lm-engine训练400M参数的层间混合语言模型,其中大部分层使用测试的亚二次序列算子,少量层保持标准自注意力。比较了xLSTM[7:1]、Mamba-2和Gated DeltaNet三种配置,在三个数据配置上训练:Nemotron-CC-Code-v1上的20B和100B token,以及Nemotron-CC-Code-v1与FineWeb-Edu的混合数据上的20B token。评估使用HumanEval的pass@k指标以及HellaSwag、PIQA、ARC-Easy、ARC-Challenge和WinoGrande等推理和常识任务的准确率。代码蒸馏实验使用Qwen3-4B-Instruct作为教师模型,将每个多头注意力块替换为混合块,结合线性注意力算子和滑动窗口注意力。训练分为两个阶段:第一阶段对齐学生和教师的输出,第二阶段最小化交叉熵和稀疏知识散度。在HumanEval、HumanEval+、MBPP和MBPP+上评估pass@1指标。时间序列基础模型预训练使用Auer等人(2025)的协议,保持语料库、patching方案、优化器和预测头不变,只改变序列混合器。在1M、4M、10M、40M和80M五个参数规模上比较Mamba-2、Gated DeltaNet和xLSTM[3:1],在GIFT-Eval上进行零样本评估,报告几何平均的MASE和CRPS。合成任务实验在长度128上训练模型,在128、512和2048上评估,包括AnBn、AnBnCn、Majority等计数任务以及Parity、Modular Arithmetic、S3等状态跟踪任务。
技术新颖性
论文的技术新颖性体现在三个方面:一是首次在代码和时间序列等复杂依赖数据域上进行xLSTM、Mamba-2和Gated DeltaNet的头对头比较,揭示了架构差异在这些任务上更加明显;二是提出了统一框架,将三种架构从原始的状态空间模型和线性注意力标记桥接到共同的形式化,使得它们在内存动力学层面上直接可比;三是基于统一框架提出了累加和有限状态跟踪的假设,并通过合成长度泛化任务提供了实验验证。这些新颖性不仅有助于理解现有亚二次架构的优势和局限,也为未来的架构设计提供了指导原则。
实验结果
在代码预训练实验中,xLSTM[7:1]在所有训练配置和所有pass@k指标上都是最佳模型。在Nemotron-CC-Code-v1的20B token上,xLSTM[7:1]的pass@64为12.13%,比次优的Gated DeltaNet的10.70%提升了1.43个百分点;在100B token上,xLSTM[7:1]的pass@64为15.84%,比Mamba-2的14.36%提升了1.48个百分点;在混合语料库上,xLSTM[7:1]的pass@64为9.50%,比Mamba-2的7.69%提升了1.81个百分点。在推理和常识任务上,xLSTM[7:1]也保持了最佳的聚合得分,但差距较小。在代码蒸馏实验中,xLSTM[1:0]在四个基准中的三个上匹配或超过了默认Gated DeltaNet,平均得分为0.768,而默认Gated DeltaNet为0.755。Gated DeltaNet[-1, 1]在HumanEval和HumanEval+上有所改善,但在MBPP和MBPP+上落后。在数学蒸馏中,xLSTM[1:0]在GSM8K和AIME 2024上领先,而Gated DeltaNet在MATH-500上略强。平均而言,xLSTM[1:0]达到0.645,Gated DeltaNet为0.625。在时间序列基础模型预训练中,xLSTM[3:1]从1M到40M参数在两个指标上都达到了最佳分数。在10M参数时,xLSTM[3:1]的MASE为0.733,CRPS为0.508,而次优的Mamba-2为0.767和0.525。在80M参数时,模型几乎收敛:xLSTM[3:1]和Mamba-2在MASE上持平,而Mamba-2在CRPS上领先0.005。在合成任务实验中,xLSTM[1:1]是唯一在两种任务类型上都能长度泛化的配置,在所有三个状态跟踪任务上达到100%准确率,同时保持有用的计数外推。xLSTM[1:0]在每个计数任务上都长度泛化,在2048长度上AnBn为0.892,AnBnCn为0.932,Majority为0.763,但在每个状态跟踪任务上都失败。混合xLSTM[1:1]在所有三个状态跟踪任务上都是精确的,同时保留有用但较弱的计数外推。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 代码生成 (HumanEval pass@64) | pass@64 (%) | 12.13 (xLSTM[7:1], 20B tokens) | 10.70 (Gated DeltaNet, 20B tokens) | +1.43 百分点 |
| 代码生成 (HumanEval pass@64) | pass@64 (%) | 15.84 (xLSTM[7:1], 100B tokens) | 14.36 (Mamba-2, 100B tokens) | +1.48 百分点 |
| 代码蒸馏 (平均准确率) | 平均准确率 | 0.768 (xLSTM[1:0]) | 0.755 (Gated DeltaNet) | +0.013 |
| 时间序列预测 (MASE, 10M参数) | MASE (越低越好) | 0.733 (xLSTM[3:1]) | 0.767 (Mamba-2) | -0.034 |
| 计数任务 (AnBn, 2048长度) | 准确率 | 0.892 (xLSTM[1:0]) | 0.241 (Mamba-2) | +0.651 |
| 状态跟踪 (Parity, 2048长度) | 准确率 | 1.000 (xLSTM[1:1]) | 0.012 (Mamba-2) | +0.988 |
局限与改进
作者承认的局限性包括:代码语言建模是在400M参数规模上进行的,蒸馏流水线只使用了一个教师模型,只有时间序列实验包含了规模扩展。未来可以扩展到更大的模型规模、更多的教师模型和更广泛的数据域。此外,论文只关注了最近的领先亚二次架构,排除了已经在Beck等人(2024)中比较过的族。更广泛的操作员调查可以进一步细化图景。我自己的观察是,论文的结论可能受到任务选择的限制,虽然在代码和时间序列上xLSTM表现更好,但在更广泛的Web文本预训练上差异可能很小。此外,论文主要关注了架构层面的差异,但没有深入讨论训练动态、优化难度和实现复杂度等因素,这些在实际应用中也很重要。
独立分析的弱点
论文的一个潜在弱点是在代码蒸馏实验中只比较了xLSTM[1:0]和Gated DeltaNet,没有包括Mamba-2,因为它的输入和遗忘门绑定的设计不适合作为即插即用的替代。这限制了蒸馏实验的全面性,特别是考虑到Mamba-2在实际的混合语言模型中被广泛使用。另一个弱点是时间序列实验只在较小的模型规模(1M-80M参数)上进行,虽然这对于时间序列基础模型来说是实用的规模范围,但可能无法反映在语言模型规模(数十亿参数)上的趋势。此外,论文的合成任务设计虽然很好地分离了累加和状态跟踪能力,但这些任务可能与实际的复杂依赖任务有一定差距,它们过于人工和理论化,不一定能直接预测在更复杂任务上的表现。
未来方向
作者提出的未来工作方向包括扩展到更大的模型规模、更多的教师模型和更广泛的数据域。基于论文的成果,可以延伸的方向包括:探索更多亚二次架构在统一框架下的比较,例如RWKV、RetNet、Linear Transformer等;研究在不同架构选择下的混合策略,例如如何在同一个模型中自适应地选择不同的序列算子;探索更复杂的合成任务来测试其他基础能力,例如层次结构建模、符号推理等;研究这些架构在计算效率和内存使用方面的权衡,特别是在不同硬件上的表现;以及探索这些架构在其他复杂依赖任务上的表现,例如图神经网络、程序合成、多模态建模等。
复现评估
论文提供了详细的实验细节,包括超参数、训练设置和评估协议,但代码和模型权重是否开源没有明确说明。训练使用了高端硬件,例如代码预训练在8xH100 GPU上进行,时间序列预训练在4xNVIDIA A100 GPU上进行,合成任务实验在单个NVIDIA H100 GPU上进行,这意味着复现这些实验需要相当大的计算资源。数据集方面,代码预训练使用Nemotron-CC-Code-v1和FineWeb-Edu,代码蒸馏使用Nemotron-Pretraining-Code-v2和Nemotron-Math-v2,时间序列预训练使用大约4750万个时间序列的语料库,这些数据集的可用性可能会影响复现的难度。总体而言,虽然论文提供了足够的细节来理解实验,但完整的复现可能需要大量的计算资源和数据获取。
论文图表
Figure 1展示了两种具有复杂依赖的任务类型。左侧是代码示例,显示了一个Python函数compute_scores,其中包含了调用图和返回语句的依赖关系,代码中的依赖存在于形式结构中:语法树、调用图、变量绑定。右侧是时间序列示例,显示了洛伦兹吸引子的轨迹,未来取决于历史中的未观测状态,时间序列中的依赖存在于部分观测的动态中。这张图通过具体例子说明了为什么代码和时间序列是评估亚二次架构的合适任务,它们都包含长程的、结构化的依赖。
这张图对理解论文的动机至关重要,它形象地展示了论文选择代码和时间序列作为评估任务的原因,这些任务包含的复杂依赖能够更好地揭示不同架构之间的差异。