From 220a4b79e4b6d742a9ca7c278a1962b870500fbc Mon Sep 17 00:00:00 2001 From: tslil Date: Thu, 10 Sep 2026 14:12:41 +0100 Subject: 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. --- src/nn_train.py | 37 +++++++++++++++++++++++++------------ 1 file changed, 25 insertions(+), 12 deletions(-) (limited to 'src/nn_train.py') 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_.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( -- cgit v1.3.1