penzai

Introduction to Penzai: Unveiling the World of Miniature Model Manipulation

Penzai, inspired by an ancient Chinese art form, takes its name from the word盆 ("pen", tray) 栽 ("zai", planting), which involves crafting beautiful miniature landscapes and trees, a precursor to the Japanese bonsai. In the world of machine learning and neural networks, Penzai is akin to this art form - it's a versatile tool designed to provide clear and functional interpretations of complex models.

What is Penzai?

Penzai is a JAX library that allows users to write their models as easy-to-understand data structures known as pytrees. What's unique about Penzai is its focus on making it simple to work with models post-training. It's an excellent choice for researchers aiming to reverse-engineer models, dissect and modify model components, inspect internal activations, and troubleshoot architectures. Although Penzai is powerful for post-training tasks, it can also build and train models.

With Penzai, your neural networks come to life as legible, interactive configurations, making them easy to visualize and modify. Below is a glimpse of what a model looks like in Penzai:

Key Features of Penzai

Penzai is composed of several modular tools, each of which can operate independently:

1. Interactive Python Pretty-Printer - Treescope

• Treescope (pz.ts): A powerful tool initially part of Penzai, now available standalone, that serves as a replacement for standard Python renderers in platforms like IPython and Colab. It makes understanding deeply-nested JAX pytrees easier by supporting visualization of multi-dimensional arrays.

2. JAX Tree and Array Manipulation Utilities

• Selectors (pz.select): This tool functions as a multipurpose 'Swiss army knife' for Penzai, broadening JAX's .at[...].set(...) operations to any type-driven pytree traversals. It's perfect for complex rewrites and quickly modifying models and data structures.

• Named Axes (pz.nx): A lightweight system allowing ordinary JAX functions to be used with named axes, facilitating a smooth transition between named and positional programming styles without learning a new array API.

3. Declarative Combinator-Based Neural Network Library

• Penzai Neural Networks (pz.nn): An alternative to libraries like Flax or Keras, this library uses declarative combinators to expose the complete structure of a model's forward pass. It allows models to be represented as JAX pytrees, enhancing visibility and modifiability, with additional support for mutable variables.

4. Transformer Architectures for Research

• Transformer Models (penzai.models.transformer): This feature supports interpretability research by enabling model surgery and training dynamics exploration. It provides a modular Transformer implementation that can load pre-trained weights for various architectures.

Getting Started with Penzai

To start using Penzai, ensure JAX is installed on your system. You can follow the installation guide in the JAX documentation. Once JAX is set up, you can install Penzai using:

pip install penzai

To utilize Penzai's powerful functionalities in a Colab or IPython notebook, it's beneficial to configure Treescope as the default pretty printer:

import treescope
treescope.basic_interactive_setup(autovisualize_arrays=True)

Visualizing and Manipulating Models

Penzai makes it remarkably straightforward to initialize and visualize neural networks. For instance, here's how you could visualize a simple multi-layer perceptron (MLP):

from penzai.models import simple_mlp
mlp = simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(0),
    feature_sizes=[8, 32, 32, 8]
)

mlp  # Automatically visualized in a Colab/IPython notebook cell

Penzai also allows easy tracking and extraction of intermediate activations, facilitating in-depth model analysis:

@pz.pytree_dataclass
class AppendIntermediate(pz.nn.Layer):
  saved: pz.StateVariable[list[Any]]
  def __call__(self, x: Any, **unused_side_inputs) -> Any:
    self.saved.value = self.saved.value + [x]
    return x

var = pz.StateVariable(value=[], label="my_intermediates")

saving_model = (
    pz.select(mlp)
    .at_instances_of(pz.nn.Elementwise)
    .insert_after(AppendIntermediate(var))
)

output = saving_model(pz.nx.ones({"features": 8}))
intermediates = var.value

For those eager to master Penzai, there are tutorials such as "How to Think in Penzai" available in the Penzai documentation.

Conclusion

Penzai emerges as a pioneering tool in the realm of model manipulation and visualization, providing researchers and developers alike with the ability to dissect and understand models as easily as viewing a bonsai. Whether you're involved in complex model surgery or simpler post-training tasks, Penzai stands ready to support your needs with elegance and efficiency.

For any further reference, detailed documentation can be found here.

By adopting Penzai for research, users can cite the following article:

@article{johnson2024penzai,
    author={Daniel D. Johnson},
    title={{Penzai} + {Treescope}: A Toolkit for Interpreting, Visualizing, and Editing Models As Data},
    year={2024},
    journal={ICML 2024 Workshop on Mechanistic Interpretability}
}

Penzai isn't officially a Google product, but it’s undoubtedly a groundbreaking tool for anyone interested in model interpretability, offering a new lens to view and refine neural networks.