The Optimal Architecture for Small Language Models

We trained 19 model configurations across 12 architecture families on 1 billion tokens each. The result? A surprising discovery about what really matters for small language models and a new architecture that's 3.8x faster with better factuality.

TL;DR

In our previous article, we found that 50% FinePDFs + 30% DCLM + 20% FineWeb-Edu is the optimal dataset mix for training GPT-2, achieving 38.50% average accuracy. But that used the standard 12-layer architecture.

What if we could do even better by changing the model itself?

We ran 19 experiments to find out:

• 7 GPT-2 variants with radically different depth-width ratios (4→64 layers)
• 12 architecture families including LLaMA3, Gemma3, Qwen3, MoE, diffusion models, and novel hybrids

Here's what surprised us:

The result: Dhara-70M, a diffusion model that sacrifices 1.33% accuracy for 3.8x throughput and superior factuality.

---

The Problem: What's the Optimal Architecture for Small Models?

Our previous work established that 50% FinePDFs + 30% DCLM + 20% FineWeb-Edu is optimal for training small models. With that dataset recipe fixed, we asked: Does model architecture matter as much as data composition?

The standard GPT-2 uses 12 layers with 768 hidden dimensions. But this was designed in 2019 for ~124M parameters. For a 70M model trained on 1B tokens, is this still optimal? And what about newer architectures like LLaMA, Gemma, MoE, or even diffusion language models?

We set out to systematically map the architecture design space.

---

Experimental Setup

To isolate the effect of architecture, we fixed everything except model design:

---

Part 1: The Depth-Width Trade-off

First, we explored how model shape affects performance by training 7 GPT-2 variants with the same ~70M parameters but radically different depth-width ratios:

Discovery #1: The Two-Tier Performance Pattern

Our first finding was completely unexpected. We expected a smooth trade-off curve—more layers for less width, or vice versa. Instead, we found a hard binary split:

The gap between tiers is substantial: 6+ percentage points separating them, while variance within each tier is only ~0.5%.

This bimodal distribution is notable: configurations either achieve the high tier (38%) or fall to the low tier (32%), with no intermediate performance levels observed.

Discovery #2: The Hidden Dimension Threshold

Why do some configurations succeed while others fail? We identified the critical factor: hidden_size >= 512.

The rule emerges: Models need either:

1. hidden_size >= 512, OR
2. Exactly 32 layers (the "Goldilocks" depth), OR
3. Extremely deep (64+ layers) to compensate

The 16L, 24L, and 48L configurations fall into a "dead zone" - their hidden dimensions are too narrow, and their depths aren't at the sweet spots that can compensate.

Discovery #3: 32 Layers is the Goldilocks Depth

With hidden=384, the 32-layer configuration achieves the best overall score (38.50%), slightly beating even the standard 12-layer design.

The 32-layer model wins on 5 out of 7 benchmarks, with particular strengths in:

• WinoGrande (+1.18%): Better pronoun resolution suggests deeper compositional reasoning
• MMLU (+1.66%): More layers help with academic knowledge retention

---

Part 2: Architecture Family Comparison

Armed with the optimal 32-layer depth, we compared 12 different architecture families:

Architectures Tested

Complete Benchmark Results

Discovery #4: Architecture Choice Has Minimal Impact at 70M Scale

Surprisingly, all 12 architecture families achieve similar benchmark accuracy:

• High tier (AR models): 32-33% average
• Low tier (Diffusion models): 31-32% average

The differences are within noise at this scale. Modern architectural improvements (RMSNorm, RoPE, GQA) are designed for 7B+ models and don't provide measurable benefits at 70M parameters.

Winner: LLaMA3-Canon (33.22%) slightly edges out GPT-2 (33.18%), but the difference is not statistically significant.

Discovery #5: dLLMs Trade Accuracy for 3.8x Throughput

The real differentiation comes from inference characteristics, not accuracy:

The trade-off is clear:

• -1.33% accuracy (31.85% vs 33.18% average)
• +3.8x throughput (183 vs 48 tok/s)
• +1.6x memory (bidirectional attention overhead)

Discovery #6: dLLMs Excel at Factuality

One of our most surprising findings: diffusion models achieve the highest TruthfulQA scores among all architectures tested.

Why might dLLMs excel at factuality? We hypothesize three contributing factors:

1. Bidirectional attention allows the model to consider full context when making predictions
2. Iterative refinement enables the model to "second-guess" its initial predictions across multiple denoising steps
3. Non-autoregressive generation may reduce the snowball effect where early hallucinations compound into larger errors

Discovery #7: Canon Layers Improve Factuality

The "Physics of Language Models" Canon layers (depthwise causal convolutions) show consistent benefits for factuality:

Canon layers add only 0.13% parameter overhead but provide meaningful TruthfulQA improvements.

Discovery #8: WSD Enables 10x Training Efficiency

We find that existing autoregressive models can be converted to diffusion with 10x less training using the Warmup-Stable-Decay (WSD) method from the LLaDA 2.0 paper, which progressively trains an AR model to handle diffusion objectives:

The efficiency gains are substantial:

WSD requires 10x fewer tokens and 10x less GPU time while achieving equivalent or better results.

Notably, WSD conversion not only matches from-scratch training but outperforms it on several benchmarks:

The AR initialization provides learned representations that benefit factual knowledge tasks, suggesting that WSD conversion preserves and potentially enhances knowledge from the source model.

---

The Result: Dhara-70M

Based on all our discoveries, we introduce Dhara-70M, available on Hugging Face.

Dhara-70M is created by taking the best autoregressive architecture (LLaMA3-Canon) and converting it to a diffusion model using the WSD method. This gives us the best of both worlds: the strong initialization from AR pretraining, plus the throughput and factuality benefits of diffusion.

Architecture

Benchmark Results

Inference Performance

---

Quick Start

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("codelion/dhara-70m")
model = AutoModelForCausalLM.from_pretrained(
    "codelion/dhara-70m",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
)

# Move to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# Generate text
prompt = "The future of artificial intelligence is"
inputs = tokenizer(prompt, return_tensors="pt").to(device)
outputs = model.generate(
    inputs.input_ids,
    max_new_tokens=50,
    temperature=0.1,
    top_p=0.5,
    top_k=5,
    repetition_penalty=1.8,
    do_sample=True,
    pad_token_id=0
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Example Output:

The future of artificial intelligence is a big challenge.
This world has the potential to improve, but this time we have no other than "theworld."
The next generation will be more exciting and its very much important for our society's
abilityto develop its

For high-throughput batch processing:

# Batch generation for maximum throughput
prompts = [
    "The future of artificial intelligence is",
    "The human brain is capable of",
    "Science has shown that",
    "Technology continues to evolve"
]

inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(device)
outputs = model.generate(
    inputs.input_ids,
    attention_mask=inputs.attention_mask,
    max_new_tokens=50,
    temperature=0.1,
    top_p=0.5,
    top_k=5,
    repetition_penalty=1.8,
    do_sample=True,
    pad_token_id=0
)

for i, output in enumerate(outputs):
    print(f"Output {i+1}: {tokenizer.decode(output, skip_special_tokens=True)}")

---

Key Takeaways

After training 19 model configurations across 12 architecture families, here are our core lessons:

1. Architecture matters less than you think at small scale: All modern architectures (LLaMA3, Qwen3, Gemma3) perform within ~1% of each other at 70M parameters.

2. Depth-width ratio matters more: The same architecture can yield 32% or 38% accuracy depending on layer/hidden ratio alone.

3. The two-tier phenomenon is real: Models don't degrade smoothly - they either work (38%) or don't (32%), with a hidden_size threshold of 512.

4. 32 layers is the Goldilocks depth: For 70M parameters, 32 layers with 384 hidden beats the standard 12-layer design by 0.35%.

5. Diffusion models excel at factuality: Despite lower average scores, dLLMs achieve the highest TruthfulQA scores (49.27%), suggesting reduced hallucination.

6. 3.8x throughput is achievable: Diffusion models offer dramatic throughput improvements for batch processing workloads.

7. WSD conversion is remarkably efficient: Converting AR to diffusion needs only 100M tokens (10x fewer than training from scratch).

8. Canon layers help factuality: Simple depthwise convolutions add 0.13% parameters but improve TruthfulQA by 1-2%.

For practitioners building small language models: start with the 50-30-20 data mix (from our previous work), use the 32-layer Goldilocks architecture, and consider diffusion for high-throughput applications where factuality matters.

---

Related Work

Our Previous Work

• The 1 Billion Token Challenge: Optimal Dataset Mixing - Finding the 50-30-20 dataset recipe
• GPT-2-70M - Our baseline GPT-2 model

Diffusion Language Models

• MDLM: Simple and Effective Masked Diffusion Language Models - The masked diffusion objective we use
• dLLM-2: Scaling Diffusion Language Models - Recent scaling study of diffusion LMs
• DREAM: Diffusion Rectification and Estimation-Adaptive Models - Training framework for diffusion models

Architecture References

• GPT-2 - Original GPT-2 architecture
• LLaMA - LLaMA architecture with RoPE and RMSNorm
• Physics of Language Models: Part 4.1 - Canon layers (depthwise causal convolutions)
• Titans - Memory-as-Context architecture

Training Efficiency

• LLaDA2.0: Scaling Up Diffusion Language Models to 100B - The WSD method for AR→Diffusion conversion

Datasets

• Pre-training Dataset Samples - 1B token dataset samples used in this work