用于守恒律的稳健基础模型：通过循环视觉转换器将上下文注入通量神经算子

# 用于守恒律的稳健基础模型：通过递归视觉 Transformer 将上下文注入通量神经算子

**来源**: https://arxiv.org/html/2605.05488

**Taeyoung Kim**
人工智能与自然科学中心
韩国高等研究院
韩国首尔 02455
[email protected]

**Joon-Hyuk Ko**${}^{1\dagger}$
人工智能与自然科学中心
韩国高等研究院
韩国首尔 02455
[email protected]

${}^{\dagger}$同等贡献。联系作者：Taeyoung Kim, Joon-Hyuk Ko.

###### 摘要

我们提出了一种架构，该架构通过基于视觉 Transformer (ViT) 的上下文注入来增强通量神经算子 (Flux NO)。Flux NO 结合了经典有限体积法 (FVM) 与神经算子。我们的模型被构建为一个超网络：它在有限的时间窗口内提取解的动态特征，使用递归视觉 Transformer 对其进行编码，并生成上下文条件化神经算子的参数。这使得模型能够在不显式访问控制方程或偏微分方程 (PDE) 系数的情况下推断并求解守恒律。实验表明，与标准神经算子相比，所提出的方法保留了 Flux NO 的稳健性、泛化能力和长期预测优势，同时在广泛的保守系统中提供了可靠的数值解，包括以前未见过的通量。我们的代码可用於 https://github.com/xx257xx/CONTEXT_FLUX_NO。

## 1 引言

基于神经网络的科学计算方法的迅速兴起已成为一个主要的研究方向，各种范式相继涌现。这种演进可以 broadly 理解为三次连续的转变。首先，物理信息神经网络 (PINN) 被提出，旨在通过直接在神经网络上优化以满足控制方程以及初始和边界条件来求解偏微分方程 (PDE) (Raissi et al., 2019)。其次，算子学习引入了一种不同的视角：与其独立求解每个 PDE 实例，不如让神经算子学习预设 PDE 族的解映射，从而能够根据输入条件直接预测正向或逆向解 (Li et al., 2021; Lu et al., 2021; Kovachki et al., 2023)。最近，受基于 Transformer 的基础模型的少样本和在上下文 (in-context) 能力启发 (Brown et al., 2020; Dosovitskiy et al., 2020)，这一观点已被扩展到科学机器学习领域，产生了 PDE 基础模型，旨在通过条件化上下文信息（如观测到的动态、方程族或域结构）来求解各类不同的 PDE (Hao et al., 2024; Herde et al., 2024; Subramanian et al., 2024)。

受这一发展脉络的启发，我们重新审视了经典的有限体积法 (FVM)，其中守恒律的演化由单元界面处的数值通量控制 (LeVeque, 2002)，并在 Flux Neural Operator (Flux NO) (Tran et al., 2024) 的精神指导下将其与神经算子相结合。基于这一公式，我们提出了一种递归 ViT 基础的上下文注入机制，将 Flux NO 提升到基础模型框架中。 resulting 模型从短解轨迹中推断底层动态，并相应地调整其数值通量算子，无需明确了解 PDE 系数或闭式通量表达式。

我们的主要贡献如下：

-   我们为参数化守恒律制定了在上下文通量学习问题，其中短的观测轨迹用于推断潜在的数值通量算子。
-   我们引入了一种上下文条件化的 Flux Neural Operator，其中递归 ViT 编码器生成紧凑的上下文代码，以此条件化有限体积通量算子。
-   我们表明，在一维守恒律基准和相关扩散型 Burgers 问题上，与通用的 PDE 基础模型基线相比，强制实行保守的通量差分更新提高了自回归稳定性和分布外 (OOD) 鲁棒性。

## 2 背景

本节回顾了我们架构的动机要素。我们要强调两点。首先，守恒律需要尊重通量差分结构的数值更新，Flux NO 编码了这种保守结构，但并非天生设计用于跨未见通量函数的在上下文适应。其次，最近的 PDE 基础模型提供了上下文条件化的适应性，但通常这样做是通过通用的预测架构，并不显式保留保守数值结构。

### 2.1 守恒律与 Flux Neural Operators

我们考虑以下形式的守恒律：

$$ \partial_t \bm{u} + \nabla \cdot \bm{F}(\bm{u}; \bm{p}) = 0, \quad (1) $$

其中 $\bm{u}(t, \bm{x}) \in \mathbb{R}^d$ 是守恒状态，$\bm{F}(\bm{u}; \bm{p})$ 是物理通量，可能由系数 $\bm{p}$ 参数化。公式 (1) 的关键结构在于时间演化由通量不平衡决定。在一维有限体积离散化中，这导致了半离散更新：

$$ \frac{d}{dt} \bar{u}_i(t) = -\frac{1}{\Delta x} \left( \hat{f}_{i+\frac{1}{2}}(t) - \hat{f}_{i-\frac{1}{2}}(t) \right), \quad (2) $$

以及时间离散化后的保守更新：

$$ \bar{u}_i^{\, n+1} = \bar{u}_i^{\, n} - \frac{\Delta t}{\Delta x} \left( \hat{f}_{i+\frac{1}{2}}^{\, n} - \hat{f}_{i-\frac{1}{2}}^{\, n} \right). \quad (3) $$

这种伸缩式通量差分结构确保在合适的边界条件下离散守恒，对于非线性双曲问题尤为重要，因为平滑解可能激波，而长期预测需要稳定的传输行为。

算子学习为近似函数空间之间的解映射提供了数据驱动的框架。如 DeepONet (Lu et al., 2021) 和 Fourier Neural Operator (Li et al., 2021) 这样的神经算子从数据中学习这些映射，并可以快速对新输入进行评估。然而，许多神经算子直接预测未来的解场，因此并不显式强制执行公式 (3) 中的保守结构。这可能导致自回归展开期间的守恒误差或不稳定的误差累积。

Flux Neural Operators 通过将神经算子与有限体积观点相结合来解决这个问题 (Kim and Kang, 2025; Kim et al., 2025)。Flux NO 不是直接预测下一个解快照，而是学习一个数值通量算子：

$$ \hat{f}_{i+\frac{1}{2}} = G_{\Theta} \left( S_{i+\frac{1}{2}}(\bm{u}^{\, n}) \right), \quad (4) $$

其中 $S_{i+\frac{1}{2}}(\bm{u}^{\, n})$ 表示界面 $i+\frac{1}{2}$ 周围的局部或非局部模板表示，而 $G_{\Theta}$ 是一个神经算子。然后将此学习到的通量代入有限体积更新以获得下一状态：

$$ \bar{u}_i^{\, n+1} = \bar{u}_i^{\, n} - \frac{\Delta t}{\Delta x} \left( G_{\Theta} \left( S_{i+\frac{1}{2}}(\bm{u}^{\, n}) \right) - G_{\Theta} \left( S_{i-\frac{1}{2}}(\bm{u}^{\, n}) \right) \right). \quad (5) $$

因此，模型受到约束，通过通量差分来演化解，赋予其与守恒律一致的归纳偏置。由于误差是通过保守残差进入而非无约束的全局预测，这种结构特别有利于稳健的长期展开和分辨率传输。

### 2.2 PDE 基础模型与上下文条件化

最近的研究已开始从单方程神经算子转向 PDE 的基础模型。目标是训练能够通过条件化上下文信息（如短观测轨迹、方程描述符、仿真元数据或类似提示的输入-输出示例）在更广泛的方程族、系数、离散化和物理机制上操作的模型。几种方法使用 Transformer 风格架构、补丁标记化、自回归序列建模或超网络条件化来启用这种跨系统适应 (Yang and Osher, 2024; Yang et al., 2025; Hao et al., 2024; Morel et al., 2025)。

这些方法提供了一种在上下文泛化的机制：单个训练好的模型可以根据观测到的任务上下文调整其行为，而无需为每个新方程实例进行显式重新训练。然而，许多 PDE 基础模型仍然是未来状态或潜在解场的通用预测器。它们的架构通常围绕序列建模或全局算子回归设计，而不是针对双曲守恒律特有的保守数值结构。因此，它们可能缺乏显式的有限体积更新规则、界面通量表示或保证的通量差分形式，这在激波主导机制、长期展开和分辨率传输中可能很重要。

我们的方法将上下文条件化适应与保守数值骨干相结合。短轨迹段被编码为上下文向量，超网络使用该向量生成 Flux NO 目标网络的参数。因此，模型不仅仅是在上下文中条件化通用预测器；它条件化数值通量算子本身。与标准神经算子相比， resulting 模型通过通量差分更新演化状态。与 Flux NO 相比，它用上下文生成的算子替换固定的通量算子。与通用的 PDE 基础模型相比，它将上下文注入到结构保留的求解器中，使其能够适应未见过的通量函数，同时保留守恒律所需的有限体积归纳偏置。

## 3 在上下文 Flux Neural Operator

### 3.1 问题设置

对于公式 (1) 中的守恒律，我们的目标是从短轨迹观测中学习上下文条件化演化算子。设 $\bm{u}(t, \bm{x})$ 为连续解，并设 $\bm{u}^n \in \mathbb{R}^{d \times N_{\bm{x}}}$ 表示其在时间 $t = n \Delta t$ 处的网格采样状态，其中 $N_{\bm{x}} := N_{x_1} \times \dots \times N_{x_n}$。给定上下文轨迹 $\bm{U}^{n-k+1:n} = (\bm{u}^{n-k+1}, \dots, \bm{u}^n) \in \mathbb{R}^{k \times d \times N_{\bm{x}}}$，我们寻求预测下一个状态 $\bm{u}^{n+1}$。

我们不是将学习此映射作为无约束的输入-输出预测器，而是将问题分解为两个阶段：首先从观测轨迹推断底层动态的潜在表示，然后使用该表示实例化上下文条件化的 Flux NO。这自然导向了一种超网络公式：

$$ \bm{c} = \mathscr{E}(\bm{U}^{n-k+1:n}), \quad \Theta = H(\bm{c}), \quad \bm{u}^{n+1} = \mathscr{F}(\bm{u}^n, \Delta t; \Theta), \quad (6) $$

其中 $\mathscr{E}$ 是上下文编码器，$H$ 将上下文向量映射到目标网络参数，$\mathscr{F}$ 是 Flux NO 目标网络。编码器未获得解析通量函数、PDE 系数或方程标签；所有条件化信息必须从观测到的解历史中推断。

### 3.2 上下文编码器与超网络

给定短轨迹段，编码器提取紧凑的上下文向量并将其映射到 Flux NO 目标网络的参数。我们施加一个信息瓶颈：

$$ \mathbb{R}^{k \times d_{\mathrm{in}} \times N_{\bm{x}}} \longrightarrow \mathbb{R}^e \longrightarrow \mathbb{R}^q, \quad e \ll q, \quad (7) $$

其中 $e$ 是上下文维度，$q$ 是生成的目标网络参数数量。当使用网格坐标时，它们被追加为额外的输入通道。

#### 时间递归混合与空间注意力

编码器旨在单独处理时间和空间轴，同时尊重时间上的因果性。因此，我们采用受 TRecViT (Patraucean et al., 2025) 启发的时间递归视觉 Transformer 设计，其中时间混合由门控线性递归单元 (De et al., 2024; Botev et al., 2024) 处理，空间混合由 Transformer 块处理。给定 $\bm{U}^{n-k+1:n} \in \mathbb{R}^{k \times d \times N_{\bm{x}}}$，编码器首先使用带有可学习位置编码的 ViT 补丁嵌入对每个时间切片进行标记化 (Dosovitskiy et al., 2020)：

$$ \bm{V}^{(0)} = \mathrm{PatchEmbed}(\bm{U}^{n-k+1:n}) \in \mathbb{R}^{k \times P \times e}, \quad (8) $$

其中 $P$ 是空间补丁的数量。每个编码器层在每个空间标记的时间递归混合和每个时间步的空间自注意力之间交替进行：

$$ \widehat{\bm{V}}^{(\ell)}_{:,p} = \mathrm{TemporalBlock}^{(\ell)}\left( \bm{V}^{(\ell)}_{:,p} \right), \quad \bm{V}^{(\ell+1)}_{t,:} = \mathrm{SpatialTransformer}^{(\ell)}\left( \widehat{\bm{V}}^{(\ell)}_{t,:} \right). \quad (9) $$

在我们的实现中，时间块是基于带有因果深度一维卷积的门控线性递归单元的剩余递归块。这种交替结构允许编码器通过观测轨迹传播信息，同时模拟每个时间步的空间相互作用。在最后一层之后，我们应用标记级层归一化并在空间标记上对最终时间状态进行平均：

$$ \bm{c} = \frac{1}{P} \sum_{p=1}^{P} \mathrm{LayerNorm}\left( \bm{V}^{(L)}_{k,p} \right) \in \mathbb{R}^e. \quad (10) $$

然后，超网络将此上下文向量映射到目标网络参数：

$$ \Theta = H(\bm{c}) \in \mathbb{R}^q. \quad (11) $$

### 3.3 Flux Neural Operator 目标网络

目标网络是一个 Flux NO，其参数由上下文向量生成。因此，与具有固定参数的原始 Flux NO 不同，我们的模型为每个推断的动态实例化不同的数值通量算子。为清晰起见，我们...