# Penzai + Treescope: A Toolkit for Interpreting, Visualizing, and Editing Models As Data

arXiv: https://arxiv.org/abs/2408.00211 · HTML: https://mlsanity.com/paper/2408.00211/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 25 references that are in the map.
- First posted on arXiv 1 Aug 2024.

## Builds on (most cited first, 16 of 16)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Towards Best Practices of Activation Patching in Language Models: Metrics and Methods](https://mlsanity.com/paper/2309.16042/) (2023)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [The Language Interpretability Tool: Extensible, Interactive Visualizations and Analysis for NLP Models](https://mlsanity.com/paper/2008.05122/) (2020)
- [pyvene: A Library for Understanding and Improving PyTorch Models via Interventions](https://mlsanity.com/paper/2403.07809/) (2024)
- [Local Interpretations for Explainable Natural Language Processing: A Survey](https://mlsanity.com/paper/2103.11072/) (2021)
- [Named Tensor Notation](https://mlsanity.com/paper/2102.13196/) (2021)

## Built on by (most cited first, 5 of 5)

- [Softmax is not Enough (for Sharp Size Generalisation)](https://mlsanity.com/paper/2410.01104/) (2024)
- [Incorporating Hierarchical Semantics in Sparse Autoencoder Architectures](https://mlsanity.com/paper/2506.01197/) (2025)
- [Scaling sparse feature circuit finding for in-context learning](https://mlsanity.com/paper/2504.13756/) (2025)
- [Jasmine: A Simple, Performant and Scalable JAX-based World Modeling Codebase](https://mlsanity.com/paper/2510.27002/) (2025)
- [It's All Just Vectorization: einx, a Universal Notation for Tensor Operations](https://mlsanity.com/paper/2607.27987/) (2026)

## Abstract (arXiv)

Much of today's machine learning research involves interpreting, modifying or visualizing models after they are trained. I present Penzai, a neural network library designed to simplify model manipulation by representing models as simple data structures, and Treescope, an interactive pretty-printer and array visualizer that can visualize both model inputs/outputs and the models themselves. Penzai models are built using declarative combinators that expose the model forward pass in the structure of the model object itself, and use named axes to ensure each operation is semantically meaningful. With Penzai's tree-editing selector system, users can both insert and replace model components, allowing them to intervene on intermediate values or make other edits to the model structure. Users can then get immediate feedback by visualizing the modified model with Treescope. I describe the motivation and main features of Penzai and Treescope, and discuss how treating the model as data enables a variety of analyses and interventions to be implemented as data-structure transformations, without requiring model designers to add explicit hooks.