前缀可扫描模型中的序列-并行二元性
机器学习
2026-03-12 v2
摘要
现代神经序列模型需满足并行训练和快速顺序推理的双重要求。近期发展催生了门控线性注意力（Gated Linear Attention, GLA）和 Mamba 等模型，实现了此类“序列-并行二元性”。这自然引出一个问题：我们能否刻画支持近常数时间并行评估和线性时间、常数空间顺序推理的神经序列模型的完整类别？我们首先描述了一类广泛模型——状态空间模型（state space models），其状态更新可通过带自定义关联聚合算子的经典并行前缀扫描算法实现。随后，我们通过放宽状态聚合算子，使其允许任意（可能非关联）函数（如 softmax 注意力），定义更一般的类——前缀可扫描模型（Prefix-Scannable Models, PSMs）。这种泛化统一了许多现有架构，包括元素级 RNN（如 Mamba）和线性变换器（如 GLA、Mamba2、mLSTM），同时引入新模型，采用 softmax 类算子，可实现每 token amortized O(1) 计算和 log(N) 内存（序列长度为 N）。我们在示意性小规模语言建模和标准合成任务上对此类模型进行经验评估，包括状态跟踪和关联记忆。经验结果表明，PSM 保留了基于 transformer 的架构表达能力，同时匹配状态空间模型的推理效率——在某些情况下，其长度广义性优于两者。
引用
@article{arxiv.2506.10918,
title = {Sequential-Parallel Duality in Prefix Scannable Models},
author = {Morris Yau and Sharut Gupta and Valerie Engelmayer and Kazuki Irie and Stefanie Jegelka and Jacob Andreas},
journal= {arXiv preprint arXiv:2506.10918},
year = {2026}
}