Deep Learning Frameworks

Question 1

(a) Derivation (8 marks)

Define intermediates:

• $z = Wx + b \in \mathbb{R}^2$
• $h = \tanh(z)$, elementwise; $h' = 1 - \tanh^2(z) = 1 - h^2$ (2)
• $\hat{y} = w^\top h$, $L = \tfrac12(\hat y - y)^2$

Let $\delta \equiv \dfrac{\partial L}{\partial \hat y} = (\hat y - y)$. (1)

w: $\hat y = w^\top h \Rightarrow \dfrac{\partial L}{\partial w} = \delta\, h$. (2)

Backprop to z: $\dfrac{\partial L}{\partial h} = \delta\, w$, and $\dfrac{\partial L}{\partial z} = \dfrac{\partial L}{\partial h}\odot (1-h^2) = \delta\,(w \odot (1-h^2))$. Call this $g_z \in \mathbb{R}^2$. (2)

b: $z = Wx+b \Rightarrow \dfrac{\partial L}{\partial b} = g_z$. (0.5)

W: $\dfrac{\partial L}{\partial W} = g_z\, x^\top$ (outer product, $2\times2$). (0.5)

(b) Numerical (6 marks)

$z = Wx = \begin{bmatrix}0.5(1)+(-0.5)(-1)\\ 1(1)+0(-1)\end{bmatrix} = \begin{bmatrix}1\\1\end{bmatrix}$. (1)

$h = \tanh(1) = 0.761594$ in both components → $h=[0.761594,\,0.761594]$. (1)

$\hat y = w^\top h = 1(0.761594)+2(0.761594)=3\times0.761594 = 2.284782$. (2)

$L = \tfrac12(\hat y - 0)^2 = \tfrac12(2.284782)^2 = 2.610135$. (1)

$\delta = \hat y - y = 2.284782$; $\dfrac{\partial L}{\partial w} = \delta h = 2.284782 \times [0.761594,0.761594] = [1.740094,\,1.740094]$. (1)

(c) Code + concepts (6 marks)

import torch
x = torch.tensor([1., -1.])
W = torch.tensor([[0.5,-0.5],[1.0,0.0]], requires_grad=True)
b = torch.zeros(2, requires_grad=True)
w = torch.tensor([1., 2.], requires_grad=True)
y = torch.tensor(0.)
 
z = W @ x + b
h = torch.tanh(z)
yhat = w @ h
L = 0.5*(yhat - y)**2
L.backward()
 
manual = (yhat.detach()-y)*h.detach()
assert torch.allclose(w.grad, manual, atol=1e-5)

(3 marks for a correct runnable graph + backward + assert)

Concepts (1 each):

• retain_graph=True: keeps the intermediate buffers after backward so a second .backward() can traverse the same graph; without it the graph is freed and a second backward raises a RuntimeError.
• .detach(): returns a tensor sharing storage but cut from the graph (no grad flows through), used here so the manual gradient computation doesn't create graph nodes.
• A second .backward() accumulates into .grad (adds), not overwrites — hence zero_grad() in training loops.

Question 2

(a) Model + loop (10 marks)

import torch, torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader
 
class LogReg(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(2, 1)
    def forward(self, x):
        return self.lin(x).squeeze(-1)   # logits, shape (B,)
 
ds = TensorDataset(X, y.float())
dl = DataLoader(ds, batch_size=32, shuffle=True)
model = LogReg()
crit = nn.BCEWithLogitsLoss()
opt = torch.optim.SGD(model.parameters(), lr=0.1)
 
for epoch in range(num_epochs):
    for xb, yb in dl:
        opt.zero_grad()
        logits = model(xb)
        loss = crit(logits, yb)
        loss.backward()
        opt.step()

Marks: class (2), Dataset/DataLoader (2), correct loss/optimizer (2), correct ordering zero_grad→forward→loss→backward→step (4).

(b) Numerical stability (4 marks)

BCEWithLogitsLoss combines sigmoid + BCE and uses the log-sum-exp trick, avoiding computing $\sigma(z)$ explicitly (which saturates to 0 or 1 causing $\log(0) = -\infty$). (2)

Identity: BCE loss for target $t$, logit $z$:

The $|z|$ form keeps the exponent $\le 0$, so $e^{-|z|}\in(0,1]$ never overflows. (2)

(c) Checkpointing (6 marks)

# save each epoch
torch.save({
    'epoch': epoch,
    'model_state': model.state_dict(),
    'opt_state': opt.state_dict(),
    'best_val': best_val,
}, 'ckpt.pt')
 
# resume
ckpt = torch.load('ckpt.pt', map_location=device)
model.load_state_dict(ckpt['model_state'])
opt.load_state_dict(ckpt['opt_state'])
start_epoch = ckpt['epoch'] + 1
best_val = ckpt['best_val']

Marks: save dict (2), resume with optimizer restore (2), start epoch/best_val (1).

Bug (1): torch.save(model) pickles the whole object including the class/module path; loading later breaks if the code/class definition moves or refactors, and it ties the checkpoint to that exact source layout. state_dict() stores only tensors → portable.

Question 3

(a) Gradient scaler / underflow (6 marks)

FP16 has a narrow dynamic range; small gradients round to 0 (underflow), stalling learning. (2)

The scaler multiplies the loss by $S$ before backward. By linearity, every gradient is multiplied by $S$: $\nabla(S\cdot L) = S\,\nabla L$. Before the optimizer step, gradients are unscaled (divided by $S$). (2)

Condition: a true gradient component $g$ survives if $S|g| \ge 6.10\times10^{-5}$, i.e.

Choosing large $S$ pushes the underflow threshold down; dynamic scaling raises $S$ until inf/nan appears then backs off. (2)

(b) AMP step (8 marks)

scaler = torch.cuda.amp.GradScaler()
model = model.to(device)             # move model ONCE, outside loop
for xb, yb in dl:
    xb, yb = xb.to(device, non_blocking=True), yb.to(device)
    opt.zero_grad()
    with torch.cuda.amp.autocast():
        logits = model(xb)
        loss = crit(logits, yb)
    scaler.scale(loss).backward()    # scaled backward
    scaler.step(opt)                 # unscales then steps (skips if inf/nan)
    scaler.update()                  # adjusts scale factor

Marks: autocast context (2), scale→backward (1), step (1), update (1), model .to once outside (1), batch .to inside loop (1), throughput reasoning (1): moving the model each iteration wastes host↔device copies; batches must move inside the loop because each is different, ideally overlapped with non_blocking=True + pinned memory.

(c) Distributed proof (6 marks)

Global batch = $KB$ samples. Per-GPU $k$ mean loss over its $B$ samples:

All-reduce averages:

Hence averaged gradient = gradient of the mean loss over the whole global batch. (4)

Learning-rate scaling: with global batch $KB$ vs. base batch $B$, scale LR by $K$ (linear scaling rule) with warmup. (1)

Breaks: at very large batch the gradient noise scale shrinks, curvature/generalization effects dominate, and linear scaling causes divergence or worse generalization; also breaks early in training (needs warmup). (1)

[
  {"claim":"yhat = 3*tanh(1) ≈ 2.284782 and L ≈ 2.610135",
   "code":"import sympy as sp\nh=sp.tanh(1)\nyhat=3*h\nL=sp.Rational(1,2)*yhat**2\nya=float(yhat); La=float(L)\nresult = abs(ya-2.284782)<1e-4 and abs(La-2.610135)<1e-4"},
  {"claim":"dL/dw = delta*h = [1.740094, 1.740094]",
   "code":"import sympy as sp\nh=float(sp.tanh(1))\nyhat=3*h\ndelta=yhat-0\ngrad=delta*h\nresult = abs(grad-1.740094)<1e-4"},
  {"claim":"Averaged per-GPU mean gradients equal global mean gradient (K=2,B=2 symbolic check)",
   "code":"import sympy as sp\na,b,c,d=sp.symbols('a b c d')\nB=2;K=2\ng1=(a+b)/B\ng2=(c+d)/B\nbar=(g1+g2)/K\nglobal_g=(a+b+c+d)/(K*B)\nresult = sp.simplify(bar-global_g)==0"},
  {"claim":"Underflow threshold: |g| >= 6.10e-5/S, e.g. S=1024 gives threshold ~5.96e-8",
   "code":"thr=6.10e-5/1024\nresult = abs(thr-5.957e-8)<1e-10"}
]