Technical Summary: Graphical Design of Interpretable Architectures

Problem Statement

Designing, implementing, and comparing interpretable AI architectures currently lacks a unified formal language. Existing representations suffer from two primary deficiencies:

1. Symbolic Equations: While mathematically precise, they obscure the global structure of an architecture. Complex tensor manipulations (e.g., feature independence) are hidden within dense algebraic expressions, requiring significant effort to parse and prone to error.
2. Probabilistic Graphical Models (PGMs) and Flowcharts: These provide high-level overviews but fail to describe actual tensor manipulations. Consequently, they hide key implementation insights, limit reproducibility, and do not specify how architectures should be concretely implemented.

There is a need for a representation that offers an immediate global view of an architecture while maintaining a direct, one-to-one correspondence with executable code (specifically PyTorch).

Methodology

The author introduces a graphical notation adapted from Penrose tensor notation (originally from physics) to design interpretable AI architectures. This notation bridges the gap between abstract mathematical concepts and concrete implementation.

Core Mechanics

• Visual Representation: Tensors are depicted as circles with "legs," where each leg represents a geometric dimension (array index).
• Operations:
  • Contraction: Connecting legs with shared labels represents summation over that index (multiplication followed by summation), effectively "canceling" the leg.
  • Free Legs: Unconnected legs become the indices of the resulting tensor.
• Code Mapping: The notation maps one-to-one onto PyTorch einsum code. A diagram can be directly translated into einsum('legs_A, legs_B -> legs_out', A, B), allowing for the design of operations as diagrams followed by immediate code generation.
• Scope: The author applies this notation to models that are interpretable by construction, including concept bottlenecks, sparse probes, prototype networks, neural additive models, and mixtures of linear models.

Key Contributions

1. A "Rosetta Stone" for Interpretable Architectures

The paper provides a comprehensive mapping for common tensor operations used in interpretability research. For each operation, it presents:

• The graphical diagram.
• The corresponding PyTorch einsum code.
• The geometric interpretation (e.g., rescaling, affinity, projection).
• The symbolic algebraic equation.
  This covers fundamental operations (dot products, normalization, convex combinations) and complex blocks (self-attention, MLPs).

2. Design of Interpretable Components

The author uses the notation to formally describe and visualize the tensor manipulations behind specific interpretable architectures:

• Concept Encoding Maps: Visualizing how latent representations are transformed into human-meaningful concepts via sparse encoders, concept bottlenecks, concept embedding bottlenecks, and prototype-based maps.
• Concept Composition Maps: Diagramming how concepts are combined to make predictions in Neural Additive Models (NAMs), concept embedding predictors, and mixtures of linear models.

3. Case Study: Steerling-8B

The notation is applied to Steerling-8B, a frontier interpretable language model. The author diagrams its full architecture, demonstrating that:

• The entire core tensor manipulation logic can be represented in a single diagram.
• The diagram translates directly into approximately 33 lines of PyTorch code.
• The visual representation reveals structural properties not immediately obvious in the original technical report, such as the fact that Steerling is a residual model where gradients flow directly from output to input.

Results and Observations

• Global Insights: The graphical notation allows researchers to see the flow of information and the independence of feature manipulations at a glance, which is difficult in symbolic equations.
• Geometric Interpretation: Every operation in the diagram has a direct geometric reading (e.g., projecting vectors, finding affinities), aiding intuition.
• Reproducibility: The one-to-one mapping to einsum ensures that the diagrams are not just illustrative but are executable specifications.
• Complexity Handling: While the notation handles linear and multilinear operations exactly, the author notes that nonlinearities (masking, top-k) require ad-hoc extensions, and full frontier architectures can become visually unwieldy.

Significance and Claims

The paper claims to be the first work to apply Penrose-style tensor notation to the forward problem of designing architectures that are interpretable by construction, rather than just for the post-hoc analysis of opaque models.

The significance lies in the potential for this notation to become a standard tool alongside Probabilistic Graphical Models. By providing a formal language that is both globally intuitive and mechanically translatable to code, the author aims to facilitate the design, comparison, and implementation of state-of-the-art interpretable architectures. The author positions this work as a step toward unifying the theoretical understanding of interpretability with practical engineering implementation.