Interactive visualization toolkit for transformer logit lens analysis.

• Try it now: Open the tutorial notebook in Colab to run Llama-8B on NDIF
• Using the code: Check out the Quick Start below and the Python API Reference
• Developing: See PLAN.md for the implementation roadmap and project structure
• Documentation: Full documentation is in the docs/ directory

The Logit Lens is an interpretability technique that decodes the internal hidden states of a language model into vocabulary probabilities. By applying the model's "unembedding" matrix (the output layer) to intermediate layers, we can see what the model "thinks" the next token is at each stage of its computation.

This technique reveals how predictions evolve layer by layer, offering insights into the model's reasoning process.

• Original Work: Interpreting GPT: the logit lens by nostalgebraist.

This library is designed to make logit lens analysis efficient and accessible, particularly for large models.

• Integration: It is built to work seamlessly with NNsight and NDIF, allowing for server-side reduction of data. This means you can visualize the internal dynamics of massive models (like Llama-70B) without downloading gigabytes of data.
• Education: This tool was created for the Neural Mechanics course to help students explore transformer internals.

• Efficient Data Collection: Server-side reduction via NDIF reduces bandwidth from ~800MB to ~400KB per prompt
• Interactive Visualization: Click to explore top-k predictions, pin tokens to track trajectories across layers
• Multi-Model Support: Works with Llama, GPT-2, GPT-J, Pythia, OLMo, and more
• Zero-Install Jupyter: Plain HTML output - no ipywidgets installation required

from nnsight import LanguageModel
from logitlenskit import collect_logit_lens, show_logit_lens

model = LanguageModel("meta-llama/Llama-3.1-8B", device_map="auto")
data = collect_logit_lens("The capital of France is", model, k=5, remote=True)

# Display widget (last row auto-pinned by default)
show_logit_lens(data, title="Llama-8B: The capital of France is")

# Customize with options
show_logit_lens(data, title="Custom view", dark_mode=True, pinned_rows=[])

var widget = LogitLensWidget("#container", data);

// With custom title
var widget = LogitLensWidget("#container", data, { title: "My Analysis" });

// Link column sizing between widgets
widget1.linkColumnsTo(widget2);

This package is intended for local development and eventual integration into nnsight.

# Python
cd python
python3 -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"

# JavaScript
cd js
npm install
npm test

Work in progress. See PLAN.md for implementation roadmap.

MIT