Pytorch implementation of paper: Scaling Laws for Neural Language Models

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$).