引言

在科学机器学习 (Scientific Machine Learning, SciML) 这个快速发展的领域中，我们正在见证一场范式的转变。研究人员不再仅仅训练神经网络来识别猫或生成文本，而是训练它们来模拟物理世界。该领域最强大的工具之一就是神经算子 (Neural Operator) 。 与将固定大小的向量映射到向量 (如将图像映射到标签) 的标准神经网络不同，神经算子学习的是函数空间之间的映射。它们可以接收物理系统的初始条件——比如流体的温度分布——并预测该函数随时间的演变，其求解偏微分方程 (PDE) 的速度比传统数值求解器快几个数量级。

然而，这里有一个问题。在高风险场景中——例如预测飓风路径、模拟桥梁结构应力或模拟气候变化——仅有速度是不够的。我们需要知道模型有多自信。如果一个神经算子预测了特定的流体流动，我们能相信它吗？如果输入数据与模型在训练期间看到的数据略有不同怎么办？

标准的神经算子，如流行的傅里叶神经算子 (FNO)，是确定性的。它们给出一个单一的答案，没有指出潜在的误差。这就是 LUNO (神经算子的线性化不确定性，Linearized Uncertainty for Neural Operators) 发挥作用的地方。

在这篇文章中，我们将深入探讨一个框架，该框架将确定性神经算子转化为概率模型。我们将探索 LUNO 如何利用函数式编程中的“柯里化 (currying)”概念和线性化技术，提供稳健、分辨率无关的不确定性估计，有效地将神经算子转化为函数值高斯过程 (Function-Valued Gaussian Processes) 。

问题所在: 不确定性缺口

要理解为什么 LUNO 是必要的，我们首先需要看看物理深度学习的现状。

神经算子

神经算子 (NO) 旨在学习 PDE 族的解算子。它们不再是求解方程的单个实例，而是学习将参数 (如初始条件或材料属性) 映射到解的数学规则。一旦训练完成，它们几乎可以瞬间评估新的输入。

然而，传统的深度学习不确定性量化 (UQ) 方法并不能直接应用于此。

  1. 无限维性: NO 的输入和输出是函数，理论上是无限维的对象。标准的贝叶斯神经网络 (BNN) 技术通常输出的是方差向量。
  2. 分辨率无关性: NO 的一个关键特性是它们可以在粗网格上训练，并在细网格上评估。不确定性度量也必须尊重这一属性；它不应该仅仅是一个“逐像素”的误差条，而应该是一个连续的不确定性场。

目标

我们需要一种能够提供输出函数概率信念 (probabilistic belief) 的方法。模型不应仅仅预测“位置 \(x\) 的温度是 \(T\)”，而是应该说: “温度函数很可能是 \(T(x)\)，并且在整个域上具有这种特定的协方差结构。”

理论: 从权重到函数空间

LUNO 框架的核心贡献是弥合了权重空间不确定性 (神经网络参数的随机性) 和函数空间不确定性 (输出函数的随机性) 之间的鸿沟。

作者提出了一个包含四个步骤的过程，涉及一个称为概率柯里化 (Probabilistic Currying) 的概念。让我们通过视觉来分解它。

图 1: LUNO 所涉步骤的图解。训练好的神经算子 F (左上) 通过 (反) 柯里化转换为输出在实数空间的等效神经网络 f (右上) 。在权重的高斯信念均值附近对 f 进行线性化，得到一个高斯过程后验 f，用于量化 f 所学函数的不确定性 (右下) 。最后，概率柯里化将 f 转换回关于神经算子 F 所学算子的函数值高斯过程后验 F (左下) 。

如图 1 所示，该过程从神经算子 \(F\) (左上) 开始，变为标准函数 \(f\) (右上) ，引入不确定性 (右下) ，最后转换回概率算子 \(\mathbf{F}\) (左下) 。

步骤 0: 神经算子

我们从一个训练好的神经算子 \(F\) 开始。在数学上，该算子将输入函数 \(a\) (来自空间 \(\mathbb{A}\)) 和参数 \(w\) (来自空间 \(\mathbb{W}\)) 映射到输出函数 \(u\) (在空间 \(\mathbb{U}\) 中) 。

\[ F: \mathbb{A} \times \mathbb{W} \rightarrow \mathbb{U} \]

在 PDE 的语境下，\(u\) 是定义在某个域 (如空间网格) 上的函数。

步骤 1: 反柯里化 (扁平化)

这是一个巧妙的理论转换。在函数式编程中，柯里化是一种技术，它将接受多个参数的函数的求值转换为求值一系列函数，每个函数只接受一个参数。

研究人员使用“反柯里化”来简化神经算子。他们不再将输出视为一个函数 \(u(x)\)，而是将空间坐标 \(x\) 仅仅视为另一个输入。

他们定义了一个新函数 \(f\): 显示反柯里化函数 f 将输入 a、x 和权重 w 映射到实向量的公式。

在这里，\(f\) 接收输入函数参数 \(a\)、特定的空间坐标 \(x\) 和权重 \(w\)，并输出一个实向量 (该点的解的值) 。在不确定性框架看来，这种转换将复杂的算子学习问题变成了一个标准的回归问题。

步骤 2: 线性化与权重不确定性

现在我们有了一个标准函数 \(f\)，我们可以应用成熟的贝叶斯深度学习技术。作者使用了线性化拉普拉斯近似 (Linearized Laplace Approximation, LLA) 。

其思想是用最优权重 (MAP 估计，\(w^*\)) 周围的二次碗来近似神经网络复杂的非线性损失景观。这意味着权重的后验分布是一个高斯分布 \(\mathcal{N}(\mu, \Sigma)\)。

为了将这种高斯信念从权重传播到输出，我们使用一阶泰勒展开对网络 \(f\) 进行线性化:

显示 f 在均值权重 mu 附近进行线性化的公式。

该方程表明，函数值近似等于均值权重处的值加上基于权重的梯度 (雅可比矩阵) 的修正项。

由于高斯随机变量的线性变换仍然是高斯分布，这种线性化将 \(f\) 的输出转化为高斯过程 (GP)。该 GP 的协方差由通过模型梯度投影的权重不确定性定义:

基于 f 的雅可比矩阵和权重协方差 Sigma 定义协方差核 K 的公式。

这个协方差核 \(K\) 告诉我们一个输入-位置对 \(((a_1, x_1))\) 处的不确定性如何与另一个 \(((a_2, x_2))\) 相关联。如果权重发生微小变化，这两点会如何共同变化？

步骤 3: 概率柯里化 (重构)

我们现在有了一个预测特定点 \(((a, x))\) 值的 GP。但我们想要的是关于函数的分布。

作者引入了概率柯里化 。 他们证明，在特定条件下，增强空间 (输入 + 坐标) 上的 GP 在数学上等同于函数值高斯过程 。

我们定义随机算子 \(\mathbf{F}\): 通过概率柯里化定义随机算子 F 的公式。

当给这个算子 \(\mathbf{F}\) 一个输入 \(a\) 时，它不会返回一个数字或向量，而是返回一个随机函数。这个随机函数是定义在空间域上的高斯过程。它的均值是原始神经算子的预测，其协方差函数源自线性化的权重:

显示函数值过程 F 的协方差的公式。

这就是 LUNO 的核心结论: 通过线性化底层函数并将空间坐标视为输入，我们获得了关于解空间的严格的、无限维的高斯信念。

实现: 用于傅里叶神经算子的 LUNO

该理论适用于任何神经算子，但作者为傅里叶神经算子 (FNO) 提供了一种具体且高效的实现。

计算完整的雅可比矩阵 (所有输出相对于所有权重的梯度) 在计算上非常昂贵。为了使其可扩展，作者使用了最后一层拉普拉斯近似 (Last-Layer Laplace Approximation) 。 他们假设早期层的权重是固定的，只将最后一层 (或最后一个块) 的权重视为概率性的。

对于 FNO，输出通常由最后一个傅里叶层之后的投影 \(q\) 生成。通过将贝叶斯推断限制在最后一层的参数上，预测分布得到了显著简化:

显示 FNO 的 LUNO 预测具体形式的公式。

这里，\(z^{(L-1)}\) 代表进入最后一层的特征。不确定性被封装在协方差核 \(K_a\) 中:

定义 FNO 特定情况下的协方差核 Ka 的公式。

这种结构允许惰性求值 。 我们不需要在训练期间为域中的每个点计算完整的协方差矩阵。我们可以高效地采样整个函数或按需查询特定点，从而保持 FNO 的分辨率无关性。

实验与结果

这个理论框架真的有效吗？作者在几个 PDE 基准测试上评估了 LUNO，包括 Burgers 方程和对流扩散方程。他们将 LUNO 与以下方法进行了比较:

  1. 输入扰动 (Input Perturbations): 向输入添加噪声以生成集成。
  2. 深度集成 (Deep Ensembles): 训练多个独立的 FNO (目前的 UQ 黄金标准) 。
  3. 基于采样的方法 (Sample-based methods): 使用权重不确定性，但对输出进行采样，而不是使用解析线性化。

小样本环境下的表现

在使用有限训练数据的 Burgers 方程 (流体力学中的一个基本 PDE) 测试中，LUNO-LA (拉普拉斯近似) 表现出了更优越的校准能力。

表 1: 在 Burgers 方程的 25 条轨迹上训练的 FNO 的 UQ 方法比较。

在表 1 中，我们查看了负对数似然 (NLL)，它衡量预测的不确定性与观察到的误差的拟合程度。数值越低越好。LUNO-LA 实现了最佳 (最低) 的 NLL (-2.0787)，显著优于集成和输入扰动方法。虽然集成的 RMSE (均值误差) 略低，但其不确定性估计 (\(\chi^2\)) 要差得多 (5.597 vs LUNO 的 1.022，其中 1.0 是理想值) 。这表明集成方法过于自信或校准不佳。

可视化不确定性

数字虽然重要，但在物理学中，我们需要看到场。 图 2 可视化了对流扩散方程 (Hyper-Diffusion equation) 的预测。

图 2: 通过几种不同方法量化的 FNO 预测不确定性。顶行显示目标、均值和样本。底行显示散布和协方差热图。

  • 顶行: 红色虚线是真实值。蓝色线是均值预测。阴影区域是置信区间。注意 LUNO-LA (最右侧) 如何生成光滑、连贯的样本 (细蓝线) ，这些样本遵循问题的物理规律。
  • 底行: 这显示了整个域的标准差 (不确定性) 。每个面板右上角的热图显示了预测协方差矩阵 。
  • 关键观察: 看看 LUNO-LA 的协方差热图。它在非对角线上显示了结构。这意味着模型“知道”点 \(x\) 处的误差与点 \(y\) 处的误差是相关的。这对于物理一致性至关重要。

分布外 (OOD) 鲁棒性

对不确定性的真正考验在于模型遇到新情况时。作者通过反转速度场 (“Flip”) 、添加热源 (“Pos”) 或组合这些变化 (“Pos-Neg-Flip”) ，为对流扩散方程创建了 OOD 数据集。

表 2: 不同 OOD 数据集上预期边际 NLL 评估。数值越低越好。

表 2 显示了在这些具有挑战性的数据集上的 NLL 得分。

  • Base (基准) : 大多数方法表现相似。
  • Pos-Neg-Flip (极端偏移) : 输入扰动彻底失败 (NLL ~494) 。集成表现不错，但 LUNO-LA 仍然非常有竞争力且稳定。

有趣的是，虽然集成具有鲁棒性，但它们的不确定性表示受到根本性的限制。

集成的“零空间”问题

为什么要选择 LUNO 而不是深度集成？集成非常昂贵 (你需要训练模型 5-10 次) 。除此之外，还有一个更微妙的数学原因。

由 \(M\) 个模型组成的集成产生的预测协方差的最大秩为 \(M-1\)。在高维输出空间 (如具有数千个网格点的流体模拟) 中，集成的不确定性在几乎所有方向上都是“扁平”的。它对其协方差“零空间”中发生的错误是视而不见的。

图 3: 比较集成 (左) 和 LUNO-LA (右) 。零空间投影面板突出了无法解释的误差。

图 3 生动地说明了这一点。

  • 左侧 (集成) : 看看标有“Null space projection (零空间投影) ”的面板。它显示了显著的结构。这代表了集成的包含不确定性估计无法解释的误差，因为其协方差矩阵是低秩的。模型在这些方向上是“不知不觉地”错误的。
  • 右侧 (LUNO-LA) : LUNO 基于权重的海森矩阵 (Hessian) 构建协方差矩阵。虽然仍是一种近似，但它通常具有高得多的秩 (受限于参数数量，而不是集成成员的数量) 。因此，LUNO 更全面地捕捉了误差结构。

自回归性能

最后，在模拟时间相关的 PDE 时，我们经常将模型的输出作为下一步的输入反馈回去 (自回归展开) 。误差会累积。

图 4: 不同 UQ 方法在 FNO 自回归展开上的平均性能。

图 4 跟踪了随时间步长的 NLL。LUNO-LA (NLL 图中最底部的紫色线) 随着时间的推移保持最低的 NLL，表明即使模拟偏离初始条件越来越远，其不确定性估计仍然保持校准。

结论与启示

LUNO 代表了神经算子在科学和工程应用中的重大进步。通过将算子视为函数值高斯过程 , LUNO 提供了:

  1. 信任: 对于安全关键型应用至关重要的可靠、经过校准的不确定性估计。
  2. 结构: 捕捉空间相关性而不仅仅是逐点误差的协方差信息。
  3. 效率: 它可以事后应用于训练好的模型 (特别是使用最后一层近似) ，无需昂贵的重新训练或管理大型集成。
  4. 分辨率无关性: 它保留了神经算子的定义特征，在任何离散化下提供不确定性估计。

LUNO 在函数式编程 (柯里化) 和贝叶斯推断 (GP) 之间建立的联系提供了一个理论上可靠的框架，其应用可能不仅限于 PDE，还有可能扩展到任何神经网络在连续空间之间进行映射的领域。对于 SciML 的学生和研究人员来说，LUNO 展示了一个未来，即我们的 AI 模拟器不仅速度快，而且能够自我感知其局限性。