Sub-Project 3b: Supervised Pre-Training + AlphaZero Refinement

Status: Design — pending approval to proceed to implementation plan. Predecessor: SP3 — Chess on the AlphaZero framework (complete; framework code shipped). Successor: SP4 — Productionize (UCI engine, 2000+ ELO target).

1. Context

SP3 shipped the chess framework end-to-end: Chess(Game), AlphaZero 8×8×73 move encoding, StockfishOpponent, parallel self-play with NN-server + workers, Trainer dispatch, CLI, configs, tests. The framework works.

However, when the user launched the SP3 training run on M4 Pro MPS:

[iter 1/80]
  parallel self-play (iter 1): 8 workers × 100 games
  checkpoint: checkpoints/iter_0001.pt (13838.9s)

That is ~3.85 hours per iteration, projecting ~13 days for the full 80-iteration run. Pure AlphaZero self-play from random initialization is wall-clock infeasible on consumer hardware for chess.

The root cause is the "cold start" problem: an untrained network produces near-uniform priors, MCTS exploration is uninformed, games are very long (hit the 75-move rule because random moves rarely deliver checkmate), and the resulting (state, π, z) tuples are weak training signal.

2. Goal

Train a chess agent that meets the SP3 kill criterion (≥50% win rate vs Stockfish ELO=1500 over 100 games) in wall-clock days, not weeks on M4 Pro by:

  1. Bootstrapping the network via supervised pre-training on a corpus of strong-engine games. After pre-training, the policy head produces meaningful priors and the value head estimates positions reasonably — eliminating the cold start.
  2. Refining the pre-trained network with a short AlphaZero loop (10 iterations instead of 80). The MCTS-improved targets push the agent past the teacher's behavior-cloning ceiling.

Wall-clock target: ~16 hours total on M4 Pro (or ~10 hours if AZ refinement runs on Colab L4).

3. Non-goals

4. What's new vs what's reused

New code (~480 lines):

Reused unchanged:

5. Architecture: three phases producing three checkpoints

Phase 1: Corpus generation
  ┌────────────────────────────────────────┐
  │  scripts/generate_chess_corpus.py      │
  │  10 parallel workers, each playing     │
  │  Stockfish-vs-Stockfish games          │
  └─────────────────┬──────────────────────┘
                    │ writes
                    ▼
        data/chess_corpus/shard_NNNN.npz  (~300 shards, ~3.8GB)

Phase 2: Supervised pre-training
  ┌────────────────────────────────────────┐
  │  python -m alphazero pretrain \         │
  │    --corpus data/chess_corpus/ \        │
  │    --config configs/chess-pretrain.toml │
  └─────────────────┬──────────────────────┘
                    │ writes
                    ▼
              pretrained.pt

Phase 3: AlphaZero refinement
  ┌────────────────────────────────────────┐
  │  python -m alphazero train \            │
  │    --game chess \                       │
  │    --config configs/chess-refine.toml \ │
  │    --resume-from pretrained.pt          │
  └─────────────────┬──────────────────────┘
                    │ writes
                    ▼
        checkpoints/iter_0001.pt ... iter_0010.pt

Each phase is independently restartable. If pre-training converges and AZ refinement disappoints, you can re-run just Phase 3 with different hyperparameters from the same pretrained.pt.

6. Phase 1: Corpus generation

6.1 Architecture: embarrassingly parallel

Unlike parallel_selfplay.py (which needs a central NN-server to batch GPU inference), corpus generation has no neural network involved. Each worker is fully independent — it spawns its own Stockfish subprocess, plays games, and writes its own shards. No queues, no shared state, no IPC beyond initial work assignment.

Main process (orchestrator)
    │
    ├─ Worker 0 ── Stockfish ── plays 5,000 games ── writes shard_0_*.npz
    ├─ Worker 1 ── Stockfish ── plays 5,000 games ── writes shard_1_*.npz
    ├─ ...
    └─ Worker 9 ── Stockfish ── plays 5,000 games ── writes shard_9_*.npz

6.2 Per-game flow

For each game: 1. Initialize board = chess.Board() 2. Random first 4 plies (diversity) — sample uniformly from board.legal_moves, push to board 3. Stockfish plays both sides for the rest of the game. At each ply: - Call engine.play(board, chess.engine.Limit(time=0.05)) - Record (board.fen(), move, current_player) into game's position log - Push the move 4. When game ends (mate / draw / 75-move rule): - Compute outcome z ∈ {-1, 0, +1} from the final result - For each recorded position: encode board via Chess.encode(board), encode move via move_to_index(board, move), assign z from mover's POV (re-uses _assign_z logic) - Append (encoded_state, move_idx, z) to worker's shard buffer 5. When shard buffer reaches 10,000 positions, flush to disk as shard_W_S.npz (W = worker id, S = shard counter)

6.3 Worker count and Stockfish configuration

10 workers on M4 Pro (10 P-cores + 4 E-cores). All P-cores fully utilized; E-cores stay free for OS responsiveness. Each Stockfish subprocess is configured with:

engine.configure({
    "Threads": 1,    # critical — prevents thread oversubscription across workers
    "Hash": 16,      # 16MB transposition table, plenty for depth ~10-12
})

Without Threads=1, each Stockfish instance would spawn multiple internal threads → 10 workers × 4-8 threads each → ~50-80 threads contending for 14 cores → severe slowdown.

6.4 Shard format

Each shard is a numpy .npz archive:

np.savez_compressed(
    shard_path,
    states=np.array(states, dtype=np.int8),       # (N, 20, 8, 8) — encoded boards
    move_indices=np.array(moves, dtype=np.int32), # (N,) — flat action indices [0, 4672)
    outcomes=np.array(zs, dtype=np.int8),         # (N,) — z values in {-1, 0, +1}
)

int8 for states and outcomes keeps shards small (~13MB each compressed). Cast to float32 at training-batch-load time.

6.5 Targets and budget

7. Phase 2: Supervised pre-training

7.1 Loss formulation

Identical to the existing _train_step: cross-entropy on policy + MSE on value, equal weight.

policy_loss = F.cross_entropy(predicted_logits, target_move_indices)
value_loss = F.mse_loss(predicted_values, target_outcomes)
total_loss = policy_loss + value_loss

The only difference vs AlphaZero training is the source of the targets: - AZ self-play: target_policy = MCTS visit distribution π, target_value = game outcome z - Pre-training: target_policy = one-hot(Stockfish's move), target_value = game outcome z

Same loss function, same checkpoint format, same network class.

7.2 Held-out validation + early stopping

To detect overfitting, split corpus shards 95/5: train on 95%, validate on 5% (the last ~15 shards). After each epoch:

  1. Run forward pass on all validation positions (no gradient updates)
  2. Compute val_loss = policy_loss + value_loss averaged over validation set
  3. If val_loss improved → save checkpoint, reset patience counter
  4. If val_loss did not improve → increment patience counter
  5. If patience counter ≥ 2 → stop training (early stopping)

Maximum epochs: 3. Patience: 2. Typical behavior: train for 2-3 epochs, val_loss flattens, stop.

7.3 Learning rate schedule: linear warmup + cosine decay

For a 15M-parameter network, naive lr=1e-3 risks first-step parameter explosion (Adam's running gradient stats are unreliable at step 0). The schedule:

def lr_schedule(step, warmup_steps=500, total_steps=N, peak_lr=5e-4, end_lr=5e-5):
    if step < warmup_steps:
        return peak_lr * (step + 1) / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return end_lr + 0.5 * (peak_lr - end_lr) * (1 + math.cos(math.pi * progress))

7.4 Hyperparameters (configs/chess-pretrain.toml)

# Network — matches what chess-refine.toml will load
n_blocks = 15
n_channels = 192

# Pre-training
num_epochs = 3                # maximum; early stopping may halt sooner
batch_size = 256
peak_lr = 5e-4
end_lr = 5e-5
warmup_steps = 500
weight_decay = 1e-4
holdout_fraction = 0.05
early_stopping_patience = 2

# Misc
seed = 42
device = "auto"
corpus_dir = "data/chess_corpus"
output_checkpoint = "pretrained.pt"
log_dir = "runs_pretrain"

7.5 New CLI subcommand

python -m alphazero pretrain --config configs/chess-pretrain.toml

The subcommand reads corpus_dir from the config, splits shards into train/val, runs the pre-training loop, writes pretrained.pt at the path specified by output_checkpoint.

pretrained.pt is in the exact same format as Trainer's iter_NNNN.pt checkpoints (keys: iteration, config, best_net, candidate_net, optimizer). This lets Phase 3 load it via the existing --resume-from flag with zero changes to Trainer.

7.6 Wall-clock estimate

Note: this is much faster than the initial ~2 hour estimate because the corpus fits in disk cache after the first epoch and SGD on MPS is fast for this network size. Will refine the estimate after the first run.

8. Phase 3: AlphaZero refinement

8.1 No code changes — just config

Phase 3 is the existing AlphaZero loop with two changes:

  1. Start from pretrained.pt (via --resume-from) instead of random initialization
  2. Use reduced hyperparameters: 10 iterations, 50 MCTS sims per move

8.2 Hyperparameters (configs/chess-refine.toml)

# Network — must match pretrained.pt
n_blocks = 15
n_channels = 192

# MCTS — reduced from chess.toml
num_simulations = 50          # down from 200; pre-trained priors compensate
c_puct = 2.5
dirichlet_alpha = 0.3
dirichlet_weight = 0.25

# Self-play
games_per_iteration = 100
temperature_threshold = 30

# Training
training_steps_per_iteration = 2000
batch_size = 256
learning_rate = 1e-3
weight_decay = 1e-4

# Replay buffer — smaller since we don't need a huge ramp from cold start
replay_buffer_capacity = 200_000
min_buffer_size = 10_000

# Eval (vs Stockfish ELO=1500)
eval_interval = 2             # eval every 2 iterations
eval_games = 100

# Schedule — reduced from 80 iters
num_iterations = 10

# Parallel self-play
num_workers = 10              # match P-core count
inference_batch_size = 64

# Misc
seed = 42
device = "auto"
checkpoint_dir = "checkpoints_refine"
log_dir = "runs_refine"

8.3 Running it

python -m alphazero train \
  --game chess \
  --config configs/chess-refine.toml \
  --resume-from pretrained.pt

8.4 Wall-clock estimate

10 iterations × ~1 hour per iter (reduced sims + 15M-param net) = ~10 hours on M4 Pro. Or ~2 hours on Colab L4.

9. End-to-end wall-clock

Phase M4 Pro local Colab L4 (refinement only)
1. Corpus generation ~4.2 hr —
2. Supervised pre-training ~0.5 hr —
3. AZ refinement ~10 hr ~2 hr
Total ~14.7 hr ~6.7 hr

Either path fits comfortably in a day or two of wall-clock time, vs the 13-day projection for cold-start AZ.

10. Success criteria

  1. Corpus generation completes 50K games without worker crashes; produces ~300 shards totaling ~3.8GB
  2. Pre-training converges: validation loss strictly decreases for at least 1 epoch; final policy accuracy on held-out set ≥ 30% (top-1 move match with Stockfish) — a reasonable baseline given Stockfish has ~30+ plausible moves per position
  3. Pre-trained checkpoint is loadable by Trainer via --resume-from; first refinement iteration completes without crashes
  4. Final agent (after 10 refinement iters) wins ≥ 50% of 100 games vs Stockfish ELO=1500 (alternating colors)

11. Testing strategy

11.1 Unit tests (tests/test_chess_corpus.py, tests/test_supervised.py)

11.2 Integration tests

11.3 Invariants worth checking

12. Risks and mitigations

Risk Likelihood Impact Mitigation
Stockfish at 0.05s plays at lower ELO than expected (~2000 instead of 2400), capping student strength Medium Medium Test corpus quality by manually evaluating a few games before full run. Bump time_per_move to 0.1s if needed (doubles wall-clock)
Behavior cloning ceiling causes student to plateau before clearing 1500 ELO bar Low-Medium High AZ refinement should push past cloning ceiling. If not, increase num_iterations to 20+ or add more corpus games
15M-param net overfits 3M positions in 3 epochs Low Medium Early stopping with val loss catches this. Reduce to 2 epochs if val loss diverges in epoch 2
MPS out-of-memory at batch 256 with 15M params Low Medium Drop batch to 128 if OOM. Will slow training ~2× but still fits in budget
Shard format incompatibility (int8 cast issues) Low Low Unit test loads a shard and verifies round-trip integrity
Random-first-4-plies still produces too many duplicate game lines Low Medium If observed, increase to first 6 random plies; or sample first move from a fixed opening book
Corpus generation worker crashes mid-run, losing partial shards Low Low Each worker writes its own shards independently; restart just the failed worker(s) with adjusted game count

13. Future work (deferred)