Replace transition theory with tangent capacity

4 files changed, 775 insertions, 104 deletions

diff --git a/.gitignore b/.gitignore
index 76f4c78..6dd4b78 100644
--- a/.gitignore
+++ b/.gitignore
@@ -24,6 +24,9 @@ checkpoints/
 wandb/
 tensorboard/
 
+# Local literature scratch files
+notes/lit_pdfs/
+
 # Data
 data/
 datasets/
diff --git a/notes/07_bp_capacity_theory.md b/notes/07_bp_capacity_theory.md
new file mode 100644
index 0000000..c7ac1f3
--- /dev/null
+++ b/notes/07_bp_capacity_theory.md
@@ -0,0 +1,372 @@
+# BP-Capacity Theory Correction
+
+The previous hard-count transition model is not acceptable as a main theory:
+
+\[
+M_{\mathrm{FA}}=P-K_{\mathrm{FA}}-Nn_L
+\]
+
+is only a crude diagnostic. It treats BP capacity as raw parameter count and
+treats FA alignment as full hard constraints. The main paper should instead use
+a capacity definition grounded in the training-set Jacobian / NTK.
+
+## Literature-Supported BP Capacity
+
+For squared loss on a finite training set, the mature local object is the output
+Jacobian
+
+\[
+J_{\mathrm{BP}}(\theta, X)
+=
+\frac{\partial f_\theta(X)}{\partial \theta}
+\in
+\mathbb R^{Nn_L\times P}.
+\]
+
+Karhadkar et al. prove the key local lemma: at a differentiable critical point,
+if the training-set Jacobian has full row rank, then the squared-loss critical
+point is a global minimizer. Their deep-ReLU results further show that many
+activation regions have full-rank Jacobian under explicit activation-pattern
+rank conditions.
+
+Therefore the hard BP interpolation capacity should be
+
+\[
+r_{\mathrm{BP}}(\theta)
+=
+\operatorname{rank} J_{\mathrm{BP}}(\theta, X),
+\]
+
+not \(P\) and not an empirical fitted multiple of \(Nn_L\).
+
+For finite-time or ill-conditioned training, use the empirical NTK
+
+\[
+K_{\mathrm{BP}}(\theta)
+=
+J_{\mathrm{BP}}(\theta,X)J_{\mathrm{BP}}(\theta,X)^\top
+\]
+
+and its effective dimension
+
+\[
+d_{\mathrm{BP}}(\sigma^2)
+=
+\operatorname{tr}
+\left[
+K_{\mathrm{BP}}
+\left(K_{\mathrm{BP}}+\sigma^2 I\right)^{-1}
+\right].
+\]
+
+This follows the NTK view: training progress on finite data is governed by the
+kernel Gram spectrum, with fast convergence in large-eigenvalue directions and
+slow convergence in small-eigenvalue directions.
+
+## FA Capacity Occupation
+
+There are two separate FA capacity notions.
+
+1. Matrix-alignment capacity: how much parameter-direction volume is occupied by
+   making forward matrices align with feedback matrices.
+2. Downstream training capacity: how much of BP's local training kernel survives
+   when BP gradients are replaced by FA surrogate gradients.
+
+The downstream experiment should use the second object as its main predictor.
+
+## FA Surrogate-Gradient Kernel
+
+Let \(m=Nn_L\). For each output coordinate basis vector \(e_i\in\mathbb R^m\),
+define the FA pseudo-gradient direction
+
+\[
+\tilde g_i^{\mathrm{FA}}(B,\theta)
+\in
+\mathbb R^P
+\]
+
+obtained by running the FA backward rule with output residual \(e_i\). Stack
+these rows into
+
+\[
+\tilde J_{\mathrm{FA}}(B,\theta)
+\in
+\mathbb R^{m\times P}.
+\]
+
+For squared loss with residual \(r=f_\theta(X)-y\), the BP update is
+
+\[
+\Delta\theta_{\mathrm{BP}}
+=
+-\eta J_{\mathrm{BP}}^\top r,
+\]
+
+and the FA update is
+
+\[
+\Delta\theta_{\mathrm{FA}}
+=
+-\eta \tilde J_{\mathrm{FA}}^\top r.
+\]
+
+The induced first-order output dynamics are therefore
+
+\[
+\Delta f_{\mathrm{BP}}
+=
+-\eta
+K_{\mathrm{BP}}r,
+\qquad
+K_{\mathrm{BP}}=J_{\mathrm{BP}}J_{\mathrm{BP}}^\top,
+\]
+
+and
+
+\[
+\Delta f_{\mathrm{FA}}
+=
+-\eta
+K_{\mathrm{FA}}r,
+\qquad
+K_{\mathrm{FA}}
+=
+J_{\mathrm{BP}}\tilde J_{\mathrm{FA}}^\top.
+\]
+
+This is the correct local object for downstream FA/BP gap. It is also the
+natural analogue of the surrogate-gradient NTK: the forward Jacobian controls
+how parameter motion changes outputs, while the surrogate Jacobian controls the
+parameter motion chosen by the learning rule.
+
+Since \(K_{\mathrm{FA}}\) can be non-symmetric, hard local solvability is governed
+by
+
+\[
+\operatorname{rank}(K_{\mathrm{FA}}),
+\]
+
+while stable finite-time learning should be measured by the spectrum of the
+symmetrized or positive-real part of the operator. A conservative symmetric
+capacity proxy is
+
+\[
+H_{\mathrm{FA}}
+=
+\frac{K_{\mathrm{FA}}+K_{\mathrm{FA}}^\top}{2},
+\]
+
+with positive eigenvalues \(\lambda_i^+\). Then
+
+\[
+d_{\mathrm{FA}}^{\mathrm{dyn}}(\sigma^2)
+=
+\sum_i
+\frac{\max(\lambda_i,0)}
+{\max(\lambda_i,0)+\sigma^2}.
+\]
+
+This is not a fitted scale. It is computed directly from the local FA learning
+operator.
+
+## Alignment-Constraint Projection
+
+Let \(C_{\mathrm{FA}}\) be the local Jacobian of the FA alignment constraints
+with respect to parameters. The hard constraint projector is
+
+\[
+\Pi_C
+=
+I
+-
+C_{\mathrm{FA}}^\top
+\left(C_{\mathrm{FA}}C_{\mathrm{FA}}^\top\right)^+
+C_{\mathrm{FA}}.
+\]
+
+The hard FA-accessible task rank is
+
+\[
+r_{\mathrm{FA}}
+=
+\operatorname{rank}
+\left(
+J_{\mathrm{BP}}\Pi_C J_{\mathrm{BP}}^\top
+\right).
+\]
+
+So the capacity loss is
+
+\[
+\Delta r
+=
+r_{\mathrm{BP}}-r_{\mathrm{FA}}.
+\]
+
+For a soft alignment burden, use
+
+\[
+K_{\mathrm{FA},\lambda}
+=
+J_{\mathrm{BP}}
+\left(M+\lambda C_{\mathrm{FA}}^\top C_{\mathrm{FA}}\right)^{-1}
+J_{\mathrm{BP}}^\top,
+\]
+
+and
+
+\[
+d_{\mathrm{FA},\lambda}(\sigma^2)
+=
+\operatorname{tr}
+\left[
+K_{\mathrm{FA},\lambda}
+\left(K_{\mathrm{FA},\lambda}+\sigma^2 I\right)^{-1}
+\right].
+\]
+
+The soft capacity loss is
+
+\[
+\Delta d_\lambda
+=
+d_{\mathrm{BP}}(\sigma^2)
+-
+d_{\mathrm{FA},\lambda}(\sigma^2).
+\]
+
+No fitted scalar is allowed here. Every quantity is computed from \(J\), \(C\),
+and a declared metric \(M\).
+
+## Correct Transition Statement
+
+For downstream memorization experiments, the FA/BP transition should first be
+predicted by the FA surrogate-gradient kernel:
+
+\[
+\operatorname{rank}(K_{\mathrm{FA}})\approx Nn_L
+\]
+
+or, for finite-time training,
+
+\[
+d_{\mathrm{FA}}^{\mathrm{dyn}}(\sigma^2)\approx Nn_L.
+\]
+
+For the matrix-alignment capacity theory, the transition can also be written as
+
+\[
+r_{\mathrm{FA}}\approx Nn_L
+\]
+
+in the hard projected-rank model, or by
+
+\[
+d_{\mathrm{FA},\lambda}(\sigma^2)\approx Nn_L
+\]
+
+in the finite-time/effective-rank model.
+
+Equivalently: the breakpoint occurs when FA-projected tangent capacity can no
+longer span the training output constraints.
+
+This is the clean version of "redundancy exhaustion":
+
+\[
+\operatorname{range}(J_{\mathrm{BP}}^\top)
+\]
+
+is the BP task-sensitive subspace. FA alignment only hurts downstream capacity
+to the extent that its constraint directions overlap this subspace.
+
+Under a generic random-subspace approximation with
+
+\[
+r_{\mathrm{BP}}=\operatorname{rank}J_{\mathrm{BP}},
+\qquad
+k=\operatorname{rank}C_{\mathrm{FA}},
+\]
+
+the expected soft overlap is
+
+\[
+\mathbb E\operatorname{tr}(P_C P_J)
+=
+\frac{kr_{\mathrm{BP}}}{P}.
+\]
+
+The hard generic rank after projection is
+
+\[
+r_{\mathrm{FA}}
+=
+\min(r_{\mathrm{BP}}, P-k).
+\]
+
+This gives a simple theorem without post-hoc parameters, but the experimentally
+plotted prediction should prefer the directly measured Jacobian/NTK quantities.
+
+## What To Change In Experiments
+
+The next downstream-capacity plot should not use only
+\(P-K_{\mathrm{FA}}-Nn_L\). It should compute, per architecture and seed:
+
+1. \(r_{\mathrm{BP}}\) or \(d_{\mathrm{BP}}(\sigma^2)\) on the full training set.
+2. \(\operatorname{rank}(K_{\mathrm{FA}})\) and
+   \(d_{\mathrm{FA}}^{\mathrm{dyn}}(\sigma^2)\) from the FA pseudo-gradient
+   kernel.
+3. \(r_{\mathrm{FA}}\) or \(d_{\mathrm{FA},\lambda}(\sigma^2)\) after applying
+   the FA alignment projector / penalty as the matrix-capacity ablation.
+4. A predicted distribution of train-gap transition from these computed spectra,
+   not from a fitted shift.
+
+The old hard margin may remain as an ablation or sanity-check axis, not as the
+main theorem.
+
+## Current Diagnostic
+
+The script
+
+`scripts/fa_tangent_kernel_capacity.py`
+
+computes \(J_{\mathrm{BP}}\), \(\tilde J_{\mathrm{FA}}\),
+\(K_{\mathrm{BP}}\), and \(K_{\mathrm{FA}}\) on the full training set for the
+existing downstream sweep.
+
+A smoke/probe run on the current random-label Adam experiment shows that the
+initial \(K_{\mathrm{FA}}\) often has a large negative symmetric spectrum. The
+continuous-time initial-kernel prediction therefore greatly overestimates FA
+failure, while the actual Adam FA trajectories recover through nonlinear
+alignment dynamics.
+
+This means:
+
+1. the FA tangent-kernel object is the right local capacity object;
+2. the current Adam random-label experiment is not a valid test of a fixed
+   initialization-kernel final-gap theorem;
+3. a quantitative final-gap theory must either use an experiment in the
+   SGD/lazy-kernel regime, or model the time-varying FA alignment trajectory.
+
+Do not present the initial-kernel prediction as a main downstream theorem for
+the existing Adam trajectories.
+
+## References To Cite
+
+- Karhadkar et al., *Mildly Overparameterized ReLU Networks Have a Favorable
+  Loss Landscape*: full-rank training-set Jacobian implies squared-loss critical
+  points are global minimizers; deep-ReLU activation regions can have full-rank
+  Jacobians under explicit rank conditions.
+- Jacot et al., *Neural Tangent Kernel*: finite-data training dynamics are
+  governed by the NTK Gram spectrum; convergence follows kernel principal
+  components.
+- Eilers et al., *A Generalized Neural Tangent Kernel for Surrogate Gradient
+  Learning*: surrogate-gradient methods can be analyzed by a generalized NTK;
+  FA should be treated as a surrogate-gradient learning rule for downstream
+  dynamics.
+- Yun et al., *Small ReLU Networks Are Powerful Memorizers*: ReLU memorization
+  capacity scales with hidden-layer parameter products, supporting the need to
+  treat memorization capacity separately from raw parameter count.
+- Bartlett et al., *Nearly-Tight VC-Dimension and Pseudodimension Bounds for
+  Piecewise Linear Neural Networks*: worst-case expressivity scales differently
+  from local trainability; useful as a background upper-bound reference, not the
+  exact transition predictor.
diff --git a/notes/07_minimal_transition_theory.md b/notes/07_minimal_transition_theory.md
deleted file mode 100644
index ced7e61..0000000
--- a/notes/07_minimal_transition_theory.md
+++ /dev/null
@@ -1,104 +0,0 @@
-# Minimal Transition Theory
-
-The cleanest correction is not to add a separate FA burden coefficient first.
-Start by calibrating the task dimension from BP alone.
-
-## Problem
-
-The hard theory used:
-
-\[
-d_{\mathrm{task}}=Nn_L=512.
-\]
-
-With this value, the smooth theory peak is near:
-
-\[
-M_{\mathrm{FA}}\approx -273,
-\]
-
-while 256 empirical FA trajectories peak at:
-
-\[
-M_{\mathrm{FA}}=-318.
-\]
-
-## BP-Only Calibration
-
-Fit the BP train curve with:
-
-\[
-L_{\mathrm{BP}}(P)
-\approx
-s\cdot \operatorname{softplus}_{\tau}(d_{\mathrm{task}}^{\mathrm{eff}}-P).
-\]
-
-The BP-only fit gives:
-
-\[
-d_{\mathrm{task}}^{\mathrm{eff}}\approx 437,
-\qquad
-\tau\approx 50.
-\]
-
-So BP is stronger than the naive parameter count suggests:
-
-\[
-d_{\mathrm{task}}^{\mathrm{eff}}
-\approx
-0.85\,Nn_L.
-\]
-
-## Minimal FA Prediction
-
-Use the same hard FA burden:
-
-\[
-K_{\mathrm{FA}}=\sum_l(D_l-1),
-\]
-
-but replace the task dimension:
-
-\[
-M_{\mathrm{FA}}^{\mathrm{min}}
-=
-P-K_{\mathrm{FA}}-d_{\mathrm{task}}^{\mathrm{eff}}.
-\]
-
-Equivalently, when plotted on the original hard-margin axis
-
-\[
-P-K_{\mathrm{FA}}-Nn_L,
-\]
-
-the theory shifts left by:
-
-\[
-Nn_L-d_{\mathrm{task}}^{\mathrm{eff}}
-\approx
-75.
-\]
-
-This explains the observed left shift without introducing a fitted FA burden
-coefficient.
-
-Generated figure:
-
-`outputs/downstream_capacity_random_main_fast/capacity_transition_bp_calibrated_p2_256.png`
-
-## Recommendation
-
-Use this as the main downstream theory:
-
-1. Matrix/capacity theory remains exact and uses the hard/log-volume definitions.
-2. Downstream performance uses a BP-calibrated effective task dimension.
-3. FA capacity loss is then tested with no extra FA-specific fitting parameter.
-
-Only add an effective FA burden coefficient
-
-\[
-K_{\mathrm{FA}}^{\mathrm{eff}}=\alpha K_{\mathrm{FA}}
-\]
-
-as an appendix ablation if needed. It improves fit but makes the theory look
-more ad hoc.
diff --git a/scripts/fa_tangent_kernel_capacity.py b/scripts/fa_tangent_kernel_capacity.py
new file mode 100644
index 0000000..ca779cf
--- /dev/null
+++ b/scripts/fa_tangent_kernel_capacity.py
@@ -0,0 +1,400 @@
+#!/usr/bin/env python3
+"""Measure BP/FA tangent-kernel capacity for existing downstream sweeps."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import math
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from scipy.sparse.linalg import expm_multiply
+
+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
+
+
+Tensor = torch.Tensor
+
+
+@dataclass(frozen=True)
+class SourceRun:
+    width: int
+    parameter_count: int
+    init_seed: int
+    feedback_seed: int
+    run_type: str
+    train_mse: float
+    bp_train_mse: float
+    train_gap_to_bp: float
+    fa_capacity_margin: int
+
+
+@dataclass(frozen=True)
+class KernelRow:
+    width: int
+    init_seed: int
+    feedback_seed: int
+    fa_capacity_margin: int
+    empirical_train_gap: float
+    predicted_bp_loss: float
+    predicted_fa_loss: float
+    predicted_train_gap: float
+    bp_rank: int
+    fa_rank: int
+    bp_d_eff: float
+    fa_d_eff: float
+    bp_lambda: float
+    fa_lambda: float
+    fa_sym_min: float
+    fa_sym_neg_mass: float
+    fa_trace: float
+    bp_trace: float
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Compute full-training-set BP/FA tangent-kernel capacity."
+    )
+    parser.add_argument(
+        "--source-outdir",
+        type=Path,
+        default=Path("outputs/downstream_capacity_random_main_fast"),
+    )
+    parser.add_argument(
+        "--outdir",
+        type=Path,
+        default=Path("outputs/fa_tangent_kernel_capacity"),
+    )
+    parser.add_argument("--max-fa-per-width", type=int, default=0)
+    parser.add_argument("--max-widths", type=int, default=0)
+    parser.add_argument("--sigma-rel", type=float, default=1e-3)
+    parser.add_argument("--rank-tol-rel", type=float, default=1e-8)
+    parser.add_argument(
+        "--prediction",
+        choices=["continuous"],
+        default="continuous",
+        help="continuous uses exp(-(lr*steps/N)K) on the initial residual.",
+    )
+    parser.add_argument("--plot", action="store_true")
+    return parser.parse_args()
+
+
+def load_config(source_outdir: Path) -> dcs.RunConfig:
+    payload = json.loads((source_outdir / "summary.json").read_text())
+    raw = payload["config"]
+    raw["outdir"] = str(source_outdir)
+    raw.setdefault("init_seed_offset", 0)
+    raw.setdefault("feedback_seed_offset", 0)
+    raw.setdefault("skip_jacobian", False)
+    return dcs.RunConfig(**raw)
+
+
+def load_runs(source_outdir: Path) -> list[SourceRun]:
+    rows: list[SourceRun] = []
+    with (source_outdir / "runs.csv").open(newline="") as handle:
+        reader = csv.DictReader(handle)
+        for raw in reader:
+            rows.append(
+                SourceRun(
+                    width=int(raw["width"]),
+                    parameter_count=int(raw["parameter_count"]),
+                    init_seed=int(raw["init_seed"]),
+                    feedback_seed=int(raw["feedback_seed"]),
+                    run_type=raw["run_type"],
+                    train_mse=float(raw["train_mse"]),
+                    bp_train_mse=float(raw["bp_train_mse"]),
+                    train_gap_to_bp=float(raw["train_gap_to_bp"]),
+                    fa_capacity_margin=int(raw["fa_capacity_margin"]),
+                )
+            )
+    return rows
+
+
+def flatten_grads(grads: list[Tensor]) -> Tensor:
+    return torch.cat([grad.reshape(-1) for grad in grads])
+
+
+def pseudo_jacobian(
+    weights: list[Tensor],
+    x: Tensor,
+    feedback: list[Tensor] | None,
+) -> Tensor:
+    activations, preacts = dcs.forward(weights, x)
+    outputs = activations[-1]
+    sample_count, output_dim = outputs.shape
+    rows: list[Tensor] = []
+
+    for index in range(sample_count * output_dim):
+        delta_out = torch.zeros_like(outputs)
+        delta_out.reshape(-1)[index] = 1.0
+        deltas: list[Tensor] = [
+            torch.empty(0, dtype=torch.float64, device=x.device) for _ in weights
+        ]
+        deltas[-1] = delta_out
+        for layer in range(len(weights) - 2, -1, -1):
+            if feedback is None:
+                back = deltas[layer + 1] @ weights[layer + 1]
+            else:
+                back = deltas[layer + 1] @ feedback[layer].T
+            deltas[layer] = back * (preacts[layer] > 0)
+        grads = [delta.T @ activations[layer] for layer, delta in enumerate(deltas)]
+        rows.append(flatten_grads(grads).detach())
+
+    return torch.stack(rows, dim=0)
+
+
+def symmetric_capacity(
+    kernel: np.ndarray,
+    sigma_rel: float,
+    rank_tol_rel: float,
+) -> tuple[int, float, float, float, float]:
+    sym = 0.5 * (kernel + kernel.T)
+    eigenvalues = np.linalg.eigvalsh(sym)
+    positive = np.clip(eigenvalues, 0.0, None)
+    max_eval = float(np.max(positive)) if positive.size else 0.0
+    lam = max(sigma_rel * max_eval, 1e-12)
+    d_eff = float(np.sum(positive / (positive + lam)))
+    rank = int(np.count_nonzero(positive > rank_tol_rel * max_eval)) if max_eval > 0 else 0
+    neg_mass = float(np.sum(np.clip(-eigenvalues, 0.0, None)))
+    min_eval = float(np.min(eigenvalues)) if eigenvalues.size else 0.0
+    trace = float(np.trace(kernel))
+    return rank, d_eff, lam, min_eval, neg_mass, trace
+
+
+def predict_loss(
+    kernel: np.ndarray,
+    residual: np.ndarray,
+    lr: float,
+    steps: int,
+    samples: int,
+) -> float:
+    scale = -(lr * steps / samples)
+    evolved = expm_multiply(scale * kernel, residual)
+    return 0.5 * float(np.dot(evolved, evolved)) / samples
+
+
+def select_fa_rows(
+    rows: list[SourceRun],
+    max_fa_per_width: int,
+    max_widths: int,
+) -> list[SourceRun]:
+    fa_rows = [row for row in rows if row.run_type == "fa"]
+    widths = sorted({row.width for row in fa_rows})
+    if max_widths > 0:
+        widths = widths[:max_widths]
+    selected: list[SourceRun] = []
+    for width in widths:
+        width_rows = [row for row in fa_rows if row.width == width]
+        width_rows.sort(key=lambda row: (row.init_seed, row.feedback_seed))
+        if max_fa_per_width > 0:
+            width_rows = width_rows[:max_fa_per_width]
+        selected.extend(width_rows)
+    return selected
+
+
+def write_kernel_rows(path: Path, rows: list[KernelRow]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(KernelRow.__annotations__.keys()))
+        writer.writeheader()
+        for row in rows:
+            writer.writerow(row.__dict__)
+
+
+def plot_results(rows: list[KernelRow], outdir: Path) -> list[Path]:
+    paths: list[Path] = []
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    scatter_path = outdir / "predicted_vs_empirical_train_gap.png"
+    plt.figure(figsize=(6.2, 5.2))
+    widths = sorted({row.width for row in rows})
+    for width in widths:
+        subset = [row for row in rows if row.width == width]
+        plt.scatter(
+            [row.predicted_train_gap for row in subset],
+            [row.empirical_train_gap for row in subset],
+            s=22,
+            alpha=0.55,
+            label=f"n={width}",
+        )
+    all_values = [
+        value
+        for row in rows
+        for value in (row.predicted_train_gap, row.empirical_train_gap)
+        if math.isfinite(value)
+    ]
+    if all_values:
+        lo = min(all_values)
+        hi = max(all_values)
+        plt.plot([lo, hi], [lo, hi], color="black", linewidth=1)
+    plt.xlabel("linearized FA/BP train-gap prediction")
+    plt.ylabel("empirical FA/BP train gap")
+    plt.title("Tangent-kernel prediction vs trajectory gap")
+    plt.legend(fontsize=8, ncols=2)
+    plt.tight_layout()
+    plt.savefig(scatter_path, dpi=180)
+    plt.close()
+    paths.append(scatter_path)
+
+    transition_path = outdir / "kernel_capacity_transition.png"
+    plt.figure(figsize=(7.5, 4.8))
+    for row in rows:
+        plt.plot(
+            [row.fa_capacity_margin, row.fa_capacity_margin],
+            [row.predicted_train_gap, row.empirical_train_gap],
+            color="0.7",
+            alpha=0.18,
+            linewidth=0.8,
+        )
+    plt.scatter(
+        [row.fa_capacity_margin for row in rows],
+        [row.predicted_train_gap for row in rows],
+        s=20,
+        alpha=0.55,
+        color="tab:blue",
+        label="kernel prediction",
+    )
+    plt.scatter(
+        [row.fa_capacity_margin for row in rows],
+        [row.empirical_train_gap for row in rows],
+        s=20,
+        alpha=0.35,
+        color="tab:orange",
+        label="trajectory",
+    )
+    plt.axvline(0.0, color="black", linewidth=1, linestyle="--")
+    plt.axhline(0.0, color="black", linewidth=1)
+    plt.xlabel("old hard FA margin, for reference")
+    plt.ylabel("FA train MSE - BP train MSE")
+    plt.title("Predicted and empirical transition")
+    plt.legend()
+    plt.tight_layout()
+    plt.savefig(transition_path, dpi=180)
+    plt.close()
+    paths.append(transition_path)
+
+    capacity_path = outdir / "fa_dynamic_effective_dimension_vs_margin.png"
+    plt.figure(figsize=(7.5, 4.8))
+    plt.scatter(
+        [row.fa_capacity_margin for row in rows],
+        [row.fa_d_eff for row in rows],
+        s=20,
+        alpha=0.55,
+        color="tab:green",
+    )
+    if rows:
+        task_dim = max(row.bp_rank for row in rows)
+        plt.axhline(task_dim, color="black", linewidth=1, linestyle="--", label="observed max BP rank")
+    plt.axvline(0.0, color="black", linewidth=1, linestyle="--")
+    plt.xlabel("old hard FA margin, for reference")
+    plt.ylabel("FA dynamic effective dimension")
+    plt.title("Measured FA tangent-kernel capacity")
+    plt.legend()
+    plt.tight_layout()
+    plt.savefig(capacity_path, dpi=180)
+    plt.close()
+    paths.append(capacity_path)
+
+    return paths
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.source_outdir)
+    if config.optimizer != "sgd":
+        print(
+            "warning: source trajectories used optimizer="
+            f"{config.optimizer!r}; continuous tangent-kernel prediction is an SGD "
+            "linearization diagnostic, not a final-gap theorem for this run.",
+            flush=True,
+        )
+    if config.torch_threads > 0:
+        torch.set_num_threads(config.torch_threads)
+
+    source_rows = load_runs(args.source_outdir)
+    fa_rows = select_fa_rows(source_rows, args.max_fa_per_width, args.max_widths)
+    x_train, y_train, _x_test, _y_test, _x_probe, _teacher = dcs.make_data(config)
+
+    result_rows: list[KernelRow] = []
+    bp_cache: dict[tuple[int, int], tuple[np.ndarray, np.ndarray, float, int, float, float, float]] = {}
+
+    for index, row in enumerate(fa_rows, start=1):
+        print(
+            f"[{index}/{len(fa_rows)}] width={row.width} init={row.init_seed} "
+            f"feedback={row.feedback_seed}",
+            flush=True,
+        )
+        initial_weights = dcs.initialize_weights(config, row.width, row.init_seed)
+        feedback = dcs.init_feedback(config, row.width, row.feedback_seed)
+        cache_key = (row.width, row.init_seed)
+        if cache_key not in bp_cache:
+            with torch.no_grad():
+                initial_residual = (dcs.predict(initial_weights, x_train) - y_train).reshape(-1)
+            j_bp = pseudo_jacobian(initial_weights, x_train, feedback=None).cpu().numpy()
+            k_bp = j_bp @ j_bp.T
+            residual = initial_residual.detach().cpu().numpy()
+            bp_loss = predict_loss(k_bp, residual, config.lr, config.steps, config.train_samples)
+            bp_rank, bp_d_eff, bp_lam, _bp_min, _bp_neg, bp_trace = symmetric_capacity(
+                k_bp, args.sigma_rel, args.rank_tol_rel
+            )
+            bp_cache[cache_key] = (
+                j_bp,
+                residual,
+                bp_loss,
+                bp_rank,
+                bp_d_eff,
+                bp_lam,
+                bp_trace,
+            )
+
+        j_bp, residual, bp_loss, bp_rank, bp_d_eff, bp_lam, bp_trace = bp_cache[cache_key]
+        j_fa = pseudo_jacobian(initial_weights, x_train, feedback=feedback).cpu().numpy()
+        k_fa = j_bp @ j_fa.T
+        fa_loss = predict_loss(k_fa, residual, config.lr, config.steps, config.train_samples)
+        fa_rank, fa_d_eff, fa_lam, fa_min, fa_neg, fa_trace = symmetric_capacity(
+            k_fa, args.sigma_rel, args.rank_tol_rel
+        )
+        result_rows.append(
+            KernelRow(
+                width=row.width,
+                init_seed=row.init_seed,
+                feedback_seed=row.feedback_seed,
+                fa_capacity_margin=row.fa_capacity_margin,
+                empirical_train_gap=row.train_gap_to_bp,
+                predicted_bp_loss=bp_loss,
+                predicted_fa_loss=fa_loss,
+                predicted_train_gap=fa_loss - bp_loss,
+                bp_rank=bp_rank,
+                fa_rank=fa_rank,
+                bp_d_eff=bp_d_eff,
+                fa_d_eff=fa_d_eff,
+                bp_lambda=bp_lam,
+                fa_lambda=fa_lam,
+                fa_sym_min=fa_min,
+                fa_sym_neg_mass=fa_neg,
+                fa_trace=fa_trace,
+                bp_trace=bp_trace,
+            )
+        )
+
+    args.outdir.mkdir(parents=True, exist_ok=True)
+    csv_path = args.outdir / "kernel_capacity_rows.csv"
+    write_kernel_rows(csv_path, result_rows)
+    print(f"rows: {csv_path}")
+    if args.plot:
+        for path in plot_results(result_rows, args.outdir):
+            print(f"plot: {path}")
+
+
+if __name__ == "__main__":
+    main()