TRIE: 随机PDE代理模型的评估框架

# TRIE: 一个用于随机偏微分方程替代模型的评估框架
来源：https://arxiv.org/html/2607.00196
Bharat Srikishan¹ & Javier E. Santos² & Nikhil Muralidhar¹ & Charles D. Young²∗
¹史蒂文斯理工学院 ²洛斯阿拉莫斯国家实验室

###### 摘要

许多科学系统由于随机强迫、未解析的自由度或不完美的观测而表现出不确定性，这使得可靠的替代预测从根本上必须是分布式的，而非逐点式的。对于此类系统，确定性神经替代模型无法捕捉统计量度和预测不确定性。我们引入了 TRIE，一个用于随机偏微分方程替代模型的评估框架，旨在检验模型是否能重现不变测度、提供可信的预测不确定性以及扩展至高效的概率生成。我们在两种平稳混沌空间扩展的随机偏微分方程（随机 Kuramoto–Sivashinsky 方程和随机 Kolmogorov 流）上，针对 11 个参数值，展示了 TRIE 框架的效果。我们的评估表明，标准逐点训练的神经替代模型能产生看似合理的短时 rollout，但在匹配长时统计结构方面却失败。近似不确定性方法，如 Monte Carlo dropout 和异方差高斯似然，能够产生随机预测，但通常在时间和空间不确定性诊断下校准不良且过于自信。综合这些标准，生成模型表现最为一致，能准确捕捉不变测度统计量，并在所有报告的概率设置中实现最低的连续排名概率分数。最后，我们展示了具有自动维度发现的隐式生成模型在保持大部分统计保真度的同时，将 Kolmogorov 推理时间减少约12×。我们发布了我们的代码¹和数据，以支持随机偏微分方程预测模型的可重复评估。

## 1 引言

现实世界的数据和过程不可避免地具有随机性，导致测量和预测中的不确定性。在某些情况下，例如低高斯测量噪声，影响很小，确定性模型对于某些应用是可靠的。通常，观测状态不包含关于随机过程的足够信息。分子系统会经历看似随机的热涨落，这可以通过快速原子尺度的自由度进行确定性解释，但后者无法被可操作地测量或计算。类似的效应在物理、化学和生物系统中很常见[19,40]，其中未知的随机强迫与动力学耦合。远离平衡态时，随机效应和确定性效应耦合驱动多样化的时空模式和转变[7]。这些概念延伸到了天气[24]、金融[41]和电网[3]中。

这里我们专注于随机偏微分方程，作为随机强迫下空间扩展系统的模型。在非线性参数区域，时间和长度尺度通常跨越多个数量级耦合，需要计算成本高昂的精细离散化和专门的数值求解器。确定性替代模型可以显著降低计算开销，同时保持短时精度[31]，并且已有改进以提升稳定性[26,32]、多步 rollout[4] 和物理约束[29]。然而，它们固有的局限性在于无法产生预测分布。近似不确定性方法，如 Monte Carlo dropout[12]、集成[23]和卡尔曼滤波[17]，可以提供不确定性估计，但可能校准不良或计算成本高昂[8,11]。最近的概率生成模型，包括基于扩散的预测器和随机插值，为分布预测提供了更直接的途径[37,42,34,13,36,2,5,1]。

我们引入了 TRustworthiness（可信度）、Invariance（不变性）和 Efficiency（效率）标准，即 TRIE，一个用于随机偏微分方程替代模型的评估框架。TRIE 检验模型是否：(i) 重现长时不变测度，(ii) 提供可信的预测不确定性，以及 (iii) 扩展至高效的概率生成。我们使用导数场联合密度和谱内容评估不变测度保真度。我们使用连续排名概率分数和空间不确定性诊断评估概率可信度。最后，我们通过壁钟推理时间评估生成扩展性，因为高质量的概率采样器可能比确定性替代模型慢得多。为了解决推理成本，我们研究了具有隐式秩最小化自编码器的隐式随机预测[20,43]。

我们在 11 个参数值下，对一维随机 Kuramoto–Sivashinsky 方程和二维随机 Kolmogorov 流展示了 TRIE 框架的应用。在确定性、近似概率、基于流和随机插值替代模型中，我们发现逐点训练的模型通常无法重现不变测度，近似不确定性方法可能过度自信或校准不良。随机插值提供了最一致的表现。隐式随机插值在保留大部分统计保真度的同时，将 Kolmogorov 推理时间减少约12×。

我们的主要贡献是：

- • 我们引入了 TRIE，一个基于不变测度保真度、概率可信度和生成规模的 SPDE 替代模型评估框架。
- • 我们展示不变测度指标和 CRPS 能够揭示逐点评估所隐藏的失败，其中分布桥接替代模型实现了最一致的表现。
- • 我们首次将自动维度发现[20,43]应用于 SPDE，在保持 TRIE 可靠性的同时，以最少的参数调整加速推理。
- • 我们发布了用于两个复杂 SPDE 系统的仿真、训练、推理和评估代码。

见图 1 标题：可信度：当预测分布偏离真实密度时，带有近似或无不确定性量化的神经替代模型通常过于自信。连续排名概率分数和空间不确定性指标能够捕捉这些常被忽视的细微差别。**不变测度**是衡量替代模型是否已学习到真实动力系统还是仅拟合了短时动态的可靠指标。**效率**：根据推理壁钟时间评估替代模型在成本-精度权衡中的表现至关重要。
## 2 相关工作

**基准测试与学习到的随机动力学。** 神经 PDE 基准测试通常强调短时预测误差、rollout 稳定性或架构比较。最近的工作已将此方向扩展到随机系统。神经 SPDE 学习基于初始状态和驱动噪声实现的解算子[38]，而 SPDEBench 评估正则和奇异 SPDE 的模型，重点在于数值格式、噪声采样、重整化以及与 FNO、NSPDE 和 DLR-Net[25] 等模型的比较。TRIE 是互补的：我们并非主要关注数值离散化或架构级别的比较，而是评估随机预测器是否重现不变测度、提供校准且锐利的预测分布，以及是否高效生成样本。这一区分对于混沌随机系统至关重要，因为逐点误差可能掩盖长时间统计结构或不确定性校准方面的失败。

**概率预测器与不确定性量化。** 概率生成模型为建模条件预测分布提供了自然方式，而非点估计。流匹配通过沿条件概率路径回归向量场来训练连续归一化流[27]，而随机插值通过概率分布之间的连续时间桥统一了基于流和基于扩散的生成模型[1]。我们在共享的一步训练和自回归 rollout 协议下，将这些方法作为单模型条件预测器进行评估。近似不确定性方法，如 Monte Carlo dropout 和异方差高斯似然，在计算上具有吸引力，而深度集成是校准不确定性估计的强有力替代方案[23]。通过 TRIE，我们定量评估这些样本在校准性方面的表现，同时结合长时 rollout 稳定性和准确性。

**长时统计与 Rollout 校正。** 对于混沌耗散系统，长时间保真度通常是统计性质的而非轨迹性质的：rollout 可能逐点发散，但仍在正确吸引子上采样。一些方法通过在自回归推理中加入去噪修正和投影来改善长时行为[28,32]。这些方法与我们的不变测度标准密切相关，但它们通过在推理时添加修正步骤来改变预测过程。相比之下，TRIE 评估的是原始的步替代模型本身。

**降阶与隐式空间生成。** 迭代概率采样器在全空间分辨率下可能计算成本高昂，这促使了降阶方法的发展。经典投影方法利用耗散系统中的低维结构[10]，而现代基于自编码器的方法学习用于动力学和物理仿真的非线性隐式坐标[22,43]。最近的隐式生成模型表明，在压缩坐标中采样可以显著降低推理成本[36]。TRIE 通过测试隐式随机插值能否在保持不变测度保真度的同时减少生成时间来评估这一成本。我们考虑隐式秩最小化自编码器[20]，该模型在确定性混沌的自动隐式大小发现中表现出色[43]。我们将其扩展到随机吸引子上的随机系统[6]，展示了在自动维度发现方面的相同成功。

## 3 方法

### 3.1 TRIE：可信度、不变性与效率

我们考虑控制状态x(t) 演化的随机偏微分方程，形式为：

∂ₜx = f(x,t) + η(x,t), t∈[0, T], x∈R^d   (1)

其中 f 表示确定性动力学，η 表示随机强迫。本文我们关注与状态无关的噪声过程 η(t)，但评估框架更广泛地适用于与状态相关的随机系统。由于在预测时刻未观测到强迫实现，替代模型无法无限期地追踪单个参考轨迹。相反，对于混沌和随机系统，长时间的 rollout 应在瞬态衰减后重现动力学的统计结构。

因此，我们通过三个互补的问题来评估 SPDE 替代模型。首先，**神经替代模型能否准确捕捉不变测度？** 这个问题至关重要，因为替代模型可能产生看似合理的短时预测，却收敛到错误的随机吸引子或分配不正确的概率质量。其次，**替代模型的预测有多可信？** 在科学应用中，预测用于指导分析、设计和下游决策。因此，一个有用的模型必须提供既校准又锐利的不确定性估计，而不仅仅是产生多样化的样本。第三，**概率替代模型能否被扩展以实现更快的生成？** 生成式预测器可以提高统计保真度，但其迭代采样过程在全空间分辨率下可能成本高昂。因此，实际的 SPDE 预测不仅需要衡量准确性，还需要衡量和改善推理时间成本。下面我们定义指标和方法来回答这些问题。

### 3.2 不变测度

对于耗散随机系统，在初始瞬态时间 τ 之后，轨迹集中在随机吸引子附近[6,35]，并在状态空间上诱导出一个不变测度。该测度描述了在确定性动力学和随机强迫共同作用下，SPDE 所访问状态的长时间分布。我们通过比较从真实模拟和模型生成轨迹计算出的经验不变测度统计量，来评估自回归替代模型 rollout 是否重现了该分布。

对于随机 Kuramoto–Sivashinsky 方程，我们使用一阶和二阶空间导数的联合概率密度，即 u_x 和 u_xx。形式上，如果 μ 表示 KS 状态 u 上的不变测度，则导数场联合密度为：

ρ_KS(a,b) = E_{u~μ} [ (1/L) ∫₀ᴸ δ(a - u_x(x)) δ(b - u_xx(x)) dx ]   (2)

该统计量总结了 KS rollout 的局部相空间几何结构。它对噪声水平的变化、替代模型收敛到错误吸引子、变异消失或扭曲斜率和曲率关系等失败情况非常敏感。我们使用 Wasserstein 距离比较真实和替代模型 rollout 的经验联合密度。

对于随机 Kolmogorov 流，我们使用时间平均的涡量谱。设 ω 为涡量场，̂ω(q) 为其在波矢量 q 处的傅里叶系数。壳平均的涡量谱为：

E_Ω(k) = E_{ω~μ} [ Σ_{q: k ≤ ‖q‖ < k+1} |̂ω(q)|² ]   (3)