🤖 machine learning

Structured Recurrent Mixers for Massively Parallelized Sequence Generation

本文介绍了结构化循环混合器（SRM），这是一种新颖的架构，能够实现并行训练与循环推理之间的代数转换，从而在训练效率、信息容量和推理吞吐量方面优于现有的线性复杂度模型，同时在标准基准测试和强化学习任务中展现出显著的性能提升。

原作者： Benjamin L. Badger

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Benjamin L. Badger

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是用简单语言和创意类比对这篇论文的解读。

核心难题：“流水线”与“图书馆”的对立

想象你正在逐字撰写一个故事。

传统模式（循环模型）： 这就像一位作家坐在书桌前。他写下一个词，记住它，再写下一个，如此循环。他在写作过程中记忆故事非常高效（内存占用低），但他一次只能写一个词。如果你需要写 100 个故事，就必须等这位作家写完一个，才能开始下一个。
现代标准（Transformer）： 这就像一座巨大的图书馆，你可以一次性查阅故事中的每一个词。这对训练非常有利，因为你可以并行阅读整本书。然而，当需要撰写故事（推理）时，图书馆就会拥堵。为了写出下一个词，图书馆必须重新扫描你迄今为止写下的整本书以寻找上下文。随着故事变长，图书馆变得越来越慢，并且需要巨大的空间（内存）来容纳所有这些书籍。

两难困境： 我们既想要图书馆在训练时的速度，又想要单人在生成文本时的效率。

解决方案：“结构化循环混合器”（SRM）

作者提出了一种名为结构化循环混合器（SRM）的新架构。你可以将 SRM 想象成一只变色龙或一位变形金刚（指超级英雄意义上的变形，而非人工智能意义上的）：它能根据正在执行的任务改变形态。

训练期间（图书馆模式）： 当模型学习时，它表现得像那座巨大的图书馆。它一次性审视整个词序列。这使得学习既快速又稳定。
推理期间（作家模式）： 当模型实际生成文本时，它会瞬间切换为那位单人作家。它不需要重新扫描整本书；它只需保留一个微小的、固定大小的“笔记本”（缓存），记录刚刚写下的内容。这使其速度极快，并允许同时处理许多故事。

其中的魔法在于，SRM 背后的数学原理允许它在代数层面上在这两种模式之间切换。这就像拥有一张蓝图，上面写着：“如果我们在建造，就使用这些砖块；如果我们住在里面，就使用这些墙壁”，而无需拆毁建筑并重建它。

为何重要：“批次”与“长度”的权衡

这篇论文对如何扩展人工智能提出了一个关键观察：

扩展长度（故事）： 论文认为，试图让这些“单人作家”模型阅读无限长的书籍是一个糟糕的主意。最终，作家的记忆（笔记本）会变得太满，开始遗忘细节。这就像试图在一张纸的笔记本里记住一本 1000 页的小说；你会丢失信息。
扩展批次（人群）： 然而，这些模型在同时处理许多不同故事方面极其出色。因为它们不需要为每个故事配备巨大的图书馆，你可以让 100 位作家并行工作，处理 100 个不同的故事，而不会互相绊倒。

类比： 想象一家咖啡店。

Transformer 就像拥有一台巨型意式浓缩咖啡机的店铺，冲泡一杯拿铁需要 10 分钟，但如果你有一个巨大的柜台，它可以一次冲泡 100 杯。随着订单变得复杂，机器会变慢。
SRM 则像是一家拥有 100 台简单、快速的手动冲泡机的店铺。每位冲泡师都能快速制作一杯咖啡，并在脑海中记住配方。你无法让单台冲泡师完成 100 杯的订单（内存太多），但你可以通过使用 1000 台不同的冲泡机，同时为 1000 位顾客提供服务。

结果：速度与容量

这篇论文测试了这种新架构，发现了一些令人印象深刻的数字：

速度： 在标准硬件上，SRM 生成文本的速度比标准 Transformer快 12 倍。
并发能力： 它能处理的并发请求（用户）数量是 Transformer 的170 倍。
准确性： 尽管速度如此之快且处理了如此多的请求，它并未丧失智能。在数学测试（GSM8k）中，在给定相同算力的情况下，它实际上比 Transformer 多解决了约**30%**的问题。

“强化学习”的转折

这篇论文还探讨了这如何有助于强化学习（通过试错来教导人工智能）。

想象你在教一只狗捡球。

传统方法： 你让狗出去一次。如果它失败了，你再试一次。
SRM 方法： 由于 SRM 速度极快，你可以同时派出50 只狗。你检查哪些狗拿到了球，只奖励获胜者。

论文发现，通过一次性生成大量样本并使用一种特殊的“重采样”技巧（确保你有足够多的“好”样本供学习），SRM 的学习效果远优于那些只尝试几次的模型。

总结

结构化循环混合器（SRM） 是一种新的 AI 设计，它集两者之长：

它像现代并行计算机一样高效地学习。
它像简单、低内存占用的作家一样高效地生成文本。
它专为同时处理大规模人群的用户而构建，随着 AI 从“给出一个宏大答案”转向“提供许多快速答案”，这一点变得越来越重要。

论文得出结论：虽然我们经常试图通过让 AI 阅读更长的书籍来使其变得更聪明，但我们反而应该专注于使其在同时处理许多不同任务时变得更快，而 SRM 正是完成这项工作的完美工具。

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

问题陈述

语言建模领域已从循环架构（RNN、LSTM）转向非循环、序列并行模型（Transformer），以实现训练效率和稳定性。然而，这一转变在推理过程中带来了显著成本：Transformer 需要 $O(n)$ 的内存和顺序令牌生成，导致与循环模型相比，推理吞吐量和并发能力较低。相反，虽然循环模型提供每个样本 $O(1)$ 的内存占用和高效的顺序生成，但它们在历史上一直难以实现训练效率和稳定性，并且由于信息压缩限制，往往无法有效地扩展到长上下文长度。

混合方法（例如 Mamba、门控 Delta 网络）试图弥合这一差距，但通常依赖专用设备内核，并且仍然面临缓存扩展问题。此外，对于输出可以快速验证的任务（例如代码生成、数学证明），“每计算单位精度”范式表明，从更小、更快的模型中生成大量样本，可能比从更大、更准确的模型中生成少量样本更高效。本文认为，循环模型不适合为信息丰富的输入扩展序列长度，但由于其每个样本具有恒定的内存占用，因此非常适合在批量（样本）维度上进行扩展。

方法论：结构化循环混合器（SRMs）

作者引入了结构化循环混合器（SRM），这是一种旨在在不依赖专用内核或特定设备内存管理的情况下，代数地将序列并行表示（用于训练）与循环表示（用于推理）相互转换的架构。

核心架构

SRM 基于掩码混合器（Masked Mixer），这是一种类似 Transformer 的架构，其中注意力机制被掩码多层感知机（MLP）所取代。其关键创新在于令牌混合矩阵的参数化：

定性特征映射：对训练后的掩码混合器的分析表明，高效的令牌混合矩阵表现出特定的定性特征：具有几乎相同值的列、具有几乎相同值的行，以及权重衰减随与主对角线距离增加而增加。
结构约束：为了实现向循环形式的代数转换，SRM 将令牌混合矩阵限制为行重复或列重复结构。
- 行重复：行包含相同的元素。
- 列重复：列包含相同的元素。
- 衰减：应用可训练的衰减常数 $\lambda$ 到对角线或 off-diagonal 元素，以管理信息保留。
代数转换：这些约束允许将矩阵乘法 $Y = XM + B $重写为固定内存的循环操作（论文中的公式 2 和 3）。循环隐藏状态（缓存）变为一个简单的求和项，无论序列长度如何，仅需恒定内存（$ O(d)$）。
混合头：该架构采用“混合”头，其中一半的头使用行重复矩阵，另一半使用列重复矩阵，并通过线性投影进行组合。这种配置被发现在训练效率和信息保留之间提供了最佳平衡。

训练与推理

训练：SRM 使用与 Transformer 类似的序列并行（令牌移位）进行训练。训练时的二次复杂度实际上被内存传输开销所掩盖，使其在标准上下文长度下的吞吐量与 Transformer 相当。
推理：模型切换到循环表示，其中令牌逐个生成，并更新固定大小的缓存。这消除了对大型 KV 缓存的需求，大幅降低了内存带宽要求。

主要贡献

定性特征映射：一种通过强制行/列重复结构和衰减，将二次复杂度模型的令牌混合操作映射到线性复杂度循环模型的方法。
双架构设计：一种在训练期间支持序列并行、在推理期间支持批量并行的架构，允许在不改变模型参数或缓存结构的情况下，在两者之间进行代数转换。
优化的推理引擎：使用 Mojo/MAX 框架实现 SRM，展示了相比标准 PyTorch 实现的显著性能提升。
用于强化学习的平衡重采样：一种用于组相对策略优化（GRPO）的重采样方法，通过在训练批次中确保“好”样本和“坏”样本的混合，以在大批量强化学习中保持探索性。

实验结果

训练效率

按参数量计算，SRM 的训练效率略低于 $O(n^2)$ Transformer 和掩码混合器，但在各种数据集（FineWeb-edu、FineMath）上比 Mamba 模型更高效。
与在旧硬件（例如 V100）上吞吐量显著下降的 Mamba 模型不同，SRM 在不同的 GPU 架构上表现一致。

基准测试性能

准确性：在计算等效基准测试（IFEval、SQuAD、GSM8k 等）中，Transformer 通常位居榜首。SRM 在大多数任务上优于 Mamba 模型，但在单样本准确性上通常落后于 Transformer。
信息保留：SRM 表现出比 Mamba 模型更高的信息保留能力和容量，与二次复杂度模型相当。然而，像所有循环模型一样，当上下文长度超过其压缩限制时，它们会面临信息丢失。

推理吞吐量与并发

吞吐量：由于减少了内存传输，SRM 展现出比 Transformer 巨大的吞吐量提升。
- 在 PyTorch (V100) 上：根据上下文长度不同，SRM 的吞吐量比 Transformer 高出 9.66 倍至 43.29 倍。
- 在 Mojo/MAX (H100) 上：SRM 的吞吐量是通过 vLLM 进行推理的 Transformer 的 12 倍。
并发：SRM 支持显著更大的批量大小。
- 在 Mojo/MAX (H100) 上：SRM 的并发能力是 Transformer 的 170 倍。
- 与 Mamba 相比：SRM 显示出比同等性能的 Mamba 模型高 >500 倍 的并发能力和约 7 倍 的吞吐量。

强化学习与可验证输出

Pass@k 扩展：对于 GSM8k，SRM 遵循与 Transformer 相似的 Pass@k 幂律。
计算效率：由于吞吐量提高了 10-30 倍，SRM 可以在相同的计算预算下生成更多样本。论文报告，在计算预算恒定的情况下，SRM 的 GSM8k Pass@k 比 Transformer 提高了 30%。
RL 训练：使用平衡重采样，SRM 在大批量（50 个样本）上使用 GRPO 进行训练，其单样本准确性和探索性优于在小批量（5 个样本）上训练的 Transformer。

意义与主张

本文结论认为，虽然循环模型不适合为信息密集的输入扩展序列长度，但它们非常适合在批量维度上进行扩展。SRM 架构成功弥合了二次模型的训练效率与循环模型的推理效率之间的差距。

作者声称，SRM 为“每计算单位精度”范式提供了一条可行路径，特别是对于具有快速验证的任务。通过利用循环模型的恒定内存占用，SRM 比 Transformer 能更有效地利用现代硬件的“内存墙”限制，在无需专用内核的情况下实现吞吐量和并发能力的数量级提升。这项工作表明，对于未来的多线程硬件趋势，像 SRM 这样最大化算术运算相对于内存访问次数的架构将变得越来越主导。