chessformer_lens

A toolkit + visualizer library for mechanistic interpretability of transformer based chess models.

Download a chessformer engine (Leela Chess Zero or Maia) then pip install chessformer_lens

Paper: Increasing Skill Level Recruits Deeper Attention Layers in a Frozen Chess Transformer (Litman, 2026). Every figure and most of the analysis in this paper was made with this library.

Case study: Knight Fork Carrier Head

Download chessformer_lens locally:

Pip install:

python3 -m venv .venv && source .venv/bin/activate
pip install https://github.com/CSSLab/maia3/archive/1e13597c42d4858b7cfd7cfdae01e297263364b2.zip --quiet
pip install chessformer_lens

Clone install:

git clone https://github.com/chessformer-lens/chessformer_lens
cd chessformer_lens
pip install -r requirements.txt

Quick interpretability demo (n=10): attention layer 5 head 5 seems to be the carrier head for knight forks

In python:

# Gather ten positions with knights forking King and Queen, then ablate every
# head in the model to see which one is most causally linked to the output.
import chess
import matplotlib.pyplot as plt
from chessformer_lens import MaiaEngine
import chessformer_lens.interp_plot as ip

eng = MaiaEngine()

knight_forks = [
{"fen": '1r2k3/3qn3/3p4/p2P1Pp1/PpP3Np/1P3Q1P/3K1PP1/7R w - - 0 37', "move":'g4f6'},
{"fen": '5b2/6p1/3p2k1/2p3pn/7r/8/P2NKPQ1/R6R b - - 2 25', "move":'h5f4'},
{"fen": '6k1/5p1p/p3p1p1/1p2n3/1P1q4/P2p2P1/3Q1N1P/6K1 b - - 2 37', "move":'e5f3'},
{"fen": 'rn1q1knr/pb2p1b1/6p1/5pN1/1pPP3p/1P6/P1B2PPP/1RQ2RK1 w - - 2 20', "move":'g5e6'},
{"fen": '5r1k/3nrBp1/b1pp1n1p/p3q2P/Pp2PN2/3PB1Q1/1PP3P1/2KR3R w - - 1 22', "move":'f4g6'},
{"fen": '6k1/1pq2ppp/p7/3p4/1P1Nn3/P2QPP2/1B4Pb/7K b - - 1 26', "move":'e4f2'},
{"fen": '4r1k1/6b1/3p2pp/3Pnp2/4N2Q/6P1/P1q2P1P/4R1K1 b - - 0 26',"move": 'e5f3'},
{"fen": '5rk1/pp5p/2p1p1p1/6N1/3PQBP1/2q4P/P7/3n3K b - - 1 25', "move":'d1f2'},
{"fen": 'r2q1rk1/ppp3p1/3ppn1B/2b1p3/3nP3/3P2QN/PPP2PPP/RN3RK1 b - - 2 11',"move": 'd4e2'},
{"fen": 'r1b2rk1/1pp2p1p/4p1p1/2PqP3/p1nPN2B/P1PQ4/6PP/R4RK1 w - - 2 21',"move": 'e4f6'}]

for pos in knight_forks:
    ip.plot_move_report(eng, chess.Board(pos["fen"]), 2400, pos["move"])
    plt.show()

There is wonderful prior chess-interp work, for instance: McGrath et al. on AlphaZero concepts, Jenner et al. on lookahead in Leela, Karvonen on chess-GPT, but there is no general infrastructure for rigorous interpretability work or interactive visualization for these models.

Chessformer_lens is built to do both, especially inspired by Neel Nanda's transformer_lens library.

Chess is unusually suited for AI interpretability because it requires complex and structured reasoning in a compact domain, and it has a hierarchy of human concepts (squares -> threats -> tactics -> strategy) which can provide deep insight into how AI carves features.

Transformer based chess models ("chessformer") are a good lens for studying LLMs because outputs have a clear ground source for the evaluation of a move and so avoid the vagueness of language, because attention functions with clearly interpretable roles on a chess board, and because informative chessformer results are straightforward to scale and investigate in LLMs.

This repo's core is one engine with three frontends: engine.py is the interp core (model + hooks + logit lens + head ablation + list of values through depth). It does all of the analysis for:

• app.py
• interp_plot.py
• interp_widget.py

Users are encouraged to read the user guides for each of these modules which can be found at the top of the respective scripts. Image from paper produced with chessformer_lens

Quickstart in colab or notebook:

Recall that FEN is the modern notation for a chess position

!pip install -q https://github.com/CSSLab/maia3/archive/1e13597c42d4858b7cfd7cfdae01e297263364b2.zip
!pip install -q chessformer_lens

import chess
from chessformer_lens import MaiaEngine
import chessformer_lens.interp_widget as iw

eng = MaiaEngine()          
# Opera Game FEN right at legendary queen sacrifice 
fen = "4kb1r/p2n1ppp/4q3/4p1B1/4P3/1Q6/PPP2PPP/2KR4 w k - 0 16" 
# [or insert another FEN string]

try:
    board = chess.Board(fen)
except ValueError as e:
    print(f"Invalid FEN: {e}")   

#Interact with the attention widget
#click between different layers and heads in the widget
iw.attention_widget(eng, board, 1500,layer=4,head=3) 

The app

To me, the app is the pièce de résistance and usage should be rather intuitive, nonetheless, a detailed guide can be found in app_README.md.

Play a transformer-based chess bot ("chessformer") and watch its move policy, its attention (both semantic QKᵀ and unique geometric GAB), and — for any move you click — its logit through depth, the heads that carry it, and the neurons that carry it. The model loads on a background thread so the window opens instantly.

#Various ways to launch, note that higher parameter models have 16 and 32 heads per layer, and 512 and 1024 dim of the model (d_model):
chessformer_lens
chessformer_lens 23m               
chessformer_lens --model maia3-79m 
CHESSFORMER_MODEL=5m chessformer_lens
chessformer_lens bt4                # Leela Chess Zero BT4 — see below   

Leela Chess Zero BT4

Weights. Download the network file BT4-1024x15x32h-swa-6147500.pb.gz from lczero.org and pass its path. The engine reads lc0's .pb.gz directly; no conversion and no lc0 install needed.

Load. One alias table covers both families; neither is a default.

from chessformer_lens import LeelaEngine, load_engine
eng = LeelaEngine(checkpoint_path="BT4-1024x15x32h-swa-6147500.pb.gz")
eng = load_engine("BT4-1024x15x32h-swa-6147500.pb.gz", device="mps")
eng = load_engine("bt4")               # looks for weights/Leela_BT4_large_model.pb.gz

chessformer_lens BT4-1024x15x32h-swa-6147500.pb.gz      # the app

Differences with Maia. BT4 has no rating input. evaluate() gains "mlh", the moves-left head's estimate of plies to game end. Depth points are emb, a0 … m14: 31 of them, and there is no enc, because the heads read the last block directly.

Code layout: all found within the chessformer_lens package

• __init__.py — the package surface: MaiaEngine, LeelaEngine, load_engine, build_cfg, pick_device, attention_widget, gab_widget, __version__.
• engine.py — the interp core (model + hooks + logit lens + head ablation + list of values through depth): ChessformerEngine holds every method, MaiaEngine and LeelaEngine supply only the model, its tokens and its move vocabulary. No UI deps; imports cleanly in a notebook.
• interp_plot.py — matplotlib figures that the app uses
• interp_widget.py — the app's two interactive panels as a self-contained notebook cell or standalone HTML page
• piece_art.py — draws pieces.py's cburnett pieces in matplotlib from PNGs
• bridge.py — the game state + the JSON API the UI calls.
• ui.py — the whole interface (HTML/CSS/JS) as one string including guide to change template.
• pieces.py — SVG piece set as data URIs.
• app.py — app launcher (native window via pywebview); the chessformer_lens command is its main().

These come with concise but detailed docstrings.

General color guide

Blue and orange is used for diverging signed values; negative is blue, positive is orange. Viridis is used for magnitudes like attention weights. Green has to do with a specific chess move. Red has to do with undesirable things. The black and white halo lens denotes a query square.

Notes

Created by David Litman

Maia-3 comes from: Chessformer / Maia-3 (Monroe et al., ICLR 2026).

BT4 comes from the Leela Chess Zero project (lczero.org); the smolgen / attention-body design it uses is described in the Chessformer paper and in lc0's own writeups.

Both Maia-3 and Leela BT4 are supported. Both treat each square as a token—which allows for beautiful board readable attention patterns. Eventually other tokenization schemes will be tackled.

Please don't hesitate to give me feedback or thoughts by email or at my personal website. I hope for this to be a useful and intuitive tool for the community.

Citing

If chessformer_lens contributes to published work, a citation is very appreciated and helps others find it! Use GitHub's "Cite this repository" button, or DOI 10.5281/zenodo.21877655.