I built a language model from scratch, then made it cheap enough to actually serve.

If you understand a Transformer decoder: multi-head self-attention, the feed-forward block, residual connections, layer normalization, causal masking, cross-entropy on next-token prediction; you have every prerequisite for this page. Everything past that point is developed here from first principles.

That is a deliberate choice about what "advanced" means. Almost nothing in a modern language model is conceptually harder than attention. What happened between 2017 and now is that a few dozen specific, individually-comprehensible modifications accumulated, each fixing a concrete problem with the 2017 design. RoPE fixes how position is encoded. RMSNorm removes a computation nobody could show was needed. Grouped-query attention fixes a memory bottleneck that only appears at inference time. Muon fixes an assumption Adam makes about weight matrices. Every one of them is a paragraph of motivation and a formula.

So this page is structured as a sequence of deltas from what you already know. Each one states the problem with the baseline, the fix, the mathematics, and what it measured in this repository. The two arcs are independent enough to read separately:

Here is the model you already know, written the way this repository writes it. A decoder-only Transformer maps a sequence of token IDs to a probability distribution over the next token at every position simultaneously. Tokens go in, an embedding table turns each into a vector, a stack of identical blocks refines those vectors, and a final linear layer projects each vector to a score for every token in the vocabulary.

Notation used throughout

V, T, d, L, H

Vocabulary size, sequence length, model width (d_model), number of layers, number of attention heads. d_head = d / H.

N, D

Parameter count and training-token count. The Chinchilla heuristic says a compute-optimal run sets D ≈ 20N.

π_θ(y | x)

The model treated as a policy: the probability it assigns to response y given prompt x. Used from the alignment section onward.

p, q

In Arc 2, p is always the big/target/teacher distribution and q the small/draft/student one. Keeping this consistent is what makes the distillation and speculative-decoding sections rhyme.

The four problems

Everything in Arc 1 is a response to one of the first three problems; everything in Arc 2 responds to the fourth. Naming them now makes the rest of the page a set of answers rather than a list of techniques.

Before any neural network runs, text has to become integers. The naive options are both bad. Characters give a tiny vocabulary but very long sequences, and attention costs O(T²), so long sequences are expensive. Words give short sequences but an unbounded vocabulary: you will always meet a word you did not see in training, and every such word collapses to a single <unk> token that destroys information.

Byte-pair encoding is the compromise, and it is genuinely a compression algorithm before it is an NLP one. Start with an alphabet. Repeatedly find the most frequent adjacent pair of symbols in the corpus, and add a new symbol that means "that pair". Stop at the target vocabulary size. Frequent sequences,  the, ing,  Congratulations , become single tokens; rare ones stay decomposed into pieces.

Byte-level is the part that matters

The critical design choice is what the starting alphabet is. If it is Unicode characters, the alphabet is effectively unbounded (there are ~150,000 assigned code points) and you are back to needing an <unk>. If it is the 256 possible byte values, the alphabet is exactly 256 and closed. Every possible string: any language, any emoji, any corrupted file, any lone surrogate, is some sequence of bytes, so it is always encodable.

This is not a soft guarantee, it is a structural one, and the repository asserts it with property-based tests that let Hypothesis generate adversarial strings. There is no <unk> token in this tokenizer because there cannot be one.

Pre-tokenization: the constraint that keeps merges sane

Left alone, BPE will happily learn a single token for " of the", because that string is extremely frequent. That is wasteful, it spends vocabulary slots on phrase-level accidents instead of morphology, and it makes the tokenization of a word depend on what precedes it. The fix is a pre-tokenization regex that first splits text into chunks (roughly: words with their leading space, runs of digits, runs of punctuation), and merges are only ever allowed inside a chunk.

// GPT-2's pattern, implemented in src/nanoscale/tokenizer/bpe.py
's|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+

This has a pleasant efficiency consequence. Because merges never cross chunk boundaries, training does not need to rescan the corpus: it can count chunk types with their frequencies, and maintain a pair &rarr, which chunks contain it inverted index so that applying a merge only touches the chunks that actually contain that pair. Encoding gets the same benefit as a cache keyed on the chunk, which is why this pure-Python tokenizer still encodes at ~14 MB/s.

What was measured

Two vocabularies are trained and committed. The nano tier uses 1,024 tokens (761 learned merges) over the synthetic corpus; the micro tier uses 16,384 tokens (16,121 merges) over TinyStories. Correctness against a reference implementation is checked as three separate claims rather than one vague "parity" number:

Longest learned tokens in the 16k TinyStories vocabulary:  enthusiastically,  Congratulations,  understandingly,  recommendations, whole words, which is what a healthy merge schedule produces.

This is the model the repository builds. Read it against the 2017 diagram above: the skeleton is identical: attention sublayer, feed-forward sublayer, residuals, stacked N times , and every difference is annotated below.

2.1  Pre-norm residuals; the trainability fix

In the 2017 block the order is x → sublayer → add → normalize. That means the residual stream itself is renormalized at every layer. Whatever the embedding wrote into it gets rescaled L times before reaching the output, and the gradient flowing backward is rescaled L times too. In practice deep post-norm stacks only train with carefully tuned learning-rate warmup; without it they diverge.

The pre-norm arrangement is x → normalize → sublayer → add. The normalization applies to the input of the branch, not to the stream. The stream itself is a pure sum:

The gradient reaching layer 0 therefore contains an undamped copy of the gradient at the output. This is why modern models are trained at depth without a warmup knife-edge, and it is the reason the residual stream is often described as a "highway" that each block reads from and writes to.

2.2  RMSNorm, delete the part nobody could justify

LayerNorm subtracts the mean, divides by the standard deviation, then applies a learned gain and bias. RMSNorm keeps only the scaling:

Zhang & Sennrich showed empirically that the mean-subtraction ("re-centering") contributes nothing measurable, and that the useful work is entirely in the rescaling. Dropping it removes one full pass over the vector, one reduction, and the bias parameters. At a 5M-parameter scale that saving is invisible; at 70B it is not, and every current open model uses RMSNorm.

2.3  RoPE, position as a rotation

This is the delta most worth understanding properly, because the trick is genuinely elegant and the implementations look opaque until you see it.

The problem with adding a learned position vector to the embedding is that it makes position absolute. The model must learn, separately for every pair of absolute offsets, that positions 5 and 7 are two apart in the same way that 105 and 107 are. And the table has a fixed number of rows, so the model cannot run on longer sequences than it was trained on.

Rotary position embedding attacks this from the observation that attention only ever consumes q and k through their dot product. So instead of adding anything, rotate them. Split each head's d_head-dimensional query and key vector into d_head/2 consecutive pairs, and treat each pair as a point in a 2-D plane. For a token at position m, rotate pair j by angle m θj:

Now the payoff. Rotation matrices compose by adding angles, and a rotation is orthogonal, so RmTRn = Rn−m. Therefore:

No parameters were added. No table can be run off the end of. And the model learns "two tokens apart" once, not once per absolute offset. The multi-frequency construction : fast rotations in the early dimensions, slow ones in the late dimensions, is the same idea as sinusoidal encodings: it gives the dot product a basis in which both short-range and long-range offsets are distinguishable.

2.4  Grouped-query attention: a fix for a problem you only see at inference

During training, attention reads the whole sequence at once and nothing is stored between steps. During generation the picture inverts. Generating token t+1 needs the keys and values of all t previous tokens, and recomputing them every step would make generation O(T²) per token. So they are cached, the KV cache , and the cache is the thing that dominates inference memory.

Put a real model in that formula. Llama-2-7B has L=32, 32 heads of d_head=128, and standard multi-head attention, so n_kv=32. At a 4,096-token context in fp16, for one single sequence:

Two gigabytes of cache for one user, against 13 GB of weights that are shared by everyone. Batch sixteen users and the cache is 34 GB; the cache, not the model, is what decides how many people your GPU can serve. That is the bottleneck GQA exists to attack.

The observation is that n_kv does not have to equal the number of query heads. Keep all H query heads; they are where attention's expressiveness lives , but let groups of them share one key/value head. With n_kv = H you have classic multi-head attention; with n_kv = 1 you have multi-query attention, which is maximally cheap and measurably hurts quality; GQA is the middle that Ainslie et al. showed keeps almost all the quality.

2.5  QK-norm, keeping the softmax out of saturation

Attention scores are q·k / √d_head. Nothing bounds the norms of q and k, and during training they tend to grow. Once the scores get large, the softmax saturates: it becomes nearly one-hot, its gradient goes to nearly zero, and that head stops learning. It is a slow, silent failure that shows up as a loss spike or a dead head, not as an error.

The fix is one line: RMS-normalize q and k per-head before the dot product, with a learned gain. The score's magnitude is now controlled by that learned gain rather than by whatever the projections happened to drift to.

2.6  SwiGLU: a gate on the feed-forward network

The 2017 feed-forward network is ReLU(xW₁)W₂: project up, threshold, project down. SwiGLU splits the up-projection into two and uses one half to gate the other:

The gate is the point. In the ReLU version the only decision available per unit is "pass or zero". In the gated version one branch computes a value and the other computes, from the same input, a smooth multiplicative weight on that value, so a unit can be attenuated by any amount rather than only switched off, and the attenuation is input-dependent. Because three matrices replace two, the hidden width is conventionally scaled by 2/3 (here to 8/3 · d rounded to a multiple of 64) so parameter count stays comparable.

Shazeer's paper introducing this is unusually honest: it offers no theoretical justification and ends "we attribute their success, as all else, to divine benevolence." It won on benchmarks and the field adopted it. The ungated ReLU² alternative from the modded-nanoGPT speedrun is implemented here too, as a togglable ablation, see §05 for what it actually measured.

2.7  Zero-init output projections and an untied head

Both sublayers' final projections (o_proj and W_down) are initialized to exactly zero. At step 0 every block therefore outputs zero, the residual stream passes the embedding through untouched, and the model computes the identity. Training starts from a provably well-conditioned point and each block "switches itself on" as its gradient grows.

This has a testable consequence that the repository uses as a canary. If the model outputs nothing but the embedding, and the embedding is small and roughly isotropic, the logits are near-uniform, so the initial loss must be exactly the entropy of a uniform distribution over the vocabulary:

If a refactor breaks the initialization, the very first logged number changes, and it changes against a value derived from theory rather than from a previous run.

Untying the LM head means the output projection is a separate matrix from the input embedding table rather than its transpose. Tying saves V×d parameters and was standard when vocabularies were small relative to model width. It also forces one matrix to serve two different jobs, "what does this token mean" and "how likely is this token here" , and at modern vocabulary sizes the parameter saving is not worth the constraint.

2.8  Two optional heads

Both are implemented and off by default, because their effect is what the ablation harness exists to measure rather than assert.

tanh logit soft-cap

logits ← c · tanh(logits / c). A smooth ceiling on logit magnitude (Gemma-2 uses c = 30), which prevents a single overconfident logit from dominating the softmax. Unlike clipping it is differentiable everywhere, so gradients keep flowing through capped logits.

multi-token prediction (MTP)

Extra small heads that predict token t+2, t+3, … alongside the main t+1 head, as an auxiliary loss. The training argument is that forcing the representation to carry information about tokens further ahead is a denser learning signal than next-token alone. The inference argument is more interesting and is why this belongs in a project with an Arc 2: those heads are already a draft model, which makes them reusable for self-speculation with no separate network: the Medusa idea in §09.

The size ladder

One architecture; only depth, width, context and token budget change. Parameter counts are asserted as exact constants in the test suite, against both the analytic formula and the built nn.Module, so an accidental architecture change fails loudly rather than quietly shifting every number downstream.

Both totals are reported because "parameter count" is ambiguous without saying whether embeddings and an untied head are included. At micro the embedding table plus head is 16.8M of the 40.4M, 42% of the model is the vocabulary interface, which is what small models look like.

Start from what Adam does, because Muon is best understood as a critique of it. Adam keeps two running averages per parameter, the gradient m and its square v, and steps by m / (√v + ε). Dividing by the root-mean-square gradient makes the step size roughly scale-invariant per coordinate: a weight whose gradients are consistently tiny still moves, and one whose gradients are huge does not blow up. AdamW's refinement is to apply weight decay directly to the weight rather than through the gradient, so the decay is not itself rescaled by v.

Notice the word per coordinate. Adam's update rule is applied elementwise, which means it makes an implicit assumption: that the parameters are exchangeable scalars whose only relevant property is their own gradient history. For a bias vector or a norm gain, that is true. For a 512×1536 weight matrix, it is false.

What a weight matrix actually does

A weight matrix is a linear map. Its meaningful structure is its singular value decomposition, W = UΣVT: a set of input directions (columns of V), a set of output directions (columns of U), and a gain σi for each pairing. What the layer does is amplify some directions and suppress others.

Raw gradients of such matrices are typically very badly conditioned: their top singular value dominates the rest by orders of magnitude. So a gradient step, and an Adam step too, moves the matrix mostly along one direction. The other directions get a step that is smaller by that same factor, which is to say they barely move at all. You are doing a low-rank update while paying for a full one.

The fix: orthogonalize the update

Muon's proposal is to take the momentum buffer B, decompose it, and throw away the singular values, keeping only the directions:

UVT is the orthogonal polar factor of B: the closest orthogonal matrix in Frobenius norm. Every direction the gradient identified now receives an equally-sized step. The update still says which way to move each direction; that information is entirely in U and V, it just refuses to let the dominant direction monopolize the step budget.

Newton–Schulz: orthogonalizing without an SVD

Computing an SVD every step for every weight matrix would be far more expensive than the forward pass. Newton–Schulz iteration approximates the polar factor using only matrix multiplications, which is exactly the operation accelerators are built for. Normalize B so its spectral norm is at most 1, then iterate five times:

This is a quintic polynomial applied to the singular values: because X and XTX share singular vectors, each σ independently follows σ ← aσ + bσ³ + cσ⁵. Those specific coefficients are not the textbook Newton–Schulz ones; Keller Jordan tuned them so the polynomial pushes small singular values up aggressively in very few iterations, deliberately trading a slight overshoot near σ=1 for speed. Five steps of three matmuls is cheap enough to run every step on every hidden matrix.

The router, which parameters get Muon

Muon applies to 2-D hidden weight matrices only. Everything else: norm gains, biases, the token embedding, and the LM head, goes to AdamW. The embedding and head are two-dimensional, so this needs a reason rather than a shape check.

The reason is that they are not linear maps between hidden representations; they are lookup tables indexed by token. Their rows are near-independent, each row's gradient is sparse (only the tokens in this batch get one), and there is no meaningful spectrum to balance. Orthogonalizing an embedding gradient mixes updates across unrelated tokens. This split is the standard one, and the repository's router asserts it by parameter name.

Muon  : 4,423,680 params  [blocks.*.attn.{q,k,v,o}_proj, blocks.*.mlp.{gate,up,down}]
AdamW :   528,384 params  [embed_tokens, lm_head, *.attn_norm, *.q_norm, *.k_norm, final_norm]

Cautious weight decay

A February 2026 refinement, implemented behind a flag: apply weight decay only to coordinates where the decay direction agrees in sign with the optimizer's update. Where they disagree, decay is actively fighting the gradient, and masking it out there empirically improves both Muon and AdamW. The repository reports the fraction of coordinates masked, so the flag's effect is observable rather than assumed.

What it measured

2.1× fewer steps to reach the same loss, and 3.5% better final loss at an equal step budget. Both arms share one seed, one data order, one schedule; they differ only in the optimizer. This reproduces, at 5M parameters, the finding that the AdamW→Muon swap was the single biggest lever in the modded-nanoGPT speedrun.

Every architectural claim in §02 came from a paper that demonstrated it at a scale 100–1000× larger than this. The honest thing to do is measure whether it reproduces here, and report the answer either way. All arms share one seed, one data order, one schedule and a fixed step budget; they differ only in the named field.

All three toggles land within 1.4% of the default on final loss, which is below the 2% threshold this project is willing to call a result from a single seed. Reported as "no measurable difference", not as a win for whichever arm happened to come out ahead.

The one signal that is above noise is in the steps-to-target column: without QK-norm the model needs 75 steps to reach the loss the default reaches in 50. The two converge to the same place at different rates, which is exactly what a stability aid should do, and exactly what a final-loss-only comparison would have missed.

6.1  SFT, and why the mask matters

A pretrained model is a text continuer. Ask it a question and a reasonable continuation is another question, because that is what documents containing questions look like. Supervised fine-tuning fixes this by training on (prompt, response) pairs formatted with role tokens.

The important detail is the loss mask. Cross-entropy is computed on the response tokens only; prompt tokens contribute nothing. Training on the prompt teaches the model to generate prompts, which is not the task, and it dilutes the gradient signal with tokens the model will always be given rather than asked to produce. The repository verifies this with a gradient check: perturbing a prompt-position label must leave the loss exactly unchanged.

6.2  DPO, derived

Classical RLHF has three stages: collect human preferences, fit a reward model, then optimize the language model against that reward with PPO. It works and it is a lot of machinery: a second network to train, and an RL loop that is notoriously sensitive. Direct Preference Optimization removes both. The derivation is short and worth following, because the conclusion is genuinely surprising.

Step 1. Write down the RLHF objective: maximize reward, but stay close to the reference policy so the model does not destroy its language ability chasing the reward.

Step 2. This constrained problem has a known closed-form solution. It is a standard result; the optimum is the reference policy reweighted by exponentiated reward:

Step 3. Now invert it. Instead of solving for the policy given the reward, solve for the reward given the policy:

This says something strange: any language model is already a reward model, read off as its log-ratio against a reference, up to a term that depends only on the prompt.

Step 4. Preferences are modelled with Bradley–Terry: the probability that yw is preferred over yl is the sigmoid of their reward difference. Substitute step 3 into it, and because both responses share the same prompt, β log Z(x) appears in both terms and cancels exactly:

Step 5. Maximize the likelihood of the observed preferences under that expression. That is the DPO loss, and it is a plain supervised objective: no reward model, no sampling, no RL:

6.3  The failure this project hit

Look at that loss again. It depends only on the difference of two log-ratios. There are two ways to make a difference bigger: raise the first term, or lower the second. Nothing in the objective prefers the first.

My DPO run looked healthy by every standard diagnostic. The implicit reward margin climbed steadily to +7.88. Preference accuracy hit 100%. Then the aligned model went head-to-head against the SFT model it started from and lost 0–7–33.

Logging the absolute log-probabilities rather than only their difference showed exactly what happened:

The margin rose because the rejected response's likelihood fell off a cliff, and the chosen response's likelihood fell too, just less far. The model was not learning to produce good answers; it was learning to produce nothing in particular, slightly less unlike good answers than bad ones. A rising margin was concealing a model getting worse at everything.

Length exploitation

The other known DPO pathology: log π(y) is a sum over tokens, each negative, so longer responses have systematically lower log-probability. The objective can be partly satisfied by changing length rather than content, and DPO-tuned models are widely observed to become verbose. The repository measures this directly: not by comparing the lengths of training responses, which only describes the data, but by measuring what the model actually generates for held-out prompts, before and after. Here mean generated length went 26.6 → 24.3 tokens: no exploitation at this scale, on this data.

6.4  SimPO and GRPO

SimPO, reference-free

Drops πref entirely, normalizes by length, and adds an explicit target margin γ: −log σ( (β/|yw|) log π(yw) − (β/|yl|) log π(yl) − γ ). Length normalization attacks the verbosity bias at its root, and removing the reference model halves the memory needed for preference tuning; you no longer hold a frozen copy of the network. It produced the largest reward margin here (+16.2) and still lost its head-to-head, which is a useful reminder that the margin is not the objective you care about.

GRPO; RL with a verifiable reward

For tasks where correctness can be checked by a program (arithmetic here), you do not need a learned reward model at all. Sample a group of G completions per prompt, score each 0/1, and use the group itself as the baseline: Ai = (ri − mean(r)) / std(r). No value network, which is what makes this cheap enough to matter. The KL penalty to the reference uses the k3 estimator, E[ρ − log ρ − 1] with ρ = πref/πθ, unbiased, and unlike the naive −log ρ it is guaranteed non-negative, so it cannot pay the policy to move away from the reference.

The judge used for head-to-head comparisons is programmatic and stated in the code: on-topic overlap with the prompt, absence of degenerate repetition, and whether the model emitted <eot> rather than running to the token cap. Those are the properties the preference labels encode, so this measures did the model learn the labels: not is the model good. It is deliberately length-insensitive, so a model that learned to game length gains nothing from it.

Generation has two phases with completely different cost characteristics, and conflating them is the most common way to reason wrongly about inference.

Prefill processes the whole prompt at once. Every weight matrix is loaded from memory once and multiplied against T token vectors. Lots of arithmetic per byte loaded ; this phase is compute-bound, and it parallelizes beautifully.

Decode generates one token at a time, and each token depends on the one before it, so it cannot be parallelized across time. To produce a single token you must read every weight in the model from memory and use each one for exactly one multiply-accumulate.

Now compare that to what the hardware wants. An A100 delivers roughly 312 TFLOP/s of bf16 arithmetic against 1.5 TB/s of memory bandwidth, a ratio of about 200 FLOPs per byte. Decode offers it one. The arithmetic units sit idle by a factor of two hundred while the memory bus does all the work.

Make that concrete. A 7B model in fp16 is 14 GB. Reading 14 GB at 1.5 TB/s takes 9.3 ms. So the hard ceiling is about 107 tokens/second, for a single sequence, no matter how fast the GPU can multiply. You are not compute-limited. You are waiting for memory.

A trained model's output distribution carries far more information than the one-hot label it was trained on. Told the next token is  cat, the label says nothing else. The teacher's distribution says  cat 0.6,  dog 0.25,  kitten 0.1, encoding that dogs are more like cats than staplers are. Hinton called these the dark knowledge in the logits, and distillation is training a student on that distribution rather than on the labels.

7.1  Forward KL: the baseline, and its problem

Write out what KL(p ‖ q) = Σ p log(p/q) punishes. The sum is weighted by p, so it only cares about places the teacher puts mass. And wherever the teacher has mass and the student has none, log(p/q) → ∞. The student is therefore forced to put probability everywhere the teacher does, including the long low-confidence tail. This is mode-covering.

For a student with comparable capacity that is fine. For a student 17× smaller it is a trap: it cannot represent the teacher's full distribution, so being punished for every gap forces it to smear probability thinly across regions it cannot model well, including the empty space between modes, where the teacher has nothing but the student now does.

7.2  Reverse KL, permission to give up

Swap the arguments. KL(q ‖ p) = Σ q log(q/p) is weighted by the student, so it only cares about places the student puts mass, and the infinite penalty now lands where the student puts mass and the teacher does not. The student may ignore as much of the teacher as it likes, provided everything it does claim is something the teacher agrees with. This is mode-seeking, and for an undersized student it is the right trade.

7.3  Why reverse KL needs a policy gradient

There is a catch that makes reverse KL harder to implement than forward KL. The expectation Σ q log(q/p) is taken over the student's own distribution, and the student is what you are differentiating. You cannot just backpropagate through a sample. This is the standard score-function situation, and MiniLLM's contribution is the machinery to make it stable:

• On-policy sampling. Trajectories come from the student, which also fixes exposure bias; the mismatch where a model trained only on teacher-generated prefixes has never seen its own mistakes and cannot recover from them.
• Reward-to-go decomposition. Credit each token with the reward accumulated from that point forward, rather than a single sequence-level score, which drastically cuts variance.
• Single-step regularization. A per-token term added directly, sidestepping the policy-gradient estimator for the part that does not need it.

7.4  What was measured

Reverse KL has worse perplexity and better generations, and that is the MiniLLM finding, reproduced. It is not a contradiction. Perplexity rewards a model for spreading probability over everything in the evaluation set, which is precisely the mode-covering behaviour reverse KL is designed to avoid. A mode-seeking student scores worse on a coverage metric and degenerates less when it actually generates.

The sharpest number in that table is the repetition rate: 0.0000 for the reverse-KL student against 0.0383 for its own teacher. Distilling on-policy against a teacher's distribution is not the same as copying its outputs, and the student can end up better-behaved than the model it learned from. Judging these three by perplexity alone would have ranked them backwards, which is exactly why the repetition diagnostic is reported beside it.

Compression: 4,952,064 → 279,168 parameters (17.7× overall, 29.9× on non-embedding parameters). Teacher and student share a tokenizer by necessity, so the embedding table and LM head are the same width in both and their cost is irreducible; the non-embedding figure is the one that describes the actual depth and width reduction.

The mechanism is simple. Pick a scale, round each weight to the nearest point on a 2k-point grid, store the integer:

One scale for a whole matrix is far too coarse: a single outlier weight stretches the grid and everything else collapses onto a few levels. So scales are computed per group of columns (64 here). That costs storage, and the repository accounts for it honestly:

8.1  Why round-to-nearest is not the best you can do

RTN rounds every weight independently to its closest grid point. That minimizes ‖W − Ŵ‖, the error in the weights. But nobody cares about the weights. A layer exists to compute WX, and what matters is:

Those are different objectives. A weight multiplying an input channel that is almost always zero can be rounded badly at no cost; a weight multiplying a high-variance channel cannot. And crucially, the rounding errors interact, if one weight rounds up, another can round down to partially cancel it in the output.

8.2  GPTQ: spend the error where it does least damage

Expand that objective and the second-order structure appears immediately. The Hessian of the reconstruction error with respect to the weights of one row is:

GPTQ (following Optimal Brain Quantization) processes columns one at a time. Quantize column q, measure the error you just introduced, and then adjust all the not-yet-quantized columns to compensate for it, using the Hessian to work out which adjustment best cancels the error in the output:

The later columns absorb the earlier columns' mistakes. By the end, the accumulated output error is far smaller than independent rounding would give, even though the individual weights have moved further from their original values.

Three engineering details make this practical rather than theoretical, and all three are in the repository: a Cholesky factorization of H−1 for numerical stability (the naive iterative inverse-update becomes indefinite and produces NaNs on real layers), lazy block updates that batch the compensation over 128 columns instead of applying it one at a time, and activation ordering: quantizing the highest-activation-magnitude columns first, while the largest pool of un-quantized columns is still available to absorb their error.

8.3  AWQ: protect the channels that matter

A different attack on the same problem. AWQ observes that a small fraction of weight channels are salient: they multiply consistently large activations, so their error dominates the output. Rather than compensating after the fact, scale those channels up by s before quantizing and divide the corresponding activation by s:

No mixed precision, no second pass over the data at inference: the scaling is folded into the weights, and s is found by a small grid search on calibration statistics.

8.4  What was measured

KV-cache quantization

The weights are not the only thing occupying memory bandwidth. As the cost model showed, the KV cache can exceed the model itself at long context and large batch. It is quantized with the same machinery: per-token, per-head scales, with independent bit-widths for keys and values, since keys are empirically more sensitive than values. The repository reports the memory saved and the perplexity cost side by side, at several context lengths, because the trade only becomes worthwhile as context grows.

Recall the cost model: a decode step reads every weight to produce one token, and the arithmetic units are 99.5% idle. That idle capacity is the opportunity. A forward pass over γ+1 positions costs almost exactly the same wall-clock time as a pass over one position, because both are bounded by the same weight read.

So: let a small draft model generate γ candidate tokens autoregressively (cheap; it is a small model). Then feed the whole candidate sequence to the target model in one forward pass. Because the sequence is known, the target can score every position in parallel, exactly as in training. Now decide which candidates to keep.

9.1  The accept rule, and why it is exactly lossless

The naive approach, accept a draft token if the target also likes it, would bias the output. The correct rule is a form of rejection sampling:

The claim is that the token you end up with is distributed exactly as p. Here is the whole proof; it is four lines.

The normalizing constant of the residual is exactly the rejection probability, so the two cancel. Notice what the proof does not assume: nothing about q at all. The draft can be a badly-trained model, a random one, or an adversarial one, and the output distribution is still exactly the target's. A bad draft costs you speed, never correctness.

9.2  What actually costs what

The expected number of tokens per target pass, for acceptance rate α and draft length γ, is the mean of a truncated geometric:

This explains the shape of the results below. Larger γ has diminishing returns ; each extra proposal is only reached if all before it were accepted, while the draft cost grows linearly in γ. There is an optimum, and it moves with α.

9.3  Medusa: speculation without a second model

Maintaining a separate draft model is operationally annoying: two checkpoints, two tokenizers to keep in sync, extra memory. Medusa's alternative is to bolt extra prediction heads onto the target itself: exactly the multi-token-prediction heads from §02.8. Head k predicts token t+k, each head proposes its top few candidates, and the cross product forms a tree of candidate continuations.

Verifying a tree in one pass needs a custom attention mask: every node must attend to its ancestors and to nothing else, so that sibling branches cannot see each other. The nodes are flattened into a sequence and the mask encodes ancestry.

9.4  What was measured

3× fewer target forward passes at γ=6, and the two levers compose; you can speculate against a quantized target and keep essentially the same acceptance rate (0.367 vs 0.362), because 4-bit quantization barely moves the target's distribution.

The serving layer is deliberately small: streaming generation with a KV cache, sampling transforms (temperature, top-k, top-p, repetition penalty) shared with the speculative path so the two cannot drift apart, stop-sequence handling, and a prefill/decode timing breakdown.

The variants table

Every Arc 2 lever, measured on the same prompt, the same token count, and the same machine, with the composed variants included.

How to read this table

The weight column is the representation size, not the tensor size

The 4-bit rows are simulated in fp32, because there is no int4 CPU kernel to run them on. Reading the footprint off the tensors would therefore report a 4-bit model as 32-bit. The figure shown is computed from the effective bit-width including the stored scales; the size the model would occupy with a real kernel.

Speculative rows include the draft's weights and cache

Speculation is not free in memory; it trades space for target passes. Reporting only the target's footprint would hide the trade.

Decode throughput does not generalize; the pass counts do

This is a CPU measurement at 5M parameters, where a forward pass is bound by Python dispatch rather than weight loading. Both compression levers therefore show their memory win and no speed win, and speculation shows a speed loss. The cost model explains exactly why, and the target-pass reduction, which is hardware-independent, is the number to carry away.

The tiny benchmark is saturated

28 hand-written questions across subject–verb agreement, coreference, schema completion and arithmetic, scored by likelihood against a stated 50% chance line. The base model gets 100%, so it is a degradation detector for Arc 2, not a quality ladder. An unchanged score means compression did not break the capabilities it probes, and nothing stronger.

The natural-language model

The micro tier, 40.4M parameters, trained on TinyStories: a corpus of short stories written with the vocabulary of a small child, built by Eldan & Li precisely to test whether sub-100M models can be coherent. That capability match is the point: the same 40M parameters given general web text learn the shape of English and produce fluent nonsense, because the distribution is far wider than the capacity.

Sampled from that checkpoint at 3.07M tokens, temperature 0.7:

Read that carefully, because the failures are as informative as the successes. The grammar is correct throughout. Tense, agreement and pronoun reference hold across sentences. Dialogue is punctuated properly. Narrative structure: setup, complication, "But then, something unexpected happened", is clearly learned. What breaks is entity tracking: the bowl becomes paper becomes a box, and Mia's mother calls her Sue. That is exactly the signature of an under-trained model: syntax is learned from local statistics and arrives early; maintaining a consistent world state across a paragraph needs far more data.

The CI/teaching tier

The quantization frontier

Everything, in one table

Debugging a machine-learning pipeline is unusual in that most failures do not raise. The model trains, the loss goes down, and the number at the end is simply worse than it should be, which is indistinguishable from "this is how well the method works" unless you have something to check against. Every entry below was found by having something to check against.

Bugs in the code

Hypotheses that were wrong: the tests, not the code

This section exists because the project's original evaluation was its weakest part: token perplexity plus a 28-question quiz the base model saturated. Neither can support a claim. What replaced them is four metrics that each answer a question the others cannot.

Bits per byte; the only number that compares across models

Perplexity is per token, and token boundaries are a property of the tokenizer, not of the text. A model with a 50,257-token vocabulary needs fewer tokens for a sentence than one with 16,384, so each of its tokens carries more information and its per-token perplexity looks worse: even if it predicts the underlying text better. Comparing perplexity across tokenizers is meaningless, and it is done constantly.

Both models score the same held-out strings, each with its own tokenizer, and the results become directly comparable. This is what the Pile and Chinchilla papers use, and it is what makes the comparison below a measurement rather than an artefact.

Minimal pairs: a benchmark with headroom

The replacement for the saturated quiz is a BLiMP-style forced choice (Warstadt et al. 2020). Each item is two sentences differing in exactly one place, one grammatical and one not; the model scores both and is correct if it prefers the grammatical one.

856 items, generated from templates over a lexicon rather than hand-written, so the suite scales and is reproducible from a seed. Difficulty is controllable: an agreement item with an intervening attractor noun is far harder than one without, because the model has to track the true subject across a nearer, disagreeing noun.

Calibration, is the confidence earned?

A model can be accurate and badly miscalibrated: 99% confident on predictions that are right 70% of the time. That matters for a small model specifically, because overconfidence is the mechanism behind degenerate generation, a model certain of a wrong continuation cannot recover from it. Expected calibration error bins predictions by confidence and measures the gap between confidence and accuracy in each bin.

Diversity: catching what likelihood rewards

Likelihood metrics reward a model for putting probability where the evaluation set is, which is exactly satisfied by a model that produces the same safe thing every time. Distinct-n and self-BLEU detect that; perplexity structurally cannot.

The single-seed ablations reported earlier used a flat rule: any gap below 2% was "no measurable difference". That was an honest guess and not a measurement, with one run per arm there is no estimate of run-to-run variance, so there is nothing to compare a gap against. Running five seeds per arm makes the comparison a real two-sample test.

The evaluation reports one number per phenomenon at the end of training. That is a photograph of a process, and the process is the more interesting object, especially given the finished model sits at 94% on simple agreement and 44% on agreement across an attractor. Was it ever better? Is 44% a floor, or a trajectory caught mid-climb?

So: one training run, the full 856-item suite evaluated every 250 steps, sixteen probes spanning 23M tokens. The cost is one training run plus a few seconds per probe.

This is the sharpest argument for the project's small scale being a feature rather than an apology. The whole experiment, sixteen full benchmark evaluations interleaved with a complete training run, took under two hours on a laptop. The same study on a model where one training run costs thousands of GPU-hours is a research programme; here it is an afternoon, and it can be re-run with a different architecture, optimizer or data mixture by changing one config field.

Everything above is a demonstration that the pipeline is correct. This section is the argument that the artifact it produces is useful, and the argument rests on an identity that is usually left as theory.

Feed a model's next-token distribution to an arithmetic coder and its cross-entropy stops being a benchmark number and becomes a byte count you can check on disk. The repository does exactly this, and every figure below survived a byte-identical round trip.

The same forward pass is an anomaly detector

Per-token surprisal is the compressor's cost function, un-summed. A line the model finds expensive to encode is a line unlike its training distribution; an unsupervised anomaly score needing no labels, no rules, and one threshold expressed as a percentile of normal traffic.

The economics, which decide whether any of this is worth doing

The model must be stored alongside the archive, so a neural codec only pays back above a break-even volume. This is the calculation that is usually skipped, and it is the one that determines whether the idea is real.

Each input byte costs 0.708 bits with the model against 2.458 with the best classical coder, saving 0.219 bytes per byte. Below break-even, use xz. Above it, the neural codec wins and keeps winning, and this is exactly where §08's quantization stops being an exercise and starts paying rent, because shrinking the model shrinks the break-even point by 7×.

Where this fits

What this project does not show

• These are 5-seed results at 5–40M parameters. Five seeds is enough to measure run-to-run variance and enough to detect a large effect; it is not enough to detect a small one. The optimizer's mean-loss advantage is a real example: d = 1.24 is a large effect and p = 0.121 still fails to clear α=0.05, because n=5 with one divergent seed is underpowered. These remain directional confirmations of findings obtained at 100–1000× this scale.
• A lever that matters at scale can be invisible here. QK-norm exists because attention logits blow up over long runs at large width. 400 steps at d=256 never gets there. "This experiment cannot see the effect" and "the effect is not real" are different statements.
• The wall-clock inference numbers do not generalize at all. Every Arc 2 technique targets memory bandwidth, and there is no memory-bandwidth bottleneck at this scale on a CPU. This is stated everywhere the numbers appear.
• The minimal-pair suite is templated, not naturalistic. 856 items generated from a lexicon of about 50 words, so it measures whether specific grammatical contrasts are represented: not general linguistic competence, and certainly not reasoning. Its vocabulary overlaps the training distribution by construction. The older 28-question quiz is retained only as a compression-damage detector, where saturation is the desired property.
• There is no MMLU number here, and there should not be one. A 40M-parameter model trained on children's stories would score at chance on any knowledge benchmark. Reporting chance-level results on a famous benchmark to look thorough would be padding.
• The alignment data is synthetic. Preference pairs are generated programmatically, and the head-to-head judge is a program that checks the properties those pairs encode. It measures did the model learn the labels, not is the model good.

What the same code does with more compute

What would make this publishable

Being direct, because it matters for what to do next: a 40M model on TinyStories will not beat state of the art on anything anyone tracks, and the bits-per-byte win over GPT-2 is a specialization result, not a capability one. Neither is a contribution on its own.

What is closer to a contribution is the methodology, and two specific findings that fall out of it:

• The agreement-attraction result. 94% on simple agreement, 44%: below chance : with an attractor. That is a clean, quantified demonstration that a 40M model has learned linear recency rather than syntactic agreement, at a scale where the whole training run is three hours and fully reproducible.
• Muon's variance reduction. 176× lower run-to-run spread than AdamW (F=176.4, p=0.0002), while the mean-loss difference is not significant. Optimizer papers overwhelmingly report mean final loss; that this is the weaker signal, and stability the stronger one, is worth saying.
• The emergence trajectories. Sixteen full benchmark evaluations interleaved with one training run, showing an ordering that is not the order of apparent difficulty, and two phenomena that never move at all, negation at ρ = +0.00, p = 1.000. The methodological point is that this study costs two laptop-hours here and is a research programme at frontier scale, which is the strongest version of the "small scale is a feature" argument.
• The codec result. 11.3× lossless compression against xz's 3.3× on in-domain text, with the break-even economics worked out. This is the piece that makes the whole thing an artifact someone could deploy rather than only read about.

The realistic framing is a workshop paper or a reproducibility report: which small-scale replications of 2023–2026 LM findings hold, which do not, and what it costs to check , with every claim backed by a committed script and a significance test. That is a real contribution and an honest one. Claiming a state-of-the-art result would not be.

Nothing above requires a redesign. The size ladder is one architecture with different hyperparameters; the streaming data path already handles corpora far larger than disk; the checkpoint format is resumable across sessions. Scaling this is a matter of compute, which is exactly the claim the project set out to demonstrate.

git clone https://github.com/<user>/nanoscale-lm && cd nanoscale-lm
make install
make smoke     # tokenizer → pretrain → SFT → DPO → quantize → speculate

make smoke runs the entire pipeline on CPU and asserts a sanity metric at every stage, loss starts at ln(V), SFT reduces the masked loss, DPO raises the reward margin, 4-bit stays coherent, GPTQ does not lose to RTN, greedy speculation is token-identical to greedy autoregressive decoding. Measured end to end: 45 seconds.

Repository map

src/nanoscale/
├── config/       16 frozen pydantic models + the size-ladder presets
├── tokenizer/    byte-level BPE: train, encode, decode, chat template
├── model/        attention (GQA · RoPE · QK-norm · KV cache), norms, MLPs, MTP, numerics
├── optim/        Muon (Newton–Schulz), AdamW, cautious decay, parameter router
├── train/        data pipeline, schedules, checkpointing, the loop
├── align/        SFT, DPO / SimPO, GRPO-RLVR
├── distill/      forward-KL, SeqKD, on-policy reverse KL
├── quantize/     RTN, GPTQ, AWQ, KV-cache quantization
├── specdec/      accept rule, draft–target sampling, Medusa + tree attention
├── serve/        streaming generation, stop sequences, timing breakdown
├── eval/         perplexity with error bars, tiny benchmark, preference judge
└── bench/        ablation harness, throughput / latency / memory harness

tests/{unit,property,dynamics,e2e}/   563 tests
results/                              committed figures, tables, JSON: the source of truth
notebooks/                            two Colab notebooks, CPU and GPU paths
demo/                                 Gradio app: generate, compare, inspect, race

▶ Live demo All results Source on GitHub Notebooks

Apache-2.0. Every algorithm implemented from scratch in PyTorch; no high-level trainer library in src/nanoscale/. Primary references: Vaswani et al. 2017 (Transformer); Su et al. 2021 (RoPE); Zhang & Sennrich 2019 (RMSNorm); Shazeer 2020 (SwiGLU); Ainslie et al. 2023 (GQA); Jordan et al. 2024 (Muon); Rafailov et al. 2023 (DPO); Meng et al. 2024 (SimPO); Shao et al. 2024 (GRPO); Hinton et al. 2015 (KD); Kim & Rush 2016 (SeqKD); Gu et al. 2024 (MiniLLM); Frantar et al. 2023 (GPTQ); Lin et al. 2023 (AWQ); Leviathan et al. 2023 (speculative sampling); Cai et al. 2024 (Medusa); Hoffmann et al. 2022 (Chinchilla); Eldan & Li 2023 (TinyStories); Miller 2024 (evaluation error bars).