跳到正文
OpenReview · State space models·· 2026-09-07精选AI 评分82

SketchSSM:写入完整状态,从紧凑 sketch 读取

SketchSSM: Write to the Full State, Read from a Compact Sketch

AI 导读

基于摘要分析。本文研究混合注意力语言模型解码时的循环状态访存瓶颈,提出 SketchSSM:保留完整状态更新,通过紧凑 sketch 近似后续查询的状态读取,以减少重复访问完整状态的开销。主要研究对象是语言模型解码,而非 EEG/BCI 信号建模。 核心机制:混合注意力模型以线性注意力替代多数 softmax attention 层,降低 KV-cache 增长,但在较大的解码批次下,循环状态访问可能成为瓶颈。ReplaySSM 通过缓存 keys 和 values 摊销状态更新,不过状态更新之间的每个新 query 仍需读取完整状态。作者观察到,低秩的状态加权 query 近似能够较好保留状态读取输出,而且用于近似 query 的基向量可离线固定。SketchSSM 在每次状态更新时读取一次完整状态,预计算这些基向量对应的输出并存入紧凑 sketch;后续解码步骤以 query 相关系数组合 sketch 向量,重构输出而不再读取完整状态。其关键区别是近似读取,而非直接压缩或替换完整写入状态。 作者报告:在四个基于 Mamba-2、GDN 和 KDA 的模型上,状态访问流量约降低 10 倍,同时基本保持四个解码基准上的平均准确率,以及四项 RULER 检索任务上的召回表现。在单张 NVIDIA B300 上,相对标准 vLLM 基线,Mamba-2、GDN、KDA 的线性注意力内核最高加速分别为 7.78 倍、5.22 倍和 5.20 倍;Nemotron 3 Super 的解码吞吐量最高提升 2.64 倍。内核加速与整体吞吐量是不同指标,不能互换。 尚需核对离线基向量的构建数据与方法、近似秩、更新频率、误差分布,以及具体批次和序列长度。实验协议、消融及统计信息尚未验证;摘要未提供各基准的准确率变化或完整复现条件。 平台推测(待验证):若流式 EEG 模型同样依赖较大的循环状态,且多次查询之间状态保持不变,该读取策略可能降低在线推理访存;可复用的是 sketch 预计算与组合机制,需改造的是信号模型的查询表示、基向量构建和更新调度。低信噪比、跨被试或通道变化可能使固定基失效,小数据可能不足以估计稳定基,频繁状态更新则可能抵消收益。可先在同一流式 EEG 模型上比较完整读取与 sketch 读取的输出误差、任务指标及延迟,检验是否确有访存瓶颈和收益。已有 EEG 相关工作尚未检索确认。

阅读价值

值得核对其“完整状态写入、紧凑 sketch 读取”的近似机制,以及读取误差与访存收益之间的权衡;作者报告了特定语言模型解码场景下的加速,但 EEG/BCI 适用性尚未验证。

来源:OpenReview · State space models · openreview.net