09catho/axon

AXON is a tool for watching a language model think in real time. It runs a small version of GPT-2, the language model from OpenAI, and as the model generates each token of text, it shows you which internal concepts inside the model are firing. The result is a live, interactive 3D graph in your browser where each glowing point is a concept and connections appear between concepts that activate together. The technique behind this comes from an area called mechanistic interpretability. A Sparse Autoencoder, or SAE, is a separate small neural network that has been trained to break a language model's internal state into a long list of features that are each meant to stand for one human-readable concept. AXON hooks into GPT-2 at layer 8, takes the model's hidden state for each new token, passes it through an SAE from a library called SAELens, picks the top features, and looks up labels for them from a public website called Neuronpedia. It then streams the token, the features, and their activation strengths to the browser over a WebSocket. The interface, built with Three.js and 3d-force-graph, draws features as nodes and co-activating features as edges. Node size and brightness grow with activation strength and slowly fade. Clicking a node flies the camera to it and opens an inspector showing the feature index, label, current and peak activation, neighbors, and a link to its Neuronpedia page. A sidebar exposes generation knobs like temperature, top-K, top-P, repetition penalty, max tokens, feature top-K, and an activation threshold. To run it you need Python 3.11 and at least 8 GB of RAM, with a GPU strongly recommended. You clone the repo, create a virtual environment, install requirements, and run server.py, then open http://127.0.0.1:8000. The first run downloads GPT-2 and the SAE from HuggingFace. There are instructions for installing the CUDA build of PyTorch, a timing table showing roughly 35 ms per token on an RTX 4050 down to 900 ms on a CPU, and a guide for swapping in other supported models like GPT-2 medium, large, XL, Pythia, Gemma, Mistral, and LLaMA-2 by editing a few constants at the top of server.py.