记忆决定学习方向:状态空间模型中基于梯度的优化理论¶
会议: NeurIPS 2026
arXiv: 2510.00563
代码: https://github.com/state-spaces/s4/tree/simple
领域: 时间序列
关键词: 状态空间模型, 记忆容量, 学习动力学, 范德蒙矩阵, 储备池计算
一句话总结¶
本文通过记忆函数(Memory Function)形式化分析了线性状态空间模型(SSM)的梯度优化动力学,证明初始状态矩阵特征值张成的范德蒙投影算子决定了反向传播的学习方向,从而在理论上确立了 S4 与 S4D 的等价性,并提出在保留高阶长程记忆初始化的前提下直接冻结递归权重的储备池优化策略。
研究背景与动机¶
以结构化状态空间模型(S4、S4D、Mamba)为代表的线性循环神经网络在长程时间序列建模中取得了匹敌甚至超越 Transformer 的性能。为了捕获极远距离的时间依赖,HiPPO 框架采用连续时间勒贝格测度正交多项式推导出了最优长程投影矩阵,而后 S4D 进一步通过对角化特征值近似化简了系统结构。然而,该领域长期缺乏对 SSM 学习动力学的严谨理论解释:为何特征值的特定初始化(如 S4Dinv 与 S4Dlin)会对最终性能产生决定性影响?为何在记忆容量理论存在“长度与精度不可兼得”权衡的前提下,网络必须不惜牺牲短期重构精度也要强行拉长记忆周期?此外,将满秩循环矩阵强行简化为对角矩阵的 S4D 与原始 S4 在理论上是否等价,此前也未得到严格证明。
这种理论空白直接制约了 SSM 的优化机制设计与训练效率。现有的循环网络动力学分析往往对循环矩阵施加人为的随机正交限制,或者简单忽略特征向量结构,无法推演梯度反向传播过程中历史信息的真实衰减与流动。在多层网络反向传播中,当梯度穿过单个 SSM 序列层时,如果递归核无法将远距离时间步的误差信号有效保留并映射回参数空间,反向传播便无法感知远距离输入与目标标签的依存关联,导致学习停滞或陷入严重的过拟合。
针对上述难题,本文引入动态系统的信息处理容量(Information Processing Capacity, IPC)及其线性退化形式——记忆函数(Memory Function, MF),从底层对离散线性 SSM 的梯度更新方程进行了严格的解析推导。研究发现,单层 SSM 线性读出层的参数梯度完全由递归矩阵特征值所构成的归一化范德蒙投影矩阵所调制。核心 idea:SSM 中历史信息的记忆衰减函数(MF)直接决定了梯度的学习方向,若初始记忆在远距离衰减为零,教师信号将遭遇不可逆的信息丢失而无法驱动有效更新;因此必须优先保证初始记忆的长度跨度,并在具备良好长程初始结构时直接冻结循环权重进行储备池计算(RC),以消除过拟合并加速收敛。
方法详解¶
整体框架¶
本文的方法分析框架涵盖三个演进阶段:首先,利用记忆函数(MF)和信息处理容量(IPC)重新刻画离散线性 SSM 的输入-状态表征,证明系统的信息处理能力仅由状态转移矩阵的特征值决定,从理论上推导 S4 与对角化 S4D 的完全等价性;其次,在均方误差(MSE)目标下精确展开单层 SSM 读出权重的参数梯度表达式,揭示归一化范德蒙投影矩阵对教师信号的加权滤波本质;最后,基于反向传播难以自主扩展零记忆区域的理论缺陷,提出冻结循环权重特征值、仅优化读出与投影层的储备池计算(RC)训练范式。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入时间序列 u_t"] --> B["特征值谱分解与记忆函数建模<br/>证明 S4 与 S4D 在信息处理容量上等价"]
B --> C["范德蒙投影算子梯度解析<br/>导出参数梯度 G(y) 与记忆权重滤波机制"]
C --> D["初始记忆长度充要条件分析<br/>若远距离 MF 归零则教师信息不可逆丢失"]
D --> E["储备池计算优化策略<br/>冻结循环特征值仅更新读出与投射权重"]
E --> F["下游长序列建模任务推理与收敛"]
关键设计¶
1. 记忆函数等价性与 S4/S4D 理论统一:消除对循环转移矩阵特征向量的人为约束 在线性时不变连续状态空间模型经零阶保持器(ZOH)离散化后,系统的状态更新方程可写为 \(\bm{x}_{k+1} = \bm{A}\bm{x}_k + \bm{B}\bm{u}_k\)。以往研究为了降低 S4 的柯西核卷积计算复杂度,通过特征值对角化经验性提出了 S4D,但未解释对角化是否会损失表达能力。本文指出,线性循环系统的任意隐藏状态本质上均是过去输入的线性组合 \(\bm{x}_t = \sum_{\tau=0}^\infty \bm{\alpha}_\tau u_{t-\tau}\)。根据信息处理容量理论,延迟输入 \(u_{t-\tau}\) 能否从当前状态线性重构的度量即为记忆函数 \(M[u_{t-\tau}]\)。解析证明表明,\(M[u_{t-\tau}] = \bm{V}_\tau^\top (\bm{V}\bm{V}^\top)^{-1}\bm{V}_\tau\),其中 \(\bm{V}\) 为由矩阵 \(\bm{A}\) 的特征值组成的范德蒙矩阵(Vandermonde matrix),\(\bm{V}_\tau = (\lambda_1^\tau, \dots, \lambda_N^\tau)^\top\)。由于状态空间中蕴含的有效历史信息完全且仅由特征值谱 \(\{\lambda_i\}_{i=1}^N\) 决定,特征向量变换矩阵与输入权重 \(\bm{B}\) 均可任意指定而不改变总记忆容量。因此,共享相同 HiPPO 特征值谱的稠密 S4 与对角化 S4D 在理论信息处理能力上完全等价。
2. 范德蒙投影矩阵驱动的梯度更新机制:揭示记忆结构对参数更新方向的调制本质 为厘清梯度优化究竟受何种内在结构约束,本文推导了线性读出层权重 \(\bm{\theta}\) 在均方误差目标下的精确梯度展开式。设理想输出序列由历史输入的线性组合表示为 \(\bm{y} = \sum_{k=0}^{T-1} \alpha_k \bm{U}_k = \hat{\bm{U}}\bm{\alpha}\),在白噪声输入假设下,协方差矩阵精简为 \(\bm{X}^\top \bm{X} = T \langle u^2 \rangle \bm{Q}\bm{V}\bm{V}^\top \bm{Q}^\top\)。进而解析求得梯度表达式为: $\(G(\bm{y}) = 2\langle u^2 \rangle B \bm{V} \left( B \bm{V}^\top \bm{\theta} - \bm{V}^\top (\bm{V}\bm{V}^\top)^{-1} \bm{V} \bm{\alpha} \right)\)$ 在该式中,矩阵 \(\bm{V}^\top (\bm{V}\bm{V}^\top)^{-1} \bm{V}\) 正是归一化的范德蒙投影矩阵,其对角线元素严格对应于各延迟阶数的记忆函数值 \(M[u_{t-\tau}]\)。这意味着该算子直接充当了针对教师目标系数向量 \(\bm{\alpha}\) 的加权衰减滤波器:第 \(\tau\) 步历史信息对参数梯度的贡献幅度,严格受制于该时间步在当前系统内的记忆存留强度。对于引入非线性激活 \(f(\cdot)\)(如 GELU/GLU)的输出层,梯度公式修正为以经过激活导数调制的目标 \(\bar{\bm{y}} = \bm{y} \odot f'(\bm{X}\bm{\theta})\) 的等价系数 \(\bar{\bm{\alpha}}\) 进行投影,其由记忆函数调控更新方向的核心数学机理完全保持一致。
3. 初始长程记忆对教师信号保真的充要条件:长程覆盖优先于短期重构精度 根据 Jaeger 记忆容量定理,离散线性系统的总记忆容量存在理论上限 \(\sum_{\tau=0}^\infty M[u_{t-\tau}] \le N\)(\(N\) 为状态维度)。若追求短延迟处的极高重构精度,远距离的记忆函数值必迅速跌落至零。梯度公式明确表明:一旦任务目标中依赖的时间步 \(\tau\) 处系统记忆衰减为 0,即 \(\bm{V}^\top (\bm{V}\bm{V}^\top)^{-1} \bm{V} \bm{\alpha} \approx \bm{0}\),反向传播计算出的梯度将完全退化为与目标标签无关的无效向量,学习过程陷入虚假停滞。反之,只要远距离记忆保持非零,哪怕由于容量守恒导致精度数值微弱,教师信号依然能够穿透归一化算子进入梯度向量,从而驱动参数朝着正确方向迭代。这一分析从数学层面彻底阐明了为何基于 HiPPO 的初始化(如 S4Dinv 与 S4Dlin)必须将特征值谱设置在单位圆内侧边缘以实现长而平缓的衰减。
4. 冻结递归权重的储备池优化策略 (RC Setting):消除过拟合并加速全局优化 在常规 SSM 训练中,循环矩阵 \(\bm{A}\) 的特征值通常随梯度一同更新。然而,反向传播缺乏主动拓展“零记忆边界”的显式机理;且在梯度驱动下改变特征值极易扰动 HiPPO 经过精心设计的正交多项式谱分布,造成部分时域记忆非故意塌缩。此外,为了维持长程记忆,特征值模长必须极其接近 1,可允许的调整空间极为狭窄,反向微调极易诱发谱半径 \(>1\) 导致数值发散。为此,本文提出将循环权重完全固定的储备池计算(Reservoir Computing, RC)设定:在初始化阶段赋予 S4Dinv 或 S4Dlin 结构化长程特征值后,训练过程中完全冻结 \(\bm{\Lambda}\) 与步长 \(\Delta\),仅更新线性投射层、前馈层及双向聚合权重。该方法直接裁剪掉了全模型中约 10% 的可训练参数,彻底隔绝了循环核在训练中的退化风险。
一个完整示例¶
考虑一个状态维度 \(N=4\)、序列长度 \(T=1024\) 的单层 SSM。若采用均匀分布随机初始化特征值(random),由于特征值模长显著小于 1,记忆函数在延迟 \(\tau > 40\) 时便衰减归零,对应的归一化范德蒙投影矩阵后部子块数值皆坍缩为 0。当给定长程关联任务(例如标签依赖于 \(\tau = 500\) 处的输入 \(u_{t-500}\),即 \(\alpha_{500} \neq 0\))时,梯度的驱动项 \(\bm{V}^\top (\bm{V}\bm{V}^\top)^{-1} \bm{V} \bm{\alpha}\) 在第 500 维的分量计算结果恒为 0,模型参数 \(\bm{\theta}\) 接收到的更新梯度仅包含噪声自相关项,根本无法感知标签中的远距离因果信号。
相反,在采用 S4Dinv 初始化时,特征值定义为 \(\lambda_m = -\frac{1}{2} + i \frac{N}{\pi} (\frac{N}{2m+1} - 1)\),经过 ZOH 离散化后特征值紧贴单位圆内缘。此时记忆函数在 \(\tau = 500\) 处依然保有微弱但大于机器精度的数值(如 \(M[u_{t-500}] \approx 0.003\))。尽管重构绝对精度低于短期窗口,但投影向量保留了非零响应,使得反向传播能够准确捕获该时间步与当前输出的关联合成项,引导读出权重 \(\bm{\theta}\) 迅速向最小二乘解 \(\bm{\theta}^* = \frac{1}{B} (\bm{V}\bm{V}^\top)^{-1}\bm{V}\bm{\alpha}\) 逼近。
损失函数 / 训练策略¶
在单层理论推导中,损失函数采用均方误差: $\(L_{\text{MSE}} = \frac{1}{T}\sum_{i=1}^T \left(f(\bm{\theta}^\top \bm{x}_i) - y_i\right)^2\)$ 在端到端深度网络(以 6~8 层堆叠的 S4D 为骨架)的实际训练中,最终任务输出采用标准交叉熵损失。优化器选用 Adam,全局基础学习率设置为 0.0005 至 0.01 不等。在传统的“可学习特征值”设置中,各层分别配置专门的步长学习率(lr_dt 为 0.0005~0.01);而在本文提出的 RC 设定中,所有循环矩阵特征值参数的学习率直接置 0,模型完全依赖各层的双向线性混合层与非线性激活(GELU、GLU)调整特征表达,从而在参数规模缩减 10% 的同时达成更高鲁棒性。
实验关键数据¶
主实验¶
论文在极具挑战性的长序列基准 Long Range Arena (LRA) 上全面评测了不同特征值配置与训练策略。各任务使用 6 个不同的随机种子运行,统计测试准确率的区间(最大值与最小值),并与官方全量可微 S4 基线进行横向对比。
| 任务 | 序列长度 | S4 官方基线 | 可学习 S4Dinv | 可学习 S4Dlin | RC S4Dinv (本文) | RC S4Dlin (本文) |
|---|---|---|---|---|---|---|
| Listops | 2,000 | 0.595 | 0.594 / 0.589 | 0.598 / 0.587 | 0.585 / 0.570 | 0.593 / 0.577 |
| Text | 4,000 | 0.865 | 0.853 / 0.839 | 0.848 / 0.830 | 0.857 / 0.846 | 0.850 / 0.843 |
| Retrieval | 4,000 | 0.910 | 0.863 / 0.760 | 0.843 / 0.750 | 0.906 / 0.901 | 0.903 / 0.898 |
| Image | 1,024 | 0.885 | 0.864 / 0.859 | 0.868 / 0.860 | 0.873 / 0.865 | 0.877 / 0.871 |
| Pathfinder | 1,024 | 0.940 | 0.910 / 0.877 | 0.896 / 0.865 | 0.899 / 0.889 | 0.892 / 0.872 |
| PathX | 16,384 | 0.960 | 0.922 / 0.903 | 0.926 / 0.912 | 0.932 / 0.923 | 0.933 / 0.922 |
(注:加粗表示 RC 设定超越对应特征值可学习设定的结果。)
消融实验¶
为了验证初始长程结构与特征值冻结策略的因果关系,论文对非结构化特征值(均匀随机 random、实数线性 linear、共轭步进 step)在特征值可学习与 RC 两种模式下的表现进行了严格对照消融。
| 特征值类型 | 学习模式 | Retrieval (max / min) | Image (max / min) | Pathfinder (max / min) | PathX (max / min) | 记忆函数有效长度 |
|---|---|---|---|---|---|---|
| random | 可学习 | 0.903 / 0.899 | 0.820 / 0.815 | 0.870 / 0.846 | 0.877 / 0.841 | 中等(训练后轻微拓展) |
| random | RC (冻结) | 0.901 / 0.896 | 0.733 / 0.698 | 0.824 / 0.777 | 0.820 / 0.773 | 短(受限于初始随机衰减) |
| linear | 可学习 | 0.899 / 0.894 | 0.782 / 0.768 | 0.797 / 0.754 | 0.820 / 0.503 | 短(仅实轴实根分布) |
| linear | RC (冻结) | 0.894 / 0.889 | 0.686 / 0.634 | 0.769 / 0.506 | 0.669 / 0.504 | 极短(高阶奇异值迅速截断) |
| step | 可学习 | 0.883 / 0.841 | 0.789 / 0.779 | 0.858 / 0.824 | 0.840 / 0.824 | 中等(复平面线性延展) |
| step | RC (冻结) | 0.903 / 0.897 | 0.745 / 0.740 | 0.832 / 0.793 | 0.813 / 0.504 | 短(缺乏正交多项式基) |
| S4Dinv | 可学习 | 0.863 / 0.760 | 0.864 / 0.859 | 0.910 / 0.877 | 0.922 / 0.903 | 极长(理论最优多项式谱) |
| S4Dinv | RC (冻结) | 0.906 / 0.901 | 0.873 / 0.865 | 0.899 / 0.889 | 0.932 / 0.923 | 极长且稳定无畸变 |
关键发现¶
- RC 设定显著抑制过拟合并提高测试稳定性:在 Retrieval 任务上,可学习 S4Dinv 的最小准确率仅为 0.760(最大 0.863),跨种子波动达 10% 以上;而冻结特征值的 RC S4Dinv 准确率跃升至 0.906,且最低准确率稳定在 0.901。在长达 16,384 步的 PathX 极限任务中,RC 设定同样在均值与上限上全面压制可学习基线。
- 训练过程难以无中生有地拓展记忆:监测非结构化特征值(random、linear、step)在训练前后的逐层最大记忆函数(Supremum MF)发现,即使经过数百轮次反向传播,其特征值自适应后扩展的记忆长度仍旧显著逊于 S4D 初始状态。这从实验上印证了理论断言——梯度下降缺乏自发拓展零记忆区域的机制。
- 收敛速度大幅提升:在 Image、Retrieval 与 PathX 任务中,RC 设定下的测试损失下降斜率显著陡峭于可学习模型,在更早的轮次即达到低稳态,证明移除循环矩阵的低效微调有助于优化器聚焦于前馈投影权重的快速对齐。
亮点与洞察¶
- 记忆函数即梯度滤波器:首次在数学上将循环神经网络的状态记忆衰减函数(MF)与参数更新梯度的投影核统一起来,证明了范德蒙投影算子是连接动态系统理论与深度学习反向传播的关键桥梁。
- 颠覆特征值必须端到端微调的思维定势:揭示了只要初始特征值具备长程覆盖,微调循环矩阵反而容易引入参数冗余与长程正交基的退化;冻结特征值的 RC 策略不仅省去 10% 参数量,更成为兼顾泛化性与稳定性的强正则手段。
- 解密 Transformers 与 SSM 的容量差异根源:在附录扩展讨论中,作者将注意力机制映射为仅含单个特征值 \(\lambda=1\)(节点数 \(N=1\))的特殊线性 RNN;由于所有注意力头共享单位特征值,其范德蒙矩阵退化为秩 1,从而无法在单序列层内形成 SSM 那样正交完备的高阶信息处理容量。
局限与展望¶
- 理论推导基于线性循环与单层白噪声假设:推导假定输入为独立同分布白噪声且单层读出为均方误差目标,尽管多层网络与交叉熵任务在经验上完全吻合,但更严格的深度非线性堆叠耦合理论仍有待形式化完善。
- 极端短序列或局部模式任务不宜使用 RC 冻结:若特定任务完全仅依赖局部几步历史信息(如超短序列控制),固定长程衰减的 HiPPO 特征值会白白浪费有限的记忆容量,此时主动缩短 MF 的特征值微调仍具优势。
- 推广至时变选择性模型(如 Mamba)的技术挑战:Mamba 引入了由输入驱动的时变步长与选择性矩阵,属于非线性时变系统,其严格的输入处理能力分析需进一步拓展到高阶非线性信息处理容量(IPC)框架。
相关工作与启发¶
- vs HiPPO & S4 (Gu et al., 2020a, 2022b):HiPPO 纯粹从连续函数逼近论出发证明勒贝格测度下的最优记忆,S4 侧重硬件高效的快速卷积实现;本文则从学习动力学与梯度流动角度,首度解释了为什么 HiPPO 初始化对后续梯度更新具有不可替代的导向作用。
- vs 循环网络梯度消失/爆炸理论 (Zucchet & Orvieto, 2024):传统 BPTT 理论聚焦于深层时间步累积相乘导致的发散;本文所揭示的问题发生在单层序列到序列模型内部,属于“因记忆衰减至零而导致目标标签投影项消失”,二者机理截然不同。
- vs 经典储备池计算 (Jaeger, 2001b; Nakajima & Fischer, 2021):经典 RC 通常随机固定非线性循环核,难以稳定求解极长依赖任务;本文结合现代 SSM 证明,以最优多项式谱初始化的确定性线性循环核在 RC 设定下能够兼获理论最优长程容量与极佳优化效率。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立记忆函数与梯度更新方向的解析等价性,从第一性原理统一了 S4/S4D 并提出 RC 训练范式。
- 实验充分度: ⭐⭐⭐⭐☆ 在完整 LRA 基准上进行了多随机种子(6 组)的严格对比,消融实验详实,充分覆盖了逐层 MF 动态与损失轨迹。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽,附录展开详尽,理论直觉与实验证据契合度极高。
- 价值: ⭐⭐⭐⭐⭐ 为状态空间模型及现代循环网络的初始化设计、结构精简与训练策略提供了不可多得的理论基石。