Check directional predictor across N

2 files changed, 124 insertions, 0 deletions

diff --git a/notes/14_early_predictor_probe.md b/notes/14_early_predictor_probe.md
index 5cfe058..43b9b6f 100644
--- a/notes/14_early_predictor_probe.md
+++ b/notes/14_early_predictor_probe.md
@@ -173,3 +173,81 @@ T\lambda_0+\frac{T(T-1)}{2}\dot\lambda_0
 \]
 
 This is more directly tied to loss dynamics than global kernel overlap.
+
+## Cross-N Directional Predictor Check
+
+We also tested a no-fit directional-gain extrapolation across
+
+\[
+N\in\{64,96,128,160,192,224,256,320\}
+\]
+
+with width fixed at \(64\), target horizon \(T=50\), 2 init seeds, and 4
+feedback seeds per init.
+
+The attempted predictor used
+
+\[
+\lambda_t\approx\lambda_0+t\dot\lambda_0,
+\qquad
+\dot\lambda_0\approx\frac{\lambda_s-\lambda_0}{s},
+\]
+
+and
+
+\[
+L_T
+\approx
+L_0
+\exp\left[
+-\frac{2\eta}{N}
+\left(
+T\lambda_0+\frac{T(T-1)}{2}\dot\lambda_0
+\right)
+\right].
+\]
+
+This failed as a finite-\(T\) gap predictor. Across \(N\), the directional-gap
+prediction had low correlation and large negative bias:
+
+| early step | fixed \(K(0)\) corr | fixed MAE | directional corr | directional MAE |
+|---:|---:|---:|---:|---:|
+| 1 | 0.979 | 0.0241 | 0.124 | 0.2838 |
+| 2 | 0.979 | 0.0241 | 0.138 | 0.2964 |
+| 5 | 0.979 | 0.0241 | 0.110 | 0.3686 |
+
+So the simple directional exponential formula should not be used.
+
+The reason is likely that a scalar gain
+
+\[
+\lambda_t=\frac{r_t^\top K_t r_t}{\|r_t\|^2}
+\]
+
+does not preserve residual direction changes. BP and FA rotate the residual in
+different ways, and the gap is a difference between two residual trajectories,
+not just a difference between two scalar decay rates.
+
+The more viable predictive object is therefore not a scalar exponential loss
+model. It is a compressed operator model, for example:
+
+\[
+K_t^{FA}
+\approx
+K_0^{FA}
++
+\alpha_t
+\left(K_0^{BP}-K_0^{FA}\right),
+\]
+
+or a low-rank/residual-subspace version. The scalar order parameter can control
+operator interpolation, but the residual update must remain vector-valued:
+
+\[
+r_{t+1}
+=
+\left(I-\frac{\eta}{N}K_t^{compressed}\right)r_t.
+\]
+
+This preserves direction changes and is closer to the measured \(K(t)\) product
+that worked.
diff --git a/scripts/early_kernel_predictors.py b/scripts/early_kernel_predictors.py
index e791a9d..3837fc6 100644
--- a/scripts/early_kernel_predictors.py
+++ b/scripts/early_kernel_predictors.py
@@ -33,6 +33,9 @@ class PredictorRow:
     fixed_bp_loss: float
     fixed_fa_loss: float
     fixed_gap: float
+    directional_bp_loss: float
+    directional_fa_loss: float
+    directional_gap: float
     kfa_fro_drift: float
     kbp_fro_drift: float
     fa_bp_overlap_0: float
@@ -143,6 +146,10 @@ def fixed_prediction(
     return 0.5 * float(rb @ rb) / samples, 0.5 * float(rf @ rf) / samples
 
 
+def loss_from_residual(residual: np.ndarray, samples: int) -> float:
+    return 0.5 * float(residual @ residual) / samples
+
+
 def overlap(a: np.ndarray, b: np.ndarray) -> float:
     denom = np.linalg.norm(a, "fro") * np.linalg.norm(b, "fro")
     if denom == 0:
@@ -157,6 +164,24 @@ def directional_gain(kernel: np.ndarray, residual: np.ndarray) -> float:
     return float(residual @ (kernel @ residual) / denom)
 
 
+def directional_loss_prediction(
+    residual: np.ndarray,
+    lambda_0: float,
+    lambda_s: float,
+    early_step: int,
+    lr: float,
+    target_steps: int,
+    samples: int,
+) -> float:
+    slope = (lambda_s - lambda_0) / early_step
+    cumulative = target_steps * lambda_0 + 0.5 * target_steps * (target_steps - 1) * slope
+    # Keep the early linear extrapolation from producing an unstable negative
+    # integrated rate; this is a stability projection, not a fitted coefficient.
+    cumulative = max(cumulative, 0.0)
+    initial_loss = loss_from_residual(residual, samples)
+    return initial_loss * float(np.exp(-(2.0 * lr / samples) * cumulative))
+
+
 def train_to_steps(
     weights: list[torch.Tensor],
     x: torch.Tensor,
@@ -211,6 +236,24 @@ def run_one(
     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)
+    directional_bp = directional_loss_prediction(
+        r0,
+        bp_directional_gain_0,
+        bp_directional_gain_s,
+        early_step,
+        config.lr,
+        config.steps,
+        config.train_samples,
+    )
+    directional_fa = directional_loss_prediction(
+        r0,
+        fa_directional_gain_0,
+        fa_directional_gain_s,
+        early_step,
+        config.lr,
+        config.steps,
+        config.train_samples,
+    )
 
     return PredictorRow(
         init_seed=init_seed,
@@ -223,6 +266,9 @@ def run_one(
         fixed_bp_loss=fixed_bp,
         fixed_fa_loss=fixed_fa,
         fixed_gap=fixed_fa - fixed_bp,
+        directional_bp_loss=directional_bp,
+        directional_fa_loss=directional_fa,
+        directional_gap=directional_fa - directional_bp,
         kfa_fro_drift=kfa_fro_drift,
         kbp_fro_drift=kbp_fro_drift,
         fa_bp_overlap_0=fa_bp_overlap_0,