ReplaySSM:缓存 SSM 输入,而非状态
ReplaySSM: Cache SSM Inputs, Not State
基于摘要分析。该研究针对混合模型解码中状态空间模型(SSM)逐步读写循环状态造成的内存流量瓶颈,提出 ReplaySSM:展开递推,将状态表示为检查点与近期 SSM 输入缓冲区,分离输出计算和检查点更新。主要贡献是减少状态写回、简化推测解码回滚,并隔离低精度状态舍入误差对后续递推的影响。 机制上,ReplaySSM 周期性更新检查点,使多数解码步无需写回状态;作者报告,主导性的状态内存流量因此约减半。保存近期输入后,推测解码回滚可通过移动缓冲区指针完成,并支持不显式生成中间状态的仅输出验证。精度方面,多数步骤使用低精度检查点计算输出,而循环更新使用较高精度检查点,使低精度副本的舍入误差不进入后续更新。 作者在 vLLM 中实现了面向 Mamba-2 与 Gated DeltaNet 的 ReplaySSM。作者报告,在服务批大小为 256、覆盖四个混合模型的标准解码评估中,相对调优后的 vLLM 基线获得 1.20–1.47 倍加速;在两个 120B 级模型的推测解码评估中,吞吐量约为 vLLM 推测解码的 1.8 倍。作者还报告,分离状态精度相较共享低精度检查点可降低输出 KL divergence 和下游 Accuracy 损失,但摘要未给出具体数值。硬件、检查点更新周期、缓冲区开销、精度配置、任务与实验协议、消融及统计信息尚未验证。 平台推测(待验证):若流式 EEG 解码器采用兼容的 SSM 递推结构,且瓶颈确实在状态读写而非信号预处理,该机制可能用于降低在线推理开销;这一假设依赖可展开的递推形式和可缓存的连续输入。可复用的是检查点与输入缓冲管理,需改造的是 EEG 流式窗口、状态重置及数值精度策略。低信噪比、跨被试或通道变化可能放大量化敏感性,小数据条件也不能由推理优化解决。一个可证伪的小实验是在同一已训练 EEG SSM 上,以相同输入流、硬件与精度对照逐步更新和周期性更新,测量端到端延迟、峰值内存及任务指标变化;小批量下收益可能不成立。已有 EEG 相关工作尚未检索确认。
值得核对其以输入缓存替代逐步状态写回的推理机制,以及在相同服务协议下的加速与精度代价;摘要提供了混合模型结果,但尚不能据此确认 EEG/BCI 收益。
来源:OpenReview · State space models · openreview.net