面向自回归动力系统预测的可扩展单步生成建模

# 面向自回归动力系统预测的可扩展单步生成模型

来源: https://arxiv.org/html/2605.05540

Tianyue Yang  
The Center for Computational Science  
University College London  

& Xiao Xue  
The Center for Computational Science  
University College London  
[email protected]

###### 摘要

高维物理动力学的快速代理建模不仅需要较低的短期误差：实用的模型必须能够高效地进行滚动预测（roll out），同时保持长轨迹的统计结构。神经算子虽然能提供廉价的自回归预测，但在湍流 regime 中容易发生漂移；而滚动扩散模型和潜在空间生成代理虽然能够表征随机转移过程，但代价是需要多步去噪、噪声调度设计或辅助压缩模型。我们提出了 MeanFlow Long-term Invariant Spatiotemporal Consistency Autoregressive Models（MeLISA），这是一种基于像素空间 MeanFlow 的无潜在空间自回归生成代理。MeLISA 定义了一个块级随机转移核，通过单次模型评估即可生成每个预测块，从而在推理阶段避免了潜在编码器和迭代式扩散求解器的使用。为了稳定长视界滚动预测，MeLISA 结合了*窗口一致性 MeanFlow*（Window-Consistency MeanFlow）目标——从部分观测的时间窗口中学习条件时空生成，以及*时间增量一致性*（Time Increment Consistency）损失——约束多滞后有限增量并针对时间相关性结构。我们在两个高分辨率基准上评估了 MeLISA，分别使用紧凑型 UNet 和可扩展的 DiT 骨干网络：扩展的 2D Kolmogorov 流动（$256\times 256$）和湍流通道流切片（$192\times 192$）。MeLISA 在短期预测准确性和长视界统计指标（包括能量谱、湍流动能以及与混合率相关的动力学）上均优于神经算子基线，同时实现了与神经算子相当甚至更快的推理速度。紧凑型 3.7–5.7M 参数变体已经展现了强大的参数效率，而 DiT 变体则提供了可扩展至 1.5 亿参数的路径。总体而言，MeLISA 同时提升了滚动预测效率和长视界统计准确性。

## 1 引言

对复杂动力系统既准确又高效的模拟仍然是计算物理学中的一个核心挑战。这些系统通常由非线性偏微分方程（PDEs）[12] 控制，获得精确解往往是不切实际的。因此，人们开发了各种数值技术，包括直接数值模拟（DNS）[33, 66, 13, 60]。虽然这些方法可以通过解析细尺度结构来实现高保真度，但往往伴随着巨大的计算成本，使得它们在许多实际场景中难以应用。

为了减轻这种计算负担，人们提出了各种降低保真度的方法，以牺牲物理分辨率为代价来换取效率，包括雷诺平均 Navier-Stokes（RANS）[1] 和大涡模拟（LES）[57]。然而，在需要精细空间分辨率的 regime 中（例如，在必须解析陡峭梯度的固体边界附近），这些方法仍然可能成本高昂。

受这些局限性的驱动，并得益于深度学习的最新进展，数据驱动的代理模型作为一种有希望的替代方案应运而生。这些代理旨在直接从数据中逼近底层动力学，并已成功地应用于广泛的科学领域，包括天气预报 [15, 4]、量子化学 [69, 73]、材料科学 [48] 和流体动力学 [76]。

确定性神经算子已成为用于 PDE 控制和时空预测问题的数据驱动代理的广泛使用类别 [35, 45, 59, 43]。当自回归使用时，它们将有限的上下文窗口映射到未来状态，并能高效地进行滚动预测。然而，长视界自回归预测仍然具有挑战性：微小的一步误差会反复作为输入反馈，导致在滚动过程中误差累积和分布偏移 [50]。这个问题在湍流或混沌 regime 中尤为突出，其中高频内容或时间相关性中的微小偏差会导致轨迹级统计量（如能量谱、湍流动能与不变测度相关量）的漂移 [25, 55, 27]。

一个互补的方向是使用生成模型，包括扩散模型 [22, 65]、流匹配 [40] 以及单步变体如一致性模型 [64, 44]，这些方法将样本生成视为从基础分布到数据分布的传输过程。这些方法最近被采用为物理动力学的自回归代理 [31, 62, 6, 76, 58, 15]，并显式地对未来状态的条件分布进行建模。然而，目前它们通常：(i) 在每一步滚动时都需要多步去噪或 SDE/ODE 积分，导致推理延迟较长；(ii) 依赖逐帧渐进式噪声调度 [6, 61] 以稳定滚动；(iii) 在由辅助 VAE 或编码器诱导的潜在空间中运行 [61, 31]，引入了额外的训练和推理复杂性。

在本工作中，我们基于最近提出的像素 MeanFlow（p-MF）框架 [46]，这是一种直接在像素空间运行并避免多步求解器和潜在编码器的单步生成模型。为了将 p-MF 转变为一种自回归代理，在保持快速滚动速度的同时保留长视界统计结构，我们解决了两个具体问题：如何将单帧 MeanFlow 扩展为*窗口条件时空*生成器，以便在多帧时间上下文下使单步生成变得非平凡，以及如何在训练期间仅观测短窗口时显式地强制长视界时间相关性和混合行为？

我们的答案——*窗口一致性 MeanFlow* 和 *时间增量一致性*——共同产生了一种随机自回归代理，每个预测块只需一次函数评估（每块 1 NFE，其中 NFE 表示函数评估次数），将直接自回归滚动与生成方法的条件分布建模相结合。关于详细讨论见附录 B。

我们的主要贡献总结如下：

- 我们介绍了 MeanFlow Long-term Invariant Spatiotemporal Consistency Autoregressive Models（MeLISA），这是一种随机自回归代理，每块只需 1-NFE，将直接自回归滚动与生成方法的条件分布建模相结合，无需渐进式噪声调度或多步 SDE 求解器。
- 我们提出了*窗口一致性 MeanFlow*，这是首次将像素 MeanFlow 从单帧生成扩展到窗口条件时空转移核，其中掩码时间上下文使得单步生成预测变得非平凡。
- 我们提出了*时间增量一致性*，这是一种有限滞后正则化器，直接约束滚动预测的时间协方差和混合结构，提供点对点状态重构损失显然无法提供的长视界约束。
- MeLISA 直接在像素空间进行生成预测，分辨率高达 $256\times 256$，去除了当前基于扩散的科学代理通常依赖的 VAE / 潜在编码器 / 保真度增强组件。
- 在扩展的 2D Kolmogorov 流动（$256\times 256$）和投影的湍流通道流可观测量（$192\times 192$）上，MeLISA 在相当的滚动速度下匹配或超过了神经算子基线的短期准确性，同时显著改善了对能量谱、湍流动能以及与混合率相关的动力学的恢复。紧凑型 3.7–5.7M 参数变体已经具有参数效率，而 DiT 实例化可扩展至 1.5 亿参数。

## 2 相关工作

#### 动力系统的生成模型。

扩散模型及其潜在空间变体——包括潜在扩散模型（LDM）[62, 6, 9] 和潜在流匹配（LFM）[37, 52]——已被广泛用于动力系统预测，并取得了最先进的性能。然而，在没有蒸馏或相关加速技术的情况下，这些方法在推理时通常需要大量的采样步骤，导致巨大的计算开销。相比之下，我们的方法建立在单步生成公式之上，消除了对专用多步采样程序的需求。此外，我们直接在像素空间进行生成建模，完全避免了降维。这使得我们能够生成高达 $256\times 256$ 的高分辨率图像，而无需训练辅助潜在空间组件，如变分自编码器（VAE）或保真度增强模块 [63, 55, 39, 80]。这种设计也使我们的方法更加通用和稳健，因为它减少了实践中所需的特定于系统的架构和训练选择数量。特别是，我们模型的单步性质消除了指定渐进式噪声调度的需求，这是大多数滚动或自回归扩散方法依赖于多步采样的组件 [6, 61, 26, 53]。

#### 单步生成模型。

单步（或少步）生成模型最近获得了动力，其动机是将生成建模视为分布传输 [65]。代表性示例包括一致性模型（CM/sCM）[64, 44]、归纳矩匹配（iMM）[82] 和捷径扩散 [14]。在图像生成领域，MeanFlow（MF）家族，包括原始 MF [16]、改进的 MeanFlow（i-MF）[17] 和像素 MeanFlow（p-MF）[46]，在条件和非条件生成任务上都表现出特别强劲的性能。用于图像和视频合成的现代生成模型通常通过蒸馏大型预训练扩散模型来实现少步或单步采样。例如，一致性蒸馏（CD/sCD）[64, 44] 已成功应用于预训练图像生成器 [47] 以及视频扩散模型 [49, 71]。相比之下，在类视频数据上*从头训练*单步生成模型受到的关注相对较少；据我们所知，这一设置尚未得到系统探索，也未在物理动力系统的背景下进行研究。

## 3 背景

#### 问题设置。

令序列数据集表示为 $\mathcal{D} \in \mathbb{R}^{B \times T \times F}$，其中 $B$ 是轨迹数量，$T$ 是轨迹长度，$F$ 是特征维度。我们将离散物理时间 $\tau \in \mathbb{Z}^+$ 到 $\tau+W$ 采样的窗口表示为 $\bar{x}^{\tau:\tau+W} := (x^\tau, \ldots, x^{\tau+W-1}) \in \mathbb{R}^{W \times F}$，并在含义清晰时使用简写 $\bar{x}_W^\tau$。在概率预测中，给定输入窗口 $\bar{x}_{W_{\mathrm{in}}}^\tau$，目标是预测未来窗口 $\hat{\bar{x}}_{W_{\mathrm{out}}}^{\tau+W_{\mathrm{in}}}$，其中 $W_{\mathrm{in}}$ 和 $W_{\mathrm{out}}$ 分别表示输入和输出窗口长度，可以任意选择。此预测问题由预测分布形式化：

$$
\hat{\bar{X}}_{W_{\mathrm{out}}}^{\tau+W_{\mathrm{in}}} \sim p\!\left( \hat{\bar{X}}_{W_{\mathrm{out}}}^{\tau+W_{\mathrm{in}}} \mid \bar{x}_{W_{\mathrm{in}}}^\tau \right). \tag{1}
$$

#### 扩散模型与流匹配。

扩散模型可以被视为分布传输方法，将简单的基础分布 $p_{\mathrm{base}}$（通常为高斯分布）映射到目标数据分布 $p_{\mathrm{data}}$。在去噪扩散概率模型（DDPM）[22] 中，这种传输是通过由扩散时间 $t$ 索引的马尔可夫扩散过程实现的。神经去噪器 $D_\theta$ 被训练以预测每个时间步注入的噪声：

$$
\mathcal{L}_{\mathrm{DDPM}} = \mathbb{E}_{t,x,\epsilon} \left[ \| D_\theta(z_t, t) - \epsilon \|_2^2 \right], \tag{2}
$$

其中 $z_t$ 表示时间 $t$ 沿扩散路径的噪声样本 [22]。此正向-反向扩散过程也可以等效地通过随机微分...