📄 Architectural Trade-offs in Small Language Models Under Compute Constraints

Summary

The paper asks which architectural choices pay off when a language model must stay tiny and its training budget is fixed. The author builds a ladder of next-token predictors, adding one ingredient at each step: a linear model over concatenated embeddings, a two-layer ReLU MLP, a single self-attention block, and then a stack of transformer blocks. Each is trained on character-level Tiny Shakespeare. The best design is then reused for word-level modeling on PTB and WikiText-2. Models are compared on test NLL, parameter count, and a rough FLOP estimate. Attention-based models give the lowest loss and use far less compute than the MLP. Adding depth or context without more training can make results worse. Swapping learned positional embeddings for RoPE gives slightly worse results at this scale. The broader point is that habits from large-model practice should be tested in the regime where they will actually be used, not adopted by default.

Key points

• Tiny Shakespeare test NLL drops along the ladder: linear 3.05, MLP 2.32, single self-attention block 2.13, 3-layer transformer 2.01.
• Efficiency differs sharply between families:
  • The MLP has about 4.29M parameters and uses roughly 1.6 × 10^14 FLOPs.
  • The 3-layer transformer has about 431K parameters, uses roughly 2.2 × 10^13 FLOPs, and still scores better.
• Sweeps on a fixed budget give these best settings:
  • 4 attention heads beat 1 or 2.
  • 3 layers beat 2 or 4.
  • Hidden width 256 is the efficient choice, because going to 512 brings only small gains.
  • The linear model gets worse as context grows from 32 to 128.
• On word-level data the 3-layer transformer (d=128, 4 heads, FF 256, T=64) reaches:
  • PTB: test NLL 6.19 with a 9,999-word vocabulary.
  • WikiText-2: test NLL 6.89 with a 33,277-word vocabulary.
  • Validation loss reaches its lowest point near epoch 2 and climbs afterward, so the early-stopped checkpoint is used.
• Training compute is estimated as FLOPs ≈ 2 × parameters × training tokens. On Tiny Shakespeare, RoPE scores slightly below learned positional embeddings.

Interesting ideas

• I1. The study adds one component at a time, from a linear baseline upward, so each gain can be credited to a single change; this kind of ablation ladder is a clean way to justify complexity in any pipeline.
• I2. Using a crude cost estimate (2 × params × tokens) as a second axis turns "which model is better" into "which is better per unit of compute," a comparison that is fair across very different model families.
• I3. The RoPE experiment changes only the positional encoding and keeps everything else equal, which shows how to test whether a borrowed best practice really helps in your own regime.
• I4. Deeper or wider-context models lost to smaller ones because their extra parameters got too few updates, so model capacity should be sized to the optimization budget rather than to ambition.
• I5. The author checks that generated samples become more coherent as NLL falls, using reading the output as a sanity check on the numeric metric.

For Mova

How these ideas could serve Mova, our small open language model — its files, stages and tests:

• I1 (learn): Evaluation could rebuild each learned component as a ladder (rules only, perceptron, then transformer), so every added layer has to show a measured gain over the simpler rung.
• I2 (tree): The perceptron parser and gparse could be plotted as LAS against estimated training and inference cost, to decide where the GPU parser is actually worth running.
• I3 (tree): Any trick imported from large parsers into gparse should go through a one-variable swap on the same treebanks before it is adopted, since the paper shows such tricks can fail at small scale.
• I4 (learn): When training on a limited pool of silver data agreed by the witness parsers, model size and depth should be swept against the available data and steps instead of defaulting to the biggest configuration.
• I5 (back): Generation and the text → graph → text compressor could pair their scores with a small fixed set of inspected outputs, to catch cases where the metric improves but the text does not.