Appendix A Optimization Methods

In this section

Steepest Descent
Vanilla Gradient Descent for Smooth Problems Preconditioned Gradient Descent for Badly-Conditioned Problems Proximal Gradient Descent for Non-Smooth Problems Stochastic Gradient Descent for Large-Scale Problems Putting Everything Together: Adam
Computing Gradients via Automatic Differentiation
Differentials Automatic Differentiation Back Propagation
Game Theory and Minimax Optimization
Learning Stackelberg Equilibria Practical Considerations when Learning Stackelberg Equilibria
Exercises

“Since the building of all the universe is perfect and is created by the wisdom creator, nothing arises in the universe in which one cannot see the sense of some maximum or minimum.”

– L. Euler

In this chapter, we give a brief introduction to some of the most basic but important optimization algorithms used in this book. The purpose is only to help the reader apply these algorithms to problems studied in this book, not to gain a deep understanding about these algorithms. Hence, we will not provide a thorough justification for the algorithms introduced, in terms of performance guarantees.

A.1 Steepest Descent

Optimization is concerned with the question of how to find where a function, say $L(\theta)$ , reaches its minimum value. Mathematically, this is stated as a problem:

\operatorname*{arg\ min}_{\theta\in\Theta}\mathcal{L}(\theta),

(A.1.1)

where $\Theta$ represents a domain to which the argument $\bm{x}$ is confined. Often (and unleess otherwise mentioned, in this chapter) $\Theta$ is simply $\mathbb{R}^{n}$ . Without loss of generality, we assume that here the function $\mathcal{L}(\theta)$ is smooth¹¹1In case the function $\mathcal{L}$ is not smooth, we replace its gradient with a so-called subgradient..

The efficiency of finding the (global) minima depends on what information we have about the function $\mathcal{L}$ . For most optimization problems considered in this book, the dimension of $\theta$ , say $n$ , is very large. That makes computing or accessing local information about $\ell$ expensive. In particular, since the gradient $\nabla\mathcal{L}$ has $n$ entries, it is often reasonable to compute; however, the Hessian $\nabla^{2}\mathcal{L}$ has $n^{2}$ entries which is often wildly impractical to compute (and the same goes for higher-order derivatives). Hence, it is typical to assume that we have the zeroth-order information, i.e., we are able to evaluate $\mathcal{L}(\theta)$ , and the first-order information, i.e., we are able to evaluate $\nabla\mathcal{L}(\theta)$ . Optimization theorists may rephrase this as saying we have a “first-order oracle.” All optimization algorithms that we introduce in this section only use a first-order oracle.²²2We refer the readers to the book by [WM22] for a more systematic introduction to optimization algorithms in a high-dimensional space, including algorithms assuming higher-order oracles.

A.1.1 Vanilla Gradient Descent for Smooth Problems

The simplest and most widely used method for optimization is gradient descent (GD). It was first introduced by Cauchy in 1847. The idea is very simple: starting from an initial state, we iteratively take small steps such that each step reduces the value of the function $\mathcal{L}(\theta)$ .

Suppose that the current state is $\theta$ . We want to take a small step, say of distance $h$ , in a direction, indicated by a vector $\bm{v}$ , to reach a new state $\theta+h\bm{v}$ such that the value of the function decreases:

\mathcal{L}(\theta+h\bm{v})\leq\mathcal{L}(\theta).

(A.1.2)

To find such a direction $\bm{v}$ , we can approximate $\mathcal{L}(\theta+h\bm{v})$ through a Taylor expansion around $h=0$ :

\mathcal{L}(\theta+h\bm{v})=\mathcal{L}(\theta)+h\langle\nabla\mathcal{L}(\theta),\bm{v}\rangle+o(h),

(A.1.3)

where the inner product here (and in this chapter) will be the $\ell^{2}$ inner product, i.e., $\langle\bm{x},\bm{y}\rangle=\bm{x}^{\top}\bm{y}$ . To find the direction of steepest descent, we attempt to minimize this Taylor expansion among unit vectors $\bm{v}$ . If $\nabla\mathcal{L}(\theta)=\bm{0}$ , then the second term above is $0$ regardless of the value of $\bm{v}$ , so we cannot attempt to make progress, i.e., the algorithm has converged. On the other hand, if $\nabla\mathcal{L}(\theta)\neq\bm{0}$ then it holds

\operatorname*{arg\ min}_{\begin{subarray}{c}\bm{v}\in\mathbb{R}^{d}\\ \|\bm{v}\|_{2}=1\end{subarray}}[\mathcal{L}(\theta)+h\langle\nabla\mathcal{L}(\theta),\bm{v}\rangle]=\operatorname*{arg\ min}_{\begin{subarray}{c}\bm{v}\in\mathbb{R}^{d}\\ \|\bm{v}\|_{2}=1\end{subarray}}\langle\nabla\mathcal{L}(\theta),\bm{v}\rangle=-\frac{\nabla\mathcal{L}(\theta)}{\|\nabla\mathcal{L}(\theta)\|_{2}},

(A.1.4)

In words, this means that the value of $\mathcal{L}(\theta+h\bm{v})$ decreases the fastest along the direction $\bm{v}=-\nabla\mathcal{L}(\theta)/\|\nabla\mathcal{L}(\theta)\|_{2}$ , for small enough $h$ . This leads to the gradient descent method: From the current state $\theta_{k}$ ( $k=0,1,\ldots$ ), we take a step of size $h$ in the direction of the negative gradient to reach the next iterate,

\theta_{k+1}=\theta_{k}-h\nabla\mathcal{L}(\theta_{k}).

(A.1.5)

The step size $h$ is also called the learning rate in machine learning contexts.

Step-Size Selection

The remaining question is what the step size $h$ should be? If we choose $h$ to be too small, the value of the function may decrease very slowly, as shown by the plot in the middle in Figure A.1. If $h$ is too large, the value might not even decrease at all, as shown by the plot on the right in Figure A.1.

Figure A.1 : The effect of step size h h italic_h on the convergence of the gradient descent method. — Figure A.1: The effect of step size $h$ on the convergence of the gradient descent method.

So the step size $h$ should be chosen based on the landscape of the function $\mathcal{L}(\theta_{k})$ . Ideally, to choose the best step size $h$ , we can solve the following optimization problem over a single variable $h$ :

h=\operatorname*{arg\ min}_{h\geq 0}\mathcal{L}(\theta_{k}-h\nabla\mathcal{L}(\theta_{k})).

(A.1.6)

This method of choosing the step size is called line search. However, when the function $L(\theta_{k})$ is complicated, which is usually the case for training a deep neural network, this one-dimensional optimization is very difficult to solve at each iteration of gradient descent.

Then how should we choose a proper step size $h$ ? One common and classical approach is to try to obtain a good approximation of the local landscape around the current state $\theta$ based on some knowledge about the overall landscape of the function $\mathcal{L}(\theta)$ .

Figure A.2: Majorization-minimization scheme and intuition. A function

\mathcal{L}\colon\Theta\to\mathbb{R}

has a global upper bound

u\colon\Theta\to\mathbb{R}

which meets

L

at at least one point

\theta_{0}

. Then, finding the

\theta_{1}

which minimizes

u

will improve the value of

\mathcal{L}

from

\mathcal{L}(\theta_{0})

. Note that similar results can be shown about local upper bounds.

Common conditions for the landscape of $\mathcal{L}(\theta)$ include:

•

$\alpha$ -Strong Convexity. Recall that $\mathcal{L}$ is $\alpha$ -strongly convex if its graph lies above a global quadratic lower bound of slope $\alpha$ , i.e.,

\mathcal{L}(\theta)\geq l_{\theta_{0},\alpha}(\theta)\doteq\mathcal{L}(\theta_{0})+\langle\nabla\mathcal{L}(\theta_{0}),\theta-\theta_{0}\rangle+\frac{\alpha}{2}\|\theta-\theta_{0}\|_{2}^{2}

(A.1.7)

for any “base point” $\theta_{0}$ . We say that $\mathcal{L}$ is convex if it is $0$ -strongly convex, i.e., its graph lies above its tangents. It is easy to show (proof as exercise) that strongly convex functions have unique global minima. Another important fact (proof as exercise) is that $\alpha$ -strongly convex twice-differentiable functions $\mathcal{L}$ have (symmetric) Hessians $\nabla^{2}\mathcal{L}$ whose minimum eigenvalue is $\geq\alpha$ . For $\alpha>0$ this implies the Hessian is symmetric positive definite, and for $\alpha=0$ (i.e., $\mathcal{L}$ is convex) this implies that the Hessian is symmetric positive semidefinite.

•

$\beta$ -Lipschitz Gradient (also called $\beta$ -Smoothness). Recall that $\mathcal{L}$ has $\beta$ -Lipschitz gradient if $\nabla\mathcal{L}$ exists and is $\beta$ -Lipschitz, i.e.,

\|\nabla\mathcal{L}(\theta)-\nabla\mathcal{L}(\theta_{0})\|_{2}\leq\beta\|\theta-\theta_{0}\|_{2}.

(A.1.8)

for any “base point” $\theta_{0}$ . It is easy to show (proof as exercise) that this is equivalent to $\mathcal{L}$ having a global quadratic upper bound of slope $\beta$ , i.e.,

\mathcal{L}(\theta)\leq u_{\theta_{0},\beta}(\theta)\doteq\mathcal{L}(\theta_{0})+\langle\nabla\mathcal{L}(\theta_{0}),\theta-\theta_{0}\rangle+\frac{\beta}{2}\|\theta-\theta_{0}\|_{2}^{2}.

(A.1.9)

for any “base point” $\theta_{0}$ . Another important fact (proof as exercise) is that convex $\beta$ -Lipschitz gradient twice-differentiable functions have (symmetric) Hessians $\nabla^{2}\mathcal{L}$ whose largest eigenvalue is $\leq\beta$ .

First, let us suppose that $\mathcal{L}$ has $\beta$ -Lipschitz gradient (but is not necessarily even convex). We will use this occasion to introduce a common optimization theme: to minimize $\mathcal{L}$ , we can minimize an upper bound on $L$ , which is justified by the following lemma visualized in Figure A.2.

Lemma A.1 (Majorization-Minimization).

Suppose that $u\colon\Theta\to\mathbb{R}$ is a global upper bound on $\mathcal{L}$ , namely $\mathcal{L}(\theta)\leq u(\theta)$ for all $\theta\in\Theta$ . Suppose that they meet with equality at $\theta_{0}$ , i.e., $\mathcal{L}(\theta_{0})=u(\theta_{0})$ . Then

\theta_{1}\in\operatorname*{arg\ min}_{\theta\in\Theta}u(\theta)\implies\mathcal{L}(\theta_{1})\leq u(\theta_{1})\leq u(\theta_{0})=\mathcal{L}(\theta_{0}).

(A.1.10)

We will use this lemma to show that we can use the Lipschitz gradient property to ensure that each gradient step cannot worsen the value of $\mathcal{L}$ . Indeed, at every base point $\theta_{0}$ , we have that $u_{\theta_{0},\beta}$ is a global upper bound on $\mathcal{L}$ , and $u_{\theta_{0},\beta}(\theta_{0})=\mathcal{L}(\theta_{0})$ . Hence by Lemma A.1

\text{if $\theta$ minimizes $u_{\theta_{0},\beta}$ then}\quad\mathcal{L}(\theta)\leq u_{\theta_{0},\beta}(\theta)\leq u_{\theta_{0},\beta}(\theta_{0})=\mathcal{L}(\theta_{0}).

(A.1.11)

This motivates us, when finding an update to obtain $\theta_{k+1}$ from $\theta_{k}$ , we can instead minimize the upper bound $u_{\theta_{k},\beta}$ over $\theta$ and set that to be $\theta_{k+1}$ . By minimizing $u_{\theta_{k},\beta}$ (proof as exercise) we get

\theta_{k+1}=\theta_{k}-\frac{1}{\beta}\nabla\mathcal{L}(\theta_{k})\implies\mathcal{L}(\theta_{k+1})\leq\mathcal{L}(\theta_{k}).

(A.1.12)

This implies that a step size $h=1/\beta$ is a usable learning rate, but it does not provide a convergence rate or certify that $L(\theta_{k})$ actually converges to $\min_{\theta}\mathcal{L}(\theta)$ . This requires a little more rigor, which we now pursue.

Now, let us suppose that $\mathcal{L}$ is $\alpha$ -strongly convex, has $\beta$ -Lipschitz gradient, and has global optimum $\theta^{\star}$ . We will show that $\theta_{k}$ will converge directly to the unique global optimum $\theta^{\star}$ , which is a very strong form of convergence. In particular, we will bound $\|\theta^{\star}-\theta_{k}\|_{2}$ using both strong convexity and Lipschitzness of the gradient of $\mathcal{L}$ , i.e., taking a look at the neighborhood around $\theta_{k}$ :³³3In this proof the $\beta$ -Lipschitz Gradient invocation step is a little non-trivial. We also leave this step as an exercise, with the hint to plug in $\theta=\theta_{0}-h\nabla\mathcal{L}(\theta_{0})$ into the Lipschitz gradient identity.

$\displaystyle\\|\theta^{\star}-\theta_{k+1}\\|_{2}^{2}$	$\displaystyle\leq\\|\theta^{\star}-\theta_{k}+h\nabla\mathcal{L}(\theta_{k})\\|_{2}^{2}$	(A.1.13)
	$\displaystyle=\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}+2h\langle\nabla\mathcal{L}(\theta_{k}),\theta^{\star}-\theta_{k}\rangle+h^{2}\\|\nabla\mathcal{L}(\theta_{k})\\|_{2}^{2}$	(A.1.14)
	$\displaystyle\leq\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}+2h\left(\mathcal{L}(\theta^{\star})-\mathcal{L}(\theta_{k})-\frac{\alpha}{2}\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}\right)+h^{2}\\|\nabla\mathcal{L}(\theta_{k})\\|_{2}^{2}\quad\text{($\alpha$-SC)}$	(A.1.15)
	$\displaystyle=\left(1-\alpha h\right)\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}+2h(\mathcal{L}(\theta^{\star})-L(\theta_{k}))+h^{2}\\|\nabla\mathcal{L}(\theta_{k})\\|_{2}^{2}$	(A.1.16)
	$\displaystyle\leq\left(1-\alpha h\right)\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}+2h(\mathcal{L}(\theta^{k})-\mathcal{L}(\theta^{\star}))+2h^{2}\beta(\mathcal{L}(\theta_{k})-\mathcal{L}(\theta^{\star}))\quad\text{($\beta$-LG)}$	(A.1.17)
	$\displaystyle=\left(1-\alpha h\right)\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}-2h(1-\beta h)(\mathcal{L}(\theta_{k})-\mathcal{L}(\theta^{\star})).$	(A.1.18)

In order to ensure that the gradient descent iteration makes progress we must pick the step size so that $1-\beta h\geq 0$ , i.e., $h\leq 1/\beta$ . If such a setting occurs, then

	$\displaystyle\\|\theta^{\star}-\theta_{k+1}\\|_{2}^{2}$	$\displaystyle\leq(1-\alpha h)\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}\leq(1-\alpha h)^{2}\\|\theta^{\star}-\theta_{k-1}\\|_{2}^{2}\leq\cdots$		(A.1.19)
		$\displaystyle\leq(1-\alpha h)^{k+1}\\|\theta^{\star}-\theta_{0}\\|_{2}^{2}.$		(A.1.20)

In order to minimize the right-hand side, we can set $h=1/\beta$ , which obtains

\|\theta^{\star}-\theta_{k+1}\|_{2}^{2}\leq(1-\alpha/\beta)^{k+1}\|\theta^{\star}-\theta_{0}\|_{2}^{2},

(A.1.21)

showing convergence to global optimum with exponentially decaying error. Notice that here we used a convergence rate to obtain a favorable step size of $h=1/\beta$ . This motif will re-occur in this section.

We end this section with a caveat: learning a global optimum is (usually) impractically hard. Under certain conditions, we can ensure that the gradient descent iterates converge to a local optimum. Also, under more relaxed conditions, we can ensure local convergence, i.e., that the iterates converge to a (global or local) optimum if the sequence is initialized close enough to the optimum.

A.1.2 Preconditioned Gradient Descent for Badly-Conditioned Problems

Figure A.3 : The negative gradient − ∇ ℒ λ -\nabla\mathcal{L}_{\lambda} - ∇ caligraphic_L start_POSTSUBSCRIPT italic_λ end_POSTSUBSCRIPT and pre-conditioned (Newton’s method step) vector field − [ ∇ 2 ℒ λ ] − 1 [ ∇ ℒ λ ] -[\nabla^{2}\mathcal{L}_{\lambda}]^{-1}[\nabla\mathcal{L}_{\lambda}] - [ ∇ start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT caligraphic_L start_POSTSUBSCRIPT italic_λ end_POSTSUBSCRIPT ] start_POSTSUPERSCRIPT - 1 end_POSTSUPERSCRIPT [ ∇ caligraphic_L start_POSTSUBSCRIPT italic_λ end_POSTSUBSCRIPT ] where λ = 19 \lambda=19 italic_λ = 19 . There is a section of the space where following the negative gradient vector field makes very little progress towards finding the minimum, but in all cases following the Newton’s method vector field achieves equal speed of progress towards the optimum since the gradient is whitened. Since the Hessian here is diagonal, adaptive learning rate algorithms (e.g. Adam, as will be discussed later in the section) can make similar progress as Newton’s method, but a non-axis-aligned Hessian may even prevent Adam from succeeding quickly. — Figure A.3: The negative gradient $-\nabla\mathcal{L}_{\lambda}$ and pre-conditioned (Newton’s method step) vector field $-[\nabla^{2}\mathcal{L}_{\lambda}]^{-1}[\nabla\mathcal{L}_{\lambda}]$ where $\lambda=19$ . There is a section of the space where following the negative gradient vector field makes very little progress towards finding the minimum, but in all cases following the Newton’s method vector field achieves equal speed of progress towards the optimum since the gradient is whitened. Since the Hessian here is diagonal, adaptive learning rate algorithms (e.g. Adam, as will be discussed later in the section) can make similar progress as Newton’s method, but a non-axis-aligned Hessian may even prevent Adam from succeeding quickly.

Newton’s Method

There are some smooth problems and strongly convex problems on which gradient descent nonetheless does quite poorly. For example, let $\lambda\geq 0$ and let $\mathcal{L}_{\lambda}\colon\mathbb{R}^{2}\to\mathbb{R}$ of the form

\mathcal{L}_{\lambda}(\theta)=\mathcal{L}_{\lambda}\left(\begin{bmatrix}\theta_{1}\\ \theta_{2}\end{bmatrix}\right)\doteq\frac{1}{2}\left\{(1+\lambda)\theta_{1}^{2}+\theta_{2}^{2}\right\}=\frac{1}{2}\theta^{\top}\begin{bmatrix}1+\lambda&0\\ 0&1\end{bmatrix}\theta.

(A.1.22)

This problem is $1$ -strongly convex and has $(1+\lambda)$ -Lipschitz gradient. The convergence rate is then geometric with rate $1-1/(1+\lambda)$ . For large $\lambda$ , this is still not very fast. In this section, we will introduce a class of optimization problems which can successfully optimize such badly-conditioned functions.

The key lies in the objective’s curvature, which is given by the Hessian. Suppose that (as a counterfactual) we had a second-order oracle which would allow us to compute $\mathcal{L}(\theta)$ , $\nabla\mathcal{L}(\theta)$ , and $\nabla^{2}\mathcal{L}(\theta)$ . Then, instead of picking a descent direction $\bm{v}$ to optimize the first-order Taylor expansion around $\theta$ , we could optimize the second-order Taylor expansion instead. Intuitively this would allow us to incorporate curvature information into the update.

Let us carry out this computation. The second-order Taylor expansion of $\mathcal{L}(\theta+h\bm{v})$ around $h=0$ is

\mathcal{L}(\theta+h\bm{v})=\mathcal{L}(\theta)+h\langle\nabla\mathcal{L}(\theta),\bm{v}\rangle+\frac{1}{2}h^{2}\langle[\nabla^{2}\mathcal{L}(\theta)]\bm{v},\bm{v}\rangle+o(h^{2}).

(A.1.23)

Then we can compute the descent direction:

\displaystyle\operatorname*{arg\ min}_{\begin{subarray}{c}\bm{v}\in\mathbb{R}^{n}\\ \|\bm{v}\|_{2}=1\end{subarray}}\left[\mathcal{L}(\theta)+h\langle\nabla\mathcal{L}(\theta),\bm{v}\rangle+\frac{1}{2}h^{2}\langle[\nabla^{2}\mathcal{L}(\theta)]\bm{v},\bm{v}\rangle\right]

\displaystyle=\operatorname*{arg\ min}_{\begin{subarray}{c}\bm{v}\in\mathbb{R}^{n}\\ \|\bm{v}\|_{2}=1\end{subarray}}\left[\langle\nabla\mathcal{L}(\theta),\bm{v}\rangle+\frac{1}{2}h\langle[\nabla^{2}\mathcal{L}(\theta)]\bm{v},\bm{v}\rangle\right].

(A.1.24)

This optimization problem is a little difficult to solve because of the constraint $\|\bm{v}\|_{2}=1$ . But in practice we never normalize the descent direction $\bm{v}$ and use the step size $h$ to control the size of the update. So let us just solve the above problem over all vectors $\bm{v}\in\mathbb{R}^{n}$ :⁴⁴4If $\nabla^{2}\mathcal{L}(\theta)$ is not invertible, then we can replace $[\nabla^{2}\mathcal{L}(\theta)]^{-1}$ with the Moore-Penrose pseudoinverse of $\nabla^{2}\mathcal{L}(\theta)$ .

\operatorname*{arg\ min}_{\bm{v}\in\mathbb{R}^{n}}\left[\langle\nabla\mathcal{L}(\theta),\bm{v}\rangle+\frac{1}{2}h\langle[\nabla^{2}\mathcal{L}(\theta)]\bm{v},\bm{v}\rangle\right]=-\frac{1}{h}[\nabla^{2}\mathcal{L}(\theta)]^{-1}[\nabla\mathcal{L}(\theta)].

(A.1.25)

We can thus use the steepest descent iteration

\theta_{k+1}=\theta_{k}-[\nabla^{2}\mathcal{L}(\theta_{k})]^{-1}[\nabla\mathcal{L}(\theta_{k})],

(A.1.26)

(this is the celebrated Newton’s method), or

\theta_{k+1}=\theta_{k}-h[\nabla^{2}\mathcal{L}(\theta_{k})]^{-1}[\nabla\mathcal{L}(\theta_{k})],

(A.1.27)

(which is called underdamped Newton’s method). Since the second-order quadratic $\mathcal{L}_{\lambda}$ is equal to its second-order Taylor expansion, if we run Newton’s method for one step, we will achieve the global minimum in one step no matter how large $\lambda$ is. Figure A.3 gives some intuition about poorly conditioned functions and the gradient steps versus Newton’s steps.

PGD

In practice, we do not have a second-order oracle which allows us to compute $\nabla^{2}\mathcal{L}(\theta)$ . Instead, we can attempt to learn an approximation to it alongside the parameter update $\theta_{k+1}$ from $\theta_{k}$ .

How do we learn an approximation to it? We shall find some equations which the Hessian’s inverse satisfies and then try to update our approximation so that it satisfies the equations. Namely, taking the Taylor series of $\nabla\mathcal{L}(\theta+\delta_{\theta})$ around point $\theta$ , we obtain

\underbrace{\nabla L(\theta+\delta_{\theta})-\nabla\mathcal{L}(\theta)}_{\doteq\delta_{\bm{g}}}=[\nabla^{2}\mathcal{L}(\theta)]\delta_{\theta}+o(\|\delta_{\theta}\|_{2}).

(A.1.28)

In this case we have

\delta_{\bm{g}}\approx[\nabla^{2}\mathcal{L}(\theta)]\delta_{\theta}\implies\delta_{\theta}\approx[\nabla^{2}\mathcal{L}(\theta)]^{-1}\delta_{\bm{g}}

(A.1.29)

We can now try to learn a symmetric positive semidefinite pre-conditioner $P\in\mathbb{R}^{n\times n}$ such that

\delta_{\theta}\approx P\delta_{\bm{g}},

(A.1.30)

updating it at each iteration along with $\theta_{k}$ . Namely, we have the PSGD iteration

	$\displaystyle P_{k}$	$\displaystyle=\mathrm{PreconditionerUpdate}(P_{k-1};\theta_{k},\nabla\mathcal{L}(\theta_{k}))$		(A.1.31)
	$\displaystyle\theta_{k+1}$	$\displaystyle=\theta_{k}-hP_{k}\nabla\mathcal{L}(\theta_{k}).$		(A.1.32)

This update has two problems: how can we even use $P$ (since we already said we cannot store an $n\times n$ matrix) and how can we update $P$ at each iteration? The answers are very related; we can never materialize $P$ in computer memory, but we can represent it using a low-rank factorization (or comparable methods such as Kronecker factorization which is particularly suited to the form of deep neural networks). Then the preconditioner update step is designed to exploit the structure of the preconditioner representation.

We end this subsection with a caveat: in deep learning, for example, $\mathcal{L}$ is not a convex function and so Newton’s method (and approximations to it) do not make sense. In this case we look at the geometric intuition of Newton’s method on convex functions, say from Figure A.3: the inverse Hessian whitens the gradients. Thus instead of a Hessian-approximating preconditioner, we can adjust the above procedures to learn a more general whitening transformation for the gradient. This is the idea behind the original proposal of PSGD [Li17], which contains more information about how to store and update the preconditioner, and more modern optimizers like Muon [LSY+25].

A.1.3 Proximal Gradient Descent for Non-Smooth Problems

Even in very toy problems, however, such as LASSO or dictionary learning, the problem is not strongly convex but rather just convex, and the objective is no longer just smooth but rather the sum of a smooth function and a non-smooth regularizer (such as the $\ell^{1}$ norm). Such problems are solved by proximal optimization algorithms, which generalize steepest descent to non-smooth objectives.

Formally, let us say that

\mathcal{L}(\theta)\doteq\mathcal{S}(\theta)+\mathcal{R}(\theta)

(A.1.33)

where $\mathcal{S}$ is smooth, say with $\beta$ -Lipschitz gradient, and $\mathcal{R}$ is non-smooth (i.e., rough). The proximal gradient algorithm generalizes the steepest descent algorithm, by using the majorization-minimization framework (i.e., Lemma A.1) with a different global upper bound. Namely, we construct such an upper bound by asking: what if we take the Lipschitz gradient upper bound of $S$ but leave $R$ alone? Namely, we have

\mathcal{L}(\theta_{1})=\mathcal{S}(\theta_{1})+\mathcal{R}(\theta_{1})\leq u_{\theta_{0},\beta}(\theta_{1})\doteq\mathcal{S}(\theta_{0})+\langle\nabla\mathcal{S}(\theta_{0}),\theta_{1}-\theta_{0}\rangle+\frac{\beta}{2}\|\theta_{1}-\theta_{0}\|_{2}^{2}+\mathcal{R}(\theta_{1}).

(A.1.34)

Note that (proof as exercise)

\operatorname*{arg\ min}_{\theta_{1}}u_{\theta_{0},\beta}(\theta_{1})=\operatorname*{arg\ min}_{\theta_{1}}\left[\frac{\beta}{2}\left\|\theta_{1}-\left(\theta_{0}-\frac{1}{\beta}\nabla\mathcal{S}(\theta_{0})\right)\right\|_{2}^{2}+\mathcal{R}(\theta_{1})\right].

(A.1.35)

Now if we try to minimize the upper bound $u_{\theta_{0},\beta}$ , we are picking a $\theta_{1}$ that:

•

is close to the gradient update $\theta_{0}-\frac{1}{\beta}\nabla\mathcal{S}(\theta_{0})$ ;
•

has a small value of the regularizer $\mathcal{R}(\theta_{1})$

and trades off these properties according to the smoothness parameter $\beta$ . Accordingly, let us define the proximal operator

\operatorname{prox}_{h,\mathcal{R}}(\theta)\doteq\operatorname*{arg\ min}_{\theta_{1}}\left[\frac{1}{2h}\|\theta_{1}-\theta\|_{2}^{2}+\mathcal{R}(\theta)\right].

(A.1.36)

Then, we can define the proximal gradient descent iteration which, at each iteration, minimizes the upper bound $u_{\theta_{k},h^{-1}}$ , i.e.,

\theta_{k+1}=\operatorname{prox}_{h,\mathcal{R}}(\theta_{k}-h\nabla\mathcal{S}(\theta_{k})).

(A.1.37)

Convergence proofs are possible when $h\leq 1/\beta$ , but we do not give any in this section.

One remaining question is: how can we compute the proximal operator? At first glance, it seems like we have traded one intractable minimization problem for another. Since we have not made any assumption on $\mathcal{R}$ so far, the framework works even when $\mathcal{R}$ is a very complex function (such as a neural network loss), which would require us to solve a neural network training problem in order to compute a single proximal operator. However, in practice, for simple regularizers $\mathcal{R}$ such as those we use in this manuscript, there exist proximal operators which are easy to compute or even in closed-form. We give a few below (the proofs are an exercise).

Example A.1.

Let $\Gamma\subseteq\Theta$ be a set, and let $\chi_{\Gamma}$ be the characteristic function on $\Gamma$ , i.e.,

\chi_{\Gamma}(\theta)\doteq\begin{cases}0,&\text{if}\ \theta\in\Gamma\\ +\infty,&\text{if}\ \theta\notin\Gamma.\end{cases}

(A.1.38)

Then the proximal operator of $\chi_{\Gamma}$ is a projection, i.e.,

\operatorname{prox}_{h,\chi_{\Gamma}}(\theta)=\operatorname*{arg\ min}_{\theta_{1}\in\Gamma}\frac{1}{2}\|\theta_{1}-\theta\|_{2}^{2}=\operatorname*{arg\ min}_{\theta_{1}\in\Gamma}\|\theta_{1}-\theta\|_{2}.

(A.1.39)

$\blacksquare$

Example A.2.

The $\ell^{1}$ norm has a proximal operator which performs soft thresholding:

S_{h}(\theta)\doteq\operatorname{prox}_{h,\lambda\|\cdot\|_{1}}(\theta)=\operatorname*{arg\ min}_{\theta_{1}}\left[\frac{1}{2h}\|\theta_{1}-\theta\|_{2}^{2}+\lambda\|\theta\|_{1}\right]

(A.1.40)

then $S_{h}(\theta)$ is defined by

S_{h}(\theta)_{i}=\begin{cases}\theta_{i}-h\lambda,&\text{if}\ \theta_{i}\geq h\lambda\\ 0,&\text{if}\ \theta_{i}\in[-h\lambda,h\lambda]\\ \theta_{i}+h\lambda,&\text{if}\ \theta_{i}\leq-h\lambda\end{cases}=\begin{cases}\max\{\lvert\theta_{i}\rvert-h\lambda,0\}\operatorname{sign}(\theta_{i}),&\text{if}\ \lvert\theta_{i}\rvert\geq h\lambda\\ 0,&\text{if}\ \lvert\theta_{i}\rvert<h\lambda.\end{cases}

(A.1.41)

The proximal gradient operation with the smooth part of the objective being least-squares and the non-smooth part being the $\ell^{1}$ norm (hence using this soft thresholding proximal operator) is called the Iterative Shrinkage-Thresholding Algorithm (ISTA). $\blacksquare$

Example A.3.

In Chapter 4 we use a proximal operator corresponding to the $\ell^{1}$ norm plus the characteristic function for the positive orthant $\mathbb{R}_{+}^{n}\doteq\{\bm{x}\in\mathbb{R}^{n}\colon x_{i}\geq 0\ \forall i\}$ , namely

T_{h}(\theta)\doteq\operatorname{prox}_{h,\lambda\|\cdot\|_{1}+\chi_{\mathbb{R}_{+}^{n}}}(\theta)=\operatorname*{arg\ min}_{\theta_{1}\in\mathbb{R}_{+}^{n}}\left[\frac{1}{2h}\|\theta_{1}-\theta\|_{2}^{2}+\lambda\|\theta\|_{1}\right],

(A.1.42)

then $T_{h}$ is defined as

T_{h}(\theta)_{i}\doteq\max\{\theta_{i}-h\lambda,0\}.

(A.1.43)

This proximal operator yields the non-negative ISTA that is invoked in Chapter 4 and beyond. $\blacksquare$

A.1.4 Stochastic Gradient Descent for Large-Scale Problems

In deep learning, the objective function $\mathcal{L}$ usually cannot be computed exactly, and instead at each optimization step it is estimated using finite samples (say, using a mini-batch). A common way to model this situation is to define a stochastic loss function $\mathcal{L}_{\omega}(\theta)$ where $\omega$ is some “source of randomness”. For example, $\omega$ could contain the indices of the samples in a batch over which to compute the loss. Then, we would like to minimize $\mathcal{L}(\theta)\doteq\operatorname{\mathbb{E}}_{\omega}[\mathcal{L}_{\omega}(\theta)]$ over $\theta$ , given access to a stochastic first-order oracle: given $\theta$ , we can sample $\omega$ and compute $\mathcal{L}_{\omega}(\theta)$ and $\nabla_{\theta}\mathcal{L}_{\omega}(\theta)$ . This minimization problem is called a stochastic optimization problem.

The basic first-order stochastic algorithm is stochastic gradient descent: at each iteration $k$ we sample $\omega_{k}$ , define $\mathcal{L}_{k}\doteq\mathcal{L}_{\omega_{k}}$ , and perform a gradient step on $\mathcal{L}_{k}$ , i.e.,

\theta_{k+1}=\theta_{k}-h\nabla\mathcal{L}_{k}(\theta_{k}).

(A.1.44)

However, even for very simple problems we cannot expect the same type of convergence as we obtained in gradient descent. For example, suppose that there are $m$ possible values for $\omega\in\{1,\dots,m\}$ which it takes with equal probability, and there are $m$ possible targets $\xi_{1},\dots,\xi_{m}$ , such that the loss function $\mathcal{L}_{\omega}$ is

\mathcal{L}_{\omega}(\theta)\doteq\frac{1}{2}\|\theta-\xi_{\omega}\|_{2}^{2}.

(A.1.45)

Then $\operatorname*{arg\ min}_{\theta}\operatorname{\mathbb{E}}[\mathcal{L}_{\omega}(\theta)]=\frac{1}{m}\sum_{i=1}\xi_{i}$ , but stochastic gradient descent can “pinball” around the global optimum value, and not converge, as visualized in Figure A.4.

Figure A.4 : Stochastic gradient descent may not converge, even for very benign objectives, but Nesterov gradient converges. For even simple quadratic objectives, stochastic gradient descent iterates may pinball around the global optimum, whereas Nesterov gradients align to point to the optimal value. — Figure A.4: Stochastic gradient descent may not converge, even for very benign objectives, but Nesterov gradient converges. For even simple quadratic objectives, stochastic gradient descent iterates may pinball around the global optimum, whereas Nesterov gradients align to point to the optimal value.

In order to fix this, we can either average the parameters $\theta_{k}$ or average the gradients $\nabla\mathcal{L}_{k}(\theta_{k})$ over time. If we average the parameters $\theta_{k}$ , then (using Figure A.4 as a mental model) the issue of pinballing is straightforwardly not possible, since the average iterate will grow closer to the center. As such, most theoretical convergence proofs consider the convergence of the average iterate $\frac{1}{k}\sum_{i=0}^{k}\theta_{i}$ to the global minimum. If we average the gradients, we will eventually learn an average gradient $\frac{1}{k}\sum_{i=0}^{k}\nabla\mathcal{L}_{k}(\theta_{k})$ which does not change much at each iteration and therefore does not pinball.

In practice, instead of using an arithmetic average, we take an exponentially moving average (EMA) of the parameters (this is called Polyak momentum) or of the gradients (this is called Nesterov momentum). Nesterov momentum is more popular and we will study it here.

A stochastic gradient descent iteration with Nesterov momentum is as follows:

	$\displaystyle\bm{g}_{k}$	$\displaystyle=\beta\bm{g}_{k-1}+(1-\beta)\nabla\mathcal{L}_{k}(\theta_{k})$		(A.1.46)
	$\displaystyle\theta_{k+1}$	$\displaystyle=\theta_{k}-h\bm{g}_{k}.$		(A.1.47)

We do not go through a convergence proof (see Chapter 7 of [GG23] for an example). However, Nesterov momentum handles our toy case in Figure A.4 easily (see the right-hand figure): it stops pinballing and eventually converges to the global optimum.

We end with a caveat: one can show that Polyak momentum and Nesterov momentum are equivalent, for certain choices of parameter settings. Then it is also possible to show that a decaying learning rate schedule (i.e., the learning rate $h$ depends on the iteration $k$ , and its limit is $h_{k}\to 0$ as $k\to\infty$ ) with plain SGD (or PSGD) can mimic the effect of momentum. Namely, [DCM+23] shows that if the SGD algorithm lasts $K$ iterations, the gradient norms are bounded $\|\nabla\mathcal{L}(\theta_{k})\|_{2}\leq G$ , and we define $D\doteq\|\theta_{0}-\theta^{\star}\|_{2}$ , then plain SGD iterates $\theta_{k}$ satisfy the rate $\operatorname{\mathbb{E}}[\mathcal{L}(\theta_{k})-\mathcal{L}(\theta^{\star})]\leq DG/\sqrt{K}$ — but only so long as the learning rate $h_{k}=(D/[G\sqrt{K}])(1-k/K)$ decays linearly with time. This matches learning rate schedules used in practice. Indeed, surprisingly, such a theory of convex optimization can predict many empirical phenomena in deep networks [SHT+25], despite deep learning optimization being highly non-convex and non-smooth in the worst case. It is so far unclear why this is the case.

A.1.5 Putting Everything Together: Adam

The gradient descent scheme proposes an iteration of the form

\theta_{k+1}=\theta_{k}+h\bm{v}_{k},

(A.1.48)

where (recall) $\bm{v}_{k}$ is chosen to be (proportional to) the steepest descent vector in the Euclidean norm:

\bm{v}_{k}=-\frac{\nabla\mathcal{L}(\theta_{k})}{\|\nabla\mathcal{L}(\theta_{k})\|_{2}}\in\operatorname*{arg\ min}_{\begin{subarray}{c}\bm{v}\in\mathbb{R}^{n}\\ \|\bm{v}\|_{2}=1\end{subarray}}\langle\nabla\mathcal{L}(\theta_{k}),\bm{v}\rangle.

(A.1.49)

However, in the context of deep learning optimization, there is absolutely nothing which implies that we have to use the Euclidean norm; indeed the “natural geometry” of the space of parameters is not well-respected by the Euclidean norm, since small changes in the parameter space can lead to very large differences in the output space, for a particular fixed input to the network. If we were instead to use a generic norm $\|\cdot\|$ on the parameter space $\mathbb{R}^{n}$ , we would get some other quantity corresponding to the so-called dual norm:

\bm{v}_{k}\in\operatorname*{arg\ min}_{\begin{subarray}{c}\bm{v}\in\mathbb{R}^{n}\\ \|\bm{v}\|=1\end{subarray}}\langle\nabla\mathcal{L}(\theta_{k}),\bm{v}\rangle.

(A.1.50)

For instance, if we were to use the $\ell^{\infty}$ -norm, it is possible to show that

\bm{v}_{k}=-\operatorname{sign}(\nabla\mathcal{L}(\theta_{k}))\in\operatorname*{arg\ min}_{\begin{subarray}{c}\bm{v}\in\mathbb{R}^{n}\\ \|\bm{v}\|_{\infty}=1\end{subarray}}\langle\nabla\mathcal{L}(\theta_{k}),\bm{v}\rangle,

(A.1.51)

where $\operatorname{sign}(\bm{x})_{i}=\operatorname{sign}(x_{i})\in\{-1,0,1\}$ . Thus if we were so-inclined, we could use the so-called sign-gradient descent:

\theta_{k+1}=\theta_{k}-h\operatorname{sign}(\nabla\mathcal{L}(\theta_{k})).

(A.1.52)

From sign-gradient descent, we can derive the famous Adam optimization algorithm. Note that for a scalar $x\in\mathbb{R}$ we can write

\operatorname{sign}(x)=\frac{x}{\lvert x\rvert}=\frac{x}{\sqrt{x^{2}}}.

(A.1.53)

Similarly, for a vector $\bm{x}\in\mathbb{R}^{n}$ we write (where $\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}}$ and $\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}}$ are element-wise multiplication and division)

\operatorname{sign}(\bm{x})=\bm{x}\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}}[\bm{x}^{\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}}2}]^{\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}}(1/2)}.

(A.1.54)

Using this representation we can write (A.1.52) as

\theta_{k+1}=\theta_{k}-h([\nabla\mathcal{L}(\theta_{k})]\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}}[\nabla\mathcal{L}(\theta_{k})^{\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}}2}]^{\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}}\frac{1}{2}}).

(A.1.55)

Now consider the stochastic regime where we are optimizing a different loss $L_{k}$ at each iteration. In SGD, we “tracked” (i.e., took an average of) the gradients using Nesterov momentum. Here, we can track both the gradient and the squared gradient using momentum, i.e.,

$\displaystyle\bm{g}_{k}$	$\displaystyle=\beta^{1}\bm{g}_{k-1}+(1-\beta^{1})\nabla\mathcal{L}_{k}(\theta_{k})$	(A.1.56)
$\displaystyle\bm{s}_{k}$	$\displaystyle=\beta^{2}\bm{s}_{k-1}+(1-\beta^{2})[\nabla\mathcal{L}_{k}(\theta_{k})]^{\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}}2}$	(A.1.57)
$\displaystyle\theta_{k+1}$	$\displaystyle=\theta_{k}-h\bm{g}_{k}\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\oslash$}}{\scalebox{0.8}{$\textstyle\oslash$}}{\scalebox{0.8}{$\scriptstyle\oslash$}}{\scalebox{0.8}{$\scriptscriptstyle\oslash$}}$}}}\bm{s}_{k}^{\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}}\frac{1}{2}},$	(A.1.58)

where $\beta^{i}\in[0,1]$ are the momentum parameters. The algorithm presented by this iteration is the celebrated Adam optimizer,⁵⁵5In order to avoid division-by-zero errors, we divide by $\bm{s}_{k}^{\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}}(1/2)}+\varepsilon\bm{1}_{n}$ where $\varepsilon$ is small, say on the order of $10^{-8}$ . which is the most-used optimizer in deep learning. While convergence proofs of Adam are more involved, it falls out of the same steepest descent principle we used so far, and so we should expect that given a small enough learning rate, each update should improve the loss.

Another way to view Adam, which partially explains its empirical success, is that it dynamically updates the learning rates for each parameter based on the squared gradients. In particular, notice that we can write

\theta_{k+1}=\theta_{k}-\eta_{k}\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}}\bm{g}_{k}\qquad\text{where}\qquad\eta_{k}=h\bm{s}_{k}^{\mathbin{\mathchoice{\raisebox{1.3pt}{$\displaystyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{1.3pt}{$\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.75pt}{$\scriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}{\raisebox{0.6pt}{$\scriptscriptstyle\mathchoice{\scalebox{0.8}{$\displaystyle\odot$}}{\scalebox{0.8}{$\textstyle\odot$}}{\scalebox{0.8}{$\scriptstyle\odot$}}{\scalebox{0.8}{$\scriptscriptstyle\odot$}}$}}}(-\frac{1}{2})}

(A.1.59)

where $\eta_{k}$ is the parameter-wise adaptively set learning rate. This scheme is called adaptive because if the gradient of a particular parameter is large up to iteration $k$ , then the learning rate for this parameter becomes smaller to compensate, and vice versa, as can be seen from the above equation.

A.2 Computing Gradients via Automatic Differentiation

Above, we discussed several optimization algorithms for deep networks which assumed access to a first-order oracle, i.e., a device which would allow us to compute $\mathcal{L}(\theta)$ and $\nabla\mathcal{L}(\theta)$ . For simple functions $\mathcal{L}$ , it is possible to do this by hand. However, for deep neural networks, this quickly becomes tedious, and hinders rapid experimentation. Hence we require a general algorithm which would allow us to efficiently compute the gradients of arbitrary (sub)differentiable network architectures.

In this section, we introduce the basics of automatic differentiation (AD or autodiff), which is a computationally efficient way to compute gradients and Jacobians of general functions $f:\mathbb{R}^{m}\to\mathbb{R}^{n}$ . We will show how this leads to the backpropagation algorithm for computing gradients of loss functions involving neural networks. A summary of the structure of this section is as follows:

1.

We introduce differentials, a convenient formalism for calculating and organizing the derivatives of functions between high-dimensional parameter spaces (which may themselves be products of other spaces involving matrices, tensors, etc.).
2.

We describe the basics of forward-mode and reverse-mode automatic differentiation, which involves considerations that are important for efficient computation of gradients/Jacobians for different kinds of functions arising in machine learning applications.
3.

We describe backpropagation in the special case of a loss function applied to a stack-of-layers neural network as an instantiation of reverse-mode automatic differentiation.

Our treatment will err on the mathematical side, to give the reader a deep understanding of the underlying mathematics. The reader should ensure to couple this understanding with a strong grasp of practical aspects of automatic differentiation for deep learning, for example as manifested in the outstanding tutorial of [Kar22a].

A.2.1 Differentials

A full accounting of this subsection is given in the excellent guide [BEJ25]. To motivate differentials, let us first consider the simple example of a differentiable function $\mathcal{L}\colon\mathbb{R}\to\mathbb{R}$ acting on a parameter $\theta$ . We can write

\mathcal{L}(\theta)-\mathcal{L}(\theta_{0})=\mathcal{L}^{\prime}(\theta_{0})\cdot(\theta-\theta_{0})+o(\lvert\theta-\theta_{0}\rvert).

(A.2.1)

If we take $\delta\theta\doteq\theta-\theta_{0}$ and $\delta\mathcal{L}\doteq\mathcal{L}(\theta_{0}+\delta\theta)-\mathcal{L}(\theta_{0})$ , we can write

\delta\mathcal{L}=\mathcal{L}^{\prime}(\theta_{0})\cdot\delta\theta+o(\lvert\delta\theta\rvert).

(A.2.2)

We will (non-rigorously) define $\mathrm{d}\theta$ and $\mathrm{d}\mathcal{L}$ , i.e., the differentials of $\theta$ and $\mathcal{L}$ , to be infinitesimally small changes in $\theta$ and $\mathcal{L}$ . Think of them as what one gets when $\delta\theta$ (and therefore $\delta\mathcal{L}$ ) are extremely small. The goal of differential calculus, in some sense, is to study the relationships between the differentials $\mathrm{d}\theta$ and $\mathrm{d}\mathcal{L}$ , namely, seeing how small changes in the input of a function change the output. We should note that the differential $\mathrm{d}\theta$ is the same shape as $\theta$ , and the differential $\mathrm{d}\mathcal{L}$ is the same shape as $\mathcal{L}$ . In particular, we can write

\mathrm{d}\mathcal{L}=\mathcal{L}^{\prime}(\theta)\cdot\mathrm{d}\theta,

(A.2.3)

whereby we have that all higher powers of $\lvert\mathrm{d}\theta\rvert$ , such as $(\mathrm{d}\theta)^{2}$ , are $0$ .

Let’s see how this works for a higher dimensions, i.e., $\mathcal{L}\colon\mathbb{R}^{n}\to\mathbb{R}$ . Then we still have

\mathrm{d}\mathcal{L}=\mathcal{L}^{\prime}(\theta)\cdot\mathrm{d}\theta

(A.2.4)

for some notion of a derivative $\mathcal{L}^{\prime}(\theta)$ . Since $\theta$ (hence $\mathrm{d}\theta$ ) is a column vector here and $\mathcal{L}$ (hence $\mathrm{d}\mathcal{L}$ ) is a scalar, we must have that $\mathcal{L}^{\prime}(\theta)$ is a row vector. In this case, $\mathcal{L}^{\prime}(\theta)$ is the Jacobian of $\mathcal{L}$ w.r.t. $\theta$ . Here notice that we have set all higher powers and products of coordinates of $\mathrm{d}\theta$ to $0$ . In sum,

All products and powers $\geq 2$ of differentials are equal to $0$ .

Now consider a higher-order tensor function $\mathcal{L}\colon\mathbb{R}^{m\times n}\to\mathbb{R}^{p\times q}$ . Then our basic linearization equation is insufficient for this case: $\mathrm{d}\mathcal{L}=\mathcal{L}^{\prime}(\theta)\cdot\mathrm{d}\theta$ does not make sense since $\theta$ is an $m\times n$ matrix but $\mathrm{d}\mathcal{L}$ is a $p\times q$ matrix, so there is no possible vector or matrix shape for $\mathcal{L}^{\prime}(\theta)$ that works in general (as no matrix can multiply a $m\times n$ matrix to form a $p\times q$ matrix unless $m=p$ ). So we must have a slightly more advanced interpretation.

Namely, we consider $\mathcal{L}^{\prime}(\theta)$ as a linear transformation whose input is $\theta$ -space and whose output is $\mathcal{L}$ -space, which takes in a small change in $\theta$ and outputs the corresponding small change in $\mathcal{L}$ . Namely, we can write

\mathrm{d}\mathcal{L}=\mathcal{L}^{\prime}(\theta)[\mathrm{d}\theta].

(A.2.5)

In the previous cases, $\mathcal{L}^{\prime}(\theta)$ was first a linear operator $\mathbb{R}\to\mathbb{R}$ whose action was to multiply its input by the scalar derivative of $\mathcal{L}$ with respect to $\theta$ , and then a linear operator $\mathbb{R}^{n}\to\mathbb{R}$ whose action was to multiply its input by the Jacobian derivative of $\mathcal{L}$ with respect to $\theta$ . In general $\mathcal{L}^{\prime}(\theta)$ is the “derivative” of $\mathcal{L}$ w.r.t. $\theta$ . Think of $\mathcal{L}^{\prime}$ as a generalized version of the Jacobian of $\mathcal{L}$ . As such, it follows some simple derivative rules, most crucially the chain rule.

Theorem A.1 (Differential Chain Rule).

Suppose $\mathcal{L}=f\circ g$ where $f$ and $g$ are differentiable. Then

\mathrm{d}\mathcal{L}=f^{\prime}(g(\theta))g^{\prime}(\theta)[\mathrm{d}\theta],

(A.2.6)

where (as usual) multiplication indicates composition of linear operators. In particular,

\mathcal{L}^{\prime}(\theta)=f^{\prime}(g(\theta))g^{\prime}(\theta)

(A.2.7)

in the sense of equality of linear operators.

It is productive to think of the multivariate chain rule in functorial terms: composition of functions gets ‘turned into’ matrix multiplication of Jacobians (composition of linear operators!). We illustrate the power of this result and this perspective through several examples.

Example A.4.

Consider the function $f(\bm{X})=\bm{W}\bm{X}+\bm{b}\bm{1}^{\top}$ . Then

\mathrm{d}f=f(\bm{X}+\mathrm{d}\bm{X})-f(\bm{X})=[\bm{W}(\bm{X}+\mathrm{d}\bm{X})+\bm{b}\bm{1}^{\top}]-[\bm{W}\bm{X}+\bm{b}\bm{1}^{\top}]=\bm{W}\mathrm{d}\bm{X}.

(A.2.8)

Thus the derivative of an affine function w.r.t. its input is

f^{\prime}(\bm{X})[\mathrm{d}\bm{X}]=\bm{W}\mathrm{d}\bm{X}\implies f^{\prime}(\bm{X})=\bm{W}.

(A.2.9)

Notice that $f^{\prime}$ is constant. On the other hand, consider the function $g(\bm{W},\bm{b})=\bm{W}\bm{X}+\bm{b}\bm{1}^{\top}$ . Then

	$\displaystyle\mathrm{d}g$	$\displaystyle=g(\bm{W}+\mathrm{d}\bm{W},\bm{b}+\mathrm{d}\bm{b})-g(\bm{W},\bm{b})=[(\bm{W}+\mathrm{d}\bm{W})\bm{X}+(\bm{b}+\mathrm{d}\bm{b})\bm{1}^{\top}]-[\bm{W}\bm{X}+\bm{b}]$		(A.2.10)
		$\displaystyle=(\mathrm{d}\bm{W})\bm{X}+(\mathrm{d}\bm{b})\bm{1}^{\top}=g^{\prime}(\bm{W},\bm{B})[\mathrm{d}\bm{W},\mathrm{d}\bm{b}].$		(A.2.11)

Notice that this derivative is constant in $\bm{W},\bm{b}$ (which makes sense since $g$ itself is linear) and linear in the differential inputs $\mathrm{d}\bm{W},\mathrm{d}\bm{b}$ . $\blacksquare$

Example A.5.

Consider the function $f=gh$ where $g,h$ are differentiable functions whose outputs can multiply together. Then $f=p\circ v$ where $v=(g,h)$ and $p(a,b)=ab$ . Applying the chain rule we have

\mathrm{d}f=p^{\prime}(v(x))v^{\prime}(x)[\mathrm{d}x].

(A.2.12)

To compute $v^{\prime}(x)$ we can compute

\mathrm{d}v=v^{\prime}(x)[\mathrm{d}x]=v(x+\mathrm{d}x)-v(x)=\begin{bmatrix}g(x+\mathrm{d}x)-g(x)\\ h(x+\mathrm{d}x)-h(x)\end{bmatrix}=\begin{bmatrix}g^{\prime}(x)[\mathrm{d}x]\\ h^{\prime}(x)[\mathrm{d}x]\end{bmatrix}.

(A.2.13)

To compute $p^{\prime}$ we can compute

	$\displaystyle\mathrm{d}p$	$\displaystyle=p^{\prime}(a,b)[\mathrm{d}a,\mathrm{d}b]=p(a+\mathrm{d}a,b+\mathrm{d}b)-p(a,b)=(a+\mathrm{d}a)(b+\mathrm{d}b)-ab$		(A.2.14)
		$\displaystyle=(\mathrm{d}a)b+a(\mathrm{d}b)+(\mathrm{d}a)(\mathrm{d}b)=(\mathrm{d}a)b+a(\mathrm{d}b),$		(A.2.15)

where (recall) the product of the differentials $\mathrm{d}a$ and $\mathrm{d}b$ is set to $0$ . Therefore

p^{\prime}(a,b)[\mathrm{d}a,\mathrm{d}b]=(\mathrm{d}a)b+(\mathrm{d}b)a.

(A.2.16)

Putting these together, we find

	$\displaystyle f^{\prime}(x)[\mathrm{d}x]$	$\displaystyle=p^{\prime}(v(x))v^{\prime}(x)[\mathrm{d}x]=p^{\prime}(g(x),h(x))[g^{\prime}(x)[\mathrm{d}x],h^{\prime}(x)[\mathrm{d}x]]$		(A.2.17)
		$\displaystyle=(g^{\prime}(x)[\mathrm{d}x])h(x)+g(x)(h^{\prime}(x)[\mathrm{d}x]).$		(A.2.18)

This gives

f^{\prime}(x)[\mathrm{d}x]=(g^{\prime}(x)[\mathrm{d}x])h(x)+g(x)(h^{\prime}(x)[\mathrm{d}x]).

(A.2.19)

If for example we say that $f,g,h\colon\mathbb{R}\to\mathbb{R}$ then everything commutes so

f^{\prime}(x)[\mathrm{d}x]=(g^{\prime}(x)h(x)+g(x)h^{\prime}(x))[\mathrm{d}x]\implies f^{\prime}(x)=g^{\prime}(x)h(x)+g(x)h^{\prime}(x)

(A.2.20)

which is the familiar product rule. $\blacksquare$

Example A.6.

Consider the function $f(\bm{A})=\bm{A}^{\top}\bm{A}\bm{B}\bm{A}$ where $\bm{A}$ is a matrix and $\bm{B}$ is a constant matrix. Then, letting $f=gh$ where $g(\bm{A})=\bm{A}^{\top}\bm{A}$ and $h(\bm{A})=\bm{B}\bm{A}$ , we can use the product rule to obtain

	$\displaystyle f^{\prime}(\bm{A})[\mathrm{d}\bm{A}]$	$\displaystyle=(g^{\prime}(\bm{A})[\mathrm{d}\bm{A}])h(\bm{A})+g(\bm{A})(h^{\prime}(\bm{A})[\mathrm{d}\bm{A}])$		(A.2.21)
		$\displaystyle=((\mathrm{d}\bm{A})^{\top}\bm{A}+\bm{A}^{\top}(\mathrm{d}\bm{A}))\bm{B}\bm{A}+\bm{A}^{\top}\bm{A}\bm{B}(\mathrm{d}\bm{A}).$		(A.2.22)

$\blacksquare$

Example A.7.

Consider the function $f\colon\mathbb{R}^{m\times n\times k}\to\mathbb{R}^{m\times n}$ given by

f(\bm{A})_{ij}=\sum_{t=1}^{k}A_{ijt}.

(A.2.23)

We cannot write a (matrix-valued) Jacobian or gradient for this function. But we can compute its differential just fine:

\mathrm{d}f_{ij}=[f(\bm{A}+\mathrm{d}\bm{A})-f(\bm{A})]_{ij}=\sum_{t=1}^{k}\mathrm{d}\bm{A}_{ijt}=\bm{1}_{k}^{\top}(\mathrm{d}\bm{A})_{ij}.

(A.2.24)

(f^{\prime}(\bm{A})[\mathrm{d}\bm{A}])_{ij}=\bm{1}_{k}^{\top}(\mathrm{d}\bm{A})_{ij},

(A.2.25)

which represents a higher-order tensor multiplication operation that is nonetheless well-defined. $\blacksquare$

This gives us all the technology we need to compute differentials of everything. The last thing we cover in this section is a method to compute gradients using the differential. Namely, for a function $\mathcal{L}$ whose output is a scalar, the gradient $\nabla\mathcal{L}$ is defined as

\mathrm{d}\mathcal{L}=\mathcal{L}^{\prime}(\theta)[\mathrm{d}\theta]=\langle\nabla\mathcal{L}(\theta),\mathrm{d}\theta\rangle,

(A.2.26)

where the inner product here is the “standard” inner product for the specified objects (i.e., for vectors it’s the $\ell^{2}$ inner product, whereas for matrices it’s the Frobenius inner product, and for higher-order tensors it’s the analogous sum-of-coordinates inner product). This definition is the correct generalization of the ‘familiar’ example of the gradient of a function from $\mathbb{R}^{n}$ to $\mathbb{R}$ as the vector of partial derivatives—a version of Taylor’s theorem for general functions $f:\mathbb{R}^{m}\to\mathbb{R}^{n}$ makes this connection rigorous. So one way to compute the gradient $\nabla\mathcal{L}$ is to compute the differential $\mathrm{d}\mathcal{L}$ and rewrite it in the form $\langle\text{something},\mathrm{d}\theta\rangle$ , then that “something” is the gradient.

A.2.2 Automatic Differentiation

The main idea of AD is to compute the chain rule efficiently. The basic problem we need to cope with is the following. In the optimization section of the appendix, we considered that the parameter space $\Theta$ was an abstract Euclidean space like $\mathbb{R}^{n}$ . In practice the parameters are really some collection of vectors, matrices, and higher-order objects: $\Theta=\mathbb{R}^{m\times n}\times\mathbb{R}^{n}\times\mathbb{R}^{r\times q\times p}\times\mathbb{R}^{r\times q}\times\cdots$ . While in theory this is the same thing as a large parameter space $\mathbb{R}^{n^{\prime}}$ for some (very large) $n^{\prime}$ , computationally efficient algorithms for differentiation must treat these two spaces differently. Forward and reverse mode automatic differentiation are two different schemes for performing this computation.

Let us do a simple example to start. Let $\mathcal{L}$ be defined by $\mathcal{L}=a\circ b\circ c$ where $a,b,c$ are differentiable. Then the chain rule gives

\mathcal{L}^{\prime}(\theta)=a^{\prime}(b(c(\theta)))b^{\prime}(c(\theta))c^{\prime}(\theta).

(A.2.27)

To compute $\mathcal{L}(\theta)$ , we first compute $c(\theta)$ then $b(c(\theta))$ then $a(b(c(\theta)))$ , and store them all. There are two ways to compute $\mathcal{L}^{\prime}(\theta)$ . The forward-mode AD will compute

c^{\prime}(\theta)\implies b^{\prime}(c(\theta))c^{\prime}(\theta)\implies a^{\prime}(b(c(\theta)))b^{\prime}(c(\theta))c^{\prime}(\theta)

(A.2.28)

i.e., computing the derivatives “from the bottom-up”. The reverse mode AD will compute

a^{\prime}(b(c(\theta)))\implies a^{\prime}(b(c(\theta)))b^{\prime}(c(\theta))\implies a^{\prime}(b(c(\theta)))b^{\prime}(c(\theta))c^{\prime}(\theta),

(A.2.29)

i.e., computing the derivatives “from the top down”. To see why this matters, suppose that $f\colon\mathbb{R}^{p}\to\mathbb{R}^{s}$ is given by $f=a\circ b\circ c$ where $a\colon\mathbb{R}^{r}\to\mathbb{R}^{s}$ , $b\colon\mathbb{R}^{q}\to\mathbb{R}^{r}$ , $c\colon\mathbb{R}^{p}\to\mathbb{R}^{q}$ . Then the chain rule is:

f^{\prime}(\bm{x})=a^{\prime}(b(c(\bm{x})))b^{\prime}(c(\bm{x}))c^{\prime}(\bm{x})

(A.2.30)

where (recall) $f^{\prime}$ is the derivative, in this case the Jacobian (since the input and output of each function are both vectors). Assuming that computing each Jacobian is trivial and the only cost is multiplying the Jacobians together, forward-mode AD has the following computational costs (assuming that multiplying $A\in\mathbb{R}^{m\times n},B\in\mathbb{R}^{n\times k}$ takes $\mathcal{O}(mnk)$ time):

	$\displaystyle\text{computing}\ c^{\prime}(\bm{x})\in\mathbb{R}^{q\times p}\ \text{takes negligible time}$		(A.2.31)
	$\displaystyle\text{computing}\ b^{\prime}(c(\bm{x}))c^{\prime}(x)\in\mathbb{R}^{r\times p}\ \text{takes $\mathcal{O}(pqr)$ time}$		(A.2.32)
	$\displaystyle\text{computing}\ a^{\prime}(b(c(\bm{x})))b^{\prime}(c(\bm{x}))c^{\prime}(\bm{x})\in\mathbb{R}^{s\times p}\ \text{takes $\mathcal{O}(pqr+prs)$ time.}$		(A.2.33)

Meanwhile, doing reverse-mode AD has the following computational costs:

	$\displaystyle\text{computing}\ a^{\prime}(b(c(\bm{x})))\in\mathbb{R}^{s\times r}\ \text{takes negligible time}$		(A.2.34)
	$\displaystyle\text{computing}\ a^{\prime}(b(c(\bm{x})))b^{\prime}(c(\bm{x}))\in\mathbb{R}^{s\times q}\ \text{takes $\mathcal{O}(qrs)$ time}$		(A.2.35)
	$\displaystyle\text{computing}\ a^{\prime}(b(c(\bm{x})))b^{\prime}(c(\bm{x}))c^{\prime}(\bm{x})\in\mathbb{R}^{s\times p}\ \text{takes $\mathcal{O}(qrs+pqs)$ time.}$		(A.2.36)

In other words, the forward-mode AD takes $\mathcal{O}(p(qr+rs))$ time, and the reverse-mode AD takes $\mathcal{O}(s(pq+qr))$ time. These take a different amount of time!

More generally, suppose that $f=f^{L}\circ\cdots\circ f^{1}$ where each $f^{\ell}\colon\mathbb{R}^{d^{\ell-1}}\to\mathbb{R}^{d^{\ell}}$ , so that $f\colon\mathbb{R}^{d^{0}}\to\mathbb{R}^{d^{L}}$ . Then the forward-mode AD takes $\mathcal{O}(d^{0}(\sum_{\ell=2}^{L}d^{\ell-1}d^{\ell}))$ time while the reverse-mode AD takes $\mathcal{O}(d^{L}(\sum_{\ell=1}^{L-1}d^{\ell-1}d^{\ell}))$ time. From the above rates, we see that all else equal:

•

If the function to optimize has more outputs than inputs (i.e., $d^{L}>d^{0}$ ), use forward-mode AD.
•

If the function to optimize has more inputs than outputs (i.e., $d^{0}>d^{L}$ ), use reverse-mode AD.

In a neural network, we compute the gradient of a loss function $\mathcal{L}\colon\Theta\to\mathbb{R}$ , where the parameter space $\Theta$ is usually very high-dimensional. So in practice we always use reverse-mode AD for training neural networks. Reverse-mode AD, in the context of training neural networks, is called backpropagation.

A.2.3 Back Propagation

In this section, we will discuss algorithmic backpropagation using a simple yet completely practical example. Suppose that we fix an input-label pair $(\bm{X},\bm{y})$ , and fix a network architecture $f_{\theta}=f_{\theta}^{L}\circ\cdots\circ f_{\theta}^{1}\circ f_{\theta}^{\mathrm{emb}}$ where $\theta=(\theta^{\mathrm{emb}},\theta^{1},\dots,\theta^{m},\theta^{\mathrm{head}})$ and task-specific head $h_{\theta^{\mathrm{head}}}$ , and write

	$\displaystyle\bm{Z}_{\theta}^{1}(\bm{X})\doteq f_{\theta}^{\mathrm{emb}}(\bm{X}),$		(A.2.37)
	$\displaystyle\bm{Z}_{\theta}^{\ell+1}(\bm{X})\doteq f_{\theta}^{\ell}(\bm{Z}_{\theta}^{\ell}(\bm{X})),\quad\forall\ell\in\{1,\dots,L\},$		(A.2.38)
	$\displaystyle\hat{\bm{y}}_{\theta^{\mathrm{head}}}(\bm{X})\doteq h_{\theta}(\bm{Z}_{\theta}^{L+1}(\bm{X})).$		(A.2.39)

Then, we can define the loss on this one term by

\mathcal{L}(\theta)\doteq\mathsf{L}(\bm{y},\hat{\bm{y}}_{\theta}(\bm{X})),

(A.2.40)

where $\mathsf{L}$ is a differentiable function of its second argument. Then the backward-mode AD instructs us to compute the derivatives in the order $\theta^{\mathrm{head}},\theta^{L},\dots,\theta^{1},\theta^{\mathrm{emb}}$ .

To carry out this computation, let us make some changes to the notation.

•

First, let us change the notation to emphasize the dependency structure between the variables. Namely,

	$\displaystyle\bm{Z}^{1}\doteq f^{\mathrm{emb}}(\bm{X},\theta^{\mathrm{emb}})$		(A.2.41)
	$\displaystyle\bm{Z}^{\ell+1}\doteq f^{\ell}(\bm{Z}^{\ell},\theta^{\ell}),\qquad\forall\ell\in\{1,\dots,L\},$		(A.2.42)
	$\displaystyle\hat{\bm{y}}\doteq h(\bm{Z}^{L+1},\theta^{\mathrm{head}}),$		(A.2.43)
	$\displaystyle\mathcal{L}\doteq\mathsf{L}(\bm{y},\hat{\bm{y}}).$		(A.2.44)

•

Then, instead of having the derivative be $f^{\prime}$ , we explicitly notate the independent variable and write the derivative as $\frac{\mathrm{d}f}{\mathrm{d}\theta^{1}}$ , for example. This is because there are many variables in our model and we only care about one at a time.

We can start by computing the appropriate differentials. First, for $\theta^{\mathrm{head}}$ we have

$\displaystyle\mathrm{d}\mathcal{L}$	$\displaystyle=\mathrm{d}\mathsf{L}$	(A.2.45)
	$\displaystyle=\frac{\mathrm{d}\mathsf{L}}{\mathrm{d}\hat{\bm{y}}}\cdot\mathrm{d}\hat{\bm{y}}$	(A.2.46)
	$\displaystyle=\frac{\mathrm{d}\mathsf{L}}{\mathrm{d}\hat{\bm{y}}}\cdot\mathrm{d}{(h(\bm{Z}^{L+1},\theta^{\mathrm{head}}))}$	(A.2.47)
	$\displaystyle=\frac{\mathrm{d}\mathsf{L}}{\mathrm{d}\hat{\bm{y}}}\left[\frac{\mathrm{d}h}{\mathrm{d}\bm{Z}^{L+1}}\cdot\mathrm{d}\bm{Z}^{L+1}+\frac{\mathrm{d}h}{\mathrm{d}\theta^{\mathrm{head}}}\cdot\mathrm{d}\theta^{\mathrm{head}}\right].$	(A.2.48)

Now since $\bm{Z}^{L+1}$ does not depend on $\theta^{\mathrm{head}}$ , we have $\mathrm{d}\bm{Z}^{L+1}=0$ , so in the end it holds (using the fact that the gradient is the transpose of the derivative for a function $\mathsf{L}$ whose codomain is $\mathbb{R}$ ):

$\displaystyle\mathrm{d}\mathcal{L}$	$\displaystyle=\frac{\mathrm{d}\mathsf{L}}{\mathrm{d}\hat{\bm{y}}}\cdot\frac{\mathrm{d}h}{\mathrm{d}\theta^{\mathrm{head}}}\cdot\mathrm{d}\theta^{\mathrm{head}}$	(A.2.49)
	$\displaystyle=[\nabla_{\hat{\bm{y}}}\mathsf{L}]^{\top}\frac{\mathrm{d}h}{\mathrm{d}\theta^{\mathrm{head}}}\cdot\mathrm{d}\theta^{\mathrm{head}}$	(A.2.50)
	$\displaystyle=\left\langle\nabla_{\hat{\bm{y}}}\mathsf{L},\frac{\mathrm{d}h}{\mathrm{d}\theta^{\mathrm{head}}}\cdot\mathrm{d}\theta^{\mathrm{head}}\right\rangle$	(A.2.51)
	$\displaystyle=\left\langle\left(\frac{\mathrm{d}h}{\mathrm{d}\theta^{\mathrm{head}}}\right)^{*}\nabla_{\hat{\bm{y}}}\mathsf{L},\mathrm{d}\theta^{\mathrm{head}}\right\rangle$	(A.2.52)
	$\displaystyle=\langle\nabla_{\theta^{\mathrm{head}}}\mathcal{L},\mathrm{d}\theta^{\mathrm{head}}\rangle.$	(A.2.53)

Thus to compute $\nabla_{\theta^{\mathrm{head}}}\mathcal{L}$ , we compute the gradient $\nabla_{\hat{\bm{y}}}\mathsf{L}$ and the adjoint⁶⁶6The adjoint is like a generalized transpose for more general linear transformations. Particularly, for a given pair of inner product spaces and linear transformation $T$ between those spaces, the adjoint $T^{*}$ is defined by the identity $\langle T\bm{x},\bm{y}\rangle=\langle\bm{x},T^{*}\bm{y}\rangle$ . In finite dimensions (i.e., all cases relevant to this book) the adjoint exists and is unique. of the derivative $\frac{\mathrm{d}h}{\mathrm{d}\theta^{\mathrm{head}}}$ and multiply (i.e., apply the adjoint linear transformation to the gradient). In practice, both derivatives can be computed by hand, but many modern computational frameworks can automatically define the derivatives (and/or their adjoints) given code for the “forward pass,” i.e., the loss function computation. While extending this automatic derivative definition to as many functions as possible is an area of active research, the resource [BEJ25] describes one basic approach to do it in some detail. By the way, backpropagation is also called the adjoint method for this reason — i.e., that we use adjoints derivatives to compute the gradient.

Now let us compute the differentials w.r.t. some $\theta^{\ell}$ :

$\displaystyle\mathrm{d}\mathcal{L}$	$\displaystyle=\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\bm{Z}^{\ell+1}}\cdot\mathrm{d}\bm{Z}^{\ell+1}$	(A.2.54)
	$\displaystyle=\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\bm{Z}^{\ell+1}}\cdot\mathrm{d}{(f^{\ell}(\bm{Z}^{\ell},\theta^{\ell}))}$	(A.2.55)
	$\displaystyle=\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\bm{Z}^{\ell+1}}\left[\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\bm{Z}^{\ell}}\cdot\mathrm{d}\bm{Z}^{\ell}+\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\theta^{\ell}}\cdot\mathrm{d}\theta^{\ell}\right]$	(A.2.56)
	$\displaystyle=\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\bm{Z}^{\ell+1}}\cdot\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\theta^{\ell}}\cdot\mathrm{d}\theta^{\ell}\qquad\text{(b/c~{}$\bm{Z}^{\ell}$ isn't fn.~{}of $\theta^{\ell}$ so $\mathrm{d}\bm{Z}^{\ell}=0$)}$	(A.2.57)
	$\displaystyle=[\nabla_{\bm{Z}^{\ell+1}}\mathcal{L}]^{\top}\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\theta^{\ell}}\cdot\mathrm{d}\theta^{\ell}$	(A.2.58)
	$\displaystyle=\left\langle\nabla_{\bm{Z}^{\ell+1}}\mathcal{L},\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\theta^{\ell}}\cdot\mathrm{d}\theta^{\ell}\right\rangle$	(A.2.59)
	$\displaystyle=\left\langle\left(\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\theta^{\ell}}\right)^{*}\nabla_{\bm{Z}^{\ell+1}}\mathcal{L},\mathrm{d}\theta^{\ell}\right\rangle$	(A.2.60)
	$\displaystyle=\langle\nabla_{\theta^{\ell}}\mathcal{L},\mathrm{d}\theta^{\ell}\rangle.$	(A.2.61)

Thus to compute $\nabla_{\theta^{\ell}}\mathcal{L}$ we compute $\nabla_{\bm{Z}^{\ell+1}}\mathcal{L}$ then apply the adjoint derivative $\left(\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\theta^{\ell}}\right)^{*}$ to it. Since $\nabla_{\theta^{\ell}}\mathcal{L}$ depends on $\nabla_{\bm{Z}^{\ell+1}}\mathcal{L}$ , we also want to be able to compute the gradients w.r.t. $\bm{Z}^{\ell}$ . This can be computed in the exact same way:

$\displaystyle\mathrm{d}\mathcal{L}$	$\displaystyle=\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\bm{Z}^{\ell+1}}\cdot\mathrm{d}\bm{Z}^{\ell+1}$	(A.2.62)
	$\displaystyle=\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\bm{Z}^{\ell+1}}\cdot\mathrm{d}{(f^{\ell}(\bm{Z}^{\ell},\theta^{\ell}))}$	(A.2.63)
	$\displaystyle=\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\bm{Z}^{\ell+1}}\left[\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\bm{Z}^{\ell}}\cdot\mathrm{d}\bm{Z}^{\ell}+\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\theta^{\ell}}\cdot\mathrm{d}\theta^{\ell}\right]$	(A.2.64)
	$\displaystyle=\frac{\mathrm{d}\mathcal{L}}{\mathrm{d}\bm{Z}^{\ell+1}}\cdot\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\bm{Z}^{\ell}}\cdot\mathrm{d}\bm{Z}^{\ell}\qquad\text{(b/c~{}~{}$\theta^{\ell}$ isn't fn.~{}of $\bm{Z}^{\ell}$ so $\mathrm{d}\theta^{\ell}=0$ this time)}$	(A.2.65)
	$\displaystyle=\left\langle\left(\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\bm{Z}^{\ell}}\right)^{*}\nabla_{\bm{Z}^{\ell+1}}\mathcal{L},\mathrm{d}\bm{Z}^{\ell}\right\rangle\qquad\text{(same machine)}$	(A.2.66)
	$\displaystyle=\langle\nabla_{\bm{Z}^{\ell}}\mathcal{L},\mathrm{d}\bm{Z}^{\ell}\rangle.$	(A.2.67)

Thus to compute $\nabla_{\bm{Z}^{\ell}}\mathcal{L}$ we compute $\nabla_{\bm{Z}^{\ell+1}}\mathcal{L}$ then apply the adjoint derivative $\left(\frac{\mathrm{d}f^{\ell}}{\mathrm{d}\bm{Z}^{\ell}}\right)^{*}$ to it. So we have a recursion to compute $\nabla_{\bm{Z}^{\ell}}\mathcal{L}$ for all $\ell$ , with base case $\nabla_{\bm{Z}^{L+1}}\mathcal{L}$ , which is given by

$\displaystyle\mathrm{d}\mathcal{L}$	$\displaystyle=\mathrm{d}\mathsf{L}$	(A.2.68)
	$\displaystyle=\frac{\mathrm{d}\mathsf{L}}{\mathrm{d}\hat{\bm{y}}}\cdot\mathrm{d}\hat{\bm{y}}$	(A.2.69)
	$\displaystyle=\frac{\mathrm{d}\mathsf{L}}{\mathrm{d}\hat{\bm{y}}}\cdot\mathrm{d}{h(\bm{Z}^{L+1},\theta^{\mathrm{head}})}$	(A.2.70)
	$\displaystyle=\frac{\mathrm{d}\mathsf{L}}{\mathrm{d}\hat{\bm{y}}}\left[\frac{\mathrm{d}h}{\mathrm{d}\bm{Z}^{L+1}}\cdot\mathrm{d}\bm{Z}^{L+1}+\frac{\mathrm{d}h}{\mathrm{d}\theta^{\mathrm{head}}}\cdot\mathrm{d}\theta^{\mathrm{head}}\right]$	(A.2.71)
	$\displaystyle=\frac{\mathrm{d}\mathsf{L}}{\mathrm{d}\hat{\bm{y}}}\cdot\frac{\mathrm{d}h}{\mathrm{d}\bm{Z}^{L+1}}\cdot\mathrm{d}\bm{Z}^{L+1}$	(A.2.72)
	$\displaystyle=\left\langle\left(\frac{\mathrm{d}h}{\mathrm{d}\bm{Z}^{L+1}}\right)^{*}\nabla_{\hat{\bm{y}}}\mathsf{L},\mathrm{d}\bm{Z}^{L+1}\right\rangle$	(A.2.73)
	$\displaystyle=\langle\nabla_{\bm{Z}^{L+1}}\mathcal{L},\mathrm{d}\bm{Z}^{L+1}\rangle.$	(A.2.74)

Thus we have the recursion:

$\displaystyle\nabla_{\bm{Z}^{L+1}}\mathcal{L}$	$\displaystyle=\left(\frac{\mathrm{d}h}{\mathrm{d}\bm{Z}^{L+1}}\right)^{*}\nabla_{\hat{\bm{y}}}\mathsf{L}$	(A.2.75)
$\displaystyle\nabla_{\bm{Z}^{L}}\mathcal{L}$	$\displaystyle=\left(\frac{\mathrm{d}f^{L}}{\mathrm{d}\bm{Z}^{L}}\right)^{*}\nabla_{\bm{Z}^{L+1}}\mathcal{L}$	(A.2.76)
$\displaystyle\nabla_{\theta^{L}}\mathcal{L}$	$\displaystyle=\left(\frac{\mathrm{d}f^{L}}{\mathrm{d}\theta^{L}}\right)^{*}\nabla_{\bm{Z}^{L+1}}\mathcal{L}$	(A.2.77)
$\displaystyle\vdots$		(A.2.78)
$\displaystyle\nabla_{\bm{Z}^{1}}\mathcal{L}$	$\displaystyle=\left(\frac{\mathrm{d}f^{1}}{\mathrm{d}\bm{Z}^{1}}\right)^{*}\nabla_{\bm{Z}^{2}}\mathcal{L}$	(A.2.79)
$\displaystyle\nabla_{\theta^{1}}\mathcal{L}$	$\displaystyle=\left(\frac{\mathrm{d}f^{1}}{\mathrm{d}\theta^{1}}\right)^{*}\nabla_{\bm{Z}^{2}}\mathcal{L}$	(A.2.80)
$\displaystyle\nabla_{\theta^{\mathrm{emb}}}\mathcal{L}$	$\displaystyle=\left(\frac{\mathrm{d}f^{\mathrm{emb}}}{\mathrm{d}\theta^{\mathrm{emb}}}\right)^{*}\nabla_{\bm{Z}^{1}}\mathcal{L}.$	(A.2.81)

This gives us a computationally efficient algorithm to find all gradients in the whole network.

We’ll finish this section by computing the adjoint derivative for a simple layer.

Example A.8.

Consider the “linear” (affine) layer $f^{\ell}$

f^{\ell}(\bm{Z},\bm{W}^{\ell},\bm{b}^{\ell})\doteq\bm{W}^{\ell}\bm{Z}+\bm{b}^{\ell}\bm{1}^{\top}=\begin{bmatrix}\bm{W}^{\ell}&\bm{b}^{\ell}\end{bmatrix}.

(A.2.82)

We can compute the differential w.r.t. both parameters as

	$\displaystyle\mathrm{d}f^{\ell}$	$\displaystyle=[(\bm{W}^{\ell}+\mathrm{d}\bm{W}^{\ell})\bm{Z}+(\bm{b}^{\ell}+\mathrm{d}\bm{b}^{\ell})\bm{1}^{\top}]-[\bm{W}^{\ell}\bm{Z}+\bm{b}^{\ell}]$		(A.2.83)
		$\displaystyle=(\mathrm{d}\bm{W}^{\ell})\bm{Z}+(\mathrm{d}\bm{b}^{\ell})\bm{1}^{\top}.$		(A.2.84)

Thus the derivative of this transformation is

\frac{\mathrm{d}f^{\ell}}{\mathrm{d}(\bm{W}^{\ell},\bm{b}^{\ell})}[\mathrm{d}\bm{W}^{\ell},\mathrm{d}\bm{b}^{\ell}]=(\mathrm{d}\bm{W}^{\ell})\bm{Z}+(\mathrm{d}\bm{b}^{\ell})\bm{1}^{\top},

(A.2.85)

namely, representing the following linear transformation from $\mathbb{R}^{m\times d}\times\mathbb{R}^{m}$ to $\mathbb{R}^{m\times n}$ :

T[\bm{A},\bm{u}]=\bm{A}\bm{Z}+\bm{u}\bm{1}^{\top}.

(A.2.86)

We calculate the adjoint $T^{*}\colon\mathbb{R}^{m\times n}\to\mathbb{R}^{m\times d}\times\mathbb{R}^{m}$ w.r.t. the sum-over-coordinates (Frobenius) inner product by the following procedure:

$\displaystyle\langle T[\bm{A},\bm{u}],\bm{B}\rangle_{\mathbb{R}^{m\times n}}$	$\displaystyle=\operatorname{tr}((\bm{A}\bm{Z}+\bm{u}\bm{1}^{\top})\bm{B}^{\top})$	(A.2.87)
	$\displaystyle=\operatorname{tr}(\bm{A}\bm{Z}\bm{B}^{\top}+\bm{u}\bm{1}^{\top}\bm{B}^{\top})$	(A.2.88)
	$\displaystyle=\operatorname{tr}(\bm{A}\bm{Z}\bm{B}^{\top})+\operatorname{tr}(\bm{u}\bm{1}^{\top}\bm{B}^{\top})$	(A.2.89)
	$\displaystyle=\operatorname{tr}(\bm{B}\bm{Z}^{\top}\bm{A}^{\top})+\operatorname{tr}(\bm{1}^{\top}\bm{B}^{\top}\bm{u})$	(A.2.90)
	$\displaystyle=\langle\bm{B}\bm{Z}^{\top},\bm{A}\rangle_{\mathbb{R}^{m\times d}}+\langle\bm{B}\bm{1},\bm{u}\rangle_{\mathbb{R}^{m}}$	(A.2.91)
	$\displaystyle=\langle\bm{B}(\bm{Z}^{\top},\bm{1}),(\bm{A},\bm{u})\rangle_{\mathbb{R}^{m\times d}\times\mathbb{R}^{m}}$	(A.2.92)

So $T^{*}\bm{B}=\bm{B}(\bm{Z}^{\top},\bm{1})$ . $\blacksquare$

Note that as a simple application of chain rule, both backpropagation and automatic differentiation work over general “computational graphs”, i.e., compositions of (simple) functions. We give all examples as neural network layers because this is the most common example in practice.

A.3 Game Theory and Minimax Optimization

In certain cases, such as in Chapter 5, a learning problem cannot be reduced to a single optimization problem but rather represents multiple potentially opposing components of the system try to each minimize their own objective. Examples of this paradigm include distribution learning via generative adversarial networks (GAN) and closed-loop transcription (CTRL). We will denote such a system as a two-player game, where we have two “players” (i.e., components) called Player 1 and Player 2 trying to minimize their objectives $\mathcal{L}^{1}$ and $\mathcal{L}^{2}$ respectively. Player 1 picks parameters $\theta\in\Theta$ and Player 2 picks parameters $\eta\in\mathrm{H}$ . In this book we consider the special case of zero-sum games, i.e., defining a common objective $\mathcal{L}$ such that $\mathcal{L}=-\mathcal{L}^{1}=\mathcal{L}^{2}$ .

Our first, very preliminary example is as follows. Suppose that there exists functions $u(\theta)$ and $v(\eta)$ such that

\mathcal{L}(\theta,\eta)=-u(\theta)+v(\eta).

(A.3.1)

Then both players’ objectives are independent of the other player, and the players should try to achieve their respective optima:

\theta^{\star}\in\operatorname*{arg\ min}_{\theta\in\Theta}u(\theta),\qquad\eta^{\star}\in\operatorname*{arg\ min}_{\eta\in\mathrm{H}}v(\eta).

(A.3.2)

The pair $(\theta^{\star},\eta^{\star})$ is a straightforward special case of an equilibrium: a situation where neither player will want to move, given the chance, since moving will end up making their own situation worse. However, not all games are so trivial; many have more complicated objectives and information structures.

In this book, the relevant game-theoretic formalism is a Stackelberg game (variously called sequential game). In this formalism, one player (without loss of generality Player 1, and also described as a leader) picks their parameters before the other (i.e., Player 2, also described as a follower), and the follower can use the full knowledge of the leader’s choice to make their own choice. The correct notion of equilibrium for a Stackelberg game is a Stackelberg equilibrium. To explain this equilibrium, note that since Player 2 (i.e., the follower) can choose $\eta$ reactively to the choice $\theta_{1}$ made by Player 1 (i.e., the leader), Player 2 would of course choose the $\eta$ which minimizes $\mathcal{L}(\theta_{1},\cdot)$ . But of course a rational Player 1 would realize this, and so pick a $\theta_{1}$ such that the worst-case $\eta$ picked by Player 2 according to this rule is not too bad. More formally, let $\mathcal{S}(\theta)\doteq\operatorname*{arg\ min}_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)$ be the set of $\eta$ minimizing $\mathcal{L}(\theta,\eta)$ , i.e., the set of all $\eta$ which Player 2 is liable to play given that Player 1 has played $\theta$ . Then $(\theta^{\star},\eta^{\star})$ is a Stackelberg equilibrium if

\theta^{\star}\in\operatorname*{arg\ max}_{\theta\in\Theta}\min_{\eta\in\mathcal{S}(\theta)}\mathcal{L}(\theta,\eta),\qquad\eta^{\star}\in\operatorname*{arg\ min}_{\eta\in\mathrm{H}}\mathcal{L}(\theta^{\star},\eta).

(A.3.3)

Actually (proof as exercise), one can show that in the context of two-player zero-sum Stackelberg games, $(\theta^{\star},\eta^{\star})$ is a Stackelberg equilibrium if and only if

\theta^{\star}\in\operatorname*{arg\ max}_{\theta\in\Theta}\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta),\qquad\eta^{\star}\in\operatorname*{arg\ min}_{\eta\in\mathrm{H}}\mathcal{L}(\theta^{\star},\eta),

(A.3.4)

(note that the notation $\mathcal{S}(\theta)$ is not used nor needed).

Proof.

Note that

\mathcal{S}(\theta)=\operatorname*{arg\ min}_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta),

(A.3.5)

so it holds

\min_{\eta\in\mathcal{S}(\theta)}\mathcal{L}(\theta,\eta)=\min_{\eta\in\operatorname*{arg\ min}_{\eta^{\prime}\in\mathrm{H}}\mathcal{L}(\theta,\eta^{\prime})}\mathcal{L}(\theta,\eta)=\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta).

(A.3.6)

∎

In the rest of the section, we will briefly discuss some algorithmic approaches to learn Stackelberg equilibria. The intuition you should have is that learning an equilibrium is like letting the different parts of the system automatically figure out tradeoffs between the different objectives they want to optimize.

We end this section with a caveat: in two-player zero-sum games, if it holds that

\max_{\theta\in\Theta}\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)=\min_{\eta\in\mathrm{H}}\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta)

(A.3.7)

then every Stackelberg equilibrium is a saddle point,⁷⁷7Famously called a Nash equilibrium. i.e.,

\theta^{\star}\in\operatorname*{arg\ max}_{\theta\in\Theta}\mathcal{L}(\theta,\eta^{\star}),\qquad\eta^{\star}\in\operatorname*{arg\ min}_{\eta\in\mathrm{H}}\mathcal{L}(\theta^{\star},\eta),

(A.3.8)

and vice versa, and furthermore each Stackelberg equilibrium has the (same) objective value

\max_{\theta\in\Theta}\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta).

Proof.

Suppose that indeed

\max_{\theta\in\Theta}\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)=\min_{\eta\in\mathrm{H}}\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta).

(A.3.9)

First suppose that $(\theta^{\star},\eta^{\star})$ is a saddle point. We will show it is a Stackelberg equilibrium. By definition we have

\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta^{\star},\eta)=\mathcal{L}(\theta^{\star},\eta^{\star})=\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta^{\star}).

(A.3.10)

It then holds for any $\theta$ and $\eta$ that

\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)\leq\mathcal{L}(\theta,\eta^{\star})\leq\mathcal{L}(\theta^{\star},\eta^{\star}).

(A.3.11)

Therefore

\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta)\leq\mathcal{L}(\theta^{\star},\eta^{\star}).

(A.3.12)

Completely symmetrically,

\mathcal{L}(\theta^{\star},\eta^{\star})\leq\mathcal{L}(\theta^{\star},\eta)\leq\max_{\theta\in\Theta}\mathcal{L}(\theta^{\star},\eta)\implies\mathcal{L}(\theta^{\star},\eta^{\star})\leq\min_{\eta\in\mathrm{H}}\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta).

(A.3.13)

Therefore since $\max\min=\min\max$ we have

\max_{\theta\in\Theta}\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)=\mathcal{L}(\theta^{\star},\eta^{\star})=\min_{\eta\in\mathrm{H}}\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta).

(A.3.14)

In particular, it holds that

\theta^{\star}\in\operatorname*{arg\ max}_{\theta\in\Theta}\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta).

(A.3.15)

From the saddle point condition we have $\eta^{\star}\in\operatorname*{arg\ min}_{\eta\in\mathrm{H}}\mathcal{L}(\theta^{\star},\eta)$ . So $(\theta^{\star},\eta^{\star})$ is a Stackelberg equilibrium. Furthermore we have also proved that all saddle points obey (A.3.14).

Now let $(\theta^{\star},\eta^{\star})$ be a Stackelberg equilibrium. We claim that it is a saddle point, which completes the proof. By the definition of minimax equilibrium,

\max_{\theta\in\Theta}\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)=\mathcal{L}(\theta^{\star},\eta^{\star}).

(A.3.16)

Then by the $\min\max=\max\min$ assumption we have

\max_{\theta\in\Theta}\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)=\mathcal{L}(\theta^{\star},\eta^{\star})=\min_{\eta\in\mathrm{H}}\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta).

(A.3.17)

This proves that all minimax equilibria have the desired objective value $\mathcal{L}(\theta^{\star},\eta^{\star})$ . Now we want to show that

\theta^{\star}\in\operatorname*{arg\ max}_{\theta\in\Theta}\mathcal{L}(\theta,\eta^{\star}),\qquad\eta^{\star}\in\operatorname*{arg\ min}_{\eta\in\mathrm{H}}\mathcal{L}(\theta^{\star},\eta).

(A.3.18)

Indeed the latter assertion holds by definition of the minimax equilibrium, so only the former need be proved. Namely, we will show that

\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta^{\star})=\mathcal{L}(\theta^{\star},\eta^{\star}).

(A.3.19)

To show this note that by definition of the $\max$

\mathcal{L}(\theta^{\star},\eta^{\star})\leq\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta^{\star}),

(A.3.20)

meanwhile we have $\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)\leq\mathcal{L}(\theta,\eta^{\star})$ so

\mathcal{L}(\theta^{\star},\eta^{\star})=\max_{\theta\in\Theta}\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)\leq\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta^{\star}).

(A.3.21)

Therefore it holds

\mathcal{L}(\theta^{\star},\eta^{\star})=\max_{\theta\in\Theta}\mathcal{L}(\theta,\eta^{\star}),

(A.3.22)

and the proof is complete. ∎

Conditions under which the $\min\max=\max\min$ equality holds are given by so-called minimax theorems; the most famous of these is a theorem of von Neumann. However, in the cases we think about, this property usually does not hold.

A.3.1 Learning Stackelberg Equilibria

How can we learn Stackelberg equilibria via GDA? In general this is clearly impossible, since learning Stackelberg equilibria via GDA is obviously at least as hard as computing a global minimizer of a loss function (say by setting the shared objective $\mathcal{L}(\theta,\eta)$ to only be a function of $\eta$ ). As such, we can achieve two types of convergence guarantees:

•

When $\mathcal{L}$ is (strongly) concave in the first argument $\theta$ and (strongly) convex in the second argument $\eta$ (as well as having Lipschitz gradients in both arguments), we can achieve exponentially fast convergence to a Stackelberg equilibrium.
•

When $\mathcal{L}$ is not concave or convex in either argument, we can achieve local convergence guarantees: namely, if we initialize the parameter values near a (local) Stackelberg equilibrium and the optimization geometry is good then we can learn that equilibrium efficiently.

The former situation is exactly analogous to the case of single-player optimization, where we proved that gradient descent converges exponentially fast for strongly convex objectives which have Lipschitz gradient. The latter situation is also analogous to the case of single-player optimization, although we did not cover it in depth due to technical difficulty; indeed there exist local convergence guarantees for nonconvex objectives which have locally nice geometry.

The algorithm in these two cases is the same algorithm, called Gradient Descent-Ascent (GDA). To motivate GDA, suppose we are trying to learn $\theta^{\star}$ . We could do gradient ascent on the function $\theta\mapsto\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)$ . But then we would need to take the derivative in $\theta$ of this function. To see how to do this, suppose that $\mathcal{L}$ is strongly convex in $\eta$ so that there is one minimizer $\eta^{\star}(\theta)$ of $\mathcal{L}(\theta,\cdot)$ . Then, Danskin’s theorem says that

\nabla_{\theta}\left[\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)\right]=\nabla_{\theta}\mathcal{L}(\theta,\texttt{stop\_grad}(\eta^{\star}(\theta))),

(A.3.23)

where the gradient is only with respect to the first argument (i.e., not a total derivative which would require computing the Jacobian of $\eta^{\star}(\theta)$ with respect to $\theta$ ), indicated by the stop-gradient operator.⁸⁸8In the case that $\mathcal{L}$ is not strongly convex in $\eta$ but rather just convex, Danskin’s theorem can be stated in terms of subdifferentials. If $\mathcal{L}$ is not convex at all in $\eta$ , then this derivative may not be well-defined, but one can obtain (local) convergence guarantees for the resulting algorithm anyways. Hence we use Danskin’s theorem as a motivation and not a justification for our algorithms. Danskin’s thoerems can be generalized into more relaxed circumstances by the so-called envelope theorems. In order to take the derivative in $\theta$ , we need to set up a secondary process to also optimize $\eta$ to obtain an approximation for $\eta^{\star}(\theta)$ . We can do this through the following algorithmic template:

	$\displaystyle\eta_{k+1}=\eta_{k+1}^{T+1};\qquad\eta_{k+1}^{t+1}=\eta_{k+1}^{t}-h\nabla_{\eta}\mathcal{L}(\theta_{k},\eta_{k+1}^{t}),\quad\forall t\in[T];\qquad\eta_{k+1}^{1}=\eta_{k}$		(A.3.24)
	$\displaystyle\theta_{k+1}=\theta_{k}+h\nabla_{\theta}\mathcal{L}(\theta_{k},\eta_{k+1}).$		(A.3.25)

That is, we take $T$ steps of gradient descent to update $\eta_{k}$ (hopefully to the minimizer of $\mathcal{L}(\theta_{k},\cdot)$ ), and then take a gradient ascent step to update $\theta_{k}$ . As a bonus, on top of estimating $\theta_{K}\approx\operatorname*{arg\ max}_{\theta}\min_{\eta}\mathcal{L}(\theta,\eta)$ , we also learn an $\eta_{K}\approx\operatorname*{arg\ min}_{\eta}\mathcal{L}(\theta_{K},\eta)$ — this is an approximate Stackelberg equilibrium.

This method is often not done in practice, as it requires $T+1$ total gradient descent iterations to update $\theta$ once. Instead, we use the so-called (simultaneous) Gradient Descent-Ascent (GDA) iteration

	$\displaystyle\theta_{k+1}$	$\displaystyle=\theta_{k}+h\nabla_{\theta}\mathcal{L}(\theta_{k},\eta_{k})$		(A.3.26)
	$\displaystyle\eta_{k+1}$	$\displaystyle=\eta_{k}-Th\nabla_{\eta}\mathcal{L}(\theta_{k},\eta_{k}),$		(A.3.27)

which can be implemented efficiently via a single gradient step on $(\theta,\eta)$ . The crucial idea here is, to make our method close to the inefficient iteration above, we use an $\eta$ update which is $T$ times faster than the $\theta$ update (these can be seen as nearly the same by taking a linearization of the dynamics).

It is crucial to pick $T$ sensibly. How can we do that? In the sequel, we discuss two configurations of $T$ which lead to convergence of GDA to a Stackelberg equilibrium under different assumptions.

Convergence of One-Timescale GDA to Stackelberg Equilibrium

If $T=1$ (i.e., named one-timescale because both $\theta$ and $\eta$ updates are of the same scale), then the GDA algorithm becomes

\theta_{k+1}=\theta_{k}+h\nabla_{\theta}\mathcal{L}(\theta_{k},\eta_{k}),\qquad\eta_{k+1}=\eta_{k}-h\nabla_{\eta}\mathcal{L}(\theta_{k},\eta_{k}).

(A.3.28)

If $\mathcal{L}$ has Lipschitz gradients in $\theta$ and $\eta$ , and is strongly concave in $\theta$ and strongly convex in $\eta$ , and is coercive (i.e., $\lim_{\|\theta\|_{2},\|\eta\|_{2}\to\infty}\mathcal{L}(\theta,\eta)=\infty$ ), then all saddle points are Stackelberg equilibria and vice versa. The work [ZAK24] shows that GDA (again, with $T=1$ ) with sufficiently small step size $h$ converges to a saddle point (hence Stackelberg equilibrium) exponentially fast if one exists, analogously to gradient descent for strongly convex functions with Lipschitz gradient.⁹⁹9We do not provide the step size or exponential base since they are complicated functions of the strong convexity/concavity/Lipschitz constant of the gradient. Of course, the paper [ZAK24] provides the precise parameter values. To our knowledge, this flavor of results constitute the only known rigorous justification for single-timescale GDA.

Local Convergence of Two-Timescale GDA to Stackelberg Equilibrium

Strong convexity/concavity is a global property, and none of the games we look into in this book have objectives which are globally strongly concave/strongly convex. In this case, the best we can hope for is local convergence to Stackelberg equilibria: if the parameters are initialized close to a Stackelberg equilibrium, then GDA can converge onto it, given an appropriate step size $h$ and timescale $T$ .

In fact, our results also hold for a version of the local Stackelberg equilibrium called the differential Stackelberg equilibrium, which was introduced in [FCR19] (though we use the precise definition in [LFD+22]), and which we define as follows. A point $(\theta^{\star},\eta^{\star})$ is a differential Stackelberg equilibrium if:

•

$\nabla_{\eta}\mathcal{L}(\theta^{\star},\eta^{\star})=\bm{0}$ ;
•

$\nabla_{\eta}^{2}\mathcal{L}(\theta^{\star},\eta^{\star})$ is symmetric positive definite;
•

$\nabla_{\theta}\mathcal{L}(\theta^{\star},\eta^{\star})=\bm{0}$ ;
•

$(\nabla_{\theta}^{2}\mathcal{L}+[\frac{\mathrm{d}}{\mathrm{d}\theta}\nabla_{\eta}\mathcal{L}][\nabla_{\eta}^{2}\mathcal{L}]^{-1}[\frac{\mathrm{d}}{\mathrm{d}\eta}\nabla_{\theta}\mathcal{L}])(\theta^{\star},\eta^{\star})$ is symmetric negative definite.

Notice that the last condition asks for the (total) Hessian

\nabla^{2}\mathcal{L}(\theta,\eta)=\begin{bmatrix}\nabla_{\theta}^{2}\mathcal{L}(\theta,\eta)&\frac{\mathrm{d}}{\mathrm{d}\eta}\nabla_{\theta}^{2}\mathcal{L}(\theta,\eta)\\ \frac{\mathrm{d}}{\mathrm{d}\theta}\nabla_{\eta\theta}^{2}\mathcal{L}(\theta,\eta)&\nabla_{\eta}^{2}\mathcal{L}(\theta)\end{bmatrix}

(A.3.29)

or, equivalently, its Schur complement to be negative definite. If we look at the computation of $\nabla_{\theta}[\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)]$ furnished by Danskin’s theorem, the last two criteria are actually constraints on the gradient and Hessian of the function $\theta\mapsto\min_{\eta\in\mathrm{H}}\mathcal{L}(\theta,\eta)$ , ensuring that the gradient is $0$ and the Hessian is negative semidefinite. This intuition tells us that we can expect that each Stackelberg equilibrium is a differential Stackelberg equilibrium; [FCR20] confirms this rigorously (up to some technical conditions).

Analogously to the notion of strict local optimum in single-player optimization (where we require $\nabla^{2}\mathcal{L}(\theta^{\star})$ to be positive semidefinite), the definition of differential Stackelberg equilibrium implies that $\mathcal{L}(\theta,\cdot)$ is locally (strictly) convex in a neighborhood of the equilibrium, and that $\min_{\eta\in\mathrm{H}}\mathcal{L}(\cdot,\eta)$ is locally (strictly) concave in the same region.

In this context, we present the result from [LFD+22]. Let $(\theta^{\star},\eta^{\star})$ be a differential Stackelberg equilibrium. Suppose that $\mathcal{L}$ has Lipschitz gradients, i.e.,

\max\left\{\|\nabla_{\theta}^{2}\mathcal{L}(\theta^{\star},\eta^{\star})\|_{2},\left\|\nabla_{\eta}^{2}\mathcal{L}(\theta^{\star},\eta^{\star})\right\|_{2},\left\|\frac{\mathrm{d}}{\mathrm{d}\eta}\nabla_{\theta}\mathcal{L}(\theta^{\star},\eta^{\star})\right\|_{2}\right\}\leq\beta.

(A.3.30)

Further define the local strong convexity/concavity parameters of $\mathcal{L}(\theta,\cdot)$ and $\min_{\eta}\mathcal{L}(\cdot,\eta)$ respectively as

\mu_{\eta}=\lambda_{\min}(\nabla_{\eta}^{2}\mathcal{L}(\theta^{\star},\eta^{\star})),\qquad\mu_{\theta}=\min\left\{\beta,-\left(\nabla_{\theta}^{2}\mathcal{L}+\left[\frac{\mathrm{d}}{\mathrm{d}\theta}\nabla_{\eta}\mathcal{L}\right]\left[\nabla_{\eta}^{2}\mathcal{L}\right]^{-1}\left[\frac{\mathrm{d}}{\mathrm{d}\eta}\nabla_{\theta}\mathcal{L}\right]\right)(\theta^{\star},\eta^{\star})\right\}.

(A.3.31)

Then define the local condition numbers as

\kappa_{\eta}=L/\mu_{\eta},\qquad\kappa_{\theta}=L/\mu_{\theta}.

(A.3.32)

The paper [LFD+22] says that if we take step size $h=\frac{1}{4\beta T}$ and take $T\geq 2\kappa$ , so that the algorithm is

	$\displaystyle\theta_{k+1}$	$\displaystyle=\theta_{k}+\frac{1}{4\beta T}\nabla_{\theta}\mathcal{L}(\theta_{k},\eta_{k}),$		(A.3.33)
	$\displaystyle\eta_{k+1}$	$\displaystyle=\eta_{k}-\frac{1}{4\beta}\nabla_{\eta}\mathcal{L}(\theta_{k},\eta_{k}),$		(A.3.34)

the total Hessian $\nabla^{2}\mathcal{L}(\theta^{\star},\eta^{\star})$ is diagonalizable, and we initialize $(\theta_{0},\eta_{0})$ close enough to $(\theta^{\star},\eta^{\star})$ , then there are positive constants $c_{0},c_{1}>0$ such that

\|(\theta_{k},\eta_{k})-(\theta^{\star},\eta^{\star})\|_{2}\leq c_{0}\left(1-\frac{c_{1}}{T\kappa_{\theta}}\right)^{k}\|(\theta_{0},\eta_{0})-(\theta^{\star},\eta^{\star})\|_{2},

(A.3.35)

implying exponential convergence to the differential Stackelberg equilibrium.

A.3.2 Practical Considerations when Learning Stackelberg Equilibria

In practice, we do not know how to initialize parameters close to a (differential) Stackelberg equilibrium. Due to symmetries within the objective, including those induced by overparameterization of the neural networks being trained, one can (heuristically) expect that most initializations are close to a Stackelberg equilibrium. Also, we do not know how to compute the step size $h$ or the timescale $T$ , since they are dependent on properties of the loss $\mathcal{L}$ at the equilibrium. In practice, there are some common approaches:

•

Take $T=1$ (equal step-sizes), and use updates for $\theta$ and $\eta$ that are derived from a learning-rate-adaptive optimizer like Adam (as opposed to vanilla GD). Here, you hope (but do not know) that the optimizer can adjust the learning rates to learn a good equilibrium.
•

Take $T$ to be some constant like $T=10^{6}$ which implies that $\eta$ equilibrates $10^{6}$ times as fast as $\theta$ . Here you can also use Adam-style updates, and hope that it fixes the time scale.
•

Let $T$ depend on the iteration $k$ , and let $T_{k}\to\infty$ as $k\to\infty$ . This schedule was studied (also in the case of noise) by Borkar in [Bor97].

For example, you can use this while training CTRL-style models (see Chapter 5), where the encoder is Player $1$ and the decoder is Player $2$ . Some theory about CTRL is given in Theorem 5.1.

A.4 Exercises

Exercise A.1.

We have shown that for a smooth function $f$ , gradient descent converges linearly to the global optimum if it is strongly convex. However, in general nonconvex optimization, we do not have convexity, let alone strong convexity. Fortunately, in some cases, $f$ satisfies the so-called $\mu$ -Polyak-Lojasiewicz (PL) inequality, i.e., there exists a constant $\mu>0$ such that for all $\theta$ ,

\displaystyle\frac{1}{2}\|\nabla f(\theta)\|_{2}^{2}\geq\mu\left(f(\theta)-f(\theta^{\star})\right),

where $\theta^{*}$ is a minimizer of $f$ .

Please show that under the PL inequality and the assumption that $f$ is $\beta$ -smooth, gradient descent (A.1.12) converges linearly to $\theta^{*}$ .

Exercise A.2.

Compute the differential and adjoint derivative of the softmax function, defined as follows.

\operatorname{\mathrm{softmax}}\left(\begin{bmatrix}x_{1}\\ \vdots\\ x_{n}\end{bmatrix}\right)=\frac{1}{\sum_{i=1}^{n}e^{x_{i}}}\begin{bmatrix}x_{1}\\ \vdots\\ x_{n}\end{bmatrix}.

(A.4.1)

Exercise A.3.

Carry through the backpropagation computation for a $L$ -layer MLP, as defined in Section 7.2.3.

Exercise A.4.

Carry through the backpropagation computation for a $L$ -layer transformer, as defined in Section 7.2.3.

Exercise A.5.

Carry through the backpropagation computation for an autoencoder with $L$ encoder layers and $L$ decoder layers (without necessarily specifying an architecture).

$\displaystyle\\|\theta^{\star}-\theta_{k+1}\\|_{2}^{2}$	$\displaystyle\leq\\|\theta^{\star}-\theta_{k}+h\nabla\mathcal{L}(\theta_{k})\\|_{2}^{2}$	(A.1.13)
	$\displaystyle=\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}+2h\langle\nabla\mathcal{L}(\theta_{k}),\theta^{\star}-\theta_{k}\rangle+h^{2}\\|\nabla\mathcal{L}(\theta_{k})\\|_{2}^{2}$	(A.1.14)
	$\displaystyle\leq\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}+2h\left(\mathcal{L}(\theta^{\star})-\mathcal{L}(\theta_{k})-\frac{\alpha}{2}\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}\right)+h^{2}\\|\nabla\mathcal{L}(\theta_{k})\\|_{2}^{2}\quad\text{($\alpha$-SC)}$	(A.1.15)
	$\displaystyle=\left(1-\alpha h\right)\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}+2h(\mathcal{L}(\theta^{\star})-L(\theta_{k}))+h^{2}\\|\nabla\mathcal{L}(\theta_{k})\\|_{2}^{2}$	(A.1.16)
	$\displaystyle\leq\left(1-\alpha h\right)\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}+2h(\mathcal{L}(\theta^{k})-\mathcal{L}(\theta^{\star}))+2h^{2}\beta(\mathcal{L}(\theta_{k})-\mathcal{L}(\theta^{\star}))\quad\text{($\beta$-LG)}$	(A.1.17)
	$\displaystyle=\left(1-\alpha h\right)\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}-2h(1-\beta h)(\mathcal{L}(\theta_{k})-\mathcal{L}(\theta^{\star})).$	(A.1.18)

	$\displaystyle\\|\theta^{\star}-\theta_{k+1}\\|_{2}^{2}$	$\displaystyle\leq(1-\alpha h)\\|\theta^{\star}-\theta_{k}\\|_{2}^{2}\leq(1-\alpha h)^{2}\\|\theta^{\star}-\theta_{k-1}\\|_{2}^{2}\leq\cdots$		(A.1.19)
		$\displaystyle\leq(1-\alpha h)^{k+1}\\|\theta^{\star}-\theta_{0}\\|_{2}^{2}.$		(A.1.20)