Level 5 — MasteryTraining Deep Networks

Training Deep Networks

90 minutes60 marksprintable — key stays hidden on paper

Time limit: 90 minutes Total marks: 60 Instructions: Answer all three questions. Show full derivations. Partial credit is awarded for correct reasoning even with arithmetic slips. Use ...... notation for mathematics. Code may be written in Python/NumPy pseudocode where requested.


Question 1 — Adam, AdamW and the geometry of decoupled decay (22 marks)

Consider the Adam update with hyperparameters β1,β2[0,1)\beta_1, \beta_2 \in [0,1), learning rate η\eta, and stabilizer ϵ\epsilon. At step tt the raw gradient of the loss (excluding regularization) is gtg_t.

(a) Write the full Adam update equations for mt,vtm_t, v_t, their bias-corrected forms m^t,v^t\hat m_t, \hat v_t, and the parameter update θt\theta_t. Explain why the bias correction factors 11β1t\frac{1}{1-\beta_1^t} and 11β2t\frac{1}{1-\beta_2^t} are needed, by computing E[mt]\mathbb{E}[m_t] under the assumption that gig_i are i.i.d. with E[gi]=g\mathbb{E}[g_i]=g. (6)

(b) In classic Adam, L2 regularization is applied by adding λθ\lambda\theta to the gradient: gtgt+λθt1g_t \leftarrow g_t + \lambda\theta_{t-1}. Show mathematically that this does not produce the same effective per-parameter shrinkage as true weight decay, and explain precisely how AdamW fixes it. Give the AdamW update line explicitly. (8)

(c) Numerical trace. Take β1=0.9, β2=0.999, η=0.1, ϵ=0\beta_1=0.9,\ \beta_2=0.999,\ \eta=0.1,\ \epsilon=0, starting from m0=v0=0m_0=v_0=0, and a constant gradient g1=g2=2g_1=g_2=2. Compute the parameter change Δθ=θ2θ0\Delta\theta = \theta_2-\theta_0 (sum of the two update steps). Report to 4 decimal places. (8)


Question 2 — Batch normalization: gradients, physics analogy, and inference (20 marks)

For a mini-batch of mm scalar pre-activations x1,,xmx_1,\dots,x_m, batch norm computes μ=1mixi,σ2=1mi(xiμ)2,x^i=xiμσ2+ϵ,yi=γx^i+β.\mu=\tfrac1m\sum_i x_i,\quad \sigma^2=\tfrac1m\sum_i (x_i-\mu)^2,\quad \hat x_i=\frac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}},\quad y_i=\gamma\hat x_i+\beta.

(a) Derive x^ixj\frac{\partial \hat x_i}{\partial x_j} (with ϵ0\epsilon\to 0) and hence show that the batch-normalized activations satisfy ix^i=0\sum_i \hat x_i = 0 and iLx^ix^i\sum_i \frac{\partial \mathcal L}{\partial \hat x_i}\,\hat x_i-type terms appear in the backward pass. State the complete formula for Lxi\frac{\partial \mathcal L}{\partial x_i} in terms of the upstream gradients Lx^i\frac{\partial \mathcal L}{\partial \hat x_i}. (8)

(b) Explain, using a dimensional/scale argument (a "physics of signal propagation" view), why batch norm makes the loss landscape less sensitive to the scale of the weights in the preceding linear layer. Formally show that scaling the weights WcWW\to cW leaves the layer's output yiy_i unchanged, and derive how the gradient w.r.t. WW scales. (6)

(c) At inference we cannot use batch statistics. Describe the running-average mechanism, and explain the failure mode when training uses batch size 1 or when train/test distributions differ. Contrast this with layer normalization, stating one concrete architecture (e.g. Transformers) where LN is preferred and why. (6)


Question 3 — Build & prove: a robust training loop (18 marks)

You must train a deep net that suffers from occasional exploding gradients and a noisy loss curve.

(a) Write NumPy-style pseudocode for a training step that combines: mini-batch SGD with momentum, global-norm gradient clipping at threshold cc, and L2 weight decay λ\lambda applied decoupled (AdamW-style, but on SGD). Clearly show the order of operations. (8)

(b) Prove that global-norm clipping preserves gradient direction: if gg is the concatenated gradient with g>c\|g\|>c, the clipped gradient g=cg/gg' = c\,g/\|g\| satisfies g=αgg' = \alpha g for some scalar α(0,1)\alpha\in(0,1), and that g,g>0\langle g',g\rangle>0 (i.e. it remains a descent direction). Then show clipping is equivalent to solving minggg2\min_{g'} \|g'-g\|^2 subject to gc\|g'\|\le c. (6)

(c) You add learning-rate warmup for the first WW steps (linear from 00 to ηmax\eta_{\max}) followed by cosine decay to 00 over the remaining TWT-W steps. Write the closed-form schedule η(t)\eta(t) and justify in 2–3 sentences why warmup specifically helps when using large batches / adaptive optimizers. (4)

Answer keyMark scheme & solutions

Question 1

(a) Adam equations (3 marks) + bias correction reasoning (3 marks)

mt=β1mt1+(1β1)gt,vt=β2vt1+(1β2)gt2m_t=\beta_1 m_{t-1}+(1-\beta_1)g_t,\qquad v_t=\beta_2 v_{t-1}+(1-\beta_2)g_t^2 m^t=mt1β1t,v^t=vt1β2t,θt=θt1ηm^tv^t+ϵ\hat m_t=\frac{m_t}{1-\beta_1^t},\qquad \hat v_t=\frac{v_t}{1-\beta_2^t},\qquad \theta_t=\theta_{t-1}-\eta\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}

Why bias correction: unrolling with m0=0m_0=0, mt=(1β1)i=1tβ1tigi.m_t=(1-\beta_1)\sum_{i=1}^{t}\beta_1^{t-i}g_i. Taking expectation with E[gi]=g\mathbb E[g_i]=g: E[mt]=g(1β1)i=1tβ1ti=g(1β1)1β1t1β1=g(1β1t).\mathbb E[m_t]=g(1-\beta_1)\sum_{i=1}^t\beta_1^{t-i}=g(1-\beta_1)\frac{1-\beta_1^t}{1-\beta_1}=g(1-\beta_1^t). So mtm_t is biased toward 0 (by factor 1β1t1-\beta_1^t, large early). Dividing by 1β1t1-\beta_1^t gives E[m^t]=g\mathbb E[\hat m_t]=g, an unbiased estimate. Same argument for vtv_t. (2 for derivation, 1 for conclusion)

(b) L2 vs decoupled decay (8 marks)

With L2-in-gradient, gtgt+λθt1g_t\leftarrow g_t+\lambda\theta_{t-1}, the decay term λθ\lambda\theta enters mtm_t and vtv_t, so it gets divided by v^t+ϵ\sqrt{\hat v_t}+\epsilon. The effective shrinkage on parameter jj becomes ηλθjv^t,j+ϵ,\eta\frac{\lambda\theta_j}{\sqrt{\hat v_{t,j}}+\epsilon}, which is inversely scaled by each parameter's gradient magnitude. Parameters with large historical gradients (large v^\hat v) get less decay; those with small gradients get more. This couples regularization strength to gradient statistics — not the intended uniform shrinkage. (4 marks)

AdamW decouples decay from the adaptive step: gradient (no λθ\lambda\theta) drives mt,vtm_t,v_t, and decay is applied directly to the weights: θt=θt1η(m^tv^t+ϵ+λθt1).\theta_t=\theta_{t-1}-\eta\left(\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}+\lambda\theta_{t-1}\right). Now every parameter is shrunk by the same factor (1ηλ)(1-\eta\lambda) regardless of v^\hat v, recovering true weight decay. (4 marks: 2 for the fixed update line, 2 for explanation)

(c) Numerical trace (8 marks)

β1=0.9,β2=0.999,η=0.1,ϵ=0,g=2\beta_1=0.9,\beta_2=0.999,\eta=0.1,\epsilon=0,g=2.

Step 1:

  • m1=0.12=0.2m_1=0.1\cdot2=0.2, m^1=0.2/(10.9)=2\hat m_1=0.2/(1-0.9)=2.
  • v1=0.0014=0.004v_1=0.001\cdot4=0.004, v^1=0.004/(10.999)=4\hat v_1=0.004/(1-0.999)=4, =2\sqrt{}=2.
  • update1=0.12/2=0.1_1=-0.1\cdot 2/2=-0.1.

Step 2:

  • m2=0.90.2+0.12=0.38m_2=0.9\cdot0.2+0.1\cdot2=0.38, m^2=0.38/(10.81)=0.38/0.19=2\hat m_2=0.38/(1-0.81)=0.38/0.19=2.
  • v2=0.9990.004+0.0014=0.003996+0.004=0.007996v_2=0.999\cdot0.004+0.001\cdot4=0.003996+0.004=0.007996, v^2=0.007996/(10.998001)=0.007996/0.001999=4.0000...\hat v_2=0.007996/(1-0.998001)=0.007996/0.001999=4.0000..., =2.0\sqrt{}=2.0.
  • update2=0.12/2=0.1_2=-0.1\cdot2/2=-0.1.

Δθ=0.10.1=0.2000.\Delta\theta = -0.1-0.1 = -0.2000.

(With constant gradient and ϵ=0\epsilon=0, each bias-corrected step gives exactly η=0.1-\eta=-0.1; total 0.2000\boxed{-0.2000}.) (4 marks per-step, 4 for final value; award full if they note the elegant cancellation.)


Question 2

(a) Backward pass (8 marks)

With ϵ0\epsilon\to0, x^i=(xiμ)/σ\hat x_i=(x_i-\mu)/\sigma. Directly ix^i=1σi(xiμ)=0\sum_i\hat x_i=\frac1\sigma\sum_i(x_i-\mu)=0 since i(xiμ)=0\sum_i(x_i-\mu)=0. (1)

Jacobian: using μ/xj=1/m\partial\mu/\partial x_j=1/m and σ2/xj=2m(xjμ)\partial\sigma^2/\partial x_j=\frac2m(x_j-\mu),

=\frac1\sigma\Big(\delta_{ij}-\tfrac1m-\tfrac1m\hat x_i\hat x_j\Big).$$ *(3)* Chaining $\frac{\partial\mathcal L}{\partial x_i}=\sum_j\frac{\partial\mathcal L}{\partial\hat x_j}\frac{\partial\hat x_j}{\partial x_i}$ gives the standard result: $$\boxed{\frac{\partial\mathcal L}{\partial x_i}=\frac{1}{m\sigma}\left(m\frac{\partial\mathcal L}{\partial\hat x_i}-\sum_j\frac{\partial\mathcal L}{\partial\hat x_j}-\hat x_i\sum_j\frac{\partial\mathcal L}{\partial\hat x_j}\hat x_j\right).}$$ The two subtracted terms remove the mean and the projection onto $\hat x$ — hence the appearance of $\sum_j(\partial\mathcal L/\partial\hat x_j)\hat x_j$. *(4)* **(b) Scale invariance (6 marks)** Let pre-activation $x_i=w^\top a_i$. Under $W\to cW$: $x_i\to cx_i$, $\mu\to c\mu$, $\sigma\to |c|\sigma$, so $\hat x_i=(cx_i-c\mu)/(|c|\sigma)=\text{sign}(c)\,\hat x_i$ — for $c>0$ exactly unchanged, hence $y_i$ unchanged. *(3)* Thus BN removes the degree of freedom of weight scale, flattening that direction of the loss landscape (the loss is invariant to radial scaling of $W$). *(1)* Gradient: since output is invariant, $\partial\mathcal L/\partial(cW)=\frac1c\partial\mathcal L/\partial W$ — gradient magnitude scales as $1/c$, so large weights get proportionally smaller gradients, an automatic scale-stabilizing effect. *(2)* **(c) Inference / LN contrast (6 marks)** Running averages: during training maintain $\mu_{run}\leftarrow(1-\rho)\mu_{run}+\rho\mu_{batch}$ and similarly $\sigma^2_{run}$; at inference use these fixed statistics so output is deterministic and batch-independent. *(2)* Failure modes: batch size 1 → $\sigma^2=0$/undefined and extremely noisy running stats; train/test distribution shift → stored stats mismatch test data, degrading accuracy. *(2)* LN normalizes across features *within a single example* (no batch dependence), so it is preferred in Transformers / RNNs where sequence lengths and batch composition vary and where per-example stability matters; identical behavior train vs test. *(2)* --- ## Question 3 **(a) Pseudocode (8 marks)** ```python # state: v (momentum buffer, init 0) def train_step(params, grads_list, v, lr, mu, c, wd): # 1. global-norm clip on raw gradients total_norm = np.sqrt(sum(np.sum(g**2) for g in grads_list)) scale = min(1.0, c / (total_norm + 1e-12)) grads = [g * scale for g in grads_list] # 2. momentum update v = [mu * vi + gi for vi, gi in zip(v, grads)] # 3. parameter update: gradient step + DECOUPLED weight decay params = [p - lr * (vi + 0.0) - lr * wd * p for p, vi in zip(params, v)] return params, v ``` Marks: clip before momentum (2), momentum buffer correct (2), decoupled decay `-lr*wd*p` applied to weights not folded into grad (3), correct ordering & return of state (1). **(b) Direction preservation (6 marks)** If $\|g\|>c$: $g'=c\,g/\|g\|=\alpha g$ with $\alpha=c/\|g\|$. Since $c>0$ and $c<\|g\|$, we have $0<\alpha<1$. *(2)* Then $\langle g',g\rangle=\alpha\|g\|^2>0$, so $g'$ has positive inner product with $g$ and $-g'$ remains a descent direction (same direction as $-g$). *(2)* Optimization equivalence: minimize $f(g')=\|g'-g\|^2$ s.t. $\|g'\|\le c$. If $\|g\|\le c$ the unconstrained min $g'=g$ is feasible → no clipping. If $\|g\|>c$, the minimizer lies on the boundary $\|g'\|=c$; the closest point on a ball of radius $c$ to external point $g$ is the radial projection $c\,g/\|g\|$ (Lagrange: $g'-g=\mu g' \Rightarrow g'\parallel g$). Hence clipping = Euclidean projection onto the $c$-ball. *(2)* **(c) Warmup + cosine (4 marks)** $$\eta(t)=\begin{cases}\eta_{\max}\dfrac{t}{W}, & t\le W\\[2mm] \dfrac{\eta_{\max}}{2}\left(1+\cos\!\left(\pi\dfrac{t-W}{T-W}\right)\right), & W<t\le T.\end{cases}$$ *(2 marks)* Justification: early in training gradient/second-moment estimates are unreliable and weights are near random init; large steps (especially with adaptive optimizers whose $\hat v$ is not yet stable, and large batches with low-noise gradients) can diverge. Warmup lets statistics settle and keeps early updates small, then cosine decay anneals for fine convergence. *(2 marks)* ```verify [ {"claim": "Adam with constant grad g=2, eps=0 gives per-step update of exactly -eta, total -0.2 over 2 steps", "code": "b1,b2,eta,g=Rational(9,10),Rational(999,1000),Rational(1,10),2\nm=0;v=0;theta=0\nfor t in [1,2]:\n m=b1*m+(1-b1)*g\n v=b2*v+(1-b2)*g**2\n mh=m/(1-b1**t)\n vh=v/(1-b2**t)\n theta=theta-eta*mh/(sqrt(vh))\nresult =