arXiv:2604.17221eess.SYcs.LG2026-04

改进Mamba模型的双线性计算与记忆保留能力

Bilinear Input Modulation for Mamba: Koopman Bilinear Forms for Memory Retention and Multiplicative Computation

  • 引入可分解的双线性输入调制,增强状态与输入的交互
  • 双线性变体在延迟摆和NARMA-10任务中显著提升性能
  • 适合需要强记忆与复杂计算的任务,如序列建模

选择性状态空间模型(SSMs),尤其是Mamba,采用对角状态转移,限制了记忆保持能力和双线性计算容量。本文提出一种分解的双线性输入调制,通过引入状态-输入乘积,可解释为有限维Koopman双线性形式。在跨通道共享状态(耦合SSM)的基础上,该调制有三种实现方式:耦合双线性输入调制(seq-BIM)在输入侧保留完整双线性乘积,但需串行计算;耦合门控调制(GM)将其线性化为门控调制,兼容并行扫描;并行双线性输入调制(p-BIM)将相同双线性乘积置于状态转移中,仍保持并行可扫描性。在多输入延迟摆(记忆保留)和NARMA-10(双线性计算)任务上的实验揭示了明显差异:GM显著提升记忆保留但不影响双线性计算,而seq-BIM和p-BIM同时改善两者。路径消融实验确认双线性信号的两条下游路径具有互补作用。统计上,双线性变体始终优于其他变体,且仅双线性变体随状态维度增大而表现提升,表明双线性机制是唯一能有效利用更大状态空间的机制。

原文摘要 · Abstract (English)

Selective State Space Models (SSMs), notably Mamba, employ diagonal state transitions that limit both memory retention and bilinear computational capacity. We propose a factorized bilinear input modulation that augments the SSM with a state-input product, interpretable as a finite-dimensional Koopman bilinear form. After introducing a shared state across channels (Coupled SSM), the modulation admits three implementations. Coupled Bilinear Input Modulation (seq-BIM) retains the full bilinear product on the input side at the cost of sequential computation, Coupled Gated Modulation (GM) linearizes it into a gate modulation that is compatible with the parallel scan, and Parallel Bilinear Input Modulation (p-BIM) places the same bilinear product on the state transition while remaining parallel-scannable. Experiments on a multiple input-delay pendulum (memory retention) and NARMA-10 (bilinear computation) reveal a clear dissociation. GM substantially improves memory retention but not bilinear computation, while both seq-BIM and p-BIM improve both. A pathway ablation confirms that the two downstream routes of the bilinear signal serve complementary roles. The improvement is statistically robust, with the bilinear variants consistently outperforming the other variants on bilinear computation. Furthermore, only the bilinear variants benefit from increasing the SSM state dimension, while coupling or gate modulation alone show no improvement, establishing the bilinear mechanism as uniquely capable of exploiting larger state spaces.

Mamba双线性状态空间记忆保留

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。