Symmetry-group augmentation

Deep dive

How it works

In: a sample (x, y) with x on an N×N grid and y a per-cell target (move, mask) or a scalar, plus a group G of grid permutations · Out: the sample (g·x, g·y), or a symmetrized prediction · State: none (permutation tables built once) · Cost: O(N²) per transform, one table lookup per cell; 8 forward passes for test-time averaging

1. Check the group. A square gives D4 = 4 rotations × {mirror, no mirror} = 8 elements. A rectangle has 4 (identity, 180° turn, two mirrors). Confirm the rules are invariant (below).
2. Build one permutation per element. p[i] = the cell where cell i goes. Outputs without a position (a pass move) map to themselves.
3. Transform input and target together. out[p[i]] = x[i] for every board plane and every per-cell target. Scalars (win probability, class) stay unchanged.
4. Pick how to use the group.
   • Random per sample: draw g each time a sample is loaded. Epoch cost unchanged; the usual choice.
   • Full expansion: emit all |G| versions. |G|× the compute per epoch, and the copies are strongly correlated.
5. Test-time averaging. For each g: run the model on g·x, apply g⁻¹ to the output, average over G. The result is exactly invariant (scalar) or equivariant (per-cell) even if the network is not. Cheaper variant: one random g per query, so errors average over a search tree.
6. Share caches. Canonicalize a position to its smallest hash over G so all versions hit one entry (Zobrist hashing).

Which symmetries are real.

• Go, Gomoku: no direction is special, all of D4 holds.
• Chess: pawns move one way, so rotations and the up-down flip are not symmetries. Only the left-right mirror (file a↔h) is, and only with no castling rights, since the two castling sides differ. Move targets mirror with it.
• Chess colour flip: a different symmetry: swap colours and side to move, mirror the ranks. Usually built into the input encoding (always show the side to move) rather than used as augmentation.
• AlphaZero uses no augmentation: the rules of chess and shogi are asymmetric.

Code

D4 as index permutations on an N×N grid, a policy permuted with the board (last slot: a pass move that stays put), and asserts for the group axioms, commutation, and test-time averaging.

/// Transform t in 0..8: mirror columns if t >= 4, then rotate 90 degrees clockwise (t % 4) times.
fn dest(n: usize, t: usize, mut r: usize, mut c: usize) -> (usize, usize) {
    if t >= 4 { c = n - 1 - c; }
    for _ in 0..t % 4 { (r, c) = (c, n - 1 - r); }
    (r, c)
}

/// p[i] = where cell i goes; `extra` trailing slots (a pass move) map to themselves.
fn perm(n: usize, t: usize, extra: usize) -> Vec<usize> {
    let mut p: Vec<usize> = (0..n * n)
        .map(|i| { let (r, c) = dest(n, t, i / n, i % n); r * n + c })
        .collect();
    p.extend(n * n..n * n + extra);
    p
}

/// out[p[i]] = x[i]: boards, policy targets and policy outputs all use this.
fn apply<T: Copy + Default>(x: &[T], p: &[usize]) -> Vec<T> {
    let mut out = vec![T::default(); x.len()];
    for (i, &v) in x.iter().enumerate() { out[p[i]] = v; }
    out
}
fn compose(a: &[usize], b: &[usize]) -> Vec<usize> { a.iter().map(|&i| b[i]).collect() } // a, then b
fn invert(p: &[usize]) -> Vec<usize> {
    let mut q = vec![0; p.len()];
    for (i, &j) in p.iter().enumerate() { q[j] = i; }
    q
}

/// Toy policy: a cell scores the stones on its 4 neighbours (symmetric, so equivariant).
/// `bias` > 1 over-weights the right neighbour and breaks the symmetry. Last slot = pass.
fn model(b: &[u8], n: usize, bias: f32) -> Vec<f32> {
    let at = |r: i32, c: i32| if r < 0 || c < 0 || r >= n as i32 || c >= n as i32 { 0.0 }
        else { b[r as usize * n + c as usize] as f32 };
    let mut pol = vec![0.0; n * n + 1];
    for i in 0..n * n {
        let (r, c) = ((i / n) as i32, (i % n) as i32);
        pol[i] = at(r - 1, c) + at(r + 1, c) + at(r, c - 1) + bias * at(r, c + 1);
    }
    pol[n * n] = b.iter().map(|&v| v as f32).sum();
    pol
}

/// Test-time averaging: run each transformed board, map outputs back with the inverse, average.
fn tta(b: &[u8], n: usize, bias: f32) -> Vec<f32> {
    let mut sum = vec![0.0f32; n * n + 1];
    for t in 0..8 {
        let out = model(&apply(b, &perm(n, t, 0)), n, bias);
        for (s, v) in sum.iter_mut().zip(apply(&out, &invert(&perm(n, t, 1)))) { *s += v / 8.0; }
    }
    sum
}

fn main() {
    let n = 5;
    let all: Vec<Vec<usize>> = (0..8).map(|t| perm(n, t, 1)).collect();
    let id: Vec<usize> = (0..n * n + 1).collect();
    for i in 0..8 { for j in 0..i { assert_ne!(all[i], all[j]); } } // 8 distinct
    for a in &all {
        for b in &all { assert!(all.contains(&compose(a, b))); }    // closure
        assert!(all.contains(&invert(a)) && compose(a, &invert(a)) == id); // inverse
    }
    let mut board = vec![0u8; n * n];
    for &i in &[1, 7, 8, 12, 19] { board[i] = 1 }
    let mv = 13; // played move, as a one-hot policy target
    let mut tgt = vec![0.0f32; n * n + 1]; tgt[mv] = 1.0;
    for p in &all {
        let tb = apply(&board, &p[..n * n]);
        assert_eq!(apply(&tgt, p)[p[mv]], 1.0); // the target moves with the board
        assert_eq!(model(&tb, n, 1.0), apply(&model(&board, n, 1.0), p)); // policy commutes
    }
    let p1 = &all[1];
    assert_ne!(model(&apply(&board, &p1[..n * n]), n, 2.0), apply(&model(&board, n, 2.0), p1));
    let (avg, avg_t) = (tta(&board, n, 2.0), tta(&apply(&board, &all[3][..n * n]), n, 2.0));
    for (a, b) in avg_t.iter().zip(apply(&avg, &all[3])) { assert!((a - b).abs() < 1e-4); }
    println!("ok: group of 8, policy commutes, tta is equivariant");
}

Compiled and run with stable Rust (std only); prints ok: group of 8, policy commutes, tta is equivariant.

Complexity and constants

• One transform is O(N²): a lookup and a store per cell. The tables are 8 × (N²+1) indices; for 19×19 that is 2,896 entries, about 5.8 KB as u16, resident in L1.
• Test-time averaging: 8 forward passes; as one batch of 8 the latency is that of a larger batch.
• Full expansion multiplies epoch cost by |G|; random-per-sample leaves it unchanged.
• Orbit size = |G| / |stabilizer|: a position fixed by some transforms has fewer than 8 distinct images (the empty board has 1), so full expansion duplicates samples.

Efficiency

• Cost: a permuted copy, memory-bound and small next to a network forward pass. The real cost is test-time averaging (|G| forward passes) or full expansion (|G|× samples per epoch). Bound: memory bandwidth for the copy, compute for the model.
• Levers:
  • random transform per sample instead of full expansion — all orientations seen over many epochs at 1× epoch cost — default for training
  • one random transform per search evaluation instead of 8 — 1× evaluation cost; the tree averages the errors (Silver et al., arXiv 2017) — tree search with a network
  • batch the |G| test-time copies — latency of one larger batch — inference
  • canonical form (smallest hash over G) as cache key — up to |G|× more hits on symmetric positions — transposition tables, feature caches
  • bake the symmetry into the layers — no augmentation and no |G| passes, at the price of more complex layers (Cohen and Welling 2016) — known fixed group
• Naive → standard → state of the art:
  • ignore the symmetry · each orientation learned separately
  • random-per-sample augmentation, optional test-time averaging · 1× training, up to 8× inference
  • equivariant layers · one set of weights for all orientations

Implementation notes

• Transform the whole input stack with one permutation. Planes that are not symmetric (a "direction of play" plane, absolute coordinates, a positional encoding) must be transformed in meaning too, or dropped.
• Do it in the data-loader workers, or as a gather by an index tensor on the GPU; draw g from the sample's seeded RNG for reproducible runs.
• No augmentation in validation; map outputs back with g⁻¹ before averaging, and average logits or probabilities consistently.

Pitfalls

• Transforming the board but not the policy target: the model trains on wrong moves. The commutation assert above catches it.
• Treating a nearly symmetric domain as symmetric: chess with castling rights, images with text, boards with a side-dependent rule.
• Inputs that reveal the orientation (coordinate planes): the net learns to use them and the symmetry is lost.
• Splitting train and test after full expansion: the 8 copies of one position leak across the split. Split first.

Variants and alternatives

In production

• Board games: AlphaGo and AlphaGo Zero trained on the 8 symmetries and evaluated tree-search positions under a random rotation or reflection; AlphaZero dropped both for chess and shogi (AlphaZero self-play, Monte Carlo tree search). Leela Zero (Ensemble::RANDOM_SYMMETRY) and KataGo (nnRandomize) evaluate under a random orientation.
• Medical imaging: rotation-equivariant CNNs for histopathology (Veeling et al. 2018).
• Molecules and point clouds: e3nn builds rotation, translation and reflection equivariance into the model (Point-cloud processing).
• Aerial imagery: augmentation libraries ship the 8 square symmetries as one transform.

Implementations

Sources

• David Silver et al., Mastering the game of Go with deep neural networks and tree search, Nature 529, 2016. paper
• David Silver et al., Mastering the game of Go without human knowledge, Nature 550, 2017. paper
• David Silver et al., Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm, arXiv 2017 (states which systems used the 8 symmetries). paper
• Taco Cohen, Max Welling, Group Equivariant Convolutional Networks, ICML 2016. paper
• Bastiaan S. Veeling et al., Rotation Equivariant CNNs for Digital Pathology, MICCAI 2018. paper
• Mario Geiger, Tess Smidt, e3nn: Euclidean Neural Networks, arXiv 2022. paper · code