credit Mitchell Riley: we can drop 65 parameters for about the same strength; also train for longer

The idea is that the current model is M_theta(board, ply) = "score for player 0", but we can instead create a different model M'_theta(board) = "score for player 0 assuming that it's player 0's move" and for a certain pair of functions sigma, normalise, recover M_theta(board, ply) = sigma(ply) M'_theta(normalise(board, ply)).

This pair of functions is sigma(ply) = +/- 1 (parity), and normalise being a conditional interchange of pieces on the board depending on the parity of the ply. The only things stopping this from being completely straightforward is that nn_train.py evaluates M'_theta itself so we need to pass sigma(ply) to the training loop.

6x6 is still untrained.

3 files changed, 34 insertions, 20 deletions

diff --git a/src/match.c b/src/match.c
index 6ae9fa7..459fab5 100644
--- a/src/match.c
+++ b/src/match.c
@@ -40,7 +40,7 @@ static inline float five_cand(tak_state_p state) {
   nn_prepare_input(cur_board, state, cand_five_embedding);
   return eval_net(cur_board, FIVE_INP_NUM, &cand_five_dense_weights[0][0],
                   cand_five_dense_biases, cand_five_logit_weights, cand_five_logit_bias,
-                  FIVE_DENSE_NUM);
+                  FIVE_DENSE_NUM, state->ply & 1 ? 1.0f : -1.0f);
 }
 
 static inline float six_cand(tak_state_p state) {
@@ -49,7 +49,7 @@ static inline float six_cand(tak_state_p state) {
   nn_prepare_input(cur_board, state, cand_six_embedding);
   return eval_net(cur_board, SIX_INP_NUM, &cand_six_dense_weights[0][0],
                   cand_six_dense_biases, cand_six_logit_weights, cand_six_logit_bias,
-                  SIX_DENSE_NUM);
+                  SIX_DENSE_NUM, state->ply & 1 ? 1.0f : -1.0f);
 }
 
 
diff --git a/src/nn_train.py b/src/nn_train.py
index 9e2b482..0e08da7 100755
--- a/src/nn_train.py
+++ b/src/nn_train.py
@@ -59,8 +59,8 @@ from torch import Tensor
 
 # (input dim, hidden dim) per board size; must match include/weights_<size>.h
 NET_DIMS: dict[int, tuple[int, int]] = {
-    5: (28, 64),  # FIVE_INP_NUM, FIVE_DENSE_NUM
-    6: (39, 64),  # SIX_INP_NUM, SIX_DENSE_NUM
+    5: (27, 64),  # FIVE_INP_NUM, FIVE_DENSE_NUM
+    6: (38, 64),  # SIX_INP_NUM, SIX_DENSE_NUM
 }
 
 
@@ -121,12 +121,13 @@ class SelfPlayRecords:
     results: Tensor  # [n] final played result from black's perspective
     terminals: Tensor  # [n] true when the recorded position is a terminal leaf
     terminal_values: Tensor  # [n] exact value of a terminal leaf, otherwise 0
+    parity_signs: Tensor  # [n] ±1 converting the net's output to "absolute" perspective
 
     @classmethod
     def load_selfplay(
         cls, fn: Path, size: int, device: torch.device
     ) -> SelfPlayRecords:
-        # [net input] [game_id] [played result] [terminal flag] [leaf result]
+        # [net input] [game_id] [played result] [terminal flag] [leaf result] [parity sign]
         df = pd.read_csv(fn, header=None)
         inp_dim = NET_DIMS[size][0]
         return cls(
@@ -154,6 +155,11 @@ class SelfPlayRecords:
                 dtype=torch.float32,
                 device=device,
             ),
+            parity_signs=torch.tensor(
+                df.iloc[:, inp_dim + 4].to_numpy(dtype=np.float32),
+                dtype=torch.float32,
+                device=device,
+            ),
         )
 
     @classmethod
@@ -167,14 +173,15 @@ class SelfPlayRecords:
             results=torch.concat([spr.results for spr in sprs]),
             terminals=torch.concat([spr.terminals for spr in sprs]),
             terminal_values=torch.concat([spr.terminal_values for spr in sprs]),
+            parity_signs=torch.concat([spr.parity_signs for spr in sprs]),
         )
 
     def make_training_data(
         self, model: Module, td_lambda: float
-    ) -> tuple[Tensor, Tensor]:
+    ) -> tuple[Tensor, Tensor, Tensor]:
         model.eval()
         with torch.no_grad():
-            values = model(self.positions)
+            values = model(self.positions) * self.parity_signs
 
         targets = torch.empty_like(values)
 
@@ -191,7 +198,7 @@ class SelfPlayRecords:
             game_targets[game_terminals] = game_terminal_values[game_terminals]
             targets[in_game] = game_targets
 
-        return self.positions, targets
+        return self.positions, targets, self.parity_signs
 
 
 # ===================================================================
@@ -208,10 +215,10 @@ class Net(Module):
         self.output = Linear(hidden, 1, bias=True)
 
     def forward(self, x: Tensor) -> Tensor:
-        # x: [n, 3 + n_cells]; the first 3 columns are the fixed head
+        # x: [n, 2 + n_cells]; the first 2 columns are the fixed head
         # inputs, the rest are raw cell states (0..6)
         n = x.shape[0]
-        inp = torch.cat([x[:, :3], self.embedding(x[:, 3:].long()).view(n, -1)], dim=1)
+        inp = torch.cat([x[:, :2], self.embedding(x[:, 2:].long()).view(n, -1)], dim=1)
         hidden = silu(self.hidden(inp))
         return tanh(self.output(hidden)).squeeze(-1)
 
@@ -226,7 +233,9 @@ class Net(Module):
         lr: float,
     ) -> tuple[float, float]:
         # data
-        positions, targets = self_play_records.make_training_data(self, td_lambda)
+        positions, targets, parity_signs = self_play_records.make_training_data(
+            self, td_lambda
+        )
         n = len(positions)
         perm = torch.randperm(n, device=positions.device)
         n_val = max(1, int(n * val_frac))
@@ -235,8 +244,10 @@ class Net(Module):
 
         val_positions = positions[val_idx]
         val_targets = targets[val_idx]
+        val_parity_signs = parity_signs[val_idx]
         train_positions = positions[train_idx]
         train_targets = targets[train_idx]
+        train_parity_signs = parity_signs[train_idx]
 
         # training widgets
         optimizer = torch.optim.AdamW(self.parameters(), lr=lr)
@@ -253,14 +264,16 @@ class Net(Module):
             for start in range(0, len(train_positions), batch):
                 batch_idx = order[start : start + batch]
                 optimizer.zero_grad()
-                predictions = self(train_positions[batch_idx])
+                predictions = (
+                    self(train_positions[batch_idx]) * train_parity_signs[batch_idx]
+                )
                 loss = loss_fn(predictions, train_targets[batch_idx])
                 loss.backward()
                 optimizer.step()
 
             with torch.no_grad():
                 self.eval()
-                predictions = self(val_positions)
+                predictions = self(val_positions) * val_parity_signs
                 mse = loss_fn(predictions, val_targets).item()
                 corr = torch.corrcoef(torch.stack([predictions, val_targets]))[
                     0, 1
@@ -532,7 +545,7 @@ def main() -> None:
     _ = p.add_argument(
         "--iters",
         type=int,
-        default=100,
+        default=20,
         help="max loop iterations",
     )
     _ = p.add_argument(
diff --git a/src/selfplay.c b/src/selfplay.c
index 5ffa6f6..42d82d2 100644
--- a/src/selfplay.c
+++ b/src/selfplay.c
@@ -38,10 +38,11 @@
 
 static tak_state_p state;
 
-static float row_head[MAX_PLY][3];
+static float row_head[MAX_PLY][2];
 static uint8_t row_cells[MAX_PLY][MAX_CELLS];
 static uint8_t row_term[MAX_PLY];
 static float row_terminal_value[MAX_PLY];
+static float row_parity_sign[MAX_PLY];
 static uint32_t row_game_id[MAX_PLY];
 static uint32_t n_rows;
 
@@ -90,6 +91,7 @@ static int play_game(const uint8_t min_ply) {
       memcpy(row_cells[n_rows], leaf->cells, sizeof(row_cells[n_rows]));
       row_term[n_rows] = leaf->terminal;
       row_terminal_value[n_rows] = leaf->terminal_value;
+      row_parity_sign[n_rows] = leaf->parity_sign;
       row_game_id[n_rows] = state->current_colour == C_BLACK ? game_id_black : game_id_white;
       n_rows++;
     }
@@ -114,12 +116,11 @@ static int play_game(const uint8_t min_ply) {
   const float R = win_value(state->won);
 
   for (uint32_t k = 0; k < n_rows; k++) {
-    fprintf(stdout, "%.8f,%.8f,%.8f,", row_head[k][0], row_head[k][1],
-            row_head[k][2]);
+    fprintf(stdout, "%.8f,%.8f,", row_head[k][0], row_head[k][1]);
     for (uint8_t c = 0; c < n_cells; c++)
       fprintf(stdout, "%u,", (unsigned)row_cells[k][c]);
-    fprintf(stdout, "%u,%.0f,%u,%.0f\n", row_game_id[k], R,
-            (unsigned)row_term[k], row_terminal_value[k]);
+    fprintf(stdout, "%u,%.0f,%u,%.0f,%.0f\n", row_game_id[k], R,
+            (unsigned)row_term[k], row_terminal_value[k], row_parity_sign[k]);
   }
 
   return (int)n_rows;
@@ -129,7 +130,7 @@ const char *usage =
   "selfplay, one game of asymmetric self-play training data.\n\n\
 Usage: selfplay <size> <depth> <min_ply>\n\n\
 Plays a single game of <size>x<size> self-play, with <min_ply> moves made randomly. Only the deterministic plies are written to stdout, one CSV row per recorded search, from ply <min_ply> onwards:\n\n\
-\tparity,white_flats,black_flats,<cell ids...>,game_id,R,terminal,leaf_R\n\n\
+\tmover_flats,opponent_flats,<cell ids...>,game_id,R,terminal,leaf_R,parity_sign\n\n\
 nn_train.py runs many of these in parallel to generate training data.\n\n\
 Copyright (C) 2026, tslil clingman\n\n\
 This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.";