#!/usr/bin/env python3 """Match FA/BP trajectory gap distributions against a BP-path bridge predictor. This is a first distributional bridge for realized FA loss gaps. For each feedback seed B, it compares: 1. Empirical gap: train the actual FA trajectory and measure L(theta_T^FA) - L(theta_T^BP). 2. Predicted bridge gap: integrate the FA/BP gradient mismatch along the fixed BP trajectory, delta_T(B) = -eta sum_t [g_FA(theta_t^BP; B) - g_BP(theta_t^BP)], then evaluate L(theta_T^BP + delta_T(B)) - L(theta_T^BP). This is not a final theorem. It is the first test of whether a trajectory-level bridge can predict a full gap distribution, not merely correlations. """ from __future__ import annotations import argparse import csv import json from dataclasses import asdict, dataclass from pathlib import Path import matplotlib.pyplot as plt import numpy as np from scipy import stats import trajectory_mlp_fa as tm @dataclass(frozen=True) class RunConfig: input_dim: int hidden_widths: list[int] output_dim: int samples: int steps: int lr: float data_seed: int init_seed: int feedback_seed_start: int feedback_runs: int feedback_init: str feedback_scale: str noise_std: float outdir: str plot: bool @dataclass(frozen=True) class GapRow: feedback_seed: int empirical_gap: float predicted_bridge_gap: float empirical_final_loss: float predicted_bridge_loss: float bp_final_loss: float initial_hidden_gradient_cosine: float final_hidden_gradient_cosine: float initial_q_mean: float final_q_mean: float bridge_delta_norm: float empirical_delta_norm: float @dataclass(frozen=True) class MatchSummary: runs: int bp_final_loss: float empirical_mean: float predicted_mean: float empirical_std: float predicted_std: float empirical_q01: float predicted_q01: float empirical_q50: float predicted_q50: float empirical_q99: float predicted_q99: float ks_2sample_statistic: float ks_2sample_pvalue: float standardized_ks_statistic: float standardized_ks_pvalue: float moment_matched_ks_statistic: float moment_matched_ks_pvalue: float moment_matched_wasserstein_distance: float wasserstein_distance: float pearson_r: float pearson_p: float spearman_r: float spearman_p: float def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="FA/BP trajectory gap distribution matching." ) parser.add_argument("--input-dim", type=int, default=16) parser.add_argument("--hidden-widths", type=int, nargs="+", default=[16, 16]) parser.add_argument("--output-dim", type=int, default=4) parser.add_argument("--samples", type=int, default=256) parser.add_argument("--steps", type=int, default=120) parser.add_argument("--lr", type=float, default=0.02) parser.add_argument("--data-seed", type=int, default=20) parser.add_argument("--init-seed", type=int, default=30) parser.add_argument("--feedback-seed-start", type=int, default=5_000) parser.add_argument("--feedback-runs", type=int, default=500) parser.add_argument( "--feedback-init", choices=["gaussian", "rademacher"], default="gaussian", ) parser.add_argument( "--feedback-scale", choices=["relu", "fan-in", "unit"], default="relu", ) parser.add_argument("--noise-std", type=float, default=0.01) parser.add_argument( "--outdir", type=Path, default=Path("outputs/trajectory_gap_distribution"), ) parser.add_argument("--plot", action="store_true") return parser.parse_args() def parse_config(args: argparse.Namespace) -> RunConfig: return RunConfig( input_dim=args.input_dim, hidden_widths=args.hidden_widths, output_dim=args.output_dim, samples=args.samples, steps=args.steps, lr=args.lr, data_seed=args.data_seed, init_seed=args.init_seed, feedback_seed_start=args.feedback_seed_start, feedback_runs=args.feedback_runs, feedback_init=args.feedback_init, feedback_scale=args.feedback_scale, noise_std=args.noise_std, outdir=str(args.outdir), plot=args.plot, ) def tm_config(config: RunConfig) -> tm.RunConfig: return tm.RunConfig( input_dim=config.input_dim, hidden_widths=config.hidden_widths, output_dim=config.output_dim, samples=config.samples, steps=config.steps, lr=config.lr, eval_every=config.steps, data_seed=config.data_seed, init_seed=config.init_seed, feedback_seed_start=config.feedback_seed_start, feedback_runs=config.feedback_runs, feedback_init=config.feedback_init, feedback_scale=config.feedback_scale, noise_std=config.noise_std, outdir=config.outdir, plot=False, ) def validate_config(config: RunConfig) -> None: run_config = tm_config(config) tm.validate_config(run_config) def zeros_like(weights: list[tm.Array]) -> list[tm.Array]: return [np.zeros_like(weight) for weight in weights] def copy_weights(weights: list[tm.Array]) -> list[tm.Array]: return [weight.copy() for weight in weights] def add_scaled(target: list[tm.Array], source: list[tm.Array], scale: float) -> None: for target_array, source_array in zip(target, source): target_array += scale * source_array def add_weights(a: list[tm.Array], b: list[tm.Array]) -> list[tm.Array]: return [left + right for left, right in zip(a, b)] def weights_delta(a: list[tm.Array], b: list[tm.Array]) -> list[tm.Array]: return [left - right for left, right in zip(a, b)] def norm_weights(weights: list[tm.Array]) -> float: return float(np.linalg.norm(tm.flatten(weights))) def train_bp_path( initial_weights: list[tm.Array], x: tm.Array, y: tm.Array, lr: float, steps: int ) -> tuple[list[list[tm.Array]], list[list[tm.Array]], list[tm.Array], float]: weights = copy_weights(initial_weights) path: list[list[tm.Array]] = [] bp_grads: list[list[tm.Array]] = [] for _step in range(steps): path.append(copy_weights(weights)) grads, _ = tm.gradients(weights, x, y, feedback=None) bp_grads.append(copy_weights(grads)) tm.sgd_step(weights, grads, lr) bp_final_loss = tm.mse_loss(tm.predict(weights, x), y) return path, bp_grads, weights, bp_final_loss def bridge_delta( bp_path: list[list[tm.Array]], bp_grads: list[list[tm.Array]], feedback: list[tm.Array], x: tm.Array, y: tm.Array, lr: float, ) -> list[tm.Array]: delta = zeros_like(bp_path[0]) for weights, bp_grad in zip(bp_path, bp_grads): fa_grad, _ = tm.gradients(weights, x, y, feedback=feedback) mismatch = [fa - bp for fa, bp in zip(fa_grad, bp_grad)] add_scaled(delta, mismatch, -lr) return delta def run_one_feedback( config: RunConfig, run_config: tm.RunConfig, initial_weights: list[tm.Array], bp_path: list[list[tm.Array]], bp_grads: list[list[tm.Array]], bp_final_weights: list[tm.Array], bp_final_loss: float, x: tm.Array, y: tm.Array, feedback_seed: int, ) -> GapRow: widths = tm.layer_widths(run_config) feedback = tm.init_feedback( widths, feedback_seed, config.feedback_init, config.feedback_scale ) delta = bridge_delta(bp_path, bp_grads, feedback, x, y, config.lr) bridge_weights = add_weights(bp_final_weights, delta) bridge_loss = tm.mse_loss(tm.predict(bridge_weights, x), y) predicted_gap = bridge_loss - bp_final_loss fa_final_weights, trajectory, _layer_metrics = tm.train_fa( initial_weights, feedback, feedback_seed, x, y, run_config ) empirical_final_loss = trajectory[-1].loss empirical_gap = empirical_final_loss - bp_final_loss initial = trajectory[0] final = trajectory[-1] empirical_delta = weights_delta(fa_final_weights, bp_final_weights) return GapRow( feedback_seed=feedback_seed, empirical_gap=empirical_gap, predicted_bridge_gap=predicted_gap, empirical_final_loss=empirical_final_loss, predicted_bridge_loss=bridge_loss, bp_final_loss=bp_final_loss, initial_hidden_gradient_cosine=float(initial.hidden_gradient_cosine), final_hidden_gradient_cosine=float(final.hidden_gradient_cosine), initial_q_mean=float(initial.q_mean), final_q_mean=float(final.q_mean), bridge_delta_norm=norm_weights(delta), empirical_delta_norm=norm_weights(empirical_delta), ) def summarize(rows: list[GapRow], bp_final_loss: float) -> MatchSummary: empirical = np.array([row.empirical_gap for row in rows], dtype=np.float64) predicted = np.array([row.predicted_bridge_gap for row in rows], dtype=np.float64) ks = stats.ks_2samp(predicted, empirical) predicted_std = np.std(predicted, ddof=1) empirical_std = np.std(empirical, ddof=1) predicted_z = (predicted - np.mean(predicted)) / predicted_std empirical_z = (empirical - np.mean(empirical)) / empirical_std standardized_ks = stats.ks_2samp(predicted_z, empirical_z) moment_matched = np.mean(empirical) + predicted_z * empirical_std moment_matched_ks = stats.ks_2samp(moment_matched, empirical) pearson = stats.pearsonr(predicted, empirical) spearman = stats.spearmanr(predicted, empirical) return MatchSummary( runs=len(rows), bp_final_loss=bp_final_loss, empirical_mean=float(np.mean(empirical)), predicted_mean=float(np.mean(predicted)), empirical_std=float(empirical_std), predicted_std=float(predicted_std), empirical_q01=float(np.quantile(empirical, 0.01)), predicted_q01=float(np.quantile(predicted, 0.01)), empirical_q50=float(np.quantile(empirical, 0.50)), predicted_q50=float(np.quantile(predicted, 0.50)), empirical_q99=float(np.quantile(empirical, 0.99)), predicted_q99=float(np.quantile(predicted, 0.99)), ks_2sample_statistic=float(ks.statistic), ks_2sample_pvalue=float(ks.pvalue), standardized_ks_statistic=float(standardized_ks.statistic), standardized_ks_pvalue=float(standardized_ks.pvalue), moment_matched_ks_statistic=float(moment_matched_ks.statistic), moment_matched_ks_pvalue=float(moment_matched_ks.pvalue), moment_matched_wasserstein_distance=float( stats.wasserstein_distance(moment_matched, empirical) ), wasserstein_distance=float(stats.wasserstein_distance(predicted, empirical)), pearson_r=float(pearson.statistic), pearson_p=float(pearson.pvalue), spearman_r=float(spearman.statistic), spearman_p=float(spearman.pvalue), ) def write_csv(path: Path, rows: list[object]) -> None: if not rows: return path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", newline="") as handle: first = asdict(rows[0]) # type: ignore[arg-type] writer = csv.DictWriter(handle, fieldnames=list(first.keys())) writer.writeheader() for row in rows: writer.writerow(asdict(row)) # type: ignore[arg-type] def write_outputs( config: RunConfig, rows: list[GapRow], summary: MatchSummary, outdir: Path ) -> None: outdir.mkdir(parents=True, exist_ok=True) write_csv(outdir / "gap_distribution.csv", rows) write_csv(outdir / "summary.csv", [summary]) payload = { "config": asdict(config), "summary": asdict(summary), "runs": [asdict(row) for row in rows], } (outdir / "summary.json").write_text( json.dumps(payload, indent=2, sort_keys=True) + "\n" ) def save_plots(rows: list[GapRow], outdir: Path) -> list[Path]: outdir.mkdir(parents=True, exist_ok=True) paths: list[Path] = [] empirical = np.array([row.empirical_gap for row in rows], dtype=np.float64) predicted = np.array([row.predicted_bridge_gap for row in rows], dtype=np.float64) predicted_z = (predicted - np.mean(predicted)) / np.std(predicted, ddof=1) empirical_z = (empirical - np.mean(empirical)) / np.std(empirical, ddof=1) moment_matched = np.mean(empirical) + predicted_z * np.std(empirical, ddof=1) hist_path = outdir / "gap_histogram_overlay.png" bins = np.histogram_bin_edges(np.concatenate([empirical, predicted]), bins=60) plt.figure(figsize=(7, 4.5)) plt.hist(predicted, bins=bins, density=True, alpha=0.45, label="bridge predicted") plt.hist(empirical, bins=bins, density=True, alpha=0.45, label="actual FA") plt.xlabel("final loss gap to BP") plt.ylabel("density") plt.title("Predicted vs empirical FA/BP gap distribution") plt.legend() plt.tight_layout() plt.savefig(hist_path, dpi=180) plt.close() paths.append(hist_path) matched_hist_path = outdir / "gap_histogram_moment_matched.png" bins = np.histogram_bin_edges(np.concatenate([empirical, moment_matched]), bins=60) plt.figure(figsize=(7, 4.5)) plt.hist(moment_matched, bins=bins, density=True, alpha=0.45, label="bridge moment-matched") plt.hist(empirical, bins=bins, density=True, alpha=0.45, label="actual FA") plt.xlabel("final loss gap to BP") plt.ylabel("density") plt.title("Moment-matched bridge vs empirical gap distribution") plt.legend() plt.tight_layout() plt.savefig(matched_hist_path, dpi=180) plt.close() paths.append(matched_hist_path) cdf_path = outdir / "gap_cdf_overlay.png" plt.figure(figsize=(7, 4.5)) for values, label in [(predicted, "bridge predicted"), (empirical, "actual FA")]: sorted_values = np.sort(values) probs = (np.arange(1, len(values) + 1) - 0.5) / len(values) plt.plot(sorted_values, probs, label=label) plt.xlabel("final loss gap to BP") plt.ylabel("CDF") plt.title("Gap distribution CDF") plt.legend() plt.tight_layout() plt.savefig(cdf_path, dpi=180) plt.close() paths.append(cdf_path) qq_path = outdir / "gap_qq_plot.png" probs = (np.arange(1, len(empirical) + 1) - 0.5) / len(empirical) pred_q = np.quantile(predicted, probs) emp_q = np.quantile(empirical, probs) min_value = float(min(pred_q[0], emp_q[0])) max_value = float(max(pred_q[-1], emp_q[-1])) plt.figure(figsize=(5, 5)) plt.scatter(pred_q, emp_q, s=8, alpha=0.35) plt.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1) plt.xlabel("bridge predicted quantile") plt.ylabel("actual FA quantile") plt.title("Gap distribution Q-Q") plt.tight_layout() plt.savefig(qq_path, dpi=180) plt.close() paths.append(qq_path) standardized_qq_path = outdir / "gap_standardized_qq_plot.png" pred_q = np.quantile(predicted_z, probs) emp_q = np.quantile(empirical_z, probs) min_value = float(min(pred_q[0], emp_q[0])) max_value = float(max(pred_q[-1], emp_q[-1])) plt.figure(figsize=(5, 5)) plt.scatter(pred_q, emp_q, s=8, alpha=0.35) plt.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1) plt.xlabel("standardized bridge quantile") plt.ylabel("standardized actual FA quantile") plt.title("Standardized gap distribution Q-Q") plt.tight_layout() plt.savefig(standardized_qq_path, dpi=180) plt.close() paths.append(standardized_qq_path) scatter_path = outdir / "paired_gap_scatter.png" plt.figure(figsize=(5, 5)) plt.scatter(predicted, empirical, s=12, alpha=0.45) min_value = float(min(np.min(predicted), np.min(empirical))) max_value = float(max(np.max(predicted), np.max(empirical))) plt.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1) plt.xlabel("bridge predicted gap") plt.ylabel("actual FA gap") plt.title("Paired feedback seeds") plt.tight_layout() plt.savefig(scatter_path, dpi=180) plt.close() paths.append(scatter_path) return paths def main() -> None: args = parse_args() config = parse_config(args) validate_config(config) run_config = tm_config(config) widths = tm.layer_widths(run_config) x, y = tm.make_synthetic_regression(run_config) initial_weights = tm.init_weights(widths, config.init_seed) bp_path, bp_grads, bp_final_weights, bp_final_loss = train_bp_path( initial_weights, x, y, config.lr, config.steps ) rows: list[GapRow] = [] for run_index in range(config.feedback_runs): feedback_seed = config.feedback_seed_start + run_index row = run_one_feedback( config, run_config, initial_weights, bp_path, bp_grads, bp_final_weights, bp_final_loss, x, y, feedback_seed, ) rows.append(row) if (run_index + 1) % max(1, config.feedback_runs // 10) == 0: print( f"completed {run_index + 1}/{config.feedback_runs}: " f"emp_gap={row.empirical_gap:.6g}, pred_gap={row.predicted_bridge_gap:.6g}" ) summary = summarize(rows, bp_final_loss) outdir = Path(config.outdir) write_outputs(config, rows, summary, outdir) plot_paths = save_plots(rows, outdir) if config.plot else [] print(f"widths: {widths}") print(f"bp_final_loss: {bp_final_loss:.8g}") print( "empirical_gap: " f"mean={summary.empirical_mean:.8g}, std={summary.empirical_std:.8g}" ) print( "predicted_bridge_gap: " f"mean={summary.predicted_mean:.8g}, std={summary.predicted_std:.8g}" ) print( "distribution_match: " f"KS={summary.ks_2sample_statistic:.8g}, " f"W1={summary.wasserstein_distance:.8g}, " f"Spearman={summary.spearman_r:.8g}" ) print( "shape_match_after_standardization: " f"KS={summary.standardized_ks_statistic:.8g}, " f"p={summary.standardized_ks_pvalue:.8g}, " f"moment_matched_W1={summary.moment_matched_wasserstein_distance:.8g}" ) print(f"gap_distribution: {outdir / 'gap_distribution.csv'}") print(f"summary: {outdir / 'summary.csv'}") for path in plot_paths: print(f"plot: {path}") if __name__ == "__main__": main()