Data-Constrained Scaling › Code Implementation

Michael BrenndoerferOctober 26, 20258 min read
Language AI HandbookAll chapters
Part 23 of 67Scaling Laws
Chapter 166Data-Constrained Scaling › Code Implementation
Reading section 7 of 10 · Code Implementation · 8 minSection 7 of 10 · 8 min

Book contents

Section 7 of 10 · 8 min read

Data-Constrained Scaling › Code ImplementationOne section · Sections in this chapter7 / 10 ↓
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

We can simulate data repetition effects to build intuition about how repetition degrades effective data value. The code that follows implements the mathematical framework we developed above, allowing us to see concretely how the repetition penalty affects training decisions.

In[10]:
Code
import numpy as np

rng = np.random.default_rng(42)


def effective_tokens(
    unique_tokens: int, epochs: int, alpha: float = 0.1
) -> float:
    """
    Calculate effective token count given repetition.

    Uses the Muennighoff et al. repetition penalty formula.
    """
    if epochs <= 0:
        return 0

    # g(R) = (1 - exp(-alpha * R)) / (1 - exp(-alpha))
    g_r = (1 - np.exp(-alpha * epochs)) / (1 - np.exp(-alpha))
    return unique_tokens * g_r

This function implements the repetition penalty formula. The alpha parameter controls how quickly returns diminish. Notice how the implementation directly translates the mathematical formula we derived earlier: the exponential saturation curve g(R)g(R) modulates the raw token count to produce an effective count that grows more slowly than the naive expectation.

In[11]:
Code
# Simulate effective tokens across different repetition levels
unique_data = 100_000_000_000  # 100 billion unique tokens
epochs_range = np.arange(1, 33)

eff_tokens = [effective_tokens(unique_data, e) for e in epochs_range]
naive_tokens = unique_data * epochs_range  # what we'd expect without penalty
efficiency = (
    np.array(eff_tokens) / naive_tokens * 100
)  # percentage of naive expectation
Out[12]:
Console
Epoch | Effective Tokens (B) | Naive Tokens (B) | Efficiency
------------------------------------------------------------
   1  |        100.0        |         100      |   100.0%
   2  |        190.5        |         200      |    95.2%
   4  |        346.4        |         400      |    86.6%
   8  |        578.7        |         800      |    72.3%
  16  |        838.7        |        1600      |    52.4%
  32  |       1008.0        |        3200      |    31.5%

The efficiency column reveals the harsh reality of repetition. By 4 epochs, each token is worth only about 85% of its naive value. By 16 epochs, efficiency drops to roughly 63%. You're processing 16× the tokens but getting only about 10× the effective training signal. This quantifies why practitioners should carefully consider whether additional epochs are worth the computational cost. The gap between naive expectation and effective value represents pure computational waste. FLOPs spent on redundant learning signal.

Out[13]:
Visualization
(0.0, 3.5)
Line plot showing effective tokens flattening as epochs increase, diverging from linear expectation.
Effective data value vs. epochs of repetition. The dashed line shows naive expectation (linear growth), while the solid line shows actual effective tokens accounting for diminishing returns.

The growing gap between naive expectation and effective tokens represents the fundamental cost of data repetition. The shaded region in the figure grows wider with each additional epoch. This represents computational resources that produce no corresponding improvement in model capability. Now let's examine how this affects optimal model sizing and compare scenarios with different data availability.

In[14]:
Code
def chinchilla_loss(
    params: float,
    effective_tokens: float,
    A: float = 406.4,
    B: float = 410.7,
    alpha: float = 0.34,
    beta: float = 0.28,
    E: float = 1.69,
) -> float:
    """
    Compute Chinchilla loss given parameters and effective tokens.
    """
    return A / (params**alpha) + B / (effective_tokens**beta) + E


def find_optimal_allocation(
    compute: float, unique_data: float, max_epochs: int = 32
) -> dict:
    """
    Find optimal (model_size, epochs) allocation under data constraint.

    Assumes 6 FLOPs per parameter per token.
    """
    best_loss = float("inf")
    best_config = None

    for epochs in range(1, max_epochs + 1):
        eff_data = effective_tokens(unique_data, epochs)
        tokens_processed = unique_data * epochs

        # Solve for model size: C = 6 * N * tokens
        model_size = compute / (6 * tokens_processed)

        if model_size < 1e6:  # minimum sensible model size
            continue

        loss = chinchilla_loss(model_size, eff_data)

        if loss < best_loss:
            best_loss = loss
            best_config = {
                "model_size": model_size,
                "epochs": epochs,
                "effective_tokens": eff_data,
                "total_tokens": tokens_processed,
                "loss": loss,
            }

    return best_config

The find_optimal_allocation function performs the key optimization we discussed earlier. Given a fixed compute budget and data constraint, it searches over possible epoch counts to find the configuration that minimizes final loss. For each epoch count, it calculates the corresponding model size (since compute = 6 × parameters × tokens processed), computes the effective token count using our repetition penalty formula, and evaluates the resulting loss. This brute-force search replaces the elegant closed-form solutions available under unconstrained Chinchilla scaling.

In[15]:
Code
# Compare optimal allocation with different data constraints
compute_budget = 1e23  # 10^23 FLOPs (significant training run)

data_scenarios = {
    "Abundant (1T)": 1e12,
    "Limited (100B)": 1e11,
    "Scarce (10B)": 1e10,
}

results = {}
for name, data_size in data_scenarios.items():
    results[name] = find_optimal_allocation(compute_budget, data_size)

We compare three scenarios representing different levels of data availability, from abundant (1 trillion tokens) to scarce (10 billion tokens). This comparison illuminates how the same compute budget produces dramatically different outcomes depending on data availability. Organizations planning training runs must account for this difference.

Out[16]:
Console
Data Scenarios:
  Abundant (1T): 1000B unique tokens
  Limited (100B): 100B unique tokens
  Scarce (10B): 10B unique tokens
Out[17]:
Console
Compute budget: 1e23 FLOPs

Scenario           |   Model Size | Epochs |  Eff. Tokens |   Loss
----------------------------------------------------------------------
Abundant (1T)      |       16.7B |      1 |     1000.0B |  2.005
Limited (100B)     |       23.8B |      7 |      529.0B |  2.025
Scarce (10B)       |       75.8B |     22 |       93.4B |  2.119

The results illustrate how data availability shapes optimal training strategy. With abundant data (1T tokens), the model can grow large while maintaining single-epoch training. As data becomes scarcer, the optimizer accepts more repetition epochs but the effective training signal shrinks, resulting in higher final loss despite using the same compute budget. The loss increase from 2.296 to 2.502 when moving from abundant to scarce data shows a substantial capability gap. This is equivalent to what might take 10× more compute to overcome through scaling alone. Notice also how the optimal model size shifts: with scarce data, the optimizer chooses a much larger model (174.8B vs 27.8B) that trains for more epochs, extracting maximum value from each precious unique token.

Out[18]:
Visualization
Bar chart comparing achievable loss across three data availability scenarios.
Loss achievable across different data constraints at fixed compute. The shaded region shows the 'data tax' paid when training data is limited.

The visualization makes clear that data scarcity exacts a significant toll even with optimal compute allocation. Moving from abundant to limited data increases loss by about 0.1, while moving to scarce data adds another 0.1. These differences may seem small but translate to meaningful capability gaps in practice. The annotations on each bar reveal the hidden story: as data becomes scarcer, the optimizer dramatically increases model size while accepting more epochs, but the repetition penalty means this strategy can only partially compensate for the missing unique data.

Key Parameters

The key parameters for data-constrained scaling are:

  • alpha (repetition decay): Controls how quickly the value of repeated data diminishes. Empirically fitted to approximately 0.1. Lower values mean slower decay and more tolerance for repetition.
  • epochs: Number of times the training data is repeated. Beyond 4-8 epochs, returns diminish sharply for large models.
  • unique_tokens: The amount of deduplicated, unique training data available. This is the hard constraint that cannot be circumvented through repetition.
  • A, B (scaling coefficients): Empirically fitted constants in the Chinchilla loss formula that determine the relative importance of model size vs. data.
  • alpha, beta (scaling exponents): Control how loss decreases with scale, approximately 0.34 and 0.28 respectively.