repro-bundle/scripts/independent_scaling.py

#!/usr/bin/env python3
"""Independent finite-sample checks for Theorems 3 and 5 of arXiv:2606.01292.

Unlike the first logbook pass, this script does not evaluate the theorem's
closed-form risk proxy.  It draws Gaussian training examples, runs minibatch
SGD for every sample size and seed, and estimates population risks directly
from the learned coefficient vectors.
"""

from __future__ import annotations

import csv
import json
from pathlib import Path

import numpy as np


OUT = Path("independent_results")
SEED = 260601292
REPS = 64


def write_csv(name: str, rows: list[dict]) -> None:
    OUT.mkdir(parents=True, exist_ok=True)
    with (OUT / name).open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def train_semantic_regressor(
    *,
    alpha: float,
    target: np.ndarray,
    samples: int,
    reps: int,
    noise_sd: float,
    seed: int,
    batch: int = 32,
    eta: float = 0.12,
    paper_schedule: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
    """Run fresh finite-sample SGD and return semantic prediction vectors."""
    rng = np.random.default_rng(seed)
    d = len(target)
    lam = np.arange(1, d + 1, dtype=np.float64) ** (-alpha)
    roots = np.sqrt(lam)
    # Model is <u, sqrt(Lambda) z>; target is <target, z>.
    weights = np.zeros((reps, d), dtype=np.float64)
    updates = (samples + batch - 1) // batch
    if paper_schedule:
        # Algorithm 1 uses K=N/log2(N) samples between halvings.  For
        # sample-wise SGD this is exactly the schedule below.
        decay_every = max(1, int(samples / np.log2(samples) / batch))
        decay = 0.5
    else:
        decay_every = max(1, updates // 5)
        decay = 0.55
    for update in range(updates):
        size = min(batch, samples - update * batch)
        if size <= 0:
            break
        z = rng.normal(size=(reps, size, d))
        features = z * roots
        labels = np.einsum("rbd,d->rb", z, target)
        if noise_sd:
            labels += noise_sd * rng.normal(size=(reps, size))
        prediction = np.einsum("rbd,rd->rb", features, weights)
        gradient = np.mean(features * (prediction - labels)[..., None], axis=1)
        learning_rate = eta * (decay ** (update // decay_every))
        weights -= learning_rate * gradient
    semantic = weights * roots
    risk = 0.5 * np.sum((semantic - target[None, :]) ** 2, axis=1)
    return semantic, risk


def claim2_der() -> dict:
    """Estimate DER from independently trained teacher/student SGD models."""
    d = 256
    target = np.ones(d, dtype=np.float64) / np.sqrt(d)
    sample_grid = np.array([256, 512, 1024, 2048, 4096, 8192, 16384])
    alpha_teacher, alpha_student, beta = 1.5, 3.0, 0.0
    expected = (alpha_teacher - 1 - beta) * (
        1 / alpha_teacher - 1 / alpha_student
    )
    rows: list[dict] = []
    ders = []
    for idx, n in enumerate(sample_grid):
        teacher, teacher_risk = train_semantic_regressor(
            alpha=alpha_teacher,
            target=target,
            samples=int(n),
            reps=REPS,
            noise_sd=0.15,
            seed=SEED + 1000 * idx,
        )
        _, student_risk = train_semantic_regressor(
            alpha=alpha_student,
            target=target,
            samples=int(n),
            reps=REPS,
            noise_sd=0.15,
            seed=SEED + 1000 * idx + 1,
        )
        # With sufficient unlabeled transfer samples, a full-rank compatible
        # student copies the teacher's semantic predictor.  Its population
        # risk is therefore measured from the fitted teacher, not a formula.
        transfer_risk = 0.5 * np.sum((teacher - target[None, :]) ** 2, axis=1)
        der_seed = student_risk / transfer_risk
        ders.append(float(np.mean(der_seed)))
        rows.append(
            {
                "N": int(n),
                "repetitions": REPS,
                "direct_student_risk_mean": float(np.mean(student_risk)),
                "direct_student_risk_se": float(np.std(student_risk, ddof=1) / np.sqrt(REPS)),
                "teacher_asymptotic_transfer_risk_mean": float(np.mean(transfer_risk)),
                "teacher_asymptotic_transfer_risk_se": float(np.std(transfer_risk, ddof=1) / np.sqrt(REPS)),
                "DER_mean_of_seed_ratios": float(np.mean(der_seed)),
                "DER_median": float(np.median(der_seed)),
                "fraction_DER_above_one": float(np.mean(der_seed > 1)),
            }
        )
    slope = float(np.polyfit(np.log(sample_grid[-4:]), np.log(ders[-4:]), 1)[0])
    write_csv("claim2_independent_sgd.csv", rows)
    return {
        "expected_asymptotic_kappa": expected,
        "finite_sample_tail_slope": slope,
        "DER_first": ders[0],
        "DER_last": ders[-1],
        "all_grid_mean_DER_above_one": bool(all(value > 1 for value in ders)),
        "last_grid_fraction_seed_DER_above_one": rows[-1]["fraction_DER_above_one"],
        "note": "Finite d=256 is expected to flatten before the infinite-dimensional asymptote.",
    }


def claim4_pgr() -> dict:
    """Estimate W2S PGR by SGD teacher fitting plus student early stopping."""
    d, k_dagger = 100, 10
    alpha_teacher, alpha_student = 1.5, 2.0
    target = np.zeros(d, dtype=np.float64)
    target[:k_dagger] = 0.25
    sample_grid = np.array([250, 500, 1000, 2000, 4000, 8000, 16000])
    student_lam = np.arange(1, d + 1, dtype=np.float64) ** (-alpha_student)
    # Integrated learning-rate grid controls the student's spectral cutoff.
    times = np.geomspace(0.05, 2.0e5, 500)
    filters = 1.0 - np.exp(-times[:, None] * student_lam[None, :])
    rows: list[dict] = []
    gaps = []
    for idx, n in enumerate(sample_grid):
        teacher, teacher_risk = train_semantic_regressor(
            alpha=alpha_teacher,
            target=target,
            samples=int(n),
            reps=REPS,
            noise_sd=1.0,
            seed=SEED + 50000 + 1000 * idx,
            batch=1,
            eta=0.08,
            paper_schedule=True,
        )
        # Evaluate every early-stopping filter on each independently fitted
        # teacher and select one global checkpoint by mean population risk.
        error = filters[:, None, :] * teacher[None, :, :] - target[None, None, :]
        transfer_risk_by_time_seed = 0.5 * np.sum(error**2, axis=2)
        mean_curve = np.mean(transfer_risk_by_time_seed, axis=1)
        best_idx = int(np.argmin(mean_curve))
        transfer_seed = transfer_risk_by_time_seed[best_idx]
        # Fully ground-truth-trained expressive student has zero population
        # approximation error in this synthetic model, so 1-PGR=R_T2S/R_T.
        gap_seed = transfer_seed / teacher_risk
        gaps.append(float(np.mean(gap_seed)))
        rows.append(
            {
                "N": int(n),
                "repetitions": REPS,
                "teacher_risk_mean": float(np.mean(teacher_risk)),
                "optimal_student_time": float(times[best_idx]),
                "transfer_risk_mean": float(np.mean(transfer_seed)),
                "one_minus_PGR_mean_of_seed_ratios": float(np.mean(gap_seed)),
                "PGR_mean": float(1 - np.mean(gap_seed)),
                "fraction_student_beats_teacher": float(np.mean(transfer_seed < teacher_risk)),
            }
        )
    empirical_delta = float(-np.polyfit(np.log(sample_grid[-4:]), np.log(gaps[-4:]), 1)[0])
    expected_delta = 2 * alpha_student / (
        alpha_teacher * (2 * alpha_student + 1)
    )
    write_csv("claim4_independent_pgr.csv", rows)
    return {
        "expected_asymptotic_delta": expected_delta,
        "finite_sample_tail_delta": empirical_delta,
        "PGR_first": 1 - gaps[0],
        "PGR_last": 1 - gaps[-1],
        "last_grid_fraction_student_beats_teacher": rows[-1]["fraction_student_beats_teacher"],
        "all_grid_positive_PGR": bool(all(value < 1 for value in gaps)),
    }


def main() -> None:
    metrics = {
        "design": {
            "repetitions": REPS,
            "seed": SEED,
            "fresh_examples_for_every_N_and_repetition": True,
            "risk": "exact population risk from fitted coefficient vectors",
        },
        "claim2": claim2_der(),
        "claim4": claim4_pgr(),
    }
    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()