NNUE: re-fit raw-to-score calibration, NNU4 weight format, super-lazy default-on, lazy probe measured net-negative

Fits an affine (a + b*raw) correction into every exported .nnue file
(export_nnue_weights.py's CalibrateScoreMapping) so the engine's
consumed centipawn output tracks Stockfish's real score scale rather
than the net's nominal raw WDL-sigmoid output -- training variance
alone shifts this mapping across retrains even with an identical
recipe, so it's re-fit and baked in per export rather than assumed
constant. Weight format bumped to "NNU4" to carry the two new floats;
older NNU1/2/3 files are no longer readable.

Super-lazy (material-only) pre-filter promoted to default-on
(g_fUseNnueSuperLazy = TRUE) with a recalibrated per-material-bucket
margin table (measured against the recalibrated net, 16-shard
cross_sectional_200k.ep_ sample) -- real, measured nps win across all
three curated suites. The separate lazy-probe gate (accumulator-based
early exit) is measured net-negative on real search traffic at this
net's probe-vs-forward latency ratio and stays off by default
(g_fUseNnueLazyProbe = FALSE); code kept live since margin tuning
could still flip this, per docs/NNUE.md.

Also: direct per-call latency instrumentation for NnueProbeScore/
NnueForward (root.c, EVAL_TIME build) to make the lazy-probe
breakeven math measurable directly instead of inferred; QSearch
futility margins given real-data-derived NNUE_FULL/NNUE_LAZY
starting points.

Co-Authored-By: Claude Sonnet 5 <[email protected]>

1 files changed, 94 insertions, 45 deletions

diff --git a/annotate_pgns/nnue_train.py b/annotate_pgns/nnue_train.py
index ea3d0a9..e413aa9 100644
--- a/annotate_pgns/nnue_train.py
+++ b/annotate_pgns/nnue_train.py
@@ -42,17 +42,22 @@ class CsrTensors:
         self.black_off = np.load(tensor_dir / "black_off.npy")
         self.stm = np.load(tensor_dir / "stm.npy")
         self.target = np.load(tensor_dir / "target.npy")
-        # material.npy added 2026-10-17 (see nnue_model.py's
-        # material_probe) -- optional for backward compat with
-        # tensor dirs built before this field existed.
-        material_path = tensor_dir / "material.npy"
-        self.material = (np.load(material_path) if material_path.exists()
-                          else np.zeros_like(self.target))
+        # score_probe_val.npy/score_probe_valid.npy added 2026-10-20
+        # (see nnue_model.py's score_probe, replacing material_probe)
+        # -- optional for backward compat with tensor dirs built before
+        # this field existed.
+        score_path = tensor_dir / "score_probe_val.npy"
+        valid_path = tensor_dir / "score_probe_valid.npy"
+        self.score_probe_val = (np.load(score_path) if score_path.exists()
+                                 else np.zeros_like(self.target))
+        self.score_probe_valid = (np.load(valid_path) if valid_path.exists()
+                                   else np.zeros_like(self.target, dtype=np.uint8))
         self.n = len(self.target)
         assert len(self.white_off) == self.n + 1
         assert len(self.black_off) == self.n + 1
         assert len(self.stm) == self.n
-        assert len(self.material) == self.n
+        assert len(self.score_probe_val) == self.n
+        assert len(self.score_probe_valid) == self.n
 
     def __len__(self):
         return self.n
@@ -93,7 +98,8 @@ def iterate_batches(data: CsrTensors, batch_size, shuffle, rng):
         b_idx, b_off = _gather_batch(data.black_idx, data.black_off, rows_sorted)
         stm = data.stm[rows_sorted]
         target = data.target[rows_sorted]
-        material = data.material[rows_sorted]
+        score_probe_val = data.score_probe_val[rows_sorted]
+        score_probe_valid = data.score_probe_valid[rows_sorted]
         yield (
             torch.from_numpy(w_idx.astype(np.int64, copy=False)),
             torch.from_numpy(w_off),
@@ -101,7 +107,8 @@ def iterate_batches(data: CsrTensors, batch_size, shuffle, rng):
             torch.from_numpy(b_off),
             torch.from_numpy(stm.astype(np.int64, copy=False)),
             torch.from_numpy(target.copy()),
-            torch.from_numpy(material.astype(np.float32, copy=True)),
+            torch.from_numpy(score_probe_val.astype(np.float32, copy=True)),
+            torch.from_numpy(score_probe_valid.astype(np.float32, copy=True)),
         )
 
 
@@ -118,34 +125,63 @@ def powered_loss(pred, target, exponent, reduction="mean"):
     return err.sum() if reduction == "sum" else err.mean()
 
 
-def evaluate(model, data, batch_size, k, loss_exponent, device="cpu"):
-    """Returns (wdl_loss, material_mse) -- kept separate rather than
-    pre-combined so the log/checkpoint can show whether the material
-    probe is actually converging, independently of whatever weight
+def score_loss_fn(raw, score_val, score_valid, score_loss_type, huber_delta):
+    if score_loss_type == "huber":
+        return masked_huber_sum(raw, score_val, score_valid, huber_delta)
+    return masked_mse_sum(raw, score_val, score_valid)
+
+
+def evaluate(model, data, batch_size, k, loss_exponent, device="cpu",
+             score_loss_type="mse", huber_delta=1.0):
+    """Returns (wdl_loss, score_mse) -- kept separate rather than
+    pre-combined so the log/checkpoint can show whether `raw` is
+    actually tracking the real score, independently of whatever weight
     the caller gives it in the training loss."""
     model.eval()
     total_wdl_loss = 0.0
-    total_material_loss = 0.0
+    total_score_loss = 0.0
     total_n = 0
+    total_score_valid = 0
     rng = np.random.default_rng(0)
     with torch.no_grad():
-        for w_idx, w_off, b_idx, b_off, stm, target, material in iterate_batches(
+        for w_idx, w_off, b_idx, b_off, stm, target, score_val, score_valid in iterate_batches(
                 data, batch_size, shuffle=False, rng=rng):
-            w_idx, w_off, b_idx, b_off, stm, target, material = _to_device(
-                (w_idx, w_off, b_idx, b_off, stm, target, material), device)
-            raw, material_pred = model.forward_with_material(w_idx, w_off, b_idx, b_off, stm)
+            w_idx, w_off, b_idx, b_off, stm, target, score_val, score_valid = _to_device(
+                (w_idx, w_off, b_idx, b_off, stm, target, score_val, score_valid), device)
+            raw = model.forward(w_idx, w_off, b_idx, b_off, stm)
             pred = torch.sigmoid(raw / k)
             wdl_loss = powered_loss(pred, target, loss_exponent, reduction="sum")
-            material_loss = nn_mse_sum(material_pred, material)
+            score_loss = score_loss_fn(raw, score_val, score_valid, score_loss_type, huber_delta)
             total_wdl_loss += wdl_loss.item()
-            total_material_loss += material_loss.item()
+            total_score_loss += score_loss.item()
             total_n += len(target)
+            total_score_valid += int(score_valid.sum().item())
     model.train()
-    return total_wdl_loss / max(total_n, 1), total_material_loss / max(total_n, 1)
+    return (total_wdl_loss / max(total_n, 1),
+            total_score_loss / max(total_score_valid, 1))
+
+
+def masked_mse_sum(pred, target, valid):
+    """Sum of squared error over only the rows where valid != 0 --
+    exact/mate rows carry no real score target (see
+    build_training_tensors.py's score_probe_target) and must not pull
+    `raw` toward an arbitrary 0.0 placeholder."""
+    return ((pred - target).pow(2) * valid).sum()
 
 
-def nn_mse_sum(pred, target):
-    return (pred - target).pow(2).sum()
+def masked_huber_sum(pred, target, valid, delta):
+    """Sum of Huber loss over only the rows where valid != 0 (same
+    masking as masked_mse_sum). Quadratic for |error| < delta
+    (preserves fine resolution on close, contested positions -- the
+    ones that actually matter for move choice), linear beyond it
+    (bounds how much a handful of large-magnitude blowout positions
+    can dominate the loss/gradient, unlike plain MSE). delta is in the
+    same pawn units as `pred`/`target`."""
+    err = (pred - target).abs()
+    quad = 0.5 * err.pow(2)
+    lin = delta * (err - 0.5 * delta)
+    loss = torch.where(err < delta, quad, lin)
+    return (loss * valid).sum()
 
 
 def main():
@@ -164,16 +200,26 @@ def main():
                           "loss ratio between small and large WDL-space "
                           "errors, pushing harder to fix decisively-wrong "
                           "predictions. See powered_loss()'s docstring.")
-    ap.add_argument("--material-loss-weight", type=float, default=0.0,
-                     help="weight on the auxiliary material_probe MSE "
-                          "term (plain pawns^2, NOT WDL-space -- so this "
-                          "is typically orders of magnitude larger in raw "
+    ap.add_argument("--score-loss-type", choices=["mse", "huber"], default="mse",
+                     help="loss shape for the direct raw-vs-score term. "
+                          "'mse': plain squared error (simple, but a "
+                          "handful of large-magnitude blowout positions "
+                          "can dominate it). 'huber': quadratic below "
+                          "--huber-delta, linear beyond it -- bounds "
+                          "blowout influence while keeping fine "
+                          "resolution on close/contested positions.")
+    ap.add_argument("--huber-delta", type=float, default=1.0,
+                     help="Huber loss transition point, in pawns -- "
+                          "only used when --score-loss-type huber.")
+    ap.add_argument("--score-loss-weight", type=float, default=0.0,
+                     help="weight on a direct MSE(raw, clipped Stockfish "
+                          "score) term, straight on fc3's own output "
+                          "(plain pawns^2, NOT WDL-space -- so this is "
+                          "typically orders of magnitude larger in raw "
                           "scale than the WDL loss; start small, e.g. "
-                          "0.001-0.01, and check sanity_check_net.py "
-                          "rather than assume a value works). 0.0 "
-                          "(default) disables it entirely -- the probe "
-                          "head still exists but gets no gradient, "
-                          "equivalent to not having added it.")
+                          "0.001, and check both holdout losses in "
+                          "train_log.jsonl rather than assume a value "
+                          "works). 0.0 (default) disables it entirely.")
     ap.add_argument("--threads", type=int, default=16,
                      help="torch CPU thread count -- kept below this box's "
                           "24 cores (not maxed out) since this is a shared "
@@ -277,16 +323,17 @@ def main():
             epoch_n = 0
             t_epoch = time.time()
 
-            for w_idx, w_off, b_idx, b_off, stm, target, material in iterate_batches(
+            for w_idx, w_off, b_idx, b_off, stm, target, score_val, score_valid in iterate_batches(
                     train_data, args.batch_size, shuffle=True, rng=rng):
-                w_idx, w_off, b_idx, b_off, stm, target, material = _to_device(
-                    (w_idx, w_off, b_idx, b_off, stm, target, material), args.device)
+                w_idx, w_off, b_idx, b_off, stm, target, score_val, score_valid = _to_device(
+                    (w_idx, w_off, b_idx, b_off, stm, target, score_val, score_valid), args.device)
                 optimizer.zero_grad()
-                raw, material_pred = model.forward_with_material(w_idx, w_off, b_idx, b_off, stm)
+                raw = model.forward(w_idx, w_off, b_idx, b_off, stm)
                 pred = torch.sigmoid(raw / args.k)
                 wdl_loss = powered_loss(pred, target, args.loss_exponent)
-                material_loss = nn_mse_sum(material_pred, material) / len(target)
-                loss = wdl_loss + args.material_loss_weight * material_loss
+                score_loss = score_loss_fn(raw, score_val, score_valid,
+                                            args.score_loss_type, args.huber_delta) / len(target)
+                loss = wdl_loss + args.score_loss_weight * score_loss
                 loss.backward()
                 optimizer.step()
 
@@ -296,34 +343,36 @@ def main():
                 global_step += 1
 
                 if global_step % args.eval_every_steps == 0:
-                    holdout_wdl, holdout_material = evaluate(
+                    holdout_wdl, holdout_score = evaluate(
                         model, holdout_data, args.batch_size, args.k,
-                        args.loss_exponent, args.device)
+                        args.loss_exponent, args.device,
+                        args.score_loss_type, args.huber_delta)
                     elapsed = time.time() - t_start
                     rec = {
                         "step": global_step, "epoch": epoch,
                         "train_mse_running": epoch_loss / max(epoch_n, 1),
                         "holdout_mse": holdout_wdl,
-                        "holdout_material_mse": holdout_material,
+                        "holdout_score_probe_mse": holdout_score,
                         "elapsed_s": round(elapsed, 1),
                         "positions_seen": epoch_n + epoch * len(train_data),
                     }
                     print(f"  step {global_step} epoch {epoch} "
                           f"train_mse={rec['train_mse_running']:.5f} "
                           f"holdout_mse={holdout_wdl:.5f} "
-                          f"holdout_material_mse={holdout_material:.4f} "
+                          f"holdout_score_probe_mse={holdout_score:.4f} "
                           f"elapsed={elapsed:.0f}s")
                     log_f.write(json.dumps(rec) + "\n")
                     log_f.flush()
 
-            holdout_wdl, holdout_material = evaluate(
+            holdout_wdl, holdout_score = evaluate(
                 model, holdout_data, args.batch_size, args.k,
-                args.loss_exponent, args.device)
+                args.loss_exponent, args.device,
+                args.score_loss_type, args.huber_delta)
             holdout_mse = holdout_wdl
             epoch_elapsed = time.time() - t_epoch
             print(f"Epoch {epoch} done: train_mse={epoch_loss / max(epoch_n, 1):.5f} "
                   f"holdout_mse={holdout_wdl:.5f} "
-                  f"holdout_material_mse={holdout_material:.4f} ({epoch_elapsed:.0f}s)")
+                  f"holdout_score_probe_mse={holdout_score:.4f} ({epoch_elapsed:.0f}s)")
 
             ckpt_path = outdir / f"ckpt_epoch{epoch:03d}.pt"
             torch.save({