% ============================================================================
%  Artificial Neural Networks — Methods, Equations and Graphical
%  Representations
%  Author  : Simon-Pierre Boucher — contact@spboucher.ai
%  Chapter 2 : Training Neural Networks (chapters/02-training.tex)
% ============================================================================
\chapter{Training Neural Networks}\label{chap:training}

The previous chapter established what a feed-forward network computes; this
chapter establishes how its parameters are found. Training is cast as the
minimization of an empirical risk, and virtually every modern network is
trained by the same recipe: compute the gradient of the loss with respect to
every parameter by \emph{backpropagation}~\cite{rumelhart1986}, then descend
along that gradient with a first-order update rule. We derive the
backpropagation equations in full, present the algorithm in pseudocode, and
then survey the family of optimizers — from plain stochastic gradient
descent to Adam~\cite{kingma2015adam} — together with the learning-rate
schedules and weight-initialization schemes that make deep training stable
in practice.

% ============================================================================
\section{The Learning Problem}

\begin{definition}[Empirical risk minimization]\label{def:train-erm}
Let $f(\cdot\,;\vect{\theta})$ be a network with parameters $\vect{\theta}$,
let $\Loss(\hat{\vect{y}},\vect{y})$ be a per-example loss, and let
$\mathcal{D}$ be the data distribution. The \emph{(true) risk} and its
Monte-Carlo estimate on a training set
$\{(\vect{x}_i,\vect{y}_i)\}_{i=1}^{n}$, the \emph{empirical risk}, are
\begin{equation}\label{eq:train-risk}
  R(\vect{\theta})
    = \E_{(\vect{x},\vect{y})\sim\mathcal{D}}
      \bigl[\Loss\bigl(f(\vect{x};\vect{\theta}),\vect{y}\bigr)\bigr],
  \qquad
  \widehat{R}(\vect{\theta})
    = \frac{1}{n}\sum_{i=1}^{n}
      \Loss\bigl(f(\vect{x}_i;\vect{\theta}),\vect{y}_i\bigr).
\end{equation}
Training solves
$\vect{\theta}^{\star} \in \argmin_{\vect{\theta}} \widehat{R}(\vect{\theta})$.
\end{definition}

The choice of $\Loss$ encodes the task. We review the two workhorses —
squared error for regression, cross-entropy for classification — and the
single most useful gradient identity in deep learning.

\subsection{Regression: mean squared error}

For scalar targets $y_i \in \R$ and predictions $\hat{y}_i$, the
\emph{mean squared error} and its gradient are
\begin{equation}\label{eq:train-mse}
  \Loss_{\mathrm{MSE}}
    = \frac{1}{n}\sum_{i=1}^{n}\bigl(y_i - \hat{y}_i\bigr)^{2},
  \qquad
  \frac{\partial \Loss_{\mathrm{MSE}}}{\partial \hat{y}_i}
    = -\frac{2}{n}\,\bigl(y_i - \hat{y}_i\bigr).
\end{equation}
Minimizing~\eqref{eq:train-mse} is maximum-likelihood estimation under
Gaussian observation noise; the quadratic growth makes it sensitive to
outliers, which motivates robust alternatives such as the mean absolute
error $\frac{1}{n}\sum_i |y_i-\hat y_i|$ (Laplace likelihood, whose
minimizer is the conditional median) and the Huber loss, quadratic for
small residuals and linear in the tails.

\subsection{Classification: cross-entropy}

For binary labels $y\in\{0,1\}$ with $\hat{y}=\sigma(z)$ produced by a
sigmoid over the logit $z$, the \emph{binary cross-entropy} is
\begin{equation}\label{eq:train-bce}
  \Loss_{\mathrm{BCE}}
    = -\frac{1}{n}\sum_{i=1}^{n}
      \Bigl[\,y_i\ln\hat{y}_i + (1-y_i)\ln\bigl(1-\hat{y}_i\bigr)\Bigr],
  \qquad
  \frac{\partial \Loss_{\mathrm{BCE}}}{\partial z} = \hat{y}-y .
\end{equation}
For $K$-way classification with one-hot target $\vect{y}$ and
$\hat{\vect{y}}=\softmax(\vect{z})$, the \emph{categorical cross-entropy}
is
\begin{equation}\label{eq:train-ce}
  \Loss_{\mathrm{CE}}
    = -\sum_{k=1}^{K} y_k \ln \hat{y}_k ,
\end{equation}
which equals, up to the (constant) entropy of $\vect{y}$, the
Kullback--Leibler divergence $\KL(\vect{y}\,\|\,\hat{\vect{y}})$.

\begin{property}[Softmax--cross-entropy gradient]\label{prop:train-ce-grad}
With $\hat{\vect{y}}=\softmax(\vect{z})$ and $\Loss_{\mathrm{CE}}$ as
in~\eqref{eq:train-ce},
\begin{equation}\label{eq:train-ce-grad}
  \frac{\partial \Loss_{\mathrm{CE}}}{\partial \vect{z}}
    \;=\; \hat{\vect{y}} - \vect{y}.
\end{equation}
\end{property}

\begin{proof}[Sketch]
The softmax Jacobian is
$\partial\hat{y}_i/\partial z_j = \hat{y}_i(\delta_{ij}-\hat{y}_j)$ with
$\delta_{ij}$ the Kronecker delta. Chaining it against
$\partial\Loss/\partial\hat{y}_i = -y_i/\hat{y}_i$ gives
$\partial\Loss/\partial z_j
 = \sum_i (-y_i/\hat{y}_i)\,\hat{y}_i(\delta_{ij}-\hat{y}_j)
 = -y_j + \hat{y}_j\sum_i y_i = \hat{y}_j - y_j$,
since $\sum_i y_i = 1$.
\end{proof}

\begin{remark}
The cancellation in~\eqref{eq:train-ce-grad} is not a coincidence: it holds
for every matched pair of a canonical link and its exponential-family
negative log-likelihood (sigmoid with binary cross-entropy
in~\eqref{eq:train-bce}, identity with MSE). The practical consequence is
the absence of saturating $\sigma'$ factors at the output layer: a
confidently wrong prediction still receives a large gradient.
\end{remark}

% ============================================================================
\section{Backpropagation}

Consider the $L$-layer network of Chapter~\ref{chap:foundations},
\begin{equation}\label{eq:train-forward}
  \vect{z}^{(\ell)} = \mat{W}^{(\ell)}\vect{a}^{(\ell-1)}+\vect{b}^{(\ell)},
  \qquad
  \vect{a}^{(\ell)} = \varphi\bigl(\vect{z}^{(\ell)}\bigr),
  \qquad \ell = 1,\dots,L,
\end{equation}
with $\vect{a}^{(0)}=\vect{x}$ and a scalar loss
$\Loss\bigl(\vect{a}^{(L)},\vect{y}\bigr)$. Backpropagation is reverse-mode
automatic differentiation applied to this composition: it computes
$\partial\Loss/\partial\vect{\theta}$ for \emph{all} parameters at the cost
of a constant number of forward passes. The central object is the
\emph{error signal} of layer $\ell$,
\begin{equation}\label{eq:train-delta-def}
  \vect{\delta}^{(\ell)}
    \;\equiv\; \frac{\partial \Loss}{\partial \vect{z}^{(\ell)}}
    \;\in\; \R^{n_\ell}.
\end{equation}

\begin{theorem}[Backpropagation equations~\cite{rumelhart1986}]
\label{thm:train-backprop}
For the network~\eqref{eq:train-forward}, the error signals and parameter
gradients satisfy
\begin{align}
  \vect{\delta}^{(L)}
    &= \nabla_{\vect{a}^{(L)}}\Loss
       \,\odot\, \varphi'\bigl(\vect{z}^{(L)}\bigr),
  \label{eq:train-bp-out}\\[2pt]
  \vect{\delta}^{(\ell)}
    &= \Bigl(\mat{W}^{(\ell+1)\transp}\,\vect{\delta}^{(\ell+1)}\Bigr)
       \odot \varphi'\bigl(\vect{z}^{(\ell)}\bigr),
  \quad \ell = L-1,\dots,1,
  \label{eq:train-bp-rec}\\[2pt]
  \frac{\partial \Loss}{\partial \mat{W}^{(\ell)}}
    &= \vect{\delta}^{(\ell)}\,\vect{a}^{(\ell-1)\transp},
  \qquad
  \frac{\partial \Loss}{\partial \vect{b}^{(\ell)}}
    = \vect{\delta}^{(\ell)} .
  \label{eq:train-bp-w}
\end{align}
\end{theorem}

\begin{proof}[Sketch]
Equation~\eqref{eq:train-bp-out} is the chain rule through
$\vect{a}^{(L)}=\varphi(\vect{z}^{(L)})$. For~\eqref{eq:train-bp-rec},
each $z^{(\ell+1)}_k$ depends on $z^{(\ell)}_j$ through
$W^{(\ell+1)}_{kj}\varphi'(z^{(\ell)}_j)$, so summing over downstream
paths,
\[
\delta^{(\ell)}_j
  = \sum_k \delta^{(\ell+1)}_k W^{(\ell+1)}_{kj}\,
    \varphi'\bigl(z^{(\ell)}_j\bigr),
\]
which is~\eqref{eq:train-bp-rec} in matrix form.
For~\eqref{eq:train-bp-w}, note
$\partial z^{(\ell)}_j / \partial W^{(\ell)}_{ji} = a^{(\ell-1)}_i$ and
$\partial z^{(\ell)}_j / \partial b^{(\ell)}_j = 1$.
\end{proof}

The name of the algorithm is visible in~\eqref{eq:train-bp-rec}: the error
is propagated \emph{backwards} through the transposes of the forward weight
matrices. With softmax output and cross-entropy loss,
Property~\ref{prop:train-ce-grad} replaces~\eqref{eq:train-bp-out} directly
by $\vect{\delta}^{(L)} = \hat{\vect{y}}-\vect{y}$.
Figure~\ref{fig:train-compgraph} shows the flow of both passes on the
computational graph, and Algorithm~\ref{alg:train-backprop} states the full
procedure for one mini-batch.

\begin{figure}[htbp]
  \centering
  \begin{tikzpicture}
    % ---- forward row -------------------------------------------------
    \node[blocinput,  minimum width=1.15cm] (x)  at (0,0)    {$\vect{x}$};
    \node[blochidden, minimum width=1.15cm] (z1) at (2.35,0) {$\vect{z}^{(1)}$};
    \node[blochidden, minimum width=1.15cm] (a1) at (4.70,0) {$\vect{a}^{(1)}$};
    \node[blochidden, minimum width=1.15cm] (z2) at (7.05,0) {$\vect{z}^{(2)}$};
    \node[blocoutput, minimum width=1.15cm] (a2) at (9.40,0) {$\hat{\vect{y}}$};
    \node[bloc,       minimum width=1.15cm] (Ls) at (11.75,0){$\Loss$};
    % ---- parameter and target nodes ------------------------------------
    \node[mem, minimum width=1.5cm] (W1) at (2.35,1.9)
      {$\mat{W}^{(1)},\vect{b}^{(1)}$};
    \node[mem, minimum width=1.5cm] (W2) at (7.05,1.9)
      {$\mat{W}^{(2)},\vect{b}^{(2)}$};
    \node[blocinput, minimum width=1.15cm] (y) at (11.75,1.9) {$\vect{y}$};
    % ---- forward arrows ------------------------------------------------
    \draw[fleche] (x)  -- (z1);
    \draw[fleche] (z1) -- node[above, etiquette] {$\varphi$} (a1);
    \draw[fleche] (a1) -- (z2);
    \draw[fleche] (z2) -- node[above, etiquette] {$\varphi$} (a2);
    \draw[fleche] (a2) -- (Ls);
    \draw[fleche] (W1) -- (z1);
    \draw[fleche] (W2) -- (z2);
    \draw[fleche] (y)  -- (Ls);
    % ---- backward (dashed) arrows, routed below ------------------------
    \draw[flechep, coutput] (Ls.south) to[bend left=35]
      node[below, etiquette] {$\nabla_{\hat{\vect{y}}}\Loss$} (a2.south);
    \draw[flechep, coutput] (a2.south) to[bend left=35]
      node[below, etiquette] {$\vect{\delta}^{(2)}$} (z2.south);
    \draw[flechep, coutput] (z2.south) to[bend left=35]
      node[below, etiquette] {$\mat{W}^{(2)\transp}\vect{\delta}^{(2)}$}
      (a1.south);
    \draw[flechep, coutput] (a1.south) to[bend left=35]
      node[below, etiquette] {$\vect{\delta}^{(1)}$} (z1.south);
    % ---- gradients to parameters ---------------------------------------
    \draw[flechep, coutput] (z1.north east) to[bend right=30]
      node[right, etiquette, xshift=2pt]
      {$\vect{\delta}^{(1)}\vect{a}^{(0)\transp}$} (W1.east);
    \draw[flechep, coutput] (z2.north east) to[bend right=30]
      node[right, etiquette, xshift=2pt]
      {$\vect{\delta}^{(2)}\vect{a}^{(1)\transp}$} (W2.east);
  \end{tikzpicture}
  \caption{Computational graph of a two-layer network. Solid arrows:
    forward pass~\eqref{eq:train-forward}. Dashed red arrows: backward
    pass — the error signals $\vect{\delta}^{(\ell)}$ of
    Theorem~\ref{thm:train-backprop} flow from the loss back through the
    layers, branching off to the parameter
    gradients~\eqref{eq:train-bp-w}.}
  \label{fig:train-compgraph}
\end{figure}

\begin{algorithm}[htbp]
  \caption{Backpropagation with mini-batch gradient descent}
  \label{alg:train-backprop}
  \begin{algorithmic}[1]
    \Require mini-batch $\{(\vect{x}_i,\vect{y}_i)\}_{i=1}^{m}$, parameters
      $\{\mat{W}^{(\ell)},\vect{b}^{(\ell)}\}_{\ell=1}^{L}$, learning rate
      $\eta$
    \For{$i = 1$ \textbf{to} $m$} \Comment{forward pass}
      \State $\vect{a}^{(0)} \gets \vect{x}_i$
      \For{$\ell = 1$ \textbf{to} $L$}
        \State $\vect{z}^{(\ell)} \gets
          \mat{W}^{(\ell)}\vect{a}^{(\ell-1)} + \vect{b}^{(\ell)}$;
          \quad $\vect{a}^{(\ell)} \gets \varphi(\vect{z}^{(\ell)})$
      \EndFor
      \State $\vect{\delta}^{(L)} \gets
        \nabla_{\vect{a}^{(L)}}\Loss \odot \varphi'(\vect{z}^{(L)})$
        \Comment{backward pass; $\hat{\vect{y}}-\vect{y}$ for softmax+CE}
      \For{$\ell = L-1$ \textbf{down to} $1$}
        \State $\vect{\delta}^{(\ell)} \gets
          \bigl(\mat{W}^{(\ell+1)\transp}\vect{\delta}^{(\ell+1)}\bigr)
          \odot \varphi'(\vect{z}^{(\ell)})$
      \EndFor
      \State accumulate
        $\Delta\mat{W}^{(\ell)} \mathrel{+}=
           \vect{\delta}^{(\ell)}\vect{a}^{(\ell-1)\transp}$,
        \; $\Delta\vect{b}^{(\ell)} \mathrel{+}= \vect{\delta}^{(\ell)}$
        \textbf{for all} $\ell$
    \EndFor
    \For{$\ell = 1$ \textbf{to} $L$} \Comment{gradient step}
      \State $\mat{W}^{(\ell)} \gets
        \mat{W}^{(\ell)} - \dfrac{\eta}{m}\,\Delta\mat{W}^{(\ell)}$;
        \quad
        $\vect{b}^{(\ell)} \gets
        \vect{b}^{(\ell)} - \dfrac{\eta}{m}\,\Delta\vect{b}^{(\ell)}$
    \EndFor
  \end{algorithmic}
\end{algorithm}

\begin{remark}[Cost]
One backward pass costs the same order as one forward pass,
$O\bigl(\sum_\ell n_\ell n_{\ell-1}\bigr)$: the full gradient of a scalar
with respect to $P$ parameters is obtained for $O(1)$ — not $O(P)$ —
forward-pass equivalents. This efficiency of reverse-mode differentiation
is what makes deep learning computationally feasible.
\end{remark}

\begin{remark}[Vanishing and exploding gradients]\label{rem:train-vanish}
The recursion~\eqref{eq:train-bp-rec} multiplies a $\varphi'$ factor and a
weight matrix at every layer. With sigmoid activations,
$\sigma'(z)\le 1/4$, so error signals shrink at least geometrically with
depth; with large weights they can instead grow without bound. This
vanishing/exploding behaviour motivates ReLU-family activations, careful
initialization (Section~\ref{sec:train-init}), normalization layers and
residual connections, treated in later chapters.
\end{remark}

% ============================================================================
\section{First-Order Optimizers}

Throughout this section $\vect{\theta}_t$ denotes the parameters at step
$t$, $\vect{g}_t = \nabla_{\vect{\theta}}\Loss(\vect{\theta}_t)$ the
mini-batch gradient, $\eta$ the learning rate, and all operations on
vectors are elementwise.

\subsection{Stochastic gradient descent and momentum}

\emph{Stochastic gradient descent} (SGD) applies the elementary update
\begin{equation}\label{eq:train-sgd}
  \vect{\theta}_{t+1} = \vect{\theta}_t - \eta\,\vect{g}_t .
\end{equation}
Classical stochastic-approximation theory guarantees convergence when the
step sizes satisfy $\sum_t \eta_t = \infty$ and $\sum_t \eta_t^2 < \infty$
\cite{goodfellow2016book}.

\emph{Momentum} (Polyak's heavy ball) accumulates an exponentially weighted
velocity,
\begin{equation}\label{eq:train-momentum}
  \vect{v}_t = \beta\,\vect{v}_{t-1} + \vect{g}_t ,
  \qquad
  \vect{\theta}_{t+1} = \vect{\theta}_t - \eta\,\vect{v}_t ,
\end{equation}
with $\beta \approx 0.9$. Directions in which successive gradients agree
are amplified by up to $1/(1-\beta)$, while oscillating components cancel —
precisely the geometry of Figure~\ref{fig:train-trajectories}.
\emph{Nesterov's accelerated gradient} evaluates the gradient at a
look-ahead point,
\begin{equation}\label{eq:train-nesterov}
  \vect{v}_t = \beta\,\vect{v}_{t-1}
    + \nabla_{\vect{\theta}}
      \Loss\bigl(\vect{\theta}_t - \eta\beta\,\vect{v}_{t-1}\bigr),
  \qquad
  \vect{\theta}_{t+1} = \vect{\theta}_t - \eta\,\vect{v}_t ,
\end{equation}
letting the update ``see'' where it is heading and correct in advance; for
smooth convex objectives it attains the optimal $O(1/t^2)$ convergence
rate.

\begin{figure}[htbp]
  \centering
  \begin{tikzpicture}
    \begin{axis}[
        width=0.88\textwidth, height=6.2cm,
        xmin=-10, xmax=1.5, ymin=-2.4, ymax=2.4,
        xlabel={$\theta_1$}, ylabel={$\theta_2$},
        xlabel near ticks, ylabel near ticks,
        tick label style={font=\scriptsize},
        label style={font=\small},
        legend style={font=\scriptsize, at={(0.98,0.04)},
                      anchor=south east, draw=black!30},
        legend cell align=left,
      ]
      % loss contours of f = x^2/20 + y^2 (anisotropic bowl)
      \addplot[domain=0:360, samples=91, smooth, black!25, forget plot]
        ({2.828*cos(x)}, {0.632*sin(x)});
      \addplot[domain=0:360, samples=91, smooth, black!25, forget plot]
        ({4.899*cos(x)}, {1.095*sin(x)});
      \addplot[domain=0:360, samples=91, smooth, black!25, forget plot]
        ({6.928*cos(x)}, {1.549*sin(x)});
      \addplot[domain=0:360, samples=91, smooth, black!25, forget plot]
        ({8.944*cos(x)}, {2.000*sin(x)});
      \addplot[domain=0:360, samples=91, smooth, black!25, forget plot]
        ({11.662*cos(x)}, {2.608*sin(x)});
      % SGD trajectory (zigzag)
      \addplot[coutput, thick, mark=*, mark size=1.1pt] coordinates {
        (-8.500,1.800) (-7.777,-1.260) (-7.116,0.882) (-6.512,-0.617)
        (-5.958,0.432) (-5.452,-0.303) (-4.988,0.212) (-4.564,-0.148)
        (-4.176,0.104) (-3.821,-0.073) (-3.496,0.051) (-3.199,-0.036)
        (-2.927,0.025) (-2.679,-0.017) (-2.451,0.012) (-2.243,-0.009)
        (-2.052,0.006)
      };
      \addlegendentry{SGD~\eqref{eq:train-sgd}}
      % Momentum trajectory (smooth)
      \addplot[cinput, thick, mark=*, mark size=1.1pt] coordinates {
        (-8.500,1.800) (-8.381,1.296) (-8.157,0.480) (-7.840,-0.390)
        (-7.446,-1.063) (-6.987,-1.371) (-6.476,-1.264) (-5.925,-0.815)
        (-5.347,-0.182) (-4.751,0.439) (-4.149,0.874) (-3.548,1.022)
        (-2.959,0.868) (-2.386,0.487) (-1.838,0.007) (-1.318,-0.426)
        (-0.832,-0.697)
      };
      \addlegendentry{Momentum~\eqref{eq:train-momentum}}
      % optimum
      \addplot[only marks, mark=star, mark size=3.2pt, black]
        coordinates {(0,0)};
      \addlegendentry{minimum $\vect{\theta}^{\star}$}
    \end{axis}
  \end{tikzpicture}
  \caption{Sixteen steps of SGD and momentum on the anisotropic quadratic
    $\Loss(\vect{\theta}) = \theta_1^2/20 + \theta_2^2$ (grey level sets),
    both computed numerically from the same starting point. SGD (red)
    oscillates across the narrow valley while creeping along the shallow
    direction; momentum (blue) damps the oscillation and accelerates along
    the valley floor.}
  \label{fig:train-trajectories}
\end{figure}

\subsection{Adaptive methods: AdaGrad, RMSProp, Adam}

Adaptive methods give every coordinate its own effective learning rate,
scaled by the history of gradient magnitudes. \emph{AdaGrad} accumulates
the squared gradients,
\begin{equation}\label{eq:train-adagrad}
  \vect{G}_t = \vect{G}_{t-1} + \vect{g}_t^{2},
  \qquad
  \vect{\theta}_{t+1}
    = \vect{\theta}_t
      - \frac{\eta}{\sqrt{\vect{G}_t}+\epsilon}\odot\vect{g}_t ,
\end{equation}
so rarely active (sparse) coordinates receive large steps. Because
$\vect{G}_t$ grows monotonically, however, the effective step size decays
to zero. \emph{RMSProp} repairs this by replacing the sum with an
exponential moving average,
\begin{equation}\label{eq:train-rmsprop}
  \E[\vect{g}^2]_t
    = \rho\,\E[\vect{g}^2]_{t-1} + (1-\rho)\,\vect{g}_t^{2},
  \qquad
  \vect{\theta}_{t+1}
    = \vect{\theta}_t
      - \frac{\eta}{\sqrt{\E[\vect{g}^2]_t}+\epsilon}\odot\vect{g}_t ,
\end{equation}
with $\rho\approx 0.9$, so that stale gradients are forgotten.

\emph{Adam}~\cite{kingma2015adam} (\emph{adaptive moment estimation})
combines the momentum idea~\eqref{eq:train-momentum} — an exponential
moving average of the gradient, the first moment — with the RMSProp
idea~\eqref{eq:train-rmsprop} — an exponential moving average of its
square, the second moment — and corrects the initialization bias of both:
\begin{align}
  \vect{m}_t &= \beta_1\,\vect{m}_{t-1} + (1-\beta_1)\,\vect{g}_t ,
  \qquad
  \vect{v}_t = \beta_2\,\vect{v}_{t-1} + (1-\beta_2)\,\vect{g}_t^{2},
  \label{eq:train-adam-moments}\\[2pt]
  \hat{\vect{m}}_t &= \frac{\vect{m}_t}{1-\beta_1^{\,t}} ,
  \qquad
  \hat{\vect{v}}_t = \frac{\vect{v}_t}{1-\beta_2^{\,t}} ,
  \label{eq:train-adam-bias}\\[2pt]
  \vect{\theta}_{t+1}
    &= \vect{\theta}_t
       - \eta\,\frac{\hat{\vect{m}}_t}{\sqrt{\hat{\vect{v}}_t}+\epsilon} .
  \label{eq:train-adam-update}
\end{align}
The standard defaults are $\eta = 10^{-3}$, $\beta_1 = 0.9$,
$\beta_2 = 0.999$ and $\epsilon = 10^{-8}$.

\begin{remark}[Why the bias correction matters]
Since $\vect{m}_0=\vect{v}_0=\vect{0}$, the raw averages
in~\eqref{eq:train-adam-moments} are biased toward zero for small $t$:
$\E[\vect{m}_t]\approx(1-\beta_1^{\,t})\,\E[\vect{g}_t]$. Dividing by
$(1-\beta_1^{\,t})$ and $(1-\beta_2^{\,t})$
in~\eqref{eq:train-adam-bias} removes this bias exactly; without it the
first updates would be far too small — dramatically so for
$\vect{v}_t$, whose decay rate $\beta_2 = 0.999$ makes the bias persist
for roughly a thousand steps. The ratio
$\hat{\vect{m}}_t/\sqrt{\hat{\vect{v}}_t}$ acts as a per-coordinate
signal-to-noise estimate, and the magnitude of each update is bounded by
approximately $\eta$ regardless of the gradient scale, making
Adam invariant to gradient rescaling. The decoupled-weight-decay variant
AdamW, which applies the $L_2$ shrinkage outside the adaptive rescaling,
is the default optimizer for modern Transformer models.
\end{remark}

% ============================================================================
\section{Learning-Rate Schedules}

The learning rate is the single most important hyperparameter of the
updates above, and it is rarely held constant. Writing $\eta_t$ for the
rate at step $t$ over a horizon of $T$ steps, the three standard schedules
are
\begin{align}
  \eta_t &= \eta_0\,\gamma^{\lfloor t/s \rfloor},
  \qquad 0<\gamma<1
  &&\text{(step decay, factor $\gamma$ every $s$ steps)},
  \label{eq:train-step-decay}\\[2pt]
  \eta_t &= \frac{\eta_0}{2}
    \Bigl(1+\cos\frac{\pi t}{T}\Bigr)
  &&\text{(cosine annealing)},
  \label{eq:train-cosine}\\[2pt]
  \eta_t &=
  \begin{cases}
    \eta_0\, t/T_w & t \le T_w\\[2pt]
    \dfrac{\eta_0}{2}
      \Bigl(1+\cos\dfrac{\pi (t-T_w)}{T-T_w}\Bigr) & t > T_w
  \end{cases}
  &&\text{(linear warmup, then cosine)}.
  \label{eq:train-warmup}
\end{align}
Warmup~\eqref{eq:train-warmup} protects the early phase of training — when
Adam's second-moment estimate $\hat{\vect{v}}_t$ is still noisy — from
destructively large steps, and is standard practice for Transformers; the
three profiles are compared in Figure~\ref{fig:train-schedules}.

\begin{figure}[htbp]
  \centering
  \begin{tikzpicture}
    \begin{axis}[
        width=0.8\textwidth, height=5.2cm,
        xmin=0, xmax=100, ymin=0, ymax=0.115,
        xlabel={training step $t$ (\% of horizon $T$)},
        ylabel={$\eta_t$},
        xlabel near ticks, ylabel near ticks,
        tick label style={font=\scriptsize},
        label style={font=\small},
        legend style={font=\scriptsize, draw=black!30},
        legend cell align=left,
      ]
      \addplot[coutput, thick, const plot, domain=0:100, samples=201]
        {0.1 * 0.5^(floor(x/30))};
      \addlegendentry{step decay~\eqref{eq:train-step-decay}}
      \addplot[cinput, thick, domain=0:100, samples=201]
        {0.05*(1+cos(deg(pi*x/100)))};
      \addlegendentry{cosine annealing~\eqref{eq:train-cosine}}
      \addplot[cgate, thick, domain=0:10, samples=21, forget plot]
        {0.1*x/10};
      \addplot[cgate, thick, domain=10:100, samples=181]
        {0.05*(1+cos(deg(pi*(x-10)/90)))};
      \addlegendentry{warmup + cosine~\eqref{eq:train-warmup}}
    \end{axis}
  \end{tikzpicture}
  \caption{Learning-rate schedules with $\eta_0=0.1$: step decay
    ($\gamma=0.5$, $s=0.3\,T$), cosine annealing, and linear warmup over
    the first $10\%$ of training followed by cosine annealing.}
  \label{fig:train-schedules}
\end{figure}

% ============================================================================
\section{Weight Initialization}\label{sec:train-init}

Remark~\ref{rem:train-vanish} showed that signals are multiplied by a
weight matrix at every layer; initialization must therefore keep the
variance of activations and of backpropagated gradients approximately
constant with depth~\cite{goodfellow2016book}. For a layer with
$n_{\mathrm{in}}$ inputs and $n_{\mathrm{out}}$ outputs,
\emph{Xavier/Glorot} initialization — appropriate for symmetric, roughly
linear-around-zero activations such as $\tanh$ — balances both passes:
\begin{equation}\label{eq:train-xavier}
  \operatorname{Var}\bigl(W_{ij}\bigr)
    = \frac{2}{n_{\mathrm{in}}+n_{\mathrm{out}}},
  \qquad\text{e.g.}\quad
  W_{ij} \sim
  \mathcal{U}\!\left[
    -\sqrt{\tfrac{6}{n_{\mathrm{in}}+n_{\mathrm{out}}}},\;
    \sqrt{\tfrac{6}{n_{\mathrm{in}}+n_{\mathrm{out}}}}
  \right].
\end{equation}
\emph{He/Kaiming} initialization corrects for the fact that ReLU zeroes
half of its inputs, which halves the activation variance at each layer:
\begin{equation}\label{eq:train-he}
  \operatorname{Var}\bigl(W_{ij}\bigr) = \frac{2}{n_{\mathrm{in}}},
  \qquad
  W_{ij} \sim
  \mathcal{N}\!\Bigl(0,\; \tfrac{2}{n_{\mathrm{in}}}\Bigr),
\end{equation}
and is the default for ReLU-family networks. Biases are initialized to
zero in both schemes. Together with the schedules of
Figure~\ref{fig:train-schedules} and an adaptive optimizer such as
Adam~\eqref{eq:train-adam-moments}--\eqref{eq:train-adam-update}, these
choices form the standard modern training recipe on which the
regularization techniques of the next chapter are layered.