Scaling & Efficient Architectures

Question 1

(a) [8] Lagrangian $\mathcal{L} = E + AN^{-\alpha} + BD^{-\beta} + \lambda(6ND - C)$. Stationarity:

• $\partial_N: -\alpha A N^{-\alpha-1} + 6\lambda D = 0$ → $\alpha A N^{-\alpha} = 6\lambda ND$ [2]
• $\partial_D: -\beta B D^{-\beta-1} + 6\lambda N = 0$ → $\beta B D^{-\beta} = 6\lambda ND$ [2]

Right-hand sides equal ($=6\lambda C$), so $\alpha A N^{-\alpha} = \beta B D^{-\beta}. \quad\checkmark \ [4]$ Why: both marginal loss reductions must balance the same shadow cost of compute.

(b) [8] From (a): $A N^{-\alpha} = \frac{\beta}{\alpha}B D^{-\beta}$ (∗). Using $D = C/(6N)$: $A N^{-\alpha} = \tfrac{\beta}{\alpha} B (6N/C)^{\beta} = \tfrac{\beta}{\alpha}B\,6^{\beta} C^{-\beta} N^{\beta}.$ So $N^{\alpha+\beta} \propto C^{\beta}$, giving $N_{\text{opt}} \propto C^{\beta/(\alpha+\beta)} \Rightarrow a = \frac{\beta}{\alpha+\beta}. \ [3]$ Symmetrically (substitute $N=C/(6D)$): $D_{\text{opt}} \propto C^{\alpha/(\alpha+\beta)} \Rightarrow b = \frac{\alpha}{\alpha+\beta}. \ [3]$ $a + b = \frac{\beta+\alpha}{\alpha+\beta} = 1$. [2] Why it must hold: $C = 6ND \Rightarrow \log C = \text{const} + \log N + \log D$; taking $N,D \propto C^{a},C^{b}$ forces $a+b=1$ for dimensional consistency.

(c) [6] $\alpha+\beta = 0.62$. $a = 0.28/0.62 = 0.4516 \approx 0.452$ [2] $b = 0.34/0.62 = 0.5484 \approx 0.548$ [2] They are close to $0.5$ each ⇒ Chinchilla guideline: scale $N$ and $D$ roughly equally with compute, giving an approximately fixed ~20 tokens per parameter rule of thumb. [2]

(d) [8] For $10\times$ compute: $N$ factor $= 10^{a} = 10^{0.452} = 2.831$ [3] $D$ factor $= 10^{b} = 10^{0.548} = 3.532$ [3] Check product $2.831 \times 3.532 = 9.999 \approx 10$ ✓ (must equal compute factor). [2]

(e) [6] The $C=6ND$ model assumes all $N$ parameters are active per token. In top-$k$ MoE only $k$ of $E$ experts fire, so active params $N_{\text{act}} \approx N_{\text{shared}} + \frac{k}{E}N_{\text{experts}} \ll N_{\text{total}}$. [3] Training/inference FLOPs scale with $N_{\text{act}}$, not $N_{\text{total}}$, so an MoE gets more "loss-reducing capacity" (total params) per unit compute — the compute-optimal frontier shifts, and one should use $N_{\text{act}}$ in $C=6N_{\text{act}}D$ while counting quality against $N_{\text{total}}$. [3]

Question 2

(a) [6] $L_{\text{aux}} = \alpha_{\text{aux}} E \sum_e f_e P_e$. Since it is (up to constant) a dot product of two probability vectors on the simplex, minimizing it drives mass away from correlating high-load and high-probability experts. The uniform point $f_e = P_e = 1/E$ gives $\sum_e f_e P_e = E\cdot\frac{1}{E^2} = 1/E$, so $L_{\text{aux}}^{\min} = \alpha_{\text{aux}}\cdot E\cdot\frac1E = \alpha_{\text{aux}}$. [4] The differentiable surrogate uses the soft $P_e$ (gradient path) with the hard $f_e$ (detached), so gradients push $P_e$ toward uniform when experts are overloaded. [2]

(b) [8] Let $f, P$ be probability vectors ($\sum f_e = \sum P_e = 1$). By Cauchy–Schwarz, $\Big(\sum_e f_e P_e\Big)\Big(\sum_e 1\Big) \geq \Big(\sum_e \sqrt{f_e P_e}\Big)^2$ is not tight enough; instead use the tighter route: [2] Consider $\sum_e (f_e - 1/E)(P_e - 1/E)$. Actually the clean bound: minimize $\sum f_e P_e$ subject to $f=P$ (as they converge under the balancing pressure). Then $\sum f_e^2 \geq (\sum f_e)^2/E = 1/E$ by Cauchy–Schwarz / power-mean, [4] with equality iff all $f_e$ equal $=1/E$. Since at balance $P_e=f_e$, $\sum f_e P_e = \sum f_e^2 \geq 1/E$, equality iff $f_e=P_e=1/E$. [2]

(c) [8] Capacity $C = \lceil (4096/8)\times 1.25\rceil = \lceil 512\times1.25\rceil = \lceil 640\rceil = 640$. [3] Expert with 900 tokens: drops $900 - 640 = 260$ tokens. [3] Raising $c_f$ reduces dropped tokens (better quality/less info loss) but increases per-expert buffer size → more compute & memory and padding waste (lower throughput). [2]

(d) [10] Pseudocode with shapes:

# x: [T, d]
logits = router(x)               # [T, E]
probs  = softmax(logits, axis=-1)# [T, E]
expert = argmax(probs, axis=-1)  # [T]   top-1 choice
gate   = max(probs, axis=-1)     # [T]   gating weight

# load-balancing stats
f = bincount(expert, E) / T      # [E]  hard fraction
P = mean(probs, axis=0)          # [E]  soft mass
L_aux = alpha_aux * E * sum(f * P)          # scalar

# capacity with dropping
cap = ceil((T/E) * c_f)
y = zeros([T, d])
for e in range(E):
    idx = where(expert == e)[:cap]          # keep first `cap`, drop rest
    out = experts[e](x[idx])                # [<=cap, d]
    y[idx] = gate[idx][:,None] * out        # scaled combine
return y, L_aux

Marks: routing+softmax [3], f/P & aux loss [3], capacity/drop logic [2], gated combine + shapes annotated [2].

Question 3

Part A

(a) [6] Unrolling: $h_1=\bar B x_1$, $h_2=\bar A\bar B x_1+\bar B x_2$, and by induction $h_t = \sum_{j=1}^t \bar A^{t-j}\bar B x_j$. Then $y_t=Ch_t=\sum_{j=1}^t C\bar A^{t-j}\bar B x_j$. [3] This is a convolution $y = \bar K * x$ with kernel $\bar K_k = C\bar A^{k}\bar B$. Training: compute the full kernel and use FFT/parallel convolution over the whole sequence at once (parallel). Inference: keep the $n$-dim state $h$ and update recurrently $O(1)$ per token (sequential). [3]

(b) [6] $\bar A = e^{\Delta A} = e^{-0.05} = 0.9512$. [2] $\bar B = (\bar A - 1)A^{-1}B = (0.9512-1)/(-0.5) = (-0.0488)/(-0.5)= 0.0975$. [2] Stability: $|\bar A| < 1$; here $0.9512 < 1$ ✓, so the impulse response decays and output is bounded. [2]

(c) [4] SSM: memory $O(n)$ (fixed state, independent of $L$), compute $O(n)$ per token → $O(1)$ in $L$. [2] Attention + KV-cache: memory $O(Ld)$ (cache grows with length), compute $O(Ld)$ per token. [2] SSM wins asymptotically as $L\to\infty$.

Part B

(d) [8] Number of accepted drafted tokens $=k$ if the first $k$ are accepted and the $(k{+}1)$-th rejected: $P(k)=\beta^k(1-\beta)$ for $k=0..\gamma-1$; all $\gamma$ accepted with prob $\beta^\gamma$. On any rejection the target supplies 1 corrected token; if all accepted, 1 bonus token. So tokens $= k+1$. [3] $\mathbb E=\sum_{k=0}^{\gamma-1}(k+1)\beta^k(1-\beta) + (\gamma+1)\beta^\gamma.$ This telescopes; equivalently the number of tokens is $1+\sum_{k=1}^{\gamma}\mathbf 1[\text{first }k\text{ accepted}]$ with $E=\sum_{