Pretraining & Fine-Tuning LLMs

Question 1 (12)

(a) $h = W_0 x + \frac{\alpha}{r} B A x$. Trainable: $A$ and $B$ only; $W_0$ frozen. (2: forward 1, trainable 1)

(b) Params per LoRA pair $= r(d+k) = 16(4096+4096) = 16 \times 8192 = 131072$. Frozen params $= d\times k = 4096^2 = 16{,}777{,}216$. Percentage $= \frac{131072}{16777216}\times100 = 0.78125\%$. (4: formula 1, LoRA count 1, frozen count 1, % 1)

(c) $\frac{\alpha}{r} = \frac{32}{16} = 2$. $B$ is initialised to zero so that at the start $\Delta W = BA = 0$, meaning the fine-tuned model exactly equals the pretrained model at step 0 (stable start, no shock to activations). $A$ is Gaussian so gradients can flow and break symmetry; if both were zero no learning signal would propagate. (3: value 1, B=0 reason 1, A random reason 1)

(d) Since $\Delta W = \frac{\alpha}{r}BA$ is just another $d\times k$ matrix, merge it: $W' = W_0 + \frac{\alpha}{r}BA$. Then $h = W'x$ is a single matmul with no extra latency; the adapter is folded into the base weight. (3: merge idea 1, algebra 1, zero-overhead statement 1)

Question 2 (10)

(a) $\mathcal{L}_{CLM} = -\sum_{t=1}^{T} \log P_\theta(x_t \mid x_{<t})$ (autoregressive negative log-likelihood; each token conditioned only on preceding tokens). (3)

(b) $\mathcal{L}_{MLM} = -\sum_{i \in M} \log P_\theta(x_i \mid x_{\setminus M})$ where $M$ is the masked set. Recipe: choose ~15% of tokens; of those, 80% → [MASK], 10% → random token, 10% → unchanged. This reduces train/inference mismatch (no [MASK] at inference). (4: loss 2, recipe 2)

(c) Causal LM is trained with a triangular attention mask so position $t$ cannot attend to $t{+}1\dots T$ — needed to prevent label leakage during autoregression. T5 span corruption masks contiguous spans replaced by a single sentinel, so one forward pass predicts multiple tokens per example (more supervised signal per token/step) → higher sample efficiency. (3: causal 1.5, span 1.5)

Question 3 (12)

(a) FP16 = 2 bytes/param: $7\times10^9 \times 2 = 14\times10^9$ B $= 14$ GB. INT4 = 0.5 bytes/param: $7\times10^9 \times 0.5 = 3.5\times10^9$ B $= 3.5$ GB. (4: FP16 2, INT4 2)

(b) Per-parameter FP32 (4 bytes) accounting for full FT + Adam:

• weights (master) FP32: 4 B
• gradients FP32: 4 B
• Adam $m$: 4 B
• Adam $v$: 4 B Total = 16 B/param. $7\times10^9 \times 16 = 112\times10^9$ B $= 112$ GB. (Accept ~112 GB; if FP16 weights+grads + FP32 states counted as 2+2+4+4+4=16 with FP32 master copy, still ≈112 GB.) (5: byte accounting 3, arithmetic 2)

(c) (1) 4-bit NormalFloat (NF4) quantization of the frozen base weights — information-theoretically optimal for normally-distributed weights, cutting base memory ~4×. (2) Double quantization — quantizing the quantization constants themselves to save further memory; plus paged optimizers to handle memory spikes. LoRA adapters are trained in FP16/BF16 on top of the frozen NF4 base. (3: NF4 1.5, double-quant/paged 1.5)

Question 4 (10)

(a) $p_i(\tau) = \dfrac{\exp(z_i/\tau)}{\sum_{j} \exp(z_j/\tau)}$. $\mathcal{L} = \lambda \,\tau^2\, \mathrm{KL}\!\big(p^T(\tau)\,\Vert\,p^S(\tau)\big) + (1-\lambda)\,\mathrm{CE}(y, p^S(1))$. (4: softmax 2, combined loss 2)

(b) For KL with soft targets the well-known result: $\frac{\partial \mathcal{L}_{soft}}{\partial z^S_i} = \tau^2 \cdot \frac{1}{\tau}\big(p^S_i(\tau) - p^T_i(\tau)\big) = \tau\,\big(p^S_i - p^T_i\big).$ Because differentiating softmax w.r.t. logits brings a $1/\tau$ factor, and the leading $\tau^2$ cancels one power → net $\tau(p^S_i-p^T_i)$. The $\tau^2$ multiplier is included so that gradient magnitudes stay comparable to the hard-label CE term (whose gradients scale as $1$), keeping the two losses balanced as $\tau$ changes. (4: derivative 2, τ² rationale 2)

(c) High $\tau$ softens the distribution so relative probabilities of incorrect classes (the "dark knowledge" about class similarity) become non-negligible and are transferred to the student. (2)

Question 5 (8)

(a) Catastrophic forgetting: when a neural network trained sequentially on a new task overwrites weights important for previously learned tasks, causing sharp performance loss on the old task/distribution. (2)

(b) Any three (2 each = 6):

• PEFT (LoRA/adapters) — freeze base weights, so original knowledge in $W_0$ is preserved; only small modules change.
• Rehearsal / data mixing — mix general-instruction data into the medical fine-tuning set so gradients don't drift entirely.
• Lower LR + fewer epochs / early stopping — smaller weight movement reduces overwriting.
• Regularisation (EWC / L2-to-pretrained) — penalise moving weights important to old tasks.
• Knowledge distillation from the original model — keep student close to teacher's general behaviour. (≥1 must be PEFT-based.) (6)

Question 6 (8)

import torch
import torch.nn as nn
import math
 
class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r=8, alpha=16):
        super().__init__()
        self.base = base
        for p in self.base.parameters():   # freeze base
            p.requires_grad = False
        in_f, out_f = base.in_features, base.out_features
        self.r = r
        self.scaling = alpha / r
        # A: (r, in), B: (out, r)
        self.A = nn.Parameter(torch.randn(r, in_f) / math.sqrt(in_f))
        self.B = nn.Parameter(torch.zeros(out_f, r))   # zero init
 
    def forward(self, x):
        base_out = self.base(x)
        lora_out = (x @ self.A.t()) @ self.B.t()
        return base_out + self.scaling * lora_out

Mark scheme: freeze base 2; A random init 1; B zero init 1; correct shapes 2; forward with scaling 2. (8)

[
  {"claim":"LoRA trainable params r(d+k)=131072","code":"r,d,k=16,4096,4096; result = (r*(d+k)==131072)"},
  {"claim":"LoRA percentage of frozen ~0.78125%","code":"r,d,k=16,4096,4096; result = (Rational(r*(d+k),d*k)*100 == Rational(78125,100000))"},
  {"claim":"7B FP16 = 14GB and INT4 = 3.5GB","code":"n=7*10**9; result = (n*2==14*10**9 and n*Rational(1,2)==Rational(35,10)*10**9)"},
  {"claim":"Full FT Adam FP32 memory 16 bytes/param = 112GB for 7B","code":"n=7*10**9; result = (n*16==112*10**9)"},
  {"claim":"LoRA effective scaling alpha/r=2","code":"result = (Rational(32,16)==2)"}
]