Search papers, labs, and topics across Lattice.
This paper introduces a novel graphical notation for designing interpretable AI architectures, addressing limitations of existing representations that either lack a global overview or obscure tensor manipulations. By adapting Penrose tensor notation, the authors provide a comprehensive framework that maps directly onto PyTorch einsum code, facilitating reproducibility and clarity. The notation is applied to various interpretable architectures, including Steerling-8B, revealing insights into their structure and operations while enabling straightforward implementation in code.
A new graphical notation reveals the inner workings of interpretable AI architectures, translating complex designs into clear, reproducible PyTorch code.
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.
