NeurIPS 2026 Optimization learning
Masked Sobolev Training for Feasibility-Reliable Optimization Proxies
Assessing generalization in feasibility and optimality for optimization proxies learned in Sobolev space.
1 Georgia Tech2 PSR3 Los Alamos National Laboratory
15–33%
Lower mean infeasibility
Across three AC power flow networks
2.5×
Fewer training iterations
To 95% reference performance in Acrobot
The problem
Accurate predictions are only part of the story.
Optimization proxies make repeated decision-making practical: a neural network replaces an expensive optimization solve with a fast forward pass. Their usefulness, however, depends on how well they generalize in objective quality and constraint satisfaction. A small prediction error can conceal a serious feasibility violation.
Constrained solution maps can pass through a combinatorial number of active-set regimes. Covering their behavior with value observations alone can demand many samples. Sobolev supervision adds local sensitivities, teaching the proxy both a decision and how that decision changes with the problem input.
For large problems, a full Jacobian introduces many simultaneous matching conditions and substantial training cost. A sparse, changing mask controls how many derivative conditions enter an update, while keeping the normalized derivative loss unbiased.
Results show that sensitivity supervision can improve generalization when local derivatives remain informative between samples. Sparse matching limits simultaneous training conditions, while vanishing or highly localized sensitivities offer weak guidance. The gains depend on the geometry of the decision map.
The method
Values, sensitivities,
and a changing mask.
Reference decisions and sensitivities supply local information. Automatic differentiation provides the proxy’s corresponding derivatives; a fresh mask selects which ones to match.
Decisions and sensitivities
The offline solver provides g(p) and Dg(p). AD supplies the proxy’s corresponding derivatives.
A fresh sparse mask
Uniform sampling limits the derivative conditions matched together. Value supervision remains complete.
Density and weight are separate
Normalization keeps the Jacobian loss unbiased. Density controls the number of conditions; λ controls their weight.
Problems and evidence
From regular maps
to changing active sets.
The experiments assess generalization under progressively different conditions: smooth local information, many coupled constraints, and highly localized sensitivities.
Acrobot / Optimal control
Generalization under regularity
and smooth mapping.
Finite-horizon control problem
The input p is the initial state. States contain joint angles and velocities; u is the torque at the second joint. The dynamics d reconstruct the state trajectory from the controls, and f is the trajectory cost.
A smooth teacher isolates the learning problem.
We first amortize the control policy with PPO. Both proxies then imitate this differentiable teacher using the same DAgger-style pipeline; Sobolev also receives selected state-to-action derivatives through automatic differentiation.
This removes numerical solver-gradient reconstruction as a source of noise. With feasibility enforced through bounded actions and dynamics, generalization is assessed in objective performance: reaching a common performance threshold sooner makes its benefit observable during learning.
Results from the paper
2.5×
Fewer iterations to the threshold
The first reconstructed threshold crossings occur at 1,349 versus 3,446 iterations, using 95% of reference performance.
About 16 vs. 40 minutes, despite a 2% increase in per-iteration cost.
Bounded actions and states reconstructed through the dynamics ensure feasibility, so objective performance provides the common measure of generalization. Both proxies use the same DAgger-style imitation pipeline; Sobolev also matches selected teacher derivatives obtained by automatic differentiation. The threshold uses a 100-step moving average of mean tip height. Feasibility by construction alone does not explain the gain: the regularity of the control map also matters.
AC optimal power flow / Energy dispatch
Generalization in a general
constrained optimization problem.
AC-OPF formulation
S = P + 𝑖Q denotes complex power, V voltage, Y branch admittance, and Yc branch charging. The star denotes complex conjugation. Bounds apply to the corresponding buses, generators and directed branches; generation is zero at buses without a generator. Demand Sd is the varying input. This is the network formulation presented in the paper.
One input change affects many coupled decisions.
Changing demand affects cost, nodal balance, voltages, generator limits and line loadings together. We reconstruct local sensitivities of the solver-selected solution through implicit differentiation of its optimality conditions.
Matching these sensitivities can introduce competing directions in the proxy’s parameter space. Uniform masking limits the number of concurrent derivative conditions; it does not choose which physical requirements matter. Evaluation therefore tracks both optimality and constraint violations.
Results from the paper
Smaller worst violations
across the test instances.
Sobolev matches or improves each instance’s maximum violation in over 85% of cases. Median normalized improvements are particularly large on PEGASE and RTE.
Reading these percentages. Let vi be the worst constraint violation in test instance i. The paper uses a shared Sobolev scale:
This is a normalized improvement, not a percentage reduction relative to each instance’s baseline. Values above 100% are possible on this shared scale.
The objective tradeoff
| Network | Value-only | Sobolev |
|---|---|---|
| IEEE-300 | 0.02% | 0.22% |
| PEGASE-1k | 0.07% | 0.15% |
| RTE-6k | 0.07% | 0.07% |
Across test cases, Sobolev matches or improves maximum constraint violations in over 85% of cases.
Sparse supervision matters.
On IEEE-300, dense matching has 5.2× higher INF than the 5% mask, despite a smaller objective gap.
These are separate configurations from the three-network comparison above.
| Density | MSE | Gap | INF ↓ |
|---|---|---|---|
| 5% | 0.0065 | 0.22% | 0.0099 |
| 10% | 0.0088 | 0.33% | 0.0094 |
| 25% | 0.0087 | 0.29% | 0.0141 |
| 100% | 0.0120 | 0.17% | 0.0517 |
Mean–variance portfolios / Allocation under risk
Impacts of uninformative
sensitivities.
Mean–variance allocation problem
x gives the asset allocations, μ the expected returns, Σ the positive semidefinite covariance matrix, B the budget, and σmax the standard-deviation risk cap. Inputs μ and σmax vary while covariance remains fixed. Unspent budget is allowed.
Convexity does not make every derivative informative.
Optimal portfolios can contain many zero allocations. Their corresponding local derivatives may also vanish; small input changes can activate different assets and change the Jacobian sharply.
Both experimental proxies enforce feasibility by construction. This setting tests what remains when derivative information is highly localized: a valid sensitivity at one sample may provide little guidance elsewhere. The role of regularity is separate from whether a feasibility-preserving reconstruction exists.
Results from the paper
The benefit changes with the risk regime
Portfolios are often sparse, leaving many sensitivities close to zero. Small parameter changes can alter the active assets or sharply change the interaction of risk, budget, and bound constraints.
Both methods enforce feasibility by construction and train against the portfolio objective. Sobolev adds solver sensitivities. The objective-only baseline has the smaller mean gap under tight risk caps; outside that region, Sobolev performs better in this experiment.
| Risk cap / budget | Objective-only | Sobolev |
|---|---|---|
| ≤ 10% | 14.7 ± 21.2 | 29.5 ± 26.9 |
| > 10% | 18.9 ± 20.2 | 8.7 ± 9.5 |
Lower is better. This result motivates using local sensitivity information selectively by regime, rather than assuming a universal advantage.
A theoretical perspective
Local information can change the error order.
In a regular region, exact value-and-Jacobian interpolation improves the bound on uniform value error as the covering radius δ decreases.
Idealized regional bounds under the stated regularity assumptions, including Lipschitz Jacobians for the joint bound. They do not guarantee convergence of masked training, but explain how informative sensitivities can strengthen interpolation between samples.
Assumptions and complete bounds
On a compact, convex parameter region, let δ = supp mini ‖p − pi‖ be the training set’s covering radius. Under exact interpolation at the samples:
L denotes a map’s Lipschitz constant and M its Jacobian’s Lipschitz constant. The derivative and joint bounds require continuously differentiable maps with Lipschitz Jacobians. These assumptions apply within regular regions, not across arbitrary active-set transitions; practical training need not attain exact interpolation.
Reference
Cite this work.
Camera-ready paper and code links will appear here when available.
@inproceedings{rosemberg2026masked,
title = {Masked Sobolev Training for
Feasibility-Reliable Optimization Proxies},
author = {Rosemberg, Andrew W. and Garcia, Joaquim Dias
and Bent, Russell and Van Hentenryck, Pascal},
booktitle = {Advances in Neural Information Processing Systems},
year = {2026}
}