Fit a scaling law and derive the compute-optimal allocation

CodingAdvanced⏱ 3 hoursDL Theoryscaling lawsfittingcompute optimalextrapolation

Theory: Why Does Deep Learning Optimization Work? →

Problem

Given synthetic loss data of the form $\mathcal L(N,D) = E + AN^{-\alpha}+BD^{-\beta}$ with noise, fit the five parameters and derive the compute-optimal allocation.

Fit by nonlinear least squares in log space and report parameter uncertainties.
Derive $N^\star(C)$ and $D^\star(C)$ from the fit.
Show how sensitive the allocation exponents are to the fitted $\alpha,\beta$ .
Demonstrate the extrapolation risk by fitting on a restricted range and predicting outside it.

Implement it, then check the implementation against something you can verify independently.

Write it

Solve it here10 tests · Python runs in your browser

import numpy as np

def scaling_loss(N, D, E, A, alpha, B, beta):
    """L(N, D) = E + A * N**(-alpha) + B * D**(-beta).

Accept scalars or arrays for N and D and return the same shape.
    """
    raise NotImplementedError

def fit_scaling_law(N, D, L):
    """Recover the five parameters from observed losses.

Return a dict with keys 'E', 'A', 'alpha', 'B', 'beta'.

You do not need a general nonlinear optimiser. Hold (alpha, beta) fixed
    and the model is *linear* in (E, A, B) — one least-squares solve gives
    the best three, exactly. So the search is only ever two-dimensional:
    sweep (alpha, beta), keep the pair with the smallest residual, then
    refine around it.
    """
    raise NotImplementedError

def allocation_exponents(alpha, beta):
    """The compute-optimal scaling exponents (a, b), with N* ~ C**a and
    D* ~ C**b. They must sum to 1, because C is proportional to N*D."""
    raise NotImplementedError

def compute_optimal(params, C, flops_per_token=6):
    """Minimise the fitted loss subject to the compute budget C = k*N*D,
    with k = flops_per_token. Return (N_star, D_star).

Substituting D = C/(k*N) leaves a one-dimensional convex problem whose
    stationary point has a closed form.
    """
    raise NotImplementedError

⌘/Ctrl + ↵ to run · limit 120s · draft saved in this browser

10 tests will run against your code. The first run downloads the Python runtime — several megabytes, once per browser.

Solution

Opening this before you have a real attempt on paper is the fastest way to feel like you understand something you cannot reproduce.

Problem	Field	Type	Difficulty	Time
Derive the compute-optimal allocation	LLMs	Proof	Intermediate	30 min
Implicit bias of gradient descent on underdetermined least squares	DL Theory	Proof	Intermediate	30 min
The irreducible loss is the conditional entropy	LLMs	Proof	Intermediate	30 min
What fitting random labels rules out	DL Theory	Proof	Intermediate	30 min
Rotary embeddings encode relative position exactly	Deep Learning	Proof	Advanced	30 min

Problem

Write it

Solution

Related problems