Distribution-Free Robust Functional Predict-Then-Optimize

TL;DR

This work tackles uncertainty quantification for neural operator surrogates solving PDEs by extending conformal prediction to infinite-dimensional Sobolev function spaces. It builds a robust predict-then-optimize pipeline that leverages calibrated, spectrum-aware prediction regions to bound downstream decision suboptimality. Theoretical results provide coverage guarantees across elliptic and parabolic PDE families, along with finite-projection and multi-stage optimization analyses to enable scalable, discretization-invariant procedures. Empirical results on quantum state discrimination and PDE-based collection problems demonstrate improved robustness and decision quality over nominal approaches. The approach offers distribution-free uncertainty control with practical implications for engineering design and quantum information tasks.

Abstract

The solution of PDEs in decision-making tasks is increasingly being undertaken with the help of neural operator surrogate models due to the need for repeated evaluation. Such methods, while significantly more computationally favorable compared to their numerical counterparts, fail to provide any calibrated notions of uncertainty in their predictions. Current methods approach this deficiency typically with ensembling or Bayesian posterior estimation. However, these approaches either require distributional assumptions that fail to hold in practice or lack practical scalability, limiting their applications in practice. We, therefore, propose a novel application of conformal prediction to produce distribution-free uncertainty quantification over the function spaces mapped by neural operators. We then demonstrate how such prediction regions enable a formal regret characterization if leveraged in downstream robust decision-making tasks. We further demonstrate how such posited robust decision-making tasks can be efficiently solved using an infinite-dimensional generalization of Danskin's Theorem and calculus of variations and empirically demonstrate the superior performance of our proposed method over more restrictive modeling paradigms, such as Gaussian Processes, across several engineering tasks.

Figures (6)

• Figure 1: Calibration curves with (solid) and without (dashed) spectral correction factors across data of varying GRF smoothness parameters for the step-index fiber Hamiltonian. Calibration is performed for three models that act on data at different spectral truncations of the full resolution $N = 64$ data.
• Figure 2: Calibration curves with (solid) and without (dashed) spectral correction factors across data of varying GRF smoothness parameters for the GRIN fiber Hamiltonian. Calibration is performed for three models that act on data at different spectral truncations of the full resolution $N = 64$ data.
• Figure 3: Calibration curves with spectral correction factors across data of varying GRF smoothness parameters for the Poisson equation. Calibration is performed for three models that act on data at different spectral truncations of the full resolution $N = 64$ data.
• Figure 4: Calibration curves with spectral correction factors across data of varying GRF smoothness parameters for the heat equation. Calibration is performed for three models that act on data at different spectral truncations of the full resolution $N = 64$ data.
• Figure 5: Visualization of collection solutions for the nominal and robust solutions laid atop the nominal field prediction (left) and true field (right).

Theorems & Definitions (38)

• Theorem 3.2
• Corollary 3.3
• Corollary 3.4
• Theorem 3.6
• Corollary 3.7
• Lemma 3.10
• Lemma 5.1
• Remark 5.2
• Remark 5.3
• Theorem 5.4