状态空间模型内存与学习动力学理论

本文介绍一种理论框架，揭示状态空间模型（SSM）的内存容量如何决定学习方向，并证明初始化参数的重要性。研究为SSM优化提供了新的理论基础和改进策略。¹

---

1. 背景：SSM的内存问题

1.1 SSM简介

状态空间模型通过隐状态 $h_t$ 捕获序列信息：

$$h_{t+1}=A{h}_t+B{x}_t,y_t=C{h}_t+D{x}_t$$

其中 $A\in {\mathbb{R}}^{N\times N}$ 为状态转移矩阵，$B,C,D$ 为输入/输出矩阵。

关键问题：SSM能够存储多少历史信息？存储长度与准确性之间如何权衡？

1.2 HiPPO初始化

Gu等人提出的HiPPO框架为连续时间记忆提供了理论基础：

$$A_{n,k}=\{\begin{array}{ll}\sqrt{2k+1}& n=k\\ \sqrt{2n+1}& n<k\\ 0& n>k\end{array}$$

这确保了S4模型理论上具有最长的可能记忆。

1.3 理论与实践的差距

尽管HiPPO提供了良好的初始化，但：

• 训练过程中内存可能退化
• 固定权重的性能有时优于微调
• 初始化对最终性能影响巨大

本文目标：提供严格的理论解释。

---

2. 内存容量评估框架

2.1 输入存储模型

考虑SSM处理输入序列 $\{x_1,\dots ,x_T\}$。定义存储函数：

$$s_t=\varphi (x_t,x_{t-1},\dots ,x_1)$$

其中 $\varphi$ 由SSM的隐状态实现。

2.2 内存准确度与长度

定义：

• 内存长度 $L$：能够可靠重建历史输入的最大步数
• 内存准确度 $A$：重建输入的保真度

定理（内存权衡）：对于S4类型的SSM，存在基本权衡：

$$L\cdot A\le \kappa (N,\alpha )$$

其中 $N$ 为状态维度，$\alpha$ 为谱范数参数，$\kappa$ 为依赖SSM结构的常数。

2.3 内存容量的谱分析

SSM的内存能力由矩阵 $A$ 的谱性质决定：

定理：设 $\{{\lambda }_i{\}}_{i=1}^N$ 为 $A$ 的特征值，则内存长度满足：

$$L\approx \frac{1}{{\min }_i|\log |{\lambda }_i||}$$

• 稳定点 (${\lambda }_i\approx 1$)：长时间记忆
• 快速衰减 ($|{\lambda }_i|<1$)：短期记忆

---

3. S4与简化S4的理论等价性

3.1 简化S4

考虑具有对角状态转移矩阵的简化版本：

$$h_{t+1}=\Lambda h_t+b{x}_t,\Lambda =\text{diag}({\lambda }_1,\dots ,{\lambda }_N)$$

3.2 理论等价性定理

定理：在以下条件下，完整S4与简化S4的内存能力等价：

1. 线性区：在输入幅度受限时
2. 谱条件：$\Vert A-\Lambda {\Vert }_2\le ϵ$，其中 $ϵ$ 很小
3. 时间尺度：${\Delta }_t$ 满足特定约束

直觉：当偏离对角结构不大时，系统的内存行为主要由谱性质决定。

3.3 实践意义

这解释了为什么简化的SSM（如Mamba的线性注意力变体）能够保持竞争力的性能。

---

4. 学习动力学分析

4.1 参数化的影响

考虑SSM参数 $\theta =(A,B,C)$。定义损失函数：

$$\mathcal{L}(\theta )={\mathbb{E}}_{x,y}[(y_t-{\hat{y}}_t{)}^2]$$

4.2 梯度与内存的耦合

关键发现：梯度方向与内存结构密切相关！

定理：设 ${\theta }^{\ast }$ 为最优参数，${\theta }_0$ 为初始参数。梯度更新

$${\theta }_{k+1}={\theta }_k-\eta \nabla \mathcal{L}({\theta }_k)$$

满足：

$$\mathbb{E}[\Vert {\theta }_k-{\theta }^{\ast }\Vert ]\le \underset{\text{内存依赖衰减}}{\underbrace{\Vert {\theta }_0-{\theta }^{\ast }\Vert \cdot \prod _{i=1}^k(1-\eta {\lambda }_i)}}$$

其中 ${\lambda }_i$ 与状态矩阵 $A$ 的谱性质相关。

4.3 初始参数的关键作用

定理（初始化重要性）：设 ${\theta }_0$ 由HiPPO初始化，则

$$\mathbb{E}[\mathcal{L}({\theta }_k)]\le \mathcal{L}({\theta }^{\ast })+O\left(\frac{1}{L_0}\right)$$

其中 $L_0$ 为 ${\theta }_0$ 对应的初始内存长度。

推论：更长的初始内存 $\Rightarrow$ 更快收敛 $\Rightarrow$ 更好的最终性能。

---

5. 固定循环权重的优势

5.1 固定vs可训练

传统观点认为：可训练参数总是优于固定参数。但实验表明，固定HiPPO权重的SSM有时优于完全微调的版本。

5.2 理论分析

定理：如果真实内存需求 $L^{\ast }\le L_0$（初始内存），则固定权重是最优的：

$$\mathcal{L}({\theta }_{fixed})\le \mathcal{L}({\theta }_{finetuned})$$

当且仅当存在 ${\theta }_{finetuned}$ 使得 $L>L^{\ast }$。

5.3 实际建议

场景	推荐策略
长依赖任务	使用HiPPO初始化，允许微调
短依赖任务	固定HiPPO权重
中等依赖任务	固定权重 + 微调B, C

5.4 收敛速度比较

定理：固定权重SGD的收敛速率为：

$$\mathbb{E}[\mathcal{L}({\theta }_k)]-\mathcal{L}({\theta }^{\ast })\le \frac{C}{k}$$

而完全微调的收敛速率为：

$$\mathbb{E}[\mathcal{L}({\theta }_k)]-\mathcal{L}({\theta }^{\ast })\le \frac{C^{\prime }}{k^{\alpha }},\alpha <1$$

其中 $\alpha$ 依赖于内存长度与目标长度的比例。

---

6. 长内存任务实验

6.1 实验设置

在需要长程记忆的任务上验证理论：

• 复制任务：需要记忆 $T$ 步前的输入
• 选择性回忆：需要从长序列中检索特定位置

6.2 主要发现

任务	内存需求	固定HiPPO	微调HiPPO
复制 (T=100)	短期	✓✓✓	✓✓
复制 (T=1000)	中期	✓✓	✓✓✓
选择性回忆	长期	✓	✓✓✓

6.3 扩展内存的困难

关键发现：训练过程中扩展内存极其困难！

一旦初始内存长度 $L_0$ 确定，后续优化主要提高内存准确度，而非延长内存长度。

这强调了良好初始化的重要性。

---

7. 优化策略建议

7.1 初始化策略

def ssm_hippo_init(N, mode='legt'):
    """
    HiPPO初始化
    
    Args:
        N: 状态维度
        mode: 'legt' (LegT) 或 'lagt' (LagT)
    
    Returns:
        A: 状态转移矩阵 (N x N)
        B: 输入矩阵 (N,)
    """
    if mode == 'legt':
        # LegT (Legendre Translated)
        n = np.arange(N)
        A = np.sqrt(2 * n + 1)[:, None] * np.ones((1, N))
        A = np.tril(A) + np.diag(n + 1)
    else:
        # LagT (Laguerre Translated)
        # ... 类似构造
        pass
    
    # 确保谱性质
    eigvals = np.linalg.eigvals(A)
    assert np.all(np.abs(eigvals) <= 1 + 1e-6)
    
    return A, np.ones(N)

7.2 混合策略

class HybridSSM(nn.Module):
    def __init__(self, N, d_model, freeze_A=True):
        super().__init__()
        # HiPPO初始化A
        A_init, _ = ssm_hippo_init(N)
        self.A = nn.Parameter(A_init, requires_grad=not freeze_A)
        self.B = nn.Parameter(torch.randn(N, d_model))
        self.C = nn.Parameter(torch.randn(d_model, N))
    
    def forward(self, x):
        # 状态更新
        h = torch.tanh(F.linear(x, self.B) + F.linear(h, self.A))
        return F.linear(h, self.C)

7.3 自适应内存扩展

class AdaptiveSSM(nn.Module):
    def __init__(self, N_init, N_max):
        super().__init__()
        self.N = N_init
        self.N_max = N_max
        
        # 初始化短内存
        A_init, B_init = ssm_hippo_init(N_init)
        self.A = nn.Parameter(A_init)
        self.B = nn.Parameter(B_init)
    
    def expand_memory(self):
        """训练过程中扩展内存（如果需要）"""
        if self.N < self.N_max:
            self.N = min(self.N + 16, self.N_max)
            # 扩展A, B矩阵
            # 保持原有谱性质

---

8. 总结与展望

8.1 主要贡献

1. 内存容量理论：建立了SSM内存能力与谱性质的精确联系
2. 学习动力学：揭示了梯度优化与内存结构的耦合关系
3. 初始化重要性：提供了严格的理论解释
4. 固定权重优势：证明了在某些场景下固定权重的优越性

8.2 实践建议

• 始终使用HiPPO或类似初始化
• 对于短依赖任务，考虑固定状态转移矩阵
• 谨慎对待”训练过程中扩展内存”的期望
• 在设计新SSM变体时，考虑内存容量作为首要指标

8.3 未来方向

• 将理论扩展到选择性SSM（如Mamba）
• 研究内存容量与泛化能力的联系
• 设计内存感知的优化器

---

参考资料

Footnotes

1. Memory Determines Learning Direction: A Theory of Gradient-Based Optimization in State Space Models. ICLR 2026 Under Review. ↩