Sequential-Parallel Duality in Prefix Scannable Models

📄 arXiv: 2506.10918v1 📥 PDF

作者: Morris Yau, Sharut Gupta, Valerie Engelmayer, Kazuki Irie, Stefanie Jegelka, Jacob Andreas

分类: cs.LG

发布日期: 2025-06-12

💡 一句话要点

提出前缀可扫描模型以实现高效序列推理

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

关键词: 神经序列模型 前缀可扫描模型 并行推理 状态空间模型 语言建模 深度学习

📋 核心要点

  1. 现有神经序列模型在并行训练与序列推理效率之间存在权衡，难以同时满足两者的需求。
  2. 本文提出前缀可扫描模型（PSMs），通过放宽聚合操作符，支持任意函数以实现高效推理。
  3. 实验结果表明，PSMs在语言建模和合成任务中表现出色，推理效率与状态空间模型相当，且在长度泛化上优于现有模型。

📝 摘要（中文）

现代神经序列模型旨在实现可并行训练和快速序列推理的双重目标。本文提出了一类称为前缀可扫描模型（PSMs）的新模型，通过放宽状态聚合操作符，允许任意函数（如softmax注意力），从而统一了多种现有架构。我们在小规模语言建模和合成任务上对这些模型进行了实证评估，发现PSMs在保持表达能力的同时，推理效率与状态空间模型相匹配，甚至在某些情况下表现出更好的长度泛化能力。

🔬 方法详解

问题定义：本文旨在解决现有神经序列模型在并行训练与快速推理之间的矛盾，现有方法在这两方面的表现往往无法兼顾，导致效率低下。

核心思路：论文提出前缀可扫描模型（PSMs），通过定义更广泛的状态聚合操作符，允许使用任意函数（如softmax注意力），从而实现高效的并行评估与线性时间的序列推理。

技术框架：PSMs的整体架构包括状态空间模型的基础，结合经典的并行前缀扫描算法，利用自定义的聚合操作符进行状态更新。模型的设计允许灵活的聚合方式，支持多种现有架构的统一。

关键创新：最重要的创新在于放宽了状态聚合操作符的限制，使得PSMs能够支持更复杂的聚合函数，进而统一了多种现有模型（如Mamba和GLA），并引入了新的模型结构。

关键设计：在模型设计中，关键参数设置包括聚合操作符的选择，损失函数的设计，以及网络结构的灵活性，以确保模型在不同任务中的适应性和效率。具体的实现细节包括O(1)的每个token计算和log(N)的内存使用。

📊 实验亮点

实验结果显示，PSMs在小规模语言建模任务中表现出色，推理效率与状态空间模型相当，且在长度泛化能力上优于传统的变换器架构。具体而言，PSMs在某些任务中实现了O(1)的每token计算，显著提升了模型的实用性和效率。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、时间序列分析和机器人控制等。通过提高序列推理的效率，PSMs可以在实时系统中发挥重要作用，推动智能系统的快速响应能力和处理复杂任务的能力。未来，PSMs有望在多模态学习和大规模数据处理等领域展现更大的价值。

📄 摘要（原文）

Modern neural sequence models are designed to meet the dual mandate of parallelizable training and fast sequential inference. Recent developments have given rise to various models, such as Gated Linear Attention (GLA) and Mamba, that achieve such ``sequential-parallel duality.'' This raises a natural question: can we characterize the full class of neural sequence models that support near-constant-time parallel evaluation and linear-time, constant-space sequential inference? We begin by describing a broad class of such models -- state space models -- as those whose state updates can be computed using the classic parallel prefix scan algorithm with a custom associative aggregation operator. We then define a more general class, Prefix-Scannable Models (PSMs), by relaxing the state aggregation operator to allow arbitrary (potentially non-associative) functions such as softmax attention. This generalization unifies many existing architectures, including element-wise RNNs (e.g., Mamba) and linear transformers (e.g., GLA, Mamba2, mLSTM), while also introducing new models with softmax-like operators that achieve O(1) amortized compute per token and log(N) memory for sequence length N. We empirically evaluate such models on illustrative small-scale language modeling and canonical synthetic tasks, including state tracking and associative recall. Empirically, we find that PSMs retain the expressivity of transformer-based architectures while matching the inference efficiency of state space models -- in some cases exhibiting better length generalization than either.