#!/usr/bin/env python3 """Train BP and FA MLPs on synthetic regression and log trajectory metrics.""" from __future__ import annotations import argparse import csv import json import math from dataclasses import asdict, dataclass from pathlib import Path import matplotlib.pyplot as plt import numpy as np Array = np.ndarray @dataclass(frozen=True) class RunConfig: input_dim: int hidden_widths: list[int] output_dim: int samples: int steps: int lr: float eval_every: int 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 TrajectoryRow: run_type: str feedback_seed: int step: int loss: float gradient_cosine: float | None hidden_gradient_cosine: float | None q_mean: float | None q_min: float | None q_max: float | None @dataclass(frozen=True) class LayerMetricRow: run_type: str feedback_seed: int step: int layer: int gradient_cosine: float | None q_alignment: float | None @dataclass(frozen=True) class RunSummary: run_type: str feedback_seed: int final_loss: float final_gap_to_bp: float initial_gradient_cosine: float | None final_gradient_cosine: float | None initial_hidden_gradient_cosine: float | None final_hidden_gradient_cosine: float | None initial_q_mean: float | None final_q_mean: float | None def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Synthetic MLP trajectory validation for BP vs feedback alignment." ) parser.add_argument("--input-dim", type=int, default=16) parser.add_argument("--hidden-widths", type=int, nargs="+", default=[32, 32]) parser.add_argument("--output-dim", type=int, default=4) parser.add_argument("--samples", type=int, default=256) parser.add_argument("--steps", type=int, default=200) parser.add_argument("--lr", type=float, default=0.03) parser.add_argument("--eval-every", type=int, default=10) parser.add_argument("--data-seed", type=int, default=0) parser.add_argument("--init-seed", type=int, default=1) parser.add_argument("--feedback-seed-start", type=int, default=100) parser.add_argument("--feedback-runs", type=int, default=5) parser.add_argument( "--feedback-init", choices=["gaussian", "rademacher"], default="gaussian", ) parser.add_argument( "--feedback-scale", choices=["relu", "fan-in", "unit"], default="relu", help="relu uses sqrt(2 / n_l); fan-in uses sqrt(1 / n_l).", ) parser.add_argument("--noise-std", type=float, default=0.01) parser.add_argument( "--outdir", type=Path, default=Path("outputs/trajectory_mlp_fa"), ) 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, eval_every=args.eval_every, 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 validate_config(config: RunConfig) -> None: widths = [config.input_dim, *config.hidden_widths, config.output_dim] if any(width < 1 for width in widths): raise ValueError("All layer widths must be positive.") if len(config.hidden_widths) < 1: raise ValueError("At least one hidden layer is required for FA metrics.") if config.samples < 1: raise ValueError("--samples must be positive.") if config.steps < 1: raise ValueError("--steps must be positive.") if config.lr <= 0: raise ValueError("--lr must be positive.") if config.eval_every < 1: raise ValueError("--eval-every must be positive.") if config.feedback_runs < 1: raise ValueError("--feedback-runs must be positive.") if config.noise_std < 0: raise ValueError("--noise-std must be non-negative.") def layer_widths(config: RunConfig) -> list[int]: return [config.input_dim, *config.hidden_widths, config.output_dim] def relu(x: Array) -> Array: return np.maximum(x, 0.0) def init_weights(widths: list[int], seed: int) -> list[Array]: rng = np.random.default_rng(seed) weights: list[Array] = [] last_index = len(widths) - 2 for layer, (fan_in, fan_out) in enumerate(zip(widths[:-1], widths[1:])): if layer == last_index: scale = 1.0 / math.sqrt(fan_in) else: scale = math.sqrt(2.0 / fan_in) weights.append(rng.standard_normal((fan_out, fan_in)) * scale) return weights def feedback_layer_scale(rows: int, mode: str) -> float: if mode == "relu": return math.sqrt(2.0 / rows) if mode == "fan-in": return math.sqrt(1.0 / rows) if mode == "unit": return 1.0 raise ValueError(f"Unknown feedback scale: {mode}") def init_feedback( widths: list[int], seed: int, distribution: str, scale_mode: str ) -> list[Array]: rng = np.random.default_rng(seed) feedback: list[Array] = [] # B_i replaces W_{i+1}^T for hidden layer i, so shape is n_i x n_{i+1}. for hidden_index in range(1, len(widths) - 1): rows = widths[hidden_index] cols = widths[hidden_index + 1] scale = feedback_layer_scale(rows, scale_mode) if distribution == "gaussian": matrix = rng.standard_normal((rows, cols)) * scale elif distribution == "rademacher": matrix = rng.choice(np.array([-1.0, 1.0]), size=(rows, cols)) * scale else: raise ValueError(f"Unknown feedback distribution: {distribution}") feedback.append(matrix.astype(np.float64)) return feedback def forward(weights: list[Array], x: Array) -> tuple[list[Array], list[Array]]: activations = [x] preactivations: list[Array] = [] hidden_last = len(weights) - 2 current = x for layer, weight in enumerate(weights): preactivation = current @ weight.T preactivations.append(preactivation) if layer <= hidden_last: current = relu(preactivation) else: current = preactivation activations.append(current) return activations, preactivations def predict(weights: list[Array], x: Array) -> Array: return forward(weights, x)[0][-1] def mse_loss(prediction: Array, target: Array) -> float: error = prediction - target return float(0.5 * np.mean(np.sum(error * error, axis=1))) def gradients( weights: list[Array], x: Array, target: Array, feedback: list[Array] | None, ) -> tuple[list[Array], float]: activations, preactivations = forward(weights, x) prediction = activations[-1] loss = mse_loss(prediction, target) batch_size = x.shape[0] deltas: list[Array] = [np.empty((0, 0)) for _ in weights] deltas[-1] = (prediction - target) / batch_size for layer in range(len(weights) - 2, -1, -1): if feedback is None: back_signal = deltas[layer + 1] @ weights[layer + 1] else: back_signal = deltas[layer + 1] @ feedback[layer].T deltas[layer] = back_signal * (preactivations[layer] > 0) grads = [delta.T @ activations[layer] for layer, delta in enumerate(deltas)] return grads, loss def sgd_step(weights: list[Array], grads: list[Array], lr: float) -> None: for weight, grad in zip(weights, grads): weight -= lr * grad def flatten(arrays: list[Array]) -> Array: return np.concatenate([array.ravel() for array in arrays]) def cosine(a: Array, b: Array) -> float: denom = np.linalg.norm(a) * np.linalg.norm(b) if denom == 0: return float("nan") return float(np.dot(a, b) / denom) def squared_frobenius_cosine(a: Array, b: Array) -> float: denom = np.linalg.norm(a) * np.linalg.norm(b) if denom == 0: return float("nan") value = float(np.sum(a * b) / denom) return value * value def layer_q_alignments(weights: list[Array], feedback: list[Array]) -> list[float]: values: list[float] = [] for hidden_index, feedback_matrix in enumerate(feedback): next_weight = weights[hidden_index + 1] values.append(squared_frobenius_cosine(next_weight.T, feedback_matrix)) return values def layer_gradient_cosines(bp_grads: list[Array], fa_grads: list[Array]) -> list[float]: return [cosine(bp.ravel(), fa.ravel()) for bp, fa in zip(bp_grads, fa_grads)] def make_synthetic_regression(config: RunConfig) -> tuple[Array, Array]: rng = np.random.default_rng(config.data_seed) x = rng.standard_normal((config.samples, config.input_dim)) x = (x - x.mean(axis=0, keepdims=True)) / (x.std(axis=0, keepdims=True) + 1e-8) teacher_widths = [config.input_dim, *config.hidden_widths, config.output_dim] teacher = init_weights(teacher_widths, config.data_seed + 10_000) y = predict(teacher, x) if config.noise_std > 0: y = y + rng.standard_normal(y.shape) * config.noise_std y = y - y.mean(axis=0, keepdims=True) return x.astype(np.float64), y.astype(np.float64) def evaluate_bp(weights: list[Array], x: Array, y: Array, step: int) -> TrajectoryRow: loss = mse_loss(predict(weights, x), y) return TrajectoryRow( run_type="bp", feedback_seed=-1, step=step, loss=loss, gradient_cosine=1.0, hidden_gradient_cosine=1.0, q_mean=None, q_min=None, q_max=None, ) def evaluate_fa( weights: list[Array], feedback: list[Array], x: Array, y: Array, seed: int, step: int, ) -> tuple[TrajectoryRow, list[LayerMetricRow]]: bp_grads, loss = gradients(weights, x, y, feedback=None) fa_grads, _ = gradients(weights, x, y, feedback=feedback) grad_cos = cosine(flatten(bp_grads), flatten(fa_grads)) hidden_grad_cos = cosine(flatten(bp_grads[:-1]), flatten(fa_grads[:-1])) layer_cosines = layer_gradient_cosines(bp_grads, fa_grads) q_values = layer_q_alignments(weights, feedback) trajectory = TrajectoryRow( run_type="fa", feedback_seed=seed, step=step, loss=loss, gradient_cosine=grad_cos, hidden_gradient_cosine=hidden_grad_cos, q_mean=float(np.mean(q_values)), q_min=float(np.min(q_values)), q_max=float(np.max(q_values)), ) layer_rows: list[LayerMetricRow] = [] for layer, layer_cos in enumerate(layer_cosines): q_value = q_values[layer] if layer < len(q_values) else None layer_rows.append( LayerMetricRow( run_type="fa", feedback_seed=seed, step=step, layer=layer, gradient_cosine=layer_cos, q_alignment=q_value, ) ) return trajectory, layer_rows def train_bp( initial_weights: list[Array], x: Array, y: Array, config: RunConfig, ) -> tuple[list[Array], list[TrajectoryRow]]: weights = [weight.copy() for weight in initial_weights] trajectory: list[TrajectoryRow] = [] for step in range(config.steps + 1): if step % config.eval_every == 0 or step == config.steps: trajectory.append(evaluate_bp(weights, x, y, step)) if step == config.steps: break grads, _ = gradients(weights, x, y, feedback=None) sgd_step(weights, grads, config.lr) return weights, trajectory def train_fa( initial_weights: list[Array], feedback: list[Array], feedback_seed: int, x: Array, y: Array, config: RunConfig, ) -> tuple[list[Array], list[TrajectoryRow], list[LayerMetricRow]]: weights = [weight.copy() for weight in initial_weights] trajectory: list[TrajectoryRow] = [] layer_metrics: list[LayerMetricRow] = [] for step in range(config.steps + 1): if step % config.eval_every == 0 or step == config.steps: row, layer_rows = evaluate_fa(weights, feedback, x, y, feedback_seed, step) trajectory.append(row) layer_metrics.extend(layer_rows) if step == config.steps: break grads, _ = gradients(weights, x, y, feedback=feedback) sgd_step(weights, grads, config.lr) return weights, trajectory, layer_metrics 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 make_summaries( bp_trajectory: list[TrajectoryRow], fa_trajectories: dict[int, list[TrajectoryRow]], ) -> list[RunSummary]: bp_final = bp_trajectory[-1].loss summaries = [ RunSummary( run_type="bp", feedback_seed=-1, final_loss=bp_final, final_gap_to_bp=0.0, initial_gradient_cosine=1.0, final_gradient_cosine=1.0, initial_hidden_gradient_cosine=1.0, final_hidden_gradient_cosine=1.0, initial_q_mean=None, final_q_mean=None, ) ] for seed, trajectory in sorted(fa_trajectories.items()): first = trajectory[0] final = trajectory[-1] summaries.append( RunSummary( run_type="fa", feedback_seed=seed, final_loss=final.loss, final_gap_to_bp=final.loss - bp_final, initial_gradient_cosine=first.gradient_cosine, final_gradient_cosine=final.gradient_cosine, initial_hidden_gradient_cosine=first.hidden_gradient_cosine, final_hidden_gradient_cosine=final.hidden_gradient_cosine, initial_q_mean=first.q_mean, final_q_mean=final.q_mean, ) ) return summaries def write_outputs( config: RunConfig, trajectories: list[TrajectoryRow], layer_metrics: list[LayerMetricRow], summaries: list[RunSummary], outdir: Path, ) -> None: outdir.mkdir(parents=True, exist_ok=True) write_csv(outdir / "trajectories.csv", trajectories) write_csv(outdir / "layer_metrics.csv", layer_metrics) write_csv(outdir / "summary.csv", summaries) payload = { "config": asdict(config), "summary": [asdict(row) for row in summaries], } (outdir / "summary.json").write_text( json.dumps(payload, indent=2, sort_keys=True) + "\n" ) def save_plots(trajectories: list[TrajectoryRow], outdir: Path) -> list[Path]: outdir.mkdir(parents=True, exist_ok=True) paths: list[Path] = [] bp_rows = [row for row in trajectories if row.run_type == "bp"] fa_seeds = sorted( {row.feedback_seed for row in trajectories if row.run_type == "fa"} ) loss_path = outdir / "loss_curves.png" plt.figure(figsize=(7, 4.5)) plt.plot( [row.step for row in bp_rows], [row.loss for row in bp_rows], color="black", linewidth=2, label="BP", ) for seed in fa_seeds: rows = [ row for row in trajectories if row.run_type == "fa" and row.feedback_seed == seed ] plt.plot([row.step for row in rows], [row.loss for row in rows], alpha=0.65) plt.xlabel("step") plt.ylabel("training loss") plt.title("BP and FA synthetic regression trajectories") plt.legend() plt.tight_layout() plt.savefig(loss_path, dpi=180) plt.close() paths.append(loss_path) gamma_path = outdir / "gradient_cosine_curves.png" plt.figure(figsize=(7, 4.5)) for seed in fa_seeds: rows = [ row for row in trajectories if row.run_type == "fa" and row.feedback_seed == seed ] plt.plot( [row.step for row in rows], [row.gradient_cosine for row in rows], alpha=0.75, label=f"seed={seed}", ) plt.axhline(0.0, color="black", linewidth=1) plt.xlabel("step") plt.ylabel("cos(BP gradient, FA gradient)") plt.title("Surrogate gradient alignment") plt.tight_layout() plt.savefig(gamma_path, dpi=180) plt.close() paths.append(gamma_path) hidden_gamma_path = outdir / "hidden_gradient_cosine_curves.png" plt.figure(figsize=(7, 4.5)) for seed in fa_seeds: rows = [ row for row in trajectories if row.run_type == "fa" and row.feedback_seed == seed ] plt.plot( [row.step for row in rows], [row.hidden_gradient_cosine for row in rows], alpha=0.75, label=f"seed={seed}", ) plt.axhline(0.0, color="black", linewidth=1) plt.xlabel("step") plt.ylabel("cos(BP hidden gradient, FA hidden gradient)") plt.title("Hidden-layer surrogate gradient alignment") plt.tight_layout() plt.savefig(hidden_gamma_path, dpi=180) plt.close() paths.append(hidden_gamma_path) q_path = outdir / "q_alignment_curves.png" plt.figure(figsize=(7, 4.5)) for seed in fa_seeds: rows = [ row for row in trajectories if row.run_type == "fa" and row.feedback_seed == seed ] plt.plot( [row.step for row in rows], [row.q_mean for row in rows], alpha=0.75, label=f"seed={seed}", ) plt.xlabel("step") plt.ylabel("mean layerwise Q") plt.title("Weight-feedback alignment") plt.tight_layout() plt.savefig(q_path, dpi=180) plt.close() paths.append(q_path) return paths def main() -> None: args = parse_args() config = parse_config(args) validate_config(config) widths = layer_widths(config) x, y = make_synthetic_regression(config) initial_weights = init_weights(widths, config.init_seed) bp_weights, bp_trajectory = train_bp(initial_weights, x, y, config) del bp_weights all_trajectories = list(bp_trajectory) all_layer_metrics: list[LayerMetricRow] = [] fa_trajectories: dict[int, list[TrajectoryRow]] = {} for run_index in range(config.feedback_runs): feedback_seed = config.feedback_seed_start + run_index feedback = init_feedback( widths, feedback_seed, config.feedback_init, config.feedback_scale ) _, trajectory, layer_metrics = train_fa( initial_weights, feedback, feedback_seed, x, y, config ) fa_trajectories[feedback_seed] = trajectory all_trajectories.extend(trajectory) all_layer_metrics.extend(layer_metrics) summaries = make_summaries(bp_trajectory, fa_trajectories) outdir = Path(config.outdir) write_outputs(config, all_trajectories, all_layer_metrics, summaries, outdir) plot_paths = save_plots(all_trajectories, outdir) if config.plot else [] bp_final = summaries[0].final_loss fa_gaps = [row.final_gap_to_bp for row in summaries if row.run_type == "fa"] fa_final_gammas = [ row.final_gradient_cosine for row in summaries if row.run_type == "fa" ] fa_final_hidden_gammas = [ row.final_hidden_gradient_cosine for row in summaries if row.run_type == "fa" ] print(f"widths: {widths}") print(f"bp_final_loss: {bp_final:.8g}") print( "fa_gap_to_bp: " f"mean={np.mean(fa_gaps):.8g}, " f"min={np.min(fa_gaps):.8g}, " f"max={np.max(fa_gaps):.8g}" ) print( "fa_final_gradient_cosine: " f"mean={np.mean(fa_final_gammas):.8g}, " f"min={np.min(fa_final_gammas):.8g}, " f"max={np.max(fa_final_gammas):.8g}" ) print( "fa_final_hidden_gradient_cosine: " f"mean={np.mean(fa_final_hidden_gammas):.8g}, " f"min={np.min(fa_final_hidden_gammas):.8g}, " f"max={np.max(fa_final_hidden_gammas):.8g}" ) print(f"summary: {outdir / 'summary.csv'}") print(f"trajectories: {outdir / 'trajectories.csv'}") print(f"layer_metrics: {outdir / 'layer_metrics.csv'}") for path in plot_paths: print(f"plot: {path}") if __name__ == "__main__": main()