Transformer policy-value network

Deep dive

How it works

Builds on: Transformer block, Embeddings, Attention variants, Rotary position embedding (RoPE), FlashAttention, Weight quantization. The legal mask comes from legal move generation; the consumer is MCTS.

Parts:

• input encoding — per square: piece one-hots for the current and recent positions, plus castling, en-passant, rule-50 and repetition features; the board is flipped so the side to move always looks the same
• token embedding — linear projection to width d, plus a learned [64, d] position embedding
• encoder stack — L blocks of multi-head self-attention and a small MLP with residuals and LayerNorm; optional learned attention bias per square pair
• policy head — query and key projections of the final tokens; logit(from, to) = q_from · k_to / √d
• value head — pool the tokens, project to three logits (win, draw, loss)
• legal mask — boolean [4096] from the move generator; illegal pairs get probability 0

Operations: encode, embed, encode layers, policy, value, search use (steps below).

1. Encode. Lc0 uses 64 tokens whose order is flipped with the side to move; each token's input has 8 one-hot vectors of dimension 12 (the piece on that square for the current and last 7 plies), castling and en-passant information, the rule-50 count and repetition flags (Lc0 blog).
2. Embed. Linear projection of the input plus a learned position embedding (Lc0 blog).
3. Layers. Attention lets every square see every other square in one layer; a convolution needs one layer per step across the board. Lc0 reports post-LayerNorm with DeepNet initialisation, mish in the MLP, and small head depths working best (Lc0 blog).
4. Policy. logits[f*64+t] = q_f · k_t / √d, set illegal entries to −∞, softmax. Promotions need extra outputs, for example one logit per promotion piece on moves reaching the last rank.
5. Value. Softmax over three logits; the scalar for search is v = P(win) − P(loss), from the side to move's view.
6. Search and training. At a leaf the policy becomes the child priors P of PUCT and v is backed up; training minimises cross-entropy of the policy against a target distribution (search visit counts in AlphaZero-style self-play, or engine moves) plus a loss on the value (training loop).

64 square features ─► embed + pos ─► [ attention ─► MLP ] × L ─┬─► q·kᵀ (64×64) ─ legal mask ─ softmax ─► P(move)
                                                              └─► mean pool ─► 3 logits ─ softmax ─► P(win,draw,loss)

Code

/// y[n,o] = x[n,i] · w[i,o] + b[o]; all row-major.
fn linear(x: &[f32], n: usize, i: usize, w: &[f32], b: &[f32], o: usize) -> Vec<f32> {
    let mut y = vec![0.0; n * o];
    for r in 0..n {
        let yr = &mut y[r * o..(r + 1) * o];
        yr.copy_from_slice(b);
        for k in 0..i {
            let xv = x[r * i + k];
            for c in 0..o { yr[c] += xv * w[k * o + c]; }
        }
    }
    y
}

fn layer_norm(x: &mut [f32], d: usize, g: &[f32], b: &[f32]) {
    for row in x.chunks_mut(d) {
        let mean = row.iter().sum::<f32>() / d as f32;
        let var = row.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / d as f32;
        let inv = 1.0 / (var + 1e-5).sqrt();
        for (j, v) in row.iter_mut().enumerate() { *v = (*v - mean) * inv * g[j] + b[j]; }
    }
}

fn gelu(x: f32) -> f32 { 0.5 * x * (1.0 + (0.797_884_6 * (x + 0.044715 * x * x * x)).tanh()) }

/// Softmax over the entries where `keep` is true; others get probability 0.
fn masked_softmax(logits: &[f32], keep: &[bool]) -> Vec<f32> {
    let m = logits.iter().zip(keep).filter(|(_, k)| **k).map(|(l, _)| *l).fold(f32::NEG_INFINITY, f32::max);
    let mut p: Vec<f32> = logits.iter().zip(keep).map(|(&l, &k)| if k { (l - m).exp() } else { 0.0 }).collect();
    let s: f32 = p.iter().sum();
    p.iter_mut().for_each(|v| *v /= s);
    p
}

pub struct Layer {
    wqkv: Vec<f32>, bqkv: Vec<f32>, wo: Vec<f32>, bo: Vec<f32>,
    w1: Vec<f32>, b1: Vec<f32>, w2: Vec<f32>, b2: Vec<f32>,
    ln1: (Vec<f32>, Vec<f32>), ln2: (Vec<f32>, Vec<f32>),
}

pub struct Net {
    pub d: usize, pub heads: usize, pub ff: usize,
    tok: Vec<f32>,  // [13, d]: empty + 6 piece types x 2 colours
    pos: Vec<f32>,  // [64, d] learned position embedding
    layers: Vec<Layer>,
    wq: Vec<f32>, wk: Vec<f32>, // policy head projections [d, d]
    wv: Vec<f32>, bv: Vec<f32>, // value head [d, 3]
}

/// Tiny deterministic initialiser (LCG, uniform in ±scale) so the example needs no crates.
struct Rng(u64);
impl Rng {
    fn vec(&mut self, n: usize, scale: f32) -> Vec<f32> {
        (0..n).map(|_| {
            self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            ((self.0 >> 40) as f32 / (1u64 << 24) as f32 * 2.0 - 1.0) * scale
        }).collect()
    }
}

impl Net {
    pub fn new(d: usize, heads: usize, ff: usize, n_layers: usize, seed: u64) -> Net {
        let mut r = Rng(seed);
        let s = 1.0 / (d as f32).sqrt();
        let layers = (0..n_layers).map(|_| Layer {
            wqkv: r.vec(d * 3 * d, s), bqkv: vec![0.0; 3 * d], wo: r.vec(d * d, s), bo: vec![0.0; d],
            w1: r.vec(d * ff, s), b1: vec![0.0; ff], w2: r.vec(ff * d, 1.0 / (ff as f32).sqrt()), b2: vec![0.0; d],
            ln1: (vec![1.0; d], vec![0.0; d]), ln2: (vec![1.0; d], vec![0.0; d]),
        }).collect();
        Net { d, heads, ff, tok: r.vec(13 * d, 0.5), pos: r.vec(64 * d, 0.5), layers,
              wq: r.vec(d * d, s), wk: r.vec(d * d, s), wv: r.vec(d * 3, s), bv: vec![0.0; 3] }
    }

    /// squares[s] in 0..13 (0 = empty). Returns (policy over 64*64 from-to pairs, [win, draw, loss]).
    /// `legal[from*64+to]` marks the legal moves; everything else gets probability exactly 0.
    pub fn forward(&self, squares: &[u8; 64], legal: &[bool]) -> (Vec<f32>, [f32; 3]) {
        let (d, h, n) = (self.d, self.heads, 64);
        let dh = d / h;
        let mut x = vec![0.0f32; n * d];
        for s in 0..n {
            for j in 0..d { x[s * d + j] = self.tok[squares[s] as usize * d + j] + self.pos[s * d + j]; }
        }
        for l in &self.layers {
            let qkv = linear(&x, n, d, &l.wqkv, &l.bqkv, 3 * d);
            let mut att = vec![0.0f32; n * d];
            for hd in 0..h {
                let off = hd * dh;
                for qi in 0..n {
                    let mut sc = [0.0f32; 64];
                    for ki in 0..n {
                        let dot: f32 = (0..dh).map(|j| qkv[qi * 3 * d + off + j] * qkv[ki * 3 * d + d + off + j]).sum();
                        sc[ki] = dot / (dh as f32).sqrt();
                    }
                    let p = masked_softmax(&sc, &[true; 64]); // no padding: every square attends to every square
                    for ki in 0..n {
                        for j in 0..dh { att[qi * d + off + j] += p[ki] * qkv[ki * 3 * d + 2 * d + off + j]; }
                    }
                }
            }
            let proj = linear(&att, n, d, &l.wo, &l.bo, d);
            for i in 0..n * d { x[i] += proj[i]; }
            layer_norm(&mut x, d, &l.ln1.0, &l.ln1.1); // post-LN, as in Lc0's encoder
            let mut hid = linear(&x, n, d, &l.w1, &l.b1, self.ff);
            hid.iter_mut().for_each(|v| *v = gelu(*v));
            let out = linear(&hid, n, self.ff, &l.w2, &l.b2, d);
            for i in 0..n * d { x[i] += out[i]; }
            layer_norm(&mut x, d, &l.ln2.0, &l.ln2.1);
        }
        // policy head: logit(from, to) = q_from · k_to / sqrt(d)
        let zero = vec![0.0; d];
        let q = linear(&x, n, d, &self.wq, &zero, d);
        let k = linear(&x, n, d, &self.wk, &zero, d);
        let mut logits = vec![0.0f32; n * n];
        for f in 0..n { for t in 0..n {
            logits[f * n + t] = (0..d).map(|j| q[f * d + j] * k[t * d + j]).sum::<f32>() / (d as f32).sqrt();
        } }
        let policy = masked_softmax(&logits, legal);
        // value head: mean over squares, then 3 logits (win/draw/loss) from the side to move's view
        let pooled: Vec<f32> = (0..d).map(|j| (0..n).map(|s| x[s * d + j]).sum::<f32>() / n as f32).collect();
        let wdl = masked_softmax(&linear(&pooled, 1, d, &self.wv, &self.bv, 3), &[true; 3]);
        (policy, [wdl[0], wdl[1], wdl[2]])
    }
}

/// Scalar value for search: expected score in [-1, 1].
pub fn value(wdl: [f32; 3]) -> f32 { wdl[0] - wdl[2] }

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn heads_and_masking() {
        let net = Net::new(16, 4, 32, 2, 7);
        let mut sq = [0u8; 64];
        sq[4] = 6; sq[60] = 12; sq[0] = 4; // white king e1, black king e8, white rook a1
        let mut legal = vec![false; 64 * 64];
        for &(f, t) in &[(0usize, 1usize), (0, 8), (4, 3), (4, 5)] { legal[f * 64 + t] = true; }
        let (p, wdl) = net.forward(&sq, &legal);
        assert!((p.iter().sum::<f32>() - 1.0).abs() < 1e-5);
        assert!(p.iter().zip(&legal).all(|(&pi, &l)| l || pi == 0.0));
        assert!(p.iter().zip(&legal).filter(|(_, l)| **l).all(|(pi, _)| *pi > 0.0));
        assert!((wdl.iter().sum::<f32>() - 1.0).abs() < 1e-5);
        assert!(value(wdl).abs() <= 1.0);
        // different position, different output
        sq[0] = 0;
        assert_ne!(net.forward(&sq, &legal).1, wdl);
    }
}

The weights are random, so the test checks structure only: probabilities sum to 1, illegal moves get exactly 0, legal ones are positive, and the value triple is a distribution.

Complexity and constants

• Per layer parameters: 4d² (attention) + 2·d·ff (MLP). Per forward pass, with 64 tokens: about 2 × params × 64 FLOPs for the dense layers plus 4·64²·d per layer for attention scores and mixing (derived).
• Example (derived): d = 256, ff = 256, L = 10 gives about 3.9 M parameters and order of 0.5 GFLOP per position.
• Attention scores are 64×64 per head, small enough to stay in cache; unlike text there is no long sequence to make attention dominate.
• Policy head output 4,096 floats per position; input as bit-packed planes is much smaller than the float tensor.

Efficiency

• Cost: order of 0.5 GFLOP per position for a 10-layer d=256 net (derived above); on a GPU the call is latency-bound at small batches and compute-bound at large ones, so throughput is batch size over batch latency.
• Levers:
  • batch many leaves per call — GPU stays busy — see batched MCTS inference
  • narrow MLP (ff = d instead of 4d) — fewer FLOPs for little loss in chess (Lc0 blog) — when training from scratch
  • half precision or int8 weights — less memory traffic, more throughput — see weight quantization; validate policy agreement after
  • cache evaluations by position hash — skips repeated positions — transpositions in search
  • fused attention kernels — less memory traffic — mostly useful for long sequences; with 64 tokens the gain is small (order of)
• Naive → standard → state of the art:
  • convolutional residual tower (AlphaZero) · local receptive field, depth needed to cross the board
  • plain transformer on square tokens · global attention in every layer
  • transformer with learned square-pair attention bias and dynamic bias (Lc0's smolgen) · BT4, about 270 Elo above the best convolutional net in raw policy (Lc0 blog)

Implementation notes

• Keep the legal mask as a dense [bool; 4096] per position and apply it as −∞ before the softmax; computing the softmax first and zeroing afterwards leaves a distribution that does not sum to 1.
• Flip the board for Black and flip the move indices back when reading the policy; test with a position and its mirror.
• Subtract the max before exp, and run the softmax in fp32 even when the layers run in fp16.
• Parameter layout: one [d, 3d] matrix for Q, K and V so a single matmul feeds all heads.
• For 2D structure, a learned [64, 64, h] attention bias or learned position embeddings suit the fixed 8×8 grid; 1D rotary embeddings were designed for sequences.
• Fixed batch sizes (padded) let an inference engine compile once.

Pitfalls

• Value from the wrong point of view: the search backs up with a sign flip per ply; use the side to move's view throughout.
• Forgetting repetition and rule-50 inputs: the network cannot see draws that the rules would declare.
• Post-LayerNorm transformers are unstable at initialisation without learning-rate warm-up (Xiong et al. 2020); use warm-up or Pre-LN.
• Training on legal-masked logits but searching with unmasked ones (or the reverse): priors then differ between training and play.
• Mixing up from/to order or promotion encoding between the move generator and the head: games look fine until a promotion appears.

Variants and alternatives

In production

• Leela Chess Zero: transformer networks; BT4 is the strongest (called Leela-CF in the ICLR 2026 paper). At the TCEC Superfinal Season 23 it lost 19 game pairs and won 2; in Seasons 24 and 25 it lost 9 and won 5 (Lc0 blog).
• AlphaZero-style systems for Go, chess and shogi use the same policy-plus-value interface with convolutional towers.

Implementations

Sources

• Leela Chess Zero team, Transformer Progress, lczero.org blog, 2024. page
• Chessformer: A Unified Architecture for Chess Modeling, ICLR 2026. OpenReview
• David Silver et al., Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm, arXiv:1712.01815, 2017. paper
• Ashish Vaswani et al., Attention Is All You Need, NeurIPS 2017. paper
• Ruibin Xiong et al., On Layer Normalization in the Transformer Architecture, ICML 2020. paper
• Anian Ruoss et al., Amortized Planning with Large-Scale Transformers: A Case Study on Chess, arXiv:2402.04494, 2024. paper