Sobolev Training of End-to-End Optimization Proxies

2025-05-16

AI Summary

This work addresses the insufficient accuracy and reliability of end-to-end surrogate models in parametric optimization. We propose a sensitivity-enhanced learning framework grounded in Sobolev training, which— for the first time—incorporates first-order sensitivity information (directional derivatives) of the optimizer into surrogate training, yielding a hybrid model combining full supervision and self-supervision. We theoretically establish that aligning first-order sensitivities ensures a uniform approximation error bound and characterize generalization performance via Lipschitz continuity analysis. Empirical evaluation on AC optimal power flow (across three standard benchmarks) and mean-variance portfolio optimization demonstrates: supervised training reduces mean squared error by 56% and improves the median worst-case constraint violation by 400%; self-supervised training halves the average optimality gap in the medium-risk regime.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Optimization
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: User privacy protection in personalized systems

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.

Problem

• Improving optimization proxies via Sobolev training for faster inference
• Aligning proxy solutions and derivatives with optimizer outputs
• Reducing errors and constraint violations in critical optimization tasks

Innovation

• Sobolev training integrates solver sensitivities into proxies
• Supervised proxies use exact solver outputs and sensitivities
• Self-supervised proxies rely on objective and constraint structure