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

TL;DR

This paper introduces Penzai and Treescope, tools that treat models as data structures for easier interpretation, visualization, and editing, enabling flexible model analysis and modification without complex hooks.

Contribution

The paper presents Penzai, a neural network library with declarative combinators, and Treescope, an interactive visualization tool, facilitating model manipulation and understanding as data structures.

Findings

Penzai models use named axes for semantic clarity.

Users can insert or replace model components interactively.

Immediate visualization feedback supports model editing.

Abstract

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.