Chinchilla Scaling Laws › Code Implementation
Language AI HandbookAll chapters Chinchilla Scaling Laws › Code ImplementationOne section · Sections in this chapter7 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement functions to compute Chinchilla-optimal allocations and compare them with Kaplan allocations. The implementation will help you build intuition for how the two frameworks diverge as compute scales, and give you a practical tool for planning training runs under either framework. We will start with the core allocation functions, then compute and visualize the results across a wide range of compute budgets to see how the predictions diverge at large scales.
import numpy as np
# Set random seed for reproducibility
np.random.seed(42)
def chinchilla_optimal_allocation(compute_flops: float) -> dict:
"""
Compute Chinchilla-optimal model size and dataset size.
Uses the approximation C = 6*N*D and D/N ≈ 20.
"""
# C = 6 * N * D = 6 * N * 20 * N = 120 * N^2
# Therefore N = sqrt(C / 120)
optimal_params = np.sqrt(compute_flops / 120)
optimal_tokens = 20 * optimal_params
return {
"parameters": optimal_params,
"tokens": optimal_tokens,
"tokens_per_param": optimal_tokens / optimal_params,
"compute_flops": compute_flops,
}
def kaplan_optimal_allocation(compute_flops: float) -> dict:
"""
Compute Kaplan-style allocation with N scaling faster than D.
Uses approximate exponents a=0.73, b=0.27 from original paper.
Calibrated so both methods agree at C = 10^20 for comparison.
"""
# Calibration constants (tuned to match similar scale)
alpha_n = 0.73
alpha_d = 0.27
# Reference point for calibration
c_ref = 1e20
n_ref = np.sqrt(c_ref / 120) # Match Chinchilla at reference
d_ref = 20 * n_ref
k_n = n_ref / (c_ref**alpha_n)
k_d = d_ref / (c_ref**alpha_d)
optimal_params = k_n * (compute_flops**alpha_n)
optimal_tokens = k_d * (compute_flops**alpha_d)
return {
"parameters": optimal_params,
"tokens": optimal_tokens,
"tokens_per_param": optimal_tokens / optimal_params,
"compute_flops": compute_flops,
}Now let's compute allocations across a range of compute budgets and visualize the differences:
# Define compute budgets from 10^18 to 10^25 FLOPs
compute_budgets = np.logspace(18, 25, 50)
chinchilla_allocations = [
chinchilla_optimal_allocation(c) for c in compute_budgets
]
kaplan_allocations = [kaplan_optimal_allocation(c) for c in compute_budgets]
# Extract arrays for plotting
chinchilla_params = np.array([a["parameters"] for a in chinchilla_allocations])
chinchilla_tokens = np.array([a["tokens"] for a in chinchilla_allocations])
kaplan_params = np.array([a["parameters"] for a in kaplan_allocations])
kaplan_tokens = np.array([a["tokens"] for a in kaplan_allocations])
The left panel shows that as compute budgets increase, Kaplan scaling recommends much larger models than Chinchilla scaling. At FLOPs, the gap is nearly an order of magnitude. The right panel reveals the consequence: Kaplan allocates far fewer tokens to training, resulting in undertrained models. Chinchilla's balanced approach maintains the 20:1 token-to-parameter ratio across all scales.
The divergence becomes more pronounced at higher compute budgets. We can quantify this by looking at specific examples:
def format_number(n: float) -> str:
"""Format large numbers with appropriate suffixes."""
if n >= 1e12:
return f"{n / 1e12:.1f}T"
elif n >= 1e9:
return f"{n / 1e9:.1f}B"
elif n >= 1e6:
return f"{n / 1e6:.1f}M"
else:
return f"{n:.0f}"
# Compare at several compute scales
compute_examples = [1e20, 1e22, 1e24]
comparison_results = []
for c in compute_examples:
chin = chinchilla_optimal_allocation(c)
kap = kaplan_optimal_allocation(c)
comparison_results.append((c, chin, kap))Chinchilla vs Kaplan Optimal Allocations ---------------------------------------------------------------------- Compute Budget: 1e+20 FLOPs Method Parameters Tokens Tokens/Param ------------------------------------------------------ Chinchilla 912.9M 18.3B 20.0 Kaplan 912.9M 18.3B 20.0 Compute Budget: 1e+22 FLOPs Method Parameters Tokens Tokens/Param ------------------------------------------------------ Chinchilla 9.1B 182.6B 20.0 Kaplan 26.3B 63.3B 2.4 Compute Budget: 1e+24 FLOPs Method Parameters Tokens Tokens/Param ------------------------------------------------------ Chinchilla 91.3B 1.8T 20.0 Kaplan 759.3B 219.5B 0.3
At FLOPs, Kaplan scaling would suggest training a model with over 20x more parameters than Chinchilla recommends, but on far fewer tokens. The Chinchilla model would likely outperform significantly because it receives adequate training data.
The comparison table reveals a clear pattern. While both methods scale up with compute, they diverge sharply in how they allocate resources. At FLOPs, Kaplan scaling recommends a 2.4T parameter model trained on only 700B tokens (0.3 tokens per parameter), while Chinchilla recommends a 91B parameter model trained on 1.8T tokens (20 tokens per parameter). The Kaplan-style allocation would produce a severely undertrained model.
Now we visualize the tokens-per-parameter ratio across compute scales:
The flat blue line at 20 tokens per parameter represents the Chinchilla prescription: regardless of how much compute you have, maintain the same data intensity. The declining red line shows the problematic Kaplan implication: as compute grows, train on proportionally less data per parameter. At FLOPs, Kaplan scaling would suggest training on fewer than 5 tokens per parameter, a recipe for severe underfitting.
This visualization makes the fundamental difference clear: Chinchilla scaling maintains a constant data intensity regardless of scale, while Kaplan scaling implies training larger models on proportionally less data.