Prototype Transformer: Towards Language Model Architectures Interpretable by Design

While state-of-the-art language models (LMs) surpass the vast majority of humans in certain domains, their reasoning remains largely opaque, reducing trust and risking deception and hallucination. In this work, we introduce the Prototype Transformer (ProtoT)—an autoregressive LM architecture that replaces the quadratic-cost self-attention in the transformer with a linear-cost module based on prototypes (parameter vectors). In ProtoT, the prototypes create communication channels aggregating contextual information at different time scales. We show that this leads to the prototypes automatically capturing nameable concepts (e.g. “woman”) during training, and it provides the potential to interpret the model’s reasoning and do targeted edits of its behavior. Compared to baselines, ProtoT scales well with model and data size, shows robustness to input perturbations, and performs well on text generation and downstream tasks (GLUE). These results suggest that ProtoT is a promising step toward autoregressive language models that are more interpretable by design.

Advanced language models such as ChatGPT can write fluent text, answer questions, and solve difficult problems, but the true reasons behind their answers are often opaque. This makes them harder to trust, especially when they hallucinate, behave deceptively, or give different answers to changes that should not matter. We introduce the Prototype Transformer (ProtoT), a language model that generates text step by step while exposing more of its internal computation. Instead of using the standard transformer’s self-attention mechanism, which compares many pieces of text with one another and becomes expensive on long inputs, ProtoT uses learned prototypes. These prototypes act like shared concept channels: as the model reads text, they gather and pass contextual information across short and long ranges. During training, prototypes naturally come to represent nameable ideas, such as “woman”, making parts of the model’s reasoning easier to connect to human-understandable concepts. Because these concepts are represented by prototypes inside the model, we can also tweak the prototypes to change the model’s behavior in targeted ways. In experiments, ProtoT scales well, performs well on text generation and a wide range of language-understanding tasks, and behaves predictably when inputs are changed in ways that either preserve or alter meaning. Our results suggest that ProtoT is a promising step toward language models that are both powerful and designed to be more transparent.