Below: two small, self-contained exercises. First, fit to a handful of (compute, loss) points using nonlinear least squares — this is literally what appears in Figure 1 of every scaling-law paper. Second, given the fitted from the two-variable form, solve the constrained-optimization problem for the compute-optimal (N, D) split at a target budget.
python
import numpy as np
from scipy.optimize import curve_fit
# Synthetic but realistic: compute (FLOPs) and final validation loss
# from a family of small training runs.
compute = np.array([1e18, 3e18, 1e19, 3e19, 1e20, 3e20, 1e21])
loss = np.array([3.90, 3.55, 3.22, 2.95, 2.70, 2.49, 2.31])
def power_law(C, a, alpha, L_inf):
return a * C ** (-alpha) + L_inf
# Fit in log-space is more numerically stable, but curve_fit with
# decent initial guesses works fine here.
p0 = [50.0, 0.05, 1.6]
params, cov = curve_fit(power_law, compute, loss, p0=p0, maxfev=20000)
a, alpha, L_inf = params
print(f"a={a:.3f} alpha={alpha:.4f} L_inf={L_inf:.3f}")
# Extrapolate to a training run 1000x bigger than anything we measured.
C_target = 1e24
predicted_loss = power_law(C_target, *params)
print(f"predicted loss at C={C_target:.0e}: {predicted_loss:.3f}")
# Sanity check: residuals should look like noise, not a curved pattern,
# or the functional form is wrong (e.g. missing a second regime).
residuals = loss - power_law(compute, *params)
print("residuals:", np.round(residuals, 4))