The 2020 paper “Scaling Laws for Neural Language Models” by Kaplan et al. (OpenAI) is one of the most significant theoretical contributions to modern AI. It shifted the industry focus from “clever architectures” to “compute-optimal scaling.”
The paper found that language model performance (measured by cross-entropy loss) follows a power law relationship with three primary variables: the number of parameters ($N$), the size of the dataset ($D$), and the amount of compute ($C$).
1. The Core Empirical Findings
The researchers discovered that for a Transformer-based model, the test loss $L$ can be predicted using the following power-law formula:
$$L(N, D) \approx \left[ \left( \frac{N_c}{N} \right)^{\frac{\alpha_N}{\alpha_L}} + \left( \frac{D_c}{D} \right)^{\frac{\alpha_D}{\alpha_L}} \right]^{\alpha_L}$$
Where:
- $N$: Number of non-embedding parameters.
- $D$: Number of training tokens.
- Performance is bottlenecked by the smaller of $N$ or $D$. If you have a massive model but a tiny dataset, you will not see scaling benefits.
2. PyTorch Implementation: Parameter Counting
A critical part of obeying scaling laws is precisely calculating $N$ (non-embedding parameters). Kaplan et al. noted that embedding parameters do not contribute significantly to the “intelligence” (loss reduction) of the model as it scales.
Python
import torch
import torch.nn as nn
def count_non_embedding_params(model):
"""
As per Kaplan et al. (2020), scaling laws are driven by
the model's internal parameters, excluding the embedding table.
"""
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
# Identify and subtract embedding parameters
emb_params = 0
for module in model.modules():
if isinstance(module, nn.Embedding):
emb_params += sum(p.numel() for p in module.parameters())
non_emb_params = total_params - emb_params
return non_emb_params
3. Compute Estimation ($C$)
The paper provides a simple “rule of thumb” for the floating-point operations (FLOPs) required for a forward-backward pass of a Transformer. For a model with $N$ parameters and a dataset of $D$ tokens:
$$C \approx 6ND$$
Python
def estimate_training_compute(non_emb_params, tokens_trained):
"""
Estimates total FLOPs for training.
6 * N * D covers:
- 2 ND for the forward pass
- 4 ND for the backward pass
"""
return 6 * non_emb_params * tokens_trained
# Example: A 1B parameter model trained on 20B tokens
params = 1e9
tokens = 20e9
total_flops = estimate_training_compute(params, tokens)
print(f"Total Compute: {total_flops:.2e} FLOPs")
4. Implementation: The Log-Linear Scheduler
The paper suggests that learning rate schedules and batch sizes should scale with the model size. Specifically, they found that the Critical Batch Size ($B^*$) increases as the loss decreases, following its own power law.
Python
class ScalingLawScheduler:
"""
Simplified representation of scaling the learning rate
based on the model's 'd_model'.
"""
def __init__(self, optimizer, d_model, warm_up_steps=3000):
self.optimizer = optimizer
self.d_model = d_model
self.warm_up_steps = warm_up_steps
self.current_step = 0
def step(self):
self.current_step += 1
# Scaling law insight: Larger models require smaller learning rates
# Base LR is often scaled by (d_model)^-0.5
lr = (self.d_model ** -0.5) * min(self.current_step ** -0.5,
self.current_step * (self.warm_up_steps ** -1.5))
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr
5. Implications for Model Design
| Finding | Implementation Action |
| Architecture Irrelevance | Performance depends on total $N$, not the specific shape (width vs. depth). |
| Overfitting is Rare | On sufficiently large datasets, performance continues to improve; “more is better.” |
| The “Compute Optimal” Point | For every 10x increase in compute, increase model size by ~5.5x and dataset size by ~1.8x. |
Why it Matters
This paper provided the “green light” for the creation of massive models like GPT-3 (175B). It proved that we hadn’t hit a ceiling; we just needed to keep scaling $N$, $D$, and $C$ in the correct proportions.
Note: These laws were later refined by the Chinchilla paper (Hoffmann et al., 2022), which argued that we were actually under-training our models and needed much more data ($D$) per parameter ($N$).