Source-linked AI summary
Graphical Design of Interpretable Architectures
Pietro Barbiero
TL;DR
Interpretable architectures lack a representation that combines a global architectural view with concrete tensor-level implementation. This paper introduces Penrose-style tensor diagrams that map directly to PyTorch einsum code and uses them to describe interpretable models and reveal that Steerling-8B is a residual model.
Problem
Existing equations lack an at-a-glance global view, while graphical models and flowcharts omit concrete tensor manipulations needed for implementation.
Method
The paper introduces Penrose-adapted tensor diagrams for interpretable architectures and applies them to interpretable-by-construction models and Steerling-8B.
Results
The notation provides global architectural insights and maps directly to PyTorch code, showing at a glance that Steerling-8B is a residual model.
Takeaways & Limitations
Tensor diagrams offer an expressive formalism for designing, comparing, and implementing interpretable architectures alongside coarser graphical tools.
Takeaways & Limitations
Tensor diagrams require ad hoc extensions for nonlinearities, masking, and discrete operations, and become unwieldy as frontier architectures grow more complex.
Abstract
from arXiv · showhide
Designing, implementing, and comparing interpretable architectures requires a formal language to represent them. The most common representations fall short in one of two ways. Symbolic equations give no global view of an architecture at a glance. Probabilistic graphical models and flowcharts do not describe actual tensor manipulations, thus hiding key insights and limiting reproducibility. To close this gap, we introduce a graphical notation for designing interpretable AI architectures, adapted from Penrose tensor notation. This graphical notation gives a global view of an architecture and maps one to one onto PyTorch einsum code. We first use this notation to describe architectures that are interpretable by construction, including concept bottlenecks, sparse probes, prototype networks, neural additive models, and mixtures of linear models. We then diagram the key architectural components of Steerling-8B, a frontier interpretable language model. The diagram yields global insights into the architecture (e.g., showing that Steerling is a residual model), a geometric interpretation of each individual operation, and a direct translation into 33 lines of PyTorch code.
1 Introduction
The paper introduces a Penrose-inspired graphical notation for tensor architectures that combines global architectural insight with concrete implementation specificity. It is designed for interpretable-by-construction models and maps one to one onto PyTorch einsum code.
- Motivation: Frontier AI models require representations that make their high-dimensional tensor manipulations understandable, comparable, and designable.Their many tensor operations make both intuitive and formal representations important for architectural understanding.
- Limitations of existing representations: Symbolic equations can obscure global architectural insights and make feature-wise tensor operations laborious and error-prone to identify.The passage notes that recognizing independent action on individual features requires working through the whole expression.
- Limitations of existing representations: Graphical models and flowcharts provide high-level architectural overviews but conceal concrete tensor manipulations, limiting reproducibility and obscuring innovations.These diagrams compensate for symbolic representations’ shortcomings but do not specify concrete implementation.
- Contribution: The proposed Penrose-adapted graphical notation provides global insights into interpretable AI architectures and maps one to one onto PyTorch einsum code.The notation is intended to bridge architectural understanding and implementation.
- Contribution: The notation describes interpretable-by-construction architectures, including concept bottlenecks, sparse probes, prototype networks, neural additive models, and mixtures of linear models.These examples demonstrate the notation’s intended scope for interpretable architecture design.
2 Graphical Einstein-inspired notation in PyTorch
This section introduces a graphical tensor notation adapted from Penrose diagrams for representing tensor operations and translating them directly into PyTorch einsum code. It defines tensors geometrically and explains how contractions and free indices make operations visually clear, formal, and unambiguous.
- Tensor representation: Tensors are represented as circles with legs, where each leg denotes one geometric dimension or array index.A tensor of order k has k indices; scalars, vectors, and matrices have orders 0, 1, and 2, respectively.
- Motivation: Penrose-style diagrams make complex tensor operations clearer and easier to parse than PyTorch einsum strings alone.PyTorch supports rearrange for reordering axes and einsum for combining tensors, but einsum is difficult to inspect at a glance.
- Graphical conventions: Shared labels connect tensor legs and specify contractions, while unconnected legs become the output’s free indices.Contraction multiplies entries for each shared-index value and sums over that index.
- Implementation: Each diagram translates directly into working PyTorch code through an einsum expression specifying input and output leg labels.The notation supports designing tensor operations visually before converting them into implementation.
3 Graphical design of simple neural models
This section demonstrates how the graphical notation designs increasingly complex neural models, from linear models and MLPs to self-attention, while mapping tensor operations directly to einsum code.
- Linear models: A linear model applies a matrix-vector product followed by an activation function and remains an important interpretable-machine-learning baseline.The input features are represented by x ∈ R^d, and learnable parameters by W ∈ R^(h×d).
- Multi-layer perceptrons: Stacking batched linear transformations produces a multi-layer perceptron with successive activated einsum contractions.The MLP computes z1, z2, through y by applying weight matrices W0, W1, through WL to batch tensor X.
- Self-attention: Self-attention projects an input token sequence into query, key, and value embeddings, scores token-pair relevance, and combines values using those scores.The projections are q = einsum('td,de->te', Z, W_q), k = einsum('pd,de->pe', Z, W_k), and v = einsum('pd,de->pe', Z, W_v).
- Self-attention: The self-attention computation decomposes into affinity scoring, softmax normalization, and a convex combination of value embeddings.The resulting probabilities determine how much each token uses each value embedding.
- Notation and implementation: Minimal diagrams retain only contraction indices, while PyTorch ellipsis notation generalizes the same operations across arbitrary batch dimensions.For example, einsum('btqij,btqjk->btqik', A, B) becomes einsum('...ij,...jk->...ik', A, B).
4 Graphical design of interpretable architectures
Interpretable architectures can be represented through three components—backbone, concept encoding, and concept composition—whose tensor operations define recurring design patterns. The section compares concept encoders and composition maps across interpretable model families.
- Interpretable architectures comprise a backbone mapping x to z, a concept encoding map mapping z to human-meaningful c, and a composition map mapping c to y.
- Concept encoding maps: Concept encoding maps include probes, sparse autoencoders, concept bottlenecks, concept embedding bottlenecks, and prototype-based maps, ordered by increasing tensor-manipulation complexity.
- Concept encoding maps: Sparse encoders and concept bottlenecks use the same tensor operation, but sparse probes recover semantics post-hoc whereas bottlenecks impose them through ground-truth concept annotations in the loss.
- Concept encoding maps: Prototype-based concept maps compare normalized embeddings with labelled prototypes, convert similarities into probabilities, and combine those probabilities with prototype labels to predict concepts.
- Concept composition maps: Concept composition is usually linear, while neural additive models transform each concept independently before summing, concept embedding predictors rescale embeddings, and mixtures weight multiple linear predictions.
5 Case study: frontier interpretable language models
The section applies the graphical notation to Steerling-8B, showing its residual structure, faithful code correspondence, and geometric interpretation of every operation.
- Case study: Steerling-8B is diagrammed as the largest publicly available interpretable-by-design language model at the time of writing.The case study focuses on Steerling-8B’s architecture.
- Architectural insight: The diagram reveals at a glance that Steerling-8B is a residual model, with gradients able to flow directly from output to the first input.This global view exposes a key architectural property without inspecting individual equations.
- Reproducibility: The notation maps one to one onto PyTorch einops and activation functions, supporting faithful reproduction of the model.With batch dimensions omitted and one attention head shown, the essential tensor manipulations require about 30 lines of code.
- Geometric interpretation: Each diagrammed operation has a direct geometric interpretation in linear algebra, allowing the computation to be understood as transformations in space.This geometric reading helps build intuition for the underlying computation.
6 Discussion
The paper extends tensor diagrams from post-hoc model analysis to designing interpretable-by-construction architectures, with a direct path from diagrams to PyTorch code. It demonstrates this approach through reusable correspondences, frontier-model modules, and complementary graphical formalisms.
- Novelty: The paper applies tensor notation to the forward problem of designing architectures interpretable by construction, with a direct, mechanical path from diagram to PyTorch code.This extends notation previously used for post-hoc analysis of trained models.
- Limitations: Tensor diagrams are exact for multilinear operations but require ad hoc extensions for nonlinearities, masking, and discrete operations such as top-k.Their implementation-level detail can also make diagrams unwieldy for full frontier architectures.
- Contributions: A “Rosetta stone” aligns diagrams, PyTorch code, geometric interpretations, and symbolic equations for common tensor manipulations in interpretability research.The alignment is intended to help readers from different backgrounds compare and understand these operations.
- Contributions: The authors also diagrammed and implemented key modules of a frontier interpretable-by-design language model, extending the method beyond standard tensor manipulations.This is presented as a harder demonstration case for the graphical notation.
- Implications: Because tensor diagrams are expressive, formal, and directly map onto PyTorch code, the authors suggest they could become a standard tool alongside probabilistic graphical models.The paper presents the two formalisms as complementary: tensor diagrams capture tensor manipulations, while coarser graphical tools remain useful for larger architectures.