Self-Ablating Transformers: More Interpretability, Less Sparsity

TL;DR

The paper tackles the challenge of aligning interpretability with sparsity in transformers by introducing a self-ablation mechanism that enforces a $k$-winner-takes-all constraint during training, while keeping inference identical to standard transformers. By training on TinyStories with a GPT-Neo backbone and exploring local and global ablation variants, the authors show that self-ablation fosters more localized circuits and specialized neurons, with higher $L1$ activation norms indicating reduced global sparsity. These interpretability gains are supported by multiple analyses, including ACDC circuit pruning, sparse autoencoders, automated neuron explanations, and NeuronGraph metrics, while language modelling performance incurs only a modest perplexity cost. The work suggests that interpretability can emerge from structured, localized computation and specialization rather than simply reducing overall activity, offering a promising ante-hoc design principle for more transparent NLP models. The findings are supported by open-source code and point to future work on larger, linguistically diverse datasets and potential applications in model unlearning and safety.

Abstract

A growing intuition in machine learning suggests a link between sparsity and interpretability. We introduce a novel self-ablation mechanism to investigate this connection ante-hoc in the context of language transformers. Our approach dynamically enforces a k-winner-takes-all constraint, forcing the model to demonstrate selective activation across neuron and attention units. Unlike post-hoc methods that analyze already-trained models, our approach integrates interpretability directly into model training, promoting feature localization from inception. Training small models on the TinyStories dataset and employing interpretability tests, we find that self-ablation leads to more localized circuits, concentrated feature representations, and increased neuron specialization without compromising language modelling performance. Surprisingly, our method also decreased overall sparsity, indicating that self-ablation promotes specialization rather than widespread inactivity. This reveals a complex interplay between sparsity and interpretability, where decreased global sparsity can coexist with increased local specialization, leading to enhanced interpretability. To facilitate reproducibility, we make our code available at https://github.com/keenanpepper/self-ablating-transformers.

Figures (2)

• Figure 1: Comparison of global and local ablation mechanisms. Both models use a transformer with 8 blocks and are structurally identical during inference. Global ablation (a) uses a two-pass structure. The first pass calculates relevance scores using specialized ablation heads across the entire network. The second pass then processes the input using only the selected components based on these global scores. Local ablation (b) integrates the ablation mechanism directly into each transformer block. Each layer independently computes local relevance scores and makes ablation decisions based on its immediate context. See Appendix \ref{['sec:globalvslbl']} for a detailed comparison of processing efficiency, information access patterns, and implementation considerations between the two approaches.
• Figure 2: Self-ablation improves interpretability as shown through (left) IOI circuit simplification (baseline: 79 edges, self-ablated: 30 edges) where fewer model components are required to perform a task and (right) a shift towards higher neuron explanation scores in the self-ablated model (red) compared to the baseline (blue). These results indicate that self-ablation leads to more focused and interpretable circuits and neurons.