summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorYurenHao0426 <Blackhao0426@gmail.com>2026-06-09 15:26:47 -0500
committerYurenHao0426 <Blackhao0426@gmail.com>2026-06-09 15:26:47 -0500
commit7dafd33f6e1847af9b6cffd3082335a2b9487762 (patch)
tree53d4086595da7b71020b60e67aa3daf2e9f6b52b
parent6188c3499bbbea6b25c975de4134666021f458a7 (diff)
Add closed-form soft-ramp theorem and matched comparison
Closed-form gap law note (34) + consolidated snapshot (35), with the comparison script reading the dense run's config/gaps from its output dir and evaluating the closed form on matched data and init seed plus a multi-init band. Corrected headline: corr(log,log) 0.933 matched / 0.982 five-init geometric mean; lam_min ensemble corr 0.987. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
-rw-r--r--notes/34_closed_form_soft_ramp_theorem.md143
-rw-r--r--notes/35_current_results_consolidated.md149
-rw-r--r--scripts/closed_form_soft_ramp.py280
3 files changed, 572 insertions, 0 deletions
diff --git a/notes/34_closed_form_soft_ramp_theorem.md b/notes/34_closed_form_soft_ramp_theorem.md
new file mode 100644
index 0000000..3c3eef4
--- /dev/null
+++ b/notes/34_closed_form_soft_ramp_theorem.md
@@ -0,0 +1,143 @@
+# Closed-Form Soft-Ramp Gap Law
+
+> **Correction (2026-06-09, note 37):** the verification table below was
+> computed on mismatched data/init seeds. Matched rerun: corr(log,log) =
+> **0.933** (matched single init) / **0.982** (5-init geometric mean);
+> `corr(lam_min_mean, -log gap) = 0.987`. Use these numbers, not 0.977.
+
+This note gives a closed-form burden that quantifies the finite-time FA/BP train
+gap and explains the empirical soft ramp (notes 22, 25) as a smooth spectral
+phenomenon rather than a phase transition. It connects the exact initial-erosion
+theorem (note 29) to the finite-time gap and sits one level below the early
+operator-velocity estimator (notes 15-17): the closed form gives the *mechanism
+and shape*; the estimator supplies the drift-corrected *magnitude*.
+
+## Assumptions
+
+**A1 (lazy / frozen tangent operators).** Over the horizon the tangent operators
+are approximately frozen at initialization, so for squared loss
+`L = ||r||^2/(2N)` the BP residual evolves as
+
+```text
+r_t^BP = (I - eta K_BP / N)^t r_0, K_BP = J J^T.
+```
+
+This is the standard lazy/NTK approximation. Operators actually drift; that
+residual is exactly the `K_s - K_0` the estimator corrects (notes 13, 15).
+
+**A2 (scalar erosion model).** The FA tangent operator `K_FA = J J_tilde^T` is
+non-symmetric; only its symmetric part `S_FA = (K_FA + K_FA^T)/2` drives the
+first-order loss change. We model
+
+```text
+S_FA ~= rho K_BP, rho = output_share = r^T K_out r / r^T K_BP r in (0,1).
+```
+
+This is *exact in the residual direction* by the initial-moment theorem
+(note 29): `r^T E_B[S_FA] r = ||g_out^BP||^2 = r^T K_out r = rho * r^T K_BP r`,
+i.e. `rho = 1 - E_B[e_0]`. Treating `S_FA` as a global scalar multiple of
+`K_BP` extends this identity to all eigenmodes and is the one approximation
+beyond A1.
+
+## Theorem (closed-form FA/BP gap)
+
+Under A1-A2, in the BP eigenbasis `K_BP v_i = lam_i v_i` with `c_i = v_i^T r_0`,
+
+```text
+gap_T = L_FA(T) - L_BP(T)
+ = (1/2N) sum_i c_i^2 [ (1 - eta rho lam_i / N)^{2T}
+ - (1 - eta lam_i / N)^{2T} ].
+```
+
+**Proof.** `L_rule(T) = ||r_T^rule||^2/(2N)`. Expanding `r_0 = sum_i c_i v_i`,
+
+```text
+r_T^BP = sum_i v_i (1 - eta lam_i/N)^T c_i
+=> ||r_T^BP||^2 = sum_i c_i^2 (1 - eta lam_i/N)^{2T}.
+```
+
+Under A2, `S_FA = rho K_BP` shares eigenvectors `v_i` with eigenvalues
+`rho lam_i`, so
+
+```text
+r_T^FA = sum_i v_i (1 - eta rho lam_i/N)^T c_i
+=> ||r_T^FA||^2 = sum_i c_i^2 (1 - eta rho lam_i/N)^{2T}.
+```
+
+Subtract and divide by `2N`. ∎
+
+## Corollary 1 (soft, not hard)
+
+Because `rho < 1`, every mode decays no faster under FA than under BP, so each
+bracket is `>= 0` and `gap_T >= 0`. `gap_T` is real-analytic in `(rho, {lam_i},
+T)`. Therefore under any smooth deformation of the spectrum (e.g. width or
+capacity margin) the gap moves smoothly: it is a **soft ramp**. A discontinuous
+transition would require `rho -> 0` (FA fully stalls), which never happens since
+`rho >= output_share > 0`. The hard-margin threshold
+`Delta d_hard = max(0, k - (P - d))` is recovered only in that degenerate limit.
+
+This is the operator-level version of the "no free capacity" statement
+(note 26): nonzero `rho < 1` gives a positive contribution from every mode at
+every margin.
+
+## Corollary 2 (the ramp law: where it comes from)
+
+Let `tau_i = N / (eta lam_i)`.
+
+- **Underparameterized side** (many `tau_i >~ T`): each bracket
+ `~= 2 eta T (1 - rho) lam_i / N`, so
+
+ ```text
+ gap_T ~= (1 - rho) (eta T / N^2) sum_{slow i} c_i^2 lam_i,
+ ```
+
+ residual-mass limited; large at negative margin (ill-conditioned spectrum,
+ residual stuck in slow modes).
+
+- **Converged side** (BP fit, `L_BP(T) ~= 0`): the gap is the slowest surviving
+ FA mode,
+
+ ```text
+ gap_T ~= (c_min^2 / 2N) exp(-2 eta rho lam_min T / N),
+ i.e. log gap_T ~= const - (2 eta T / N) rho lam_min.
+ ```
+
+The soft ramp is therefore driven by `lam_min(w)` (smallest BP-NTK eigenvalue)
+growing smoothly with capacity; the two regimes cross over smoothly.
+
+## Verification (no fit)
+
+Script: `scripts/closed_form_soft_ramp.py`.
+Figure: `outputs/closed_form_soft_ramp/closed_form_vs_measured_ramp.png`.
+Rows: `outputs/closed_form_soft_ramp/closed_form_vs_measured.csv`.
+
+Frozen-initialization closed form (`rho`, `lam_i`, `c_i` all measured at init,
+no fitted parameters) versus the dense `T=30000` ramp (note 22):
+
+| quantity | value |
+|---|---|
+| `corr(log pred, log measured)` across `w=20..40` | **0.977** |
+| measured `log(gap) ~ -0.021*margin`, fit `R^2` | **0.993** (log-linear ramp) |
+| `lam_min(w)` over the sweep | 0.07 -> 0.99 |
+| `corr(lam_min, -log measured)` | **0.963** (ramp driven by `lam_min`) |
+
+## Scope / honesty
+
+- Frozen-init captures the ramp **shape** (corr 0.977) but **compresses the
+ range**: it under-predicts the large-gap (negative-margin) end and
+ over-predicts the small-gap (positive-margin) tail. That residual is operator
+ drift `K_s - K_0`; the early-velocity estimator (notes 15-17) corrects it and
+ supplies the quantitative magnitude. Closed form = mechanism + shape;
+ estimator = number.
+- A2 is exact only in the residual direction; a per-mode `rho_i` (spectral
+ erosion) refinement is the natural next step.
+- Lazy regime, full-batch SGD, synthetic random-label MLP scope, as elsewhere.
+
+## Role in the paper
+
+This is the quantitative core of contribution 2 (soft erosion). It converts the
+exact scalar burden `rho = 1 - E_B[e_0]` (note 29) into a finite-time gap and
+proves the soft ramp is a spectral consequence of `rho < 1` plus a smoothly
+growing `lam_min`, not a redundancy-exhaustion threshold. Pairs with: note 29
+(supplies `rho`), notes 15-17 (drift-corrected magnitude), notes 22/25 (the
+empirical ramp it explains).
diff --git a/notes/35_current_results_consolidated.md b/notes/35_current_results_consolidated.md
new file mode 100644
index 0000000..f22a0c1
--- /dev/null
+++ b/notes/35_current_results_consolidated.md
@@ -0,0 +1,149 @@
+# Current Results (Consolidated)
+
+Authoritative snapshot of where the project stands. Supersedes the framing in
+`00_project_overview.md` and `03_paper_outline.md` (both stale: they still say
+"redundancy-exhaustion **phase transition**" and omit the estimator). The
+current, correct framing lives here and in notes 18, 24, 26, 29, 34.
+
+## Central claim
+
+> Random feedback imposes a quantifiable, **soft** (non-threshold) operator
+> burden relative to BP that is present at every capacity margin (including
+> overparameterized) and grows smoothly as capacity tightens. The static
+> alignment geometry is exactly characterized (Beta law, prior-free minimax);
+> the initial operator erosion is exactly the hidden-layer BP speed share; a
+> closed-form spectral law turns that scalar burden into the finite-time gap and
+> explains the soft ramp; and an early operator-velocity estimator predicts the
+> finite-time FA/BP gap with no fit (corr 0.999).
+
+## The erosion ladder (organizing spine)
+
+All "burden" objects, as one ladder (note 24 is the through-line: *random
+feedback induces a distribution of residual-direction operator erosion `e(r)`;
+capacity controls that distribution; the finite-time gap is the product of the
+tangent operators*):
+
+| level | object | status | role |
+|---|---|---|---|
+| 0 static geometry | `Q_l = cos^2(W^T,B) ~ Beta(1/2,(D-1)/2)` | exact | defines alignment |
+| 1 null model | `e ~ Beta(k/2,(P-k)/2)`, `E[e]=k/P` | exact, **overpredicts** | proves nonzero cost for any `k>0` |
+| 2 exact init | `e_0 = 1 - r^T S_FA r / r^T K_BP r`, `E_B[e_0]=hidden_share` | exact, validated | the real `t=0` burden; gives `rho=1-E_B[e_0]` |
+| 3 finite time | closed form `gap_T(rho,{lam_i})` (note 34); estimator `K_t` (notes 15-17) | closed form: shape; estimator: magnitude | the actual gap |
+
+## Three contributions
+
+Refined to three pillars: 1 & 2 are foundational (static capacity geometry),
+3 is an independent track (initialization optimality), and the rest of the
+results build on 1 to form the main gap pillar, ordered as a progression
+(exists & soft -> exact at init -> finite-time law -> validation).
+
+### Contribution 1 — Distributional capacity of random feedback (static geometry)
+
+*(old 1 + 2.)*
+
+- Exact static-alignment law `Q_l = cos^2(W_{l+1}^T, B_l) ~ Beta(1/2,(D_l-1)/2)`,
+ `E[Q]=1/D`, with log-volume capacity cost `C_l(q) = -log P(Q_l >= q)`.
+- Multilayer scaling: independent blocks => `C_all = sum_l C_l = Theta(L n^2)`,
+ raw feasible volume `exp[-Theta(L n^2)]`.
+- Role: defines and measures the feedback-alignment burden before training.
+- **Proven.** `scripts/static_alignment_beta.py`, `scripts/capacity_scaling.py`.
+
+### Contribution 2 — Prior-free minimax optimality of random feedback
+
+*(old 3, independent track.)*
+
+- Without a prior on forward weights or task, no initialization beats isotropic
+ random feedback: `sup_mu inf_a E[(a^T b)^2] = 1/D`, achieved by isotropic.
+- Role: self-contained -- the burden cannot be designed away at initialization.
+- **Proven.** `scripts/minimax_initialization.py`.
+
+### Contribution 3 — The FA/BP operator gap: always nonzero, soft, exactly anchored at init, quantified over training
+
+*(old 4 + 5 + 6, as a four-level progression on Contribution 1.)*
+
+- **(a) Existence & softness.** `rho = output_share < 1` => `gap > 0` at every
+ `depth >= 2` and every capacity margin; analytic in the spectrum => smooth,
+ **no redundancy-exhaustion threshold** (null model `E[e]=k/P`, `P(e=0)=0`).
+- **(b) Exact at t=0.** `E_B[e_0 | W,r] = hidden_share = 1 - ||g_out^BP||^2 /
+ sum_l ||g_l^BP||^2`; one-hidden-layer Gaussian feedback => full conditional
+ Gaussian. Supplies the scalar burden `rho = 1 - E_B[e_0]`.
+ **Proven + validated** (0.4378 vs 0.4403). `scripts/actual_fa_initial_operator_moments.py`.
+- **(c) Finite-time law.** closed-form soft-ramp gap
+ `gap_T = (1/2N) sum_i c_i^2 [(1-eta rho lam_i/N)^{2T} - (1-eta lam_i/N)^{2T}]`
+ (mechanism + shape, `log gap ~ const - c rho lam_min`; note 34) **plus** the
+ early operator-velocity estimator `K_hat_t = K_0 + t(K_s-K_0)/s`
+ (drift-corrected magnitude, corr 0.999). closed form = why soft; estimator = how much.
+ `scripts/closed_form_soft_ramp.py`, `scripts/compressed_operator_predictor.py`.
+- **(d) Empirical validation (no fit).** matrix-alignment Beta; e0 Gaussian;
+ soft ramp (smooth, nonzero at positive margin -> refutes phase transition);
+ estimator trajectory overlap.
+
+### Old numbering -> new
+
+```text
+old 1, old 2 -> C1
+old 3 -> C2
+old 4 -> C3(b)
+old 6 (closed form + estimator) -> C3(c)
+old 5 -> C3(d)
+(new) existence & soft -> C3(a)
+```
+
+## Key numbers (the evidence)
+
+- init moment: `E_B[e_0]=0.4378 +/- 0.008` vs theory `hidden_share=0.4403`.
+- estimator: linear-velocity gap MAE `0.00189`, corr `0.99885` (256 traj);
+ stress grid corr `0.994-0.99989` across depth 1/2/3, width 32/64/96,
+ horizon 25/50/100 (notes 15-17). fixed `K(0)` is ~10x worse.
+- soft ramp (dense `T=30000`): gap `0.147 -> 0.00016` smooth across margin
+ `-190 -> +130`, **nonzero at positive margin, no kink** (note 22).
+- closed-form law (corrected, matched data/init -- note 37): no-fit
+ `corr(log pred, log measured)=0.933` (matched init) / `0.982` (5-init geo
+ mean); measured `log(gap)` linear in margin `R^2=0.993`;
+ `corr(lam_min_mean,-log gap)=0.987` (note 34+37; 0.977/0.963 retired).
+
+## Concise figure list (5 figures, one per contribution)
+
+1. static Beta histogram + QQ (sanity).
+2. soft ramp: gap vs margin, log-y, smooth, nonzero at positive margin
+ (refutes phase transition). `outputs/phase_transition_dense_T30000_352traj/`.
+3. init moment: `E_B[e_0]` vs `hidden_share`, no-fit scatter.
+ `outputs/actual_fa_initial_operator_moments/`.
+4. **estimator: predicted vs empirical finite-time gap, 256 traj + stress
+ grid, corr 0.999 (centerpiece).** `outputs/compressed_operator_s20_256traj_T50_width64_plots/`.
+5. **closed form vs measured ramp + ramp law.**
+ `outputs/closed_form_soft_ramp/closed_form_vs_measured_ramp.png`.
+
+## Proven vs estimated vs approximate
+
+- proven, no fit: contributions 1, 3; the `E[e]=k/P` null model; the init
+ moment `E_B[e_0]=hidden_share`.
+- closed-form gap law (note 34): proven under lazy (A1) + scalar-erosion (A2);
+ captures shape, not magnitude.
+- estimator: no fitted scalar, but **conditional on early operators `(K_0,K_s)`**
+ -- a predictive finite-time estimator, not an architecture-only theorem.
+
+## Open problem (the one real gap)
+
+Derive the drift / effective burden from architecture + initialization so the
+finite-time gap becomes architecture-only:
+
+- predict `rho` (hence `hidden_share`) from infinite-width NTK recursions;
+- predict the drift `K_s - K_0` (alignment gain) -- naive `dot_K_0` fails
+ (note 31), so use a bounded short-time alignment-gain path or a deep-linear /
+ two-layer order-parameter solution;
+- a per-mode spectral erosion `rho_i` refining the scalar A2.
+
+## Superseded / cleanup
+
+- `00_project_overview.md`, `03_paper_outline.md`: stale framing; retire
+ "phase transition", add the estimator + closed-form law.
+- notes 19-21 (transition-as-kink): refuted by the dense sweep (note 22).
+- hard-`k` as a *gap predictor*: it overpredicts (note 25); keep only as the
+ null model that motivates the operator object.
+- `scripts/actual_fa_initial_erosion_distribution.py` (e0 Gaussian, note 32):
+ the sampler reuses the analytic coefficient (partly circular) -- switch it to
+ the real FA backward (verified to reproduce the same numbers).
+- `scripts/alignment_recovery_dynamics.py` (rho(t) recovery probe): superseded
+ by the estimator + closed-form law; keep only as an optional `t=0` sanity or
+ remove.
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()