1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
|
#!/usr/bin/env python3
"""The Phase A prototype net: a small, known HalfKAv2_hm architecture
per NNUE.md's architecture decisions -- 256x2-32-32-1-ish, a single
shared feature-transformer weight matrix applied independently to
each perspective's active-feature set (that sharing, not two
separate tables, is what makes the two 256-wide "accumulator" halves
comparable before concatenation), own-perspective-half-first
concatenation keyed by side-to-move so the net's single scalar
output is side-to-move-relative by construction (per the "Output
convention" architecture decision), and a sigmoid-WDL loss against
the already-fitted K=1.3123 (per "WDL-sigmoid target conversion
built and run", fit_k_and_label.py) rather than raw centipawn MSE.
Deliberately NOT doing yet (explicitly deferred, see NNUE.md): clipped
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 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.
2026-10-20: replaced by `score_probe` (same mechanism, different
target -- see that section below) and removed. Once `score_probe`
regresses Stockfish's actual raw score directly, `material_probe`'s
target (piece count) is a strict special case of what `score_probe`
already has to recover -- material dominates Stockfish's own score in
almost any position, so a net that fits the real score well has
implicitly recovered material too. Kept as one objective instead of
two overlapping ones; re-add material_probe alongside score_probe only
if a future sanity check shows score_probe alone is still
material-blind in some case it doesn't already cover.
2026-10-20: tried `score_probe`, same mechanism as material_probe
(small Linear head on the same accumulator, trained jointly, dropped
before export) targeting Stockfish's own raw per-position score
directly via plain MSE -- motivated by a measured failure of a
different kind from material_probe's: in clearly-losing positions,
the main WDL head's score magnitude did not reflect how lost the
position actually was (-2 for "lost" and -3 for "about to be mated",
no meaningful gradation), because the WDL loss only ever supervises
`raw` through `sigmoid(raw/K)` against a win-probability target, which
saturates hard once a position is clearly decided (K=2.0659 on the
2026-10-20 corpus -- sigmoid(8/K)=0.9796 vs sigmoid(13/K)=0.9982, a
target spread of under 0.02 for a 5-pawn difference in the real
score). **Measured, then reverted**: a medium-scale (500k positions,
~15 epochs) controlled comparison (score_probe-weighted run vs. an
otherwise-identical control) showed only a small improvement in how
well `forward()`'s own raw output tracked the real score at the
losing extremes -- compression was still clearly present in both.
Root cause: score_probe's gradient only ever reaches the *shared
accumulator* (`ft`/`ft_p`/`ft_bias`), never `fc1`/`fc2`/`fc3` --
nothing forces the actually-exported path to widen its own range, it
can only optionally make use of a richer upstream representation if
doing so happens to help the (still-saturating) WDL loss too.
2026-10-20: replaced score_probe with a direct, unsaturated regression
term on `raw` itself: `MSE(raw, clipped_score)`, added straight
alongside the WDL loss in nnue_train.py, no auxiliary head at all.
This puts gradient directly on `fc1`/`fc2`/`fc3` -- the actual
exported path -- rather than hoping a richer accumulator happens to
propagate through. Not a competing objective with the WDL loss, either:
since `p_search = sigmoid(score/K)`, the WDL loss's own implied ideal
`raw` is already approximately `score` (especially at lambda close to
1.0) -- this term is an unsaturated restatement of the same target,
not a different one, so no separate probe/model-surface-area is
needed to get its benefit. See nnue_train.py's --score-loss-weight.
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, PS_NB
L1 = 256 # feature-transformer output width (per perspective)
L2 = 32
L3 = 32
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,
ps_nb=PS_NB):
super().__init__()
# Single shared feature transformer: one [dimensions, l1]
# weight matrix, applied once per perspective to that
# perspective's own active-feature set. mode="sum" matches
# Stockfish's own feature-transformer semantics (accumulate
# one weight row per active feature).
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 accumulator(self, white_idx, white_off, black_idx, black_off, stm):
"""The ReLU'd, own-perspective-first concatenated accumulator
(shape [B, 2*l1]) -- input to fc1. 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)
stm = stm.view(-1, 1).float()
# 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
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 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)
|
