Probe early kernel predictors

2 files changed, 466 insertions, 0 deletions

diff --git a/notes/14_early_predictor_probe.md b/notes/14_early_predictor_probe.md
new file mode 100644
index 0000000..5cfe058
--- /dev/null
+++ b/notes/14_early_predictor_probe.md
@@ -0,0 +1,175 @@
+# Early Predictor Probe
+
+Measured time-varying \(K(t)\) is an oracle diagnostic. To get a predictive
+finite-\(T\) theory, we need quantities computable from initialization or the
+first few steps.
+
+## Probe Setup
+
+- architecture \(16\to64\to64\to4\);
+- \(N=128\);
+- target horizon \(T=50\);
+- SGD, learning rate \(10^{-3}\);
+- 4 init seeds and 8 feedback seeds;
+- early probes \(s\in\{1,2,5\}\).
+
+Candidate predictors:
+
+- fixed \(K(0)\) gap prediction;
+- FA kernel Frobenius drift:
+  \[
+  \|K_s^{FA}-K_0^{FA}\|_F/\|K_0^{FA}\|_F;
+  \]
+- BP kernel Frobenius drift:
+  \[
+  \|K_s^{BP}-K_0^{BP}\|_F/\|K_0^{BP}\|_F;
+  \]
+- FA/BP kernel overlap:
+  \[
+  \rho_s
+  =
+  \frac{\langle K_s^{FA},K_s^{BP}\rangle_F}
+  {\|K_s^{FA}\|_F\|K_s^{BP}\|_F};
+  \]
+- residual-direction gains:
+  \[
+  \lambda_s^{FA}
+  =
+  \frac{r_s^\top K_s^{FA}r_s}{\|r_s\|^2},
+  \qquad
+  \lambda_s^{BP}
+  =
+  \frac{r_s^\top K_s^{BP}r_s}{\|r_s\|^2}.
+  \]
+
+## Result
+
+Output:
+
+`outputs/early_kernel_predictors_T50_N128_32runs/early_kernel_predictors.csv`
+
+The strongest single predictor of the \(T=50\) empirical gap remains the fixed
+\(K(0)\) predicted gap:
+
+\[
+\operatorname{corr}
+\left(
+\Delta L_{50}^{fixedK},
+\Delta L_{50}^{emp}
+\right)
+\approx
+0.951.
+\]
+
+The fixed prediction is biased high:
+
+\[
+\operatorname{MAE}
+\left(
+\Delta L_{50}^{fixedK},
+\Delta L_{50}^{emp}
+\right)
+=
+0.02063.
+\]
+
+If we fit only a linear calibration of fixed gap on this single setting, the
+in-sample MAE drops to:
+
+\[
+0.00517.
+\]
+
+Adding an early drift predictor can reduce the in-sample residual further. The
+best early predictors were:
+
+- \(s=1\), fixed gap + FA directional drift:
+  \[
+  \operatorname{MAE}\approx0.00296;
+  \]
+- \(s=5\), fixed gap + BP Frobenius drift:
+  \[
+  \operatorname{MAE}\approx0.00307.
+  \]
+
+## Interpretation
+
+This probe is useful but not yet a theorem-ready predictor.
+
+The good news:
+
+1. fixed \(K(0)\) already ranks trajectories well at \(T=50\);
+2. early drift features explain part of the finite-\(T\) residual;
+3. this supports the finite-\(T\) picture:
+   \[
+   \Delta L_T
+   =
+   F(K_0)
+   +
+   \text{early drift correction}
+   +
+   \text{higher-order residual}.
+   \]
+
+The caution:
+
+The linear combinations above are in-sample on one architecture/task setting.
+They are not yet valid paper predictors. The next step must test whether a
+fixed formula learned or derived from one setting generalizes across \(N\),
+width, and feedback seeds.
+
+## Next Direction
+
+The most promising non-oracle predictive object is not raw Frobenius drift. It
+is the residual-direction kernel action:
+
+\[
+\lambda_t^{rule}
+=
+\frac{r_t^\top K_t^{rule}r_t}{\|r_t\|^2}.
+\]
+
+Reason: loss decrease is controlled by
+
+\[
+L_{t+1}-L_t
+\approx
+-
+\frac{\eta}{N}
+r_t^\top K_t r_t.
+\]
+
+So a finite-\(T\) predictor should approximate the cumulative directional gain:
+
+\[
+G_T^{rule}
+=
+\sum_{t<T}
+\frac{\eta}{N}
+\lambda_t^{rule}.
+\]
+
+For a predictive version, estimate the early slope:
+
+\[
+\lambda_t
+\approx
+\lambda_0+t\dot\lambda_0,
+\qquad
+\dot\lambda_0
+\approx
+\frac{\lambda_s-\lambda_0}{s},
+\]
+
+then predict
+
+\[
+G_T
+\approx
+\frac{\eta}{N}
+\left(
+T\lambda_0+\frac{T(T-1)}{2}\dot\lambda_0
+\right).
+\]
+
+This is more directly tied to loss dynamics than global kernel overlap.
diff --git a/scripts/early_kernel_predictors.py b/scripts/early_kernel_predictors.py
new file mode 100644
index 0000000..e791a9d
--- /dev/null
+++ b/scripts/early_kernel_predictors.py
@@ -0,0 +1,291 @@
+#!/usr/bin/env python3
+"""Early kernel-drift predictors for finite-time 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 PredictorRow:
+    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
+    kfa_fro_drift: float
+    kbp_fro_drift: float
+    fa_bp_overlap_0: float
+    fa_bp_overlap_s: float
+    fa_bp_overlap_slope: float
+    fa_directional_gain_s: float
+    bp_directional_gain_s: float
+    fa_directional_drift_s: float
+    bp_directional_drift_s: float
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Compute early kernel predictors.")
+    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=4)
+    parser.add_argument("--feedback-seeds", type=int, default=8)
+    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/early_kernel_predictors"),
+    )
+    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 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 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 fixed_prediction(
+    k_bp0: np.ndarray,
+    k_fa0: np.ndarray,
+    residual: np.ndarray,
+    lr: float,
+    steps: int,
+    samples: int,
+) -> tuple[float, float]:
+    scale = lr / samples
+    rb = residual.copy()
+    rf = residual.copy()
+    for _ in range(steps):
+        rb = rb - scale * (k_bp0 @ rb)
+        rf = rf - scale * (k_fa0 @ rf)
+    return 0.5 * float(rb @ rb) / samples, 0.5 * float(rf @ rf) / samples
+
+
+def overlap(a: np.ndarray, b: np.ndarray) -> float:
+    denom = np.linalg.norm(a, "fro") * np.linalg.norm(b, "fro")
+    if denom == 0:
+        return 0.0
+    return float(np.sum(a * b) / denom)
+
+
+def directional_gain(kernel: np.ndarray, residual: np.ndarray) -> float:
+    denom = float(residual @ residual)
+    if denom == 0:
+        return 0.0
+    return float(residual @ (kernel @ residual) / denom)
+
+
+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 run_one(
+    config: dcs.RunConfig,
+    x: torch.Tensor,
+    y: torch.Tensor,
+    init_seed: int,
+    feedback_seed: int,
+    early_step: int,
+) -> PredictorRow:
+    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()
+
+    kbp0, kfa0 = kernel_pair(initial, x, feedback)
+    fixed_bp, fixed_fa = fixed_prediction(
+        kbp0, kfa0, r0, config.lr, config.steps, 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)
+
+    bp_early = train_to_steps(initial, x, y, config.lr, early_step, feedback=None)
+    fa_early = train_to_steps(initial, x, y, config.lr, early_step, feedback=feedback)
+    kbp_s, _kfa_unused = kernel_pair(bp_early, x, feedback)
+    _kbp_unused, kfa_s = kernel_pair(fa_early, x, feedback)
+
+    with torch.no_grad():
+        rb_s = (dcs.predict(bp_early, x) - y).reshape(-1).cpu().numpy()
+        rf_s = (dcs.predict(fa_early, x) - y).reshape(-1).cpu().numpy()
+
+    kfa_fro_drift = float(np.linalg.norm(kfa_s - kfa0, "fro") / np.linalg.norm(kfa0, "fro"))
+    kbp_fro_drift = float(np.linalg.norm(kbp_s - kbp0, "fro") / np.linalg.norm(kbp0, "fro"))
+    fa_bp_overlap_0 = overlap(kfa0, kbp0)
+    fa_bp_overlap_s = overlap(kfa_s, kbp_s)
+    fa_directional_gain_0 = directional_gain(kfa0, r0)
+    fa_directional_gain_s = directional_gain(kfa_s, rf_s)
+    bp_directional_gain_0 = directional_gain(kbp0, r0)
+    bp_directional_gain_s = directional_gain(kbp_s, rb_s)
+
+    return PredictorRow(
+        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_fa - empirical_bp,
+        fixed_bp_loss=fixed_bp,
+        fixed_fa_loss=fixed_fa,
+        fixed_gap=fixed_fa - fixed_bp,
+        kfa_fro_drift=kfa_fro_drift,
+        kbp_fro_drift=kbp_fro_drift,
+        fa_bp_overlap_0=fa_bp_overlap_0,
+        fa_bp_overlap_s=fa_bp_overlap_s,
+        fa_bp_overlap_slope=(fa_bp_overlap_s - fa_bp_overlap_0) / early_step,
+        fa_directional_gain_s=fa_directional_gain_s,
+        bp_directional_gain_s=bp_directional_gain_s,
+        fa_directional_drift_s=(fa_directional_gain_s - fa_directional_gain_0) / max(
+            abs(fa_directional_gain_0), 1e-12
+        ),
+        bp_directional_drift_s=(bp_directional_gain_s - bp_directional_gain_0) / max(
+            abs(bp_directional_gain_0), 1e-12
+        ),
+    )
+
+
+def write_rows(path: Path, rows: list[PredictorRow]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(PredictorRow.__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[PredictorRow] = []
+    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 / "early_kernel_predictors.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")
+    print(f"rows: {outdir / 'early_kernel_predictors.csv'}")
+
+
+if __name__ == "__main__":
+    main()