Scaling & Efficient Architectures

Q1 (12 marks)

(a) Minimize $L=E+A N^{-\alpha}+B D^{-\beta}$ s.t. $C=6ND$ fixed. (1) Substitute $D=C/(6N)$: $L(N)=E+A N^{-\alpha}+B(6N/C)^{\beta}.$ (1) Differentiate and set to 0: $\frac{dL}{dN}=-\alpha A N^{-\alpha-1}+\beta B (6/C)^\beta N^{\beta-1}=0.$ (2) So $\alpha A N^{-\alpha}=\beta B D^{-\beta}$, i.e. at the optimum $\boxed{\alpha \cdot \frac{A}{N^\star{}^{\alpha}} = \beta \cdot \frac{B}{D^\star{}^{\beta}}}.$ (2) (The two variable loss terms are balanced in proportion $\alpha:\beta$.)

(b) From the balance: $A N^{-\alpha}\propto B D^{-\beta}$ with $D=C/(6N)$ gives $N^{-\alpha}\propto (C/N)^{-\beta}=C^{-\beta}N^{\beta}$, so $N^{\alpha+\beta}\propto C^{\beta}$: (2) $N^\star\propto C^{\frac{\beta}{\alpha+\beta}},\qquad D^\star\propto C^{\frac{\alpha}{\alpha+\beta}}.$ (1) Thus $a=\frac{\beta}{\alpha+\beta},\ b=\frac{\alpha}{\alpha+\beta}$, and $a+b=\frac{\alpha+\beta}{\alpha+\beta}=1$. ✓ (1)

(c) With $\alpha=\beta=0.34$: $a=b=0.5$. (1) Meaning: parameters and tokens should be scaled equally with compute — roughly double data whenever you double model size (the Chinchilla correction to earlier param-heavy scaling). (1)

Q2 (12 marks)

(a) For token $x\in\mathbb{R}^d$, router weights $W_r\in\mathbb{R}^{E\times d}$: (1)

• logits $h=W_r x$ (1)
• gate probs $g=\mathrm{softmax}(h)$, $g_i=\frac{e^{h_i}}{\sum_j e^{h_j}}$ (1)
• top-$k$: keep set $\mathcal{T}=\text{TopK}(g)$; renormalize $\tilde g_i=g_i/\sum_{j\in\mathcal{T}}g_j$ (1)
• output $y=\sum_{i\in\mathcal{T}}\tilde g_i\, E_i(x)$ where $E_i$ is expert $i$'s FFN. (1)

(b) $\mathcal{L}_{aux}=\alpha\, E\sum_{i=1}^{E} f_i P_i$ (2) where $f_i=\frac{1}{|B|}\sum_{x}\mathbf{1}[\arg\max g(x)=i]$ (dispatch fraction, non-differentiable) and $P_i=\frac{1}{|B|}\sum_x g_i(x)$ (mean soft prob, differentiable). (1) Why both: $f_i$ carries the true load but has no gradient (it's an argmax count); $P_i$ is the differentiable surrogate. Multiplying $f_i P_i$ makes overloaded experts receive gradient pressure to lower their routing probability. Penalizing $f_i$ alone gives no gradient; penalizing $P_i$ alone ignores actual dispatch imbalance. (2)

(c) $\sum f_i P_i = 0.5(0.4)+0.25(0.3)+0.25(0.2)+0 = 0.2+0.075+0.05=0.325$. (1) $\mathcal{L}_{aux}=1\cdot4\cdot0.325=1.3$. (1)

Q3 (12 marks)

(a) KV-cache stores K and V for every layer, every token, every head: $\text{bytes}=2 \times B \times L \times S \times d \times (\text{bytes/elem}).$ (2) Factor 2 = keys + values; $d=n\cdot d_{head}$ so summing over heads = full hidden dim; fp16 = 2 bytes. (2) (So $\text{bytes}=2\cdot2\cdot B L S d = 4BLSd$ for fp16.)

(b) $=4 \times 1 \times 32 \times 8192 \times 4096$ bytes $=4\times32\times8192\times4096 = 4,294,967,296$ bytes $=4$ GB (using $1\text{GB}=2^{30}$). (3)

(c) Speculative decoding: a cheap draft model autoregressively proposes $k$ tokens; the target model then scores all $k$ in a single parallel forward pass. (2) Accept/reject rule: for proposed token with draft prob $q(x)$ and target prob $p(x)$, accept with probability $\min\!\left(1,\frac{p(x)}{q(x)}\right)$. (2) On the first rejection, resample from the normalized residual distribution $\frac{\max(0,p(x)-q(x))}{\sum_x\max(0,p-q)}$; this makes the marginal output distribution exactly the target's, while amortizing the target's cost over several tokens. (1)

Q4 (12 marks)

(a) Tensor parallelism splits the model within a layer — e.g. weight matrices sharded across the hidden/head dimension; requires an all-reduce per layer (high-bandwidth, intra-node). Pipeline parallelism splits across layers into stages; main cost is pipeline bubble/idle time and point-to-point activation transfers between stages. (4) (2 each)

(b) ZeRO-1 shards optimizer states; ZeRO-2 additionally shards gradients; ZeRO-3 additionally shards parameters — each rank holds only its slice and gathers others on demand. (4)

(c) In FSDP, each rank stores only a shard of a layer's params. Before the forward of that layer, an all-gather reconstructs the full weights; after use they are freed (re-gathered again in backward). This holds full weights only momentarily, cutting memory to $\sim1/N$, at the cost of extra all-gather communication each pass. (4)

Q5 (7 marks)

(a) Continuous: $h'(t)=Ah(t)+Bx(t),\ y(t)=Ch(t)+Dx(t)$. (1) Discretized (step $\Delta$): $h_t=\bar A h_{t-1}+\bar B x_t,\ y_t=C h_t$ with $\bar A=e^{\Delta A}$, $\bar B=(\Delta A)^{-1}(e^{\Delta A}-I)\Delta B$. (1) Because $A,B,C$ are time-invariant, unrolling gives $y_t=\sum_{j\ge0} C\bar A^{j}\bar B\, x_{t-j}$, a convolution with kernel (1) $\bar K=(C\bar B,\ C\bar A\bar B,\ C\bar A^2\bar B,\dots).$ (1)

(b) Mamba makes $B,C,\Delta$ input-dependent (selective), so the model can gate/forget based on content. (1.5) This makes the SSM time-varying, so $\bar A^j$ no longer factors into a fixed kernel — the convolution view breaks, and Mamba instead uses a hardware-aware parallel scan. (1.5)

Q6 (5 marks)

An emergent ability is a capability absent in smaller models that appears (often sharply) once scale crosses a threshold, not predictable by smooth extrapolation from small models. (2) Example: few-shot arithmetic / multi-step reasoning / instruction following appearing only past some parameter count. (1) Critique (Schaeffer et al., "mirage"): apparent sharp emergence can be an artifact of discontinuous/nonlinear metrics (e.g. exact-match accuracy that is 0 until the whole answer is right); switching to a smooth metric (token edit distance, log-likelihood) shows steady, predictable improvement — the ability was scaling continuously all along. (2)

[
  {"claim":"Chinchilla optimal exponents a=b=0.5 for alpha=beta=0.34, and a+b=1",
   "code":"alpha=Rational(34,100); beta=Rational(34,100); a=beta/(alpha+beta); b=alpha/(alpha+beta); result=(a==Rational(1,2)) and (b==Rational(1,2)) and (a+b==1)"},
  {"claim":"MoE aux loss = 1.3 for given f,P,E=4,alpha=1",
   "code":"f=[Rational(1,2),Rational(1,4),Rational(1,4),0]; P=[Rational(4,10),Rational(3,10),Rational(2,10),Rational(1,10)]; E=4; alpha=1; L=alpha*E*sum(fi*Pi for fi,Pi in zip(f,P)); result=(L==Rational(13,10))"},
  {"claim":"KV cache = 4 GB for L=32,d=4096,S=8192,B=1,fp16 (bytes=4*B*L*S*d)",
   "code":"B=1; L=32; S=8192; d=4096; by=4*B*L*S*d; result=(by==2**32) and (by/2**30==4)"}
]