Partially Rewriting a Transformer in Natural Language

TL;DR

This paper explores partially rewriting a large language model using natural language explanations, aiming to improve interpretability while maintaining model performance through a pipeline of approximation, explanation, and simulation.

Contribution

It introduces a novel interpretability pipeline that replaces parts of a language model with natural language explanations and simulators, advancing understanding of model internals.

Findings

Replacing model components with explanations does not significantly increase loss.

More detailed explanations are needed for better performance.

The approach maintains model behavior close to original despite modifications.

Abstract

The greatest ambition of mechanistic interpretability is to completely rewrite deep neural networks in a format that is more amenable to human understanding, while preserving their behavior and performance. In this paper, we attempt to partially rewrite a large language model using simple natural language explanations. We first approximate one of the feedforward networks in the LLM with a wider MLP with sparsely activating neurons - a transcoder - and use an automated interpretability pipeline to generate explanations for these neurons. We then replace the first layer of this sparse MLP with an LLM-based simulator, which predicts the activation of each neuron given its explanation and the surrounding context. Finally, we measure the degree to which these modifications distort the model's final output. With our pipeline, the model's increase in loss is statistically similar to entirely replacing the sparse MLP output with the zero vector. We employ the same protocol, this time using a sparse autoencoder, on the residual stream of the same layer and obtain similar results. These results suggest that more detailed explanations are needed to improve performance substantially above the zero ablation baseline.