Compute-Optimal Scaling Fits

The problem, from first principles

For a fixed compute budget you can train a small model on many tokens or a large model on few. Plotting the final loss against model size for one budget (an IsoFLOP curve) gives a U-shape: too small underfits, too large is undertrained. The bottom of the U is the compute-optimal model size for that budget. Repeating this for several budgets gives pairs of (budget, optimal size), and the relationship is a power law N_opt = a * C^b, which is a straight line on log-log axes. Fitting this line is how scaling-law papers extrapolate to budgets too large to run. Two small fits are needed: a parabola for each IsoFLOP curve and a line for the power law.

From theory to code

Implement isoflop_optimum and fit_power_law.

Constraints

• isoflop_optimum(params, losses): fit a quadratic loss = a * x**2 + b * x + c with x = ln(params) by least squares (np.polyfit, degree 2) and return the minimizing size exp(-b / (2 * a)). Require a > 0 and raise ValueError otherwise.
• fit_power_law(x, y): fit ln y = ln a + b * ln x by least squares and return (a, b) with a = exp(intercept).
• Inputs are NumPy arrays of positive numbers with enough points for the fits.

Hints