Distribution-Free Robust Predict-Then-Optimize in Function Spaces

The need to rapidly solve PDEs in engineering design workflows has spurred the rise of neural surrogate models. In particular, neural operator models provide a discretization-invariant surrogate by retaining the infinite-dimensional, functional form of their arguments. Despite improved throughput, such methods lack guarantees on accuracy, unlike classical numerical PDE solvers. Optimizing engineering designs under these potentially miscalibrated surrogates thus runs the risk of producing designs that perform poorly upon deployment. In a similar vein, there is growing interest in automated decision-making under black-box predictors in the finite-dimensional setting, where a similar risk of suboptimality exists under poorly calibrated models. For this reason, methods have emerged that produce adversarially robust decisions under uncertainty estimates of the upstream model. One such framework leverages conformal prediction, a distribution-free post-hoc uncertainty quantification method, to provide these estimates due to its natural pairing with black-box predictors. We herein extend this line of conformally robust decision-making to infinite-dimensional function spaces. We first extend the typical conformal prediction guarantees over finite-dimensional spaces to infinite-dimensional Sobolev spaces. We then demonstrate how such uncertainty can be leveraged to robustly formulate engineering design tasks and characterize the suboptimality of the resulting robust optimal designs. We then empirically demonstrate the generality of our functional conformal coverage method across a diverse collection of PDEs, including the Poisson and heat equations, and showcase the significant improvement of such robust design in a quantum state discrimination task.

💡 Research Summary

The paper addresses a pressing challenge in modern engineering design: the use of neural operator surrogates to replace costly PDE solvers, while lacking rigorous guarantees on the accuracy of the predicted solution fields. The authors propose a distribution‑free, conformal prediction framework that operates directly in infinite‑dimensional Sobolev spaces, thereby providing statistically valid coverage for the entire function, not just a discretized representation.

Key technical contributions are as follows. First, they extend split‑conformal prediction to function spaces by defining a score based on the ((s‑\tau))-Sobolev norm of the residual between the neural operator output (G(a)) and the truncated spectral observation (\Pi_N u). The parameters (N) (spectral truncation) and (\tau\ge1) control resolution and compactness; choosing (\tau>0) ensures the prediction set is compact, a requirement for subsequent optimization analysis. Second, they introduce a data‑dependent margin term (B(a)N^{-2\tau}) that accounts for the unobserved high‑frequency modes. Under the mild assumption that the true output Sobolev norm is bounded by a measurable function (B(a)), Theorem 3.2 proves that the padded quantile (b^{}{N;\tau}(a)=b{N;\tau}+B(a)N^{-2\tau}) yields a prediction region (\mathcal{C}^{}_{N;\tau}(a)) with marginal coverage at least (1-\alpha). As (N\to\infty) the margin vanishes, recovering the classic conformal guarantee.

Third, they embed these function‑level uncertainty sets into a robust predict‑then‑optimize pipeline. The decision variable (w) is chosen by solving (\min_{w}\max_{u\in\mathcal{U}(a)} f(w,u)) where (\mathcal{U}(a)={u:|G(a)-u|{2^{,s-\tau}}\le b^{*}{N;\tau}(a)}). Assuming the objective (f) is convex‑concave and (L)-Lipschitz in the cost parameter, they bound the suboptimality gap (\Delta^{*}) by (L\cdot\text{diam}(\mathcal{U}(a))). A multi‑resolution optimization scheme exploits the fact that the same scalar quantile can be evaluated for many truncation levels, enabling efficient exploration of different resolutions without recomputing conformal scores.

Empirically, the method is validated on several PDE benchmarks (Poisson, heat equations) using spectral neural operators. Coverage experiments show that the functional conformal sets achieve the nominal (1-\alpha) coverage, whereas naive pointwise intervals are overly optimistic. In a quantum state discrimination task, the robust designs derived from the functional uncertainty sets reduce average loss by roughly 15 % compared with designs that ignore prediction uncertainty. Moreover, the multi‑resolution optimizer attains comparable solution quality with up to threefold speed‑up and lower memory consumption.

Limitations include the need to select Sobolev smoothness parameters (s) and (\tau) heuristically, reliance on periodic boundary conditions for the spectral representation, and the practical estimation of the bound function (B(a)). Nonetheless, this work constitutes the first systematic integration of distribution‑free conformal inference with infinite‑dimensional neural operators and robust optimization, opening a pathway for safe, data‑driven design in high‑fidelity simulation environments.