import random
import matplotlib.pyplot as plt
import numpy as np
%load_ext autoreload
%autoreload 2
SEED = 42
random.seed(SEED)
rng = np.random.default_rng(SEED)
plt.rcParams["figure.dpi"] = 90
from ml_first_principles.llm_models import BPETokenizer, dpo_loss
Exercise 1: Hand Calculation of BPE¶
Corpus word frequencies: "hug" ×10, "pug" ×5, "hugged" ×3.
Task: Apply the word-level BPE algorithm from theory.md: split each word into characters plus the end-of-word marker </w>, count adjacent symbol pairs weighted by word frequency, and merge the most frequent pair. Derive the first 2 merges by hand, showing the pair counts of each round.
Solution¶
Initial symbol sequences:
| Word | Symbols | Frequency |
|---|---|---|
| hug | h u g </w> |
10 |
| pug | p u g </w> |
5 |
| hugged | h u g g e d </w> |
3 |
Round 1 — pair counts weighted by word frequency:
| Pair | Count |
|---|---|
(u, g) |
$10 + 5 + 3 = 18$ |
(g, </w>) |
$10 + 5 = 15$ |
(h, u) |
$10 + 3 = 13$ |
(p, u) |
$5$ |
(g, g), (g, e), (e, d), (d, </w>) |
$3$ each |
The unique maximum is (u, g). Merge 1: u, g $\rightarrow$ ug, giving h ug </w> ×10, p ug </w> ×5, h ug g e d </w> ×3.
Round 2 — recount on the merged corpus:
| Pair | Count |
|---|---|
(ug, </w>) |
$10 + 5 = 15$ |
(h, ug) |
$10 + 3 = 13$ |
(p, ug) |
$5$ |
(ug, g), (g, e), (e, d), (d, </w>) |
$3$ each |
The unique maximum is (ug, </w>). Merge 2: ug, </w> $\rightarrow$ ug</w>.
Result: merges $=$ [(u, g), (ug, </w>)], and "hug" now tokenizes as h, ug</w>.
corpus = "hug " * 10 + "pug " * 5 + "hugged " * 3
# Base symbols {h, u, g, p, e, d, </w>} -> 7, so target size 9 yields exactly 2 merges.
tokenizer = BPETokenizer(target_vocab_size=9)
tokenizer.fit(corpus)
print("Learned merges:", tokenizer.merges)
print("Encoded 'hug':", tokenizer.encode("hug"))
assert tokenizer.merges == [("u", "g"), ("ug", "</w>")]
assert tokenizer.encode("hug") == ["h", "ug</w>"]
print("Hand-derived merges confirmed by the package tokenizer.")
Learned merges: [('u', 'g'), ('ug', '</w>')]
Encoded 'hug': ['h', 'ug</w>']
Hand-derived merges confirmed by the package tokenizer.
Exercise 2: LoRA Parameter Savings¶
Task: Write a function that calculates the parameter savings of applying LoRA to a single linear layer compared to full fine-tuning.
Given $W_0 \in \mathbb{R}^{d_{in} \times d_{out}}$ and LoRA matrices with rank $r$.
Return the percentage reduction in parameters as a float between 0 and 100.
def lora_savings(d_in, d_out, r):
"""Return the percentage reduction in trainable parameters (0-100)."""
# YOUR CODE HERE
pass
lora_savings(4096, 4096, 8) # runs as a no-op until implemented
Solution¶
def lora_savings(d_in, d_out, r):
"""Return the percentage reduction in trainable parameters (0-100)."""
full_params = d_in * d_out
lora_params = (d_in * r) + (r * d_out)
return 100.0 * (1.0 - lora_params / full_params)
savings = lora_savings(4096, 4096, 8)
print(f"Savings: {savings:.4f}%")
assert np.isclose(savings, 99.609375, atol=1e-12)
assert np.isclose(lora_savings(100, 50, 10), 70.0, atol=1e-12)
Savings: 99.6094%
Exercise 3: Conceptual Analysis of Alignment¶
Question: Why is Direct Preference Optimization (DPO) highly preferred in modern LLM training over the traditional Reinforcement Learning from Human Feedback (RLHF) via PPO? What is the main mathematical assumption DPO makes to achieve this?
Answer: DPO is preferred because it is significantly more stable and easier to implement. RLHF requires training and maintaining a separate Reward Model, and then running a complex, hyperparameter-sensitive PPO loop that is notorious for diverging or experiencing mode collapse.
DPO bypasses the reward model entirely. It translates the RL objective directly into a standard cross-entropy-style classification loss (Negative Log Likelihood) over preference pairs. This means you only need to run standard supervised gradient descent, which scales predictably.
The main mathematical assumption DPO makes is the Bradley-Terry preference model. It assumes that the human probability of preferring response $y_w$ over $y_l$ is perfectly modeled by the sigmoid of the difference of their scalar rewards: $P(y_w \succ y_l) = \sigma(r(y_w) - r(y_l))$. By substituting the closed-form optimal RL policy into this Bradley-Terry formulation, DPO expresses the preference probability purely in terms of the policy model's log-likelihoods, eliminating the need for an explicit reward model $r$.
Numerical check. At zero implicit-reward margin the Bradley–Terry model must give $P(y_w \succ y_l) = \tfrac{1}{2}$, i.e. a DPO loss of exactly $\log 2$:
zeros = np.zeros(4)
loss_at_zero_margin = dpo_loss(zeros, zeros, zeros, zeros, beta=0.1)
print(f"DPO loss at zero margin: {loss_at_zero_margin:.10f} (log 2 = {np.log(2):.10f})")
assert np.isclose(loss_at_zero_margin, np.log(2), atol=1e-12)
DPO loss at zero margin: 0.6931471806 (log 2 = 0.6931471806)