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

Technical Summary: Distribution-Free Robust Predict-Then-Optimize in Function Spaces

Authors: Yash Patel and Ambuj Tewari (University of Michigan)

1. Problem Statement

The paper addresses a critical gap in the deployment of neural surrogate models (specifically neural operators) for engineering design optimization.

In modern engineering, solving Partial Differential Equations (PDEs) is computationally expensive. Neural operators are used to learn "flow maps" that approximate these PDE solutions rapidly. However, unlike classical numerical solvers, neural operators lack formal accuracy guarantees. If an optimization pipeline uses a miscalibrated neural surrogate to find an optimal design (e.g., minimizing drag on an aircraft), the resulting design may perform poorly in reality because the optimizer "overfit" to the errors in the surrogate.

While "Predict-Then-Optimize" frameworks exist for finite-dimensional vectors, they have not been extended to infinite-dimensional function spaces. Existing uncertainty quantification (UQ) methods for operators often rely on distributional assumptions that fail under misspecification, or they only provide guarantees on fixed-discretization grids, sacrificing the discretization-invariance that makes neural operators powerful.

2. Methodology

The authors propose a framework that integrates Conformal Prediction with Robust Optimization in Sobolev spaces.

A. Functional Conformal Prediction:
To provide distribution-free guarantees, the authors extend conformal prediction to infinite-dimensional spaces.

• Score Function: They define a spectral score function $s_{N;\tau}(a, u)$ based on the $(s-\tau)$-Sobolev norm residual. This allows them to quantify error using truncated spectral representations (which are observable) while making claims about the full, untruncated function.
• Margin-Padding: Because we only observe truncated functions, the authors prove that a "margin" must be added to the conformal quantile to account for the unobserved high-frequency modes. This margin decays as the spectral resolution $N$ increases.
• Theoretical Guarantees: They provide theorems proving that their prediction regions $C^*_{N;\tau}(a)$ achieve marginal frequentist coverage in Sobolev spaces for various PDE families (Elliptic and Parabolic).

B. Robust Functional Predict-Then-Optimize:
Instead of optimizing against a single predicted function $\hat{u}$, the authors formulate a min-max optimization problem:
$$\min_{w \in W} \max_{b_u \in C^*_{N;\tau}(a)} J[w, b_u]$$
where $J$ is the engineering objective (e.g., resource collection or energy dissipation) and $C^*$ is the conformal uncertainty set.

C. Multi-Stage Optimization:
Solving robust optimization in high-dimensional spectral spaces is computationally taxing. The authors introduce a Multi-Stage Robust Optimization algorithm. It starts by solving the problem at a coarse spectral resolution (low $N$) and uses the result to initialize optimization at progressively finer resolutions. They provide a convergence analysis showing that this approach is significantly more efficient than naive high-resolution optimization.

3. Key Contributions

1. Mathematical Extension: The first formal extension of conformal prediction guarantees from finite-dimensional vectors to infinite-dimensional Sobolev spaces.
2. Discretization-Invariant UQ: A method to provide uncertainty sets that remain valid regardless of the resolution used for observation.
3. Robust Design Framework: A novel pipeline for robust engineering design that propagates functional uncertainty through to the final decision variable.
4. Efficient Algorithm: A multi-stage optimization strategy that leverages the smoothness of the optimization landscape across different spectral scales.

4. Experimental Results

The authors validate their method across three distinct domains:

• PDE Calibration: They demonstrate that their "margin-padded" conformal method successfully achieves the desired coverage for the Poisson and Heat equations, whereas standard uncorrected conformal prediction fails to cover the true functions.
• Resource Collection (Spatial Optimization): In a "maximal coverage" problem (placing collectors to capture a resource), the robust approach significantly outperformed the nominal approach. The robust optimizer "hedges" its bets by spreading collectors out rather than concentrating them on predicted local peaks that might be artifacts of surrogate error.
• Quantum State Discrimination: In a quantum communication task (maximizing mutual information between sender and receiver), the robust measurement scheme significantly outperformed both the "Pretty Good Measurement" (PGM) and the nominal measurement scheme, proving the utility of the framework in high-stakes, high-dimensional physics applications.

5. Significance

This work is significant because it provides a principled, distribution-free safety net for the growing field of AI-driven engineering. By moving from "point predictions" to "robust functional sets," it allows engineers to use fast neural surrogates with the confidence that their optimized designs will remain performant even when the surrogate is imperfect. It bridges the gap between deep learning (neural operators), statistical learning (conformal prediction), and classical numerical analysis (Sobolev spaces).