Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Tutorial on amortized optimization

Brandon Amos

TL;DR

This tutorial surveys amortized optimization, a paradigm that combines learning with optimization to predict solutions to repeated, structured problems. It distinguishes fully- and semi-amortized models, discusses how to train them via regression- or objective-based losses, and explains how to differentiate through inner optimization loops, including practical tricks like TBPTT and implicit differentiation. The coverage spans foundational theory, extensions to stochastic, constrained, and non-Euclidean domains, and a broad set of applications from variational inference and sparse coding to meta-learning, fixed-point acceleration, optimal transport, and control. It also provides implementation guidance, software resources, and a candid discussion of successes, limitations, and open problems in the field.

Abstract

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.

Tutorial on amortized optimization

TL;DR

This tutorial surveys amortized optimization, a paradigm that combines learning with optimization to predict solutions to repeated, structured problems. It distinguishes fully- and semi-amortized models, discusses how to train them via regression- or objective-based losses, and explains how to differentiate through inner optimization loops, including practical tricks like TBPTT and implicit differentiation. The coverage spans foundational theory, extensions to stochastic, constrained, and non-Euclidean domains, and a broad set of applications from variational inference and sparse coding to meta-learning, fixed-point acceleration, optimal transport, and control. It also provides implementation guidance, software resources, and a candid discussion of successes, limitations, and open problems in the field.

Abstract

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.
Paper Structure (88 sections, 3 theorems, 89 equations, 21 figures, 2 tables)

This paper contains 88 sections, 3 theorems, 89 equations, 21 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Amortized optimization foundations
  3. Defining the model $\hat{y}_\theta(x)$
  4. Fully-amortized models
  5. Semi-amortized models
  6. Semi-amortized models over the domain ${\mathcal{Y}}$
  7. Semi-amortized models over a latent space ${\mathcal{Z}}$
  8. Comparing semi-amortized models with warm-starting
  9. On second-order derivatives of the objective
  10. Models based on differentiable optimization
  11. Practically choosing a model
  12. Learning the model's parameters $\theta$
  13. Choosing the objective for learning
  14. Regression-based learning
  15. Objective-based learning
  16. ...and 73 more sections

Key Result

Theorem 1

Let the roots of $g(y; \varphi)$ define an implicit mapping $Y^\star(\varphi)$ given by $Y^\star(\varphi)\vcentcolon=\{y \mid g(y;\varphi)=0\}$, where $\varphi\in\mathbb{R}^m$, $y\in\mathbb{R}^n$, and $g: \mathbb{R}^n\times\mathbb{R}^m\rightarrow\mathbb{R}^n$. Let $g$ be continuously differentiable

Figures (21)

  • Figure 1: Illustration of the parametric optimization problem in \ref{['eq:opt']}. Each context $x$ parameterizes an optimization problem that the objective $f(y; x)$ depends on. The contours show the values of the objectives where darker colors indicate higher values. The objective is then minimized over $y$ and the resulting solution $y^\star(x)$ is shown in red. In other words, each vertical slice is an optimization problem and this visualization shows a continuum of optimization problems.
  • Figure 2: An amortized optimization method learns a model $\hat{y}_\theta$ to predict the minimum of an objective$f(y;x)$ to a parameterized optimization problem, as in \ref{['eq:opt']}, which depends on a context$x$. For example, in control, the context space ${\mathcal{X}}$ is the state space of the system, e.g. angular positions and velocities describing the configuration of the system, the domain ${\mathcal{Y}}\vcentcolon={\mathcal{U}}$ is the control space, e.g. torques to apply to each actuated joint, the cost (or negated value) of a state-action pair is $f(u; x)\vcentcolon= -Q(x,u)$, and the state distribution is $p(x)$. For an encountered state $x$, many reinforcement learning policies $\pi_\theta(x)\vcentcolon=\hat{y}_\theta(x)$ amortize the solution to the underlying control problem with true solution $y^\star(x)$. This humanoid policy was obtained with the model-based stochastic value gradient in amos2021model.
  • Figure 3: Overview of amortized optimization modeling and loss choices.
  • Figure 4: Overview of key losses for optimizing the parameters $\theta$ of the amortization model $\hat{y}_\theta$. Regression-based losses optimize a distance between the model's prediction $\hat{y}_\theta(x)$ and the ground-truth $y^\star(x)$. Objective-based methods update $\hat{y}_\theta$ using local information of the objective $f$ and without access to the ground-truth solutions $y^\star$.
  • Figure 5: Contours of the regression-based amortization loss ${\mathcal{L}}_{\rm reg}$ (in black) alongside the contours of the objective (in purple where darker colors indicate higher values). This shows the inaccuracies of the regression-based loss, e.g. along a level set, may impact the overall objective differently.
  • ...and 16 more figures

Theorems & Definitions (15)

  • Definition 1
  • Definition 2
  • Definition 3
  • Remark 1
  • Remark 2
  • Remark 3
  • Remark 4
  • Definition 4
  • Theorem 1: Dini's implicit function theorem
  • Definition 5
  • ...and 5 more