Latent Matters: 突破深层状态空间模型的动力学学习瓶颈

TL;DR

深层状态空间模型 (DSSMs) 在理论上非常适合处理带有噪声的时序预测，但在实践中，单纯最大化 ELBO 往往会导致模型“偷懒”：它可能重构得很好，但动力学预测却一塌糊涂。本文提出的 EKVAE (Extended Kalman VAE) 结合了约束优化框架与改进的卡尔曼滤波，不仅解决了动力学识别不准的问题，还实现了静态位置与动态速度的完美解耦。

1. 痛点：为什么你的 ELBO 很高，预测却很差？

在学术界，训练 DSSM 的标准做法是优化变分下界 (ELBO)。然而，作者尖锐地指出，高 ELBO 指标与准确的动力学学习（System Identification）之间并不存在必然联系。

现有方法存在三大顽疾：

• 重构与压缩的失衡：模型往往过度关注单帧图像的重构（Distortion），而忽略了跨时间步的演化逻辑（Rate）。
• 打破生成模型（Broken Generative Models）：初始时刻 z1 的先验通常被简单假设为标准高斯，这与后验分布严重不匹配，导致预测起点就已偏离航线。
• RNN 的诅咒：许多 DSSM 使用 RNN 来辅助推理，但这会导致状态潜空间变成非马尔可夫（Non-Markovian）——即当前的潜变量 zt 不再包含全部状态信息，部分信息被“偷藏”在 RNN 的隐藏状态中。

2. 核心架构：EKVAE 与神经线性化

为了彻底解决上述问题，作者引入了 EKVAE。其核心直觉是：既然动力学函数是未知的，那我们就让神经网络去动态学习如何进行线性化。

神经线性化 (Neural Linearization)

传统的扩展卡尔曼滤波 (EKF) 需要对非线性函数进行一阶泰勒展开，这在深度学习中计算开销巨大。EKVAE 转而采用了一组基矩阵 (Base Matrices) 的线性组合： ${\mathbf{F}}_{\psi }({\mathbf{z}}_t,{\mathbf{u}}_t)=\sum {\alpha }^{(m)}{\mathbf{F}}^{(m)}$ 模型会根据当前状态，通过神经权重 $\alpha$ 自动合成最适合当前局部的线性转移矩阵。

图 1：EKVAE 概率图模型。红色箭头表示推理网络，其核心是结合了辅助变量 $a_t$ 的闭式贝叶斯平滑。

3. 约束优化：让训练更“硬核”

本文最显著的贡献之一是引入了约束优化 (Constrained Optimisation, CO) 框架。作者将训练目标重构为：

• 最小化 Rate（复杂度）
• 约束 Distortion（重构误差） $\le D_0$

通过拉格朗日乘子 $\lambda$，模型在训练初期专注于满足重构指标；一旦重构质量达标，$\lambda$ 自动调整，迫使模型开始攻克困难的动力学学习。这比传统的人工调整退火参数 (Annealing) 要鲁棒得多。

4. 实验见证：精准识别动力学系统

在 Pendulum（摆钟）和 Reacher（机械臂）任务中，EKVAE 展示了令人惊叹的识别能力。

图 2：在 Pendulum 图像数据上，EKVAE 能够清楚地提取出代表角速度的“桶状”潜空间结构，而传统方法往往只能学到表面的角度。

关键数据对比：

• 在 64x64 高维 RGB 图像输入下，EKVAE 对关节角度识别的 $R^2$ 达到 0.98+。
• 预测精度 (MSE) 相比传统的退火训练法提升了 10-100 倍。
• 解耦性能：EKVAE 能自动将 z 空间划分为静态特征（如位置）和动态特征（如速度）。

5. 深度洞察：为什么 RNN 成了绊脚石？

作者通过消融实验证明，KVAE 或 RSSM 等使用 RNN 处理潜变量的模型，其预测精度在平滑（Smoothing）后反而下降。原因在于 RNN 破坏了马尔可夫性，导致初始状态 z1 无法通过反向时间流获取足够的动力学信息。而 EKVAE 坚持使用纯粹的马尔可夫潜表征，确保了信息的完整。

6. 总结与启示

Latent Matters 告诉我们，在处理复杂时序数据时，数学直觉（如卡尔曼滤波的线性化）与现代优化策略（如约束优化）的结合，比单纯堆砌神经网络深度更有力。

对于致力于 Model-based RL 的研究者，EKVAE 提供了一个极其简洁的方案：通过 VHP 解决冷启动问题，通过 CO 框架保证动力学的一致性，通过 H 矩阵解耦重构目标。这不仅是一篇学术论文，更是一套可以直接落地的高维系统识别工具包。