Mechanistic interpretability experiment on double descent

What is asked. Reproduce sample-wise double descent in a small model, explain why it happens from theory, show the mechanism inside the model rather than only the curve, and propose a fix that you then verify. The deliverable is an experiment someone else can rerun plus a short explanation.

The phenomenon. Hold the model fixed and grow the training set. Test error first falls, then rises to a peak when the number of training examples n is close to the number of parameters p (the interpolation threshold), then falls again. More data making a model worse is the surprising part, and the interview is about explaining it precisely.

Clarify first. How much compute is there (a laptop: so a linear or tiny model, seconds per run)? Which axis: samples (here), model width, or training time? Is there label noise (there must be some: without noise the peak mostly disappears, which is itself a good check)? What counts as "mechanistic": the weight-space story below, plus, if time allows, the toy superposition model in the follow-ups.

Smallest setting that shows it. Linear regression with p = 50 Gaussian features, a fixed true weight vector w* with ‖w*‖ ≈ 1, and label noise σ = 0.5. Fit with minimum-norm least squares (pinv(X) @ y), which is what gradient descent from zero converges to on a linear model. Sweep n from 10 to 200, 20 seeds per point, report the median test MSE on 2,000 fresh points.

Why this setting: every quantity has a closed form, so the curve can be checked against theory, and the mechanism can be read off the weights directly. A two-layer network shows the same peak but costs more and explains less.

Controls, so the result means something.

• σ = 0: the peak should mostly vanish (the blow-up comes from fitting noise).
• Ridge λ swept: the optimal λ should remove the peak.
• Test noise floor σ² = 0.25: no method can go below it.

What was measured (median over 20 seeds, p = 50, σ = 0.5):

The peak at n = p is 35× the error at n = 25 with half the data.

For isotropic Gaussian features and the min-norm solution, the expected test error has a closed form (σ² is the noise floor):

• n < p (underdetermined): ‖w*‖² (1 − n/p) + σ² · n / (p − n − 1) + σ². The first term is bias: directions the data never touched stay at zero. The second is variance from fitting noise.
• n > p (overdetermined): σ² · p / (n − p − 1) + σ².

Both variance terms diverge at n = p. Check against the table: n = 100 gives 0.25 × (1 + 50/49) = 0.505 (measured 0.50); n = 200 gives 0.25 × (1 + 50/149) = 0.334 (measured 0.35); n = 25 gives 0.5 + 0.25 × 25/24 + 0.25 = 1.01 (measured 0.95).

Why it diverges: the smallest singular value. The learned weights split exactly into two parts:

ŵ = X⁺ y = P w* + X⁺ ε

where P projects onto the directions the data spans and X⁺ ε is the part fitted to noise. ‖X⁺ε‖² = Σᵢ (uᵢ·ε)² / sᵢ² over the singular values sᵢ of X. For a random n × p matrix, the smallest singular value is about |√n − √p|, which is near 0 when n = p. Dividing noise by a tiny singular value gives huge weights along that direction, and those weights point somewhere random relative to test inputs. Away from n = p, either there are spare dimensions to absorb the noise at small norm (n < p) or there is enough data to average it out (n > p).

"Mechanistic" means showing this inside the trained model, not just matching a formula:

1. Decompose each fitted ŵ into P w* (signal) and X⁺ ε (noise fit). Run the same fit on the noise alone. The noise component's norm should peak at n = p and account for nearly all of the excess test error.
2. Track the smallest singular value of X against n: it should dip toward 0 at n = p.
3. Look at where the weight goes: the projection of ŵ onto the bottom singular direction of X carries most of ‖ŵ‖ at the peak. Those are memorized training points, not features of the task.

The tests below check all three.

Ridge regularization. Replace pinv with (XᵀX + λI)⁻¹ Xᵀ y. Ridge adds λ to every sᵢ² in the denominator, so noise is never divided by a tiny number. The Bayes-optimal λ for this setting is σ² / (per-coordinate prior variance of w*) = 0.25 / (1/50) = 12.5. Measured at n = 50: λ = 0.1 gives 2.53, λ = 1 gives 1.07, λ = 12.5 gives 0.64, λ = 50 gives 0.70. With λ tuned, the curve is monotone: more data never hurts.

Other remedies and why they work: early stopping (gradient descent fits large-singular-value directions first; stopping early leaves the tiny ones unfitted, an implicit ridge); changing p to move away from the threshold; cleaning labels (the peak scales with σ²).

import numpy as np


def make_task(p=50, seed=0):
    rng = np.random.default_rng(seed)
    return rng.normal(size=p) / np.sqrt(p)  # ||w*|| close to 1


def sample(rng, n, w_true, noise):
    x = rng.normal(size=(n, len(w_true)))
    eps = noise * rng.normal(size=n)
    return x, x @ w_true + eps, eps


def fit(x, y, ridge=0.0):
    if ridge == 0.0:
        return np.linalg.pinv(x) @ y  # minimum-norm least squares
    p = x.shape[1]
    return np.linalg.solve(x.T @ x + ridge * np.eye(p), x.T @ y)


def sweep(ns, p=50, noise=0.5, ridge=0.0, trials=20, seed=0):
    """Rows of (n, median test MSE, median ||w||, median ||noise part||, median min singular value)."""
    w_true = make_task(p, seed)
    rng = np.random.default_rng(seed + 1)
    x_test, y_test, _ = sample(rng, 2000, w_true, noise)
    rows = []
    for n in ns:
        errs, norms, noise_norms, smin = [], [], [], []
        for _ in range(trials):
            x, y, eps = sample(rng, n, w_true, noise)
            w = fit(x, y, ridge)
            errs.append(np.mean((x_test @ w - y_test) ** 2))
            norms.append(np.linalg.norm(w))
            noise_norms.append(np.linalg.norm(fit(x, eps, ridge)))  # the part fitted to noise
            smin.append(np.linalg.svd(x, compute_uv=False).min())
        rows.append((n, *(float(np.median(v)) for v in (errs, norms, noise_norms, smin))))
    return rows


def theory_mse(n, p, noise, w_norm_sq=1.0):
    if n < p - 1:
        return w_norm_sq * (1 - n / p) + noise**2 * n / (p - n - 1) + noise**2
    if n > p + 1:
        return noise**2 * p / (n - p - 1) + noise**2
    return float("inf")

1. Neural networks. Train a two-layer ReLU network of fixed width with plain gradient descent to zero training loss and sweep n; the peak appears near n ≈ number of parameters. Train long (thousands of steps): the peak only shows once the model interpolates. Adding weight decay or early stopping flattens it, the same remedy as above.
2. Superposition and memorization. In a toy model that compresses many sparse features into fewer hidden dimensions, small datasets get stored as one direction per training example (memorization), and large datasets as one direction per feature (generalization). Measure, for each training example, how much of the hidden space its representation takes, and the same for each feature. The switch between the two regimes happens around the test-loss peak. This is a mechanistic picture of the same transition the linear model shows in closed form.
3. Model-wise double descent. Fix n and grow width: the same peak appears at p ≈ n. Plot both axes as a heatmap of test error over (n, p): the ridge of high error lies along p = n.
4. Robustness of the claim. Report seed variance (the peak is heavy-tailed; use medians and show interquartile ranges), test non-isotropic features (a decaying spectrum moves and widens the peak), and check that the curve does not depend on the test-set sample.
5. Why practitioners rarely see it. Real training uses weight decay, early stopping and data augmentation, and models are usually far from the threshold. It matters most for small fine-tuning sets where n can land near the number of trainable parameters (for example, LoRA with few examples).

Mid-level (L4)

• Sets up a clean sweep over training-set size with a fixed model, enough seeds, and held-out test data.
• Reproduces the peak and explains it as fitting noise at the interpolation threshold.
• Applies regularization and shows it helps.

Senior (L5)

• Chooses the minimal model where theory is exact, includes controls (no noise, ridge sweep), and checks numbers against the closed form.
• Explains the mechanism through the smallest singular value and the noise component of the weights, and measures both.
• Picks the regularization strength from theory and confirms it empirically.

Staff (L6+)

• Connects the linear story to neural networks, superposition and memorization, and proposes the experiment that tests the connection.
• Is careful about statistical claims (heavy tails, medians, ranges) and about what would falsify the explanation.
• Says when the phenomenon matters in practice (small fine-tuning sets) and what to do about it.