📈 ScaleTrace — Empirical Neural Scaling Laws Analysis

Investigates how transformer language models scale with model size, training data, and compute by training a family of small transformers and fitting Chinchilla-style neural scaling laws. The project explores how parameter count and dataset size influence language modeling loss, and uses empirical results to estimate compute-optimal scaling behaviour.

Why this exists

Large language models exhibit predictable scaling behaviour: as model size, dataset size, and compute increase together, validation loss follows empirical power-law relationships. Understanding these scaling laws is fundamental to designing compute-efficient models and estimating the trade-offs between training cost and performance.

Rather than reproducing billion-parameter experiments, this project builds a small-scale research benchmark to investigate whether the same qualitative scaling behaviour emerges using consumer hardware and a realistic source-code corpus.

A methodological note before the results

An earlier version of this project trained transformers on randomly generated byte sequences. Because the corpus contained essentially no learnable statistical structure, increasing model size produced little meaningful improvement and the recovered scaling exponents were unstable.

This version replaces the synthetic corpus with a dataset constructed from twelve Python standard library modules. Unlike random data, Python source code contains syntax, repeated patterns, and long-range dependencies that transformers can learn, producing substantially more realistic scaling behaviour and allowing meaningful empirical scaling-law fitting.

Headline Results

Metric	Result
Models trained	20
Parameter range	35K → 19M
Dataset sizes	10K → 500K tokens
Best validation loss	2.26
Scaling-law fit (full grid)	R² = 0.876
Scaling-law fit (high-data regime)	R² = 0.862

Recovered scaling exponents (full dataset):

• α = 0.379 (model-size exponent)
• β = 0.795 (dataset-size exponent)

Recovered scaling exponents (high-data regime, D ≥ 200K):

• α = 0.616
• β = 0.421

Reference (Chinchilla, Hoffmann et al., 2022):

• α ≈ 0.34
• β ≈ 0.28

Although the recovered exponents differ from the original Chinchilla study, the experiments reproduce the expected qualitative scaling behaviour while operating in a much smaller experimental regime.

---

Experimental Setup

• Framework: Pytorch
• Hardware: NVIDIA T4 GPU
• Corpus: Twelve Python standard library modules
• Models: Five transformer sizes (35K–19M parameters)
• Training tokens: 10K, 50K, 200K and 500K
• Evaluation: Validation loss
• Curve fitting: Nonlinear least-squares estimation of the Chinchilla scaling law

---

What the experiments show

Across every model size, increasing the training dataset consistently reduced validation loss.

Small models benefited immediately from additional data, while larger models only became effective once sufficient training tokens were available. Extremely large models trained on very small datasets performed poorly, illustrating the data-limited regime predicted by neural scaling laws.

These experiments demonstrate that both model capacity and dataset size must increase together to achieve efficient language-model scaling.

---

Repository Structure

scaletrace/
│
├── notebook/
│   └── ScaleTrace.ipynb
│
├── figures/
│   ├── scaling_heatmap.png
│   ├── portfolio_stat_card.png
│   ├── loss_vs_token_budget.png
│   ├── loss_vs_model_size.png
│   └── compute_frontier.png
    └── compute_efficiency.png
│
├── results/
│   ├── scaletrace_results.csv
│   ├── results.json
│   └── report.md
│
├── paper/
│   └──paper.md
│
├── README.md
└── requirements.txt

---

Reproducing

Run the notebook from top to bottom on a CUDA-enabled GPU.

The notebook will:

1. Build the Python source-code corpus.
2. Train transformers across multiple model sizes.
3. Evaluate validation loss.
4. Fit Chinchilla-style scaling laws.
5. Generate publication-ready figures and tables.
6. Save all outputs to the results/ and figures/ directories.

---

Key Findings

• Validation loss consistently decreased as dataset size increased.
• Larger models only became beneficial when trained on sufficiently large datasets.
• Small-scale experiments reproduced the qualitative behaviour predicted by neural scaling laws.
• Fitted scaling-law models explained approximately 88% of the observed variance across experiments.

---

Future Work

• Train larger transformer families.
• Evaluate additional programming-language corpora.
• Compare scaling behaviour across natural language and source code.
• Investigate compute-optimal training under fixed FLOP budgets.
• Extend scaling analysis to instruction-tuned language models.