SlowGuess/NeurIPS_2023_OCR huggingface dataset:
David Lindner†
János Kramár
Sebastian Farquhar
Matthew Rahtz
Thomas McGrath
Vladimir Mikulik†
We show how to "compile" human-readable programs into standard decoder-only transformer models. Our compiler, Tracr, generates models with known structure. This structure can be used to design experiments. For example, we use it to study "superposition" in transformers that execute multi-step algorithms. Additionally, the known structure of Tracr-compiled models can serve as ground-truth for evaluating interpretability methods. Commonly, because the "programs" learned by transformers are unknown it is unclear whether an interpretation succeeded. We demonstrate our approach by implementing and examining programs including computing token frequencies, sorting, and parenthesis checking. We provide an open-source implementation of Tracr at https://github.com/google-deepmind/tracr.
Large language models (LLMs) are powerful but their inner workings are poorly understood (Danilevsky et al., 2020). The development of techniques for interpreting them is held back by a lack of ground-truth explanations (Yang et al., 2019). Our "compiler", Tracr, converts human-readable programs written in RASP, a domain-specific language for transformer computations (Weiss et al., 2021), into standard decoder-only transformers.
Tracr constructs models with known computational structure, which makes it easier to conduct interpretability experiments. As an example, we study neural networks' ability to compress a large number of sparse features into fewer dimensions using superposition (Elhage et al., 2022a). Compressing Tracr models using gradient de
scent allows us to study superposition in transformers implementing multi-step algorithms.
Figure 1: Interpretability tools produce explanations for a given neural network. Tracr creates models that implement a known mechanism. We can then compare this mechanism to explanations an interpretability tool produces.
A second use of transformers that implement known computations is evaluating interpretability methods aiming to reveal facts about a model's computation. Tracr could allow future work to directly test methods including, for example, classifier probes (Belinkov, 2022), gradient-based attribution (Nielsen et al., 2022), and causal tracing (Meng et al., 2022).
$^{\dagger}$ Correspondence to dlindner@google.com, vmikulik@google.com.
is_x = (tokens == "x")  
prevs = select(indices, indices, <=)  
frac_prevs = aggregate(prevs, is_x)
Figure 2: An example RASP program (left) that computes the fraction of previous "x" tokens at each position of the input. Tracr compiles this program to a transformer model. (right) A visualisation of a forward pass through the compiled model, showing the contents the residual stream, one panel per layer. The y-axis shows the residual stream dimensions, while the x-axis of each panel corresponds to the input sequence, "<bos>xacx" (x-axis). Changes to the residual are marked in red. Attn 1 is a no-op, MLP 1 computes the indicator variable is_x, Attn 2 implements the select-aggregate operation to compute frac_prevs, and MLP 2 is a no-op again. Section 4 discusses this and other examples in more detail. A detailed, step-by-step interpretation of the figure is provided in Appendix C.
Our main contributions are to: (1) Introduce Tracr, which "compiles" RASP programs into transformer models (Section 3); (2) Showcase models produced by Tracr (Section 4); (3) Provide a case-study where we examine superposition Tracr models compressed using gradient descent (Section 5). We confirm key observations by Elhage et al. (2022b) in a new setting: compressed models drop unnecessary features, and represent less important features in superposition.
In addition to aiding interpretability research, we think compiled models are a powerful didactic tool for developing a more concrete imagination for transformer mechanisms.
We discuss the applications and limitations of Tracr in Section 7 and Appendix A, and we provide an open-source implementation of the compiler at https://github.com/google-deepmind/tracr.
This section provides an overview of key concepts our work builds on. In particular, we review the transformer model architecture and the RASP programming language.
2.1 Transformer Architecture
A transformer model consists of alternating multi-headed attention (MHA) and multi-layer perceptron (MLP) layers with residual connections. Multi-headed attention (Vaswani et al., 2017) computes attention maps on sequences of length  $N$ . A single attention head  $i$  first computes an attention pattern
$$
A ^ {i} = \operatorname {s o f t m a x} \left(\left(x W _ {Q} ^ {i}\right) \left(x W _ {K} ^ {i}\right) ^ {T} / \sqrt {d _ {k}}\right) \in \mathbb {R} ^ {N \times N}
$$
for some input  $x \in \mathbb{R}^{N \times d}$ , where  $W_{Q}^{i}, W_{K}^{i} \in \mathbb{R}^{d \times d_{k}}$  are learnable parameters. Usually, we call the entries of  $(xW_{K}^{i})$  keys, and the entries of  $(xW_{Q}^{i})$  queries. Multi-headed attention combines  $H$  attention heads heads by computing
$$
\operatorname {M H A} (x) = \operatorname {C o n c a t} \left[ A ^ {1} \left(x W _ {V} ^ {1}\right), \dots , A ^ {H} \left(x W _ {V} ^ {H}\right) \right] W _ {O}
$$
where  $W_V^i \in \mathbb{R}^{d \times d_v}$  and  $W_O \in \mathbb{R}^{Hd_v \times d}$  are another set of learnable parameters. We commonly call the entries of  $(xW_V^i)$  values.
The MLP layers in transformer models compute  $\mathrm{MLP}(x) = \sigma(xW_1)W_2$  where  $W_1 \in \mathbb{R}^{d \times h}$ ,  $W_2 \in \mathbb{R}^{h \times d}$  are learnable weights, and  $\sigma$  is a non-linear function; for simplicity, we use the Rectified Linear Unit (ReLU). In this paper, we focus on decoder-only transformers, which consists of alternating blocks of MHA and MLP.
When designing Tracr, we adopt the transformer circuits perspective, introduced by Elhage et al. (2021). This view (1) focuses on the transformer being a residual stream architecture and (2)
introduces an alternative parameterisation for attention operations. Taking this viewpoint, simplifies reasoning about the transformer architecture and will help us when assembling transformers manually.
The residual stream view. Transformers have residual connections at each attention and MLP layer. Elhage et al. (2021) consider the residual connections a core feature of the architecture and describe the model in terms of a residual stream that each layer reads from and writes to in sequence. The residual stream acts as a type of memory that earlier layers can use to pass information to later layers.
Parameterising attention as  $W_{QK}$  and  $W_{OV}$ . Following Elhage et al. (2021), we parameterise an attention head by two (low-rank) matrices  $W_{QK}^i = W_Q^i (W_K^i)^T /\sqrt{d_k}\in \mathbb{R}^{d\times d}$  and  $W_{OV}^i = W_V^i W_O^i\in \mathbb{R}^{d\times d}$  where we split  $W_{O}$  into different heads, such that  $W_{O} = [W_{O}^{1},\dots W_{O}^{H}]$ , where each  $W_O^i\in \mathbb{R}^{d_v\times d}$ . We can then write MHA as
$$
A ^ {i} = \operatorname {s o f t m a x} \left(x W _ {Q K} ^ {i} x ^ {T}\right) \quad \mathrm {M H A} (x) = \sum_ {i = 1} ^ {H} A ^ {i} x W _ {O V} ^ {i}
$$
Importantly, we can think of MHA as summing over the outputs of  $H$  independent attention heads, each parameterised by low-rank matrices  $W_{QK}$  and  $W_{OV}$ .  $W_{QK}$  acts as a bilinear operator reading from the residual stream, and  $W_{OV}$  is a linear operator both reading from and writing to the residual stream. The softmax is the only nonlinearity in an attention head.
2.2 RASP
The "Restricted Access Sequence Processing Language" (RASP) is a human-readable computational model for transformer models introduced by Weiss et al. (2021). RASP is a sequence processing language with two types of variables, sequence operations (s-ops) and selectors, and two types of instructions, elementwise and select-aggregate transformations.
Sequence operations. A sequence operation (s-op) represents sequences of values during evaluation. tokens and indices are built-in primitive s-ops that return a sequence of input tokens or their indices, respectively. For example: tokens("hello") = [h,e,l,l,o], and indices("hello") = [0,1,2,3,4]. S-ops roughly correspond to the state of the residual stream in transformers.
Elementwise operations. RASP allows arbitrary elementwise operations on s-ops. For example, we can compute (3*indices) ("hello") = [0, 3, 6, 9, 12]. Elementwise operations roughly correspond to MLP layers in transformers.
Select-aggregate operations. To move information between token positions, RASP provides select-aggregate operations which roughly correspond to attention in transformers. A selector has a graph dependency on two s-ops and evaluates on inputs of length  $N$  to a binary matrix of size  $N \times N$ . To create a selector, the select operation takes two s-ops and a boolean predicate  $p(x,y)$ . For example:
$$
\text {s e l e c t} (\text {i n d i c e s}, [ 1, 0, 2 ], <  ) \left(\text {" a b c "}\right) = \left[ \begin{array}{l l l} 1 & 0 & 0 \\ 0 & 0 & 0 \\ 1 & 1 & 0 \end{array} \right].
$$
Here,  $p(x,y) = x < y$ , where  $x$  comes from indices, and  $y$  comes from the constant s-op  $[1,0,2]$ .
The aggregate operation takes as input a selector and an s-op, and produces an s-op that averages the value of the s-op weighted by the selection matrix. For example:
$$
\text {a g g r e g a t e} \left(\left[ \begin{array}{c c c} 1 & 0 & 0 \\ 0 & 0 & 0 \\ 1 & 1 & 0 \end{array} \right], [ 1 0, 2 0, 3 0 ]\right) = [ 1 0, 0, 1 5 ].
$$
A selector roughly corresponds to an attention pattern in a transformer. Together a select-aggregate operation roughly corresponds to an attention head in transformers.
2.3 Mechanistic Interpretability and Superposition
Mechanistic interpretability (Cammarata et al., 2020; Olah, 2022) aims to produce mechanistic explanations of the inner workings of ML programs. This includes attempts to reverse engineer how neural networks implement specific behaviours (Cammarata et al., 2020).<sup>3</sup>
In Section 5, we study superposition: the ability of a neural network to approximately represent many more features than the number of dimensions of the embedding space (Elhage et al., 2022a). Despite preliminary evidence that superposition occurs in neural networks, it remains poorly understood, in part because it has only been studied in small (2-layers or less) networks that implement very simple algorithms (Elhage et al., 2022b; Scherlis et al., 2022). Understanding superposition in larger models could represent a major step forward for mechanistic interpretability (Olah, 2022).
3 Tracr: A Transformer Compiler for RASP
In this section, we provide an overview of how Tracr translate RASP programs to transformer weights. For more details on the implementation, we refer to Appendix D and our open-source implementation at https://github.com/google-deepmind/tracr including the accompanying documentation.
Tracr comes with an implementation of RASP embedded in Python. A RASP program is an expression graph which is incrementally constructed from atomic RASP operations. We make a few technical modifications to allow translating RASP to model weights: we disallow boolean combinations of selectors, enforce annotated categorical or numerical embeddings for the residual stream, and enforce the use of a beginning-of-sequence token. We discuss the motivations for each of these changes in Appendix B, where we also explain how any RASP program can be refactored to be compatible with these restrictions. In practice, we can implement programs to solve all tasks described by Weiss et al. (2021).
If RASP is the high-level language we compile, Craft is our "assembly language", offering slightly more abstraction than pure weight matrices (cf. Figure 3). Craft provides a transformer implementation using vector spaces with labelled basis dimensions and operations on them. This lets us define projections or other linear operations in terms of basis direction labels, which simplifies con
Structing model components that act on different vector spaces. As a bonus, models represented in Craft are independent of specific transformer implementations. Models compiled by Tracr can be translated into weights of any standard decoder-only transformer model (without layer norm).
Figure 3: Tracr translates RASP to Craft and then to model weights, analogous to how programming languages are first translated to assembly then to machine code.
Tracr translates RASP programs to transformer weights in six steps:
1. Construct a computational graph (Figure 4(a)).
2. Infer s-op input and output values (Figure 4(a)).
3. Independently translate s-ops into model blocks (Figure 4(b)).
4. Assign components to layers (Figure 4(c)).
5. Construct the model (Figure 4(c)).
6. Assemble weight matrices.
Let us go through these step by step. Figure 4 gives a schematic overview using an example program.
1. Construct a computational graph (Figure 4(a)). First, we trace the whole program to create a directed graph representing the computation. The graph has source nodes representing tokens and indices and a sink node for the output s-op. Each operation in the RASP program becomes a node in the computational graph.
3Our approach is complementary: we construct neural networks to implement specific behaviours.
(a) Steps 1 & 2: Graph with inferred s-op value sets.
(b) Step 3: Nodes translated to MLPs and attention heads.
(c) Steps 4 & 5: Nodes allocated to locations in a model.
Figure 4: Schematic overview of how Tracr compiles the frac_prevs program from Figure 2 with a input vocabulary {"x", "y"} and context size 3. (a) shows the computational graph with value annotations after step 2 of the compilation. (b) shows how is_x and frac_prevs are translated to model components independently in step 3. (c) shows the assembled model which has two no-op components because models blocks always need to have one attention and one MLP layer.
We use a standard decoder-only transformer implementation in Haiku (Hennigan et al., 2020), notably removing layer norms. Extending Tracr to support any other transformer implementation is straightforward by reimplementing only step 6.
Figure 5: RASP program that sorts a sequence of numbers without duplicates. Attn 1 and MLP 1 implement the selector_width primitive (cf. Appendix D) which the program uses to compute the target position for each token. Attn 2 moves the tokens to the desired position, and MLP 2 is a no-op.
We are now ready to compile models with Tracr and walk through a few example programs.
4 Exploring Compiled Transformers
In this section, we walk through two example programs to illustrate how the compiled models work. While these models are not necessarily realistic, they represent configurations of weights that could, in principle, be learned. Examining these models can therefore be a powerful didactic tool for understanding how transformers perform complex computation, which we hope will expand our collective imagination for their inner workings.
We were able to compile RASP programs for all the tasks described in Weiss et al. (2021), though we had to modify a few programs to only use features supported by Tracr. Appendix G contains more examples.
4.1 Example 1: Counting tokens
Figure 2 shows our primary running example, the frac_prevs program, that computes the fraction of previous "x" tokens.
The compiled frac_pre vs model has a 14-dimensional residual stream, but it uses 12 out of these for the input embeddings. The remaining two dimensions contain the main numerical variables used in the computation: is_x and frac_pre (the output variable). The input embeddings have a few special dimensions. In particular, tokens :bos is the beginning of sequence token which we need to implement arbitrary attention patterns, and one is an input dimension that is always 1, used as a constant, e.g., to add a bias in MLP layers.
The compiled model uses one MLP layer and one attention head. However, because our model architecture always starts with an attention layer, the compiled model has four layers, with the first and last layers being no-ops. The first MLP layer computes the indicator variable is_x based on the input tokens. The following attention layer computes a causal attention pattern and uses it to compute the faction of previous "x" tokens.
4.2 Example 2: Sorting
As a second example, let us consider sorting a sequence of numbers. Figure 5 shows a sort_unique program that sorts a sequence of unique tokens.
The program computes uses a selector to select smaller tokens for each input token, and then uses the selector_width primitive in RASP to compute the target position for each token. selector_width counts the number of elements in each row of a selector that are 1, in this case the number of elements that are smaller than a given input token. selector_width can be implemented in terms of other RASP operations (Weiss et al., 2021). However, in Tracr we treat it as a primitive
that compiles directly to an attention and MLP layer (here Attn 1 and MLP 1). See Appendix D for more details. The model then uses a second attention layer to move each token to its target position.
Weiss et al. (2021) propose a sort program that can handle duplicates (cf. their Figure 13). However, that implementation uses a composite selector
select(tokens, tokens, <) or (
    select(key, key, ==) and select(indices, indices, <))
to treat duplicates, which is not currently supported by Tracr. In Appendix G, we provide an alternative implementation of sort that handles duplicates by adding a small multiple of indices to the keys and then applying sort_unique.
4.3 More examples
Tracr can compile a wide range of RASP programs. In Appendix G, we discuss several additional examples, leading up to a program to check balanced parentheses ( $Dyck-n$ ). Our open-source Tracr implementation (https://github.com/google-deepmind/tracr) contains a library of even more example programs to compile.
5 Compressing Compiled Transformers
Superposition is an important phenomenon in large language models (see Section 2.3, Elhage et al. (2022b), and Scherlis et al. (2022)). But to the best of our knowledge, it has not yet been studied in detail for models with more than two layers or in transformer models executing multi-step algorithms. Tracr lets us examine these models, and we can force different levels of superposition by applying a gradient-descent-based compression algorithm.
In addition to helping us study superposition, compressed models could be more efficient and realistic. Tracr models can be sparse and inefficient because they reserve an orthogonal subspace of the residual stream for each s-op.
Here, we present two case studies of compressing compiled models using the frac_prevs and the sort_unique programs from Section 4. These sketch how Tracr can be practically useful in advancing interpretability research, while also giving a glimpse of how Tracr could be extended to produce more realistic models.
5.1 Gradient Descent Based Compression
We use a single linear projection  $W \in \mathbb{R}^{D \times d}$  to compress the disentangled residual stream with size  $D$  to a smaller space with dimension  $d < D$ . We modify the model to apply  $W^T$  whenever it reads from and  $W$  whenever it writes to the residual stream (see Figure 6). We freeze all other weights and train only  $W$  using stochastic gradient descent (SGD). Since vanilla Tracr models are sparse and have orthogonal features, this process can be viewed as learning the projection from a "hypothetical disentangled model" to the "observed model" described by Elhage et al. (2022b).
We want the compressed model to minimise loss under the constraint that it implements the same computation as the original model. We train  $W$  to minimise  $\mathbb{E}_x[\mathcal{L}_{\mathrm{out}}(W,x) + \mathcal{L}_{\mathrm{layer}}(W,x)]$ , where
$$
\mathcal {L} _ {\mathrm {o u t}} = \mathrm {l o s s} (f (x), \hat {f} _ {W} (x)); \mathcal {L} _ {\mathrm {l a y e r}} = \sum_ {\mathrm {l a y e r} i} \left(h _ {i} (x) - \hat {h} _ {W, i} (x)\right) ^ {2}
$$
Here,  $f(x)$  is the output of the compiled model for input  $x$ ,  $\hat{f}_W(x)$  is the output of the compressed model, and  $h_i(x)$  and  $\hat{h}_{W,i}(x)$  are the output vectors at layer  $i$  of the respective models.
For categorical outputs,  $\mathcal{L}_{\mathrm{out}}$  is the softmax cross-entropy loss, whereas, for numerical outputs, it is the mean-squared error.  $\mathcal{L}_{\mathrm{layer}}$  is a regularization term that incentivizes the compressed model to match the per-layer outputs of the original model. To minimise this loss, the compressed model will have to approximate the computation of the original model but with a smaller residual stream. We give both loss terms equal weight, but we found other weighting factors give similar results in practice.
Figure 6: Training setup for compressing a compiled transformer model. At each layer, we use the same matrix  $W \in \mathbb{R}^{D \times d}$  to embed the disentangled  $D$ -dimensional residual stream to  $d \leq D$  dimensions. We freeze the layer weights and only train  $W$  to compress the model.
(b) Output loss vs.  $d$
Figure 7: Compressing the frac_prevs model Figure 2. (a) shows the loss during training for different embedding sizes  $d$  and (b) shows the final loss for different embedding sizes  $d$ . After about  $d = 6$  the compressed model solves the task essentially as well as the original compiled model which uses  $D = 14$  dimensions. (c) shows  $W^T W$  for the compression procedure described in Section 5 with  $d = 8$  where  $W$  is the learned compression matrix. As a comparison, (d) shows the same plot for applying PCA and retaining only the first 8 components. In contrast to PCA, our compression procedure produces a compression matrix  $W$  that retains features necessary for the task (e.g., is_x and frac_prevs) and discards features that are unimportant (e.g., tokens : a).
(c) SGD Compression
(d) PCA
We could set up this compression in other ways. For example, we could use a different projection at each layer, use different matrices for embedding and unembedding, or modify weights other than  $W$  when compressing the model. These design choices come with a tradeoff between making the model more expressible and potentially more realistic and enforcing the ground truth computation. For simplicity, we use a shared  $W$  for embedding/unembedding at every layer, and we already observe a rich structure in models compressed with this procedure.
Appendix E contains more details on the training setup, hyperparameters, and resources used.
5.2 What Does the Compression Learn?
As our first case study, Figure 7 shows the example model from Figure 2, that computes the fraction of token "x". By learning an embedding matrix  $W$ , we can reduce the residual dimension from  $D = 14$  to  $d = 6$  without hurting performance (cf Figure 7(b)). Once we reduce  $d$  further, the model's performance starts to suffer.
To understand the compression better, we can study how  $W$  embeds the original  $D$  features in  $d < D$  dimensions. We can only do this because we started with a compiled model with known features. Figure 7 shows  $W^T W$  for compressing the model to  $d = 8$ . We can compare this to using principle component analysis (PCA) to compress the model. To interpret the results, we need to use our knowledge of the algorithm the model implements. The input tokens: x and the variables is_x and frac_prevs are crucial for computing the fraction of tokens that is "x", and we find that these variables mostly get separate dimensions in the compressed residual stream. The other input tokens stored in tokens: a, tokens: b, tokens: c are not necessary for solving the task, and so they are discarded in the compressed model. Other variables, such as the indices embeddings, are stored in non-orthogonal dimensions in the compressed space. This is consistent with existing findings on superposition as the indices embeddings are sparse and do not occur together (Elhage et al., 2022b).