diff options
| author | YurenHao0426 <Blackhao0426@gmail.com> | 2026-06-05 15:11:00 -0500 |
|---|---|---|
| committer | YurenHao0426 <Blackhao0426@gmail.com> | 2026-06-05 15:11:00 -0500 |
| commit | 3d691df89e5cff5d0dda69e865b08de93ef75b15 (patch) | |
| tree | ae562bf3ae289f2af7e66e04fc892269221f242d | |
| parent | e200d9179acc2282f77de00a042c6c8619813668 (diff) | |
Compare soft erosion theory distributions
| -rw-r--r-- | notes/25_soft_erosion_theory_distribution_validation.md | 214 | ||||
| -rw-r--r-- | scripts/soft_erosion_theory_vs_empirical.py | 473 |
2 files changed, 687 insertions, 0 deletions
diff --git a/notes/25_soft_erosion_theory_distribution_validation.md b/notes/25_soft_erosion_theory_distribution_validation.md new file mode 100644 index 0000000..1b078d0 --- /dev/null +++ b/notes/25_soft_erosion_theory_distribution_validation.md @@ -0,0 +1,214 @@ +# Soft-Erosion Theory Distribution Validation Attempt + +We tested whether the revised soft-erosion theory can predict the empirical +FA/BP train-gap distribution from the dense `T=30000` sweep. + +## Data + +Empirical run: + +```text +outputs/phase_transition_dense_T30000_352traj +``` + +Validation output: + +```text +outputs/soft_erosion_theory_vs_empirical +``` + +Main plots: + +```text +outputs/soft_erosion_theory_vs_empirical/gap_distribution_theory_vs_empirical.png +outputs/soft_erosion_theory_vs_empirical/local_erosion_distribution_theory_vs_empirical.png +outputs/soft_erosion_theory_vs_empirical/local_operator_scalar_gap_theory_vs_empirical.png +``` + +Script: + +```text +scripts/soft_erosion_theory_vs_empirical.py +``` + +## Theory 1: Hard-k Random-Subspace Gap Distribution + +This test uses the null model: + +```text +P = parameter count +k = hard FA constraint count +rho ~ Beta((P-k)/2, k/2) +K_FA ≈ rho K_BP +``` + +Then it predicts finite-time loss from the BP initial tangent spectrum: + +```text +r_T_BP = exp(-eta T K_BP / N) r_0 +r_T_FA = exp(-eta T rho K_BP / N) r_0 +gap_T = (||r_T_FA||^2 - ||r_T_BP||^2) / (2N) +``` + +This is not fitted. The only random quantity is the random-subspace theory +distribution for `rho`. + +Result: + +```text +gap_distribution_theory_vs_empirical.png +``` + +The hard-k random-subspace theory strongly overpredicts the empirical gap, +especially on the positive-margin side. + +Examples: + +| width | margin | hard-k theory mean | empirical mean | +|---:|---:|---:|---:| +| 40 | 130 | 0.040374 | 0.000162 | +| 36 | 66 | 0.050133 | 0.000791 | +| 32 | 2 | 0.064534 | 0.003729 | +| 28 | -62 | 0.103630 | 0.014195 | +| 24 | -126 | 0.121998 | 0.046464 | +| 20 | -190 | 0.151140 | 0.146803 | + +So the hard constraint count `k` is not the effective burden. It is too +pessimistic. + +## Theory 2: Local One-Step Erosion + +We measured the actual initial local operator erosion for each feedback seed: + +```text +speed_ratio = <g_BP, g_FA> / <g_BP, g_BP> +erosion = 1 - speed_ratio +``` + +This directly estimates: + +```text +1 - (r^T S_FA r) / (r^T K_BP r) +``` + +where: + +```text +S_FA = (K_FA + K_FA^T) / 2 +``` + +The comparison: + +```text +local_erosion_distribution_theory_vs_empirical.png +``` + +shows that real FA local erosion is much smaller than the hard-k Beta null. + +Examples: + +| width | empirical local erosion mean | hard-k Beta erosion mean | +|---:|---:|---:| +| 24 | 0.3907 | about 0.634 | +| 32 | 0.2338 | about 0.691 | +| 40 | 0.3381 | about 0.733 | + +Again, hard `k` overcounts effective burden. + +## Theory 3: Local-Operator Scalar Gap Prediction + +We then used the measured initial speed ratio as a no-fit scalar operator +prediction: + +```text +K_FA ≈ speed_ratio(B) K_BP +``` + +and propagated the BP tangent spectrum for `T=30000`. + +Result: + +```text +local_operator_scalar_gap_theory_vs_empirical.png +``` + +This improves over hard-k random subspace but is still not sufficient. + +Examples: + +| width | margin | local-operator theory mean | empirical mean | +|---:|---:|---:|---:| +| 40 | 130 | 0.002701 | 0.000162 | +| 36 | 66 | 0.004612 | 0.000791 | +| 32 | 2 | 0.005276 | 0.003729 | +| 28 | -62 | 0.026043 | 0.014195 | +| 24 | -126 | 0.047863 | 0.046464 | +| 22 | -158 | 0.037249 | 0.100526 | +| 20 | -190 | 0.071783 | 0.146803 | + +It matches around `width=24` and is roughly right near `width=32`, but: + +1. it overpredicts the positive-margin tail; +2. it underpredicts the more negative-margin regime; +3. it is not monotone enough because a single initial scalar cannot capture + training-time alignment and feature drift. + +## Conclusion + +The current result is not a successful validation of the hard-k random-subspace +theory. + +It gives a sharper diagnosis: + +```text +hard FA constraint count k is not k_eff. +``` + +The empirical gap is controlled by an effective operator erosion that evolves +during training: + +```text +K_FA,t = J_t J_tilde_t^T +``` + +not by a fixed random deletion of a `k`-dimensional subspace. + +## Revised Theory Target + +The next theory should predict a distribution over: + +```text +e_t(B) = 1 - (r_t^T S_FA,t r_t) / (r_t^T K_BP,t r_t) +``` + +or its cumulative finite-time form: + +```text +r_T_FA = product_t (I - eta K_FA,t / N) r_0 +r_T_BP = product_t (I - eta K_BP,t / N) r_0 +``` + +The right effective capacity is therefore: + +```text +k_eff(B,t) = P * e_t(B) +``` + +or a spectral version of this quantity, not the raw hard matrix constraint +count. + +## Relation to Song, Xu, and Lafferty + +`Convergence and Alignment of Gradient Descent with Random Backpropagation +Weights` is highly relevant and must be read carefully next. + +Current hypothesis: + +```text +Their overparameterized convergence theorem is compatible with our findings. +It shows FA can converge; it does not imply BP-equivalent finite-time tangent +capacity. +``` + +We need to check whether their alignment dynamics imply a decay law for +`e_t(B)` or `k_eff(B,t)`. 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() |
