Prototype Transformer: Towards Language Model Architectures Interpretable by Design

TL;DR

The paper tackles opacity and unreliability in reasoning for large autoregressive language models. It proposes ProtoT, a prototype-based mixer that replaces standard self-attention and communicates bidirectionally with input tokens, achieving linear complexity $O(n)$ versus the usual $O(n^2)$ and enabling interpretable routing across time scales. Prototypes learn nameable concepts and support targeted edits, with interpretability validated through concept hubs, intervention studies, and time-scale analyses; robustness is augmented by prototype-mediated routing and causal interventions. Empirically, ProtoT achieves competitive generation quality and GLUE performance, while offering transparent pathways to analyze and edit model reasoning, indicating potential for safer, more transparent autoregressive LMs in practical deployments.

Abstract

While state-of-the-art language models (LMs) surpass the vast majority of humans in certain domains, their reasoning remains largely opaque, undermining trust in their output. Furthermore, while autoregressive LMs can output explicit reasoning, their true reasoning process is opaque, which introduces risks like deception and hallucination. In this work, we introduce the Prototype Transformer (ProtoT) -- an autoregressive LM architecture based on prototypes (parameter vectors), posed as an alternative to the standard self-attention-based transformers. ProtoT works by means of two-way communication between the input sequence and the prototypes, and we show that this leads to the prototypes automatically capturing nameable concepts (e.g. "woman") during training. They provide the potential to interpret the model's reasoning and allow for targeted edits of its behavior. Furthermore, by design, the prototypes create communication channels that aggregate contextual information at different time scales, aiding interpretability. In terms of computation scalability, ProtoT scales linearly with sequence length vs the quadratic scalability of SOTA self-attention transformers. Compared to baselines, ProtoT scales well with model and data size, and performs well on text generation and downstream tasks (GLUE). ProtoT exhibits robustness to input perturbations on par or better than some baselines, but differs from them by providing interpretable pathways showing how robustness and sensitivity arises. Reaching close to the performance of state-of-the-art architectures, ProtoT paves the way to creating well-performing autoregressive LMs interpretable by design.

Figures (26)

• Figure 1: A single autoregressive step of the ProtoT mixer module. Prototypes $P_1, \dots, P_R$ route inputs $x_1, \dots, x_{i-1}$ (past-only -- excluding $x_i$) into $R$ channels via similarity scores at the write gate. Time-discount and time-wise mean are applied per channel. The read gate reads from each channel via the similarity between its prototype and $x_i$, followed by aggregation into the output $y_i$.
• Figure 2: histograms of the numerical scores provided by LLM aided interpretability
• Figure 3: Left: Sequences most strongly activating prototype 11 at layer 0, which encodes the concept of narrative in a scholastic context. Right: Write-gate, read-gate, and memory curves for a sequence that strongly activates prototype 23 at layer 4. Read-gate peaks precede write-gate activations, spiking on the token immediately before those that trigger write-gate routing.
• Figure 4: Visualization for prototype L9 P7
• Figure 5: Visualization for prototype L9 P18