State Space Models

参考文献：Fundamentals of Recurrent Neural Network (RNN) and Long Short-Term Memory (LSTM) Network

理论起源：微分方程视角

微分方程视角：RNN 的动力学方程可由一阶非齐次 ODE 推导而来。设 $\vec{s}(t)\in {\mathbb{R}}^d$ 为 $d$ 维状态信号向量，其随时间 $t$ 的演化可表示为：

$$\frac{d\vec{s}(t)}{dt}=\vec{f}(t)+\vec{\varphi }$$

• $\vec{x}(t)$ 是 $d$ 维输入信号向量，$s$ 为 $d$ 为状态向量；
• $\vec{f}(t)$ 是 $d$ 维向量函数，通常依赖于 $\vec{s}(t)$ 和 $\vec{x}(t)$，即 $\vec{f}(t)=\vec{h}(\vec{s}(t),\vec{x}(t))$；
• $\vec{\varphi }$ 是常数 $d$ 维向量（偏差项）。

加性模型：这里采用加性模型 (Additive Model)，将 $\vec{f}(t)$ 分解为三个独立项的线性组合：

$$\vec{f}(t)=\vec{a}(t)+\vec{b}(t)+\vec{c}(t)$$

记 ${\tau }_s,{\tau }_r,{\tau }_x$ 分别表示状态、读出和输入的延迟时间常数。

• 状态反馈项 $\vec{a}(t)$：表示状态信号自身的延迟影响（"模拟"分量）

$$\vec{a}(t)=\sum _{k=0}^{K_s-1}{\vec{a}}_k(\vec{s}(t-{\tau }_s(k)))$$

• 读出反馈项 $\vec{b}(t)$：表示经过非线性变换后的读出信号的延迟影响：

$$\vec{b}(t)=\sum _{k=0}^{K_r-1}{\vec{b}}_k(\vec{r}(t-{\tau }_r(k)))$$

• 外部输入项 $\vec{c}(t)$：表示外部输入信号的延迟影响

$$\vec{c}(t)=\sum _{k=0}^{K_x-1}{\vec{c}}_k(\vec{x}(t-{\tau }_x(k)))$$

延迟微分方程 (DDE) 系统：

$$\frac{d\vec{s}(t)}{dt}=\underset{\text{状态反馈}}{\underbrace{\sum _{k=0}^{K_s-1}{\vec{a}}_k(\vec{s}(t-{\tau }_s(k)))}}+\underset{\text{读出反馈}}{\underbrace{\sum _{k=0}^{K_r-1}{\vec{b}}_k(\vec{r}(t-{\tau }_r(k)))}}+\underset{\text{外部输入}}{\underbrace{\sum _{k=0}^{K_x-1}{\vec{c}}_k(\vec{x}(t-{\tau }_x(k)))}}+\vec{\varphi }$$

线性系数 DDE：当上述函数为线性时，可以转化为具有线性矩阵系数的非线性 DDE：

$$\frac{d\vec{s}(t)}{dt}=\sum _{k=0}^{K_s-1}{A}_k\vec{s}(t-{\tau }_s(k))+\sum _{k=0}^{K_r-1}{B}_k\vec{r}(t-{\tau }_r(k))+\sum _{k=0}^{K_x-1}{C}_k\vec{x}(t-{\tau }_x(k))+\vec{\varphi }$$

其中 $A_k,B_k,C_k$ 分别为对应项的系数矩阵。

简化模型 (The Simplified Model)：通过以下特定约束，可将通用模型简化为包含 Continuous Hopfield Network 和 Cellular Neural Network 的特例形式：

• 单项约束：设 $K_s=K_r=K_x=1$（每种反馈仅有一项）。
• 延迟约束：设 ${\tau }_s(0)=0,{\tau }_x(0)=0$（状态和输入无延迟），仅保留读出信号的单一延迟 ${\tau }_r(0)={\tau }_0$。
• 矩阵重命名：$A_0=A,B_0=B,C_0=C$。

$$\frac{d\vec{s}(t)}{dt}=A\vec{s}(t)+B\vec{r}(t-{\tau }_0)+C\vec{x}(t)+\vec{\varphi }$$

时间离散与 RNN 公式

命题(离散化公式)：使用向后欧拉离散上述方程可得到

$$\vec{s}[n]=W_s\vec{s}[n-1]+\underset{\text{循环/记忆项}}{\underbrace{((\Delta T)W_{s}B)\vec{r}[n-1]}}+\underset{\text{输入项}}{\underbrace{((\Delta T)W_{s}C)\vec{x}[n]}}+\underset{\text{偏差项}}{\underbrace{((\Delta T)W_s\vec{\varphi })}}$$

证明：设采样时间步长为 $\Delta T$，时间 $t=n\Delta T$。使用后向欧拉法近似导数：

$$\frac{d\vec{s}(t)}{dt}\approx \frac{\vec{s}(n\Delta T+\Delta T)-\vec{s}(n\Delta T)}{\Delta T}$$

令延迟时间 ${\tau }_0$ 等于单步采样时间 $\Delta T$（即 ${\tau }_0=\Delta T$）：

$$\frac{\vec{s}[n+1]-\vec{s}[n]}{\Delta T}=A\vec{s}[n+1]+B\vec{r}[n]+C\vec{x}[n+1]+\vec{\varphi }$$

这里采用了简化记号 $\vec{s}[n]\equiv \vec{s}(n\Delta T)$。注意方程右侧使用了 $n+1$ 时刻的 $\vec{s}$ 和 $\vec{x}$（后向法的特性），但使用了 $n$ 时刻的 $\vec{r}$（由延迟 ${\tau }_0$ 决定）。接着

$$(I-(\Delta T)A)\vec{s}[n+1]=\vec{s}[n]+((\Delta T)B)\vec{r}[n]+((\Delta T)C)\vec{x}[n+1]+(\Delta T)\vec{\varphi }$$

定义 $W_s=(I-(\Delta T)A{)}^{-1}$，并在方程两边左乘 $W_s$，得到：

$$\vec{s}[n+1]=W_s\vec{s}[n]+((\Delta T)W_{s}B)\vec{r}[n]+((\Delta T)W_{s}C)\vec{x}[n+1]+((\Delta T)W_s\vec{\varphi })$$

将索引向前平移一步 ($n\to n-1$) 可得结论。

标准 RNN 形式：为了进一步简化表示，定义新的权重矩阵和偏差向量：

$$\begin{array}{rl}\vec{s}[n]& =W_s\vec{s}[n-1]+W_r\vec{r}[n-1]+W_x\vec{x}[n]+{\vec{\theta }}_s\\ \vec{r}[n]& =G(\vec{s}[n])\end{array}$$

其中

$$W_r=(\Delta T)W_{s}B,W_x=(\Delta T)W_{s}C,{\vec{\theta }}_s=(\Delta T)W_s\vec{\varphi }$$

• $W_s$: 状态自身的递归权重（State-to-State）。
• $W_r$: 读出信号（即上一时刻的激活值）的反馈权重（Readout-to-State）。
• $W_x$: 当前输入的权重（Input-to-State）。

稳定性分析

稳定性条件：上述系统稳定的条件是矩阵 $\hat{W}=W_s+W_r$ 的所有特征值必须位于复平面的单位圆内。

标准 RNN 定义：在实际应用中，常作进一步简化以获得最简形式：

• 单位时间步长：设 $\Delta T=1$。
• 快速状态衰减：假设矩阵 $A$ 为对角阵且对角元为很大的负数（$a_{ii}\ll 0$），这意味着状态的衰减非常快。
• 忽略状态记忆：由此导致 $W_s\approx -A^{-1}$ 为对角阵且元素为很小的正数。在这种情况下，状态信号 $\vec{s}[n-1]$ 对当前轨迹的显式影响（第一项）可以忽略不计（尽管通过 $\vec{r}[n-1]$ 的隐式影响依然存在）。
• 忽略 $W_s$ 项（即设 $W_s\approx 0$）

我们得到最常见的标准 RNN 定义，其中 $G(s)$ 通常是 $\tanh$ 函数

$$\begin{array}{rl}\vec{s}[n]& =W_r\vec{r}[n-1]+W_x\vec{x}[n]+{\vec{\theta }}_s\\ \vec{r}[n]& =G(\vec{s}[n])\end{array}$$

在此简化下，系统的稳定性完全取决于 $W_r$ 的特征值 ${\mu }_i$。在“小信号”区域（$\Vert \vec{s}[n]\Vert \ll 1$），稳定性的充要条件是 $0<{\mu }_i<1$。

原论文

Investigation of Memory Patterns

SSM 中的记忆衰减

状态更新方程 (State Update Equation)：考虑如下更新格式

$$h_t=A\cdot h_{t-1}+B\cdot x_t,y_t=C\cdot h_t$$

• $A\in {\mathbb{R}}^{d_s\times d_s}$: 状态转移矩阵，满足 $|A|<1$ （谱半径小于 1）以保证 BIBO 稳定性。
• $h_t$: $t$ 时刻的隐藏状态。
• $x_t$: $t$ 时刻的输入。

信息贡献衰减 (Decay of Information Contribution):为了衡量输入 $x_{t-k}$ (发生在 $k$ 步之前) 对当前状态 $h_t$ 的影响，定义其贡献度：

$$\text{Contribution}(x_{t-k}\to h_t)=|A^k\cdot B\cdot x_{t-k}|\le |A^k|\cdot |B|\cdot |x_{t-k}|$$

指数级遗忘: 随着 $k$ 增加 (输入变得更久远)，$A^k$ 呈指数级衰减 ($A^k\approx e^{-\alpha k},\alpha >0$)。因此早期输入几乎被完全遗忘，导致长距离信息丢失。

Transformer 中的记忆衰减

Transformer 时间复杂度 (Time Complexity, TC)：令 $n$ 为序列长度，$L$ 是网络层数，$d$ 是特征维度

$$\text{TC}=O(L\cdot n^2\cdot d)$$

对于超长序列 (如 $n=10^5$)，$n^2$ 项导致运算量 ($10^{10}$) 超出硬件承载能力。因此在实际应用中，为了解决计算瓶颈，通常采用近似方法（如滑动窗口注意力，窗口大小 $w=512$）。

有效建模长度 (Effective Modeling Length, EML)：滑动窗口注意力引入了截断效应：

$$\text{EML}\le w\ll n⟹\text{inability to capture long-range dependencies}$$

窗口之外的信息被直接丢弃，导致无法捕捉超长依赖。

水平与垂直记忆保真度 (Horizontal and Vertical Memory Fidelity)

为了系统地量化关键信息的丢失，本文提出了 水平-垂直记忆保真度框架 (Horizontal–Vertical Memory Fidelity Framework)。

定义(Expected Token Memory Fidelity, ETMF)：Token 级的语义信息在层内递归传播过程中的保留程度，维度水平 (Horizontal)

定义(Expected Cross-Layer Memory Fidelity)：信息在跨层垂直传输过程中的保留程度，维度垂直(Vertical)

Mamba 记忆的双重挑战：(1) 记忆衰减 Memory Decay，由 ETMF 反映，长距离 Token 语义在递归中逐渐模糊。(2) 外推限制 (Extrapolation Limits)，由 ECLMF 量化，信息在深层网络传播中的退化。

Methodology: The MemMamba Network

理念与架构

动机 (Motivation): 现有的状态空间模型 (SSMs) 虽然具有线性复杂度，但在递归更新中会逐渐丢失长距离依赖信息。

灵感 (Inspiration): 模拟人类阅读长文档时的"做笔记" (Note-taking) 行为 —— 当遇到关键信息时将其记录下来，而非试图死记硬背所有内容。

解决方案: MemMamba 在有限的表示空间内动态保存关键上下文，并提供跨层和跨 Token 的长距离交互索引。

架构概述：MemMamba 由 $n$ 个堆叠的 MemMamba Block Layers 组成。每层集成三个组件：

• SSM 更新 (State Space Model updates)
• 跨 Token 注意力 (Cross-token attention)：在每一层执行，用于恢复层内被遗忘的信息。
• 跨层注意力 (Cross-layer attention)：每隔 $p$ 层触发一次 ($l\mathrm{mod}p=0$)，用于整合全局深层信息。

模块细节

笔记模块 Note Block：第 $l$ 层，时间步 $t$，对输入 $x_t^l$ 使用评分函数 ${\mathcal{I}}_{\text{token}}$ 进行评估。如果评分超过阈值 ${\tau }_1$，执行“做笔记”操作，通过降维算子 ${\mathcal{N}}^l$ (如线性投影或池化) 生成摘要 $s_t^l$ 并插入状态池 (state pool) $S_{t-1}^l$。

$$\text{if }{\mathcal{I}}_{\text{token}}(x_t^l)>{\tau }_1⟹s_t^l={\mathcal{N}}^l(x_t^l)$$

$$S_t^l=\text{Insert}(S_{t-1}^l,s_t^l)$$

策略: 状态池采用 FIFO 或基于优先级的替换策略，确保仅保留高信息量的摘要。

跨 Token 注意力 Cross-Token Attention：当检测到当前 SSM 状态可能遗忘重要信息时（即状态评分 ${\mathcal{I}}_{\text{state}}>{\tau }_2$）。在当前输入 $x_t^l$ (Query) 与状态池摘要 ${\tilde{s}}_{t-1}^l$ (Key/Value) 之间进行注意力计算。

$$c_t^{\text{token},l}=\{\begin{array}{ll}\text{Attention}(Q=h_t^l,K=s_{t-1}^l,V=s_{t-1}^l),& \text{if }{\mathcal{I}}_{\text{state}}(z_{t-1}^l)>{\tau }_2\\ 0,& \text{otherwise}\end{array}$$

跨层注意力 Cross-Layer Attention：稀疏激活，仅在特定的层 ($l\mathrm{mod}p=0$) 触发。聚合过去 $g$ 层的状态摘要形成跨层上下文 $s^{\mathcal{R}(l)}$，并进行注意力计算。

$$c_t^{\text{layer},l}=\{\begin{array}{ll}\text{Attention}(Q=x_t^l,K=s^{\mathcal{R}(l)},V=s^{\mathcal{R}(l)}),& \text{if }l\mathrm{mod}p=0\\ 0,& \text{otherwise}\end{array}$$

信息聚合：最终输出 ${\bar{x}}_t^{l+1}$ 是原始输入与两种上下文信息的融合，随后传入 SSM 更新：

$${\bar{x}}_t^{l+1}=x_t^{l+1}+{\mathcal{F}}_{\text{tok}}^l(c_t^{\text{token},l})+{\mathcal{F}}_{\text{lay}}^l(c_t^{\text{layer},l})$$

${\mathcal{F}}_{\text{tok}}^l,{\mathcal{F}}_{\text{lay}}^l$: 融合函数 (如门控机制或残差映射)。

图片讲解

MemMamba 架构

• 布局：纵向 $n$ 层 Layer (神经网络深度)；横向是时间轴，代表数据从左往右流过；Layer k+1 的输入是 Layer k 的输出；
• MemMamba Block：接收当前的 Input Token，生成隐藏状态 h2, h3, ... 并传递给下一个时间步；
• Note Block 记笔记：每个 MemMamba Block 下面有个 note block，如果输入 token 重要 (${\mathcal{I}}_{\text{token}}(x_t)>{\tau }_1$)，则会被存进 state pool (蓝色长条) 中。每隔 $p$ 层会执行一次跨层注意力。
• Note Block 查笔记：当 MemMamba Block 发现 $h_t$ 中信息不够了 (${\mathcal{I}}_{\text{state}}(h_t)>{\tau }_2$)，则立刻回去查阅蓝色池子中的笔记
• 深度整合：最底部的 Layer i 是 MemMamba Block (layer)，上面的层主要关注水平的记忆，而 Layer i 则触发垂直跨层注意力。