Architectural Trade-offs in Small Language Models Under Compute Constraints

TL;DR

This work addresses the problem of designing small language models under strict compute constraints and investigates how architectural choices interact with training budgets. It systematically compares Linear, MLP, Self-attention, and Transformer architectures across Tiny Shakespeare and word-level PTB/WikiText-2, using test NLL and approximate FLOPs as key metrics, with a clear formula $FLOPs \approx 2 \times (#parameters) \times (#training tokens)$ to quantify compute. The main contributions show that attention-based models achieve the best NLL per FLOP at small scale, that increasing context or depth without matching optimization budget can degrade performance, and that rotary positional embeddings (RoPE) do not reliably help in tiny regimes. The findings provide regime-aware guidance for designing compute-efficient small LMs and caution against blindly adopting large-model techniques in constrained settings.

Abstract

We present a systematic empirical study of small language models under strict compute constraints, analyzing how architectural choices and training budget interact to determine performance. Starting from a linear next-token predictor, we progressively introduce nonlinearities, self-attention, and multi-layer transformer architectures, evaluating each on character-level modeling of Tiny Shakespeare and word-level modeling of Penn Treebank (PTB) and WikiText-2. We compare models using test negative log-likelihood (NLL), parameter count, and approximate training FLOPs to characterize accuracy-efficiency trade-offs. Our results show that attention-based models dominate MLPs in per-FLOP efficiency even at small scale, while increasing depth or context without sufficient optimization can degrade performance. We further examine rotary positional embeddings (RoPE), finding that architectural techniques successful in large language models do not necessarily transfer to small-model regimes.

Figures (10)

• Figure 1: Training and validation NLL for Tiny Shakespeare across architectures. More expressive models achieve lower loss while still avoiding major overfitting in this setting.
• Figure 2: Linear model: test NLL vs. context length. Longer contexts increase dimensionality and parameters, but the model remains linear and appears under-trained at fixed optimization budget.
• Figure 3: MLP: test NLL vs. hidden dimension. Wider layers help, but marginal gains shrink as parameter count grows.
• Figure 4: Self-attention: test NLL vs. number of heads. Four heads yields the best performance, consistent with learning multiple attention subspaces.
• Figure 5: Transformer: test NLL vs. layers. Depth $L=3$ performs best here; $L=4$ likely suffers from insufficient optimization steps at fixed budget.