Sequential-Parallel Duality in Prefix Scannable Models

现代神经序列模型旨在同时满足并行化训练和快速顺序推理的双重需求。近期的研究发展催生了多种实现“顺序-并行二元性”的模型，例如门控线性注意力（GLA）和 Mamba。这自然引发了一个问题：我们能否刻画出一类完整的神经序列模型，它们支持接近常数时间的并行评估以及线性时间、常数空间的顺序推理？我们首先描述了一类广泛的此类模型——状态空间模型，其状态更新可以通过经典的并行前缀扫描算法计算，使用自定义的结合性聚合操作符。接着，我们通过放松状态聚合操作符的限制，允许任意（可能非结合性的）函数（如 softmax 注意力），定义了一个更通用的模型类别——可前缀扫描模型（PSMs）。这一推广统一了许多现有的架构，包括逐元素 RNN（例如 Mamba）和线性变换器（例如 GLA、Mamba2、mLSTM），同时也引入了新的模型，这些模型使用类似 softmax 的操作符，能够实现每标记摊销 O(1) 的计算复杂度和对长度为 N 的序列达到 log(N) 的内存需求。我们在说明性的小型语言建模任务和经典合成任务（包括状态跟踪和关联回忆）上对这些模型进行了实证评估。实验结果表明，PSMs 保留了基于变换器架构的表达能力，同时匹配了状态空间模型的推理效率——在某些情况下，甚至表现出比两者更好的长度泛化性能。

该论文试图解决如何设计一类神经序列模型的问题，这类模型能够在训练时支持近常数时间的并行计算，并在推理时实现线性时间和常数空间的顺序处理。这是一个相对较新的问题，尤其是在统一表达性和效率方面。

论文提出了一种称为Prefix-Scannable Models (PSMs) 的新模型类别，通过放松状态聚合算子的限制（允许非结合函数如softmax注意力），将现有的架构（如Mamba和GLA）统一起来，同时引入了具有O(1)摊销计算复杂度和log(N)内存的新模型。相比现有研究，这篇论文的关键创新在于定义了一个更广泛的模型类别，并展示了如何在保持Transformer表达能力的同时提高推理效率。

论文通过小规模语言建模和合成任务（如状态跟踪和关联回忆）验证了PSMs的有效性。实验表明，PSMs在某些情况下比状态空间模型或Transformer具有更好的长度泛化能力。此外，论文还探讨了PSMs在理论上的优势，并为未来的研究提供了方向。虽然未明确提及代码开源，但提出的理论框架可能为后续工作提供重要参考。

最近的相关研究包括Gated Linear Attention (GLA) 和Mamba等模型，这些模型均致力于实现序列-并行二元性。其他相关研究还包括线性Transformer变体（如Performer、Linformer）以及优化RNN结构的工作（如mLSTM）。一些相关标题包括《Gated Linear Transformers》、《Mamba: Fast and Efficient Sequence Models》和《Linear Transformers Are Secretly Optimal》。