提出新型稀疏状态空间模型,兼顾高效与强表达能力。
Flash PD-SSM: Memory-Optimized Structured Sparse State-Space Models

- 动态选择结构化稀疏矩阵,实现高效计算下的高表达力
- 在17,000+长度序列上达到当前最优准确率
- 可直接替换主流大模型中的SSM模块,提升速度与内存效率
状态空间模型(SSMs)在效率与表达力之间存在根本权衡,主要由转移矩阵的结构决定。非结构化转移矩阵虽具备最强表达力(可建模有限状态自动机),但计算与内存开销巨大;而多数结构化形式虽高效,却受限于表达能力。本文基于结构化稀疏SSM的最新进展,提出Flash PD-SSM:通过维护一组可训练的结构化稀疏矩阵,并在每个时间步离散选择其一,实现接近非结构化矩阵的有限状态自动机表达力,同时保持大规模训练所需的高效性。实验表明,该模型在合成机制与状态追踪任务中验证了理论表达力;在长度超17,000的多变量时间序列任务中,达到当前最优准确率;此外,作为混合大语言模型的即插即用替代方案,在自然语言状态追踪和通用语言建模中均表现更优,相较现有前沿语言模型中的SSM,具有更高吞吐量与更低内存占用。
原文摘要 · Abstract (English)
State-space models (SSMs) face a fundamental trade-off between efficiency and expressivity that is mainly dictated by the structure of the model's transition matrix. Unstructured transition matrices enable maximal expressivity, as measured by their ability to model finite-state automaton (FSA) transitions, but come at a prohibitively high compute and memory cost. In contrast, most structured transition matrix forms are highly efficient both in runtime and memory consumption, but suffer from limited expressivity. Building on recent work on structured sparse SSMs, we propose Flash PD-SSM, a novel SSM that achieves comparable throughput to widely-used structured SSMs with significantly better expressivity guarantees. Flash PD-SSM maintains a trainable set of structured sparse matrices, a single one of which is discretely selected at each time-step, enabling FSA expressiveness at the level of unstructured matrices while maintaining the efficiency required for training models at scale. First, we validate Flash PD-SSM against a suite of alternative models on synthetic mechanistic and state-tracking tasks, finding that its theoretical expressivity is achieved in practice. Second, on multivariate time-series tasks involving sequences of length over 17,000, we find that Flash PD-SSM defines a new state-of-the-art (SoTA) accuracy among competing SSM methods. Finally, we demonstrate that Flash PD-SSM is an effective drop-in replacement for hybrid LLMs, yielding improvements both in natural language state-tracking and in common language modeling scenarios. The model exhibits increased throughput and decreased memory consumption compared to SSMs widely used in frontier language models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。