One attention head carries knight forks in a chess transformer, and here's a new toolkit that found it.

A new open-source toolkit, chessformer_lens, localizes a knight fork to a single attention head in the Maia-3 chess transformer, with ablating head 5 costing 2.78 logits versus ≤0.14 for other heads in the layer. The library, built for chess transformers and inspired by Neel Nanda's transformer_lens, provides interactive visualization and analysis tools, with a demo Colab notebook and pip-installable GitHub repository.

Localize knight forks to a single head in Maia-3 with logit-lens and per-head ablation. https://colab.research.google.com/drive/1YYZBd SZbjOscRXIqJUbfCaY7rRbEzWx?usp=sharing https://colab.research.google.com/drive/1YYZBd SZbjOscRXIqJUbfCaY7rRbEzWx?usp=sharing This is a demo of the library's capabilities so the sample size is tiny... much more analysis is done in an upcoming paper, for instance we mine hundreds of forks and show that ablating head 5 costs 2.78 logits whereas every other head in the layer costs ≤0.14 The quickest way to run and reproduce the image state is: python3 -m venv .venv && source .venv/bin/activatepip install git+https://github.com/CSSLab/maia3 Maia-3 not pip installable yetpip install "chessformer lens all " then run the app at 23m, set Elo to 2400, and input FEN: 4kb1r/p2n1ppp/4q3/4p1B1/4P3/1Q6/PPP2PPP/2KR4 w k - 0 16 chessformer lens 23m Try to use move microscope bottom middle window and ablate this head top right button to determine which head is most causally linked to carrying the stunning queen sacrifice. Bonus points if you can name this legendary game -------------------------------------------------------------------------------------------------------- The github repo is https://github.com/chessformer-lens/chessformer lens https://github.com/chessformer-lens/chessformer lens , and it is pip installable. This repo's core is one engine with three frontends : engine.py is the interp core model + hooks + logit lens + head ablation + GAB decomposition + logit/policy across depth . It does all of the analysis for three frontends: a native interactive app see above , a module to plot mechanistic analyses such as the residual stream across depth, and interactive widgets for notebooks for playful experiments for hypothesis generation. chessformer lens is built especially for chess transformers that map one token to one square—allowing for beautiful board readable attention patterns. 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 shared general infrastructure for rigorous interpretability work with interactive visualization of chessformer models. chessformer lens is built for this, especially inspired by Neel Nanda's transformer lens library. The Maia-3 model interpreter is completed; other tokenization schemes will be tackled. Please don't hesitate to give me feedback or thoughts by email or at davidlitman.com . I hope for this to be a useful and intuitive tool for the community