Scaling Laws & Compute-Optimal Training

Scaling laws are the rare case where "just make it bigger" turned out to be rigorous science. The math says your model is too small, your data is too little, and your budget is never enough. At least now you can quantify the despair.

Scale, Budget Crushing AI Agent

6.3.1 The Power Law Foundation

The remarkable empirical discovery underlying all scaling laws is that language model loss follows a power law relationship with respect to model size, dataset size, and compute. This means that as you increase any of these quantities, loss decreases predictably as a straight line on a log-log plot. Formally, the relationship takes the form:

Here, $x$ is the quantity being scaled (parameters, tokens, or FLOPs), $a$ is a constant, $\alpha$ is the scaling exponent (typically between 0.05 and 0.10), and $L_{ \infty }$ is the irreducible loss (the entropy of natural language itself). The irreducible loss represents the theoretical limit: no model, regardless of size, can predict language perfectly because language is inherently stochastic.

This relationship holds across many orders of magnitude, which is what makes it practically useful. You can train a series of small models, fit a power law curve, and then extrapolate to predict the loss of a much larger model.

6.3.2 Kaplan Scaling Laws (2020)

The foundational work by Kaplan et al. at OpenAI established three key relationships. First, loss scales as a power law with model parameters $N$ (number of non-embedding parameters):

Second, loss scales as a power law with dataset size $D$ (number of tokens):

Third, loss scales as a power law with compute budget $C$ (in FLOPs (floating-point operations, a count; not to be confused with FLOPS, which measures operations per second)):

The Kaplan Compute-Optimal Recipe

A critical conclusion from the Kaplan analysis was that, given a fixed compute budget, you should prioritize increasing model size over increasing the number of training tokens. Specifically, Kaplan found that as compute increases by 10x, you should scale model size by roughly 5x but only increase data by about 2x. This led to a generation of very large models trained on relatively modest amounts of data, exemplified by GPT-3 (175B parameters trained on 300B tokens).

The Kaplan recipe dominated thinking at major labs for two years. Then a team at DeepMind asked a simple question: what happens if we actually train models to convergence before drawing conclusions? The answer upended the entire field's approach to compute allocation.

6.3.3 Chinchilla Scaling Laws (2022)

Hoffmann et al. at DeepMind revisited scaling with more careful experimental design, training over 400 models ranging from 70M to 16B parameters. Their key methodological improvement was training each model to near-convergence on its dataset, eliminating the early-stopping bias in the Kaplan analysis.

The Chinchilla result was striking: for a compute-optimal training run, the number of parameters and the number of training tokens should scale equally. The combined loss is modeled as:

where $\alpha \approx 0.34$, $\beta \approx 0.28$, $E \approx 1.69$ (the irreducible entropy), and $A$, $B$ are constants. Minimizing this loss subject to a compute constraint $C \approx 6ND$ yields the compute-optimal allocation:

The factor of 6 comes from counting floating-point operations per parameter per token. A forward pass through a weight is a multiply-accumulate, or 2 FLOPs (one multiply, one add); the backward pass costs about twice the forward, 4 FLOPs (one matmul to propagate the gradient to the layer inputs, one to compute the weight gradient). Forward plus backward is therefore $2 + 4 = 6$ FLOPs per parameter per token, so training $N$ parameters on $D$ tokens costs $C \approx 6ND$.

This means parameters and tokens should be scaled at roughly the same rate. (Hoffmann et al. reported three estimation approaches: Approaches 1 and 2 yield exponents near 0.50/0.50, while the parametric Approach 3 yields the slightly skewed $a = \alpha/(\alpha+\beta) \approx 0.55$, $b = \beta/(\alpha+\beta) \approx 0.45$. The 50/50 form is the most-cited simplification.) The practical implication is that a 70B model should be trained on approximately 1.4 trillion tokens (a ratio of about 20 tokens per parameter).

The same 20:1 rule can be checked in three lines of Python for any candidate model:

# 20-tokens-per-param sanity check for a real model.
def chinchilla_ratio(params: float, tokens: float) -> str:
    r = tokens / params
    verdict = "optimal" if 15 <= r <= 25 else ("under-trained" if r < 15 else "over-trained")
    return f"ratio = {r:.0f} tokens/param -> {verdict}"

print(chinchilla_ratio(175e9, 300e9))   # GPT-3:        ratio = 2 tokens/param -> under-trained
print(chinchilla_ratio(70e9, 1.4e12))   # Chinchilla:   ratio = 20 tokens/param -> optimal
print(chinchilla_ratio(8e9, 15e12))     # Llama-3 8B:   ratio = 1875 tokens/param -> over-trained

Algorithm: Compute-optimal (N*, D*) for fixed FLOPs budget C
Input:  total FLOPs budget C, fitted constants
        E = 1.69, A = 406.4, B = 410.7, alpha = 0.34, beta = 0.28
        FLOPs-per-token approximation: 6 N D = C
Output: optimal parameter count N*, optimal token count D*

  // 1. Parametric loss surface (Hoffmann et al., 2022)
  L(N, D) := E + A / N^alpha + B / D^beta

  // 2. Substitute the budget constraint D = C / (6 N), then minimize over N
  L(N) := E + A / N^alpha + B * (6 N / C)^beta

  // 3. Set dL/dN = 0
  d/dN [ A / N^alpha + B * (6 N / C)^beta ] = 0
  -alpha A / N^{alpha + 1} + beta B (6/C)^beta N^{beta - 1} = 0
  => N^{alpha + beta} = (alpha A) / (beta B (6/C)^beta)
  => N* = G * C^a    with    a = beta / (alpha + beta)
         D* = (C / 6) / N* = G' * C^b    with    b = alpha / (alpha + beta)

  // 4. Numerically: a ~= 0.50, b ~= 0.50 (parameters and tokens scale jointly)
  Return N*, D*

Practical rule of thumb (Hoffmann et al., 2022):
    D* / N* ~= 20   (i.e., "20 training tokens per parameter")

Figure 6.3.4: The Kaplan approach favors larger models with less data, while the Chinchilla approach recommends equal scaling of parameters and tokens.

Why Chinchilla Changed Everything

The Chinchilla result implied that many existing models were significantly undertrained. Gopher (280B parameters trained on 300B tokens) was revealed to be suboptimal: a 70B model trained on 1.4T tokens (Chinchilla) matched or exceeded Gopher on nearly every benchmark, while being 4x smaller and therefore 4x cheaper to serve at inference time (for the techniques that make smaller models fast to serve, see Chapter 9). This triggered a major shift in the industry. Post-Chinchilla models like Llama were designed with much larger data-to-parameter ratios.

6.3.4 Beyond Chinchilla: Over-Training for Inference

While Chinchilla defines the compute-optimal point for a single training run, real-world deployments face a different optimization problem. A model is trained once but serves millions of inference requests. From this total cost perspective, it can be economical to train a smaller model on far more data than is compute-optimal, paying more in training compute to reduce inference cost per query.

The Llama family exemplifies this strategy. Llama-1 7B was trained on 1 trillion tokens, giving a ratio of approximately 143 tokens per parameter, roughly 7x beyond the Chinchilla-optimal ratio. Llama-2 was trained on 2 trillion tokens. The rationale: the additional training cost is paid once, but the smaller model saves compute on every single inference call. An alternative to over-training is knowledge distillation, which transfers capabilities from a large teacher model into a smaller student.

Post-Chinchilla Scaling: Inference-Optimal Training in Practice

The post-Chinchilla era (2023 to 2025) has produced a rich empirical record of how leading labs actually allocate training budgets. The trend is clear: every major open-weight model released since 2023 has been trained on far more tokens than Chinchilla would prescribe, and the degree of over-training has increased with each generation.

Llama-3 (Meta, 2024) represents the most dramatic departure from Chinchilla ratios. The Llama-3 8B model was trained on approximately 15 trillion tokens, giving a tokens-per-parameter ratio of roughly 1,875, nearly 100x the Chinchilla-optimal ratio of approximately 20. Even the Llama-3 70B model used 15T tokens, yielding a ratio of about 214 tokens per parameter. Meta's technical report explicitly frames this as an inference-cost optimization: by investing more in training, the resulting models deliver stronger per-query performance at their given parameter count, amortizing the training cost over billions of inference calls. The Llama-3 training consumed roughly 30 million GPU-hours on H100 hardware, a cost that Meta estimated would be recovered within months of deployment given inference volume.

Inference-optimal scaling laws. Sardana and Frankle (2024) formalized this intuition into quantitative "inference-aware" scaling laws that incorporate both training cost and expected inference volume into the optimization objective. Their key finding: when the total inference compute is expected to exceed training compute by 10x or more (which is typical for any model serving thousands of users), the optimal strategy is to train a model that is 2 to 4x smaller than Chinchilla-optimal but on 5 to 20x more data. The exact ratio depends on the inference hardware and deployment pattern. For batch inference workloads (where throughput matters more than latency), the optimal model is larger; for real-time serving (where latency constraints dominate), smaller models trained on more data deliver better cost efficiency.

Algorithm: Total-cost-optimal (N*, D*) accounting for inference
Input:  expected lifetime inference tokens T_inf
        Chinchilla parameters (E, A, B, alpha, beta) from Algorithm 6.3.1
        FLOPs constants: train = 6 N D, inference = 2 N per generated token
Output: (N*, D*) minimizing total FLOPs cost for target loss L_target

  // 1. Total compute = training + inference
  C_total(N, D) := 6 N D + 2 N T_inf

  // 2. Substitute D from the loss constraint L(N, D) = L_target
  //    using L(N, D) = E + A / N^alpha + B / D^beta
  D(N) := ( B / (L_target - E - A / N^alpha) )^{1 / beta}

  // 3. Minimize C_total(N, D(N)) over N
  Solve dC_total / dN = 0 numerically (one-dimensional convex problem)

  // 4. Interpretation: the optimum trades 6 D training FLOPs against 2 T_inf
  //    inference FLOPs per parameter. As T_inf grows, the optimum N* shrinks
  //    (smaller model) and D* grows (more training tokens per parameter).

  Return (N*, D*)

Practical regime (T_inf >> training tokens):
    N* ~= 0.3 to 0.5 * N*_Chinchilla   (smaller model)
    D* ~= 5 to 20 * D*_Chinchilla     (over-train on more data)
    matches Llama-3 8B (15T tokens, ~1875 tokens/param, ~94x Chinchilla)

Diminishing returns and the data quality frontier. Over-training does not scale indefinitely. Empirical evidence from Llama-3 shows that loss continues to decrease as more tokens are processed, but the rate of improvement diminishes. Between 5T and 10T tokens, loss improves more slowly than between 1T and 5T tokens, and the gains from 10T to 15T tokens are smaller still. Critically, these diminishing returns depend on data quality. Training on 15T tokens of high-quality, deduplicated data produces substantially better models than 15T tokens of raw web crawl. Meta's data curation pipeline for Llama-3 was reported to filter out over 90% of raw web data, and they used classifier-based quality filtering trained on human preferences. This has shifted the competitive frontier from "who can get the most data" to "who can build the best data curation pipeline," a theme explored further in Section 15.1.

6.3.5 Data-Constrained Scaling

A growing concern in the LLM community is the potential exhaustion of high-quality training data. Muennighoff et al. (2023) studied what happens when the Chinchilla-optimal token count exceeds available data. Their findings suggest that repeating data up to 4 epochs causes minimal degradation in performance, but beyond that, the value of additional repetitions diminishes rapidly. For a given compute budget $C$ with a data budget $D_{\max}$, the effective token count follows:

where $R = D_{total}/D_{\max}$ is the number of epochs. This diminishing-returns formula implies that once you have exhausted your data budget, the marginal benefit of additional epochs is exponentially decaying.

Scaling laws tell us how loss decreases smoothly with more compute and data. But something unexpected happens as models grow: certain capabilities appear to materialize abruptly rather than gradually. This phenomenon has sparked one of the most active debates in the field.

6.3.6 Emergent Capabilities and Phase Transitions

One of the most debated phenomena in LLM scaling is emergence: the apparent sudden appearance of new capabilities at certain model sizes. Tasks like arithmetic, chain-of-thought reasoning, and multi-step logic appear to be absent in small models and then abruptly appear in larger ones. Chapter 42 covers the evaluation frameworks used to measure these capabilities systematically. Wei et al. (2022) catalogued over 100 such emergent tasks across the BIG-Bench benchmark suite.

The Metric Mirage Hypothesis

Schaeffer et al. (2023) challenged the notion of sharp emergence. Their key argument: whether a capability appears "emergent" depends heavily on the choice of evaluation metric. With discrete metrics like exact-match accuracy, performance looks flat at zero until a threshold is crossed, creating the illusion of a sudden phase transition. When the same tasks are measured with continuous metrics (like token-level log-likelihood), performance improves smoothly and predictably. The capability was always improving; the metric just could not detect the gradual progress.

Figure 6.3.5a: The same underlying capability can appear emergent or smoothly scaling depending on the evaluation metric chosen.

6.3.7 Mixture-of-Experts Architecture

The scaling laws we have studied so far assume a dense model: one where every parameter participates in every forward pass. But what if you could double the number of parameters without doubling the compute cost per token? That is the core idea behind Mixture-of-Experts (MoE), a sparse architecture that is reshaping how frontier models are built.

An MoE layer replaces the single feed-forward network (FFN) sublayer in a transformer block with N expert FFN layers, but activates only K of them for any given token (where K is much smaller than N, typically K=2 or K=8 out of 64 or 256 experts). The parameter count grows with N, but the FLOPs per token grow only with K. The original idea was introduced by Shazeer et al. (2017), later scaled dramatically by Fedus et al. (2022) with Switch Transformers, and has since become a standard architectural choice at frontier scale.

Router and Gating Mechanism

The routing decision is made by a small learned gating network: a linear projection from the token's hidden state into N scores, one per expert. A softmax normalizes these scores into routing probabilities. For top-K routing (hard routing), only the K experts with the highest probabilities receive the token; the others are skipped entirely. The output of the MoE layer is a weighted combination of the K activated experts' outputs, weighted by their normalized routing probabilities.

Formally, given token representation $\mathbf{x}$, the router computes:

The top-K selection produces a sparse gate vector $\tilde{g}$ with only K non-zero entries, and the layer output is:

The distinction between hard routing (top-K with exactly K winners) and soft routing (a weighted mixture of all experts with learned soft gates) has practical consequences. Hard routing is more compute-efficient because non-selected experts do not run at all, making the FLOPs savings concrete. Soft routing can be more stable during training but requires running all experts, negating the sparsity benefit during training itself.

Load Balancing: Preventing Expert Collapse

Without any constraint, routing networks collapse: the gating network discovers a few "popular" experts early in training and routes most tokens to them, while the remaining experts receive almost no gradient signal and never develop useful specializations. This is called expert collapse.

The standard fix is an auxiliary load-balancing loss added to the training objective. Define the fraction of tokens routed to expert i as $f_i$ and the average routing probability assigned to expert i as $p_i$. The auxiliary loss penalizes uneven routing:

Here $\alpha$ is a small coefficient (typically 0.01 to 0.001) and the factor $N$ normalizes by the number of experts. This auxiliary loss is differentiable with respect to $p_i$ (the soft routing probabilities), even though the hard selection of top-K experts is not itself differentiable. The practical tuning of $\alpha$ is delicate: too high and the model sacrifices quality for perfect load balance; too low and collapse re-emerges.

MoE Layer Diagram

Figure 6.3.5b: An MoE layer with 6 experts and top-2 routing. The router assigns probabilities to all experts but only activates the 2 highest-scoring ones for this token. The final output is a weighted sum of the two active experts' outputs. The other 4 experts perform no computation for this token, saving FLOPs while preserving parameter capacity.

Landmark MoE Models: Mixtral and DeepSeek-V3

The active-parameter accounting that drives the efficiency dividend for both models can be stated compactly. For an MoE layer with $N$ total expert parameters, $E$ experts, top-$K$ routing, and shared (non-expert) parameters $P_{\text{shared}}$ per layer, the active parameter count per token is:

For Mixtral 8x7B the FFN holds roughly 41B of the 47B total, so $N_{\text{active}} \approx 6 + \tfrac{2}{8} \cdot 41 \approx 16\text{B}$, dominated by the active expert FFNs. For DeepSeek-V3 the analogous computation lands at about 37B out of 671B, a 18x sparsity ratio. Training FLOPs scale with $N_{\text{active}}$, not the total, which is why DeepSeek-V3 trains for the price of a 37B dense model while delivering quality closer to a 70-100B dense model.

Mixtral 8x7B (Mistral AI, 2024) was the model that brought MoE into widespread practical use. The architecture uses 8 expert FFN layers per transformer block with top-2 routing. Despite having 47 billion total parameters (8 experts times the FFN size), only about 13 billion parameters are active for any given token, making inference cost comparable to a 13B dense model while achieving quality closer to a 70B dense model on most benchmarks. The reduction in active parameters directly translates to faster generation and lower memory bandwidth requirements at inference time.

DeepSeek-V3 (DeepSeek AI, 2024) pushed MoE to a new scale. With 671 billion total parameters across 256 experts per layer and top-8 routing, only about 37 billion parameters are active per token. The architecture also incorporates auxiliary-loss-free load balancing, a refinement over the standard auxiliary loss approach: rather than penalizing imbalance during training, the router's bias terms are adjusted dynamically based on observed load, decoupling load balancing from the loss objective and allowing the router to specialize more aggressively. DeepSeek-V3 was trained on 14.8 trillion tokens, and its technical report claims the full training run cost approximately 2.8 million H800 GPU hours, a fraction of what comparable dense models cost. The key efficiency gain comes from the MoE architecture: the 671B total parameters provide the knowledge capacity of a very large model, but the 37B active parameters mean the FLOPs budget per training step is far below what a 671B dense model would require.

# Load Mixtral 8x7B with 4-bit quantization (~25 GB vs ~94 GB at FP16).
# device_map="auto" distributes layers across all available GPUs / RAM.
# Requires: pip install transformers accelerate bitsandbytes
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",          # NormalFloat4 (best for LLM weights)
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,     # second-level quant on the constants
)

model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
)

prompt = "Explain mixture-of-experts in two sentences."
inputs = tok(prompt, return_tensors="pt").to(model.device)
output = model.generate(**inputs, max_new_tokens=200, temperature=0.7, do_sample=True)
print(tok.decode(output[0], skip_special_tokens=True))

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mixtral-8x7B-v0.1",
    device_map="auto", # spreads across available GPUs
    torch_dtype=torch.bfloat16,
    load_in_4bit=True, # BitsAndBytes 4-bit quantization
)
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-v0.1")
# Check how many parameters are loaded vs. total
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters loaded: {total_params / 1e9:.1f}B")
# Active per forward pass: ~13B out of 47B total

Loading DeepSeek-V3 follows the same pattern but requires trust_remote_code=True because the MoE routing and multi-head latent attention layers ship as custom modeling code in the repository rather than as a built-in Transformers architecture.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "deepseek-ai/DeepSeek-V3"
tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,  # required: custom MoE + MLA modeling code
)
prompt = "Explain auxiliary-loss-free load balancing in one paragraph."
inputs = tok(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=200, do_sample=False)
print(tok.decode(out[0], skip_special_tokens=True))

Expert Parallelism: Distributing MoE Across GPUs

MoE introduces a new parallelism strategy called expert parallelism. In a dense model, tensor parallelism splits each matrix across GPUs. In an MoE model, different GPUs can hold different subsets of the expert pool. Because a given token only activates K of the N experts, and those experts may reside on different GPUs, expert parallelism requires an all-to-all communication pattern: after the router decides which experts each token goes to, the tokens must be dispatched to the GPUs holding those experts, processed, and then gathered back. This all-to-all communication is a significant engineering challenge at large scale. With 256 experts across dozens of GPUs, the communication volume can become a bottleneck, particularly when routing is uneven and some GPUs receive more tokens than others in a given batch. This is a second reason load balancing matters so much in practice: not just for gradient signal, but for communication efficiency.

6.3.8 Scaling Laws for Sparse Models

Sparse MoE models do not follow the same Chinchilla-optimal ratios as dense models. Clark et al. (2022) derived scaling laws specifically for routed models, finding that MoE introduces an additional scaling axis: the number of experts E. The effective parameter count for loss prediction is not the total parameter count N but something closer to the active parameter count multiplied by a factor that grows with E. Informally, experts add capacity at a discount: doubling experts improves loss, but by less than doubling the total parameter count of a dense model.

The Effective Parameter Concept

A useful practical heuristic for MoE models is the concept of effective parameters. An MoE model with N total parameters and K/E active fraction behaves somewhere between a dense model of size K*N/E (active parameters only) and a full dense model of size N, depending on the degree of expert specialization and the number of training tokens. For training compute budgeting, the relevant quantity is the FLOPs cost, which scales with active parameters, not total parameters. For inference quality, the relevant quantity is closer to total parameters, because different tokens see different experts and the model's "effective knowledge" is spread across the full expert pool.

Concretely: DeepSeek-V3 with 671B total parameters and 37B active behaves at inference quality roughly like a 70-100B dense model on most tasks, but its training FLOPs per token are set by the 37B active parameters. This means DeepSeek-V3 was trained at roughly the compute cost of a 37B dense model while achieving the quality of a much larger one. This breaks the implicit Chinchilla assumption that quality scales with the parameters receiving gradient signal: in MoE, parameters that are not selected for a given token are still updated indirectly through the load-balancing loss and through tokens that do route to them, but they receive less gradient signal per training step.

Training Compute vs. Inference Compute

MoE delivers efficiency advantages on both fronts. Training a 671B MoE model at 37B active parameters per token costs roughly the same per-step FLOPs as training a 37B dense model. At inference time, serving the model costs roughly as much memory bandwidth as serving a 37B dense model (plus the overhead of routing and all-to-all communication), while delivering quality closer to a 70-100B dense model. The cost tradeoffs differ between deployment contexts:

• High-throughput batch inference: MoE models can be very efficient because all-to-all communication overhead is amortized over large batches, and the active-parameter FLOPs advantage is maximized.
• Low-latency single-token serving: The all-to-all communication adds latency overhead, and the full expert pool must reside in GPU memory (or be swapped, which is slow). Dense models with smaller active parameter counts may be preferable for strict latency requirements.
• Edge and on-device inference: MoE models are generally a poor fit because the full model must be loaded even though only a fraction executes per token. The memory constraint matters more than the compute constraint in this regime.

Dense vs. MoE Comparison

# Fit a power-law scaling curve L(N) = a * N^b + c to empirical loss
# data from small training runs, then extrapolate to larger model sizes.
import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
# Empirical data: (parameters, final_loss) from small training runs
params = np.array([1e6, 5e6, 2e7, 5e7, 1e8, 5e8])
losses = np.array([4.20, 3.75, 3.35, 3.15, 2.98, 2.70])
# Power law model: L(N) = a * N^(-alpha) + L_inf
def scaling_law(N, a, alpha, L_inf):
    return a * N ** (-alpha) + L_inf
# Fit the curve
popt, pcov = curve_fit(
    scaling_law, params, losses,
    p0=[100, 0.07, 1.5], # initial guesses
    bounds=([0, 0, 0], [1e6, 1.0, 5.0])
    )
a_fit, alpha_fit, L_inf_fit = popt
print(f"Fitted: a={a_fit:.2f}, alpha={alpha_fit:.4f}, L_inf={L_inf_fit:.3f}")
# Predict loss for larger model sizes
target_sizes = [1e9, 7e9, 70e9]
for size in target_sizes:
    predicted = scaling_law(size, *popt)
    print(f"  {size/1e9:.0f}B params => predicted loss: {predicted:.3f}")

Computing the Chinchilla-Optimal Allocation

This snippet computes the optimal token-to-parameter ratio for a given compute budget using the Chinchilla scaling law.

# Chinchilla-optimal allocation: given a fixed FLOPs budget,
# compute the ideal model size N and token count D (roughly D ~ 20*N).
def chinchilla_optimal(compute_budget_flops):
    """
    Given a FLOPs budget, compute the Chinchilla-optimal
    model size (N) and token count (D).
    Uses the approximation: C = 6 * N * D
    Chinchilla ratio: D = 20 * N
    Therefore: C = 6 * N * 20 * N = 120 * N^2
    """
    N_opt = (compute_budget_flops / 120) ** 0.5
    D_opt = 20 * N_opt
    return N_opt, D_opt
# Example compute budgets
budgets = {
    "Small (1e19 FLOPs)": 1e19,
    "Medium (1e21 FLOPs)": 1e21,
    "Large (1e23 FLOPs)": 1e23,
    "GPT-4 scale (1e25)": 1e25,
    }
for name, budget in budgets.items():
    N, D = chinchilla_optimal(budget)
    print(f"{name}:")
    print(f"  Optimal model size: {N/1e9:.1f}B parameters")
    print(f"  Optimal data:       {D/1e9:.0f}B tokens")
    print()

Exercises

6.3.9 Multi-Token Prediction and Scaling

Standard language model training predicts one token at a time: given the context, predict the next token. Multi-token prediction (MTP) changes this by adding additional prediction heads to the model, each trained to predict a token further in the future. Gloeckle et al. (Meta, 2024) showed that this architectural change improves both sample efficiency and downstream performance, with the gains scaling favorably with model size.

Architecture: Multiple Prediction Heads

The standard MTP implementation adds d additional small transformer layers, each producing a prediction of the token d steps ahead. The architecture during training looks like:

• The main model body processes the sequence and produces hidden states.
• Head 0 (the standard next-token head) predicts token t+1.
• Head 1 (additional head) predicts token t+2, using the hidden state from head 0 concatenated with the embedding of the true token t+1.
• Head 2 predicts token t+3, using head 1's hidden state concatenated with the embedding of the true token t+2.
• And so on for the remaining d-2 future heads.

Each future head contributes its own cross-entropy loss term to the total training objective, with a weighting coefficient (typically 0.1 to 0.3 per future head) to avoid dominating the main next-token signal. During inference, only head 0 is used for standard auto-regressive generation; the future heads are discarded.

Why MTP Improves Sample Efficiency

The intuition is that predicting multiple future tokens provides richer gradient signal per forward pass. A single forward pass over a sequence of length L with d future prediction heads produces d times as many loss terms as standard training. Put differently, the model receives signal about global structure (several tokens ahead) in addition to local next-token prediction. This is especially valuable for tasks requiring planning: code generation, structured text, and mathematical reasoning, where the correct next token depends on knowing what the overall structure should be several steps later.

From a scaling law perspective, MTP effectively shifts the loss curve downward without changing the scaling exponent. A model trained with MTP achieves the same validation loss as a standard model trained on more data. Gloeckle et al. (2024) report that the gains from MTP are more pronounced at larger model scales: for models below 1B parameters, MTP shows modest improvements, but for models at 7B and above, MTP consistently delivers meaningful gains in coding and reasoning benchmarks at the same training compute budget.

Connection to Speculative Decoding

MTP creates an interesting synergy with speculative decoding, covered in Section 10.3. In speculative decoding, a small "draft" model generates several tokens quickly, and the larger target model verifies them in parallel. If the model was trained with MTP, the future prediction heads can serve as the draft model directly, allowing the same model to speculate on its own output. This is called self-speculative decoding. The draft and target are the same weights; the future heads propose tokens that the main model verifies. DeepSeek-V3 reports using MTP-based self-speculation to improve throughput at inference time, getting meaningful speedups without a separate small draft model.

Now that we have covered the theory behind scaling laws and their practical implications, let us get hands-on. The following lab walks you through fitting your own scaling law curves to empirical data and using them to predict performance at larger scales.

6.3.10 Summary Table: Scaling Regimes