Self-Ablating Transformers: More Interpretability, Less Sparsity

TL;DR

This paper introduces a self-ablation mechanism in language transformers that enhances interpretability by promoting feature localization and neuron specialization, while also decreasing overall sparsity, challenging previous assumptions about sparsity and interpretability.

Contribution

The paper presents a novel self-ablation method that integrates interpretability into transformer training, resulting in more localized and specialized neural circuits without sacrificing performance.

Findings

Increased neuron specialization and feature localization.

Decreased overall sparsity with improved interpretability.

Enhanced interpretability through dynamic self-ablation.

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.