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.

1 files changed, 25 insertions, 12 deletions

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(