Predicting Model Performance › Uncertainty Quantification

import numpy as np
from scipy.optimize import curve_fit

# Re-define necessary variables for this cell
training_runs = [
    (10, 2, 3.45),
    (10, 5, 3.28),
    (10, 10, 3.15),
    (50, 10, 2.95),
    (50, 25, 2.78),
    (50, 50, 2.65),
    (100, 20, 2.75),
    (100, 50, 2.58),
    (100, 100, 2.48),
    (300, 60, 2.52),
    (300, 150, 2.35),
    (300, 300, 2.25),
]
N_params = np.array([r[0] * 1e6 for r in training_runs])
D_tokens = np.array([r[1] * 1e9 for r in training_runs])
losses = np.array([r[2] for r in training_runs])


def chinchilla_loss(X, E, A, B, alpha, beta):
    N, D = X
    return E + A / np.power(N, alpha) + B / np.power(D, beta)


X_data = np.array([N_params, D_tokens])
p0 = [1.5, 400, 400, 0.34, 0.28]
bounds = ([0, 0, 0, 0.1, 0.1], [3, 1e6, 1e6, 0.8, 0.8])
params, covariance = curve_fit(
    chinchilla_loss, X_data, losses, p0=p0, bounds=bounds, maxfev=10000
)

# Standard errors from covariance matrix diagonal
# The diagonal of the covariance matrix contains variances; sqrt gives standard errors
param_stderrs = np.sqrt(np.diag(covariance))
param_names = ["E", "A", "B", "alpha", "beta"]

def predict_with_uncertainty(
    N_millions, D_billions, params, covariance, n_samples=1000
):
    """
    Predict loss with uncertainty via Monte Carlo sampling from the parameter posterior.

    The idea: instead of using point estimates for scaling law parameters,
    we sample from their joint distribution (approximated as multivariate normal)
    and compute loss for each sample. This propagates parameter uncertainty
    through to prediction uncertainty.

    Mathematically, if $\theta \sim \mathcal{N}(\hat{\theta}, \Sigma)$ represents our uncertain parameters, then:

    $$L(N,D) \approx \mathbb{E}[f(N, D; \theta)] \pm \text{std}[f(N, D; \theta)]$$

    where $f$ is the scaling law and the expectation/std are over parameter samples.

    Args:
        N_millions: Model size in millions of parameters
        D_billions: Training data in billions of tokens
        params: Fitted parameter estimates [E, A, B, α, β]
        covariance: Parameter covariance matrix from curve fitting
        n_samples: Number of Monte Carlo samples (more = tighter estimates)

    Returns dict with mean, std, and 95% confidence interval bounds.
    """
    N = N_millions * 1e6
    D = D_billions * 1e9

    # Sample parameters from multivariate normal
    param_samples = np.random.multivariate_normal(params, covariance, n_samples)

    # Compute loss for each parameter sample
    loss_samples = []
    for p in param_samples:
        E, A, B, alpha, beta = p
        # Skip unphysical samples
        if alpha <= 0 or beta <= 0 or A < 0 or B < 0:
            continue
        loss = E + A / np.power(N, alpha) + B / np.power(D, beta)
        if loss > 0:  # Must be positive
            loss_samples.append(loss)

    loss_samples = np.array(loss_samples)

    return {
        "mean": np.mean(loss_samples),
        "std": np.std(loss_samples),
        "ci_lower": np.percentile(loss_samples, 2.5),
        "ci_upper": np.percentile(loss_samples, 97.5),
    }

import numpy as np
from scipy.optimize import curve_fit


def holdout_validation(
    training_runs, holdout_fraction=0.3, p0=p0, bounds=bounds
):
    """
    Validate scaling law by holding out largest runs.

    This tests whether scaling laws fitted on smaller models can predict
    performance of larger models - the main use case for extrapolation.

    Args:
        training_runs: List of (params_millions, tokens_billions, loss) tuples
        holdout_fraction: Fraction of largest runs to hold out for testing

    Returns:
        Dict with train/test counts, individual errors, MAE, and bias
    """
    # Sort by scale (parameters × tokens)
    sorted_runs = sorted(training_runs, key=lambda r: r[0] * r[1])

    # Hold out largest runs
    n_holdout = max(1, int(len(sorted_runs) * holdout_fraction))
    train_runs = sorted_runs[:-n_holdout]
    test_runs = sorted_runs[-n_holdout:]

    # Fit on training subset
    N_train = np.array([r[0] * 1e6 for r in train_runs])
    D_train = np.array([r[1] * 1e9 for r in train_runs])
    L_train = np.array([r[2] for r in train_runs])

    try:
        fit_params, _ = curve_fit(
            chinchilla_loss,
            np.array([N_train, D_train]),
            L_train,
            p0=p0,
            bounds=bounds,
            maxfev=10000,
        )

        # Predict on holdout
        errors = []
        for r in test_runs:
            N, D, actual = r[0] * 1e6, r[1] * 1e9, r[2]
            predicted = chinchilla_loss(np.array([[N], [D]]), *fit_params)[0]
            errors.append(predicted - actual)

        return {
            "train_runs": len(train_runs),
            "test_runs": len(test_runs),
            "errors": errors,
            "mae": np.mean(np.abs(errors)),
            "bias": np.mean(errors),
        }
    except:
        return None