Wire NNUE eval into search, build the training pipeline, and add a fast

move/eval-agreement benchmark against Stockfish (tests/lop)

Implements the architecture decisions in src/docs/NNUE.md: HalfKAv2_hm
feature set, int16/int8 quantization, per-ply accumulator backup for
king moves, and a lazy-eval-ready integration into search.c/eval.c.
Adds the full annotate_pgns/ training pipeline (extraction, K-fit/
labeling, quiet filtering, tensor building, training, export) plus an
end-to-end orchestrator (run_nnue_pipeline.py) and validation tooling
(sanity_check_net.py, compare_vs_stockfish_net.py). Adds tests/lop:
a 10k-position list-of-positions benchmark comparing engine eval rank
and move choice against pre-computed Stockfish/lc0 references, much
cheaper than a full self-play match for day-to-day progress checks.

First sd10-matched self-play result against classic HCE: LLR +2.28
(H1=+2.94), -72=119+169 over 360 games, confirming the net is
"smarter per node" at matched depth. A real-time (--st 4) match is
currently weaker (candidate loses ~3 plies of depth to a ~3.3-3.7x
nps gap against an unrecalibrated search), consistent with
NNUE.md's adoption-criterion section -- pruning/reduction constants
still need retuning for the new eval's score scale before a fair
real-time comparison is possible.

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

1 files changed, 90 insertions, 8 deletions

diff --git a/annotate_pgns/nnue_model.py b/annotate_pgns/nnue_model.py
index acdc8fd..c483ce7 100644
--- a/annotate_pgns/nnue_model.py
+++ b/annotate_pgns/nnue_model.py
@@ -16,11 +16,54 @@ activations / quantization-aware scaling (this is a float32-only
 Phase A prototype, no export-for-C-inference step exists yet), the
 lazy-eval PSQT-term fast path, and the super-lazy material-only
 prefilter. All irrelevant to proving the training pipeline works.
+
+2026-10-17: added `material_probe`, a small auxiliary linear head on
+the same concatenated accumulator `fc1` reads, trained jointly with
+the main WDL loss (see nnue_train.py's --material-loss-weight) against
+each position's raw material balance (mover-relative pawns, no
+sigmoid). Added after a real, measured failure: a net trained with
+plain MSE (and separately, exponent=2.6) against WDL-sigmoid targets
+alone could not tell a lost pawn from a lost queen in crowded midgame
+positions (see sanity_check_net.py) -- directionally correct but
+magnitude-blind. The probe supplies a non-saturating gradient straight
+at material, forcing `ft`'s embedding weights to keep material
+linearly recoverable even where the main WDL loss has little pressure
+to care. It's training-only scaffolding, like the SF doc's "virtual
+features" factorization trick -- dropped before export, zero
+inference-time or quantization-format cost.
+
+2026-10-18: added `ft_p`, the actual Stockfish-doc "P" factor (see
+NNUE.md's "feature factorization" status entry) -- a second,
+king-bucket-INDEPENDENT EmbeddingBag over PS_NB=704 rows
+(piece_square x piece_type x piece_color only), summed into the same
+accumulator `ft` feeds. Motivation: HalfKAv2_hm's real 22528-row
+table is king-bucket-conditioned, so a (piece, square, color) tuple
+that's common in general chess but rare in a specific king bucket
+gets its own mostly-untrained row -- measured directly on the
+Stockfish self-play training pool: 9.9% of rows never fire at all,
+23.6% see fewer than 50 examples. `ft_p`'s 704 rows see ~32x more
+observations each (one per king bucket collapsed into one), supplying
+dense gradient to a king-bucket-independent piece-value prior that
+`ft`'s sparser rows can then specialize away from. Zero-initialized
+(nn.init.zeros_) specifically so it contributes nothing at step 0 --
+a default (unit-variance) init would silently double the
+accumulator's effective scale at init, confounding the factorization
+effect with an incidental learning-rate change. Deliberately no bias
+of its own -- `ft_bias` stays the only bias on this sum, which is
+also what keeps the post-training coalescing step (see
+export_nnue_weights.py) exact rather than approximate. Training-only,
+like material_probe: coalesced into `ft.weight` and dropped before
+export, since `idx % PS_NB` (the P-factor's index for real feature
+`idx`) is recoverable from the real index alone, needing no change to
+build_training_tensors.py's cached tensors. Deliberately NOT adding
+Stockfish's other two documented factors ("K", "HalfRelativeKP") --
+HalfKAv2_hm already encodes the king directly via the shared PS_KING
+plane, making a K factor redundant here; one variable at a time.
 """
 import torch
 import torch.nn as nn
 
-from nnue_feature_spec import DIMENSIONS
+from nnue_feature_spec import DIMENSIONS, PS_NB
 
 L1 = 256   # feature-transformer output width (per perspective)
 L2 = 32
@@ -30,7 +73,8 @@ DEFAULT_K = 1.3123   # pawns; fit_k_and_label.py's Phase-A train-split fit
 
 
 class NnueNet(nn.Module):
-    def __init__(self, dimensions=DIMENSIONS, l1=L1, l2=L2, l3=L3):
+    def __init__(self, dimensions=DIMENSIONS, l1=L1, l2=L2, l3=L3,
+                 ps_nb=PS_NB):
         super().__init__()
         # Single shared feature transformer: one [dimensions, l1]
         # weight matrix, applied once per perspective to that
@@ -40,15 +84,37 @@ class NnueNet(nn.Module):
         self.ft = nn.EmbeddingBag(dimensions, l1, mode="sum")
         self.ft_bias = nn.Parameter(torch.zeros(l1))
 
+        # Training-only "P" factorization table -- see module
+        # docstring's 2026-10-18 note. King-bucket-independent
+        # (ps_nb=704 rows, vs ft's full 22528), summed into the same
+        # accumulator. Zero-init so it contributes nothing at step 0.
+        self.ft_p = nn.EmbeddingBag(ps_nb, l1, mode="sum")
+        nn.init.zeros_(self.ft_p.weight)
+        self.ps_nb = ps_nb
+
         self.fc1 = nn.Linear(2 * l1, l2)
         self.fc2 = nn.Linear(l2, l3)
         self.fc3 = nn.Linear(l3, 1)
 
-    def forward(self, white_idx, white_off, black_idx, black_off, stm):
-        """stm: float/long tensor of shape [B], 1.0 if white to move
-        else 0.0. Returns raw score in pawn-ish units, shape [B]."""
-        white_acc = self.ft(white_idx, white_off) + self.ft_bias
-        black_acc = self.ft(black_idx, black_off) + self.ft_bias
+        # Auxiliary, training-only head -- see module docstring's
+        # 2026-10-17 note. Reads the same concatenated accumulator
+        # fc1 does; not used at inference/export time.
+        self.material_probe = nn.Linear(2 * l1, 1)
+
+    def accumulator(self, white_idx, white_off, black_idx, black_off, stm):
+        """The ReLU'd, own-perspective-first concatenated accumulator
+        (shape [B, 2*l1]) -- shared input to both fc1 (the real eval
+        head) and material_probe (the auxiliary head). ft_p's
+        contribution uses the identical bag offsets as ft -- same
+        positions, same active-feature counts, just a second lookup
+        of each index modulo ps_nb into the smaller, denser table."""
+        white_idx_p = white_idx % self.ps_nb
+        black_idx_p = black_idx % self.ps_nb
+
+        white_acc = (self.ft(white_idx, white_off)
+                     + self.ft_p(white_idx_p, white_off) + self.ft_bias)
+        black_acc = (self.ft(black_idx, black_off)
+                     + self.ft_p(black_idx_p, black_off) + self.ft_bias)
 
         white_acc = torch.relu(white_acc)
         black_acc = torch.relu(black_acc)
@@ -57,13 +123,29 @@ class NnueNet(nn.Module):
         # own-perspective-first concatenation, per side to move.
         own = stm * white_acc + (1.0 - stm) * black_acc
         enemy = stm * black_acc + (1.0 - stm) * white_acc
-        x = torch.cat([own, enemy], dim=1)
+        return torch.cat([own, enemy], dim=1)
 
+    def forward(self, white_idx, white_off, black_idx, black_off, stm):
+        """stm: float/long tensor of shape [B], 1.0 if white to move
+        else 0.0. Returns raw score in pawn-ish units, shape [B]."""
+        x = self.accumulator(white_idx, white_off, black_idx, black_off, stm)
         x = torch.relu(self.fc1(x))
         x = torch.relu(self.fc2(x))
         x = self.fc3(x)
         return x.view(-1)
 
+    def forward_with_material(self, white_idx, white_off, black_idx, black_off, stm):
+        """Like forward(), but also returns the auxiliary material
+        probe's prediction (mover-relative pawns, no sigmoid) off the
+        same accumulator -- one extra Linear(2*l1, 1), training-only.
+        Returns (raw_score, material_pred), both shape [B]."""
+        x = self.accumulator(white_idx, white_off, black_idx, black_off, stm)
+        material_pred = self.material_probe(x).view(-1)
+        y = torch.relu(self.fc1(x))
+        y = torch.relu(self.fc2(y))
+        y = self.fc3(y)
+        return y.view(-1), material_pred
+
     def win_probability(self, white_idx, white_off, black_idx, black_off, stm, k=DEFAULT_K):
         raw = self.forward(white_idx, white_off, black_idx, black_off, stm)
         return torch.sigmoid(raw / k)