19 — Solomon: The Transformer, Built From the Paper
Solomon 1 of 6 · Solomon/solomon/model
Stack: PyTorch only — no HuggingFace, no transformers, no pretrained anything
Reference: The Annotated Transformer (Harvard NLP) — Vaswani et al. in pure PyTorch
Where this sits. The deepest "do you actually understand it" artifact in the tree. Every component is either verbatim from the reference or has its deviation justified in writing, component by component, in
PLAN.md.
The component map — written before the code
| Reference component | Status | Note |
|---|---|---|
attention(q,k,v,mask,dropout) | verbatim | scaled dot-product, masked_fill(mask==0, -1e9), softmax, dropout |
MultiHeadedAttention | verbatim | 4 linears, view/transpose to (b,h,t,d_k), self-attention only |
PositionwiseFeedForward | verbatim | w_2(dropout(relu(w_1(x)))) |
LayerNorm | verbatim | a_2, b_2, eps 1e-6 |
SublayerConnection | verbatim | x + dropout(sublayer(norm(x))) |
Embeddings | verbatim | lut(x) * sqrt(d_model) |
PositionalEncoding | verbatim | sinusoidal, div_term = exp(arange(0,d,2) · −ln(10000)/d) |
subsequent_mask | verbatim | triu(ones, k=1) == 0 |
Generator | verbatim | linear + log_softmax |
NoamOpt | verbatim | factor · d_model^−0.5 · min(step^−0.5, step·warmup^−1.5) |
LabelSmoothing | verbatim | KLDivLoss, confidence/(V−2), padding zeroed |
| Xavier init | verbatim | xavier_uniform_ on dim > 1 |
Encoder, EncoderLayer, src-attn | dropped | no encoder in a causal LM |
EncoderDecoder | adapted → SolomonLM | embed → N decoder layers (causal mask) → final LayerNorm → Generator |
Batch, run_epoch, SimpleLossCompute | adapted | LM batches: y = x shifted by 1; ntokens excludes pad |
greedy_decode | adapted | + beam, temperature/top-k/top-p — decode-time only, no architecture change |
The three deviations, each defended
- Decoder-only. Required by the mission: a generative LM doing next-token prediction has no source sequence to encode.
- Weight tying (embedding ↔ generator). Saves ~2.1M of ~7M params; on a small-corpus CPU model the regularization is worth more than the capacity (Press & Wolf 2017). The reference itself notes shared embeddings as an option.
- Pre-norm — not actually a deviation. The reference code already applies norm before each sublayer ("for code simplicity the norm is first"). Saying so, rather than claiming credit for it, is the point.
Configuration
vocab 8192 · d_model 256 · N = 6 layers · h = 8 heads · d_ff 1024 · seq 256 · dropout 0.1 · tied weights → 6.84M parameters.
Sized deliberately: it trains to a real checkpoint on 16 CPU cores in hours, and the loop fits the ~13 GB of RAM left over after the production services on the same host take theirs. The model was sized by the machine's spare capacity, and that constraint is documented rather than hidden.
Decoding
Greedy, beam, and temperature / top-k / top-p — all decode-time, no architectural change. KV-cached decode was proven exactly equal to full recompute, not assumed. That proof is the thing to mention: caching is where generation quietly goes wrong, and "exactly equal" is a testable claim.
The tests that prove understanding
- Causal-mask leak test — does information from position t+1 reach t?
- Positional-encoding formula test — the actual sinusoid, not a shape check.
- Overfit-one-batch — the canonical "is my training loop wired at all" test.
- Greedy determinism — same input, same output, every time.
Interview surface this opens
- Attention at the tensor level: shapes, masking, why
−1e9and not−inf - Why pre-norm trains more stably than post-norm
- Weight tying: the parameter saving vs the representational cost
- KV caching and how you prove equivalence rather than assert it
- Sinusoidal vs learned vs rotary positional encoding