path: root/scripts/trajectory_mlp_fa.py

#!/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()