用于动力系统重构的循环神经网络的时间并行训练

# 并行时间训练递归神经网络用于动力系统重建 来源：https://arxiv.org/html/2605.12683 Florian Hess¹,² & Florian Götz¹,³ & Daniel Durstewitz¹,²,⁴ ¹德国曼海姆中央精神卫生研究所理论神经科学系 ²德国海德堡大学物理与天文学院 ³德国海德堡大学数学与计算机科学学院 ⁴德国海德堡大学科学计算跨学科中心 (IWR) ###### 摘要 从数据中重建非线性动力系统（DSR）是科学与工程中的一项基本挑战，但其本质上依赖于序贯模型。近年来，序贯模型的突破性进展产生了沿序列长度T实现计算并行化的算法，达到对数时间复杂度O(log T)。由于经典的时间反向传播算法线性运行复杂度O(T)实际限制了序列长度，这为DSR开辟了新途径。本文针对该任务研究了两类重要的并行时间算法，它们均以并行关联扫描为核心计算原语。第一类包括线性但非自治动力学及非线性读出的模型，例如现代状态空间模型（SSM）；第二类为通用非线性模型，可通过DEER框架实现并行化。我们发现，第一类模型的线性训练时间循环施加了局限性，常常阻碍学习精确的非线性动力学。为解决此问题，我们在更通用的非线性框架内，用广义教师强制（GTF）增强DEER，这是一种新变体，可确保在任意序列长度上稳定有效地学习非线性动力学。利用GTF-DEER，我们研究了在极长序列（T > 10⁴）上训练对DSR的益处。我们的结果表明，若数据具有长时间尺度，访问此类长轨迹能显著改善DSR。本工作将GTF-DEER确立为数据驱动发现的有力工具，并强调了长序列学习在建模复杂动力学系统中的巨大未被开发潜力。 ## 1 引言 理解和预测复杂非线性系统的行为——从神经回路和气候动力学到流体流动和生态网络——是自然科学与工程科学的核心目标 (Bhalla & Iyengar, 1999; Kalnay, 2003; Tziperman et al., 1997; Izhikevich, 2007)。一个尤为雄心勃勃的目标是从观测时间序列中直接*重建*底层动力系统（DS），即所谓的动力系统重建（DSR）。除了短期预测外，DSR要求学习的模型能忠实地再现真实系统的长期统计与几何特性，如吸引子几何、功率谱和李雅普诺夫指数。DS定义的核心是流算子，它提供了DS随时间演化的递归规则 (Strogatz, 2024; Katok & Hasselblatt, 1995)。近似它的DSR方法本质上是递归的 (Goring et al., 2024; Mikhaeil et al., 2022)，例如RNN，通常通过时间反向传播（BPTT; Werbos, 1990）训练。尽管BPTT的一些已知病态——最显著的是混沌动力学下的梯度爆炸 (Mikhaeil et al., 2022)——可以通过控制论训练算法（如稀疏和广义教师强制 STF/GTF）成功缓解 (Mikhaeil et al., 2022; Hess et al., 2023)，但BPTT的计算成本在序列长度T上仍然是线性的，这使得在具有长固有时间尺度的问题上训练成本过高。事实上，DSR应用历史上仅限于适中的序列长度以保持训练可行性 (Brenner et al., 2022; Hess et al., 2023; Vlachas et al., 2024)，留下了“从更长序列中学习是否对重建质量有益”这一开放问题。 近期并行时间序列建模的进展提供了解决途径，我们考察了两种常见范式在DSR中的性能。具有非线性读出的线性训练时间循环（如现代SSM所实现的）通过线性扫描轻松实现并行化，并通过构造避免混沌引起的梯度爆炸 (Mikhaeil et al., 2022; Orvieto et al., 2023; Zucchet et al., 2024)。然而，通过揭示这些线性SSM与非线性RNN之间的*对偶性*，我们表明这种便利性是有代价的：训练中使用的常见对角参数化线性循环施加了结构限制，常常阻止模型学习精确的非线性动力学。此外，其训练遭受暴露偏差，降低了测试时自回归展开的质量 (Bengio et al., 2015; Ranzato et al., 2016)。 第二种范式——通过DEER并行化的通用非线性RNN (Lim et al., 2024)——原则上更适合DSR，但在混沌数据上朴素应用会失败，因为驱动DEER牛顿更新的雅可比乘积在底层动力学具有正李雅普诺夫指数时会发散 (Mikhaeil et al., 2022; Gonzalez et al., 2025)。我们的主要方法学贡献是通过将DEER与广义教师强制（GTF）结合来解决这一矛盾 (Hess et al., 2023)。得到的算法GTF-DEER继承了DEER的长序列可扩展性，同时保留了GTF的混沌驯服特性：GTF在训练期间将模型转变为稳定DS，使得GTF-DEER享有Gonzalez等人 (2025) 建立的平均情况O((log T)²)缩放，即便底层系统是混沌的。经验上，GTF-DEER相比序贯训练实现了高达870倍的速度提升，同时匹配或改善了重建质量。配备GTF-DEER后，我们重新审视核心实证问题：*在更长序列上训练是否确实改善了DSR？*利用在长度T > 10⁴的轨迹上稳定训练的能力，我们发现如果数据中存在长时间尺度，访问长轨迹使DSR模型能有效捕获这些尺度，在长期统计量上带来显著增益，这是等效线性SSM模型无法匹敌的。综合来看，我们的结果将GTF-DEER确立为序贯训练的直接替代方案，同时强调了长序列训练作为DSR一个尚未充分利用的杠杆。 ## 2 相关工作 ##### 动力系统重建 用于DSR的数据驱动方法大致分为两类：第一类近似支配观测动力学的*向量场*，假设过程由微分方程控制。稀疏非线性动力学识别（SINDy）及其变体 (Brunton et al., 2016; Brunton et al., 2019; Champion et al., 2019; Messenger & Bortz, 2021; Cortiella et al., 2021) 在物理科学中尤为流行，通过最小二乘回归实现快速训练，但依赖预定义函数库，在含噪、非平稳和部分观测的实证数据中效果不佳。神经常微分/偏微分方程方法 (Chen et al., 2018; Karlsson & Svanberg, 2019; Alvarez et al., 2020; Ko et al., 2023; Aka et al., 2025) 具有通用逼近能力，并可结合物理先验 (Raissi et al., 2019; Li et al., 2022)，但实际训练困难。第二类更大类的模型通过基于神经网络的黑箱通用逼近器直接近似*流算子*，包括神经算子 (Li et al., 2020; Lu et al., 2021)、库普曼算子 (Lusch et al., 2018; Otto & Rowley, 2019; Brunton et al., 2021; Naiman & Otto, 2021; Azencot et al., 2020; Wang et al., 2022) 及其混合体、储层计算 (Pathak et al., 2017; Pathak et al., 2018; Verzelli et al., 2021; Platt et al., 2023; Patel et al., 2023; Gauthier et al., 2021) 以及通过BPTT训练的RNN (Vlachas et al., 2018; Vlachas et al., 2024; Brenner et al., 2022; Hess et al., 2023; Rusch & Mishra, 2022; Brenner et al., 2024)，通常伴有专门的优化控制技术以减轻混沌下的梯度爆炸 (Mikhaeil et al., 2022; Brenner et al., 2022; Hess et al., 2023; Sağtekin et al., 2025)。后一种方法尤其在一系列基准系统上达到最先进性能，甚至在挑战性的实证系统上也表现良好 (Hess et al., 2023; Volkmann et al., 2024; Brenner et al., 2024)，导致其被采纳为DSR基础模型的骨干 (Hemmer et al., 2025)。 ##### 高效序列建模 序列建模的核心挑战是稳定捕获和检索长期依赖关系 (Bengio et al., 1994; Hochreiter & Schmidhuber, 1997; Gu et al., 2022; Gu & Dao, 2024; Orvieto et al., 2023; Zucchet et al., 2024)。使用自回归模型如SSM或RNN时，一个实际瓶颈是BPTT的O(T)运行时间 (Werbos, 1990)。序列模型领域的最新发展通过在并行加速器（如GPU或TPU）上使用并行关联扫描 (Blelloch, 1990; Martin & Cundy, 2018; Smith et al., 2023) 将本质上的序贯操作并行化。并行扫描在O(log T)时间内评估*线性*循环，这常被用作现代SSM的核心模块 (Orvieto et al., 2023; Smith et al., 2023; Gu & Dao, 2024)。此外，Lim等人 (2024) 引入DEER框架，通过将前向传播重新表述为不动点迭代问题来并行化通用*非线性*序列模型，每次迭代可用并行扫描求解。尽管最坏情况运行复杂度为O(T log T) (Gonzalez et al., 2024)，但实践中对于呈现收缩动力学的模型，平均情况缩放为O((log T)²) (Gonzalez et al., 2025)，引发了非线性RNN在长期序列建模中的复兴。虽然DEER在多种ML问题中展示了有前景的性能，但由于评估混沌RNN时保证的最坏情况缩放，其在DSR领域的应用仍不实用。 ## 3 理论基础 ### 3.1 自回归模型的动力系统重建 给定来自某个底层物理过程的观测时间序列数据 \(\bm{X} \in \mathbb{R}^{T_{\mathrm{obs}} \times N}\)，DSR旨在学习一个生成模型，既能进行精确的短期预测，又能再现观测系统的长期行为。为解决此任务，我们考虑参数化的状态空间模型，形式为：
\[
\bm{z}_t = F_{\boldsymbol{\theta}}(\bm{z}_{t-1}, \bm{x}_{t-1}, \bm{s}_t), \quad \bm{\hat{x}}_t = G_{\boldsymbol{\psi}}(\bm{z}_t), \tag{1}
\]
其中 \(\boldsymbol{\theta}\) 和 \(\boldsymbol{\psi}\) 是参数向量，\(\bm{z}_t\) 是 M 维状态向量，\(\bm{\hat{x}}_t\) 是 N 维预测观测，\(\bm{x}_{t-1}\) 是训练中可选的教学信号，\(\bm{s}_t\) 是 K 维可选外部输入。\(F_{\boldsymbol{\theta}}\) 是离散时间通用DS，建模潜过程，通过 \(G_{\boldsymbol{\psi}}\) 与观测（数据）耦合。训练目标是学习 \(\{\boldsymbol{\theta},\boldsymbol{\psi}\}\) 使得SSM近似底层DS的流算子，训练后我们有：
\[
\bm{x}_t \approx G_{\boldsymbol{\psi}}(F_{\boldsymbol{\theta}}(\dots F_{\boldsymbol{\theta}}(F_{\boldsymbol{\theta}}(\bm{z}_0, \bm{s}_1), \bm{s}_2) \dots, \bm{s}_t)) \eqqcolon G_{\boldsymbol{\psi}}(F_{\boldsymbol{\theta}}^{\circ t}(\bm{z}_0, \bm{s}_{1:t})). \tag{2}
\]
在本工作中，我们将研究方程 (1) 的两种一般参数化形式。

#### 3.1.1 线性训练时间循环
在一种设置中，\(F_{\boldsymbol{\theta}}\) 严格线性，这将捕获数据中非线性的负担转移到非线性观测函数 \(G_{\boldsymbol{\psi}}\) 上，该函数在测试时反馈回系统：
\[
\bm{z}_t = \bm{A} \bm{z}_{t-1} + \bm{U} \bm{x}_{t-1} + \bm{C} \bm{s}_t + \bm{h}, \quad \hat{\bm{x}}_t = \bm{B} \phi(\bm{V} \bm{z}_t + \bm{b}), \tag{3}
\]
其中 \(\bm{A} \in \mathbb{R}^{M \times M}\)（通常选为对角阵），\(\bm{U} \in \mathbb{R}^{M \times N}\)，\(\bm{C} \in \mathbb{R}^{M \times K}\)，\(\bm{h} \in \mathbb{R}^M\)，\(\bm{B} \in \mathbb{R}^{N \times L}\)，\(\bm{V} \in \mathbb{R}^{L \times M}\)，\(\bm{b} \in \mathbb{R}^L\)，\(\phi\) 是非线性函数如 ReLU(\(\bm{z}\)) = max(0, \(\bm{z}\))。这基本上是现代SSM的架构设置，其中线性循环后跟非线性逐点序列变换 (Orvieto et al., 2023; Orvieto et al., 2024; Smith et al., 2023; Gu & Dao, 2024)，通常由MLP实现。在我们的情况下，MLP是简单的单隐层神经网络，隐层大小为 L。在下文中，我们将方程 (3) 定义的模型称为 “LSSM”。一个关键见解是，尽管方程 (3) 中的循环在训练期间构成一个线性、非自治DS，但模型在评估时可通过将教学信号 \(\bm{x}_{t-1}\) 替换为预测 \(\hat{\bm{x}}_{t-1}\) 来产生非线性动力学。实际上，在训练后的轨迹生成中，方程 (3) 的递归变为：
\[
\bm{z}_t = \bm{A} \bm{z}_{t-1} + \overline{\bm{W}} \phi(\bm{V} \bm{z}_{t-1} + \bm{b}) + \bm{C} \bm{s}_t + \bm{h}, \tag{4}
\]
其中 \(\bm{U} \bm{B} =: \overline{\bm{W}} \in \mathbb{R}^{M \times L}\) 且 \(\operatorname{rank}(\overline{\bm{W}}) \leq \min(N, M, L)\)。这使得SSM在自回归生成过程中能够表现出固有非线性特征，如混沌动力学和多稳态。虽然方程 (3) 的主要特征是可以使用并行扫描在对数时间内计算前向传播，但另一个见解是线性循环不会遭受混沌引起的梯度爆炸 (Mikhaeil et al., 2022)，因此通过设计缓解了相关的训练不稳定性 (Mikhaeil et al., 2022; Orvieto et al., 2023; Zucchet et al., 2024)；详见附录A (https://arxiv.org/html/2605.12683#A1)。LSSM的训练使用来自序列建模领域的教师强制变体，其中上一个时间步的地面真实数据 \(\bm{x}_{t-1}\) 通过专用输入层 (\(\bm{U}\)) 输入 (Bengio et al., 2015; Goodfellow et al., 2016)；通用训练算法见 Alg.1 (https://arxiv.org/html/2605.12683#A1)。