Self-Ablating Transformers: More Interpretability, Less Sparsity (2505.00509v1)

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 LLMling 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.

• The paper introduces a self-ablation mechanism that enforces a kWTA constraint during training to enhance interpretability through localized, specialized pathways.
• It employs a dual residual stream with a straight-through estimator to balance masked and clean outputs, ensuring effective end-to-end training.
• Experimental results demonstrate improved interpretability metrics and modest performance trade-offs, validating practical benefits in model explainability.

This paper introduces "Self-Ablating Transformers," a novel method designed to enhance the interpretability of LLMs during training rather than relying solely on post-hoc analysis. The core idea is to integrate a mechanism that encourages the model to develop more localized and specialized computational pathways from inception, while preserving the standard transformer architecture and efficiency during inference.

The proposed self-ablation mechanism enforces a k-winner-takes-all (kWTA) constraint dynamically during the forward pass of training. This is achieved through learnable gating weights applied to neuron and attention units. The mechanism aims to ensure that for a given input, only the k most relevant components within a layer remain active. This selection is based on learned relevance scores derived from the gating weights.

To implement the non-differentiable top-$k$ selection for backpropagation, the authors employ a straight-through estimator. During the forward pass, a hard binary mask selects the top-$k$ activations. However, during backpropagation, continuous weights derived from a softmax function (applied to the activation values relative to a dynamic threshold and temperature) are used to compute gradients. This allows end-to-end training of both the base transformer weights and the auxiliary gating weights.

A dual residual stream is introduced during training to facilitate the calculation of the training loss. One stream processes the input with the self-ablation masks applied (ablated stream), while the other processes the input without any masks (clean stream). The total loss is a combination of the standard cross-entropy loss on the clean output and the cross-entropy loss on the ablated output. This design ensures that the model learns representations that are robust to the enforced sparsity while the gating weights learn to identify crucial components.

Two architectural variants of the self-ablation mechanism were explored:

1. Local Ablation: The ablation mechanism is integrated within each transformer block, making decisions based on the immediate context and output of the previous layer. This allows for fine-grained, layer-specific control and has lower computational overhead (single integrated pass) compared to the global approach during training.
2. Global Ablation: An initial forward pass through the entire network is performed to compute relevance scores globally. A second pass then processes the input using only the components selected based on these global scores. This approach leverages information from all layers for ablation decisions but requires a double forward pass during training.

The efficacy of self-ablation was evaluated by training small GPT-Neo models (TinyStories-3M config) on the TinyStories dataset and assessing interpretability using several methods:

• Automatic Circuit Discovery (ACDC): Used to identify the Indirect Object Identification (IOI) circuit. A reduced number of edges in the discovered circuit indicates better localization and interpretability.
• Sparse Autoencoders (SAEs): Trained on neuron activations to assess the disentanglement and sparsity of learned features (measured by L0 norm and reconstruction Cross-Entropy score). Lower L0 and higher CE suggest more interpretable features.
• Automated Neuron Explainability: Using a separate LLM (GPT-4o mini) to generate and score natural language explanations for neuron behavior. Higher scores suggest more consistently explainable neurons.
• Neuron to Graph (N2G): Analyzed neuron connectivity patterns (graph density, transitivity, out-degree) and specialization (token diversity, entropy) to understand how self-ablation affects neuron behavior at scale.

Key findings from the experiments include:

• Improved Interpretability: Self-ablated models consistently demonstrated improved interpretability metrics. ACDC revealed significantly sparser IOI circuits (up to 62% fewer edges), SAE analysis showed lower L0 norms (sparser feature representations) and higher CE scores (better reconstruction), automated neuron explainability showed a shift towards higher explanation scores, and N2G analysis indicated sparser graph connectivity, reduced transitivity, and increased neuron specialization, particularly in later layers.
• Modest Performance Trade-off: The improved interpretability came with only a slight increase in LLMling perplexity compared to the baseline model. Local ablation variants generally maintained perplexity closer to the baseline.
• Less Overall Sparsity: Surprisingly, self-ablated models exhibited higher L1 norms, indicating less overall weight sparsity compared to the baseline. This challenges the intuition that sparsity directly leads to interpretability, suggesting that self-ablation achieves interpretability through increased functional specialization and localization rather than widespread inactivity.
• Local vs. Global: Local ablation generally performed better across interpretability metrics, potentially due to more granular control and direct feedback within each layer.

A significant practical advantage of the self-ablation mechanism is that the auxiliary gating weights are only active during training. During inference, they are deactivated, leaving the model structurally identical to a standard transformer. This ensures that the interpretability benefits are gained during the learning phase without impacting inference speed or compatibility with existing deployment frameworks.

The authors acknowledge limitations, primarily the use of the simple TinyStories dataset and small model scale, which may limit the generalizability of the findings to more complex language tasks and larger models. Future work includes evaluating self-ablation on more diverse datasets and larger models, exploring its application in AI safety contexts (like controlled unlearning), and investigating its use in fine-tuning paradigms for existing LLMs. The code for self-ablating transformers is made publicly available on GitHub to facilitate reproducibility.