🤖 machine learning

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

本文介绍了 Flash PD-SSM，这是一种内存优化的结构化稀疏状态空间模型，它从一组可训练的稀疏矩阵中动态选择，在实现非结构化模型高表达力的同时，保持大规模训练所需的计算效率，并在长序列任务上达到最先进的性能。

原作者： Aleksandar Terzić, Francesco Carzaniga, Nicolas Menet, Yannick Biehl, Michael Hersche, Thomas Hofmann, Abbas Rahimi

发布于 2026-05-20

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Aleksandar Terzić, Francesco Carzaniga, Nicolas Menet, Yannick Biehl, Michael Hersche, Thomas Hofmann, Abbas Rahimi

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试构建一个超级智能的机器人，它能够阅读长篇故事、记住细节，并预测接下来会发生什么。为了实现这一目标，该机器人需要一个“大脑”，能够在阅读过程中保留信息。

很长一段时间以来，胜任这项工作的最佳“大脑”是Transformer（许多当前 AI 聊天机器人背后的技术）。它们极其智能，但非常笨重且缓慢，就像一辆耗油巨大、需要巨大车库的豪华轿车。

随后，工程师们发明了状态空间模型（SSMs）。可以将它们想象成电动滑板车。它们速度快得多，且能耗（内存）低得多，非常适合长途旅行。然而，早期的滑板车存在一个问题：它们过于简单。它们可以处理直路，但如果路径变得曲折或需要复杂的逻辑（例如记住针对特定角色的特定规则），它们就会迷路。

这篇论文介绍了一项新发明，称为FLASH PD-SSM。它就像将那辆电动滑板车升级，配备了一套智能、模块化的导航系统，使其在处理复杂曲折方面与笨重的豪华轿车一样出色，同时又不损失其速度或燃油效率。

以下是他们如何实现这一点的分解说明，以简单的概念呈现：

1. 问题：“过于沉重”的地图

这些智能滑板车的先前版本（例如名为PD-SSM的模型）试图通过在旅程的每一步携带一张巨大且详细的地图来变得超级智能。

问题所在： 在长途旅行中携带这张巨大的地图占用了太多内存，导致滑板车在走远之前就会耗尽电量。对于实际应用而言，它过于笨重而不切实际。
结果： 其他模型（如Mamba）选择携带一张微小、简单的地图。它们快速且轻便，但无法解决复杂的谜题或记住复杂的规则。

2. 解决方案：“魔法开关”

本文的作者FLASH PD-SSM想出了一个巧妙的技巧。他们不是在每一步都携带一张巨大且详细的地图，而是构建了一个由预制的专用地图组成的工具箱。

工作原理： 在旅程的每一步，机器人都会查看当前的情况，并切换开关，从工具箱中挑选出最适合该时刻的唯一一张完美地图。
类比： 想象你在开车。与其每次转弯时都绘制一张全新的、详细的整个城市地图（这需要很长时间并消耗大量纸张），不如你拥有一套 100 张预先绘制好的“局部地图”。你只需挑选下一街区所需的那一张即可。
优势： 这个开关切换起来极其迅速，且几乎不占空间。这使得机器人能够拥有豪华轿车的复杂性（它可以解决复杂的逻辑谜题），同时保持滑板车的速度和轻便性。

3. 他们证明了什么

团队通过三种主要方式测试了这种新型“智能滑板车”：

逻辑测试（FSA 模拟）： 他们给机器人一系列逻辑谜题（例如计算奇偶性或穿越迷宫）。FLASH PD-SSM 解决了几乎所有谜题（准确率为 96%），而较简单的模型（如 Mamba2）则表现挣扎。这证明了机器人实际上能够“思考”复杂的规则。
长记忆测试（时间序列）： 他们要求机器人分析极其长的数据流（超过 17,000 步）。FLASH PD-SSM 在预测接下来会发生什么方面最为准确，击败了所有其他竞争对手。
语言测试： 他们将这个新大脑放入一个语言模型（一个撰写文本的机器人）中。
- 状态跟踪： 当被要求跟踪故事中物体在盒子之间的移动时（例如，“红球从盒子 A 移动到了盒子 B"），新模型在记住最终位置方面比以前的模型要好得多。
- 写作： 当训练其像人类一样写作时，该模型的混合版本（将新大脑与现有技术相结合）比仅使用旧式简单大脑的模型写得更好、更快。

4. 核心结论

该论文声称，FLASH PD-SSM解决了当前 AI 中最大的权衡：速度与智能。

以前： 你必须在“快但‘笨’"（简单模型）和“‘聪明’但慢且昂贵”（复杂模型）之间做出选择。
现在： FLASH PD-SSM 就像一列既能高速运行又能翻山越岭的高速列车。它运行速度与当前顶级模型（Mamba2）一样快，但能处理更复杂的任务，同时使用更少的内存。

简而言之，他们找到了一种方法，使 AI 模型能够同时变得更轻、更快、更智能，而无需构建更大、更昂贵的计算机来运行它们。

技术摘要：FLASH PD-SSM

问题陈述

状态空间模型（SSM）为 Transformer 提供了一种可扩展的替代方案，在推理过程中提供线性计算扩展和恒定内存消耗。然而，受限于其转移矩阵的结构，它们面临着效率与表达能力之间的根本权衡。

非结构化矩阵提供最大表达能力（能够模拟任何有限状态自动机，即 FSA），但会产生过高的计算和内存成本。
结构化矩阵（例如对角矩阵、缩放单位矩阵）效率极高，但表达能力有限，理论上无法模拟具有不可解变换半群的 FSA。
结构化稀疏矩阵（PD-SSM）： 一项最新进展（PD-SSM）利用结构化稀疏矩阵实现了 FSA 级别的表达能力，同时支持高效的并行计算。然而，PD-SSM 存在一个关键瓶颈：生成输入依赖的稀疏矩阵会产生 $O(LN^2)$ 的内存成本（其中 $L$ 为序列长度， $N$ 为隐藏状态大小）。这一开销使得 PD-SSM 比 Mamba2 等高效竞争对手慢 40 倍且更消耗内存，导致其不适合大规模部署。

方法论：FLASH PD-SSM

作者提出了FLASH PD-SSM，这是一种内存优化的结构化稀疏 SSM，在保留其理论表达能力保证的同时，解决了 PD-SSM 的生成瓶颈。

1. 新颖的选择机制（硬选择）

核心创新是用硬选择机制取代了 PD-SSM 的软选择机制：

原始 PD-SSM： 生成稠密字典矩阵的凸组合，并应用列向量的 hardmax 来稀疏化结果。这需要在每个时间步实例化稠密的 $N \times N$ 矩阵，导致 $O(LN^2)$ 的内存消耗。
FLASH PD-SSM： 预先计算包含 $K$ 个稀疏矩阵的字典。在每个时间步，输入 $u_t$ 通过线性变换生成 $K$ 个分数。模型选择分数最高的单个矩阵（argmax）并直接使用它。
复杂度降低： 由于选中的矩阵本身是稀疏的（one-hot 列选择器），内存成本从 $O(LN^2)$ 降至**$O(LN) $**。矩阵生成的计算成本从$ O(LN^2K) $降至$ O(LNK)$。

2. 梯度传播

由于选择涉及离散操作（argmax/hardmax），模型采用带有斜率退火的**直通估计器（STE）**进行反向传播：

在前向传播中，使用硬选择（argmax）。
在反向传播中，硬选择通过 tempered softmax 进行近似，以允许梯度流经选择权重和字典矩阵。
论文中的命题 2 提供了代理梯度表达式，当温度参数 $\tau \to 0$ 时，这些表达式变得精确。

3. 硬件感知实现

作者实现了一个自定义的三阶段 CUDA 内核（封装在 Triton 中），以利用 GPU 并行性：

阶段 A： 计算序列块内的独立顺序递归。
阶段 B： 聚合块间依赖关系（前缀扫描）以计算进位状态。
阶段 C： 将进位校正应用于块状态。
该设计最大限度地减少了内存流量并最大化了带宽利用率，实现了接近峰值的 GPU 内存带宽。

4. 架构集成

FLASH PD-SSM 嵌入在具有多头 SSM 层的标准预归一化架构中。多头设计允许摊销稠密字典矩阵的参数成本。作者证明，将字典大小 $K$ 按 $O(\sqrt{D})$ （其中 $D$ 为嵌入维度）扩展，可以在不“膨胀”模型大小的情况下保持参数数量可控。

主要贡献

算法改进： 通过离散选择机制，将结构化稀疏 SSM 生成的内存复杂度从 $O(LN^2)$ 降低到$O(LN)$，实现了大规模训练。
理论表达能力： 证明（命题 1）FLASH PD-SSM 保留了非结构化 SSM 的理论保证：它可以使用单层大小为 $N$ 的模型精确表示任何具有 $N$ 个状态的确定性 FSA。
硬件优化： 提供自定义 CUDA 实现，实现了最先进的吞吐量和内存效率。
混合模型集成： 证明 FLASH PD-SSM 可作为混合大语言模型（LLM）中的有效即插即用替代品，在保持吞吐量的同时提升状态跟踪能力。

实验结果

合成状态跟踪与表达能力

FSA 模拟： 在合成任务（奇偶校验、模计数、循环导航）上，FLASH PD-SSM 实现了96.3% 的平均准确率，显著优于对角 SSM（Mamba2: 58.9%, Mamba3: 90.7%）和 DPLR 模型（门控 DeltaProduct: 80.7%）。这证实了其理论表达能力的实际实现。
机制能力（MAD 基准）： FLASH PD-SSM 得分为69.2%，与 Mamba（69.3%）相当，表明其具有强大的上下文回忆和记忆能力。

运行时与内存效率

吞吐量： 在序列长度为 5120 时，FLASH PD-SSM 比 DeltaNet 快18%，比 Mamba2 快4%。
内存： 在序列长度超过 4000 时，与 DeltaNet 和 Mamba2 相比，其峰值内存消耗降低了**>12%**。
可扩展性： 该模型在不同序列长度下保持高吞吐量，解决了此前阻碍 PD-SSM 的瓶颈。

自然语言状态跟踪

在Box 数据集（扩展至 26 个状态转换）上使用混合模型（Granite4, Nemotron-H）进行评估。
插入这些模型中的 FLASH PD-SSM 层，将分布外（OOD）准确率比 Mamba2 提高了多达10%，比 DeltaNet 提高了8%，同时保持了相当的吞吐量。

语言建模

DELTA-PD： 一个结合 FLASH PD-SSM 和 DeltaNet（比例 1:3）的混合模型，在 FineWeb100B 上进行了预训练。
性能： DELTA-PD 在准确率（平均 42.3% vs 41.3%）和吞吐量（253k tok/s vs 168k tok/s，提升 1.5 $\times$ ）方面均超越了纯 DeltaNet。
它也优于纯 FLASH PD-SSM，表明混合化平衡了不同类型 SSM 的优势。

时间序列建模

在 UEA 多变量时间序列分类任务（序列长度 >17,000）上实现了**最先进（SoTA）**的准确率。
在 PPG-DaLiA 预测任务上达到了接近最先进（SoTA）的均方误差（MSE）。

意义与主张

该论文声称，FLASH PD-SSM 成功打破了限制 SSM 发展的效率 - 表达能力权衡。通过降低结构化稀疏矩阵的内存开销，它使得在满足现代语言模型所需规模下使用高表达能力模型（能够模拟 FSA）成为可能。

作者将 FLASH PD-SSM 定位为不仅仅是理论上的改进，而是前沿混合模型中现有 SSM 层的实用即插即用替代品。它在性能格局中占据独特位置：它是图 1 中“右上象限”中唯一被识别的 SSM，意味着它同时具备高表达能力和快速/内存高效的特性。

论文承认了使用直通估计器的局限性，指出其在高度离散序列方法中有效性的理论解释仍是未来研究的开放领域。然而，在合成、机制和自然语言任务上的实证结果验证了该方法作为可扩展、高表达能力序列建模的稳健解决方案。