Add synthetic FA trajectory validation

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()