PaperlayerView exact revision on arXiv

Source-linked AI summary

Tutorial on amortized optimization

Brandon Amos

arXiv:2202.00665v5cs.LGcs.AImath.OC

TL;DR

Repeatedly solving similar optimization problems creates a computational challenge. This tutorial presents amortized optimization, which learns models to predict solutions; one example is 25,000 times faster than solving the same-quality batch from scratch.

  • Problem

    The tutorial addresses how machine learning can improve repeated solves of parametric optimization problems.

    Source: Introduction · p. 3

  • Method

    Amortized optimization trains models to rapidly predict solutions to optimization problems.

    Source: Introduction · p. 4

  • Results

    25,000 times faster: the amortized variational autoencoder component takes 0.4 ms versus 8 seconds to solve 1,024 same-quality optimization problems from scratch.

    Source: Introduction · pp. 5–6

  • Takeaways & Limitations

    Amortized optimization can significantly reduce computation on a focused subset of problems by learning their solution mapping.

    Source: Introduction · p. 4

    Amortized optimization struggles in some settings, where learning the model can cost more than solving the original problems and introduce instabilities atop potentially inaccurate solutions.

    Source: §5.4 · p. 71

Abstract

from arXiv · show

Optimization is a ubiquitous modeling tool and is often deployed in settings which repeatedly solve similar instances of the same problem. Amortized optimization methods use learning to predict the solutions to problems in these settings, exploiting the shared structure between similar problem instances. These methods have been crucial in variational inference and reinforcement learning and are capable of solving optimization problems many orders of magnitudes faster than traditional optimization methods that do not use amortization. This tutorial presents an introduction to the amortized optimization foundations behind these advancements and overviews their applications in variational inference, sparse coding, gradient-based meta-learning, control, reinforcement learning, convex optimization, optimal transport, and deep equilibrium networks. The source code for this tutorial is available at https://github.com/facebookresearch/amortized-optimization-tutorial.

Introduction

This tutorial presents amortized optimization as learning to predict solutions across related optimization contexts, using shared structure to reduce repeated computational costs, and develops a unified view of modeling choices and applications such as variational inference, control, and reinforcement learning. It also explains differentiation through semi-amortized models, selecting models for application contexts and solution spaces, and implementation choices for higher-order derivatives and variable-sized inputs or outputs.

Introduction

  • Amortized optimization learns to predict solutions from context, exploiting a solution mapping that classical optimization methods usually do not assume access to.The tutorial uses this unified perspective to connect applications including variational inference, meta-learning, policy learning, sparse coding, convex optimization, optimal transport, and deep equilibrium networks.Source: Introduction · pp. 4–5
  • For MNIST variational autoencoders, amortization is 25,000 times faster—0.4ms versus 8 seconds—than solving 1,024 problems from scratch at the same solution quality.For millions of training iterations, this difference can separate a process finishing in hours from one taking a month.Source: Introduction · pp. 5–6

Amortized optimization foundations

2.1 Defining the model ˆyθ(x)

  • The model maps contexts to predicted solutions, and its design balances expressivity and computational efficiency while determining how much objective information it uses.The tutorial distinguishes black-box fully-amortized models from semi-amortized models that access the optimization objective.Source: §2.1 · pp. 8–9

2.1.1 Fully-amortized models

  • Fully-amortized models map contexts to solutions without accessing the objective, and work best when solutions are predictable, the domain is modest-sized, and the context distribution is not too large.They can predict solutions without querying the objective after training.Source: §2.1.1 · p. 9

2.1.2 Semi-amortized models

  • Semi-amortized models iteratively access the objective, integrating optimization procedures over the solution domain or a latent space that can represent solution structure.Latent spaces can let the optimizer search an optimal region rather than the entire solution space.Source: §2.1.2 · p. 9, Semi-amortized models over a latent… · pp. 10–11
  • Unlike warm-starting, semi-amortized learning can train end-to-end through the final prediction, but unrolling iterative models can require costly second-order derivatives.For K steps, the objective Hessian is required at every iteration, increasing computational burden.Source: Comparing semi-amortized models with… · pp. 11–12, On second-order derivatives of the… · pp. 12–13
  • First-order MAML omits second-order terms, while MAML++ and Reptile offer related approximations; memory-efficient unrolling can also reduce training memory use.First-order MAML works well in some cases, including MiniImagenet; MAML++ [Antoniou et al., 2019] adds second-order information later in training, while Reptile [Nichol et al., 2018] also avoids it. Memory-efficient methods for backpropagation and unrolling are described by Gruslys et al. [2016] and Chen et al. [2016].Source: On second-order derivatives of the… · p. 13
  • When the objective is learned, its second-order derivatives can pose issues; Belanger et al. [2017] use softplus activations to keep them non-zero.They learn both an initial iterate and an energy-function objective parameterized by a neural network.Source: On second-order derivatives of the… · p. 13

2.1.3 Models based on differentiable optimization

  • Semi-amortized models differentiate through optimization updates to obtain parameter derivatives, typically involving second-order derivatives; differentiable optimization tools enable end-to-end learning with more general methods.When updates are objective-based, differentiation through the sequence of gradient steps supplies the derivatives needed to optimize the loss.Source: §2.1.3 · p. 14

2.1.4 Practically choosing a model

  • Model choice depends on speed, expressibility, context and solution spaces, and variable-length requirements; VeLO handles varying parameter counts by predicting smaller MLPs over parameter groups.For image-valued solutions, the passage gives U-nets, dilated convolutional networks, and fully convolutional networks as suitable model examples.Source: §2.1.4 · p. 14

2.2 Learning the model’s parameters θ

After specifying the amortization model, learning its parameters is a major design choice, often framed as an outer optimization over the model and an inner problem over the solution domain. Loss design depends on whether learning uses ground-truth solutions, objective information, or iterative model computations.

2.2 Learning the model’s parameters θ

  • Parameter learning is often bilevel: an outer loss trains ˆyθ(x) to solve the inner optimization problem over Y.The best loss is application-specific.Source: §2.2 · p. 14

2.2.1 Choosing the objective for learning

  • Regression losses minimize distance to ground-truth solutions, whereas objective-based losses use local information from f without access to y⋆.Regression-based learning typically differentiates the loss with respect to model parameters; objective-based learning can use first-order gradients when differentiation is available.Source: Figure 2.2 · p. 15, Regression-based learning · p. 15, Objective-based learning · p. 16
  • Regression can distill known solutions and avoid evaluating costly objectives, but depends on tractable ground truths; objective-based learning avoids y⋆ but may need gradients and get stuck in local optima.Regression can also fail when solutions are nonunique or too intractable to compute.Source: Regression-based learning · pp. 15–16, Comparing the regression- and… · p. 18
  • Equal ℓ2-regression loss can correspond to substantially different objective values, and combining regression with objective losses may help when only a few contexts are labeled.Reweighting or normalizing regression dimensions to reflect objective curvature is another proposed approach.Source: Comparing the regression- and… · pp. 17–18, Figure 2.3 · p. 17

2.2.2 Learning iterative semi-amortized models

  • Truncated backpropagation through time (TBPTT) improves compute and memory efficiency by differentiating over only the most recent iterates, but introduces biased gradients.This bias is especially damaging for MAML, which parameterizes only the first iterate.Source: Truncated backpropagation through time… · pp. 19–20
  • Implicit MAML adds a penalty around the initial iterate to make the inner problem globally solvable, then uses implicit differentiation instead of unrolling.Implicit differentiation requires an exact solution satisfying the implicit function theorem; the penalty changes the objective, and solving the Jacobian system may be costly.Source: Semi-amortized learning with shrinkage… · p. 21

2.2.3 Learning with zeroth-order methods and RL

  • Zeroth-order methods use local perturbations when gradients are impossible or unstable; reinforcement-learning formulations treat iterates as actions, losses as rewards, and the model as a policy.The context and previous iterates or losses typically form the states.Source: §2.2.3 · p. 22, Reinforcement learning · p. 22
  • Gaussian smoothing can reduce poor local minima, but may shift the global minimum; decaying the smoothing during training is one proposed remedy.The tutorial describes fixed-variance smoothing with reparameterization gradients or zeroth-order evolutionary methods.Source: Loss smoothing and optimization with… · pp. 22–23, Figure 2.6 · p. 23

2.3 Extensions

The tutorial starts from deterministic, unconstrained Euclidean problems with a fixed context distribution, then surveys extensions to the optimization domain, model, and learning setup. These extensions include stochastic and constrained problems, alternative model types, uncertainty-aware methods, and additional learnable parameters.

2.3 Extensions

  • The starting scope is deterministic, unconstrained, finite-dimensional Euclidean optimization with a fixed context distribution; the section presents extensions that may arise in practice.The tutorial is mostly scoped to differentiable parametric models that are learned end to end.Source: §2.3 · pp. 23–24, §2.3.2 · p. 29

2.3.1 Extensions of the domain Y

  • Optimization over stochastic domains can treat the domain as distributions; for parametric families, optimizing their parameters can return the problem to Euclidean optimization.More general density optimization may lack a closed-form solution, but this tutorial does not study amortizing those problems.Source: Deterministic →stochastic optimization · pp. 24–25
  • Differentiable projections can recast constrained optimization as an unconstrained problem, while numerically computed projections can often be differentiated implicitly.Closed-form examples include ReLU, sigmoid, and softargmax; the defined projections are not generally idempotent.Source: Unconstrained →constrained optimization · pp. 26, 28
  • The domain extensions also cover differentiable projections onto convex cones and amortization over manifolds, including spherical problems transformed through ambient-space projections.Some cone projections have closed forms, while the exponential-cone projection can be computed with Newton's method and differentiated implicitly.Source: Unconstrained →constrained optimization · pp. 27–28, Euclidean →non-Euclidean optimization · p. 29

2.3.2 Extensions of the model ˆyθ

  • Symbolic models search interpretable update rules, and learned optimizers have exhibited behaviors including momentum, gradient clipping, and learning-rate adaptation.The tutorial contrasts this direction with neural-network amortization models, which it describes as uninterpretable.Source: Symbolic models: uncovering… · pp. 29–30

2.3.3 Uncertainty-aware and Bayesian optimization

  • Amortized Bayesian optimization predicts solutions to acquisition functions, while Bayesian learning to optimize studies uncertainty about whether an amortization model is the best optimizer.Applications described include protein sequence design, contextual bandits, and few-shot learning.Source: §2.3.3 · p. 30

2.3.4 Settings with additional learnable contexts φ

  • When parameters φ affect the context distribution or objective, a higher-level loss on solutions can define a bilevel learning problem.The solution derivative can be computed implicitly using an optimality system; the objective value can also be differentiated under mild assumptions using the envelope theorem.Source: §2.3.4 · p. 30, Learning φ by differentiating the… · pp. 31–32

Applications of amortized optimization

This chapter reviews applications of amortized optimization to highlight shared ideas across topics. Its subsections are designed to be read independently, but provide relevant context only for the amortization components.

Applications of amortized optimization

3.1 Variational inference and variational autoencoders

Variational inference uses optimization to approximate intractable densities and integrals. In a VAE, amortized variational inference predicts variational parameters that maximize an evidence lower bound for each input sample.

3.1 Variational inference and variational autoencoders

  • Variational inference uses optimization to approximate intractable densities and integrals.This section focuses on amortized variational inference in VAEs and related generative models.Source: §3.1 · p. 33

3.1.1 The variational autoencoder (VAE) by Kingma and Welling [2014]

  • A VAE uses a tractable evidence lower bound because computing the likelihood integral over the latent space is typically intractable.The ELBO combines expected log likelihood with a KL-divergence penalty against the prior.Source: §3.1.1 · p. 34
  • Amortized variational inference learns a model that predicts the variational parameters maximizing the ELBO for each sample, rather than explicitly solving that optimization for each sample.The prediction model is usually a feedforward neural network and is trained with an objective-based loss.Source: §3.1.1 · pp. 34–35
  • Semi-amortization methods such as SAVAE and IVAE address the amortization gap, while more expressive variational distributions can address approximation gaps.Source: §3.1.1 · p. 35
  • The section fixes the model parameters to isolate the amortized optimization component, although full VAE training also optimizes the ELBO over the model parameters.Source: §3.1.1 · p. 35

3.2 Sparse coding

Sparse coding reconstructs an input from a sparse combination of basis vectors, traditionally using ISTA or FISTA. Predictive Sparse Decomposition and LISTA amortize this optimization in distinct ways, with LISTA incorporating iterative updates into its model.

3.2 Sparse coding

  • Sparse coding recovers a code by reconstructing an input from a sparse linear combination of basis vectors, traditionally using ISTA or momentum-accelerated FISTA.ISTA methods have also motivated learning-based approaches to speed up sparse coding.Source: §3.2 · p. 36

3.2.1 Predictive Sparse Decomposition (PSD) by Kavukcuoglu et al. [2010]

  • Predictive Sparse Decomposition predicts the best sparse code with a fully amortized model trained by regression to the true code.PSD also iteratively approximates the code while learning the model, rather than solving for it directly with (F)ISTA.Source: §3.2.1 · p. 36

3.2.2 Learned ISTA (LISTA) by Gregor and LeCun [2010]

  • LISTA proposes a semi-amortized model that integrates ISTA's iterative updates and soft-thresholding operator, then predicts a sparse code after K steps.Its parameters are optimized with a regression-based loss to ground-truth sparse codes.Source: §3.2.2 · pp. 36–37

3.3 Multi-task learning and meta-learning

Multi-task and meta-learning methods amortize parameter optimization by exploiting shared structure across tasks. The section contrasts fully amortized prediction with iterative refinement and surveys applications and related methods.

3.3.1 Fully-amortized hyper networks (HyperNets)

  • HyperNets predict model parameters from a task’s data without accessing its loss, learning through an objective-based loss.Tasks pair a data sample with its classification or regression target.Source: §3.3.1 · pp. 37–38

3.3.2 Learning to optimize (LM) by Li and Malik [2017a]

  • Li and Malik [2017a] refine dataset-level task parameters iteratively, using recent iterates, objectives, gradients, and objective improvements to predict updates.Their experiments use H = 25 and typically run 100 model updates; they frame learning as an MDP solved with guided policy search.Source: §3.3.2 · p. 38

3.3.3 Model-agnostic meta-learning (MAML) by Finn et al. [2017]

  • MAML [Finn et al., 2017] predicts an initial task-specific parameter iterate, then applies gradient updates and optimizes an objective-based loss through the final prediction.This makes MAML a semi-amortized method in this tutorial’s framing.Source: §3.3.3 · p. 38

3.3.4 Protein MSA modeling with the Neural Potts Model

  • Sercu et al. [2021] find amortization yields better Potts-model parameter solutions than classical optimization on finite MSAs, attributing the gain to information sharing across proteins.Their fully amortized attention-based model predicts Potts parameters from each protein’s MSA and optimizes pseudolikelihood.Source: §3.3.4 · p. 39

3.3.5 Other relevant multi-task and meta-learning work

  • Ravi and Larochelle [2017] use recurrent update prediction; LEO [Rusu et al., 2019] and CAVIA [Zintgraf et al., 2019] amortize over a learned latent space.Other examples include differentiable SVM and ridge-regression layers in MetaOptNet and R2D2.Source: §3.3.5 · pp. 39–40
  • Later work studies optimizer stability and scaling, including Jacobian-spectrum issues [Metz et al., 2021] and learned updates for models with millions of parameters [Metz et al., 2022].Other directions include atomic-structure optimization [Merchant et al., 2021] and architecture search with graph-based HyperNets [Zhang et al., 2019a, Knyazev et al., 2021].Source: §3.3.5 · pp. 40–41

3.4 Fixed-point computations and convex optimization

Fixed-point problems connect to unconstrained optimization through residual minimization, and classical iterations can be accelerated using learned updates. The section reviews learned acceleration for cone programs, matrix factorization, deep equilibrium models, and quadratic programs.

3.4 Fixed-point computations and convex optimization

  • A fixed point satisfies g(y⋆) = y⋆; fixed-point problems and unconstrained optimization can be transformed into one another, and classical iterations include direct updates and Anderson acceleration.Anderson acceleration combines past iterates, while Broyden’s method treats the problem as root finding.Source: §3.4 · pp. 41–42

3.4.1 Neural fixed-point acceleration (NeuralFP) and conic optimization with the splitting cone solver (NeuralSCS)

  • NeuralFP learns an initial iterate and recurrent updates to amortize fixed-point residual minimization; NeuralSCS applies this approach to convex cone programs.SCS represents the cone program as a fixed-point computation using a homogeneous self-dual embedding.Source: §3.4.1 · p. 43

3.4.2 Neural acceleration for matrix factorization (NeuralNMF)

  • NeuralNMF [Sjölund and Bånkestad, 2022] adds transformer-based initialization and acceleration to ADMM iterations for nonnegative low-rank matrix factorization.The learned model is trained with an objective-based loss and unrolls through the ADMM iterations.Source: §3.4.2 · p. 44

3.4.3 HyperAnderson acceleration and deep equilibrium models (HyperDEQ)

  • HyperAnderson [Bai et al., 2022] learns an initial point and Anderson-update coefficients, using a summed regression loss that increasingly penalizes later iterates that miss the fixed point.The method is used to semi-amortize deep equilibrium models.Source: §3.4.3 · pp. 44–45

3.4.4 Comparing NeuralFP and HyperAA

  • NeuralFP learns the entire update with an objective-based loss, whereas HyperAA learns Anderson coefficients with a regression-based loss.Both methods learn an initial iterate and can semi-amortize fixed-point computations.Source: §3.4.4 · p. 45

3.4.5 RLQP by Ichnowski et al. [2021]

  • RLQP [Ichnowski et al., 2021] learns OSQP’s step size as a policy over solver states, training with TD3 and a reward tied to primal and dual residual thresholds.The policy controls the step size; the remaining updates follow OSQP.Source: §3.4.5 · pp. 45–46

3.5 Optimal transport

Optimal transport includes repeatedly solved optimization problems whose solutions or subproblems can be amortized. The section surveys methods for predicting transport solutions, accelerating convex conjugates, and optimizing sliced Wasserstein projections.

3.5 Optimal transport

  • The Kantorovich formulation optimizes over admissible couplings, while its dual uses potentials; discrete measures permit matrix and vector representations and linear-programming or Sinkhorn solutions.With negative inner-product cost in Euclidean spaces, the c-transform is the Legendre-Fenchel transform, or convex conjugate.Source: §3.5 · pp. 46–47

3.5.1 Amortizing solutions to optimal transport problems

  • Amortized OT methods predict optimal couplings or dual solutions across repeated transport problems, while WGAN potentials are warm starts for changing problems rather than amortized solutions.Meta OT [Amos et al., 2022] predicts dual variables from measures and costs, then warm-starts standard solvers; CondOT [Bunne et al., 2022] conditions on patient and drug context to predict effects without the target measure.Source: §3.5 · p. 47, §3.5.1 · pp. 47–48

3.5.2 Amortized convex conjugates (AmorConj) and c-transforms

  • Amortized convex conjugates address repeated c-transform optimization, which can be costly in continuous settings when each Monte Carlo estimate requires solving it from scratch.Methods amortize the conjugate operation, and numerical fine-tuning can improve accuracy of the dual objective estimate.Source: §3.5.2 · pp. 48–49

3.5.3 Amortized Sliced Wasserstein (A-SW) by Nguyen and Ho [2022]

  • Amortized Sliced Wasserstein [Nguyen and Ho, 2022] learns the max-sliced optimization over mini-batches, projecting measures onto one-dimensional slices for more efficient distance computation.The one-dimensional 1-Wasserstein distance has a closed-form solution.Source: §3.5.3 · p. 49

3.6 Policy learning for control and reinforcement learning

This section examines how control and reinforcement-learning methods amortize optimization solutions into policies, from behavioral cloning to semi-amortized policy updates, including settings with unknown dynamics or Q functions. It also covers learning value estimates that satisfy the Bellman equation and using DCEM to learn latent control-sequence spaces for semi-amortized controllers.

3.6.1 Background

  • For unknown dynamics, reinforcement learning typically alternates policy evaluation from system samples with policy improvement, while stochastic policies can optimize expected Q or KL distance to a Q-distribution.In the fully observable continuous-MDP setting, deterministic policies can instead be obtained by solving the max-Q problem; Dirac stochastic policies recover that deterministic solution.Source: §3.6.1 · pp. 51–52

3.6.2 Behavioral cloning and imitation learning

  • Behavioral cloning is regression-based amortization that distills an expert’s state-to-action behavior, but it does not require knowing the expert’s original objective or cost.The objective is assumed to exist for the amortization framing; inverse control and reinforcement learning that recover it are outside this section’s scope.Source: §3.6.2 · pp. 53–54

3.6.3 The deterministic policy gradient

  • Deterministic policy-gradient methods, including DPG, DDPG, and TD3, perform objective-based amortization of the max-Q optimization through interleaved policy evaluation and improvement.They estimate the current policy’s state-action value Q and improve the policy to optimize the max-Q objective.Source: §3.6.3 · p. 54

3.6.4 The stochastic value gradient and soft actor-critic

  • SVG- and SAC-based policy updates perform objective-based amortization of stochastic control optimization, with SAC adding entropy regularization to soften policy and value estimates.These methods optimize stochastic policies over a state distribution using expected-Q or KL-to-Q objectives; SVG methods span model-free and model-based value estimation.Source: §3.6.4 · pp. 54–55

3.6.5 PILCO by Deisenroth and Rasmussen [2011]

  • PILCO optimizes a deterministic RBF policy using gradients through a Gaussian-process estimate of finite-horizon Q, while GPS fits an approximate dynamics model and amortizes local MPC controller solutions.GPS uses controller samples to optimize KL divergence between the controller and learned policy; the supplied passage identifies this as its summary formulation.Source: §3.6.5 · p. 55, §3.6.6 · p. 56

3.6.7 POPLIN by Wang and Ba [2020]

  • POPLIN extends behavioral cloning by using a learned policy to initialize control optimization or by adapting its parameters, while DCEM enables differentiable semi-amortized control in a learned latent space.POPLIN-A initializes action-space CEM with a predicted sequence; POPLIN-P searches policy-parameter perturbations. DCEM smooths CEM’s top-k operation, allowing decoder parameters to receive gradients through latent optimization.Source: §3.6.7 · p. 56, §3.6.8 · pp. 57–58

3.6.9 Iterative amortized policy optimization (IAPO) by Marino et al. [2021]

  • IAPO replaces fully amortized DPG/SVG policies with iterative semi-amortized policies that take gradient steps on actions in the underlying control optimization.The approach is explored in both model-free and model-based reinforcement learning; iterative policies can worsen exploitation of optimistic inaccuracies in the value estimator.Source: §3.6.9 · p. 58

3.6.10 Learning the value function

  • Actor-critic methods jointly learn a policy that amortizes the current Q estimate and a critic fitted to system data; the section begins describing a Q-based replacement for the value estimate.The supplied passage introduces replacing Vπ with a Q-value estimate but does not provide the resulting relationship.Source: §3.6.10 · pp. 58–59
  • The Bellman equation [Bellman, 1966, Sutton and Barto, 2018] should hold across all system states and actions, and a parameterized value estimate can be learned to satisfy it.A common approach is to optimize the residual of the equation.Source: §3.6.10 · p. 59
  • DCEM [Amos and Yarats, 2020] learns a latent space that generates smoother, more cyclic control-sequence distributions than CEM’s initially temporally uninformed samples.In the cartpole visualization, DCEM’s distributions become more feasible across iterations.Source: Figure 3.3 · p. 60

Implementation and software examples

The section examines how amortization software is built and evaluated, using existing variational-autoencoder and control implementations plus a spherical optimization example.

Implementation and software examples

4.1 Amortization in the wild: a deeper look

The section inspects amortization code and evaluates its speed in a VAE and two humanoid-control settings. The amortized solutions are reported to be substantially faster than solving each optimization problem from scratch.

4.1 Amortization in the wild: a deeper look

  • The section evaluates the amortization components of existing VAE and control implementations, alongside code for a spherical optimization problem.The VAE and policy-learning systems also perform larger tasks, but the evaluation focuses on their amortization components.Source: §4.1 · p. 61, Table 4.1 · p. 62

4.1.1 The variational autoencoder (VAE)

  • The VAE code defines an encoder that predicts an ELBO solution for image data with a Bernoulli likelihood and optimizes it over training images.The walkthrough uses an implementation based on standard PyTorch and evaluates a pretrained MNIST model.Source: §4.1.1 · pp. 62–63
  • 25000 times faster: the VAE encoder predicts the ELBO solution than running 2k Adam iterations on a batch of 1024 samples.The comparison is for MNIST; the encoder’s solution quality is assessed alongside explicitly optimized solutions and decoded samples.Source: Figure 4.1 · p. 62, §4.1.1 · p. 64, Figure 4.3 · p. 64

4.1.2 Control with a model-free value estimate

  • The model-free control example learns a deterministic policy mapping humanoid states to actions that maximize a known action-conditional value estimate.The humanoid MDP has 45 states and 17 actions; the policy is compared with maximizing the value estimate from scratch for encountered states.Source: §4.1.2 · pp. 64–65
  • 1000 times faster: the learned humanoid policy solves the model-free value-estimate optimization than solving it from scratch.Values are normalized and clamped so the expert action is 0 and the zero action is -1; the discussion flags possible approximation errors in the Q estimate.Source: §4.1.2 · pp. 65–66, Figure 4.5 · p. 66

4.1.3 Control with a model-based value estimate

  • The model-based control comparison uses a short-horizon optimization whose action sequence can be approximated by interleaving a model-free policy with the dynamics.Adam is used as a gradient-based shooting baseline in this short-horizon setting.Source: §4.1.3 · p. 66
  • 700 times faster: the learned policy solves the humanoid short-horizon model-based control problem than solving it from scratch.This comparison uses horizon H = 3; model and model-based value-estimate approximation errors are noted.Source: §4.1.3 · p. 66, Figure 4.6 · p. 67

4.2 Training an amortization model on a sphere

The tutorial demonstrates amortized optimization on a new class of sphere-constrained problems, training a neural network to predict their solutions.

Amortization objective

  • The sphere constraint can be handled by projecting an unconstrained model’s output onto the sphere.The experiment treats sphere optimization as a transformation to an unconstrained problem and composes the model with the sphere projection.Source: §4.2 · pp. 67–68

c-convex functions on the sphere

  • The demonstration constructs synthetic sphere problems from c-convex functions, with contexts sampled using points uniform on the sphere and normally distributed coefficients.The context also uses a soft minimization operator.Source: §4.2 · pp. 67–68

Optimizing the gradient-based loss

  • A single neural network trained with a gradient-based loss predicts solutions to the sphere problems from their contexts.Figure 4.7 shows predictions from the untrained model through the trained model.Source: §4.2 · p. 68, Figure 4.7 · p. 68

4.3 Other useful software packages

The tutorial surveys software packages that support differentiable optimization, semi-amortized models, and related computational operations.

Differentiable optimization packages

  • cvxpylayers exports CVXPY optimization problems to PyTorch, JAX, and TensorFlow as differentiable optimization layers.jaxopt supports differentiable optimization and fixed-point computations with implicit derivatives.Source: §4.3 · pp. 68–69

Semi-amortized optimization

  • higher enables unrolling PyTorch modules and optimizers as semi-amortized models, while TorchOpt provides a functional differentiable optimizer.TorchOpt has higher performance than higher in some cases.Source: §4.3 · p. 69

Related tools

  • functorch provides composable batching and derivative transforms, while DiffOpt.jl brings differentiable optimization to Julia’s JuMP.Torchmeta and learn2learn support meta-learning tasks and data loading, while hypertorch computes hypergradients with efficient approximations.Source: §4.3 · p. 69

Discussion

The discussion adds further considerations about amortized optimization alongside application-specific tradeoffs and limitations.

Discussion

  • The tutorial’s discussion complements application-specific tradeoffs with additional considerations about amortized optimization.It points to broader treatments of learning-to-optimize paradigms by Chen et al. [2021a] and Metz et al. [2021].Source: Discussion · p. 70

5.1 Surpassing the convergence rates of classical methods

Amortized optimization can outperform classical methods on trained context distributions, but its value depends on generalization, convergence, performance measurement, and application conditions.

5.1 Surpassing the convergence rates of classical methods

  • Amortized methods can surpass classical worst-case convergence rates by tuning initialization and updates to best-case contexts in the training distribution.Some fully amortized variational-inference and actor-critic models solve their optimization problems with one prediction, without evaluating the objective.Source: §5.1 · p. 70

5.2 Generalization and convergence guarantees

  • Amortized optimization can struggle to generalize and converge to reasonable solutions, although checking residuals can allow poor instances to be flagged and re-solved with a standard solver.Semi-amortized models based on gradient descent are described as more stable than those based on sequential models; layer normalization can also improve stability.Source: §5.2 · pp. 70–71

5.3 Measuring performance

  • Low objective values after a few iterations do not establish a useful speedup if reaching the same accuracy takes longer than a specialized algorithm or another amortized model.Performance measures are often tied to problem-specific metrics.Source: §5.3 · p. 71

5.4 Successes and limitations of amortized optimization

  • Despite successes in variational inference, reinforcement learning, and meta-learning, amortization may cost more to train and introduce instability into higher-level processes using inaccurate solutions.The tutorial frames these as limitations that arise in some settings, not all applications.Source: §5.4 · p. 71

Characteristics of successful applications

  • Successful applications favor smooth objectives with unique solutions, tolerate suboptimal early predictions, and restrict contexts to a well-scoped problem class.Examples include amortizing VAE objectives over their training dataset and RL control problems over a specific MDP’s replay buffer.Source: Characteristics of successful… · p. 72
  • Semi-amortized models often learn an initial iterate or selected update components rather than parameterizing an entire update rule.Examples include MAML’s learned initial iterate and LISTA’s parameterized components {F, G, β}.Source: Characteristics of successful… · p. 72
  • Pure sequence models for updating iterates can be unstable and overfit, with recurrent models observed to induce chaotic behavior and exploding gradients [Metz et al., 2021].Source: Characteristics of successful… · p. 72

Limitations and failures

Amortized optimization is most useful when similar problems must be solved repeatedly, but it can struggle when the context or solution mapping is difficult to learn or when high accuracy is essential. The tutorial identifies open research directions, distinguishes its scope from related reviews, and notes areas outside its coverage.

Limitations and failures

  • Amortization is suited to repeated, structurally similar solves, but can struggle with hard-to-learn context–solution mappings and domains requiring accuracy beyond what it typically achieves.Non-unique solutions may be easier to handle with carefully designed objective-based losses; safety-critical model-free reinforcement learning also requires optimizing value estimates that encode safety properties.Source: Limitations and failures · p. 73

5.5 Some open problems and under-explored directions

  • Objective smoothing and stochastic policies offer related approaches to overcoming poor local minima, while GECO adjusts an ELBO Lagrange multiplier to target conditional log-likelihood.The tutorial suggests generalizing these approaches to other amortization settings.Source: §5.5 · pp. 73–74
  • Amortized convex solvers remain an open direction because current work does not scale to general online use without imposing substantial learning overhead.The cited directions study adding amortization to off-the-shelf optimization packages.Source: §5.5 · p. 74
  • Open problems also include improving training-time gains for implicit and differentiable models, measuring the amortization gap, and transferring implicit-differentiation and shrinkage ideas across settings.Existing amortization work in implicit and differentiable settings often does not improve wall-clock training time from scratch; the amortization gap measures approximation to the true solution.Source: §5.5 · p. 74
  • The tutorial flags distribution shift and out-of-distribution generalization, and describes semi-amortized control and reinforcement learning as budding areas for sample-efficient controllers.Marino et al. [2021] show better amortization of solutions with gradient-descent-based semi-amortization than standard fully-amortized models.Source: §5.5 · pp. 74–75

5.6.1 Other tutorials, reviews, and discussions on amortized optimization

  • This tutorial organizes amortized optimization by modeling and learning approaches; related tutorials and reviews cover adjacent topics including algorithm unrolling and meta-learning.Monga et al. [2021] review algorithm unrolling beyond amortized optimization, while Hospedales et al. [2020] connect meta-representations to amortized optimization problems.Source: §5.6.1 · pp. 75–76

5.6.2 Amortized optimization over discrete domains

  • Discrete and mixed discrete-continuous optimization is outside the tutorial’s scope because many methods differ substantially and derivative information often becomes unavailable.The passage points readers to work on combinatorial spaces and discrete control, including learned latent action spaces and mixed-integer control policies.Source: §5.6.2 · pp. 76–77

5.6.3 Learning-augmented and amortized algorithms beyond optimization

  • Learning can also improve algorithms unrelated to optimization, while continuation and homotopy methods share amortization’s motivation of solving related problems across contexts.Continuation and homotopy typically solve problems along X = [0, 1] without learning; the tutorial suggests those problem classes could also benefit from shared amortization models.Source: §5.6.3 · p. 77, §5.6.4 · p. 77

Related papers

Loading 2202.00665v5…