diff options
Diffstat (limited to 'scripts')
| -rw-r--r-- | scripts/README.md | 22 | ||||
| -rwxr-xr-x | scripts/trajectory_mlp_fa.py | 613 |
2 files changed, 635 insertions, 0 deletions
diff --git a/scripts/README.md b/scripts/README.md index f24c117..42633d0 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -111,3 +111,25 @@ and: \] Outputs are written under `outputs/functional_capacity_overlap/`. + +## Synthetic MLP FA/BP Trajectories + +Run: + +```bash +python scripts/trajectory_mlp_fa.py --samples 128 --hidden-widths 24 24 --steps 80 --lr 0.02 --eval-every 10 --feedback-runs 3 --data-seed 3 --init-seed 4 --feedback-seed-start 50 --plot +``` + +This trains one BP baseline and several FA runs from the same initial weights on +a synthetic regression task. At each checkpoint, the script records: + +- training loss; +- full-model cosine between the BP gradient and the FA surrogate gradient at the FA weights; +- layerwise \(Q_l=\cos^2(W_{l+1}^{\top},B_l)\). + +Outputs are written under `outputs/trajectory_mlp_fa/`: + +- `summary.csv` +- `trajectories.csv` +- `layer_metrics.csv` +- diagnostic plots when `--plot` is set. diff --git a/scripts/trajectory_mlp_fa.py b/scripts/trajectory_mlp_fa.py new file mode 100755 index 0000000..ce20fa5 --- /dev/null +++ b/scripts/trajectory_mlp_fa.py @@ -0,0 +1,613 @@ +#!/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 + 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_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, + 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)) + 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, + 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_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_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) + + 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" + ] + 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(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() |