Sobolev Training of End-to-End Optimization Proxies

TL;DR

Optimization proxies enable millisecond-scale inference by learning the solution operator $g(p)=\arg\min_{x}\{ f(x;p) : c(x;p)=0, x\ge 0\}$. The paper inserts solver-derived directional derivatives into a Sobolev loss, aligning both values and local derivatives via a masked Jacobian term, and analyzes supervised and self-supervised variants. It proves uniform value-and-gradient error bounds under Lipschitz regularity and validates the approach on large AC-OPF benchmarks and a mean–variance portfolio task, achieving substantial MSE reductions, major improvements in constraint satisfaction, and competitive or improved optimality gaps. The results suggest Sobolev training yields fast, reliable surrogates for safety-critical, large-scale optimization, with a practical path toward mixture-of-experts strategies depending on constraint tightness and availability of ground-truth derivatives.

Abstract

Optimization proxies - machine learning models trained to approximate the solution mapping of parametric optimization problems in a single forward pass - offer dramatic reductions in inference time compared to traditional iterative solvers. This work investigates the integration of solver sensitivities into such end to end proxies via a Sobolev training paradigm and does so in two distinct settings: (i) fully supervised proxies, where exact solver outputs and sensitivities are available, and (ii) self supervised proxies that rely only on the objective and constraint structure of the underlying optimization problem. By augmenting the standard training loss with directional derivative information extracted from the solver, the proxy aligns both its predicted solutions and local derivatives with those of the optimizer. Under Lipschitz continuity assumptions on the true solution mapping, matching first order sensitivities is shown to yield uniform approximation error proportional to the training set covering radius. Empirically, different impacts are observed in each studied setting. On three large Alternating Current Optimal Power Flow benchmarks, supervised Sobolev training cuts mean squared error by up to 56 percent and the median worst case constraint violation by up to 400 percent while keeping the optimality gap below 0.22 percent. For a mean variance portfolio task trained without labeled solutions, self supervised Sobolev training halves the average optimality gap in the medium risk region (standard deviation above 10 percent of budget) and matches the baseline elsewhere. Together, these results highlight Sobolev training whether supervised or self supervised as a path to fast reliable surrogates for safety critical large scale optimization workloads.

Figures (5)

• Figure 1: End-to-end pipeline for Sobolev training of optimization proxies. The upper half shows the online training loop driven by automatic differentiation (AD); the lower half is the offline oracle that supplies ground-truth solutions $x^\star$ and (masked) sensitivities via a Differentiable Optimization (DiffOpt) layer. Solid arrows are forward evaluations, dashed arrows are Jacobian/Gradient calculations.
• Figure 2: Comparison of NN Sobolev vs Benchmark across the Three Test Xases.
• Figure 2: Test-set gap (mean $\pm$ std)
• Figure 3: Change in Relative Optimality Gap ($\text{Gap}_{\text{Sobolev}}-\text{Gap}_{\text{Benchmark}}$). Negative Values Indicate Improvement.
• Figure 4: Test-set gap

Theorems & Definitions (7)

• Theorem 1: Value Matching Only
• Proof 1
• Theorem 2: Jacobian Matching Only
• Proof 2
• Theorem 3: Sobolev Guarantee
• Proof 3
• Remark 1