Add early-kernel operator predictor

3 files changed, 854 insertions, 0 deletions

diff --git a/notes/15_early_kernel_operator_predictor.md b/notes/15_early_kernel_operator_predictor.md
new file mode 100644
index 0000000..eacf852
--- /dev/null
+++ b/notes/15_early_kernel_operator_predictor.md
@@ -0,0 +1,154 @@
+# Early-Kernel Operator Predictor
+
+This note records the first predictive finite-\(T\) operator-level result.
+
+## Problem
+
+The fixed initialization tangent-kernel prediction
+
+\[
+r_{t+1}
+=
+\left(I-\frac{\eta}{N}K_0\right)r_t
+\]
+
+works at very short horizon, but at \(T=50\) it over-predicts the FA/BP train
+gap. The time-varying diagnostic in `notes/13_time_varying_kernel_diagnostic.md`
+showed that the error is caused mostly by kernel drift \(K_t-K_0\), not by a
+large second-order output residual.
+
+The goal here is a predictive, non-oracle correction: use only an early
+observation window, not the final loss gap.
+
+## Predictors
+
+Let \(s<T\) be an early probe step. For each BP or FA rule, compute
+
+\[
+K_0^{rule},
+\qquad
+K_s^{rule}.
+\]
+
+The strongest predictor in this probe is the early kernel velocity model:
+
+\[
+K_t^{rule}
+\approx
+K_0^{rule}
++
+t\frac{K_s^{rule}-K_0^{rule}}{s},
+\qquad
+0\le t<T.
+\]
+
+The residual prediction is then
+
+\[
+\hat r_{t+1}^{rule}
+=
+\left(
+I-\frac{\eta}{N}\hat K_t^{rule}
+\right)
+\hat r_t^{rule}.
+\]
+
+No fitted scale, offset, or final-gap calibration is used.
+
+We also tested two weaker alternatives:
+
+- FA-only compressed interpolation:
+  \[
+  K_t^{FA}
+  \approx
+  K_0^{FA}+\alpha_t(K_0^{BP}-K_0^{FA}),
+  \]
+  with \(\alpha_t\) inferred from the early projection of
+  \(K_s^{FA}-K_0^{FA}\).
+- early re-tangent:
+  run \(s\) empirical steps, then freeze \(K_s\) for the remaining \(T-s\)
+  steps.
+
+## Experiment
+
+- task: random-label regression;
+- architecture: \(16\to64\to64\to4\);
+- optimizer: full-batch SGD;
+- learning rate: \(10^{-3}\);
+- target horizon: \(T=50\);
+- early probe: \(s=20\);
+- train samples:
+  \[
+  N\in\{64,96,128,160,192,224,256,320\};
+  \]
+- 4 initialization seeds and 8 feedback seeds per \(N\);
+- 256 total trajectory points.
+
+Output directory:
+
+`outputs/compressed_operator_s20_256traj_T50_width64_plots`
+
+Key files:
+
+- `combined_rows.csv`
+- `predictor_metrics.csv`
+- `T50_s20_velocity_theory_band_empirical_trajectories.png`
+- `T50_s20_empirical_trajectories_only.png`
+- `T50_s20_prediction_scatter.png`
+- `T50_s20_prediction_error_by_margin.png`
+
+## Results
+
+Across all 256 trajectory points:
+
+| predictor | MAE | bias | corr |
+|---|---:|---:|---:|
+| fixed \(K(0)\) | 0.018487 | 0.018487 | 0.984460 |
+| FA compressed | 0.014915 | 0.014915 | 0.990472 |
+| early kernel velocity | 0.001893 | -0.001491 | 0.998848 |
+| early re-tangent | 0.004868 | 0.004868 | 0.999077 |
+
+Per-\(N\) summaries for the early kernel velocity model:
+
+| \(N\) | hard FA margin | fixed MAE | velocity MAE |
+|---:|---:|---:|---:|
+| 64 | 770 | 0.025444 | 0.001621 |
+| 96 | 642 | 0.021466 | 0.002578 |
+| 128 | 514 | 0.020630 | 0.001886 |
+| 160 | 386 | 0.017452 | 0.001879 |
+| 192 | 258 | 0.016190 | 0.001556 |
+| 224 | 130 | 0.015798 | 0.002354 |
+| 256 | 2 | 0.015586 | 0.001548 |
+| 320 | -254 | 0.015333 | 0.001725 |
+
+## Interpretation
+
+This is the first usable finite-\(T\) predictive version.
+
+The fixed \(K(0)\) theory gets the ordering and rough shape right, but its
+gap prediction is uniformly too high because it misses kernel drift. The early
+operator-velocity model removes almost all of that bias without using the final
+loss gap.
+
+The scalar directional-gain predictor in `notes/14_early_predictor_probe.md`
+failed because it compressed the dynamics too aggressively: it discarded
+changes in residual direction and in the full operator. The successful object
+is the operator \(K_t\), not a single scalar \(\lambda_t\).
+
+The early re-tangent model is conservative: it remains biased high. The linear
+velocity model is slightly biased low, so together they may define a clean
+no-fit prediction interval around the empirical gap.
+
+## Paper Use
+
+This should be folded into contribution 4:
+
+1. static distribution and local \(K(0)\) predictions;
+2. finite-\(T\) failure mode from kernel drift;
+3. predictive early-kernel correction;
+4. trajectory-level distribution validation.
+
+The current version is not yet a closed-form architecture-only theorem. It is a
+predictive finite-time theorem conditional on the early kernel pair
+\((K_0,K_s)\). That is acceptable if clearly separated from the prior-free
+initialization/capacity bounds.
diff --git a/scripts/compressed_operator_predictor.py b/scripts/compressed_operator_predictor.py
new file mode 100644
index 0000000..ce52c7a
--- /dev/null
+++ b/scripts/compressed_operator_predictor.py
@@ -0,0 +1,389 @@
+#!/usr/bin/env python3
+"""Compressed-operator finite-time predictor for FA/BP gaps."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import sys
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPT_DIR))
+
+import downstream_capacity_sweep as dcs  # noqa: E402
+from fa_tangent_kernel_capacity import pseudo_jacobian  # noqa: E402
+
+
+@dataclass(frozen=True)
+class CompressedRow:
+    init_seed: int
+    feedback_seed: int
+    target_steps: int
+    early_steps: int
+    empirical_bp_loss: float
+    empirical_fa_loss: float
+    empirical_gap: float
+    fixed_bp_loss: float
+    fixed_fa_loss: float
+    fixed_gap: float
+    compressed_fa_loss: float
+    compressed_gap: float
+    linear_bp_loss: float
+    linear_fa_loss: float
+    linear_gap: float
+    retangent_bp_loss: float
+    retangent_fa_loss: float
+    retangent_gap: float
+    alpha_s: float
+    alpha_slope: float
+    fixed_gap_error: float
+    compressed_gap_error: float
+    linear_gap_error: float
+    retangent_gap_error: float
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Compressed operator predictor.")
+    parser.add_argument("--input-dim", type=int, default=16)
+    parser.add_argument("--output-dim", type=int, default=4)
+    parser.add_argument("--width", type=int, default=64)
+    parser.add_argument("--train-samples", type=int, default=128)
+    parser.add_argument("--test-samples", type=int, default=512)
+    parser.add_argument("--target-steps", type=int, default=50)
+    parser.add_argument("--early-steps", type=int, nargs="+", default=[1, 2, 5])
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--init-seeds", type=int, default=2)
+    parser.add_argument("--feedback-seeds", type=int, default=4)
+    parser.add_argument("--data-seed", type=int, default=123)
+    parser.add_argument(
+        "--feedback-scale",
+        choices=["relu", "fan-in", "unit"],
+        default="relu",
+    )
+    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
+    parser.add_argument("--torch-threads", type=int, default=8)
+    parser.add_argument(
+        "--outdir",
+        type=Path,
+        default=Path("outputs/compressed_operator_predictor"),
+    )
+    return parser.parse_args()
+
+
+def make_config(args: argparse.Namespace) -> dcs.RunConfig:
+    return dcs.RunConfig(
+        task="random",
+        input_dim=args.input_dim,
+        output_dim=args.output_dim,
+        teacher_rank=4,
+        teacher_width=64,
+        teacher_hidden_layers=2,
+        normalize_targets=False,
+        widths=[args.width],
+        train_samples=args.train_samples,
+        test_samples=args.test_samples,
+        probe_samples=8,
+        steps=args.target_steps,
+        lr=args.lr,
+        optimizer="sgd",
+        init_seeds=args.init_seeds,
+        feedback_seeds=args.feedback_seeds,
+        init_seed_offset=0,
+        feedback_seed_offset=0,
+        data_seed=args.data_seed,
+        noise_std=0.0,
+        feedback_scale=args.feedback_scale,
+        capacity_q=0.01,
+        jacobian_lambda_rel=1e-3,
+        skip_jacobian=True,
+        device=args.device,
+        torch_threads=args.torch_threads,
+        outdir=str(args.outdir),
+        plot=False,
+    )
+
+
+def loss_from_residual(residual: np.ndarray, samples: int) -> float:
+    return 0.5 * float(residual @ residual) / samples
+
+
+def kernel_pair(
+    weights: list[torch.Tensor],
+    x: torch.Tensor,
+    feedback: list[torch.Tensor],
+) -> tuple[np.ndarray, np.ndarray]:
+    j_bp = pseudo_jacobian(weights, x, feedback=None).cpu().numpy()
+    j_fa = pseudo_jacobian(weights, x, feedback=feedback).cpu().numpy()
+    return j_bp @ j_bp.T, j_bp @ j_fa.T
+
+
+def sgd_step(
+    weights: list[torch.Tensor],
+    x: torch.Tensor,
+    y: torch.Tensor,
+    lr: float,
+    feedback: list[torch.Tensor] | None,
+) -> list[torch.Tensor]:
+    grads = dcs.gradients(weights, x, y, feedback)
+    return [weight - lr * grad for weight, grad in zip(weights, grads)]
+
+
+def train_to_steps(
+    weights: list[torch.Tensor],
+    x: torch.Tensor,
+    y: torch.Tensor,
+    lr: float,
+    steps: int,
+    feedback: list[torch.Tensor] | None,
+) -> list[torch.Tensor]:
+    current = dcs.clone_weights(weights)
+    for _ in range(steps):
+        current = sgd_step(current, x, y, lr, feedback)
+    return current
+
+
+def fixed_rollout(
+    kernel: np.ndarray,
+    residual: np.ndarray,
+    lr: float,
+    steps: int,
+    samples: int,
+) -> np.ndarray:
+    current = residual.copy()
+    scale = lr / samples
+    for _ in range(steps):
+        current = current - scale * (kernel @ current)
+    return current
+
+
+def compressed_rollout(
+    k_fa0: np.ndarray,
+    k_bp0: np.ndarray,
+    alpha_slope: float,
+    residual: np.ndarray,
+    lr: float,
+    steps: int,
+    samples: int,
+) -> np.ndarray:
+    current = residual.copy()
+    scale = lr / samples
+    direction = k_bp0 - k_fa0
+    for t in range(steps):
+        alpha_t = float(np.clip(alpha_slope * t, 0.0, 1.0))
+        kernel = k_fa0 + alpha_t * direction
+        current = current - scale * (kernel @ current)
+    return current
+
+
+def linear_velocity_rollout(
+    kernel0: np.ndarray,
+    kernel_s: np.ndarray,
+    early_step: int,
+    residual: np.ndarray,
+    lr: float,
+    steps: int,
+    samples: int,
+) -> np.ndarray:
+    current = residual.copy()
+    scale = lr / samples
+    velocity = (kernel_s - kernel0) / early_step
+    for t in range(steps):
+        kernel = kernel0 + t * velocity
+        current = current - scale * (kernel @ current)
+    return current
+
+
+def projection_alpha(k_fa_s: np.ndarray, k_fa0: np.ndarray, k_bp0: np.ndarray) -> float:
+    direction = k_bp0 - k_fa0
+    denom = float(np.sum(direction * direction))
+    if denom <= 0:
+        return 0.0
+    alpha = float(np.sum((k_fa_s - k_fa0) * direction) / denom)
+    return float(np.clip(alpha, 0.0, 1.0))
+
+
+def run_one(
+    config: dcs.RunConfig,
+    x: torch.Tensor,
+    y: torch.Tensor,
+    init_seed: int,
+    feedback_seed: int,
+    early_step: int,
+) -> CompressedRow:
+    initial = dcs.initialize_weights(config, config.widths[0], init_seed)
+    feedback = dcs.init_feedback(config, config.widths[0], feedback_seed)
+    with torch.no_grad():
+        r0 = (dcs.predict(initial, x) - y).reshape(-1).cpu().numpy()
+
+    k_bp0, k_fa0 = kernel_pair(initial, x, feedback)
+    fixed_bp_residual = fixed_rollout(k_bp0, r0, config.lr, config.steps, config.train_samples)
+    fixed_fa_residual = fixed_rollout(k_fa0, r0, config.lr, config.steps, config.train_samples)
+
+    fa_early = train_to_steps(initial, x, y, config.lr, early_step, feedback=feedback)
+    _kbp_unused, k_fa_s = kernel_pair(fa_early, x, feedback)
+    alpha_s = projection_alpha(k_fa_s, k_fa0, k_bp0)
+    alpha_slope = alpha_s / early_step
+    compressed_fa_residual = compressed_rollout(
+        k_fa0,
+        k_bp0,
+        alpha_slope,
+        r0,
+        config.lr,
+        config.steps,
+        config.train_samples,
+    )
+
+    bp_early = train_to_steps(initial, x, y, config.lr, early_step, feedback=None)
+    k_bp_s, _kfa_unused = kernel_pair(bp_early, x, feedback)
+    linear_bp_residual = linear_velocity_rollout(
+        k_bp0,
+        k_bp_s,
+        early_step,
+        r0,
+        config.lr,
+        config.steps,
+        config.train_samples,
+    )
+    linear_fa_residual = linear_velocity_rollout(
+        k_fa0,
+        k_fa_s,
+        early_step,
+        r0,
+        config.lr,
+        config.steps,
+        config.train_samples,
+    )
+    with torch.no_grad():
+        bp_early_residual = (dcs.predict(bp_early, x) - y).reshape(-1).cpu().numpy()
+        fa_early_residual = (dcs.predict(fa_early, x) - y).reshape(-1).cpu().numpy()
+    retangent_bp_residual = fixed_rollout(
+        k_bp_s,
+        bp_early_residual,
+        config.lr,
+        config.steps - early_step,
+        config.train_samples,
+    )
+    retangent_fa_residual = fixed_rollout(
+        k_fa_s,
+        fa_early_residual,
+        config.lr,
+        config.steps - early_step,
+        config.train_samples,
+    )
+
+    bp_target = train_to_steps(initial, x, y, config.lr, config.steps, feedback=None)
+    fa_target = train_to_steps(initial, x, y, config.lr, config.steps, feedback=feedback)
+    empirical_bp = dcs.mse(bp_target, x, y)
+    empirical_fa = dcs.mse(fa_target, x, y)
+    fixed_bp = loss_from_residual(fixed_bp_residual, config.train_samples)
+    fixed_fa = loss_from_residual(fixed_fa_residual, config.train_samples)
+    compressed_fa = loss_from_residual(compressed_fa_residual, config.train_samples)
+    linear_bp = loss_from_residual(linear_bp_residual, config.train_samples)
+    linear_fa = loss_from_residual(linear_fa_residual, config.train_samples)
+    retangent_bp = loss_from_residual(retangent_bp_residual, config.train_samples)
+    retangent_fa = loss_from_residual(retangent_fa_residual, config.train_samples)
+    empirical_gap = empirical_fa - empirical_bp
+    fixed_gap = fixed_fa - fixed_bp
+    compressed_gap = compressed_fa - fixed_bp
+    linear_gap = linear_fa - linear_bp
+    retangent_gap = retangent_fa - retangent_bp
+
+    return CompressedRow(
+        init_seed=init_seed,
+        feedback_seed=feedback_seed,
+        target_steps=config.steps,
+        early_steps=early_step,
+        empirical_bp_loss=empirical_bp,
+        empirical_fa_loss=empirical_fa,
+        empirical_gap=empirical_gap,
+        fixed_bp_loss=fixed_bp,
+        fixed_fa_loss=fixed_fa,
+        fixed_gap=fixed_gap,
+        compressed_fa_loss=compressed_fa,
+        compressed_gap=compressed_gap,
+        linear_bp_loss=linear_bp,
+        linear_fa_loss=linear_fa,
+        linear_gap=linear_gap,
+        retangent_bp_loss=retangent_bp,
+        retangent_fa_loss=retangent_fa,
+        retangent_gap=retangent_gap,
+        alpha_s=alpha_s,
+        alpha_slope=alpha_slope,
+        fixed_gap_error=fixed_gap - empirical_gap,
+        compressed_gap_error=compressed_gap - empirical_gap,
+        linear_gap_error=linear_gap - empirical_gap,
+        retangent_gap_error=retangent_gap - empirical_gap,
+    )
+
+
+def write_rows(path: Path, rows: list[CompressedRow]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(CompressedRow.__annotations__.keys()))
+        writer.writeheader()
+        for row in rows:
+            writer.writerow(asdict(row))
+
+
+def main() -> None:
+    args = parse_args()
+    if args.torch_threads > 0:
+        torch.set_num_threads(args.torch_threads)
+    config = make_config(args)
+    x_train, y_train, _x_test, _y_test, _x_probe, _teacher = dcs.make_data(config)
+
+    rows: list[CompressedRow] = []
+    total = args.init_seeds * args.feedback_seeds * len(args.early_steps)
+    count = 0
+    for init_index in range(args.init_seeds):
+        init_seed = 10_000 + init_index
+        for feedback_index in range(args.feedback_seeds):
+            feedback_seed = 100_000 + init_index * 1000 + feedback_index
+            for early_step in args.early_steps:
+                count += 1
+                print(
+                    f"[{count}/{total}] init={init_seed} feedback={feedback_seed} "
+                    f"early={early_step}",
+                    flush=True,
+                )
+                rows.append(
+                    run_one(config, x_train, y_train, init_seed, feedback_seed, early_step)
+                )
+
+    outdir = Path(args.outdir)
+    outdir.mkdir(parents=True, exist_ok=True)
+    write_rows(outdir / "compressed_operator_rows.csv", rows)
+    payload = {
+        "config": {
+            key: str(value) if isinstance(value, Path) else value
+            for key, value in vars(args).items()
+        },
+        "rows": [asdict(row) for row in rows],
+    }
+    (outdir / "summary.json").write_text(json.dumps(payload, indent=2) + "\n")
+    fixed = np.array([row.fixed_gap_error for row in rows])
+    compressed = np.array([row.compressed_gap_error for row in rows])
+    linear = np.array([row.linear_gap_error for row in rows])
+    retangent = np.array([row.retangent_gap_error for row in rows])
+    print(f"rows: {outdir / 'compressed_operator_rows.csv'}")
+    print(f"fixed MAE={np.mean(np.abs(fixed)):.6g}, bias={fixed.mean():.6g}")
+    print(
+        f"compressed MAE={np.mean(np.abs(compressed)):.6g}, "
+        f"bias={compressed.mean():.6g}"
+    )
+    print(f"linear MAE={np.mean(np.abs(linear)):.6g}, bias={linear.mean():.6g}")
+    print(
+        f"retangent MAE={np.mean(np.abs(retangent)):.6g}, "
+        f"bias={retangent.mean():.6g}"
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/plot_compressed_operator_results.py b/scripts/plot_compressed_operator_results.py
new file mode 100644
index 0000000..c23fa14
--- /dev/null
+++ b/scripts/plot_compressed_operator_results.py
@@ -0,0 +1,311 @@
+#!/usr/bin/env python3
+"""Plot early-kernel finite-time FA/BP gap predictions."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+from dataclasses import dataclass
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+
+@dataclass(frozen=True)
+class MetricsRow:
+    predictor: str
+    rows: int
+    mae: float
+    bias: float
+    corr: float
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--glob",
+        default="outputs/compressed_operator_s20_256traj_T50_width64_N*/compressed_operator_rows.csv",
+    )
+    parser.add_argument("--input-dim", type=int, default=16)
+    parser.add_argument("--output-dim", type=int, default=4)
+    parser.add_argument("--width", type=int, default=64)
+    parser.add_argument(
+        "--outdir",
+        type=Path,
+        default=Path("outputs/compressed_operator_s20_256traj_T50_width64_plots"),
+    )
+    return parser.parse_args()
+
+
+def read_rows(pattern: str, input_dim: int, output_dim: int, width: int) -> pd.DataFrame:
+    frames = []
+    for path in sorted(Path().glob(pattern)):
+        dirname = path.parent.name
+        try:
+            train_samples = int(dirname.rsplit("N", 1)[1])
+        except (IndexError, ValueError) as exc:
+            raise ValueError(f"Cannot infer train sample count from {dirname}") from exc
+        frame = pd.read_csv(path)
+        frame["train_samples"] = train_samples
+        frames.append(frame)
+    if not frames:
+        raise FileNotFoundError(f"No rows found for pattern: {pattern}")
+    df = pd.concat(frames, ignore_index=True)
+    parameter_count = input_dim * width + width * width + width * output_dim
+    constraint_rank = max(width * width - 1, 0) + max(width * output_dim - 1, 0)
+    df["hard_fa_capacity_margin"] = (
+        parameter_count - constraint_rank - output_dim * df["train_samples"]
+    )
+    df["trajectory_id"] = (
+        df["init_seed"].astype(str) + ":" + df["feedback_seed"].astype(str)
+    )
+    return df.sort_values(["hard_fa_capacity_margin", "trajectory_id"]).reset_index(drop=True)
+
+
+def predictor_metrics(df: pd.DataFrame) -> list[MetricsRow]:
+    rows: list[MetricsRow] = []
+    for predictor, column in [
+        ("fixed K(0)", "fixed_gap"),
+        ("FA compressed", "compressed_gap"),
+        ("linear velocity", "linear_gap"),
+        ("early re-tangent", "retangent_gap"),
+    ]:
+        error = df[column].to_numpy() - df["empirical_gap"].to_numpy()
+        corr = float(np.corrcoef(df[column], df["empirical_gap"])[0, 1])
+        rows.append(
+            MetricsRow(
+                predictor=predictor,
+                rows=len(df),
+                mae=float(np.mean(np.abs(error))),
+                bias=float(np.mean(error)),
+                corr=corr,
+            )
+        )
+    return rows
+
+
+def write_metrics(path: Path, rows: list[MetricsRow]) -> None:
+    with path.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(MetricsRow.__annotations__.keys()))
+        writer.writeheader()
+        for row in rows:
+            writer.writerow(row.__dict__)
+
+
+def summarize_by_margin(df: pd.DataFrame, column: str) -> pd.DataFrame:
+    records = []
+    for margin, group in df.groupby("hard_fa_capacity_margin"):
+        values = group[column].to_numpy()
+        records.append(
+            {
+                "margin": margin,
+                "mean": float(values.mean()),
+                "q05": float(np.quantile(values, 0.05)),
+                "q25": float(np.quantile(values, 0.25)),
+                "q75": float(np.quantile(values, 0.75)),
+                "q95": float(np.quantile(values, 0.95)),
+            }
+        )
+    return pd.DataFrame(records).sort_values("margin")
+
+
+def plot_theory_band_with_empirical(df: pd.DataFrame, outdir: Path) -> Path:
+    theory = summarize_by_margin(df, "linear_gap")
+    empirical = summarize_by_margin(df, "empirical_gap")
+    fig, ax = plt.subplots(figsize=(11.5, 6.4), dpi=180)
+
+    for _tid, group in df.groupby("trajectory_id"):
+        group = group.sort_values("hard_fa_capacity_margin")
+        ax.plot(
+            group["hard_fa_capacity_margin"],
+            group["empirical_gap"],
+            color="#c65f16",
+            alpha=0.16,
+            linewidth=1.0,
+        )
+
+    ax.fill_between(
+        theory["margin"],
+        theory["q05"],
+        theory["q95"],
+        color="#1f5a9d",
+        alpha=0.16,
+        linewidth=0,
+        label="velocity theory 5-95%",
+    )
+    ax.fill_between(
+        theory["margin"],
+        theory["q25"],
+        theory["q75"],
+        color="#1f5a9d",
+        alpha=0.28,
+        linewidth=0,
+        label="velocity theory 25-75%",
+    )
+    ax.plot(
+        theory["margin"],
+        theory["mean"],
+        color="#174f91",
+        linewidth=2.3,
+        marker="o",
+        markersize=4.5,
+        label="velocity theory mean",
+    )
+    ax.plot(
+        empirical["margin"],
+        empirical["mean"],
+        color="#a84708",
+        linewidth=2.1,
+        marker="o",
+        markersize=4.2,
+        label="empirical mean",
+    )
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
+    ax.set_title("Early-kernel velocity prediction vs empirical FA/BP gap")
+    ax.set_xlabel("hard FA capacity margin")
+    ax.set_ylabel("FA train MSE - BP train MSE")
+    ax.legend(loc="upper left", frameon=True)
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    outpath = outdir / "T50_s20_velocity_theory_band_empirical_trajectories.png"
+    fig.savefig(outpath)
+    plt.close(fig)
+    return outpath
+
+
+def plot_empirical_only(df: pd.DataFrame, outdir: Path) -> Path:
+    fig, ax = plt.subplots(figsize=(11.5, 6.4), dpi=180)
+    for _tid, group in df.groupby("trajectory_id"):
+        group = group.sort_values("hard_fa_capacity_margin")
+        ax.plot(
+            group["hard_fa_capacity_margin"],
+            group["empirical_gap"],
+            color="#c65f16",
+            alpha=0.20,
+            linewidth=1.05,
+        )
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
+    ax.set_title("Empirical FA/BP gap trajectories")
+    ax.set_xlabel("hard FA capacity margin")
+    ax.set_ylabel("FA train MSE - BP train MSE")
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    outpath = outdir / "T50_s20_empirical_trajectories_only.png"
+    fig.savefig(outpath)
+    plt.close(fig)
+    return outpath
+
+
+def plot_prediction_scatter(df: pd.DataFrame, outdir: Path) -> Path:
+    fig, ax = plt.subplots(figsize=(7.0, 6.2), dpi=180)
+    empirical = df["empirical_gap"].to_numpy()
+    predictors = [
+        ("fixed K(0)", "fixed_gap", "#7f7f7f", 0.38),
+        ("linear velocity", "linear_gap", "#1f5a9d", 0.58),
+        ("early re-tangent", "retangent_gap", "#2b8a3e", 0.45),
+    ]
+    max_value = float(
+        max(
+            empirical.max(),
+            *(df[column].max() for _label, column, _color, _alpha in predictors),
+        )
+    )
+    min_value = float(
+        min(
+            empirical.min(),
+            *(df[column].min() for _label, column, _color, _alpha in predictors),
+        )
+    )
+    for label, column, color, alpha in predictors:
+        ax.scatter(
+            df[column],
+            empirical,
+            s=19,
+            alpha=alpha,
+            color=color,
+            label=label,
+            edgecolors="none",
+        )
+    ax.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1.0)
+    ax.set_title("Predicted vs empirical FA/BP train-gap")
+    ax.set_xlabel("predicted gap")
+    ax.set_ylabel("empirical gap")
+    ax.legend(loc="upper left", frameon=True)
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    outpath = outdir / "T50_s20_prediction_scatter.png"
+    fig.savefig(outpath)
+    plt.close(fig)
+    return outpath
+
+
+def plot_error_by_margin(df: pd.DataFrame, outdir: Path) -> Path:
+    fig, ax = plt.subplots(figsize=(10.5, 5.8), dpi=180)
+    for label, column, color in [
+        ("fixed K(0)", "fixed_gap", "#7f7f7f"),
+        ("linear velocity", "linear_gap", "#1f5a9d"),
+        ("early re-tangent", "retangent_gap", "#2b8a3e"),
+    ]:
+        records = []
+        for margin, group in df.groupby("hard_fa_capacity_margin"):
+            error = group[column].to_numpy() - group["empirical_gap"].to_numpy()
+            records.append(
+                {
+                    "margin": margin,
+                    "mean": float(error.mean()),
+                    "q25": float(np.quantile(error, 0.25)),
+                    "q75": float(np.quantile(error, 0.75)),
+                }
+            )
+        summary = pd.DataFrame(records).sort_values("margin")
+        ax.plot(summary["margin"], summary["mean"], color=color, marker="o", label=label)
+        ax.fill_between(
+            summary["margin"],
+            summary["q25"],
+            summary["q75"],
+            color=color,
+            alpha=0.15,
+            linewidth=0,
+        )
+    ax.axhline(0, color="black", linewidth=1.0)
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
+    ax.set_title("Prediction error by capacity margin")
+    ax.set_xlabel("hard FA capacity margin")
+    ax.set_ylabel("predicted gap - empirical gap")
+    ax.legend(loc="upper left", frameon=True)
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    outpath = outdir / "T50_s20_prediction_error_by_margin.png"
+    fig.savefig(outpath)
+    plt.close(fig)
+    return outpath
+
+
+def main() -> None:
+    args = parse_args()
+    args.outdir.mkdir(parents=True, exist_ok=True)
+    df = read_rows(args.glob, args.input_dim, args.output_dim, args.width)
+    df.to_csv(args.outdir / "combined_rows.csv", index=False)
+    metrics = predictor_metrics(df)
+    write_metrics(args.outdir / "predictor_metrics.csv", metrics)
+    paths = [
+        plot_theory_band_with_empirical(df, args.outdir),
+        plot_empirical_only(df, args.outdir),
+        plot_prediction_scatter(df, args.outdir),
+        plot_error_by_margin(df, args.outdir),
+    ]
+    print(f"combined rows: {args.outdir / 'combined_rows.csv'}")
+    print(f"metrics: {args.outdir / 'predictor_metrics.csv'}")
+    for row in metrics:
+        print(
+            f"{row.predictor}: rows={row.rows}, MAE={row.mae:.6g}, "
+            f"bias={row.bias:.6g}, corr={row.corr:.6g}"
+        )
+    for path in paths:
+        print(f"plot: {path}")
+
+
+if __name__ == "__main__":
+    main()