tl;dr

I trained OthelloGPT-small, a 600k param model that predicts legal next moves in 6x6 Othello with 99.3% accuracy
I trained attention SAEs and transcoders on every module output with 94%+ loss recovered
I used the learned features and the static model weights to generate a computational graph of the model based on a linear upstream relevance metric
I analysed the computational graph for u_A1 and found decision trees that built up compound board state patterns as well as an inhibitory feature that seemed to perform a non-linear operation
This provides evidence that patch-free methods can produce interpretable computational graphs at scale

Setup

OthelloGPT-small

OthelloGPT is a language model that predicts the next (legal) move in strings representing Othello games. It was first proposed in 2022 by Kenneth Li et al. Some great posts to read on the topic are:

Kenneth Li, "Do Large Language Models learn world models or just surface statistics?", The Gradient, 2023.
1. Non-linear probes can extract the board state at each text position
2. The probes can be used to instruct interventions on the residual stream that change the model's internal board state representation, evidenced by corresponding changes in the model's predictions
Neel Nanda, "Actually, Othello-GPT Has A Linear Emergent World Model", 2023
1. Linear probes can extract the board state when transformed from (black/empty/white) to (theirs/empty/mine)
2. The probes correspond to directions that can be linearly rescaled to perform interventions
3. A bunch of other mech interp techniques can be applied: DLA, weight probing, activation patching, max activating datasets, etc.
Zhengfu He et al., "Dictionary Learning Improves Patch-Free Circuit Discovery in Mechanistic Interpretability: A Case Study on Othello-GPT", 2024
1. SAEs can be applied to transformer module outputs that rediscover all previous probes in an unsupervised manner
2. Approximate latent attributions can be used to find modular circuits

I trained my own OthelloGPT model on the smaller 6x6 variant of the game, on the premise that a quadratic reduction in vocab size and context length could reduce inference times by $O (n^{6})$ while maintaining sufficient task complexity. I managed to get the model size down to 600k params (3 layers, 128 dims) with 99.3% top-1 accuracy in predicting legal next moves, compared to the original model with 25m params (8 layers, 512 dims) and ...

Posts

Wikitag Contributions

Comments