Data-Constrained Scaling › Code Implementation
Language AI HandbookAll chapters 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.
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_rThis 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 modulates the raw token count to produce an effective count that grows more slowly than the naive expectation.
# 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 expectationEpoch | 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.
(0.0, 3.5)

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.
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_configThe 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.
# 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.
Data Scenarios: Abundant (1T): 1000B unique tokens Limited (100B): 100B unique tokens Scarce (10B): 10B unique tokens
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.

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.