Lecture 6 — Training Dynamics

Project connection. Project Step 6 trains the 816,640-parameter model of Lecture 5 for 5,000 steps with batch size 64. The empirical table in Section 6 is the reference run for $(V,T_{\max},d,H,L)=(65,64,128,4,4)$.

Chapter overview. Initialization, adaptive preconditioning, learning- rate schedules, and gradient clipping can be formulated precisely and tested directly. The broader question—why an overparametrized non-convex model both optimizes and generalizes—has only partial answers. This chapter separates exact calculations and algorithm definitions from empirical observations and theoretical regimes.

0. Setup

Let $\theta\in\mathbb R^n$ be the parameter vector and $\mathcal L(\theta)$ the empirical loss. A training method specifies an initial random vector $\theta_0$, a gradient estimator $g_t$, and an update rule for producing $\theta_t$ from $\theta_{t-1}$.

For a random scalar $X$ with finite second moment,

$$\operatorname{Var}(X)=\mathbb E[(X-\mathbb EX)^2].$$

If $X$ and $Y$ are independent, then $\mathbb E[XY]=\mathbb E[X]\mathbb E[Y]$; if they are also centered, the variance of their product is $\operatorname{Var}(XY)=\mathbb E[X^2]\mathbb E[Y^2]$.

1. Initialization as variance bookkeeping

Consider one linear layer

$$y=Wx, \qquad W\in\mathbb R^{m\times n}, \quad x\in\mathbb R^n.$$

Assume the entries $W_{ij}$ are independent with mean zero and variance $\sigma_W^2$, the coordinates $x_j$ are independent with mean zero and variance $\sigma_x^2$, and $W$ is independent of $x$.

Lemma 1.1 (variance through a linear layer). For every output coordinate,

$$\operatorname{Var}(y_i)=n\sigma_W^2\sigma_x^2.$$

Proof. Since $y_i=\sum_{j=1}^nW_{ij}x_j$, its summands are centered and pairwise uncorrelated. Each has variance $\mathbb E[W_{ij}^2]\mathbb E[x_j^2]=\sigma_W^2\sigma_x^2$. Variances therefore add. $\square$

To preserve forward variance, choose $\sigma_W^2=1/n$, called fan-in scaling. In the backward relation $\bar x=W^\top\bar y$, the analogous calculation gives $\sigma_W^2=1/m$, called fan-out scaling. Xavier, or Glorot, initialization compromises with

$$\sigma_W^2=\frac{2}{n+m}.$$

For a symmetric centered input, ReLU sets half of the values to zero and halves the second moment. He initialization compensates by taking $\sigma_W^2=2/n$.

Remark 1.2 (depth). If each layer multiplies activation variance by a factor $c$, a depth-$L$ composition multiplies it by $c^L$. A small per-layer scale error therefore becomes exponential in depth. Nonlinear activation means and dependencies make this calculation approximate, but it remains a useful initialization diagnostic.

In a residual network, suppose the stream is updated by $2L$ approximately centered, uncorrelated contributions, each having variance $v$. Then

$$\operatorname{Var}(X^{(L)}) \approx\operatorname{Var}(X^{(0)})+2Lv.$$

Scaling each branch’s output projection by $1/\sqrt{2L}$ replaces $v$ by $v/(2L)$ and keeps the total added variance of order one. The GPT-2-style initialization applies this scaling to each attention $W_O$ and MLP $W_2$.

Project observation 1.3. The reference model begins at validation loss 4.3039, close to the uniform value $\ln65=4.1744$. A much larger initial loss indicates logits with excessive spread. The small positive gap comes from the particular unembedding initialization used by PyTorch.

2. Adam and AdamW

For vectors, $u^{\odot2}$ denotes coordinatewise squaring, and square roots and divisions in this section are also coordinatewise. Let $g_t$ be a minibatch gradient evaluated at $\theta_{t-1}$, choose $0\leq\beta_1,\beta_2<1$, and initialize $m_0=v_0=0$.

Definition 2.1 (adaptive moment estimation; Adam). Adam forms exponential moving averages

$$\begin{aligned} m_t&=\beta_1m_{t-1}+(1-\beta_1)g_t, &\widehat m_t&=\frac{m_t}{1-\beta_1^t},\\ v_t&=\beta_2v_{t-1}+(1-\beta_2)g_t^{\odot2}, &\widehat v_t&=\frac{v_t}{1-\beta_2^t}, \end{aligned}$$

and updates

$$\theta_t =\theta_{t-1}-\eta_t \frac{\widehat m_t}{\sqrt{\widehat v_t}+\varepsilon},$$

where $\varepsilon>0$ prevents division by zero.

The vector $m_t$ estimates a first moment and supplies momentum; $v_t$ estimates a coordinatewise second moment and supplies an adaptive diagonal preconditioner. Here a diagonal preconditioner means a coordinatewise rescaling of the update direction.

Proposition 2.2 (loss-scale invariance). Set $\varepsilon=0$ and assume $\widehat v_{t,i}>0$ for every coordinate under consideration. If the loss and hence every gradient are multiplied by a constant $c>0$, the Adam update direction is unchanged.

Proof. The first-moment estimate is multiplied by $c$, and the second-moment estimate by $c^2$. Thus $\widehat m_t/\sqrt{\widehat v_t}$ is unchanged. $\square$

Remark 2.3 (coordinatewise preconditioning). A coordinate with a persistently small gradient also tends to have a small second-moment denominator. Adam therefore reduces sensitivity to coordinatewise gradient scale, though it does not make Adam a true inverse-Hessian method.

If a random gradient has time-independent mean $\mu$, then

$$\mathbb E[m_t]=(1-\beta_1^t)\mu.$$

Similarly, if $\mathbb E[g_t^{\odot2}]$ is time-independent, the same factor $1-\beta_2^t$ appears for $v_t$. The denominators in Definition 2.1 are therefore exact bias corrections under these stationary-moment assumptions.

Definition 2.4 (AdamW). Given a weight-decay coefficient $\lambda\geq0$, AdamW uses

$$\theta_t =(1-\eta_t\lambda)\theta_{t-1} -\eta_t\frac{\widehat m_t}{\sqrt{\widehat v_t}+\varepsilon}.$$

This is decoupled weight decay: shrinkage is applied directly to the weights rather than being included in the gradient before adaptive preconditioning.

Remark 2.5 ($L^2$ regularization versus AdamW). For plain SGD, adding $(\lambda/2)\|\theta\|_2^2$ to the loss produces the same first-order update as multiplicative decay. Under Adam, the added gradient $\lambda\theta$ is divided by a coordinate-dependent denominator, so it no longer gives uniform shrinkage. AdamW preserves the latter interpretation.

Hyperparameter aside 2.5. The project uses $\beta_2=0.99$ instead of 0.999. An exponential moving average has an effective horizon on the order of $1/(1-\beta_2)$, so the smaller value adapts its second-moment estimate more quickly at the cost of greater sampling noise. This choice is an empirical convention rather than a theorem.

3. Warmup and cosine decay

Fix endpoints $0<t_w<t_f$ and learning rates $0\leq\eta_{\min}\leq\eta_{\max}$. Define the linear-warmup, cosine-decay schedule for $0\leq t\leq t_f$ by

$$\eta(t)= \begin{cases} \eta_{\max}\,t/t_w,&0\leq t<t_w,\\[4pt] \eta_{\min} +\dfrac{\eta_{\max}-\eta_{\min}}2 \left(1+\cos\dfrac{\pi(t-t_w)}{t_f-t_w}\right), &t_w\leq t\leq t_f. \end{cases}$$

If training continues beyond $t_f$, set $\eta(t)=\eta_{\min}$. The schedule is continuous, reaches $\eta_{\max}$ at $t_w$, and reaches $\eta_{\min}$ at $t_f$.

Remark 3.1 (warmup). At the beginning of training, moment estimates are based on few samples and model activations may change rapidly. Warmup limits the size of early steps. Its practical value is well established, but no single general theorem explains the best warmup length.

Remark 3.2 (decay and stochastic noise). Near a minimizer, random gradient error may dominate the mean gradient. For a quadratic model with constant step size, the iterates fluctuate in a stationary neighborhood whose size decreases with the step size. Learning-rate decay reduces this noise floor. A finite cosine schedule resembles the noise-control role of Robbins—Monro decay but does not satisfy its infinite-horizon conditions.

The reference run uses $\eta_{\max}=3\times10^{-3}$, $t_w=100$, $t_f=5000$, and $\eta_{\min}=3\times10^{-4}$.

4. Gradient clipping

Let $\|g\|_2=(\sum_i g_i^2)^{1/2}$ be the Euclidean norm.

Definition 4.1 (global norm clipping). For a threshold $c>0$, replace a gradient $g$ by

$$\operatorname{clip}_c(g) =\begin{cases} g,&\|g\|_2\leq c,\\ c\,g/\|g\|_2,&\|g\|_2>c. \end{cases}$$

Thus clipping preserves direction and bounds the norm by $c$.

Remark 4.2 (bias). Even if $g$ is an unbiased stochastic gradient, the nonlinear random vector $\operatorname{clip}_c(g)$ is generally not unbiased. Clipping therefore violates a hypothesis of the basic Robbins—Monro theorem. It is used as protection against rare large updates rather than as an exact estimator.

Project observation 4.3. In the reference run, pre-clipping norms stay between 0.298 and 0.414 and never reach the threshold 1.0. Clipping is inactive in that run. Logging the norm is what makes this conclusion possible.

5. Partial theories of deep-network training

The Hessian of a twice differentiable loss is $\nabla^2\mathcal L(\theta)$, the matrix of second partial derivatives. A point with zero gradient is stationary. A stationary point with Hessian having both positive and negative eigenvalues is a saddle point. A model is overparametrized relative to a dataset when it has enough degrees of freedom to fit, and often exactly interpolate, the training observations.

5.1 Landscape observations

Remark 5.1 (high-dimensional saddles). In high dimension, requiring every Hessian eigenvalue to be positive is a strong condition. Many non-minimizing stationary points are saddles, and stochastic noise can help an optimizer leave directions of negative curvature. This geometric heuristic does not by itself prove convergence for a transformer.

Remark 5.2 (mode connectivity). Empirically, independently trained overparametrized networks can often be connected by low-loss curves in parameter space. Global minimizers may form large sets rather than isolated points. Parameter symmetries, such as hidden-unit permutations, guarantee some multiplicity; the observed connectivity is a stronger phenomenon.

5.2 Neural tangent kernel

For fixed training inputs $x_1,\ldots,x_N$ and a scalar network output $f_\theta(x)$, define the empirical neural tangent kernel (NTK) by

$$K_\theta(i,j) =\left\langle\nabla_\theta f_\theta(x_i), \nabla_\theta f_\theta(x_j)\right\rangle.$$

It is positive semidefinite because, for every $u\in\mathbb R^N$, $u^\top K_\theta u=\|\sum_i u_i\nabla_\theta f_\theta(x_i)\|^2\geq0$. A symmetric matrix $K$ is positive definite if $u^\top Ku>0$ for every nonzero $u$. Gradient flow is the continuous-time equation $\dot\theta(t)=-\nabla_\theta\mathcal L(\theta(t))$.

Theorem 5.3 (neural-tangent-kernel dynamics). Let $y\in\mathbb R^N$ be the target vector and use squared loss

$$\mathcal L(\theta)=\frac1{2N}\sum_{i=1}^N (f_\theta(x_i)-y_i)^2.$$

Suppose that along gradient flow the kernel is a fixed matrix $K_{\theta(t)}=K$. If $f(t)=(f_{\theta(t)}(x_i))_{i=1}^N$, then

$$\dot f(t)=-\frac1N K(f(t)-y).$$

If $K$ is positive definite with smallest eigenvalue $\lambda_{\min}>0$, then

$$\mathcal L(\theta(t)) \leq e^{-2\lambda_{\min}t/N}\mathcal L(\theta(0)),$$

so the training predictions converge to $y$.

Proof. The chain rule and the gradient-flow equation give

$$\dot f_i =\langle\nabla_\theta f_i,\dot\theta\rangle =-\frac1N\sum_jK(i,j)(f_j-y_j).$$

For $e=f-y$,

$$\frac{d}{dt}\|e\|^2=-(2/N)e^\top Ke \leq-(2\lambda_{\min}/N)\|e\|^2.$$

Integrating this differential inequality proves the bound. $\square$

Remark 5.4 (infinite-width limit and limitation). For certain fixed-depth network families with inverse-square-root width scaling, random initialization, and regular activations, the initial empirical NTK converges as width tends to infinity to a deterministic kernel and changes negligibly on finite training intervals. Theorem 5.3 then describes the limiting output dynamics. Because the kernel is fixed, features remain close to initialization; this rigorous regime does not explain practical feature learning such as the changing embedding geometry in Step 3.

5.3 Double descent

The interpolation threshold is the capacity at which a model first achieves essentially zero training error.

Observation 5.2 (double descent). In many model families, test error first follows a classical decreasing-then-increasing curve as capacity grows, peaks near the interpolation threshold, and decreases again in the overparametrized regime.

Remark 5.5 (implicit regularization). Implicit regularization is a preference induced by the optimization algorithm without an explicit penalty in the loss. In linear regression, gradient methods can select a minimum-norm interpolating solution, and increasing dimension supplies additional low-norm interpolants. This gives a precise double-descent analysis in that setting. The analogous explanation for nonlinear feature-learning networks remains incomplete.

5.4 Edge of stability

For $\lambda>0$, consider the quadratic function $\mathcal L(\theta)=\tfrac12\lambda\theta^2$, gradient descent with constant step size $\eta$ updates $\theta_t=(1-\eta\lambda)\theta_{t-1}$. It is stable exactly when $0<\eta\lambda<2$.

Empirical remark 5.3 (edge of stability). In trained networks, the largest Hessian eigenvalue often grows toward approximately $2/\eta$ and remains near that boundary while the loss decreases non-monotonically. This edge-of-stability behavior is reproducible but not explained by the one-dimensional quadratic theory.

5.5 Status of the claims

Theoretical aside 5.4. Variance calculations, Adam’s scale invariance, and stochastic-approximation convergence under stated hypotheses are exact. The NTK is a theorem about a limiting regime. Mode connectivity, double descent in nonlinear networks, edge of stability, and practical warmup behavior are empirical phenomena with only partial theories. Why the solutions selected by large-scale training generalize remains open.

6. The reference training run

The reference optimizer is AdamW with $\eta_{\max}=3\times10^{-3}$, $\lambda=0.1$, $(\beta_1,\beta_2)=(0.9,0.99)$, 100 warmup steps, cosine decay through step 5,000, and clipping threshold 1.0.

step	train loss	validation loss	validation bits/char	gap	gradient norm	learning rate
0	4.3064	4.3039	6.209	—	—	0
500	1.6556	1.8134	2.616	0.158	0.322	$2.96\times10^{-3}$
1000	1.4714	1.6605	2.396	0.189	0.310	$2.78\times10^{-3}$
1500	1.4019	1.6087	2.321	0.207	0.304	$2.49\times10^{-3}$
2000	1.3478	1.5679	2.262	0.220	0.302	$2.12\times10^{-3}$
2500	1.3028	1.5384	2.219	0.235	0.298	$1.69\times10^{-3}$
3000	1.2713	1.5209	2.194	0.250	0.325	$1.27\times10^{-3}$
3500	1.2363	1.5197	2.192	0.283	0.334	$8.78\times10^{-4}$
4000	1.2037	1.5188	2.191	0.315	0.373	$5.68\times10^{-4}$
4500	1.1848	1.5093	2.177	0.325	0.414	$3.69\times10^{-4}$
5000	1.1656	1.5178	2.190	0.352	0.399	$3.00\times10^{-4}$

Empirical observations 6.1. The initial loss is near $\ln V$. Most improvement occurs in the first 500 steps. The training—validation gap grows from 0.158 to 0.352, indicating increasing overfit. Validation loss is essentially flat after step 3,000 and is best at step 4,500 rather than at the final step; a production run would retain the best validation checkpoint. Pre-clipping gradient norms rise late while the learning rate falls. The actual step scale $\eta_t\|g_t\|$ still shrinks. The norm trend is consistent with movement into a sharper region, but gradient norm alone is not a direct Hessian measurement.

Benchmark aside 6.2. The add-one bigram obtains 3.5417 bits per character. The final transformer obtains 2.1897, and the best checkpoint obtains 2.1774. The reference computation uses 816,640 parameters and about 73 minutes of laptop CPU time, closing roughly half of the bits-per-character gap between the bigram and the approximate one-bit Shannon benchmark. Samples acquire character frequencies, word forms, line breaks, and SPEAKER: structure in that order; this qualitative ordering motivates the scaling discussion in Lecture 7.

Summary

Variance-preserving initialization controls the scale of forward signals and backward gradients; residual branches motivate the $1/\sqrt{2L}$ output scaling. Adam combines momentum with a diagonal second-moment preconditioner, and AdamW separates weight decay from that preconditioner. Warmup limits unreliable early steps, cosine decay reduces late stochastic motion, and clipping bounds rare large gradients at the cost of bias. These algorithmic statements are more complete than the theory of feature learning and generalization: NTK results describe an infinite-width regime, while double descent and edge-of-stability behavior remain only partially explained.

Exercises (paired with Step 6)

A star marks a Project Step 6 task.

1. ★ Prove Lemma 1.1, derive fan-in and fan-out scalings, and derive the $1/\sqrt{2L}$ residual correction.
2. ★ Write AdamW in full, prove Proposition 2.2, and verify the bias- correction formulas under stationary first and second moments.
3. Prove that $L^2$ regularization and decoupled weight decay agree for plain SGD to first order but differ under Adam. Express the discrepancy using $\widehat v_t$.
4. ★ Implement the warmup—cosine schedule and test its values at $t=0,t_w,t_f$.
5. ★ Compare depth and width at fixed parameter count. State carefully why the experiment neither proves nor refutes Lecture 2’s worst-case depth separation theorem.
6. For a one-dimensional quadratic loss with additive mean-zero gradient noise of variance $\sigma^2$, compute the stationary variance of constant-step SGD and determine its scaling with $\eta$.
7. Analyze double descent in minimum-norm linear regression as the feature dimension crosses the sample size.
8. For $\mathcal L(\theta)=\tfrac12a\theta^2$, prove the stability boundary $\eta a=2$ and explain why it does not itself explain network edge-of- stability behavior.

Pointers

Glorot and Bengio (2010); He et al. (2015); Kingma and Ba, Adam (2014); Loshchilov and Hutter, AdamW (2017); Jacot et al., Neural Tangent Kernel (2018); Belkin et al., Reconciling Modern Machine-Learning Practice and the Classical Bias—Variance Trade-Off (2019); Cohen et al., Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability (2021). See the resources page and Project Step 6.