Compare soft erosion theory distributions

1 files changed, 473 insertions, 0 deletions

diff --git a/scripts/soft_erosion_theory_vs_empirical.py b/scripts/soft_erosion_theory_vs_empirical.py
new file mode 100644
index 0000000..a390241
--- /dev/null
+++ b/scripts/soft_erosion_theory_vs_empirical.py
@@ -0,0 +1,473 @@
+#!/usr/bin/env python3
+"""Compare soft-erosion theory distributions with empirical FA trajectories."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+from scipy import stats
+
+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
+from fa_tangent_kernel_capacity import flatten_grads, pseudo_jacobian  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--run-dir",
+        type=Path,
+        default=Path("outputs/phase_transition_dense_T30000_352traj"),
+    )
+    parser.add_argument(
+        "--outdir",
+        type=Path,
+        default=Path("outputs/soft_erosion_theory_vs_empirical"),
+    )
+    parser.add_argument("--theory-samples", type=int, default=20_000)
+    parser.add_argument("--local-widths", type=int, nargs="+", default=[])
+    parser.add_argument("--seed", type=int, default=0)
+    return parser.parse_args()
+
+
+def load_config(run_dir: Path) -> dcs.RunConfig:
+    payload = json.loads((run_dir / "summary.json").read_text())
+    raw = payload["config"]
+    raw["outdir"] = str(run_dir)
+    return dcs.RunConfig(**raw)
+
+
+def loss_from_spectrum(
+    eigenvalues: np.ndarray,
+    coefficients: np.ndarray,
+    rho: np.ndarray,
+    lr: float,
+    steps: int,
+    samples: int,
+) -> np.ndarray:
+    tau = lr * steps / samples
+    positive = np.clip(eigenvalues, 0.0, None)
+    # Shape: theory_samples x modes.
+    decay = np.exp(-2.0 * tau * rho[:, None] * positive[None, :])
+    return 0.5 * (decay @ (coefficients * coefficients)) / samples
+
+
+def bp_spectrum_for_width(
+    config: dcs.RunConfig,
+    x_train: torch.Tensor,
+    y_train: torch.Tensor,
+    width: int,
+    init_seed: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    weights = dcs.initialize_weights(config, width, init_seed)
+    residual = (dcs.predict(weights, x_train) - y_train).reshape(-1).cpu().numpy()
+    jacobian = pseudo_jacobian(weights, x_train, feedback=None).cpu().numpy()
+    kernel = jacobian @ jacobian.T
+    eigenvalues, eigenvectors = np.linalg.eigh(0.5 * (kernel + kernel.T))
+    coefficients = eigenvectors.T @ residual
+    return eigenvalues, coefficients
+
+
+def theory_gap_distribution(
+    config: dcs.RunConfig,
+    summary: pd.Series,
+    eigenvalues: np.ndarray,
+    coefficients: np.ndarray,
+    rng: np.random.Generator,
+    samples: int,
+) -> np.ndarray:
+    p_count = int(summary["parameter_count"])
+    k_count = int(summary["fa_constraint_rank"])
+    if not 0 < k_count < p_count:
+        raise ValueError(f"Invalid random-subspace dimensions: P={p_count}, k={k_count}")
+    rho = rng.beta((p_count - k_count) / 2.0, k_count / 2.0, size=samples)
+    fa_loss = loss_from_spectrum(
+        eigenvalues,
+        coefficients,
+        rho,
+        config.lr,
+        config.steps,
+        config.train_samples,
+    )
+    bp_loss = loss_from_spectrum(
+        eigenvalues,
+        coefficients,
+        np.ones(1, dtype=np.float64),
+        config.lr,
+        config.steps,
+        config.train_samples,
+    )[0]
+    return fa_loss - bp_loss
+
+
+def summarize(values: np.ndarray) -> dict[str, float]:
+    return {
+        "mean": float(np.mean(values)),
+        "q05": float(np.quantile(values, 0.05)),
+        "q25": float(np.quantile(values, 0.25)),
+        "q50": float(np.quantile(values, 0.50)),
+        "q75": float(np.quantile(values, 0.75)),
+        "q95": float(np.quantile(values, 0.95)),
+    }
+
+
+def flatten_gradient_list(grads: list[torch.Tensor]) -> np.ndarray:
+    return flatten_grads(grads).detach().cpu().numpy()
+
+
+def local_erosion_rows(
+    config: dcs.RunConfig,
+    x_train: torch.Tensor,
+    y_train: torch.Tensor,
+    runs: pd.DataFrame,
+    widths: list[int],
+) -> pd.DataFrame:
+    records: list[dict[str, float | int]] = []
+    for width in widths:
+        width_runs = runs[(runs["width"] == width) & (runs["run_type"] == "fa")]
+        if width_runs.empty:
+            continue
+        for init_seed, init_group in width_runs.groupby("init_seed"):
+            weights = dcs.initialize_weights(config, width, int(init_seed))
+            bp_grad = flatten_gradient_list(dcs.gradients(weights, x_train, y_train, None))
+            denom = float(np.dot(bp_grad, bp_grad))
+            for _, row in init_group.iterrows():
+                feedback = dcs.init_feedback(config, width, int(row["feedback_seed"]))
+                fa_grad = flatten_gradient_list(
+                    dcs.gradients(weights, x_train, y_train, feedback)
+                )
+                ratio = float(np.dot(bp_grad, fa_grad) / denom)
+                records.append(
+                    {
+                        "width": int(width),
+                        "init_seed": int(init_seed),
+                        "feedback_seed": int(row["feedback_seed"]),
+                        "fa_capacity_margin": int(row["fa_capacity_margin"]),
+                        "speed_ratio": ratio,
+                        "erosion": 1.0 - ratio,
+                    }
+                )
+    return pd.DataFrame(records)
+
+
+def plot_gap_distributions(
+    empirical: pd.DataFrame,
+    theory: pd.DataFrame,
+    outdir: Path,
+) -> Path:
+    fig, ax = plt.subplots(figsize=(9.2, 5.6), dpi=180)
+    theory = theory.sort_values("fa_capacity_margin", ascending=False)
+    empirical = empirical.sort_values("fa_capacity_margin", ascending=False)
+    x_theory = theory["fa_capacity_margin"].to_numpy()
+    x_emp = empirical["fa_capacity_margin"].to_numpy()
+
+    ax.fill_between(
+        x_theory,
+        np.clip(theory["q05"], 1e-9, None),
+        np.clip(theory["q95"], 1e-9, None),
+        color="#174f91",
+        alpha=0.16,
+        linewidth=0,
+        label="theory 5-95%",
+    )
+    ax.fill_between(
+        x_theory,
+        np.clip(theory["q25"], 1e-9, None),
+        np.clip(theory["q75"], 1e-9, None),
+        color="#174f91",
+        alpha=0.28,
+        linewidth=0,
+        label="theory 25-75%",
+    )
+    ax.plot(
+        x_theory,
+        np.clip(theory["mean"], 1e-9, None),
+        color="#174f91",
+        marker="o",
+        linewidth=2.2,
+        label="theory mean",
+    )
+    ax.errorbar(
+        x_emp,
+        np.clip(empirical["mean"], 1e-9, None),
+        yerr=[
+            np.clip(empirical["mean"] - empirical["q05"], 0.0, None),
+            np.clip(empirical["q95"] - empirical["mean"], 0.0, None),
+        ],
+        color="#c65f16",
+        marker="o",
+        linewidth=2.0,
+        capsize=3,
+        label="empirical 5-95%",
+    )
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
+    ax.set_yscale("log")
+    ax.set_xlim(max(x_emp.max(), x_theory.max()) + 20, min(x_emp.min(), x_theory.min()) - 20)
+    ax.set_xlabel("hard FA capacity margin P - K_FA - N*out")
+    ax.set_ylabel("FA train MSE - BP train MSE")
+    ax.set_title("Soft-erosion random-subspace theory vs empirical FA gap")
+    ax.grid(alpha=0.18, which="both")
+    ax.legend(frameon=True, ncol=2)
+    fig.tight_layout()
+    path = outdir / "gap_distribution_theory_vs_empirical.png"
+    fig.savefig(path)
+    plt.close(fig)
+    return path
+
+
+def plot_local_erosion(
+    local: pd.DataFrame,
+    summary: pd.DataFrame,
+    rng: np.random.Generator,
+    theory_samples: int,
+    outdir: Path,
+) -> Path:
+    widths = sorted(local["width"].unique())
+    fig, axes = plt.subplots(1, len(widths), figsize=(5.2 * len(widths), 4.6), dpi=180)
+    if len(widths) == 1:
+        axes = [axes]
+    for ax, width in zip(axes, widths, strict=True):
+        values = local.loc[local["width"] == width, "erosion"].to_numpy(dtype=np.float64)
+        row = summary.loc[summary["width"] == width].iloc[0]
+        p_count = int(row["parameter_count"])
+        k_count = int(row["fa_constraint_rank"])
+        theory = rng.beta(k_count / 2.0, (p_count - k_count) / 2.0, size=theory_samples)
+        bins = np.linspace(
+            min(values.min(), theory.min(), 0.0),
+            max(values.max(), theory.max(), 1.0),
+            45,
+        )
+        ax.hist(
+            values,
+            bins=bins,
+            density=True,
+            alpha=0.46,
+            color="#c65f16",
+            label="empirical FA",
+        )
+        ax.hist(
+            theory,
+            bins=bins,
+            density=True,
+            histtype="step",
+            linewidth=2.0,
+            color="#174f91",
+            label="Beta subspace",
+        )
+        ax.axvline(float(np.mean(values)), color="#c65f16", linestyle="--", linewidth=1.2)
+        ax.axvline(float(np.mean(theory)), color="#174f91", linestyle="--", linewidth=1.2)
+        ax.set_title(f"width={width}, margin={int(row['fa_capacity_margin'])}")
+        ax.set_xlabel("one-step erosion 1 - ΔL_FA/ΔL_BP")
+        ax.grid(alpha=0.16)
+        ax.legend(frameon=True, fontsize=8)
+    axes[0].set_ylabel("density")
+    fig.suptitle("Local operator erosion: real FA vs random-subspace null")
+    fig.tight_layout()
+    path = outdir / "local_erosion_distribution_theory_vs_empirical.png"
+    fig.savefig(path)
+    plt.close(fig)
+    return path
+
+
+def plot_local_scalar_gap(
+    empirical: pd.DataFrame,
+    local_theory: pd.DataFrame,
+    outdir: Path,
+) -> Path:
+    fig, ax = plt.subplots(figsize=(9.2, 5.6), dpi=180)
+    theory = local_theory.sort_values("fa_capacity_margin", ascending=False)
+    empirical = empirical.sort_values("fa_capacity_margin", ascending=False)
+    x_theory = theory["fa_capacity_margin"].to_numpy()
+    x_emp = empirical["fa_capacity_margin"].to_numpy()
+    ax.fill_between(
+        x_theory,
+        np.clip(theory["q05"], 1e-9, None),
+        np.clip(theory["q95"], 1e-9, None),
+        color="#2b8a3e",
+        alpha=0.16,
+        linewidth=0,
+        label="operator theory 5-95%",
+    )
+    ax.fill_between(
+        x_theory,
+        np.clip(theory["q25"], 1e-9, None),
+        np.clip(theory["q75"], 1e-9, None),
+        color="#2b8a3e",
+        alpha=0.28,
+        linewidth=0,
+        label="operator theory 25-75%",
+    )
+    ax.plot(
+        x_theory,
+        np.clip(theory["mean"], 1e-9, None),
+        color="#2b8a3e",
+        marker="o",
+        linewidth=2.2,
+        label="operator theory mean",
+    )
+    ax.errorbar(
+        x_emp,
+        np.clip(empirical["mean"], 1e-9, None),
+        yerr=[
+            np.clip(empirical["mean"] - empirical["q05"], 0.0, None),
+            np.clip(empirical["q95"] - empirical["mean"], 0.0, None),
+        ],
+        color="#c65f16",
+        marker="o",
+        linewidth=2.0,
+        capsize=3,
+        label="empirical 5-95%",
+    )
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
+    ax.set_yscale("log")
+    ax.set_xlim(max(x_emp.max(), x_theory.max()) + 20, min(x_emp.min(), x_theory.min()) - 20)
+    ax.set_xlabel("hard FA capacity margin P - K_FA - N*out")
+    ax.set_ylabel("FA train MSE - BP train MSE")
+    ax.set_title("Local-operator scalar theory vs empirical FA gap")
+    ax.grid(alpha=0.18, which="both")
+    ax.legend(frameon=True, ncol=2)
+    fig.tight_layout()
+    path = outdir / "local_operator_scalar_gap_theory_vs_empirical.png"
+    fig.savefig(path)
+    plt.close(fig)
+    return path
+
+
+def main() -> None:
+    args = parse_args()
+    if args.theory_samples < 1:
+        raise ValueError("--theory-samples must be positive")
+    rng = np.random.default_rng(args.seed)
+    outdir = args.outdir
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    config = load_config(args.run_dir)
+    x_train, y_train, _x_test, _y_test, _x_probe, _teacher = dcs.make_data(config)
+    runs = pd.read_csv(args.run_dir / "runs.csv")
+    summary = pd.read_csv(args.run_dir / "width_summary.csv")
+
+    theory_records = []
+    empirical_records = []
+    spectra: dict[int, tuple[np.ndarray, np.ndarray]] = {}
+    for _, width_row in summary.sort_values("width").iterrows():
+        width = int(width_row["width"])
+        init_seed = int(
+            runs.loc[(runs["width"] == width) & (runs["run_type"] == "bp"), "init_seed"].iloc[0]
+        )
+        print(f"width={width}: computing BP spectrum", flush=True)
+        eigenvalues, coefficients = bp_spectrum_for_width(
+            config, x_train, y_train, width, init_seed
+        )
+        spectra[width] = (eigenvalues, coefficients)
+        theory_gaps = theory_gap_distribution(
+            config,
+            width_row,
+            eigenvalues,
+            coefficients,
+            rng,
+            args.theory_samples,
+        )
+        t_summary = summarize(np.clip(theory_gaps, 1e-12, None))
+        t_summary.update(
+            {
+                "width": width,
+                "fa_capacity_margin": int(width_row["fa_capacity_margin"]),
+                "parameter_count": int(width_row["parameter_count"]),
+                "fa_constraint_rank": int(width_row["fa_constraint_rank"]),
+            }
+        )
+        theory_records.append(t_summary)
+
+        emp_values = runs.loc[
+            (runs["width"] == width) & (runs["run_type"] == "fa"),
+            "train_gap_to_bp",
+        ].to_numpy(dtype=np.float64)
+        e_summary = summarize(np.clip(emp_values, 1e-12, None))
+        e_summary.update(
+            {
+                "width": width,
+                "fa_capacity_margin": int(width_row["fa_capacity_margin"]),
+                "parameter_count": int(width_row["parameter_count"]),
+                "fa_constraint_rank": int(width_row["fa_constraint_rank"]),
+            }
+        )
+        empirical_records.append(e_summary)
+
+    theory_df = pd.DataFrame(theory_records)
+    empirical_df = pd.DataFrame(empirical_records)
+    theory_csv = outdir / "theory_gap_summary.csv"
+    empirical_csv = outdir / "empirical_gap_summary.csv"
+    theory_df.to_csv(theory_csv, index=False)
+    empirical_df.to_csv(empirical_csv, index=False)
+
+    local_widths = args.local_widths or sorted(summary["width"].astype(int).tolist())
+    local = local_erosion_rows(config, x_train, y_train, runs, local_widths)
+    local_csv = outdir / "local_erosion_rows.csv"
+    local.to_csv(local_csv, index=False)
+
+    local_gap_records = []
+    for width, group in local.groupby("width"):
+        eigenvalues, coefficients = spectra[int(width)]
+        bp_loss = loss_from_spectrum(
+            eigenvalues,
+            coefficients,
+            np.ones(1, dtype=np.float64),
+            config.lr,
+            config.steps,
+            config.train_samples,
+        )[0]
+        rho = group["speed_ratio"].to_numpy(dtype=np.float64)
+        # Negative local speed ratios are valid predictions of local instability,
+        # but the scalar frozen-kernel rollout can explode. Keep the raw value in
+        # local_erosion_rows.csv and clip only for this conservative distribution
+        # plot.
+        rho_for_rollout = np.clip(rho, 0.0, 2.0)
+        predicted_loss = loss_from_spectrum(
+            eigenvalues,
+            coefficients,
+            rho_for_rollout,
+            config.lr,
+            config.steps,
+            config.train_samples,
+        )
+        predicted_gap = predicted_loss - bp_loss
+        row = summary.loc[summary["width"] == int(width)].iloc[0]
+        local_summary = summarize(np.clip(predicted_gap, 1e-12, None))
+        local_summary.update(
+            {
+                "width": int(width),
+                "fa_capacity_margin": int(row["fa_capacity_margin"]),
+                "parameter_count": int(row["parameter_count"]),
+                "fa_constraint_rank": int(row["fa_constraint_rank"]),
+            }
+        )
+        local_gap_records.append(local_summary)
+    local_gap_df = pd.DataFrame(local_gap_records)
+    local_gap_csv = outdir / "local_operator_scalar_gap_summary.csv"
+    local_gap_df.to_csv(local_gap_csv, index=False)
+
+    paths = [
+        plot_gap_distributions(empirical_df, theory_df, outdir),
+        plot_local_erosion(local, summary, rng, args.theory_samples, outdir),
+        plot_local_scalar_gap(empirical_df, local_gap_df, outdir),
+    ]
+    print(f"theory summary: {theory_csv}")
+    print(f"empirical summary: {empirical_csv}")
+    print(f"local erosion rows: {local_csv}")
+    print(f"local scalar gap summary: {local_gap_csv}")
+    for path in paths:
+        print(f"plot: {path}")
+
+
+if __name__ == "__main__":
+    main()