Check directional predictor across N

1 files changed, 46 insertions, 0 deletions

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,