Discovering Interpretable Algorithms by Decompiling Transformers to RASP

TL;DR

This work tackles mechanistic interpretability of Transformers by proposing a decompilation pipeline that reconstructs human-interpretable RASP-like programs from trained models. The core idea is to reparameterize Transformers as Decompiled RASP (D-RASP) programs under a linear layer norm assumption and then prune the resulting program to a minimal, causal substructure that faithfully reproduces the model’s behavior. Empirical results on small GPT-2 style models trained on algorithmic and formal-language tasks show that length-generalizing models often implement compact, readable RASP programs (e.g., histogram-based majority, induction-head copying, Dyck language counting), while non-length-generalizing models resist such simplification. This provides direct evidence that, at least in controlled settings, Transformers encode simple, interpretable algorithms that can be automatically extracted, supporting the length-generalization conjecture and offering a pathway toward systematic interpretability of neural computation.

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.

Figures (92)

• Figure 1: Program (a) for finding the most frequent character in a string, extracted from a real transformer. An example input is BOSobroSEPo; the model is trained to predict the last token "o". Line 1 computes, at each position, the relative frequencies of symbols at preceding positions ($\texttt{a1} \in \mathbb{R}^{|\Sigma| \times N}$, in (b). x-axis = input string, y-axis = variable dims., by aggregating over $\texttt{token} \in \mathbb{R}^{|\Sigma| \times N}$ (Figure \ref{['fig:initial-variables']}b). In Line 1, s=[] indicates that no selector is used, i.e., the weights $a_{i,j}$ in aggregation are constant. Line 2 projects a1 to output logits via the matrix in (c) (rows = input dims., columns = output dims.). For ordinary tokens (covered by op=), each token receives an output logit proportional to its frequency, corresponding to a temperature-scaled identity matrix. Logits for special tokens are excluded (special_op=(uniform selection)), precluding outputting BOS/SEP on length-1 strings (App. Figure \ref{['fig:majorityat1l4h256d3lr01dropactivation']}b has logits1). Overall, the program assigns the highest output logit to the most frequent non-BOS/SEP token. The program is extracted from a 1-layer 4-head transformer (App. \ref{['majorityat1l4h256d3lr01dropsection']}); an equivalent program results from a 4-layer 4-head transformer (App. \ref{['majorityat4l4h256d3lr00dropsection']}).
• Figure 2: Initial variables on the example input used in Figure \ref{['fig:majority']}. x-axis = input string $\in \Sigma^N$; y-axis = activation dimensions
• Figure 3: An extracted program for copying a string without repetitions: The input BOS1342SEPis completed with the string 1342EOS. The primitive "k == q-1" (line 1) selects the position immediately preceding the query position; it is given by the matrix in (a) (rows = query dimensions; columns = key dimensions) and yields the selector in (b) (x-axis = key positions; y-axis = query positions). The variable a1 (c) (x-axis = positions in the input; y-axis = dimensions of the variable) stores the previous token at each position, as a one-hot vector. The variables token and a1 jointly encode the bigrams appearing in the string, which is sufficient for copying a string without repetitions. The variable a2 (f) then stores the token to be output next. It is obtained by finding the position $j$ where a1(j) equals token(i), and retrieving token(j) into a2(i). Here, the primitive "k==q" selects a matching token; it is given by the matrix in (d). In these cases, there is special behavior on the special tokens, ensuring that the transformer produces the first token in the given sequence after $\langle \text{sep}\rangle$. a2 is directly forwarded into logits1 (h) by a scaled identity matrix (g). The bias in line 6 (i) favors EOS in the absence of other outputs, ensuring EOS is output at the end. The program recapitulates the "induction head" motif; importantly, our decompilation pipeline provides it fully automatically.
• Figure 4: Causal Pruning (Step 2.1): Pareto frontiers over lines of code vs. match accuracy for models satisfying LLNA.
• Figure 5: An extracted program for sorting a list of integers: The input BOS101549SEPis completed with the string 491015EOS. Line 1 uses a select operation given by the matrix in (a). Entries are largest when $k$ is slightly larger than $q$, and then fall off. As a result, a1 collects tokens bigger than the current one, and among those favoring the smallest one (b). SEP always collects the smallest number to start counting. Line 3 performs a nonlinear operation $f : \mathbb{R}^{|\Sigma|}\rightarrow\mathbb{R}^{d}$, not replaced by a primitive. We inspect behavior on twelve sample inputs $\texttt{a1}(i) \in \mathbb{R}^{|\Sigma|}$, each denoted by a row in (c, left), together with the associated output vector $\texttt{new\_a1}(i) \in \mathbb{R}^{|\Sigma|}$ (c, right). The operation "hardens" the input, selecting the dimension for which entries are largest in the input (e.g., 9, 10, 11 in the examples), effectively making the vector one-hot. When there is no number bigger than the current one, the attention is diffuse (as shown in the matrix in (a), entries are very similar when k is $\leq$ q), making the input a1 also diffuse. The same operation in Line 3 then maps this type of inputs to EOS (Fig. \ref{['fig:sortat1l1h256d3lr01drop-EOS']} shows that EOS is promoted when none of the non-BOS entry contains decisively high value). App. \ref{['sortat1l1h256d3lr01dropsection']} has more heatmaps.

Theorems & Definitions (8)

• Theorem 2.1: Correspondence of D-RASP and C-RASP
• Theorem 3.2
• Definition 2.1
• Theorem 2.2: Restated from Theorem \ref{['thm:faithfulness']}
• proof
• Theorem 3.1: Correspondence of D-RASP and C-RASP, repeated from Theorem \ref{['thm:c-rasp-no-pos']}
• Remark 3.2
• proof : Proof of Theorem \ref{['thm:c-rasp-no-pos']}