Training Deep Networks
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 , learning rate , and stabilizer . At step the raw gradient of the loss (excluding regularization) is .
(a) Write the full Adam update equations for , their bias-corrected forms , and the parameter update . Explain why the bias correction factors and are needed, by computing under the assumption that are i.i.d. with . (6)
(b) In classic Adam, L2 regularization is applied by adding to the gradient: . 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 , starting from , and a constant gradient . Compute the parameter change (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 scalar pre-activations , batch norm computes
(a) Derive (with ) and hence show that the batch-normalized activations satisfy and -type terms appear in the backward pass. State the complete formula for in terms of the upstream gradients . (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 leaves the layer's output unchanged, and derive how the gradient w.r.t. 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 , and L2 weight decay applied decoupled (AdamW-style, but on SGD). Clearly show the order of operations. (8)
(b) Prove that global-norm clipping preserves gradient direction: if is the concatenated gradient with , the clipped gradient satisfies for some scalar , and that (i.e. it remains a descent direction). Then show clipping is equivalent to solving subject to . (6)
(c) You add learning-rate warmup for the first steps (linear from to ) followed by cosine decay to over the remaining steps. Write the closed-form schedule 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)
Why bias correction: unrolling with , Taking expectation with : So is biased toward 0 (by factor , large early). Dividing by gives , an unbiased estimate. Same argument for . (2 for derivation, 1 for conclusion)
(b) L2 vs decoupled decay (8 marks)
With L2-in-gradient, , the decay term enters and , so it gets divided by . The effective shrinkage on parameter becomes which is inversely scaled by each parameter's gradient magnitude. Parameters with large historical gradients (large ) 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 ) drives , and decay is applied directly to the weights: Now every parameter is shrunk by the same factor regardless of , recovering true weight decay. (4 marks: 2 for the fixed update line, 2 for explanation)
(c) Numerical trace (8 marks)
.
Step 1:
- , .
- , , .
- update.
Step 2:
- , .
- , , .
- update.
(With constant gradient and , each bias-corrected step gives exactly ; total .) (4 marks per-step, 4 for final value; award full if they note the elegant cancellation.)
Question 2
(a) Backward pass (8 marks)
With , . Directly since . (1)
Jacobian: using and ,
=\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 =