Discovering Interpretable Algorithms by Decompiling Transformers to RASP
Abstract
Recent work has shown that the computations of Transformers can be simulated in the RASP family of programming languages. These findings have enabled improved understanding of the expressive capacity and generalization abilities of Transformers. In particular, Transformers have been suggested to length-generalize exactly on problems that have simple RASP programs. However, it remains open whether trained models actually implement simple interpretable programs. In this paper, we present a general method to extract such programs from trained Transformers. The idea is to faithfully re-parameterize a Transformer as a RASP program and then apply causal interventions to discover a small sufficient sub-program. In experiments on small Transformers trained on algorithmic and formal language tasks, we show that our method often recovers simple and interpretable RASP programs from length-generalizing transformers. Our results provide the most direct evidence so far that Transformers internally implement simple RASP programs.
Lay Summary
LLMs can solve impressive problems, but we still have a limited understanding of the limits of their capabilities. This matters for safety: some tasks, such as copying a string of random symbols, are hard for Transformers to learn in a way that generalizes beyond the examples seen during training. As a result, we cannot always rely on models to perform such tasks correctly. Recent work has taken a step toward characterizing which tasks are difficult for Transformers to learn by showing that Transformer computations can be simulated in a family of programming languages called RASP. RASP provides a simpler way to describe sequence-processing algorithms, making it a useful tool for studying what Transformers are capable of representing. However, the fact that Transformers can implement RASP programs does not mean that trained models actually do so. In this paper, we present a general method for extracting RASP programs from trained Transformers. Our method rewrites the computations inside a Transformer and then prunes away causally irrelevant parts to find short, interpretable code. We test our approach on small Transformers trained on algorithmic and formal-language tasks. Our results provide the most direct evidence so far that some Transformers internally implement simple RASP programs.