Flash PD-SSM: Memory-Optimized Structured Sparse State-Space Models¶

会议: NeurIPS 2026
arXiv: 2605.19150
领域: 模型压缩
关键词: 状态空间模型, 结构化稀疏, 自动机模拟, 显存优化, 硬件感知算法

一句话总结¶

针对传统结构化稀疏状态空间模型生成输入依赖转移矩阵时面临的 $\mathcal{O}(LN^2)$ 显存与计算瓶颈，Flash PD-SSM 提出离散字典选择机制配合退火直通估计器，将矩阵生成复杂度骤降至 $\mathcal{O}(LN)$，并通过专用三阶段分块 CUDA 内核达到超过 Mamba2 的训练吞吐量，在合成有限状态机模拟、超长多变量时间序列和混合大语言模型推理中展现出优异的状态追踪精度。

研究背景与动机¶

状态空间模型（State-Space Models, SSMs）因其在长序列建模时具备线性计算复杂度以及自回归推理阶段具有常数级内存开销，已成为替代或混合传统 Transformer 的重要架构方向。然而，SSM 的表达能力与计算效率之间长期存在着根本性的权衡，这一权衡主要由其隐藏状态转移矩阵的数学结构所决定。在主流模型中，为了维持并行前缀扫描（Parallel Scan）的极致高效性，研究者普遍采用标量缩放单位阵或对角转移矩阵。但在时变非自治动态系统中，对角转移矩阵存在严重的理论表达局限：已有形式语言理论证明，实数或复数对角时变 SSM 均无法模拟具有非可解变换半群的确定性有限状态自动机（FSA），这导致其在代码语法校验、多步骤状态追踪以及复杂形式推理中频繁出现能力崩溃。

为了打破对角结构的表达上限，后续研究探索了对角加低秩（DPLR）及 Householder 矩阵乘积参数化（如 DeltaNet 与 DeltaProduct 系列），但在不显著堆叠深度与宽度的前提下，这类模型仍难以达到非结构化密集转移矩阵的 FSA 模拟能力。近期提出的 PD-SSM 开辟了结构化稀疏转移矩阵的新路径，证明了仅需保持每列单激活元素的一热（one-hot）列稀疏结构即可在单层隐状态维度 $N$ 下完全模拟任意 $N$ 状态的 FSA，且稀疏矩阵乘积在结合律下保持结构封闭性。然而，原始 PD-SSM 依赖于对 $K$ 个密集候选字典矩阵进行输入加权凸组合后再执行列硬极大化（column hardmax），这导致在序列长度为 $L$、状态维度为 $N$ 的并行展开过程中，中间状态张量的显存占用膨胀至 $\mathcal{O}(LN^2)$，在实际训练中比参数匹配的 Mamba2 慢 40 倍以上，完全无法投入规模化大模型预训练。

针对这一关键显存瓶颈，本文切入的角度在于重构稀疏矩阵的生成机制：既然理论上任意确定性自动机在每个时间步只需依据输入确定性地选取转移算子，那么与其在每个时间步动态合成连续的凸组合矩阵再施加稀疏化，不如直接在预先完成稀疏化的离散字典矩阵集合中进行硬路由选择。本文的核心 idea 是：将结构化稀疏转移矩阵生成重构为离散查表路由，在预先稀疏化的字典中依据输入直接选取离散算子，通过斜率退火直通估计器实现端到端梯度反传，将生成显存从 $\mathcal{O}(LN^2)$ 降为 $\mathcal{O}(LN)$，并协同定制三阶段分块硬件级 CUDA 内核打破带宽瓶颈。

方法详解¶

整体框架¶

Flash PD-SSM 旨在构建一个既保留非结构化转移矩阵的完整自动机模拟能力，又具备硬件级高吞吐和低显存占用的线性循环层。模型整体采用经典前归一化（Pre-norm）的多头状态空间架构，输入序列通过线性投影映射为隐层表示，经过 Flash PD-SSM 核心循环层完成状态更新，再经由输出投影与门控单元重组输出。Flash PD-SSM 核心循环包含三大协同模块：首先通过预稀疏化离散矩阵字典与输入硬路由机制在时间步 $t$ 生成当前转移算子；随后借助斜率退火直通估计器为离散选择过程构建精确的替代梯度；最后依托三阶段分块并行硬件感知 CUDA 内核将循环展开为局部分块扫描、全局状态传递与分块校正，实现贴近硬件峰值内存带宽的高效执行。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列 u_t"] --> B["预稀疏化离散矩阵字典与输入硬路由机制<br/>计算线性投影 k(u_t) 并离散选取算子 P_t"]
    B --> C["三阶段分块并行硬件感知 CUDA 内核<br/>局部扫描构建块内算子与隐状态"]
    C --> D["全局前缀传递与分块校正<br/>跨块状态累加输出全局隐藏状态 x_t"]
    D --> E["斜率退火直通估计器<br/>反向传播时使用带温度的连续松弛梯度更新字典与路由"]
    E --> F["输出投影与最终隐层表示"]

关键设计¶

1. 预稀疏化离散矩阵字典与输入硬路由机制：消除二次方显存爆炸

原始 PD-SSM 的主要显存瓶颈在于每次计算都需要将输入特征向量 $u_t \in \mathbb{R}^D$ 映射为权重向量，对 $K$ 个 $N \times N$ 的连续稠密矩阵进行线性加权求和，从而产生形状为 $[B, L, N, N]$ 的密集张量，开销达 $\mathcal{O}(LN^2)$。Flash PD-SSM 从离散有限状态机转移的基本代数性质出发，证明了在模型初始化或每个优化步开始前，即可对字典中的 $K$ 个稠密参数矩阵 $M_k \in \mathbb{R}^{N \times N}$ 进行列级硬极大化预处理，直接生成 $K$ 个结构化稀疏索引数组： $$M_k^{\text{sparse}}[j] = \arg\max_{n \in \{0, \dots, N-1\}} (M_k[n, j]) \in \{0, \dots, N-1\}$$ 在序列前向传播的各个时间步 $t$，输入向量 $u_t$ 仅通过线性投影矩阵 $S \in \mathbb{R}^{K \times D}$ 计算选择得分 $k(u_t) = S u_t \in \mathbb{R}^K$，并直接选取最大得分对应的字典索引 $k^*(u_t) = \arg\max_{k \in \{1, \dots, K\}} k(u_t)$。当前步的转移置换矩阵即为 $P_t = M_{k^*(u_t)}^{\text{sparse}}$。由于每个稀疏矩阵在内存中仅需存储为包含 $N$ 个整数的一维索引数组，且采用 int16 格式存储，显存开销从 $\mathcal{O}(LN^2)$ 骤降至 $\mathcal{O}(LN)$，浮点计算开销也从 $\mathcal{O}(LN^2 K)$ 缩减至 $\mathcal{O}(LNK)$。理论上，任何具有 $N$ 个状态和输入字母表大小为 $|\Sigma|=K$ 的确定性 FSA 均可被单层 Flash PD-SSM 精确表达。

2. 斜率退火直通估计器：解决双重离散化反向传播

由于字典稀疏化 $\arg\max$（列内）和字典路由 $\arg\max$（跨矩阵）均为不可导的阶跃离散操作，网络无法直接应用标准反向传播。Flash PD-SSM 在反向传播路径中构建了带有斜率退火（Slope Annealed）机制的双重直通估计器（Straight-Through Estimator, STE）。在前向传播中维持完全离散的整数索引以追求极致速度与显存压缩；在反向传播中，将离散的硬选择操作松弛为带温度参数 $\tau$ 的平滑 Softmax 算子。对于第 $k$ 个字典矩阵 $M_k$，其接收的累积梯度仅聚合在被选中的时间步集合 $\{t \mid k^*(u_t) = k\}$： $$\frac{\partial \ell}{\partial M_k} \approx \sum_{t: k^*(u_t) = k} \left( \frac{\partial \ell}{\partial x_t} (D_t x_{t-1})^T \right) \frac{\partial \operatorname{softmax}_\tau(M_k)}{\partial M_k}$$ 同理，针对路由打分向量 $k(u_t)$，梯度通过前向激活状态与对偶概率导数的乘积进行传递： $$\frac{\partial \ell}{\partial k(u_t)} \approx \left( \frac{\partial \ell}{\partial x_t} (P_t D_t x_{t-1})^T \right) \frac{\partial \operatorname{softmax}_\tau(k(u_t))}{\partial k(u_t)}$$ 当温度参数 $\tau \to 0$ 时，连续松弛精确收敛于前向离散选择，保证了梯度估计在渐近意义下的无偏性。该设计避免了在前向传播中存储密集权重张量，反向传播的额外显存消耗仅为静态字典参数本身的 $\mathcal{O}(KN^2)$，对长序列长度 $L$ 完全解耦。

3. 三阶段分块并行硬件感知 CUDA 内核：实现高带宽利用率状态更新

即便算法复杂度降至线性，朴素的 PyTorch Associative Scan 在处理稀疏置换与对角复合运算 $x^{(t)}[i] = b_t[i] + D_t[q] x^{(t-1)}[q]$（其中 $q = P_t[i]$）时，仍会因频繁的显存随机读写与核函数调用开销而受限。为此，本文设计了高度优化的三阶段分块并行 CUDA 内核。将序列长度 $L$ 划分为 $C = \lceil L / \tau_{\text{chunk}} \rceil$ 个块（默认块大小 $\tau_{\text{chunk}} = 128$）： - 阶段 A（块内局部扫描与汇聚）：每个 CUDA 线程块处理一个批次样本的一个时间块，线程块内使用寄存器和轻量共享内存（仅需 $28N$ 字节，在 $N=128$ 时仅占 3.5KB，远低于 A100 的 164KB 限制）顺序迭代计算该块内部的状态累积与转移变换算子，输出块级复合算子 $(\bar{P}_c, \bar{D}_c, \bar{b}_c)$； - 阶段 B（块间全局前缀传递）：单个线程块在时域上快速按序扫描各块的总算子，利用结合律递推得出进入每个时间块前的排他性前缀状态（Exclusive Prefix Carry）； - 阶段 C（全局前缀注入与最终状态写回）：再次启动并行线程块，各时间块以对应前缀状态为初始输入，重演块内递推并将每步隐藏状态 $x_t$ 连续写回全局显存。整个过程仅启动 3 次 CUDA 核函数，有效消除了寄存器溢出（Spilling），在 NVIDIA A100 GPU 上实测达到了理论峰值全局显存带宽的 82%，相较 PyTorch 关联扫描基线取得多达 10 到 100 倍的端到端执行加速。

实验关键数据¶

主实验¶

论文在四类极具代表性的基准上进行了评估：合成有限状态机模拟（长度外推验证）、MAD 机制能力测试、长达 17,000+ 步的 UEA 多变量时间序列分类，以及在大模型骨干中替换 SSM 层的自然语言状态追踪与下游语言建模。

在 Delétang 等人提出的 4 项合成状态追踪任务（Cycle Nav.、Even Pairs、Mod. Arith.、Parity）中，训练长度 $\le 40$，验证长度外推至 $40 \sim 256$。如下表所示，Flash PD-SSM 在验证集上达到平均 96.3% 的准确率，显著超越对角线架构和 Householder 族模型：

转移矩阵 $\mathbf{A}(u_t)$ 结构	模型	Cycle Nav. (%)	Even Pairs (%)	Mod. Arith. (%)	Parity (%)	平均准确率 (%)
无约束注意力	Transformer	24.4 ± 0.5	90.4 ± 10.4	23.6 ± 0.7	52.2 ± 0.4	47.7 ± 2.6
实数对角线	Mamba2	48.4 ± 2.2	100.0 ± 0.0	33.1 ± 6.6	54.2 ± 2.1	58.9 ± 1.8
复数对角线	Mamba3	69.0 ± 11.7	100.0 ± 0.0	46.8 ± 5.5	76.5 ± 7.8	73.0 ± 6.2
DPLR 矩阵	DeltaNet	49.8 ± 4.7	100.0 ± 0.0	42.2 ± 4.8	57.8 ± 0.8	62.5 ± 1.7
DPLR 连乘	Gated DeltaProduct	46.3 ± 6.6	100.0 ± 0.0	78.4 ± 10.9	98.0 ± 1.4	80.7 ± 3.2
结构化稀疏	PD-SSM (基线)	99.5 ± 0.7	99.7 ± 0.3	96.2 ± 3.4	99.9 ± 0.1	98.8 ± 0.9
结构化稀疏	Flash PD-SSM (本文)	92.7 ± 1.9	98.8 ± 0.8	94.7 ± 2.9	99.0 ± 1.5	96.3 ± 1.8
随机基线	Random	20.0	50.0	20.0	50.0	35.0

在自然语言复杂状态追踪基准（Box 数据集，训练转移步数 $\le 6$，OOD 测试转移步数扩展至 26）中，分别在开源混合模型 Granite4 与 NemotronH 的注意力层前插入并微调 4 层候选 SSM，结果如下表所示：

骨干模型族	规模	Mamba2 (OOD %)	DeltaNet (OOD %)	Flash PD-SSM (本文 OOD %)	相对 Mamba2 提升
Granite4	300M	53.29	49.64	53.42	+0.13%
Granite4	1B	72.81	68.59	74.95	+2.14%
Granite4	3B	85.32	84.12	90.68	+5.36%
Granite4	7B	87.65	88.63	89.79	+2.14%
NemotronH	9B	79.89	82.02	89.89	+10.00%
NemotronH	12B	82.04	84.77	89.25	+7.21%
NemotronH	30B	89.14	90.45	91.77	+2.63%

消融实验¶

在 FineWeb 100B 数据集上完成预训练后，基于 LM-Eval-Harness 对纯 DeltaNet、纯 Flash PD-SSM 以及不同混合策略的 Delta-PD（每 2 层 DeltaNet 插入 1 层 Flash PD-SSM，即 3, 6, 9, 12 层替换）进行消融与对比：

模型架构配置	参数量	FineWeb ppl ↓	PIQA acc ↑	Arc-E acc ↑	WinoGrande acc ↑	平均准确率 (%) ↑	吞吐量 (tok/s) ↑
DeltaNet (基线)	100M	2.85	64.7	42.0	49.1	37.2	350k
Delta-PD (3,6,9,12 混合)	100M	2.81	64.6	42.4	51.3	38.0	368k (+5%)
Flash PD-SSM (纯 1–12 层)	100M	3.08	63.2	42.6	51.9	34.0	-
Delta-PD (前 6 层 DeltaNet)	100M	2.84	65.5	42.1	52.9	37.4	-
Delta-PD (后 6 层 DeltaNet)	100M	2.82	65.2	42.4	50.4	37.5	-
DeltaNet (基线)	230M	2.61	68.3	46.8	51.3	41.3	168k
Delta-PD (3,6,9,12 混合)	230M	2.60	69.4	49.2	53.3	42.3	253k (+50%)
Flash PD-SSM (纯 1–12 层)	230M	2.84	67.8	46.6	51.6	38.5	-

此外，在 Granite4-3B 上微调层数的消融显示：当插入层数从 1 增加至 4 时，Flash PD-SSM 的 OOD 追踪准确率由 80.84% 跃升至 90.68%，而 Mamba2 在增加到 2 层后即出现性能饱和（85.55% 停滞于 85.32%）。改变字典大小 $K$（8 到 256）的实验证实，随着候选稀疏算子数目扩增，NemotronH-9B 的 OOD 状态追踪表现由 89.89% 单调递增至 94.44%。

关键发现¶

离散选择完全保留了有限状态自动机的模拟能力：Flash PD-SSM 相比需要复杂凸组合计算的 PD-SSM，在合成任务上保持了 96.3% 对 98.8% 的极高外推准确率，证明预稀疏化离散选择完全足以承载形式语言中的转移表。
纯 Flash PD-SSM 适合状态追踪，与联想记忆算子混合效果最佳：在 MAD 基准测试中，纯 Flash PD-SSM 强于选择性复制与压缩，但联想召回弱于 DeltaNet；当以 1:3 比例与 DeltaNet 混合构建 Delta-PD 时，不仅下游零样本评估优于纯 DeltaNet，更因其高效稀疏内核将推理吞吐提升了 50%（从 168k 增至 253k tok/s）。
显存与速度实现双重反超：在序列长度达到 5120 时，Flash PD-SSM 相比 Mamba2 训练速度快 4%、比 DeltaNet 快 18%，同时显存占用降低超过 12%，彻底消除了结构化稀疏模型无法规模化应用的工程阻碍。

亮点与洞察¶

离散预稀疏化解耦空间开销：通过将每列最大值选择提炼到模型前向初始化时执行，成功将连续加权生成的张量开销从空间瓶颈中剥离，仅通过索引查表即可完成符号级算子装载。
三阶段分块 Kernel 精准契合 GPU 存储金字塔：局部块内扫描仅利用 3.5KB 共享内存与纯寄存器累加，跨块前缀采用闭式递推，在规避昂贵全精度全局写回的同时跑出了 82% 的峰值显存带宽利用率。
即插即用的混合架构增益：在大规模语言模型中仅替换不到 10% 的层数（如 4 层），即可在不破坏模型既有常识能力的前提下，显著增强其对于长程实体流动与逻辑状态的追踪稳定性。

局限与展望¶

直通估计器的梯度偏差问题：作者指出模型高度依赖直通估计器（STE）穿透离散路由与稀疏极大化，这种梯度替代在复杂深层网络中可能引入不可忽视的估计偏差，这在一定程度上解释了为何完全由 Flash PD-SSM 构成的模型在纯语言建模任务上略微逊色于 DeltaNet。
状态追踪与联想泛化的正交性：单纯增强 FSA 状态追踪无法完全覆盖类似下推自动机（PDA）所需的栈式记忆或复杂上下文检索需求，未来需要更深入地探索其与注意力层或门控联想记忆单元的最优协同比例。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙提出预稀疏化离散选择与退火直通估计，破除了结构化稀疏 SSM 长久以来的显存爆炸瓶颈。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖合成自动机、机制探测、1.7万步长序列时间序列、百亿 Token 预训练及 30B 混合大模型微调。
写作质量: ⭐⭐⭐⭐⭐ 理论证明完备，工程分析与硬件内存层次结构论述详实深入。
价值: ⭐⭐⭐⭐⭐ 为次代线性注意力与状态空间混合架构（如 Mamba/DeltaNet 升级版）提供了极具潜力的基础运算单元。