用于大规模并行序列生成的结构化循环混合器

# 用于大规模并行序列生成的结构化循环混合器
来源: https://arxiv.org/html/2605.08696
###### 摘要

在过去的二十年中，语言建模经历了一个从主要处理训练和推理过程中逐个处理token的循环架构到非循环模型的转变，这些模型在训练期间可以并行处理序列元素，从而提高了训练效率和稳定性，但降低了推理吞吐量。在这里，我们介绍了结构化循环混合器（Structured Recurrent Mixer），这是一种允许在训练时的序列并行表示和推理时的循环表示之间进行代数转换的架构，特别之处在于不需要专门的内核或特定于设备的内存管理。实验表明，与其他线性复杂度模型相比，这种双重表示允许更高的训练效率、更大的输入信息容量以及更大的推理吞吐量和并发性。我们假设，对于语言中典型的信息丰富的输入，循环模型不适合扩展序列长度缩放，但由于每个样本的恒定内存，它们在样本（batch）维度上更适合扩展。我们提供了SRM的Mojo/MAX推理实现，与vLLM上的类似强大的Transformer相比，其吞吐量提高了12倍，并发性提高了170倍，这些增加是PyTorch实现的特征，导致计算恒定GSM8k Pass@k增加了30%。最后，我们通过展示SRM是有效的强化学习训练候选者来结束本文。

00脚注文本:作者希望感谢IBM在研究和撰写本文期间提供的支持。本文的代码可以在https://github.com/blbadger/mrm找到## 1引言

对于语言任务来说，模型速度还是准确性更重要？答案通常取决于任务，特别是它是否提供了一种快速检查样本好坏的方法。对于无法快速验证样本的任务，模型准确性通常比速度更重要，但对于许多任务，验证是可行的：例如，可以对代码样本应用单元测试或功能测试，对代理工作流程进行结果评估，通过Lean对数学定理证明生成进行自我验证。如果快速验证是可能的，那么从较小且能力较弱的模型生成许多样本通常比从较大且更准确的模型生成少数样本更有效（以每正确输出花费的最小计算为术语）\(Brown et al., 2024; Chen et al., 2025\)。

大多数功能语言基准并不关注这些术语中的效率，而是关注在一次或多次尝试后正确答案或下一个正确token的可能性，尽管一些较新的基准已经开始纳入任务成本方面的效率（例如Arc-AGI\(Foundation, 2026\)）。从循环式架构（循环神经网络\(Rumelhart et al., 1985\)和LSTMs\(Hochreiter and Schmidhuber, 1997\)）的转变发生在下一个token预测准确性的背景下，以及它们与基于模型每样本准确性而不考虑生成该样本所需计算的应用基准的关系。对于容易验证的输出，每计算应用的准确性指标似乎更为合适，较小的Transformer模型\(Vaswani et al., 2023\)在这种情况下优于较大模型的发现引出了一个问题：那些比Transformer执行更有效地进行推理的架构是否更适合这种每计算准确性的范式？

对于大规模语言建模，通常希望训练因果模型来预测上下文窗口中的每个下一个token并行地通过token移位，同时在推理过程中改变模型行为以去除这种序列并行性以便效率，这样模型一次预测一个token并使用缓存存储来自先前token的信息。序列并行模型，特别是那些具有O(n²)时间和O(n)空间复杂性的模型，训练效率高且稳定，但对于顺序token生成来说效率低下，而循环模型（在本工作中定义为具有O(n)时间和O(1)空间复杂性的模型）对于序列生成来说是有效的，但通常没有简单的方法来高效且稳定地训练，因为它们不是内在地在序列维度上可并行化的。

由于这些互补的缺陷，以及相对于GPU中存在的算术操作的全局设备内存访问的大延迟，有许多尝试将线性复杂度架构（如Mamba SSMs for Granite 4\(Mishra et al., 2024\); ibm\)和Nemotron 3\(Blakeman and et al., 2025\)，Gated Delta Nets for Qwen 3.5\(Qwen Team, 2026\)和Kimi Linear\(Team et al., 2025\)）混合到Transformer主干中。然而，混合引入了一个重大挑战：虽然这些模型获得了较小可变大小缓存的优势，但它们在缓存扩展方面表现出与Transformer相同的缺点，只是常量值较低。

这些观察结果激励我们引入一种新架构——结构化循环混合器（SRM），我们发现它表现出二次复杂度模型的一些理想的信息保留和训练效率特性，同时改进了当前循环模型的有效推理计算。从信息角度看，我们表明循环模型不适合无限序列长度扩展，但更适合应用于批处理维度扩展，并显示SRM相对于其他测试模型显示出大的吞吐量和并发性增加，无需任何设备特定内核优化并行化的努力。我们探讨了设计具有高效样本吞吐量特性的架构能否带来超越代表性Transformer的准确性，同时保持训练和推理计算恒定。最后，我们调查使用强化学习将多样本效率转化为单样本或少样本效率。

## 2我们的贡献

本文详细描述了测试两个主要见解的实验。首先是掩码混合器应该更容易线性化（在序列维度上）比Transformer，因为它们的数据无关token混合操作，允许在二次复杂度和循环模型之间建立有效桥梁。我们发现：

设计循环模型token混合操作以捕捉训练更高复杂度模型的质量特征可以恢复这些模型的某些训练效率和信息保留

其次，我们重新评估应用于语言等密集信息输入时的常数空间模型的好处，提出：

由于每个样本的连续恒定大小内存，循环语言模型更适合于批处理（样本数）维度而不是序列长度扩展

为此，SRM被证明即使在其他循环模型中也具有无与伦比的吞吐量特性，这是任何样本可能快速测试的任务中所需的特性。我们介绍了以下内容：

1. 1\.一种从二次复杂度到线性复杂度模型的质量特征映射方法
2. 2\.一种双重架构序列并行（在训练时间）和批处理并行（在推理时间）循环模型
3. 3\.一种SRM的Mojo/MAX推理引擎，其吞吐量超过PyTorch实现的7倍
4. 4\.一种重采样方法，用于对大批次应用GRPO而不会降低探索性

## 3相关工作

循环神经网络是最早应用于可变长度序列生成问题的模型之一\(Rumelhart et al., 1985\)。长短期记忆模型是一种类似于循环的模型，引入了多个隐藏状态和相关的记忆流，旨在稳定训练并更好地模拟长序列\(Hochreiter and Schmidhuber, 1997\)。Transformer相比LSTM更高的训练效率，无论是每参数损失还是每输入上下文长度损失，部分解释了为什么这些架构在大多数语言建模任务中被Transformer取代\(Kaplan et al., 2020\)。

试图直接通过将Transformer转换为类似循环架构的经历了一些有限的成功：用线性注意力替换点积注意力的直接交换\(Katharopoulos et al., 2020; Shen et al., 2018\)导致因果建模的低效训练和稳定性\(Poli et al., 2023\)，尽管通过重构线性注意力操作取得了一些进展。沿着这条路线，RWKV架构\(Peng et al., 2023\)被引入，这些模型保留键值向量但执行需要自定义GPU内核进行高效计算的内存映射操作。

最广泛研究的允许在序列并行和循环表示之间转换的模型架构存在于状态空间框架中，其中状态矢量作为循环隐藏层并通过线性操作更新。从S2/S3\(Fu et al., 2023\)模型细化而来，Mamba\(Gu and Dao, 2024\)和Mamba 2\(Dao and Gu, 2024\)SSM架构在token级别 incorporating选择性并具备复杂的内存管理，允许高效的基于平行扫描的状态更新从而在序列并行和循环形式之间转换。\Dao and Gu, 2024进一步调查了SSM作为结构化矩阵模型子集的性质，通过半分离矩阵表示，其中token混合权重矩阵可以表述为低秩子矩阵。

## 4方法

结构化循环混合器的结构使得能够将使用矩阵乘法混合token信息的序列并行表示转换为必须读取和写入恒定大小内存并对每个生成的token使用固定数量操作的循环表示。训练过程通过将token混合矩阵沿行或列扩展权重向量，应用可选的可训练衰减率，然后继续移位下一个token预测来进行。

### 4.1掩码混合器参数化

SRM基于掩码混合器，这是一种类似Transformer的架构，其中注意力被替换为掩码MLP（矩阵乘法）\(Badger, 2025\)。这些模型对序列维度应用掩码矩阵乘法（或等效于步长为1的掩码1-D卷积），可以表达为Y=XM+B其中M∈Rⁿˣⁿ且X,Y,B∈Rᵈˣⁿ。值得注意的是，M的值对于训练和推理都是三角掩码的，并且不是数据依赖的，而是由模型中的可训练值参数化的。这些架构表现出与Transformer相似的时间和空间复杂度，但在token混合操作中形成较少的激活（dd每层），导致在固定参数数量下具有更高的吞吐量，这通过其在每参数基础上较低的训练效率得到平衡。

### 4.2训练掩码混合器权重：质量特征表示

初步研究（详见表S2）表明，将token混合矩阵限制为循环形式（行重复矩阵）会导致相对低效的因果语言模型训练。这促使对训练过的掩码混合器token混合层的结构进行分析，揭示了三个显著的质量特征：包含近似相同值的列的矩阵，包含近似相同值的行的矩阵，以及随着远离主对角线距离增加的权重衰减（或这三种特征的组合），如图1所示。正如我们在下一节中所探讨的那样，构建一个循环SRM来表示这些质量特征会导致 substantially 更高的训练效率。有趣的是，并不是所有训练过的混合器权重矩阵的特征显然都是必要的：大多数层显示低于接近恒定的主对角线权重相对于非对角线权重（图1, S1），但向SRM引入独立的主对角线权重并不导致大模型的显著增加的训练效率（表S5）。

见图1：因果掩码混合器token混合矩阵权重，(a) 整个矩阵图上的标注特征，(b) 主对角线附近的权重。蓝色表示负值，红色正值。

### 4.3架构

为了保留循环表示，SRM的token混合层受到限制，我们现在说明需要什么样的限制。简而言之，我们所探讨的限制相对简单：行必须包含相同元素或列必须包含相同元素，并且可以选择一个常数因子乘到每个对角线上。本质上我们执行矩阵-向量乘法（带掩码），这就是为什么该操作可以转换为固定内存循环操作。

为了说明一个行重复混合操作的例子，考虑输入X∈Rᵈˣⁿ，其中d是隐藏维度和n是tokens，输出Y∈Rᵈˣⁿ，混合器权重矩阵W∈Rⁿˣⁿ，偏置B∈Rᵈˣⁿ，以及可训练衰减常数λ∈(0.9, 1]。我们从输入X,B,W获得Y如方程1所示，为了简单起见假设n=3。方程1可以用其循环表示在方程2中表达，其中循环隐藏状态（缓存）仅仅是求和项。

$$
\begin{pmatrix} | & | & | \\ Y_0 & Y_1 & Y_2 \\ | & | & | \end{pmatrix} = \begin{pmatrix} | & | & | \\ X_0 & X_1 & X_2 \\ | & | & | \end{pmatrix} \begin{pmatrix} \alpha_0 & \lambda\alpha_0 & \lambda^2\alpha_0 \\ b\lambda & \alpha_1 & 0 \\ 0 & \alpha_2 & 0 \end{pmatrix} + \begin{pmatrix} \beta_0 \\ \beta_1 \\ \beta_2 \end{pmatrix}
$$