跳转至

Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning

会议: NeurIPS 2026
arXiv: 2606.17513
领域: 物理/科学计算
关键词: 算子学习, 不确定性量化, 高斯过程, 几何感知深度核, 事后评估

一句话总结

针对神经算子缺乏可靠不确定性度量的问题,本文提出事后 UQ 框架 REEF-GP,直接利用冻结算子的内部几何潜特征构建深核高斯过程残差修正模型,在不损失预测精度的情况下实现媲美深度集成的高质量校准与物理可解释不确定性场。

研究背景与动机

神经算子(如 Fourier Neural Operator 与 Transolver)作为偏微分方程(PDE)的超快代理模型,已在流体力学、气象预报、固体力学与等离子体物理中得到广泛应用。然而绝大多数神经算子是确定性的,仅能给出单一的点估计预测。在诸如高风险工程决策、自适应主动数据采集等科学计算下游任务中,缺乏可靠的不确定性量化(UQ)极大地阻碍了神经算子的安全落地,尤其是在涉及几何边界发生复杂变形的非结构化网格任务中。

现有面向神经算子的 UQ 方案主要分为训练期和事后期两类。训练期方法如深度集成(Deep Ensembles)需要从头独立训练多个模型,在面对数百万网格点的大规模算子学习场景时算力开销极为高昂;MC Dropout 或概率神经算子(PNO)则常常伴随基准预测精度的显著劣化。传统的事后方法(如最后层 Laplace 近似或 LUNO)大多沿袭了标准深度学习的参数空间扰动视角,通过局部二阶展开近似权重后验。但在无限维泛函算子学习场景下,简单的参数先验假设极易因函数空间正则性失配而发生后验脆性塌陷,导致模型给出高置信度却严重错误的病态预测,且未能真正捕获算子在深层表征中所学到的几何形变流。

针对这一困境,本文改变了扰动权重参数的传统路线,转向函数回归与 Kennedy–O'Hagan(KOH)物理模型差异建模视角:既然性能优秀的神经算子在逐层逼近物理场时已经将复杂的物理几何解耦并嵌入到了内部潜特征中,为何不直接以其内部几何表征为度量空间来建模残差差异?核心 idea:将冻结的神经算子视为高保真确定性基底,提出 REEF-GP 框架,直接以算子内部多尺度几何特征构建非平稳深度核高斯过程,拟合输出残差并自适应捕捉物理激波与边界误差。

方法详解

整体框架

REEF-GP(Residual on Embedded Features Gaussian Process)将算子学习的单点预测转为泛函标量回归问题,并基于 Kennedy–O'Hagan 模型差异框架对真实物理场进行建模。将真解 \(v(\mathbf{x})\) 分解为冻结神经算子基线预测 \(\mathcal{G}_{\theta}(\mathbf{u}, \mathbf{a})(\mathbf{x})\)、确定性残差修正函数 \(f\) 与局部异方差不确定性噪声项。

整体执行管线分为三个阶段:首先,通过基于中心化核匹配(CKA)的代表层选择机制,从基准算子(如 Transolver)中抽取出跨越浅层几何输入、深层物理预测与中间代表层的空间对齐潜特征 \(\mathbf{h}\),拼装为增广状态向量;接着,通过谱归一化多层感知机对潜特征做平滑映射,与物理坐标核相乘构建几何感知非平稳乘积深度核;最后,在随机子集上通过微批次负对数边缘似然(NLL)进行参数训练,并在测试阶段采用广义专家乘积(gPoE)实现高效的分布式闭式后验推断。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入条件与几何网格<br/>u, a, x"] --> BB["冻结神经算子基底<br/>Transolver G_theta"]
    BB --> BasePred["基线预测场<br/>G_theta(u, a)(x)"]
    BB --> LSel["CKA代表层特征抽取<br/>浅层 + 中间层 + 深层特征拼接"]
    LSel --> AugState["增广状态构建<br/>z_bar = (u, x, h)"]
    AugState --> DK["几何感知深核设计<br/>谱归一化映射与乘积协方差"]
    AugState --> HN["异方差空间噪声网络<br/>输入依赖方差估计 sigma_psi^2"]
    DK & HN --> Opt["随机子集优化训练<br/>解耦网格分辨率的分批似然"]
    Opt --> PoE["广义专家乘积推断<br/>多专家闭式精度融合"]
    PoE --> PostRes["残差修正分布<br/>均值 delta_star 与 方差 sigma_star^2"]
    BasePred & PostRes --> Out["几何感知后验预测输出<br/>N(G_theta + delta_star, sigma_star^2)"]

关键设计

1. 几何感知深核设计:基于算子内部几何流构建非平稳协方差

传统高斯过程在无限维函数空间极易发生先验失配导致置信度失真,直接在欧几里得网格坐标上施加平稳各向同性核无法应对复杂流场的不连续性(如激波断裂)。针对这一问题,REEF-GP 不从零学习特征映射,而是借用冻结算子通过残差连接与物理注意力逐层形成的物理几何形变流。利用核的封闭性,将协方差定义为物理空间核、输入函数核与潜特征核的三重乘积: $\(k_{\boldsymbol{\phi}}(\bar{\mathbf{z}}, \bar{\mathbf{z}}') = \sigma^2 \cdot k_{\text{space}}(\mathbf{x}, \mathbf{x}') \cdot k_{\text{fun}}(\mathbf{u}, \mathbf{u}') \cdot k_{\text{latent}}(\rho(\mathbf{h}), \rho(\mathbf{h}'))\)$ 其中潜特征映射 \(\rho(\cdot)\) 采用带谱归一化(Spectral Normalization)的两层 MLP。谱归一化强制限制了网络映射的 Lipschitz 常数,防止深层嵌入发生特征坍缩(feature collapse)。在隐式几何流中,激波前后的网格点尽管在物理欧氏空间紧密相连,在算子特征空间中却被自然撕裂,使得核函数能在物理误差急剧聚集的局部自适应缩放相关长度,赋予先验极强的几何与物理感知能力。

2. CKA代表层特征抽取:轻量且充分的多尺度状态拼接

若将神经算子的全部隐藏层(例如 Transolver 的全部 8 层)盲目堆叠,不仅带来巨大的显存开销,更会因深层注意力块之间的特征强共线性引入大量信息冗余。REEF-GP 提出使用线性中心化核匹配(Linear Centered Kernel Alignment, CKA)自动化选取最具代表性的三层架构 \(\{L_0, l^*, L_{\text{last}}\}\)。浅层 \(L_0\) 保留原始几何与输入边界的细粒度空间拓扑;深层 \(L_{\text{last}}\) 编码已解析的高层物理场;中间层 \(l^*\) 则通过在各候选层中最小化与首尾两层的 CKA 相似度之和自动搜寻选出: $\(l^* = \arg\min_{l \in \{1, \dots, L-1\}} \big[\text{CKA}(\mathbf{H}_l, \mathbf{H}_0) + \text{CKA}(\mathbf{H}_l, \mathbf{H}_L)\big]\)$ 这种抽取策略使得特征向量维度保持在紧凑的 \(3D_h\),在消融实验中以极低显存代价完全追平了拼接全栈 8 层的标定效果。

3. 异方差空间噪声网络:解耦物理残差与局部不可解释误差

物理系统的误差往往高度非均质,激波前缘或边界层处伴随着算子拟合能力的骤降。若采用经典单标量同方差噪声项,会导致全局不确定性被均匀抹平。REEF-GP 将残差分解为潜确定性平滑函数与异方差噪声 \(\varepsilon(\bar{\mathbf{z}}) + \epsilon \sim \mathcal{N}(0, \sigma_{\boldsymbol{\psi}}^2(\bar{\mathbf{z}}) + \sigma_n^2)\)。其中局部依赖方差由轻量三层谱归一化 MLP 参数化: $\(\sigma_{\boldsymbol{\psi}}^2(\bar{\mathbf{z}}) + \sigma_n^2 = \text{Softplus}\big(g_{\boldsymbol{\psi}}(\bar{\mathbf{z}})\big) + \sigma_{\min}^2\)$ 下界 \(\sigma_{\min}^2 = 10^{-3}\) 保障了协方差矩阵求逆时的数值正定性。该网络使得高斯过程在残差高发区能自发提升噪声下界,生成高度锐化且紧密贴合误差轮廓的置信区间。

4. 随机子集优化与广义专家乘积推断:百万网格规模的可扩展性落地

在算子学习中,全样本点数 \(MN\) 经常达到 \(10^5\) 至 \(10^6\) 级,高斯过程 \(\mathcal{O}((MN)^3)\) 的计算开销无法直接承受。REEF-GP 借鉴 Kernel Flows 思想,在训练期每轮迭代仅随机采样 \(N_s \ll MN\)(默认 \(N_s = 5000\))个残差采样点,通过精确 Cholesky 分解直接优化微批次负对数边缘似然(NLL): $\(\mathcal{L}(\boldsymbol{\phi}, \boldsymbol{\psi}) = \frac{1}{2} \boldsymbol{\delta}^{\top} \mathbf{K}_{\boldsymbol{\phi}\boldsymbol{\psi}}^{-1} \boldsymbol{\delta} + \frac{1}{2}\log|\mathbf{K}_{\boldsymbol{\phi}\boldsymbol{\psi}}| + \frac{N_s}{2}\log(2\pi)\)$ 在测试期,若直接对全集条件推断依然受限,REEF-GP 构建 \(K\) 个条件在互斥随机子集上的独立 GP 专家,并在精度(逆方差)空间采用广义专家乘积(gPoE)实现闭式集成: $\(\frac{1}{\sigma_*^2(\bar{\mathbf{z}}_*)} = \frac{1}{K}\sum_{k=1}^K \frac{1}{\sigma_k^2(\bar{\mathbf{z}}_*)}, \quad \bar{\delta}_*(\bar{\mathbf{z}}_*) = \sigma_*^2(\bar{\mathbf{z}}_*) \cdot \frac{1}{K}\sum_{k=1}^K \frac{\bar{\delta}_k(\bar{\mathbf{z}}_*)}{\sigma_k^2(\bar{\mathbf{z}}_*)}\)$ 将推理显存严格限制在单个专家的 \(\mathcal{O}(N_s^2)\) 级别,彻底解耦了训练推理成本与原始网格分辨率。

实验关键数据

主实验

在 2D Transonic Airfoil(翼型气动跨音速流,含激波)和 3D ShapeNet Car(复杂车辆表面三维非结构化网格压力场)两大代表性基准上,将 REEF-GP 与基线 Transolver(Base 无 UQ)以及六种主流 UQ 方案(Deep Ensemble、MC Dropout、PNO、Input Perturbation、LUNO-LA、DVE-spatial)进行系统对比。各方法使用相同骨干网络与数据划分。

数据集 方法 rL2 (%) ↓ NLL ↓ CRPS (×10⁻²) ↓ NIS (×10⁻²) ↓ ES ↓
Airfoil (2D) Base (无 UQ) 1.24 ± 0.12 – – – –
Deep Ensemble 1.06 -3.46 28.63 3.89 0.35
MC Dropout 1.58 ± 0.21 -3.38 ± 0.15 44.88 ± 5.66 5.55 ± 0.91 0.55 ± 0.08
PNO 2.32 ± 1.40 -3.28 ± 0.65 74.27 ± 56.12 8.26 ± 6.39 0.79 ± 0.47
Perturbation 3.33 ± 0.19 -3.32 ± 0.05 81.05 ± 5.67 9.71 ± 0.66 1.08 ± 0.06
LUNO-LA (Laplace) 1.24 ± 0.12 -3.07 ± 0.82 38.70 ± 2.50 6.19 ± 0.54 0.47 ± 0.04
DVE-spatial 2.09 ± 0.25 19.98 ± 7.32 88.41 ± 12.53 25.89 ± 3.09 0.92 ± 0.10
REEF-GP (本文) 1.24 ± 0.12 -3.51 ± 0.06 39.80 ± 2.35 5.69 ± 0.47 0.45 ± 0.05
ShapeNet Car (3D) Base (无 UQ) 8.62 ± 0.13 – – – –
Deep Ensemble 7.73 4.16 1.88 30.28 194.57
MC Dropout 9.04 ± 0.59 17.21 ± 4.21 2.52 ± 0.26 67.96 ± 9.04 267.68 ± 17.66
PNO 13.76 ± 1.74 4.98 ± 1.52 3.53 ± 0.50 51.10 ± 12.17 337.16 ± 43.04
Perturbation 8.62 ± 0.13 13.44 ± 1.14 2.28 ± 0.03 50.82 ± 1.10 224.73 ± 3.97
LUNO-LA (Laplace) 8.62 ± 0.13 33.55 ± 23.43 2.47 ± 0.13 73.47 ± 14.91 271.32 ± 12.10
DVE-spatial 13.08 ± 0.26 11.22 ± 3.03 3.76 ± 0.10 96.17 ± 10.43 396.70 ± 6.79
REEF-GP (本文) 8.52 ± 0.14 2.84 ± 0.09 2.05 ± 0.03 28.81 ± 1.65 224.31 ± 4.18

(注:指标中 rL2、NLL、CRPS、NIS、ES 均越低越好。Airfoil 的 NLL/CRPS/NIS 缩放比例见表头,ShapeNet Car 为三维流场直接绝对尺度)

消融实验

在 Airfoil 基准上验证潜特征层选取策略(表 3)、训练样本规模适应性(表 4)、随机子集尺寸 \(N_s\)(表 5)及噪声模型选择(表 7)的影响。下表重点展示层选取与噪声建模消融:

配置 / 变体 rL2 (%) ↓ NLL ↓ CRPS (×10⁻⁴) ↓ NIS (×10⁻³) ↓ ES ↓ 说明
REEF-GP (CKA三层选优) 1.24 ± 0.12 -3.51 ± 0.06 39.80 ± 2.35 56.89 ± 4.66 0.45 ± 0.05 选 \(\{L_0, l^*, L_{\text{last}}\}\),开销低且表现均衡
仅浅层 \(L_0\) 1.24 ± 0.12 -2.75 ± 0.14 40.89 ± 2.46 67.72 ± 5.50 0.46 ± 0.05 仅包含原始几何输入,后验校准差,NLL显著恶化
仅深层 \(L_{\text{last}}\) 1.24 ± 0.12 -3.36 ± 0.12 40.25 ± 2.41 60.70 ± 5.10 0.46 ± 0.05 缺乏多尺度几何支撑,NIS 与 NLL 均弱于三层方案
全量堆叠 (All layers, \(L_0 \dots L_7\)) 1.24 ± 0.12 -3.54 ± 0.06 39.71 ± 2.40 55.99 ± 4.74 0.45 ± 0.05 性能与三层在统计上无显著差异,但显存开销翻倍
纯同方差噪声 (Homoscedastic) 1.24 ± 0.12 -3.41 ± 0.14 38.82 ± 2.66 58.84 ± 5.84 0.45 ± 0.05 全局使用单标量方差,不确定性场缺乏空间锐度

关键发现

  • 完全保持确定性算子的预测精度:REEF-GP 作为事后修正方法,完全没有损害冻结基底网络的逼近精度(Airfoil 上保持 1.24%,ShapeNet Car 上达到 8.52%),而 MC Dropout 和 PNO 的误差分别增加至 1.58% 和 2.32%(ShapeNet Car 上甚至退化至 13.76%)。
  • 3D 复杂几何下校准质量逆超集成与 Laplace:在非结构化 3D 点云的 ShapeNet Car 任务中,REEF-GP 的负对数似然(NLL)达到了 2.84,不仅显著优于参数空间 Laplace 近似(LUNO-LA 的 33.55),甚至击败了深度集成(4.16),区间得分 NIS 达到最佳的 28.81。
  • 物理可解释的空间不确定性聚集:定性标准差云图显示,REEF-GP 输出的不确定性极度自适应地聚焦在翼型上表面的激波锋面与车辆前保险杠边缘,与真实算子误差分布的空间轮廓高度一致,打破了同方差高斯过程平滑涂抹的局限。
  • 几何分布外(OOD)极强鲁棒性:在基于最大均值差异(MMD)划分的几何分布外四分位区间测试中(Q1近分布 → Q4远分布),LUNO-LA 与 Perturbation 的 NLL 发生剧烈发散(如 LUNO 从 -3.37 恶化至 -2.68),而 REEF-GP 的 NLL 始终稳定在 -3.61 至 -3.41 之间,随着几何形变自适应平滑放大方差。

亮点与洞察

  • 从“参数空间扰动”向“特征空间几何度量”的范式跃迁:传统 UQ 执着于对神经网络庞大权重参数进行高斯逼近,在无限维函数空间极易引发先验正则性错配。本文指出算子隐层已经内化了 PDE 的物理几何流形,在潜空间利用深核 GP 建模残差是更符合算子本质的事后 UQ 路径。
  • 巧用 CKA 降维与谱归一化防塌陷:通过 CKA 相似性最小化自动筛选首、中、尾最具代表性的三层潜特征,避免了盲目拼接导致的高维灾难与多重共线性;通过在特征投射网络中引入谱归一化,严密守住了核度量空间的几何流形。
  • 轻量开销实现生产级性价比:相比需要重训 5 倍模型的 Deep Ensembles,REEF-GP 在后处理拟合时仅需微小开销(在 RTX 6000 级工作站数分钟即可完成拟合),且支持对任意训练完毕的确定性算子即插即用。

局限与展望

  • 依赖具有点对点空间对应结构的算子架构:当前框架依赖 Transolver 这类能输出与离散网格点在物理空间显式对齐隐藏特征的 Transformer 架构,尚无法直接迁移至以全局频谱频域模式存在的 Fourier Neural Operator(FNO)。
  • 局限于稳态单标量输出问题:目前主要验证了稳态 PDE 上的标量场预测,尚未支持瞬态多时间步序列建模以及复杂多物理场耦合的向量输出场景。
  • 推断期核矩阵显存占用较高:尽管测试时间较短,但 gPoE 中各专家需缓存 \(\mathcal{O}(N_s^2)\) 的核矩阵,在大规模 3D 稠密网格上峰值显存需求依然高于简单的蒙特卡洛采样。

相关工作与启发

  • vs Deep Ensembles:集成学习需耗费 5 倍以上的巨额算力从头训练多个深度算子,REEF-GP 仅在推理后阶段对单模型进行轻量 GP 拟合,在 3D 几何任务上以 1/5 的训练代价取得了持平甚至更优的校准质量。
  • vs LUNO-LA (Laplace Approximation):LUNO 将最后一层权重参数展开为局部高斯后验,当遇到复杂非结构化点云时往往因曲率求逆受阻导致后验崩溃;REEF-GP 跳出参数空间,利用内部几何特征构建非平稳核,NLL 改善达一个数量级。
  • vs DVE-spatial (Deep Vecchia Ensemble):DVE 在原始特征空间执行近邻搜索易因稠密点嵌入高度重合导致协方差病态;REEF-GP 通过谱归一化投影、全局子集优化与 gPoE 规避了最近邻近退化,数值稳定性与准确率显著占优。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 颠覆了传统算子学习 UQ 局限于权重空间的思路,巧妙借助 KOH 差异模型与算子固有几何流实现特征空间深核建模。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 2D 翼型/管道/弹性体与 3D 车辆/Ahmed 体共 5 个代表性 PDE 几何基准,包含了完整的几何 OOD 分布迁移评估与详尽的消融实验。
  • 写作质量: ⭐⭐⭐⭐⭐ 概念阐述清晰严谨,从函数空间泛函回归到高斯过程设计的推导脉络一气呵成。
  • 价值: ⭐⭐⭐⭐⭐ 为高价值工业界科学计算代理模型提供了一个高精度、即插即用且具备物理可解释性的事后可靠性评估方案。