Homework 2 — Transformers for NLP

This homework accompanies Session 05: Transformers for NLP. You will climb the ladder of language models — from a feedforward neural LM to a working decoder-only Transformer (a GPT) — building the core pieces from scratch in PyTorch, and generating simple, coherent short stories along the way.

📘 The assignment

Do the work in the notebook homework2.ipynb — that's the assignment.

---

Setup

# from this folder (final_lectures/hw2)
python -m venv .venv
source .venv/bin/activate        # Windows: .venv\Scripts\activate
pip install -r requirements.txt
jupyter notebook homework2.ipynb

A GPU is recommended for the Transformer in Part 4 (CPU works, but training takes ~10 minutes). No GPU? The free GPU on Google Colab is plenty — upload the notebook, hw2_tests.py, and data/input.txt (plus data/arabizi_data.jsonl if you do the Part 5 bonus).

Dataset. TinyStories (Eldan & Li, 2023) — short, simple children's stories written with a small vocabulary — ~10M characters bundled as data/input.txt, modelled at the word level: a token is a word or punctuation mark (with <nl> for line breaks). We keep the ~8,000 most frequent tokens (here that covers 99.9%+; almost nothing becomes <unk>). Because the language is so simple, even the small model in Part 4 learns to write coherent little stories. Your yardstick is cross-entropy loss and perplexity (= exp(loss)) — note word-level perplexities are much larger than character-level ones, so only compare models using the same tokenizer.

---

Structure

Part	Topic	What you do
Part 1	Where we left off	Build a tokenizer (with a nod to BPE), perplexity, and a reusable feed-forward MLP; train a provided feedforward LM (Bengio-style) and hit the wall of a fixed context window.
Part 2	Attention	Build softmax, then self-attention in small steps — scores and the head (causal mask provided) — and visualize the attention matrix.
Part 3	Building the Transformer	Positional embeddings and the decoder block (Add & Norm + residuals, reusing your MLP), stacked into a GPT.
Part 4	Training a Transformer	Train your GPT, beat the Part 1 baseline on perplexity, generate a story, and create a submission file to upload.
Part 5	Fine-tuning (bonus)	Fine-tune the real Qwen3-1.7B with HuggingFace + LoRA to chat in Arabizi, on a dataset we generated for you with GPT (bundled in data/arabizi_data.jsonl).

Part 5 is a bonus capstone that switches from building-from-scratch to the real-world pretrain → fine-tune workflow. It needs a GPU (Google Colab's free tier works) — but no API key: the Arabizi training data is generated ahead of time by the course staff and bundled in data/, so students just load it. It installs extra libraries (transformers, trl, peft, datasets, accelerate) from a cell inside the notebook — they are not in requirements.txt, so Parts 1–4 stay lightweight.

The folder is intentionally minimal — the notebook is self-contained (it imports its data/training/plotting helpers from hw2_tests.py, so there is no utils.py):

hw2/
├── homework2.ipynb           <- the assignment (work here)
├── hw2_tests.py              <- the only support file: data helpers, training/
│                                generation utilities, and the self-checks
├── requirements.txt          <- packages to install
├── reading/                  <- start here: visual explainers of attention & Transformers
├── data/
│   ├── input.txt             <- TinyStories, bundled (word-level corpus)
│   └── arabizi_data.jsonl    <- Part 5 bonus: GPT-generated Arabizi chat data
└── figures/                  <- (optional) saved plots

---

What you implement

You build the core ideas from basic tensor ops; you may use nn.Linear, nn.Embedding, nn.LayerNorm, and torch.optim as building blocks (you built their equivalents by hand in HW1). You do not use torch.softmax (in your own code), nn.MultiheadAttention, or nn.Transformer. Routine or wiring-only pieces (cross_entropy_loss, the train loop, FeedForwardLM, the causal-mask helper, and each Block's and Head's __init__) are provided so you focus on the ideas and finish in ~3–4 hours.

The pieces you implement (each with a PASS/FAIL self-check): Tokenizer (build the vocabulary), perplexity, softmax, FeedForward (the reusable MLP block), self-attention in two steps — attention_scores and Head.forward (the causal mask is provided) — then Block.forward (residual + norm) and GPT.forward (embeddings + stack).

---

Self-checking your work

Throughout the notebook you'll see cells like:

from hw2_tests import check_head
check_head(Head)          # pass YOUR class

These run lightweight unit tests (shapes, known values, causality of the attention head) and print PASS / FAIL with a hint. Run everything at once at the end of Part 3:

from hw2_tests import run_all
run_all(globals())

Passing the checks means the building blocks are correct. The real payoff is Part 4: a Transformer whose validation perplexity clearly beats the feedforward baseline, generating coherent little stories — learned purely from next-word prediction.

---

Submitting Part 4

1. Near the start of Part 4, put your full name in STUDENT_NAME and run the SUBMISSION SETUP cell before you train.
2. Train your Transformer.
3. Run the Submit your work cell. It writes hw2_certificate_<yourname>.pt (~9 MB).
4. Upload two files to the Google Form: your hw2_certificate_<yourname>.pt and your completed notebook (homework2.ipynb).

If you trained on Google Colab, download the certificate file first (Files panel → right-click → Download) before uploading.