跳到正文
OpenReview · State space models· Arghadip Das; Arnab Raha; Shamik Kundu; Soumendu Kumar Ghosh; Deepak Mathaikutty; Vijay Raghunathan·· 2025-03-06精选AI 评分74

XAMBA:在资源受限的神经处理单元上实现高效状态空间模型

XAMBA: Enabling Efficient State Space Models on Resource-Constrained Neural Processing Units

AI 导读

基于摘要分析。该研究关注状态空间模型(SSMs)在资源受限商用神经处理单元(NPUs)上的部署,而非 EEG 或 BCI 建模。XAMBA 通过算子改写与激活函数近似,使 SSMs 适配现有硬件并优化执行性能,避免为每类模型重新设计专用加速器。 方法分为三个步骤:先使 SSMs 能在 NPU 上运行,再针对目标关键性能指标优化,最后以一定精度代价换取额外性能。CumBA 与 ReduBA 将 CumSum、ReduceSum 操作替换为矩阵计算,以改善执行速度与内存效率;ActiBA 则以分段线性近似将 Swish、Softplus 等计算开销较高的激活函数映射到 NPU 硬件。其贡献主要在部署与算子实现层面,摘要未提供具体矩阵改写、近似区间或误差控制方式。 作者报告,在 Intel® Core™ Ultra Series 2 AI PC 上,相较基线实现,执行延迟最高获得 4.8× 加速。摘要未明确对应模型、任务、输入长度、数值精度、基线优化程度及计时范围,因此该数字不能直接外推至其他硬件或 EEG 任务。作者称激活近似仅造成较小精度损失,但具体指标与幅度尚未验证。材料提供了实现仓库地址;环境依赖、支持算子、实验协议、消融及统计信息尚未验证。 平台推测(待验证):若 EEG 序列模型采用 SSMs,且部署瓶颈确实来自上述算子,XAMBA 的硬件映射策略可能有助于降低边缘端推理延迟。这一路径依赖具体 SSM 的计算图、序列长度和 NPU 支持能力,不要求新增监督数据;可复用的是算子替换与激活近似思路,需改造的是 EEG 模型计算图、输入通道布局及近似误差容限。低信噪比、通道变化和跨被试分布差异可能放大近似误差,小数据条件也可能使精度影响难以稳定估计。 可检验的小实验是固定一个已有 EEG SSM、权重及数据划分,在同一 NPU 上比较原实现、仅算子改写、再加入激活近似三种配置,同时记录端到端延迟、内存占用及原任务指标,并保持被试内或跨被试协议一致。已有 EEG 相关工作尚未检索确认,原领域加速不构成 BCI 有效性的证据。

阅读价值

值得阅读的是其针对现有商用 NPU 的算子适配与近似计算路径,可用于核对 SSM 部署瓶颈,但摘要中的最高加速结果尚需结合具体模型、精度和测量协议验证。

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