Add time-varying kernel diagnostic

2 files changed, 389 insertions, 0 deletions

diff --git a/notes/13_time_varying_kernel_diagnostic.md b/notes/13_time_varying_kernel_diagnostic.md
new file mode 100644
index 0000000..19cef0a
--- /dev/null
+++ b/notes/13_time_varying_kernel_diagnostic.md
@@ -0,0 +1,98 @@
+# Time-Varying Kernel Diagnostic
+
+We tested the finite-\(T\) proposal directly.
+
+## Setup
+
+- task: random-label regression;
+- architecture: \(16\to64\to64\to4\);
+- \(N=128\);
+- SGD, learning rate \(10^{-3}\);
+- horizon \(T=50\);
+- 2 initialization seeds;
+- 4 feedback seeds per initialization;
+- 8 FA trajectories total.
+
+For each trajectory, compute three predictions:
+
+1. empirical BP/FA train losses;
+2. fixed-kernel prediction using \(K(0)\);
+3. measured time-varying product using \(K(t)\) at every step.
+
+The time-varying residual recursion is
+
+\[
+r_{t+1}^{\mathrm{tv}}
+=
+\left(I-\frac{\eta}{N}K_t\right)r_t^{\mathrm{tv}},
+\]
+
+with
+
+\[
+K_t^{\mathrm{BP}}=J_tJ_t^\top,
+\qquad
+K_t^{\mathrm{FA}}=J_t\tilde J_t^\top.
+\]
+
+No fitted scale or offset is used.
+
+## Result
+
+Output:
+
+`outputs/finite_time_kernel_probe_T50_N128_8runs`
+
+Figures:
+
+- `outputs/finite_time_kernel_probe_T50_N128_8runs/T50_fixed_vs_timevarying_gap.png`
+- `outputs/finite_time_kernel_probe_T50_N128_8runs/T50_prediction_scatter.png`
+
+Gap prediction:
+
+| predictor | mean gap error | gap MAE | max abs gap error |
+|---|---:|---:|---:|
+| fixed \(K(0)\) | 0.027683 | 0.027683 | 0.037170 |
+| time-varying \(K(t)\) | 0.000273 | 0.000273 | 0.000390 |
+
+BP loss prediction:
+
+| predictor | mean BP error | BP MAE |
+|---|---:|---:|
+| fixed \(K(0)\) | -0.025404 | 0.025404 |
+| time-varying \(K(t)\) | -0.000237 | 0.000237 |
+
+FA loss prediction:
+
+| predictor | mean FA error | FA MAE |
+|---|---:|---:|
+| fixed \(K(0)\) | 0.002279 | 0.004991 |
+| time-varying \(K(t)\) | 0.000036 | 0.000060 |
+
+## Interpretation
+
+At \(T=50\), the fixed-kernel gap error is almost entirely removed by using the
+measured time-varying kernel product.
+
+Therefore, for this regime, the finite-time mismatch is dominated by kernel
+drift:
+
+\[
+K_t-K_0.
+\]
+
+The second-order output Taylor residual is small at this step size/horizon,
+because the measured \(K(t)\) product already matches empirical losses to
+roughly \(10^{-4}\) to \(10^{-3}\) absolute error.
+
+This validates the finite-time extension direction:
+
+\[
+\text{local theory: }K(0)
+\quad\rightarrow\quad
+\text{finite-time theory: }K(t).
+\]
+
+The next theoretical task is not to fit a scalar correction. It is to model the
+evolution of \(K_t^{\mathrm{BP}}\) and \(K_t^{\mathrm{FA}}\), especially the FA
+alignment-driven improvement of \(K_t^{\mathrm{FA}}\).
diff --git a/scripts/finite_time_kernel_diagnostic.py b/scripts/finite_time_kernel_diagnostic.py
new file mode 100644
index 0000000..4827e69
--- /dev/null
+++ b/scripts/finite_time_kernel_diagnostic.py
@@ -0,0 +1,291 @@
+#!/usr/bin/env python3
+"""Finite-time BP/FA diagnostic with measured time-varying tangent kernels."""
+
+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 DiagnosticConfig:
+    input_dim: int
+    output_dim: int
+    width: int
+    train_samples: int
+    test_samples: int
+    steps: int
+    lr: float
+    init_seeds: int
+    feedback_seeds: int
+    data_seed: int
+    feedback_scale: str
+    device: str
+    torch_threads: int
+    outdir: str
+
+
+@dataclass(frozen=True)
+class DiagnosticRow:
+    init_seed: int
+    feedback_seed: int
+    steps: int
+    empirical_bp_loss: float
+    empirical_fa_loss: float
+    empirical_gap: float
+    fixed_bp_loss: float
+    fixed_fa_loss: float
+    fixed_gap: float
+    tv_bp_loss: float
+    tv_fa_loss: float
+    tv_gap: float
+    fixed_gap_error: float
+    tv_gap_error: float
+    fixed_bp_error: float
+    fixed_fa_error: float
+    tv_bp_error: float
+    tv_fa_error: float
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Measure finite-time fixedK vs time-varyingK predictions."
+    )
+    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("--steps", type=int, default=50)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--init-seeds", type=int, default=1)
+    parser.add_argument("--feedback-seeds", type=int, default=2)
+    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/finite_time_kernel_diagnostic"),
+    )
+    return parser.parse_args()
+
+
+def make_run_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.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 sgd_step(
+    weights: list[torch.Tensor],
+    x_train: torch.Tensor,
+    y_train: torch.Tensor,
+    lr: float,
+    feedback: list[torch.Tensor] | None,
+) -> list[torch.Tensor]:
+    grads = dcs.gradients(weights, x_train, y_train, feedback)
+    return [weight - lr * grad for weight, grad in zip(weights, grads)]
+
+
+def kernel_for(
+    weights: list[torch.Tensor],
+    x_train: torch.Tensor,
+    feedback: list[torch.Tensor] | None,
+) -> np.ndarray:
+    j_bp = pseudo_jacobian(weights, x_train, feedback=None).cpu().numpy()
+    if feedback is None:
+        return j_bp @ j_bp.T
+    j_fa = pseudo_jacobian(weights, x_train, feedback=feedback).cpu().numpy()
+    return j_bp @ j_fa.T
+
+
+def run_one(
+    config: dcs.RunConfig,
+    x_train: torch.Tensor,
+    y_train: torch.Tensor,
+    init_seed: int,
+    feedback_seed: int,
+) -> DiagnosticRow:
+    initial_weights = 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_weights, x_train) - y_train).reshape(-1).cpu().numpy()
+
+    k_bp0 = kernel_for(initial_weights, x_train, feedback=None)
+    k_fa0 = kernel_for(initial_weights, x_train, feedback=feedback)
+
+    fixed_bp_residual = r0.copy()
+    fixed_fa_residual = r0.copy()
+    tv_bp_residual = r0.copy()
+    tv_fa_residual = r0.copy()
+    bp_weights = dcs.clone_weights(initial_weights)
+    fa_weights = dcs.clone_weights(initial_weights)
+    scale = config.lr / config.train_samples
+
+    for _step in range(config.steps):
+        fixed_bp_residual = fixed_bp_residual - scale * (k_bp0 @ fixed_bp_residual)
+        fixed_fa_residual = fixed_fa_residual - scale * (k_fa0 @ fixed_fa_residual)
+
+        k_bp_t = kernel_for(bp_weights, x_train, feedback=None)
+        k_fa_t = kernel_for(fa_weights, x_train, feedback=feedback)
+        tv_bp_residual = tv_bp_residual - scale * (k_bp_t @ tv_bp_residual)
+        tv_fa_residual = tv_fa_residual - scale * (k_fa_t @ tv_fa_residual)
+
+        bp_weights = sgd_step(bp_weights, x_train, y_train, config.lr, feedback=None)
+        fa_weights = sgd_step(fa_weights, x_train, y_train, config.lr, feedback=feedback)
+
+    empirical_bp_loss = dcs.mse(bp_weights, x_train, y_train)
+    empirical_fa_loss = dcs.mse(fa_weights, x_train, y_train)
+    fixed_bp_loss = loss_from_residual(fixed_bp_residual, config.train_samples)
+    fixed_fa_loss = loss_from_residual(fixed_fa_residual, config.train_samples)
+    tv_bp_loss = loss_from_residual(tv_bp_residual, config.train_samples)
+    tv_fa_loss = loss_from_residual(tv_fa_residual, config.train_samples)
+
+    empirical_gap = empirical_fa_loss - empirical_bp_loss
+    fixed_gap = fixed_fa_loss - fixed_bp_loss
+    tv_gap = tv_fa_loss - tv_bp_loss
+
+    return DiagnosticRow(
+        init_seed=init_seed,
+        feedback_seed=feedback_seed,
+        steps=config.steps,
+        empirical_bp_loss=empirical_bp_loss,
+        empirical_fa_loss=empirical_fa_loss,
+        empirical_gap=empirical_gap,
+        fixed_bp_loss=fixed_bp_loss,
+        fixed_fa_loss=fixed_fa_loss,
+        fixed_gap=fixed_gap,
+        tv_bp_loss=tv_bp_loss,
+        tv_fa_loss=tv_fa_loss,
+        tv_gap=tv_gap,
+        fixed_gap_error=fixed_gap - empirical_gap,
+        tv_gap_error=tv_gap - empirical_gap,
+        fixed_bp_error=fixed_bp_loss - empirical_bp_loss,
+        fixed_fa_error=fixed_fa_loss - empirical_fa_loss,
+        tv_bp_error=tv_bp_loss - empirical_bp_loss,
+        tv_fa_error=tv_fa_loss - empirical_fa_loss,
+    )
+
+
+def write_rows(path: Path, rows: list[DiagnosticRow]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(DiagnosticRow.__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_run_config(args)
+    x_train, y_train, _x_test, _y_test, _x_probe, _teacher = dcs.make_data(config)
+
+    rows: list[DiagnosticRow] = []
+    total = args.init_seeds * args.feedback_seeds
+    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
+            count += 1
+            print(
+                f"[{count}/{total}] init={init_seed} feedback={feedback_seed}",
+                flush=True,
+            )
+            rows.append(run_one(config, x_train, y_train, init_seed, feedback_seed))
+
+    outdir = Path(args.outdir)
+    outdir.mkdir(parents=True, exist_ok=True)
+    write_rows(outdir / "finite_time_kernel_rows.csv", rows)
+    payload = {
+        "config": asdict(
+            DiagnosticConfig(
+                input_dim=args.input_dim,
+                output_dim=args.output_dim,
+                width=args.width,
+                train_samples=args.train_samples,
+                test_samples=args.test_samples,
+                steps=args.steps,
+                lr=args.lr,
+                init_seeds=args.init_seeds,
+                feedback_seeds=args.feedback_seeds,
+                data_seed=args.data_seed,
+                feedback_scale=args.feedback_scale,
+                device=args.device,
+                torch_threads=args.torch_threads,
+                outdir=str(outdir),
+            )
+        ),
+        "rows": [asdict(row) for row in rows],
+    }
+    (outdir / "summary.json").write_text(json.dumps(payload, indent=2) + "\n")
+
+    fixed_errors = np.array([row.fixed_gap_error for row in rows])
+    tv_errors = np.array([row.tv_gap_error for row in rows])
+    print(f"rows: {outdir / 'finite_time_kernel_rows.csv'}")
+    print(
+        "fixed gap error mean="
+        f"{fixed_errors.mean():.6g}, mae={np.mean(np.abs(fixed_errors)):.6g}"
+    )
+    print(
+        "time-varying gap error mean="
+        f"{tv_errors.mean():.6g}, mae={np.mean(np.abs(tv_errors)):.6g}"
+    )
+
+
+if __name__ == "__main__":
+    main()