Optimization for Machine Learning

Bayesian Optimization: Theory and Applications

The foundations of Bayesian optimization trace back to work on efficient global optimization (EGO) (Jones, Schonlau, and Welch 1998), which introduced expected improvement as an acquisition function for response surface optimization. The theoretical framework using Gaussian processes draws heavily from the comprehensive treatment in (Rasmussen and Williams 2006).

Modern tutorials (Brochu, Cora, and Freitas 2010; Frazier 2018) provide accessible introductions to the field, while (Shahriari et al. 2016) offers an extensive review covering methodological advances and applications. The GP-UCB algorithm (Srinivas et al. 2012) established regret bounds connecting Bayesian optimization to the bandit literature.

Practical Bayesian optimization for machine learning hyperparameter tuning was demonstrated by (Snoek, Larochelle, and Adams 2012), showing that careful GP prior design leads to state-of-the-art performance. High-dimensional extensions include random embeddings (Z. Wang et al. 2016), local trust regions (Eriksson et al. 2019), and sparse axis-aligned subspaces (Eriksson and Jankowiak 2021).

Information-theoretic acquisition functions (Hernández-Lobato, Hoffman, and Ghahramani 2014) provide principled approaches to exploration, while portfolio methods (Brochu, Hoffman, and Freitas 2011) combine multiple acquisition strategies. Multi-fidelity methods (Kandasamy et al. 2017; Li et al. 2020; Swersky, Snoek, and Adams 2014) leverage cheap approximations to accelerate optimization. Transfer learning (Feurer et al. 2022; Swersky, Snoek, and Adams 2013) enables knowledge transfer across related optimization tasks.

Scalable Gaussian process inference (Snelson and Ghahramani 2006; Titsias 2009; Matthews et al. 2016) underpins practical BO implementations, with GPyTorch (Gardner et al. 2018) providing efficient GPU-accelerated computation. Applications span neural architecture search (Kandasamy et al. 2018; Falkner, Klein, and Hutter 2018), biological sequence design (Stanton et al. 2022), and safe exploration (Sui et al. 2015).

The following leaderboard summarises the methods this chapter dissects, organised by the regime each one targets. Read the table top-down as exploitation/exploration first, then scalability, then dimensionality, then setting-specific extensions.

Constrained Optimization in Machine Learning

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 SoTA leaderboard 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, Martins, and Niculae 2020), and to dual learning (Klamkin, Tanneau, and Hentenryck 2025) for large-scale constrained programs.

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¹². The 100% constraint-satisfaction reported in the table for these rows reflects that the forward pass is a feasible solver call; 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. The wall-clock cost of these layers (0.3-0.5x compared to a forward-only solver call) is dominated by an LU factorization on the KKT matrix at every step¹³, with BPQP reducing the constant factor by 5-10x via structure-aware factorization but not the asymptotic order; the implicit-layer family also fails to exploit problem-specific decompositions that classical solvers use for free¹⁴. Within this line, input-convex networks (Amos, Xu, and Kolter 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 pays for it with a 4.2% optimality gap at convergence (vs OptNet's 0.3% on small QPs) because the surrogate is not the true regret¹⁶. 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 all inner-loop activations, limiting the inner horizon to ~20-50 steps (Franceschi et al. 2018); Neumann-implicit methods (Lorraine, Vicol, and Duvenaud 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. The reported 99% constraint-satisfaction (vs 100% for solver-in-the-loop methods) is the price of dropping the inner solver from the training loop, recovered as an 8x training speedup; 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 trades a small infeasibility rate for a much faster training loop and better 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. First, no single method dominates: implicit-layer methods are exact but slow (0.3-0.5x), predict-then-optimize is fast but has a 1-4% optimality floor, combinatorial is fast on integer problems but hurt by relaxation bias, and dual learning is fastest but trades a percent of feasibility. 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.

Optimal Transport: Theory and Applications in Machine Learning

The Feasibility-vs-Fidelity Tension

One tension drives the entire computational history of optimal transport: exact solutions are provably correct but computationally prohibitive, while tractable approximations each pay a distinct fidelity price. Every algorithmic generation from 1955 to 2024 is best understood as a specific answer to the question "how much fidelity can we afford to trade for the compute budget we actually have?" The five major lineages (classical LP, entropic/Sinkhorn, sliced projection, unbalanced/Gromov-Wasserstein extensions, and neural/flow-matching) converge, in recent work, on one unifying principle: recover a coupling or transport map that minimizes a cost, subject to the marginal and feasibility constraints your budget can enforce.

The table below places every major method on this trade-off surface. Methods appear as responses to the limitations of their predecessors; each section below follows that logic rather than a strict chronology. Textbook coverage is in the Mathematical Foundations section; algorithmic details are in Computational Methods and Extensions.

Villani's two volumes (Villani 2003, 2009) remain the definitive references; Santambrogio (Santambrogio 2014, 2015) is the most compact entry point with a calculus-of-variations emphasis; Ambrosio-Gigli-Savaré (Ambrosio, Gigli, and Savaré 2008) develops the gradient-flow theory; and Peyré-Cuturi (Peyré and Cuturi 2020) is the standard computational reference. Rachev's 1985 monograph (Rachev 1985) introduced the unified "Monge-Kantorovich" nomenclature to the Western probability community. Each reference carries characteristic strengths and well-known limitations^{222324252627282930}, walked through method-by-method after the table.

The table groups the literature into five generations, each a response to the feasibility-vs-fidelity tension. The sub-sections below follow the same logic. The table below summarizes the structural trade-offs of each lineage; individual method numbers are in the SoTA leaderboard above.

Generation 1 (1955-2000): Exact LP and Analytic Foundations

Classical exact-LP methods (Hungarian, network simplex) solve the discrete Kantorovich problem to numerical zero error, but the $O(n^3 \log n)$ cost makes them prohibitive beyond $n \approx 10^4$ source/target points³¹; this cubic wall is the principal motivation for every subsequent computational development. The same era produced the analytic theory: Brenier's (Brenier 1991) polar factorization established that, for absolutely continuous source measures and squared-Euclidean cost, the optimal map exists, is unique, and equals the gradient of a convex potential; Gangbo-McCann (Gangbo and McCann 1996) generalized this to any strictly convex superlinear cost and identified $c$-cyclical monotonicity as the universal characterization of optimal couplings. These foundational results give exact answers in closed form for tractable cases (Gaussians via Bures-Wasserstein, 1D via quantile composition) but provide no algorithm for the high-dimensional empirical regime that ML faces.

The fidelity of exact-LP is maximal (zero relative error) but the feasibility cost is a cubic barrier. Generation 2 resolves this by accepting a tunable bias in exchange for near-linear parallel complexity.

Generation 2 (2013-2018): Entropic Regularization and the Sinkhorn Revolution

Entropic-OT via the Sinkhorn iteration (Cuturi 2013) is the algorithmic pivot of modern computational OT: the addition of a small KL-regularizer $\varepsilon \mathrm{KL}(\gamma \| \mu \otimes \nu)$ makes the problem strictly convex, the optimal coupling factors as $P^* = \mathrm{diag}(\mathbf{u}) K \mathrm{diag}(\mathbf{v})$, and matrix-scaling iterations converge linearly. The complexity drops from $O(n^3 \log n)$ to $\tilde{O}(n^2/\varepsilon^2)$ (Dvurechensky, Gasnikov, and Kroshnin 2018), with massive parallelism across rows and columns, and the iterations are differentiable through the converged dual potentials; these three properties (speed, parallelism, differentiability) are jointly responsible for OT's penetration into deep learning. The price is the entropic bias: $W_\varepsilon \neq W$ and the entropic plan is "blurred" by an amount controlled by $\varepsilon$³². Sinkhorn divergences $S_\varepsilon = W_\varepsilon(\mu,\nu) - \tfrac12 W_\varepsilon(\mu,\mu) - \tfrac12 W_\varepsilon(\nu,\nu)$ (Genevay, Peyré, and Cuturi 2018) correct this debiasing at the cost of three solver runs; epsilon-scaling (geometric annealing of $\varepsilon$) recovers near-exact OT in roughly an order of magnitude fewer total iterations.

Entropic OT still scales as $O(n^2)$ in the cost matrix, which dominates memory and compute when $n$ is large and the ground cost is expensive to evaluate. Generation 3 sidesteps the cost matrix entirely by projecting to 1D.

Generation 3 (2018-2019): Sliced OT and Dimension-Free Projections

Sliced and projection-based methods (Kolouri et al. 2018; Le et al. 2019) sidestep the $O(n^2)$ cost matrix by averaging univariate Wasserstein distances over directions of the unit sphere $S^{d-1}$: $SW_p^p(\mu,\nu) = \int_{S^{d-1}} W_p^p(\theta_\#\mu, \theta_\#\nu)\,d\theta$. Each 1D projection sorts in $O(n \log n)$, so $L$ projections cost $O(Ln \log n)$; for typical $L = 50\text{-}1000$ this is two to three orders of magnitude faster than full OT. The dimension-free $O(n^{-1/2})$ sample complexity of sliced-Wasserstein (vs the $O(n^{-1/d})$ curse of full OT³³) is its main statistical advantage and the reason it scales gracefully from 1D to 100D. The trade-off is projection bias³⁴: SW is a true metric but only equivalent to $W_2$ for Gaussian distributions; on data concentrated on low-dimensional manifolds, projection averages over informative-and-uninformative directions equally and underestimates the true distance.

Both sliced and entropic methods still assume the source and target live in the same metric space and carry equal total mass. Many practical problems violate both assumptions. Generation 4 relaxes them.

Generation 4 (2017-2022): Unbalanced OT, Gromov-Wasserstein, and Neural Potentials

Unbalanced OT (Chizat et al. 2018) relaxes the equal-mass marginal constraint of classical OT; required whenever source and target have different totals (single-cell snapshots, partial-shape correspondences, class-imbalanced datasets)³⁵. The relaxation replaces hard marginals with soft KL or TV penalties:

which admits a scaling iteration analogous to Sinkhorn. Gromov-Wasserstein (Mémoli 2011) handles the dual difficulty: source and target live in different metric spaces (graphs of different sizes, point clouds in different ambient dimensions), so there is no shared cost $c(x,y)$; GW instead matches pair-distances via a quadratic objective $\min_\gamma \sum_{ijkl} |d^X_{ik} - d^Y_{jl}|^2 \gamma_{ij}\gamma_{kl}$. The objective is a quadratic assignment problem, NP-hard in general and non-convex in the coupling; entropic GW solvers depend strongly on initialization, with random / spectral / identity warm-starts producing different local optima³⁶.

The unbalanced and GW extensions are not Pareto-dominated by classical OT; they enable problems (unequal-mass biology, graph matching) that classical OT cannot even formulate. But all of the above methods require either a discrete cost matrix or a fixed-support approximation, blocking application to the continuous high-dimensional distributions that arise in generative modeling. Generation 5 abandons the matrix entirely.

Neural-OT methods (Korotin, Selikhanovych, and Burnaev 2023; Seguy et al. 2018; Amos, Xu, and Kolter 2017) parameterize Kantorovich potentials $f, g$ (or the Brenier potential $\phi$) as neural networks and solve the dual saddle-point min-max objective by stochastic gradient ascent-descent. Input convex neural networks (ICNNs) (Amos, Xu, and Kolter 2017) enforce the convexity required by Brenier's theorem via non-negative weights and convex non-decreasing activations; this guarantees the gradient $T(x) = \nabla \phi_\theta(x)$ is a valid Brenier map. The price is the adversarial training instability of all GAN-class objectives³⁷: hyperparameter sensitivity, mode collapse, and oscillation; mitigated by careful learning-rate scheduling, gradient penalties, and warm-starting from Sinkhorn-derived initial maps. JKOnet (Bunne et al. 2022) uses ICNNs to discretize Wasserstein gradient flows via the JKO scheme, learning energy landscapes from population snapshots.

Generation 5 (2022-2024): Flow-Matching, Schrödinger Bridges, and Unification

OT flow-matching and Schrödinger-bridge methods are the 2022-2024 generative-modelling synthesis. Flow-matching trains a velocity field whose ODE integration transports noise to data; the OT-conditional variant uses mini-batch OT couplings (rather than the diffusion's stochastic noise schedule) to define straighter probability paths, which simpler velocity fields can learn and which integrate in fewer ODE solver steps³⁸. Schrödinger bridges solve the entropy-regularized OT problem in path-measure space via iterative proportional fitting (the dynamic analog of Sinkhorn), interpolating between deterministic OT ($\varepsilon \to 0$) and pure diffusion ($\varepsilon \to \infty$).

The five lineages (exact-LP, entropic, sliced, unbalanced/GW, and neural/flow) converge on one principle: recover a coupling or transport map that minimizes a transport cost under the marginal and feasibility constraints your compute budget can enforce. Brenier potentials, Schrödinger bridges, and continuous-time normalizing flows are different parameterizations of the same underlying displacement-interpolation theory introduced by Brenier (Brenier 1991) and Benamou-Brenier (Benamou and Brenier 2000) three decades earlier. The figure below maps all five lineages onto a feasibility-vs-fidelity trade-off plane.

References