Stress test early-kernel operator predictor

4 files changed, 472 insertions, 5 deletions

diff --git a/notes/16_operator_stress_grid.md b/notes/16_operator_stress_grid.md
new file mode 100644
index 0000000..7d7bf25
--- /dev/null
+++ b/notes/16_operator_stress_grid.md
@@ -0,0 +1,104 @@
+# Operator Predictor Stress Grid
+
+This note tests whether the early-kernel operator predictor from
+`notes/15_early_kernel_operator_predictor.md` survives changes in network
+depth, width, and training horizon.
+
+## Predictor
+
+For each rule, use the early operator velocity
+
+\[
+\hat K_t^{rule}
+=
+K_0^{rule}
++
+t\frac{K_s^{rule}-K_0^{rule}}{s},
+\qquad 0\le t<T,
+\]
+
+and roll out
+
+\[
+\hat r_{t+1}^{rule}
+=
+\left(I-\frac{\eta}{N}\hat K_t^{rule}\right)
+\hat r_t^{rule}.
+\]
+
+This uses \((K_0,K_s)\) only. There is no fitted scale, offset, or final-gap
+calibration.
+
+## Grid
+
+All experiments use random-label regression, input dimension 16, output
+dimension 4, full-batch SGD, and learning rate \(10^{-3}\).
+
+Stress settings:
+
+- depth: \(1,2,3\) hidden layers at width 64, \(T=50,s=20\);
+- width: \(32,64,96\) at depth 2, \(T=50,s=20\);
+- horizon: \(T=25,s=10\), \(T=50,s=20\), \(T=100,s=40\) at depth 2, width 64.
+
+Each new stress setting uses 5 train-sample values around the hard FA capacity
+margin and 8 trajectories per value. The baseline \(d=2,w=64,T=50,s=20\) uses
+the existing 256 trajectory points.
+
+Output directory:
+
+`outputs/operator_stress_grid_summary`
+
+Key files:
+
+- `stress_grid_rows.csv`
+- `stress_grid_metrics.csv`
+- `stress_grid_gap_mae_by_setting.png`
+- `stress_grid_linear_bias_by_margin.png`
+- `stress_grid_linear_prediction_scatter.png`
+
+## Results
+
+| setting | rows | fixed MAE | linear MAE | retangent MAE | linear bias | corr |
+|---|---:|---:|---:|---:|---:|---:|
+| \(d=1,w=64,T=50,s=20\) | 40 | 0.003862 | 0.000233 | 0.001162 | -0.000233 | 0.999891 |
+| \(d=2,w=32,T=50,s=20\) | 40 | 0.047110 | 0.005652 | 0.013333 | -0.005375 | 0.994077 |
+| \(d=2,w=64,T=25,s=10\) | 40 | 0.020808 | 0.002722 | 0.005084 | -0.002668 | 0.999074 |
+| \(d=2,w=64,T=50,s=20\) | 256 | 0.018487 | 0.001893 | 0.004868 | -0.001491 | 0.998848 |
+| \(d=2,w=64,T=100,s=40\) | 40 | 0.034378 | 0.002647 | 0.008100 | -0.002169 | 0.998653 |
+| \(d=2,w=96,T=50,s=20\) | 40 | 0.016912 | 0.001794 | 0.004004 | -0.001551 | 0.998634 |
+| \(d=3,w=64,T=50,s=20\) | 40 | 0.051670 | 0.005049 | 0.012450 | -0.004916 | 0.998546 |
+
+## Interpretation
+
+The predictor holds across all tested depth, width, and horizon changes.
+
+The hard cases are the narrower and deeper MLPs:
+
+- width 32 increases finite-width/kernel-drift error;
+- depth 3 increases the magnitude of kernel drift and FA/BP mismatch.
+
+Even in these cases, the early operator-velocity predictor reduces the fixed
+\(K(0)\) gap error by roughly one order of magnitude. The remaining systematic
+error is a mild negative bias, meaning the linear velocity model slightly
+over-extrapolates kernel drift.
+
+The early re-tangent predictor is consistently conservative and biased high.
+Together, early re-tangent and linear velocity form a useful no-fit bracket:
+
+\[
+\Delta L_T^{linear}
+\lesssim
+\Delta L_T^{empirical}
+\lesssim
+\Delta L_T^{retangent}
+\]
+
+in most tested settings.
+
+## Current Boundary
+
+This stress grid supports the finite-\(T\) predictor conditional on early
+operator observations \((K_0,K_s)\). It does not yet prove an architecture-only
+closed-form predictor for long training. For that stronger claim, the next
+theoretical object would need to predict \(K_s-K_0\) from initialization
+statistics rather than measuring \(K_s\).
diff --git a/scripts/compressed_operator_predictor.py b/scripts/compressed_operator_predictor.py
index ce52c7a..3f3f73b 100644
--- a/scripts/compressed_operator_predictor.py
+++ b/scripts/compressed_operator_predictor.py
@@ -23,6 +23,9 @@ from fa_tangent_kernel_capacity import pseudo_jacobian  # noqa: E402
 
 @dataclass(frozen=True)
 class CompressedRow:
+    hidden_layers: int
+    width: int
+    train_samples: int
     init_seed: int
     feedback_seed: int
     target_steps: int
@@ -54,6 +57,7 @@ def parse_args() -> argparse.Namespace:
     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("--hidden-layers", type=int, default=2)
     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)
@@ -114,6 +118,57 @@ def loss_from_residual(residual: np.ndarray, samples: int) -> float:
     return 0.5 * float(residual @ residual) / samples
 
 
+def initialize_weights(
+    config: dcs.RunConfig,
+    width: int,
+    hidden_layers: int,
+    seed: int,
+) -> list[torch.Tensor]:
+    generator = torch.Generator(device=config.device)
+    generator.manual_seed(seed)
+    dims = [config.input_dim, *([width] * hidden_layers), config.output_dim]
+    weights: list[torch.Tensor] = []
+    for layer, (fan_in, fan_out) in enumerate(zip(dims[:-1], dims[1:])):
+        scale = np.sqrt(2.0 / fan_in) if layer < len(dims) - 2 else 1.0 / np.sqrt(fan_in)
+        weights.append(
+            torch.randn(
+                fan_out,
+                fan_in,
+                generator=generator,
+                dtype=torch.float64,
+                device=config.device,
+            )
+            * scale
+        )
+    return weights
+
+
+def init_feedback(
+    config: dcs.RunConfig,
+    width: int,
+    hidden_layers: int,
+    seed: int,
+) -> list[torch.Tensor]:
+    generator = torch.Generator(device=config.device)
+    generator.manual_seed(seed)
+    shapes = [(width, width)] * max(hidden_layers - 1, 0) + [
+        (width, config.output_dim)
+    ]
+    feedback: list[torch.Tensor] = []
+    for rows, cols in shapes:
+        feedback.append(
+            torch.randn(
+                rows,
+                cols,
+                generator=generator,
+                dtype=torch.float64,
+                device=config.device,
+            )
+            * dcs.feedback_scale(rows, config.feedback_scale)
+        )
+    return feedback
+
+
 def kernel_pair(
     weights: list[torch.Tensor],
     x: torch.Tensor,
@@ -211,14 +266,16 @@ def projection_alpha(k_fa_s: np.ndarray, k_fa0: np.ndarray, k_bp0: np.ndarray) -
 
 def run_one(
     config: dcs.RunConfig,
+    hidden_layers: int,
     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)
+    width = config.widths[0]
+    initial = initialize_weights(config, width, hidden_layers, init_seed)
+    feedback = init_feedback(config, width, hidden_layers, feedback_seed)
     with torch.no_grad():
         r0 = (dcs.predict(initial, x) - y).reshape(-1).cpu().numpy()
 
@@ -296,6 +353,9 @@ def run_one(
     retangent_gap = retangent_fa - retangent_bp
 
     return CompressedRow(
+        hidden_layers=hidden_layers,
+        width=width,
+        train_samples=config.train_samples,
         init_seed=init_seed,
         feedback_seed=feedback_seed,
         target_steps=config.steps,
@@ -334,6 +394,10 @@ def write_rows(path: Path, rows: list[CompressedRow]) -> None:
 
 def main() -> None:
     args = parse_args()
+    if args.hidden_layers < 1:
+        raise ValueError("--hidden-layers must be positive.")
+    if any(early_step >= args.target_steps for early_step in args.early_steps):
+        raise ValueError("Every --early-steps value must be smaller than --target-steps.")
     if args.torch_threads > 0:
         torch.set_num_threads(args.torch_threads)
     config = make_config(args)
@@ -354,7 +418,15 @@ def main() -> None:
                     flush=True,
                 )
                 rows.append(
-                    run_one(config, x_train, y_train, init_seed, feedback_seed, early_step)
+                    run_one(
+                        config,
+                        args.hidden_layers,
+                        x_train,
+                        y_train,
+                        init_seed,
+                        feedback_seed,
+                        early_step,
+                    )
                 )
 
     outdir = Path(args.outdir)
diff --git a/scripts/plot_compressed_operator_results.py b/scripts/plot_compressed_operator_results.py
index c23fa14..ee2f472 100644
--- a/scripts/plot_compressed_operator_results.py
+++ b/scripts/plot_compressed_operator_results.py
@@ -53,8 +53,23 @@ def read_rows(pattern: str, input_dim: int, output_dim: int, width: int) -> pd.D
     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)
+    if "width" in df.columns:
+        row_width = df["width"].astype(int)
+    else:
+        row_width = pd.Series(width, index=df.index)
+    if "hidden_layers" in df.columns:
+        hidden_layers = df["hidden_layers"].astype(int)
+    else:
+        hidden_layers = pd.Series(2, index=df.index)
+    parameter_count = (
+        input_dim * row_width
+        + (hidden_layers - 1).clip(lower=0) * row_width * row_width
+        + row_width * output_dim
+    )
+    constraint_rank = (
+        (hidden_layers - 1).clip(lower=0) * np.maximum(row_width * row_width - 1, 0)
+        + np.maximum(row_width * output_dim - 1, 0)
+    )
     df["hard_fa_capacity_margin"] = (
         parameter_count - constraint_rank - output_dim * df["train_samples"]
     )
diff --git a/scripts/summarize_operator_stress_grid.py b/scripts/summarize_operator_stress_grid.py
new file mode 100644
index 0000000..5c844e2
--- /dev/null
+++ b/scripts/summarize_operator_stress_grid.py
@@ -0,0 +1,276 @@
+#!/usr/bin/env python3
+"""Summarize finite-time operator-predictor stress tests."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+
+STRESS_RE = re.compile(
+    r"stress_depth(?P<depth>\d+)_w(?P<width>\d+)_T(?P<T>\d+)_s(?P<s>\d+)_N(?P<N>\d+)"
+)
+
+
+@dataclass(frozen=True)
+class SettingMetrics:
+    setting: str
+    hidden_layers: int
+    width: int
+    target_steps: int
+    early_steps: int
+    rows: int
+    fixed_mae: float
+    fixed_bias: float
+    compressed_mae: float
+    compressed_bias: float
+    linear_mae: float
+    linear_bias: float
+    retangent_mae: float
+    retangent_bias: float
+    linear_corr: float
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input-dim", type=int, default=16)
+    parser.add_argument("--output-dim", type=int, default=4)
+    parser.add_argument("--include-baseline", action="store_true")
+    parser.add_argument(
+        "--outdir",
+        type=Path,
+        default=Path("outputs/operator_stress_grid_summary"),
+    )
+    return parser.parse_args()
+
+
+def hard_margin(
+    input_dim: int,
+    output_dim: int,
+    width: np.ndarray,
+    hidden_layers: np.ndarray,
+    train_samples: np.ndarray,
+) -> np.ndarray:
+    parameter_count = (
+        input_dim * width
+        + np.maximum(hidden_layers - 1, 0) * width * width
+        + width * output_dim
+    )
+    constraint_rank = (
+        np.maximum(hidden_layers - 1, 0) * np.maximum(width * width - 1, 0)
+        + np.maximum(width * output_dim - 1, 0)
+    )
+    return parameter_count - constraint_rank - output_dim * train_samples
+
+
+def load_stress_rows(input_dim: int, output_dim: int, include_baseline: bool) -> pd.DataFrame:
+    frames: list[pd.DataFrame] = []
+    for path in sorted(Path("outputs").glob("stress_*/compressed_operator_rows.csv")):
+        match = STRESS_RE.fullmatch(path.parent.name)
+        if not match:
+            continue
+        meta = {key: int(value) for key, value in match.groupdict().items()}
+        frame = pd.read_csv(path)
+        frame["setting"] = (
+            f"d{meta['depth']}_w{meta['width']}_T{meta['T']}_s{meta['s']}"
+        )
+        frame["hidden_layers"] = meta["depth"]
+        frame["width"] = meta["width"]
+        frame["target_steps"] = meta["T"]
+        frame["early_steps"] = meta["s"]
+        frame["train_samples"] = meta["N"]
+        frames.append(frame)
+
+    if include_baseline:
+        for path in sorted(
+            Path("outputs").glob(
+                "compressed_operator_s20_256traj_T50_width64_N*/compressed_operator_rows.csv"
+            )
+        ):
+            train_samples = int(path.parent.name.rsplit("N", 1)[1])
+            frame = pd.read_csv(path)
+            frame["setting"] = "d2_w64_T50_s20_256traj"
+            frame["hidden_layers"] = 2
+            frame["width"] = 64
+            frame["target_steps"] = 50
+            frame["early_steps"] = 20
+            frame["train_samples"] = train_samples
+            frames.append(frame)
+
+    if not frames:
+        raise FileNotFoundError("No stress rows found.")
+
+    df = pd.concat(frames, ignore_index=True)
+    df["hard_fa_capacity_margin"] = hard_margin(
+        input_dim,
+        output_dim,
+        df["width"].to_numpy(dtype=np.float64),
+        df["hidden_layers"].to_numpy(dtype=np.float64),
+        df["train_samples"].to_numpy(dtype=np.float64),
+    )
+    return df
+
+
+def summarize(df: pd.DataFrame) -> list[SettingMetrics]:
+    rows: list[SettingMetrics] = []
+    for setting, group in df.groupby("setting"):
+        empirical = group["empirical_gap"].to_numpy()
+
+        def error_stats(column: str) -> tuple[float, float]:
+            error = group[column].to_numpy() - empirical
+            return float(np.mean(np.abs(error))), float(np.mean(error))
+
+        fixed_mae, fixed_bias = error_stats("fixed_gap")
+        compressed_mae, compressed_bias = error_stats("compressed_gap")
+        linear_mae, linear_bias = error_stats("linear_gap")
+        retangent_mae, retangent_bias = error_stats("retangent_gap")
+        linear_corr = float(np.corrcoef(group["linear_gap"], empirical)[0, 1])
+        rows.append(
+            SettingMetrics(
+                setting=setting,
+                hidden_layers=int(group["hidden_layers"].iloc[0]),
+                width=int(group["width"].iloc[0]),
+                target_steps=int(group["target_steps"].iloc[0]),
+                early_steps=int(group["early_steps"].iloc[0]),
+                rows=len(group),
+                fixed_mae=fixed_mae,
+                fixed_bias=fixed_bias,
+                compressed_mae=compressed_mae,
+                compressed_bias=compressed_bias,
+                linear_mae=linear_mae,
+                linear_bias=linear_bias,
+                retangent_mae=retangent_mae,
+                retangent_bias=retangent_bias,
+                linear_corr=linear_corr,
+            )
+        )
+    return sorted(rows, key=lambda row: row.setting)
+
+
+def write_metrics(path: Path, rows: list[SettingMetrics]) -> None:
+    with path.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(SettingMetrics.__annotations__.keys()))
+        writer.writeheader()
+        for row in rows:
+            writer.writerow(row.__dict__)
+
+
+def plot_mae_bars(metrics: list[SettingMetrics], outdir: Path) -> Path:
+    labels = [row.setting for row in metrics]
+    x = np.arange(len(labels))
+    width = 0.22
+    fig, ax = plt.subplots(figsize=(13.5, 6.4), dpi=180)
+    ax.bar(x - width, [row.fixed_mae for row in metrics], width, label="fixed K(0)", color="#777777")
+    ax.bar(x, [row.linear_mae for row in metrics], width, label="linear velocity", color="#1f5a9d")
+    ax.bar(
+        x + width,
+        [row.retangent_mae for row in metrics],
+        width,
+        label="early re-tangent",
+        color="#2b8a3e",
+    )
+    ax.set_xticks(x)
+    ax.set_xticklabels(labels, rotation=35, ha="right")
+    ax.set_ylabel("gap prediction MAE")
+    ax.set_title("Operator-predictor stress grid")
+    ax.legend(frameon=True)
+    ax.grid(axis="y", alpha=0.18)
+    fig.tight_layout()
+    path = outdir / "stress_grid_gap_mae_by_setting.png"
+    fig.savefig(path)
+    plt.close(fig)
+    return path
+
+
+def plot_linear_error_by_margin(df: pd.DataFrame, outdir: Path) -> Path:
+    fig, ax = plt.subplots(figsize=(11.5, 6.4), dpi=180)
+    for setting, group in df.groupby("setting"):
+        records = []
+        for margin, by_margin in group.groupby("hard_fa_capacity_margin"):
+            error = by_margin["linear_gap"].to_numpy() - by_margin["empirical_gap"].to_numpy()
+            records.append((margin, float(np.mean(error)), float(np.mean(np.abs(error)))))
+        records = sorted(records)
+        ax.plot(
+            [record[0] for record in records],
+            [record[1] for record in records],
+            marker="o",
+            linewidth=1.7,
+            label=setting,
+        )
+    ax.axhline(0, color="black", linewidth=1.0)
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.0)
+    ax.set_xlabel("hard FA capacity margin")
+    ax.set_ylabel("linear velocity prediction bias")
+    ax.set_title("Linear-velocity bias across capacity margins")
+    ax.legend(frameon=True, fontsize=8, ncol=2)
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    path = outdir / "stress_grid_linear_bias_by_margin.png"
+    fig.savefig(path)
+    plt.close(fig)
+    return path
+
+
+def plot_linear_scatter(df: pd.DataFrame, outdir: Path) -> Path:
+    fig, ax = plt.subplots(figsize=(7.2, 6.4), dpi=180)
+    settings = sorted(df["setting"].unique())
+    cmap = plt.get_cmap("tab10")
+    for index, setting in enumerate(settings):
+        group = df[df["setting"] == setting]
+        ax.scatter(
+            group["linear_gap"],
+            group["empirical_gap"],
+            s=22,
+            alpha=0.62,
+            color=cmap(index % 10),
+            label=setting,
+            edgecolors="none",
+        )
+    lo = float(min(df["linear_gap"].min(), df["empirical_gap"].min()))
+    hi = float(max(df["linear_gap"].max(), df["empirical_gap"].max()))
+    ax.plot([lo, hi], [lo, hi], color="black", linewidth=1.0)
+    ax.set_xlabel("linear velocity predicted gap")
+    ax.set_ylabel("empirical gap")
+    ax.set_title("Stress-grid predicted vs empirical gaps")
+    ax.legend(frameon=True, fontsize=8, ncol=2)
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    path = outdir / "stress_grid_linear_prediction_scatter.png"
+    fig.savefig(path)
+    plt.close(fig)
+    return path
+
+
+def main() -> None:
+    args = parse_args()
+    args.outdir.mkdir(parents=True, exist_ok=True)
+    df = load_stress_rows(args.input_dim, args.output_dim, args.include_baseline)
+    df.to_csv(args.outdir / "stress_grid_rows.csv", index=False)
+    metrics = summarize(df)
+    write_metrics(args.outdir / "stress_grid_metrics.csv", metrics)
+    paths = [
+        plot_mae_bars(metrics, args.outdir),
+        plot_linear_error_by_margin(df, args.outdir),
+        plot_linear_scatter(df, args.outdir),
+    ]
+    print(f"rows: {args.outdir / 'stress_grid_rows.csv'}")
+    print(f"metrics: {args.outdir / 'stress_grid_metrics.csv'}")
+    for row in metrics:
+        print(
+            f"{row.setting}: rows={row.rows}, fixed={row.fixed_mae:.6g}, "
+            f"linear={row.linear_mae:.6g}, retangent={row.retangent_mae:.6g}, "
+            f"linear_bias={row.linear_bias:.6g}, corr={row.linear_corr:.6g}"
+        )
+    for path in paths:
+        print(f"plot: {path}")
+
+
+if __name__ == "__main__":
+    main()