Differentiating Through Constrained Optimization: Implicit Layers and Convex Programs

Many ML pipelines have a two-stage structure: a predictor outputs unknown parameters; a solver consumes them and returns decisions. Training the predictor with MSE or cross-entropy is agnostic to how prediction errors propagate through the solver, and small parameter errors near a constraint boundary can flip an entire decision¹². Constrained-optimization learning (variously decision-focused learning, predict-then-optimize, or differentiable-optimization) closes this loop by differentiating through the solver itself; the predictor learns to minimize regret on decisions, not error on parameters. Each subfamily comes with characteristic open problems, walked through paragraph-by-paragraph after the family table below.

The historical trajectory has three eras. Early implicit-layer work (2010-2017) established the gradient calculus for graphical-model inference (Domke 2010) and convex-program sensitivity (Barratt 2019); the differentiable-layer era (2017-2021) packaged this calculus into deep-learning frameworks via OptNet (Amos and Kolter 2017) and CvxpyLayers (Agrawal et al. 2019), with the predict-then-optimize formulation (Elmachtoub and Grigas 2022) giving consistent surrogate losses; and the modern era (2021-present) has scaled to combinatorial problems via blackbox finite differences (Vlastelica et al. 2020) and perturbed maximizers (Berthet et al. 2020), to structured prediction via Fenchel-Young losses (Blondel et al. 2020), and to dual learning (Klamkin et al. 2025) for large-scale constrained programs.

Read across the eras and one question keeps being asked in different vocabulary: an argmin is piecewise constant in its inputs, so its honest derivative is zero almost everywhere and undefined on a measure-zero set where all the information lives. Where do you get a usable gradient from? There are only four answers, and every family below is one of them. Differentiate the optimality conditions rather than the solver, which is exact where the active set is fixed and singular where it changes. Differentiate a smoothed problem, which is always defined but is the gradient of a different function than the one you solve. Take a finite difference through the solver itself, which needs no structure at all but costs a second solve and returns nothing when the perturbation fails to cross a face. Or refuse the gradient and change the loss instead, so that the training signal never has to pass through the solver. The seven families are these four answers crossed with what the inner problem happens to be, and the footnotes above are the bill each answer runs up.

The taxonomy below distinguishes seven families. Implicit-layer methods solve the forward problem with a classical solver and back-propagate by differentiating the KKT system; this is exact when the active set is fixed, but the linear system is ill-conditioned at the boundary³ and costly to assemble at every step⁴. Predict-then-optimize methods replace the MSE training signal with a regret-based surrogate (SPO+, CaVE), eliminating the loss-mismatch⁵ but inheriting the solver's non-differentiability for combinatorial problems⁶. Combinatorial methods sidestep non-differentiability via blackbox interpolation, perturb-and-MAP, or Sinkhorn relaxations, paying for the smoothing with an integer-relaxation gap⁷. Augmented-Lagrangian and prox-based methods enforce constraints by penalty rather than solver call, trading exact feasibility for a tunable violation-vs-optimality balance⁸. Bilevel methods fold the inner problem into a hyper-gradient, enabling millions of hyperparameters but requiring careful Hessian-vector treatment. Structured-prediction methods cast inference as regularized optimization with closed-form Fenchel-Young losses. Application-driven methods (MPC, portfolio, routing) tune the recipe to a problem class. A persistent diagnostic across the field is scale mismatch: an MLP predictor's natural output scale is rarely the solver's parameter scale, and naive loss combinations across the two scales train poorly⁹; another is the failure to exploit problem structure¹⁰, where a generic implicit-layer beats a problem-specific decomposition only on toy sizes.

The table groups the literature into seven families, walked through in chronological order. Implicit-layer foundations (Domke, Barratt) gave the gradient calculus for parametric MAP and convex programs by differentiating the KKT system; the construction is exact when the active set is fixed but is dominated by the cost of solving an inner linear system at every gradient step¹¹, and is numerically delicate near a constraint boundary¹². Feasibility is not at issue for these rows, because the forward pass is a call to a feasible solver; the OOD column flags whether the predictor's parameters generalize beyond the training distribution of constraints, which Domke's marginal-NLL recipe does not address while CvxpyLayers and Lorraine's hyperparameter approach do.

The differentiable-optimization-layer line (OptNet, CvxpyLayers, DDN, BPQP) extended the implicit-layer recipe to general QP and DCP problems, packaging the KKT-Jacobian backward into PyTorch and JAX. A training step with one of these layers is slower than a forward-only solver call, because the cost is dominated by a factorization of the KKT matrix at every step¹³; BPQP attacks the constant factor with a structure-aware backward but leaves the asymptotic order alone, and the implicit-layer family also fails to exploit problem-specific decompositions that classical solvers use for free¹⁴. Within this line, input-convex networks (Amos et al. 2017) take an orthogonal route: instead of differentiating a solver, they constrain the network architecture to be convex in its input, so the forward pass is an unconstrained convex optimization with closed-form gradients.

Predict-then-optimize methods (SPO+, CaVE) replace the predictor's MSE loss with a regret-based surrogate that knows the downstream solver. SPO+ is a Fisher-consistent convex upper bound on the true regret, achieved with a single extra solver call per training point and gradients computed without backprop through the solver itself; this sidesteps the non-differentiable-solver obstacle¹⁵, but because SPO+ is an upper bound on the SPO loss rather than the loss itself, its minimizer need not minimize regret and the residual regret at convergence is not driven to zero (Elmachtoub and Grigas 2022)¹⁶. CaVE extends the recipe to MILP by aligning the predicted cost vector to the optimality cone of the integer-feasible solution, the first decision-focused method that closes the gap on binary programs without smoothing the integer feasibility¹⁷.

Combinatorial methods (Blackbox-CO, Perturbed-FY, Fenchel-Young, GNN-NP, AttnTSP, SATNet) attack the harder problem of differentiating through non-differentiable integer/SAT solvers¹⁸. Blackbox differentiation runs the solver twice with perturbed parameters and uses a finite-difference gradient on the resulting decisions; perturbed-FY adds Gumbel noise to the solver inputs and treats the expected argmax as a differentiable smoothed version, with closed-form gradients via the Fenchel-Young machinery. SATNet takes a third route, replacing the solver entirely with a low-rank SDP relaxation differentiable end-to-end. The cost of all three is an integer-relaxation gap¹⁹: the relaxed gradient signal is well-defined but biased away from the true integer optimum, and the gap does not vanish as the predictor improves.

Bilevel methods (Franceschi, Lorraine, Lee) fold an inner optimization (model fit on the train set) into an outer hyper-gradient (val loss). Reverse-mode unrolling computes exact hyper-gradients but stores every inner iterate, so memory grows linearly in the number of inner steps and the affordable horizon is set by whatever fits in device memory (Franceschi et al. 2018); Neumann-implicit methods (Lorraine et al. 2020) approximate the hyper-Jacobian as a truncated Neumann series, scaling to millions of hyperparameters at the cost of approximation bias. The OOD column is high for these because the meta-loop is training for OOD generalization by construction; the inner-loop ill-conditioning is the same KKT pathology²⁰ as in the implicit-layer family.

Dual learning (Klamkin 2025) and augmented-Lagrangian methods enforce constraints via a learned dual variable rather than a solver call, training a primal predictor and a dual predictor jointly with an augmented-Lagrangian objective. Klamkin et al. do not accept an infeasibility rate at all: a dual completion layer projects the network's output onto the dual feasible set, so every prediction is dual-feasible by construction and the learned bound is a valid certificate (Klamkin et al. 2025). What is traded away is exactness rather than feasibility, and the paper reports optimality gaps below 1% together with speedups of several orders of magnitude over a commercial solver on large linear programs; the prox-update on the dual variable is what stabilizes the alternation between primal-feasibility and primal-optimality²¹. Against the predict-then-optimize family, dual learning drops the inner solver from the training loop entirely, which is what buys the scaling to large LP/QP instances.

The table separates four eras: foundational implicit gradients (2010-2017), the differentiable-layer build-out (2017-2021), combinatorial and predict-then-optimize maturation (2018-2024), and the modern dual learning and cone-aligned recipes (2024-2025). Three readings, and they are readings of the ordering in the table rather than of its digits, which are a composite. First, no single method dominates: implicit-layer methods give exact gradients but pay a linear solve per backward pass, predict-then-optimize needs only forward solver calls but optimizes a surrogate whose floor is nonzero, combinatorial methods run on integer problems but carry relaxation bias, and dual learning drops the solver from the loop and scales furthest. Second, OOD generalization is the unresolved axis: the methods rated "high" (SPO+, CvxpyLayers, Lorraine, Perturbed-FY, FY, Meta-DCO, CaVE, Dual-OL) are those whose training signal is the downstream task, while methods rated "low" (Domke, GNN-NP, AttnTSP, SATNet, Kotary) overfit to the training instance distribution. Third, the trend is toward hybrid recipes: BPQP combines structured backward with task loss, CaVE combines cone alignment with decision-focused regret, and Dual-OL combines augmented-Lagrangian with primal-dual prediction; the next generation appears to be augmented-Lagrangian + decision-focused trained end-to-end on industrial-scale LPs.

References