diff options
Diffstat (limited to 'scripts/closed_form_soft_ramp.py')
| -rw-r--r-- | scripts/closed_form_soft_ramp.py | 280 |
1 files changed, 280 insertions, 0 deletions
diff --git a/scripts/closed_form_soft_ramp.py b/scripts/closed_form_soft_ramp.py new file mode 100644 index 0000000..ac71988 --- /dev/null +++ b/scripts/closed_form_soft_ramp.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python3 +"""Closed-form soft-ramp gap law: derive and validate against the dense sweep. + +Lazy (frozen) tangent operators, squared loss L = ||r||^2/(2N). BP residual +decays as r_t = (I - eta K_BP/N)^t r_0 with K_BP = J J^T. FA uses the symmetric +part S_FA = (K_FA + K_FA^T)/2 of the non-symmetric tangent operator +K_FA = J J_tilde^T. The actual-FA initial-moment theorem gives the residual- +direction scalar burden + + rho = output_share = r^T K_out r / r^T K_BP r = 1 - E_B[e_0] in (0,1), + +and the lazy scalar model S_FA ~= rho K_BP yields the closed form + + gap_T = (1/2N) sum_i c_i^2 [ (1 - eta rho lam_i/N)^{2T} + - (1 - eta lam_i/N)^{2T} ], + +with K_BP v_i = lam_i v_i and c_i = v_i^T r_0. Each bracket is >= 0 and analytic +in (rho, lam_i, T): the gap is a smooth functional of the spectrum, so it is a +soft ramp, with a hard threshold only in the degenerate rho -> 0 limit. + +In the BP-converged regime the gap is dominated by the slowest surviving FA mode, + gap_T ~= (c_min^2/2N) exp(-2 eta rho lam_min T/N), +so log gap_T ~= const - (2 eta T/N) rho lam_min: the ramp is driven by lam_min(w) +growing smoothly with capacity. + +This script computes the no-fit closed-form gap at initialization for each width +in the dense T=30000 sweep and compares it to the empirical ramp. The dense +run's config and measured gaps are loaded from its output directory, and the +closed form is evaluated on the *same data and the same init seed* the sweep +actually used (matched comparison), plus extra init seeds on the same data to +show the init-to-init variability of the no-fit prediction (band). +""" + +from __future__ import annotations + +import argparse +import csv +import json +import sys +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import torch + +SCRIPT_DIR = Path(__file__).resolve().parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +import downstream_capacity_sweep as dcs # noqa: E402 + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Closed-form soft-ramp gap law.") + p.add_argument( + "--dense-dir", + type=Path, + default=Path("outputs/phase_transition_dense_T30000_352traj"), + help="Dense sweep output dir (summary.json + width_summary.csv).", + ) + p.add_argument("--extra-init-seeds", type=int, default=4, + help="Extra init seeds (same data) for the variability band.") + p.add_argument("--torch-threads", type=int, default=8) + p.add_argument("--outdir", type=Path, default=Path("outputs/closed_form_soft_ramp")) + return p.parse_args() + + +def load_dense_run(dense_dir: Path) -> tuple[dict, list[dict]]: + """Load the dense sweep's config and per-width measured gaps.""" + summary = json.loads((dense_dir / "summary.json").read_text()) + dense_config = summary["config"] + rows = [] + with (dense_dir / "width_summary.csv").open() as fh: + for row in csv.DictReader(fh): + rows.append({ + "width": int(row["width"]), + "margin": int(row["fa_capacity_margin"]), + "measured_gap": float(row["train_gap_mean"]), + "measured_gap_std": float(row["train_gap_std"]), + "runs_fa": int(row["runs_fa"]), + }) + rows.sort(key=lambda r: r["width"]) + return dense_config, rows + + +def make_config(dense_config: dict, torch_threads: int, outdir: Path) -> dcs.RunConfig: + # Mirror the dense run's data-generating config exactly so make_data + # reproduces the very same x_train/y_train the sweep trained on. + return dcs.RunConfig( + task=dense_config["task"], + input_dim=dense_config["input_dim"], + output_dim=dense_config["output_dim"], + teacher_rank=dense_config["teacher_rank"], + teacher_width=dense_config["teacher_width"], + teacher_hidden_layers=dense_config["teacher_hidden_layers"], + normalize_targets=dense_config["normalize_targets"], + widths=dense_config["widths"], + train_samples=dense_config["train_samples"], + test_samples=dense_config["test_samples"], + probe_samples=dense_config["probe_samples"], + steps=dense_config["steps"], + lr=dense_config["lr"], + optimizer=dense_config.get("optimizer", "sgd"), + init_seeds=dense_config["init_seeds"], + feedback_seeds=dense_config["feedback_seeds"], + init_seed_offset=dense_config["init_seed_offset"], + feedback_seed_offset=dense_config["feedback_seed_offset"], + data_seed=dense_config["data_seed"], + noise_std=dense_config["noise_std"], + feedback_scale=dense_config["feedback_scale"], + capacity_q=dense_config["capacity_q"], + jacobian_lambda_rel=dense_config["jacobian_lambda_rel"], + skip_jacobian=True, + device="cpu", + torch_threads=torch_threads, + outdir=str(outdir), + plot=False, + ) + + +def forward(weights: list[torch.Tensor], x: torch.Tensor) -> torch.Tensor: + a = x + for layer, w in enumerate(weights): + z = a @ w.T + a = torch.relu(z) if layer < len(weights) - 1 else z + return a + + +def jacobian(weights: list[torch.Tensor], x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]: + leaves = [w.clone().requires_grad_(True) for w in weights] + pred = forward(leaves, x).reshape(-1) + n_params = sum(w.numel() for w in leaves) + jac = torch.zeros(pred.numel(), n_params, dtype=torch.float64) + for i in range(pred.numel()): + grads = torch.autograd.grad(pred[i], leaves, retain_graph=True) + jac[i] = torch.cat([g.reshape(-1) for g in grads]) + return jac, pred.detach(), leaves[-1].numel() + + +def closed_form_gap(config, x, y, width: int, init_seed: int) -> dict: + weights = dcs.initialize_weights(config, width, init_seed) + jac, pred, n_out = jacobian(weights, x) + r0 = pred - y.reshape(-1) + k_bp = jac @ jac.T + jac_out = jac[:, jac.shape[1] - n_out:] + speed_bp = float((jac.T @ r0).pow(2).sum()) + rho = float((jac_out.T @ r0).pow(2).sum() / speed_bp) + lam, vecs = torch.linalg.eigh(k_bp) + coeff = vecs.T @ r0 + eta, n, t = config.lr, config.train_samples, config.steps + a_bp = 1.0 - eta * lam / n + a_fa = 1.0 - eta * rho * lam / n + # only modes the discrete recursion does not blow up (|a|<=1); the real SGD + # run was stable, so unstable modes carry negligible residual mass. + bp = torch.where(a_bp.abs() <= 1, a_bp.pow(2 * t), torch.zeros_like(lam)) + fa = torch.where(a_fa.abs() <= 1, a_fa.pow(2 * t), torch.zeros_like(lam)) + gap = float((coeff.pow(2) * (fa - bp)).sum() / (2 * n)) + pos = lam[lam > 1e-9] + return {"rho": rho, "lam_min": float(pos.min()), "lam_max": float(lam.max()), + "pred_gap": gap} + + +def main() -> None: + args = parse_args() + torch.set_num_threads(args.torch_threads) + dense_config, dense_rows = load_dense_run(args.dense_dir) + config = make_config(dense_config, args.torch_threads, args.outdir) + args.outdir.mkdir(parents=True, exist_ok=True) + + # The sweep's init seed is 10_000 + init_seed_offset + init_index + # (downstream_capacity_sweep.main); the dense run used init_seeds=1. + matched_seed = 10_000 + dense_config["init_seed_offset"] + band_seeds = [matched_seed + k for k in range(1, args.extra_init_seeds + 1)] + x, y, *_ = dcs.make_data(config) + + rows, seed_rows = [], [] + print(f"matched init seed = {matched_seed}, band seeds = {band_seeds}") + print(f"{'w':>3}{'margin':>7}{'rho':>7}{'lam_min':>9}{'pred_gap':>11}{'measured':>11}") + for dense_row in dense_rows: + width, margin = dense_row["width"], dense_row["margin"] + meas, meas_std = dense_row["measured_gap"], dense_row["measured_gap_std"] + d = closed_form_gap(config, x, y, width, matched_seed) + band = [] + for seed in band_seeds: + extra = closed_form_gap(config, x, y, width, seed) + band.append(extra["pred_gap"]) + seed_rows.append({"width": width, "margin": margin, "init_seed": seed, + "rho": extra["rho"], "lam_min": extra["lam_min"], + "pred_gap": extra["pred_gap"]}) + seed_rows.append({"width": width, "margin": margin, "init_seed": matched_seed, + "rho": d["rho"], "lam_min": d["lam_min"], + "pred_gap": d["pred_gap"]}) + all_preds = band + [d["pred_gap"]] + width_seed_rows = [r for r in seed_rows if r["width"] == width] + rows.append({"width": width, "margin": margin, "rho": d["rho"], + "lam_min": d["lam_min"], "lam_max": d["lam_max"], + "pred_gap": d["pred_gap"], "measured_gap": meas, + "measured_gap_std": meas_std, + "pred_gap_band_lo": min(all_preds), + "pred_gap_band_hi": max(all_preds), + "pred_gap_geo_mean": float(np.exp(np.mean( + [np.log(r["pred_gap"]) for r in width_seed_rows]))), + "rho_seed_mean": float(np.mean( + [r["rho"] for r in width_seed_rows])), + "lam_min_seed_mean": float(np.mean( + [r["lam_min"] for r in width_seed_rows]))}) + print(f"{width:>3}{margin:>7}{d['rho']:>7.3f}{d['lam_min']:>9.2f}" + f"{d['pred_gap']:>11.5f}{meas:>11.5f}", flush=True) + + csv_path = args.outdir / "closed_form_vs_measured.csv" + with csv_path.open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(rows[0].keys())) + w.writeheader() + w.writerows(rows) + seed_csv_path = args.outdir / "closed_form_per_seed.csv" + with seed_csv_path.open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(seed_rows[0].keys())) + w.writeheader() + w.writerows(seed_rows) + + margins = np.array([r["margin"] for r in rows]) + measured = np.array([r["measured_gap"] for r in rows]) + measured_std = np.array([r["measured_gap_std"] for r in rows]) + pred = np.array([r["pred_gap"] for r in rows]) + band_lo = np.array([r["pred_gap_band_lo"] for r in rows]) + band_hi = np.array([r["pred_gap_band_hi"] for r in rows]) + lam_min = np.array([r["lam_min"] for r in rows]) + pred_geo = np.array([r["pred_gap_geo_mean"] for r in rows]) + lam_min_mean = np.array([r["lam_min_seed_mean"] for r in rows]) + lm = np.log(measured) + fit = np.polyfit(margins, lm, 1) + corr_pred = np.corrcoef(np.log(np.clip(pred, 1e-12, None)), lm)[0, 1] + corr_geo = np.corrcoef(np.log(np.clip(pred_geo, 1e-12, None)), lm)[0, 1] + corr_lam = np.corrcoef(lam_min, -lm)[0, 1] + corr_lam_mean = np.corrcoef(lam_min_mean, -lm)[0, 1] + r2 = np.corrcoef(margins, lm)[0, 1] ** 2 + + fig, axes = plt.subplots(1, 2, figsize=(12.0, 4.8), dpi=180) + axes[0].errorbar(margins, measured, yerr=measured_std, fmt="o-", color="#174f91", + capsize=2.5, label="empirical (T=30000, 32 feedback seeds)") + axes[0].semilogy(margins, pred_geo, "s-", color="#c65f16", + label=f"closed form, {len(band_seeds) + 1}-init geo mean (no fit)") + axes[0].semilogy(margins, pred, "^:", color="#7a4a13", alpha=0.7, markersize=4, + label="closed form, matched init") + axes[0].fill_between(margins, np.clip(band_lo, 1e-12, None), band_hi, + color="#c65f16", alpha=0.18, + label="init-seed range") + axes[0].set_yscale("log") + axes[0].axvline(0, color="grey", linewidth=1, linestyle=":") + axes[0].set_xlabel("hard FA capacity margin") + axes[0].set_ylabel("FA/BP train gap") + axes[0].set_title(f"Soft ramp: closed form vs empirical\ncorr(log,log)={corr_geo:.3f} (init mean), no kink at margin 0") + axes[0].legend(fontsize=8) + axes[0].grid(alpha=0.18, which="both") + + axes[1].plot(lam_min_mean, -lm, "o", color="#174f91") + coef = np.polyfit(lam_min_mean, -lm, 1) + grid = np.linspace(lam_min_mean.min(), lam_min_mean.max(), 50) + axes[1].plot(grid, np.polyval(coef, grid), "--", color="#c65f16") + axes[1].set_xlabel("smallest BP-NTK eigenvalue, init mean lam_min(w)") + axes[1].set_ylabel("-log(empirical gap)") + axes[1].set_title(f"Ramp law: log gap ~ const - c.rho.lam_min\ncorr={corr_lam_mean:.3f}") + axes[1].grid(alpha=0.18) + fig.tight_layout() + fig_path = args.outdir / "closed_form_vs_measured_ramp.png" + fig.savefig(fig_path) + plt.close(fig) + + print(f"\nmeasured log(gap) ~ {fit[0]:.4f}*margin + {fit[1]:.2f} (R^2={r2:.3f})") + print(f"corr(log pred, log measured) = {corr_pred:.3f} (matched init)") + print(f"corr(log pred, log measured) = {corr_geo:.3f} ({len(band_seeds) + 1}-init geo mean)") + print(f"corr(lam_min, -log measured) = {corr_lam:.3f} (matched init)") + print(f"corr(lam_min, -log measured) = {corr_lam_mean:.3f} (init mean)") + print(f"rows: {csv_path}") + print(f"per-seed rows: {seed_csv_path}") + print(f"plot: {fig_path}") + + +if __name__ == "__main__": + main() |