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
|
#!/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 (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, 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)
# 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)
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 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)
|
