提出统一框架,让模型既快又省内存地处理长序列。
Sequential-Parallel Duality in Prefix Scannable Models
- 用前缀扫描算法统一建模,支持并行训练和快速推理。
- 新模型每词计算量接近常数,内存仅随长度对数增长。
- 适合需要高效长序列处理的场景,如超长文本生成。
现代神经序列模型需兼顾可并行训练与快速顺序推理。近年来的Gated Linear Attention(GLA)和Mamba等模型实现了这种‘串行-并行双重性’。本文首先定义一类广义模型——状态空间模型,其状态更新可通过经典并行前缀扫描算法结合自定义结合算子实现。随后,通过放宽状态聚合算子限制,允许任意(可能非结合)函数(如softmax注意力),提出更通用的前缀可扫描模型(Prefix-Scannable Models, PSMs)。该框架统一了多种现有架构,包括逐元素RNN(如Mamba)和线性Transformer(如GLA、Mamba2、mLSTM),同时引入新模型,使用softmax类算子实现每令牌平均计算量为O(1),序列长度为N时内存消耗为O(log N)。我们在小规模语言建模和典型合成任务(如状态追踪、关联回忆)上进行实验,结果表明PSMs在保持Transformer级表达力的同时,推理效率媲美状态空间模型,部分情况下更优的长度外推能力。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。