Tracr: Compiled Transformers as a Laboratory for Interpretability
Authors: David Lindner
Janos Kramar
Sebastian Farquhar
Matthew Rahtz
Thomas McGrath
Vladimir Mikulik
Venue: NeurIPS 2023 spotlight (2023)
Keywords: interpretability
transformers
language models
RASP
Tracr
mechanistic interpretability
TL;DR: Compiling human-readable programs into weights of a transformer model to accelerate interpretability research.

Abstract: 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.

Tracr: Compiled Transformers as a Laboratory for Interpretability
David Lindner Google Deep.ind
János Kramár Google DeepMind
Sebastian Farquhar Google DeepMind
Matthew Rahtz Google DeepMind
Thomas McGrath Google DeepMind
Vladimir Mikulik" Google DeepMind
Abstract
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 al- gorithms. Additionally, the known structure of Tracr-compiled models serve as ground-truth for evaluating interpretability methods. Commonly, be- cause 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 coom/goolee-eeee weeeeddddeeeeeaneeee..
Introduction
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 explana- tions (Yang et al., 2019). Our "compiler", Tracr, converts human-readable programs written in RASP, domain- specific language for transformer computations (Weiss al., 2021), into standard decoder-only transformers.
Noural
Interpretability
Explanation
Traor
the explanation
Known
Macharism
Figure 1: Interpretability tools produce explanations for given neural network. Tracr creates models that implement a known mechanism. We can then pare this mechanism to explanations an interpretability tool produces.
Tracr constructs models with known computational struc- ture, which makes it easier to conduct interpretability periments. As an example, we study neural networks' ability to compress 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.
A second use of transformers that implement known computations is evaluating interpretability methods aiming to reveal facts about model's computation. Tracr could allow future work to directly test methods including, for example, classifier probes (Belinkov, 2022), cradient-based attribution (Nielsen al., 2022), and causal tracing (Meng et al., 2022).
*Correspondence dlindner google.com, vmikulik goaiiikooooool.ocmmmm
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Input
tokens: peo
tolens:
CELL bos bos /0002
tolkens
1a_x (tokens =* 'x")
tokens:
prevs select (indice,, indices, <=)
one
frac_preve aggregate (preve, is_x)
indices:
indices:
Indicos:
indices:
indices:
bos bos
An example RASP program (left) that computes the fraction of previous "x" tokens at each position of the input. Tracr compiles this program 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 is no-op, MLP computes the indicator variable _x_,, Attn 2 implements the select-aggregate operation to compute prrccprres, and MLP 2 is no-op again. Section 4 discusses this and other examples in more detail. A detailed, step-by-step interpretation of the figure provided Appendix C.
Figure 2:
Our main contributions are to: (1) Introduce Tracr, which "compiles" RASP programs into trans- former models (Section 3); (2) Showcase models produced by Tracr (Section 4); (3) Provide case-study where we examine superposition Tracr models compressed using gradient descent (Sec- tion 5). We confirm key observations by Elhage et (2022b) in a new setting: compressed models drop unnecessary features, and represent less important features in superposition.
In addition aiding interpretability research, think compiled models powerful didactic tool for developing a more concrete imagination for transformer mechanisms.
We discuss the applications and limitations of Tracr in Section and Appendix A, and we provide an open-source implementation the compiler at https: / //ti//giittt github. com/ com/google wwgggge-deeemmnndeeeeooooo www-dddeeemmmaaaaaeee
2
Background
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
transformer model consists alternating mulli-healed attention (MHA) and multi-layer perceptron (MLP) layers with residual connections. Multi-headed attention (Vaswani et al., 2017) computes attention maps on sequences length N. A single attention head first computes an attention pattern
A' = softmax (zxWa) (w=W)) /vdz E N
for some input E Nxd where W W Rdxd leamable parameters. Usually, call the entries of (=WW) ) keys, and the entries of (rWb) queries. Multi-headed attention combines H attention heads heads by computing
MHA(x) Concat (onnat(A'(=w]]) [='W)))...... (xWF)] Wo
where Wi E Bdxdd and Wo E Hd. xd are another set of learnable parameters. We commonly call the entries of (xWi)) values.
The MLP layers in transformer models compute (LP oorWi where W1 gdxh We Rixdd are learnable weights, and is non-linear function; for simplicity, use the Rectified Linear Unit (ReLU). this paper, we focus on decoder-only transformers, which consists of alternating blocks of MHA and MI..
When designing Tracr, we adopt the transformer circuits perspective, introduced by Elhage et al. (2021). This view (1) focuses the transformer being 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 al. (2021) consider the residual connections 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 type of memory that earlier layers can use to pass information to later layers.
Parameterising attention as WOX and Wov. Following Elhage et al. (2021), we parameterise an attention head by two (low-rank) matrices Wok" W /Vdd E Rdxd and Wov' = Wi W E Rddd where split Wo into different heads, such that Wo [Wo Wo), where each W Rdoxd We can then write MHA
H MHA(*) Σ A'
A' = softmax (xWox x²)
Importantly, we can think of MHA summing over the outputs of H independent attention heads, each parameterised by low-rank matrices Wok and Wov. WoK acts as a bilinear operator reading from the residual stream, and Wov is linear operator both reading from and writing the residual stream. The softmax is the only nonlinearity an attention head.
2.2 RASP
The "Restricted Access Sequence Processing Language" (RASP) is 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 sequence of input tokens or their indices, respectively. For example: tokens ( ("hello") = [h, 1, o]. and indices ("hello") 1,, 4]. S-ops roughly correspond the state of the residual stream in transformers.
Elementwise operations. RASP allows arbitrary elementwise operations on s-ops. For exam- ple, we can compute (3+indi ces) ( ("hello") = [0, 3, 6,9, 12). Elementwise operations roughly correspond to MLP layers in transformers.
Select-aggregate operations. move information between token positions, RASP provides select- aggregale operations which roughly correspond to attention in transformers. A selector has graph dependency on two s-ops and evaluates on inputs of length N to binary matrix of size N x N. To create selector, the select operation takes two s-ops and boolean predicate p(a,, y). For example:
1
0
select (indices, [1, 0, 2], <) ("abc")
0
Here, p(z,y) y, where comes from indices, and y comes from the constant s-op [1, 0, The aggregate operation takes as input 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:
I
0
0
aggregate
0
0
[10, 20, 30] = [10, 0, 15).
1 0
selector roughly corresponds to an attention pattern in a transformer. Together a select-aggregate operation roughly corresponds to 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).²
In Section 5, we study superposition: the ability of a neural network to approximately represent many more features than the number dimensions of the embedding space (Elhage 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, provide 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 aattps://gite attps://gi.uub..ommggglleeee www./google-ddeemmnnddeoooooooe ttrcr including the accompanying documentation.
Tracr comes with implementation of RASP embed- ded in Python. A RASP program is an expression graph which is incrementally constructed from atomic RASP operations. We make few technical modifications to allow translating RASP to model weights: disallow boolean combinations of selectors, enforce annotated cat- egorical or numerical embeddings for the residual stream, token. We
Programming language
RASP
and enforce the use of a beginning-of-sequee discuss the motivations for each of these changes in Ap- pendix B, where we also explain how any RASP program can be refactored be compatible with these restrictions. In practice, we can implement programs to solve all tasks described by Weiss et al. (2021).
craft
Assembly
Machine code
RASP is the high-level language we compile, Craft is "assembly language", offering slightly more abstrac- tion than pure weight matrices (cf. Figure 3). Craft pro- vides 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 Tracr translates RASP to Craft and then model weights, anal- ogous to how programming languages are first translated assembly then to machine code.
Tracr translates RASP programs to transformer weights six steps:
1. Construct 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 gives schematic overview using example program.
1. Construct computational graph (Figure 4(a)). First, trace the whole program create directed graph representing the computation. The graph has source nodes representing tokens and indi ces and sink node for the output s-op. Each operation in the RASP program becomes node in the computational graph.
Our approach is complementary: construct neural networks implement specific behaviours.

---

[0.12]
FX. y)
(0.12)
tokens
Indices
tokena
atto Laool
MLP: x__
Attn: prevs
no-op mlp
0 0
M__
prevs
MLP: 1
prevs
frac_, preva
[0: % &
Attin: provs
[0.% %
Input (c) Steps & 5: Nodes allocated to locations in a model.
Output
(a) Steps 1 & 2: Graph with inferred s-op value sets.
(b) Step 3:
Nodes translated to MLPs and attention heads.
Figure 4:
Schematic overview of how Tracr compiles the frac_ program from Figure with 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 trac_preve 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.
2. Infer s-op values (Figure 4(a)). For each s-op, we need to decide how to embed it in the residual stream. To use categorical encodings, we need to know which values an s-op can take. All nodes have a finite set of output values because computations are deterministic, and we have a finite input vocabulary and context size. Therefore, in the second step, we traverse the graph and annotate each node with its possible outputs. This annotation uses simple heuristics that ensure we find superset of the values an s-op will take, though, sometimes, an output set can contain values that the s-op never takes in practice.
3. Independently translate s-ops (Figure 4(b)). Next, we consider each node in the computational graph independently and translate into model block. Elementwise operations become MLP blocks, and select-aggregate operations become attention blocks. We use library of manually engineered MLP and attention blocks to approximate arbitrary functions for numerical and categorical inputs and outputs. MLPs with categorical inputs and outputs function as lookup tables. MLPs with numerical inputs and outputs use piecewise linear approximations. For attention layers, we translate a selector into the Wo operator and the corresponding aggregate operation into the Wov operator. We only support attention with categorical inputs. We also do a few basic simplifications of RASP programs this stage. For example, we combine consecutive elementwise operations into single s-op. For more details on the MLP and attention blocks, see Appendix D.
4. Assign components to layers (Figure 4(c)). To construct transformer model, we need to allocate all model blocks in the computational graph to layers. Ideally, we want to find the smallest model to perform the desired computation. We can generally formulate this combinatorial optimization problem with several constraints: the transformer architecture has alternating attention and MLP layers, and all computations that depend reasons, we solve this problem heuristically. First, we compute the longest path from the input to given node. This path length is an upper bound for the layer number which we can allocate the node. Then we apply additional heuristics to combine layers with blocks that we can compute in parallel. This approach returns correct but sometimes suboptimal layer allocation.
each other need to be in the correct order. For scope
5. Construct the model (Figure 4(c)). We construct the residual stream space as the direct sum of all model components' input and output spaces. In other words, we embed each s-op in its orthogonal subspace, which is reserved for its sole use throughout the entire network. Now, we can traverse the computational graph in the order determined by the layer allocation and stack the components to obtain a full transformer represented in Craft.
6. Assemble weight matrices. Finally, translate the Craft representation of the model into concrete model weights. First, we combine parallel MLP layers into a single layer and parallel attention heads into single layer. In attention layers, we then factor the Wok and Wov matrices into separate Wg Ws. W.. We weight matrices. Finally, we adjust the shapes of all weights and connect them width, residual stream size, etc.) to fit the elements we have created.
transformer architecture. We can then infer the model configuration (depth, layer
We use standard decoder-only transformer implementation in Haiku (Hennigan et al., 2020), notably removing layer norms. Extending Tracr to support any other transformer implementation straightforward by reimplementing only step 6.
5

---

MLP2
Lokers: pad
smaller select (tokens, tokens, <-) target_pos vidhh (smalell)e))))r asl_sort select (targat..oo. indices, ==) aggregate (ee1_ tožens)
tokens:
tokons:
tokens:
at outpurt
target, pos
target: pas:
target_
Figure 5: RASP program that sorts sequence of numbers without duplicates. Attn and MLP 1 implement the selector_wiee (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 is a no-op.
sorti
primitive
wortz
sort:
sort:
inticass:
anticoss:
Indices:
bax
bos o'd bos
o bos
a 2 bod 7
We are now ready to compile models with Tracr and walk through a few example programs.
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 compile RASP programs for all the tasks described in Weiss 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 shows our primary running example, the frac_prevs program, that computes the fraction of previous "x" tokens.
The compiled frac_prevs prevs model has 14-dimensional residual stream, but it uses 12 of these for the input embeddings. The remaining two dimensions contain the main numerical variables used in the computation: is_x and frac_preve (the output variable). The input embeddings have few special dimensions. In particular, tokens bos the beginning of sequence token which we need implement arbitrary attention patterns, and one is an input dimension that always used as constant, e.g., to add bias in MLP layers.
The compiled model uses MLP layer and one attention head. However, because 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 based on the input tokens. The following attention layer computes a causal attention pattern and uses to compute the faction of previous "x" tokens.
4.2 Example 2: Sorting
As second example, let us consider sorting sequence of numbers. Figure shows sort_ parr_unique program that sorts a sequence of unique tokens.
The program computes uses selector to select smaller tokens for each input token, and then uses the selector_wia width primitive 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 given input token. selector, width can implemented in terms of other RASP operations (Weiss et al., 2021). However, in Tracr treat primitive

---

that compiles directly to an attention and MLP layer (here Attn 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 composite selector
se 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 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 program check balanced parentheses (Dyck-n). Our open-source Tracr implementation (https:/ htap::://gtthuie /ggtttttt....../.g com/ serglllddeeminn/treeee contains library of even more example programs compile.
Compressing Compiled Transformers
Superposition is an important phenomenon large language models (see Section 2.3, Elhage et al. (2022b), and Scherlis et al. (2022)). But the best of our knowledge, it has yet been studied in detail for models with more than two layers or in transformer models executing multi-step algorithms. Tracr lets examine these models, and we can force different levels of superposition by applying gradient-descent-base compression algorithm.
In addition to helping study superposition, compressed models could be more efficient and realistic. Tracr models can 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_preva and the port_uniqg programs from Section 4. These sketch how Tracr can be practically useful in advancing interpretability research, while also giving 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 E RDDx to compress the disentangled residual stream with size D to a smaller space with dimension
D. We modify the model to apply W 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 be viewed as learning the projection from "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 minimise E.e Cout (W, x)) + Laayer
BayerrW; (W, where
Cout = (w Llayer
Σ (h,(z) - (x))
layer
Here, f(x the output of the compiled model input vr fw (1:) is the output of the compressed model, and h2(r)) and hw.,(z) are the output vectors at layer the respective models.
For categorical outputs, Lout is the sofumax cross-entropy loss, whereas, for numerical outputs, it is mean-squared error. Llayer is regularization term that incentives 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 with smaller residual stream. We give both loss terms equal weight, but we found other weighting factors give similar results in practice.
7

---

Attn
MLP
Attn
MLP P
bx
hs
h4
W
W
W
W
W
W
W
Input
W
W
Output
Figure 6: Training setup for compressing a compiled transformer model. At each layer, we use the same matrix W E IR Pxd to embed the disentangled D-dimensional residual stream to d s D dimensions. We freeze the layer weights and only train W to compress the model.
0.06
10
outpun loo
Boot carpor
0.01
10--
0.02
0.00
10
sraining steps
eubbedding
(a) Training loss
(b) Output loss vs.
(c) SGD Compression
(d) PCA
Figure 7: Compressing the frac_prewe model Figure 2. (a) shows the loss during training for different embedding sizes d and (b) shows the final loss for different embedding sizes After about d = 6 the compressed model solves the task essentially well as the original compiled model which uses = 14 dimensions. (c) shows W for the compression procedure described in Section with d = 8 where W is the learned compression matrix. As comparison, (d) shows the same plot applying PCA and retaining only the first 8 components. contrast to PCA, our compression procedure produces compression matrix W that retains features necessary for the task (e.g., is_x and frac_prevs) and discards fcatures that are unimportant (c.g., tokens a).
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, 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 soredddinn/nemmbedddnng at every layer, and we already observe rich structure models compressed with this procedure.
Appendix E contains more details on the training setup, hyperparameters, and resources used.
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 < D dimensions. We can only do this because we started with compiled model with known features. Figure 7 shows W W for compressing the model to d = 8. We can compare this to using principle component analysis (PCA) compress the model. To interpret the results, we need use our knowledge of the algorithm the model implements. The input tokens x and the variables is and frac_preve 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 tokens: :, tokens:b tokens: c are not necessary for solving task, and so they are discarded in the compressed model. Other variables, such as the indices embeddings, are stored non-orthogonal dimensions in compressed space. This is consistent with existing findings on superposition the indices embeddings sparse and do not occur together (Elhage et al., 2022b).

---

Compited Compresoed Errr
1.0
tokens:bog bas
Figure 8: We compress sort_ sort_unique _nique (Figure 5). The plots on the right show that the compressed model achieves nearly perfect accuracy, but the layer outputs of the compressed model are different from the original compiled model. The left plot shows the average layer outputs the compiled model, the compressed model, and the squared difference. The compressed model seems to learn to use a different (nu- merical) encoding for the target_pos variable, which causes the discrepancy.
takens:
tokens:
tokens:
0.6
atto ourput:
Target pos:
0.0
target pos:
pos:
10
20
tardet
undookking size
Targer pos:
sart:
sort:
1.0
sart:
Coltte simileitty
one:
fndires:
0.5
intices:
MUP
AeS MLP
Att
MLP
MLP
0.0
10
20
cmooodking
However, our results go beyond previous work on superposition. Tracr models often have multiple variables that depend on each other and encode shared information. For example, in _pees, the is_x variable an indicator that essentially contains the same information as the input dimension tokens:x." In Figure 7, we see that the embeddings of is_x and tokens:: :x share part of the embedding space. Intuitively, this occurs because the variables encode similar information.
Future experiments could aim to further clarify the effect of shared information between variables on superposition. Tracr provides, for the first time, setting to systematically study superposition in transformer models that implement nontrivial algorithms.
5.3 Do the Compressed Models Still Implement the Same Computation?
Even if the compressed models successfully achieve a low loss, we need to check if they implement the same computation the compiled models, or else we would no longer know the ground truth mechanisms the models implement. To this end, we evaluate the average cosine similarity between the output at each layer of the two models. Values far from suggest the compressed model is structured differently from the base model.
We find that for some models the cosine similarity stays substantially below even as the compressed model gets close 100% in accuracy. For example, Figure 8 shows results from compressing the sort, unique model. Here, the compressed model achieves almost perfect accuracy on the task, but the average cosine similarity of the outputs at individual layers stays around 0.8, far shy of 1.
By inspecting the models' outputs at each layer, we can attribute the error to the target_ variable. In the compiled model, target_paa __osss is encoded as one-hot vector. However, the compiled model only uses a single dimension. This suggests that the compressed model moves the tokens to the target position with a numerical encoding of the target position rather than categorical encoding.
This difference in encodings shows that even with a fairly restrictive compression setup, compressed models may not stay faithful to the original RASP programs. This is both a setback for adding compression to the compiler-the compiler's annotations no longer serve as the exact ground truth- but also an opportunity. The ways neural networks solve algorithmic tasks regularly surprise researchers (Nanda et al., 2023). Studying such discrepancies could way to learn more about the ways NNs naturally represent certain computations without everse-engineering entire models.
Related Work
There are many approaches to interpretability in machine learning (Carvalho et al., 2019), and in language models specifically (Danilevsky et al., 2020; Belinkov and Glass, 2019; Rogers et al., 2020). In this paper, we focus on interpretability in the sense of giving a faithful (Jacovi and Goldberg. 2020) and detailed account of the mechanisms learned by model, sometimes called mechanistic interpretability (Olah, 2022) or transparency (Räukur al., 2023).
They are exactly the same because is_ only populated in a later layer.

---

Mechanistic interpretability has been used to reverse engineer circuits in state-of-the fft-the-att art vision models (Cammarata et al., 2020), small transformer models trained on toy tasks (Olsson et al., 2022; Nanda et al., 2023), and medium-sized language models ( (Wang et al., 2023). Reverse-engineered circuits can be used as more realistic alternative to compiled models. However, they are labor- intensive to identify, and our knowledge of them can be incomplete or inaccurate even when they are analysed carefully. For example, Chan et al. (2022) show that the "induction head" hypothesis by Olsson et al. (2022) needs to be modified to adequately explain in-context learning performance even in small attention-only transformers.
While Tracr is based on RASP (Weiss et al., 2021), there are potential alternatives for constructing transformer models. Wei et al. (2022) and Akyürek et al. (2023) study more general computational models for transformers. Based on this line of work, Giannou et al. (2023) propose a curine-complete model for constructing transformers, whereas RASP might have limited expressibility (Weiss et al., 2021; Merrill et al., 2022). However, the work by Giannou et al. (2023) is purely theoretical, and the practical cost-benefit trade-off between their approach and our RASP-based approach unclear.
Evaluation is a perennial topic of debate in interpretability, and there is little consensus on the best approach (Lipton, 2018; Yang et al., 2019; Mohseni et al., 2021). We hope that compiled models contribute a new perspective to this discussion and can complement other evaluation methods.
Our approach is closest to prior work trying create ground truth for evaluating interpretability. via careful manipulation of the training mechanism and dataset. Yang and Kim (2019) and Adebayo et al. (2020) introduce label correlations to the background of images, and Zhou et al. (2022) use label reassignments to achieve a similar goal. However, these approaches focus on convolutional image classification models, and they can only modify part of a model have a ground truth interpretation. Tracr, on the other hand, creates transformer models that implement fully human-readable code.
Since releasing an early version of our work, Conmy al. (2023) successfully used Tracr to evaluate a method for automatically detecting circuits transformer models, and Friedman et al. (2023) built on Tracr and studied learning transformer programs instead of manually writing them.
Discussion & Conclusion
We proposed to compile human-rcadable programs to neural network weights as a testbed for developing and evaluating interpretability tools. To this end, we introduced Tracr which compiles human-readable code to the weights of a transformer model.
Applications. Compiled transformer models can be broadly useful for accelerating interpretability research. We highlight four usecases that could be particularly useful. First, we can use Tracr to create test cases and ultimately benchmarks for interpretbility tools. This can help to confirm methods work as expected and surface potential failure modes. Second, we can measure our understanding of model by manually replacing components of it with compiled components (similar to Nanda et al. (2023)). Over time, the research community could build a library of programs that represent our understanding of what neural networks learn. Third, we can use compiled models to isolate and study phenomena that occur in real neural networks. Our study of superposition in Section 5 demonstrates the benefits of studying an isolated phenomenon in model we otherwise fully understand. Finally, compiled models can help us understand how transformers can implement certain algorithms and improve our ability to form concrete intuitions and hypotheses about models we want to interpret. Appendix A. discusses these applications in more detail.
Limitations. RASP and Tracr have important limitations in terms of expressivity, efficiency and realism compared to real transformer models. While many limitations can be overcome in future versions, some are fundamental to using compiled models. Clearly, we will likely never compile fully featured language models in Tracr. Therefore, we should interpret experiments conducted on compiled models carefully, and treat evaluations based on them minimum bar rather than full validation of a technique. Appendix A.2 discusses these limitations in detail.
Despite these limitations, we think Tracr provides a promising new approach to studying transform- ers and to evaluating interpretability tools. The current approach to doing interpretability research is similar to trying to invent microscope lens without ever being able to point it at familiar, well- understood shapes. Tracr enables researchers to point their interpretability methods at models they fully understand to calibrate, evaluate, and improve the methods.
10

---

Acknowledgements
We thank Avraham Ruderman, Jackie Kay, Michela Paganini, Tom Lieberum, and Geoffrey Irving for valuable discussions, Victoria Krakovna and Marlene Staib for collaborating early experiments with compiling RASP, and Chris Olah and Tristan Hume for feedback on an early draft of this paper. We thank the Less Wrong user "Gurkenglas" for pointing out mistake in earlier draft of the way to implement selectors combined with and described in Appendix
Author Contributions
VM proposed the initial idea for Tracr and wrote our RASP implementation. DL, VM, JK and MR designed and developed Tracr. DL designed, implemented, and ran the compression experiments in Section 5. MR wrote documentation and led the open-sourcing process. JK derived the theoretical results in Appendix F. TM and VM advised on research direction. DL, SF, and VM wrote the manuscript. DL led the project.
References
Adebayo, M. Muelly, Liccardi, and B. Kim. Debugging tests for model explanations. In Advances in Neural Information Processing Systems, 2020.
M. Aharon, M. Elad, and A. Bruckstein. K-SVD: An algorithm for designing overcomplete dic- tionaries for sparse representation. IEEE Transactions signal processing, 54(11):4311-45 2006.
E. Akyürek, Schuurmans, Andreas, Ma, and D. Zhou. What learning algorithm is in- context learning? Investigations with linear models. In International Conference Learning Representations (ICLR), 2023.
D. Bau, B. Zhou, A. Khosla, Oliva, and A. Torralba. Network dissection: Quantifying inter- pretability of deep visual representations. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2017.
Y. Belinkov. Probing classifiers: Promises, shortcomings. and advances. Computational Linguistics, 48(1):207-219, 2022.
Y. Belinkov and Glass. Analysis methods neural language processing: A survey. Transactions of the Association for Computational Linguistics, 7:49-72, 2019.
N. Cammarata, S. Carter, G. Goh, C. Olah, M. Petrov, L. Schubert, C. Voss, B. Egan, and S. K. Lim. Thread: Circuits. Distill, 2020. URL https://distill: pub/2020/caa ctt.:/it.......s...///.
D. V. Carvalho, E. M. Percira, and J. S. Cardoso. Machine learning interpretability: A survey on methods and metrics. Electronics, 8(8):832, 2019.
L. Chan, A. Carriga-Alonse N. Goldwosky- B. Shlegeris, and N. Thomas. Causal scrubbing, a method for rigorously testing interpretabil- ity hypotheses. Al Alignment Forum, 2022. https://www al awwgnmentforummmmmmmmmmmmm org/ org/posts/ CY2hh2y wttp:////wwrr/aaaaaaaaaaaraaaeee sartsascrraaabbbinnggggg.g....ge ieaiinng---eeehoe sorethod-for-rigrouuslyy
Greenblatt, J. Nitishinskaya, A. Radhakrishnan,
A. Commy, A. N. Mavor-Parker, Lynch, S. Heimersheim, and A. Garriga-Alonso. Towards automated circuit discovery for mechanistic interpretability. In Advances in Neural Information Processing Systems, 2023.
M. Danilevsky, K. explainable AI for natural language processing. AACL-IJCNLP 2020, 2020.
Qian, R. Aharonov, Y. Katsis, B. Kawas, and Sen. A survey of the state of
D. I.. Donoho. Compressed sensing. IEEE Trunsactions on information theory, 52(():1229-133 2006.

---

N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, N. DasSarma, D. Drain, D. Ganguli, Z. Hatfield-Dodds. D. Hernandez. A. Jones. J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and
C. Olah, mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. htt////aaafoomerrreeeeeeeeoeee circuits pub/2021/ http///araww/r/////eo///aaa.eeee hinpe/oorr///xxx..oome html.
URL https:
N. Elhage, T. Hume, C. Olsson, N. Nanda, TT Henighan, S. Johnston, S. ElShowk, N. Joseph, N. DasSarma, B. Mann, D. Hernandez, A. Askell, K. Ndousse, A. Jones, D. Drain, A. Chen, Kernion, Conerly, S. Kravec, S. Fort, S. Kadavath, J. Jacobson, E. Tran-Johnson, J. Kaplan, J. Clark, T. Brown, S. McCandlish, D. Amodei, and C. Olah. Softmax linear units. Transformer Circuits Thread, 2022a. URL pub/2022/ stte///sl////////eee solu/index html.
Y. Bai. D. Ganguli. L. Lovitt, Z. Hatfield-Dodds,
https: httt///aanffrmrrrrcteeeeoe
N. Elhage, T. Hume, C. Olsson, N. Schiefer, T. Henighan, S. Kravec, Z. Hatfield-Dodee R. Lasenby,
D. Drain, C. Chen, R. Grosse, S. McCandlish, J. Kaplan, D. Amodei, M. Wattenberg, and C. Olah. Toy models of superposition. Transformer Circuits Thread, 2022b. URL https: //trans camasformerreee circuits....... mubb/20222toy_mmod////// html.
D. Friedman, A. Wettig, and D. Chen. Learning transformer programs. In Advances in Neural Information Processing Systems, 2023.
A. Giannou, S. Rajput, J.-y. Sohn, K. Lee, D. Lec, and D. Papailiopoulos. Looped transformers as programmable computers. In International Conference on Machine Learning (ICML), 2023.
T. Hennigan, T. Cai, T. Norman, and Babuschkin. Haiku: Sonnet for JAX, 2020. URL http:
// //github. cwm/deemmmee htt///deeemmd////ee/eeeee dm- haiku.
A. Jacovi and Y. Goldberg.
Towards faithfully interpretable NLP systems: How should we define and evaluate faithfulness? In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, 2020.
M. L. Leavitt and A. Morcos. Towards falsifiable interpretability research. In NeurlPS Workshop:
ML Retrospectives. Surveys & Meta-Analyses (ML-RSA), 2020.
Z. C. Lipton. The mythos of model interpretability: In machine learning, the concept of interpretability is both important and slippery. Queue, 16(3):31-57, 2018.
K. Meng, D. Bau, J. Andonian, and Y. Belinkov. Locating and editing factual associations in GPT. In Advances Neural Information Processing Systems, 2022.
W. Merrill, A. Sabharwal, and N. A. Smith. Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, 10:843-800 885, 2022.
S. Mohseni, N. Zarei, and D. Ragan. A multidisciplinary survey and framework for design and evaluation of explainable AI systems. ACM Transactions on Interactive Intelligent Systems, 11 (3-4):1-45, 2021.
N. Nanda, L. Chan, T. Liberum, J. Smith, and Steinhardt. Progress measures for grokking via mechanistic interpretability. In International Conference Learning Representations (ICLR), 2023.
I. E. Nielsen, D. Dera, G. Rasool, R. P. Ramachandran, and N. C. Bouaynaya. Robust explainability: tutorial on wradient-based attribution methods for deep neural networks. IEEE Signal Processing Magazine, 39(4):33-84 84, 2022.
C. Olah. Mechanistic interpretability, variables, and the importance of interpretable bases. 2022.
C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, D. Drain, D. Ganguli, Z. HatfieldDDodss, D. Hernan- dez, S. Johnston, A. Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. URL. https: ttaps:///trasffommee circuitse pub//022/ in- context- learning and- -et--innuutioonnbnn heads/index inwex/ndexxx../// html.
12

---

T. Räukur, A. Ho, S. Casper, and D. Hadfield-Menell. Toward transparent AI: A survey on interpreting the inner structures of deep neural networks. In IEEE Conference on Secure and Trustworthy Machine Learning (SaTML.), 2023.
A. Rogers, O. Kovaleva, and A. Rumshisky, A primer in BERTology: What we know about how BERT works. Transactions of the Association for Computational Linguistics, 8:842-866, 866, 2020.
A. Scherlis, K. Sachan, A. S. Jermyn, J. Benton, and B. Shlegeris. Polysemanticity and capacity in neural networks. arXiv preprint arXiv:2210001892. 2022.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
K. Wang, A. Variengien, A. Commy, B. Shlegeris, and J. Steinhardt. Interpretability in the wild: circuit for indirect object identification GPT-2 small. In International Conference on Learning Representations (ICLR), 2023.
C. Wei, Y. Chen, and T. Ma. Statistically meaningful approximation: Turing machines with transformers. Advances in Neural Information Processing Systems, 2022.
study approximating
G. Weiss, Y. Goldberg, and E. Yahav. Thinking like transformers. In International Conference on Machine Learning (ICML), 2021.
Yang, M. Du, and X. Hu. Evaluating explanation without ground truth in interpretable machine learning, arXiv preprint
2019.
M. Yang and B. Kim. Benchmarking attribution methods with relative feature importance. arXiv preprint weXiiv:907.099000 2019.
Y. Zhou, S. Booth, M. T. Ribeiro, and Shah. feature attribution methods correctly attribute features? In AAAI Conference on Artificial Intelligence, 2022.
13

---

A Applications and Limitations of Tracr
We provide open-source implementation of Tracr because think it has many potential appli- cations in interpretability research. In this section, we discuss applications we see for Tracr and compiled transformers more generally and reflect on the current limitations of Tracr and how they can be addressed.
A.1 Applications of compiled models in interpretability research
Compilers like Tracr allow researchers to set up controlled experiments that test specific hypotheses about the computational structure of transformers. In this way, acts as laboratory for research in interpretability, enabling research that might otherwise be intractable.
Understanding model phenomena and developing new techniques. Compiled models can be used as a testbed for studying how learning affects circuits, and developing new approaches for interpreting transformer models. This is the approach where we successfully induce superposition in compressed Tracr models. Future work could analyse superposition in Tracr models, extending previous work in toy models (Elhage al., 2022b; Scherlis et al., 2022). In particular, Tracr allows studying how the structure of computation implemented by model affects which features will be stored in superposition. One goal for this line of research could be to predict how a specific Tracr model will be compressed, which features will be stored in superposition and how. A complementary approach is to try reversing the superposition induced by compression procedure, e.g., using ideas from compressed sensing and dictionary learning (Donoho, 2006; Aharon et al., 2006).
demonstrate in this work in section
Test cases for interpretability tools. Compiled models serve as a natural foundation for testing the faithfulness (Jacovi and Goldberg, 2020) of an explanation, and provide way to falsify (Leavitt and Morcos, 2020) the explanations given by interpretability techniques that aim to describe the inner workings of models.
For instance, classifier probes (Belinkov, 2022; Bau et al., 2017) aim to determine locations the model where particular features represented. A simple example of this approach is training linear classifiers using intermediate activations of a subject model as inputs. The performance of these classifiers at predicting some feature using activations from layer L is then taken as a proxy for the extent to which the feature is represented at that layer. Applying this method and correctly interpreting its results is challenging (Belinkov, 2022). Tracr- crapr-compiled models provide opportunity to what these methods say about models whose representations we understand fully, contextualising their results real models.
Ultimately, compiled models could be used to build libraries of test cases for interpretability tools, which could in turn enable quantitative evaluation metrics.
Replacing model components. Another way to evaluate our understanding of how model works is to replace parts of the model with hand-coded components. For example. Nanda et al. (2023) test their understanding of how transformer implements modular addition by replacing components of the model with their own idealised implementation and find that this can increase downstream performance, which strong evidence that the proposed explanation compiles an algorithm into full transformer model, it could be adapted to only compile part of model to replace part of a trained model. This could make easier to evaluate our understanding of a large model.
correct. While Tracr
Building intuition for algorithms implementable by transformers. Weiss al. (2021) highlight that RASP can be used to gain intuition for how transformers might implement certain tasks. Tracr natural next step in this direction, spelling out the relationship between the program and transformer implementing it in complete detail. We caution, however, that Tracr is but one approach to doing so, while real learned models could exhibit greater variety in their algorithms.
A.2 Limitations RASP and Tracr
RASP and Tracr are limited in terms of expressivity, efficiency and realism compared to real transformer models. Many of these limitations could be overcome in future versions of Tracr.
14

---

Expressivity. RASP is designed for algorithmic tasks that map an input sequence to discrete output sequence. However, current language models usually map sequence of imput tokens a probability distribution over the next token. Circuits in real models often consist of components that increase or decrease the probability of some tokens based on previous tokens (Wang et al., 2023). RASP, and hence Tracr, cannot model such "probabilistic" computation, but could potentially be extended to support it. RASP only uses binary attention patterns, which inherently limits the range of algorithms it can implement (Merrill et al., 2022). A way to extend RASP to support numeric attention patterns is discussed in Weiss et al. (2021).
Efficiency. Tracr models store all variables in orthogonal subspaces of the residual stream. Even if a variable is only used in part of the computation, Tracr reserves subspace of the residual stream for it in all layers of the model. Real models use a more compressed representation and likely reuse dimensions for multiple features. Improved versions of the compression procedure discussed in Section 5 could address this limitation, as would using constraint optimisation solver instead of heuristic for layer allocation.
Realism. Tracr constructs layers from hand-coded parameter matrices. This is both unrealistic and inefficient, but could be addressed by learning the layers in isolation, then assembling them into full model manually. Similarly, instead of manually splitting the Wok and Woy matrices, matrix factorisation could be used get more efficient solutions. Also, Tracr models align their features with the computational basis. This is unrealistic, and makes the resulting models easy to interpret just by inspecting the residual stream activations. Rotating the basis of the compiled model is straightforward way to address this if obfuscation is needed; compression would be an even more comprehensive approach.
While all of these issues could be overcome in a more sophisticated compiler, there are fundamental limitations on the role compiled models can play. Compiled models are an intermediate step between very simple toy models and real learned models. They help us understand ideas and methods, but results in compiled models do necessarily generalise to real models. Compared with real models, compiled models will always be simpler. For example, we will likely never compile full- full-fledged language models. Compiled models will be more likely to be intepretable (e.g., the axis-aligned orthogonal residual stream bases in Tracr), and more likely to fit into existing paradigms for thinking about transformers. When using them evaluate interpretability tools, we should be careful make sure that the tools do not exploit this, treating such evaluations as a minimum bar rather than full validation of technique. Conversely, some methods might conceivably rely on features present in real models but not in compiled models.
B Modifications to RASP
Disallow arbitrary selector combinations. RASP allows boolean combinations of selectors; how- ever, real transformers have natural analogue. Combining selectors with different input variables is particularly problematic. For example, in RASP we can define a selector
select (A, B, ==) and select (c, D, ==)
using four s-ops A, B, and D. However, real attention pattern only has two input vector spaces. There is no straightforward and efficient construction for representing arbitrary compositions of selectors (appendix F). Because of this, we restrict RASP selectors with only two input variables. In practice, this limitation seems not severe. In particular, we could implement programs to solve all tasks described by Weiss al. (2021).
a composite selector cannot be avoided, one can always refactor it into an atomic selector by first using s-ops create product spaces the inputs. In the example above, we'd ''d construct two s-ops whose values are pairs over values of A, and C, D respectively. Then, we could construct an atomic selector operating these composite s-ops:
select
product (A, B)
product D)
lambda b), (c, d): a==c and band
15

---

While this refactoring can be done mechanically, and would naturally generalise to arbitrary selector combinations, we chose not to include in our compiler implementation for two reasons. First, without compression it is inefficient: the s-op dimensions scale as product of the input s-op dimensions. Second, doing this automatically would break the 1-1 correspondence hetween selectors in RASP and attention heads in the compiled model: the compound s-ops require MLP blocks.
Encoding annotations. compiled model needs to pass information between layers. In a transformer, it is natural to do this in the residual stream (Elhage et al., 2021). However, our compiler must decide how to represent information in the residual stream. For simplicity, we only use categorical and numerical encodings. We encode categorical variables as one-hot vectors in dedicated subspace of the residual stream. We encode numerical variables as the magnitude of a dedicated me-dimensional subspace of the residual stream. We require each s-op to be either categorical or numerical and augment RASP annotate s-ops with the desired encoding. S-ops are categorical by default.
Even when both categorical and numerical encodings are possible for the same information, categori- cal encoding generally uses more dimensions and often requires an extra decoding step. However, some aggregate operations only work with one type of encoding. For instance, aggregation with a mean across token positions natural for numerical encodings but not categorical ones.
Beginning of sequence token. Transformers often assume any input sequence starts with a dedicated "beginning of sequence" token (BOS). We make the BOS token mandatory in RASP because it is crucial when implementing arbitrary attention patterns. In particular, RASP allows selectors that can produce all-zero rows; this is convenient when programming in RASP, but the softmax makes this be- haviour impossible in a real attention head. In these situations, we use the BOS token as a "default" po- sition attend it is attended to iff other token is. This allows the non-BOS part of the sequence to emulate the intended RASP behaviour. In our case, this choice comes from practical considerations; but, interestingly, real models sometimes show similar behaviour (e.g., see Elhage et al., 2021).
C Reading Model Output Figures
In the main paper and Appendix G, we show figures of a forward pass in a compiled model. We found that these figures can be confusing to read at first, especially as the compiled models get bigger. This section serves as a reference for how to interpret these figures.
As an example, let us walk through the figure for the frac_ model from Figure 2:
Input
Attn
Attn 2
tokens:
tokens:
pad
tokens:
tokens:
bos
tokens:
tokens:
one
indices:
indices:
indices:
indices:
indices:
frac_prevs
d x
The figure has 5 panels, each of which shows the content of the residual stream after the corresponding layer in the model. This allows us to follow what the model does step-by-step. The residual stream has size [sequence length heatmap. this example, we have a size of 5 x 14 including the BOS token position.
dimensionall
therefore visualize it as a 2-dimensional
We show forward pass for specific input sequence [bos, x, a, c, x]. On the x-axis of each panel we label the token positions with the corresponding input token. The y-axis of the plot contains the dimensions of the residual stream. Thanks to our knowledge of the program the model implements, we can label each dimension according to what it encodes. Dimensions starting with 'tokens' contain the (categorical) input embeddings. Dimensions starting with 'indices' contain the (categorical) position embeddings. Labels that contain a ": are dimensions that correspond to categorical ('one-hot') enoding and the value after the ";" is the value encoded in this dimension.
16

---

Labels without a ::" mean that this dimension encodes a numerical value. In each of the four panels we show the full residual stream content a heatmap. Entries that were changed by the layer corresponding to the panel are highlighted with a red border.
Let's go through the plot step by step and map it to the code used to compile this model:
1s_x (tokens == "x")
prevs select (indices, indices, <-)
frac_prevs aggregate (prevs, is_x)
In the leftmost panel we see the residual stream after the input embedding layer. It contains the categorical encoding of the tokens and the categorical encoding of the indices. For example at the token position of the 'a' token, the dimension tokens contains 1, and the dimension indices: 1 contains 1. The auxiliary dimension one contains at every token position.
The first attention layer is a no-op, so the residual stream afterwards (shown in the second panel) is the same as before. No entry is highlighted.
MLP computes the first line of the rasp program is_x = (tokens = "x"). It writes the result into a numerical dimension in the residual stream labelled with is_x. x. For this concrete sequence the layer writes a into both token positions that contain the token 'x' in the input.
Attn 2 computes the select-aggregate operations in lines 2 and 3 of the RASP program. It computes the fraction of previous 'x' tokens, and writes the result into a single dimension labelled frac_ _aaeessess.... It writes values between 0 and 1 in all token positions except for the BOS token position. For this example the result will be [1, 1/2, 1/3, 1/2].
The final MLP layer is no-op again and does not change anything in the residual stream. The output unembedding layer will then read the result from the srac_preva dimension.
D Tracr Implementation Details
This section highlights few more implementation details of Tracr. We describe how we construct MLP and attention blocks, how we implement the selector width primitive, and how we extend RASP and Tracr use causal attention. For the full implementation and documentation, refer the code repository at https: : /tte:://gitttt //gtttu...tt///////g//// com/ grogle-deeppine ttrarr.
D.1 MLP and Attention Blocks
For MLP layers, we distinguish between Map operations with a single input and output and SequenceMap operations with two inputs and one output. We can recursively represent functions with more than two inputs using SequenceMaps.
We translate Maps with categorical inputs and outputs to MLPs that act as lookup table. SequenceMaps with categorical inputs and outputs become MLPs where the first layer maps to an encoding of all pairs of inputs and the second layer acts as a lookup table.
For numerical inputs and outputs, we explicitly construct MLP layers as universal function approxi- mators. In these MLPs, the first layer discretises the input, and the second layer maps each discrete bucket to a corresponding output value. We know which input/output values can occur, so we can choose the discretisation around these known input values minimise the approximation error.
We now turn our attention to the attention blocks, which we construct from RASP selectors.
We first construct the Wox matrix to implement the desired attention pattern in the attention logits. We will refer to this as the direct attention matrix. This matrix has low rank, with its row space being the part of the residual stream where the query s-op is stored, and the column space being where the key s-op is stored. We adjust the direct attention matrix matrix by adding rank-one update W Woos- = BBo To thos with ßBOS 1 or 3 attended to either always, or whenever no other token is. (Nane and Trokens: here are unit vectors for the special embedding dimensions introduced in Section 4.) We then scale up the matrix by an inverse- temperature parameter -11 (100 by default), getting Wok T W Wos).. As a result,
n. ensure that the BOS token is
17

---

the attention weights Asy softmax (a.Wak exp(q[Wokk))/ Wola are
very close to # (selected ) on selected tokens and 0 elsewhere.
The Wov matrix maps the value input the corresponding output dimensions. Attention layers only support categorical key and query inputs. The value inputs can numerical or categorical. We can only use categorical values if the head never attends to more than one token.
D.2 Selector Width Primitive
RASP provides the selector width primitive, which counts the number of Is in each of a selector. It provides an alternative to aggregate for processing selectors.
Weiss et al. (2021) provide selector width implementation in pure RASP, making it not necessarily a language primitive. However, the most efficient implementation uses the BOS token, which exists Tracr but not exposed to the RASP program.
Therefore, Tracr translates selector width directly into an efficient implementation in Craft con- sisting of an attention layer and an MLP layer. The attention layer implements an attention pattern that matches the selector to compute the width of. uses the BOS token as value input, resulting in the attention head computing MLP layer then computes w = 1/x x and cleans the BOS token position.
1/(1 (1 w) where w is the desired selector width output. The next
D.3
Causal Attention
Most transformer models used in practice use causal attention, i..., they apply mask to the attention patterns that allows the model attend only previous tokens. This allows training the models autoregressively. However, RASP assumes non-causal (i.e. bidirectional) attention by default. While all models discussed in the main paper use non-causal attention, Tracr also supports causal attention.
To enable this, we extend RASP to support causal attention via a flag set during evaluation. To evaluate RASP program in the causal evaluation mode, we apply causal mask to the output of each selector. Causal evaluation changes the semantics of some RASP operations, and, in general, is necessary to adapt RASP programs function with causal attention. For example, the frac rrc__prees program no longer needs to compute a causal mask manually. However, for example, the length implementation by Weiss et al. (2021) no longer correctly computes the length of sequence because it requires attending to future tokens.
Similarly, Tracr has a flag to enable causal compilation. Most of the compilation process does not change, and we only need to ensure to compile selectors to causal attention heads.
E Compression Training Details
We implemented the compression described in Section 5 in Jax top of the Haiku transformer implementation that comes with Tracr. We train W using the AdamW optimizer (implemented in Optax) with weight decay factor of 0.1, and parameters 31 0.9, 32 0.99. We train for 3 x 10a steps with a batch size of 256. We decay the learning rate lincarly from 10-3 to 6 over the first half of training. Each compression (depending on the size the model to compress).
requires between and 4 hours of run time two CPU cores
F Theoretical Results Combining Attention Heads
The RASP language permits combining arbitrary selectors elementwise using boolean operators, such as and, or, and not. It is not immediately obvious what operators can be implemented given the way we encode selectors attention matrices WoK, described in Appendix D.1.
First, let's consider not operator for a selector select (query, key, pred) with given direct attention matrix WyR. One way to implement not select (query, key, pred) is to note that it's equivalent to select (query, key, not pred). Another is to use transformed direct attention matrix Wrat - WeK. alongside a Bos that's 0 or -
18

---

Next, let's consider the and operator on two selectors select (query-a, key_a, pred_a) and select (query_b, key_b, pred_b) b) whose direct attention matrices WỘR. W produce 0-1 attention logits. We can observe that taking W W + Wox results in attention logits taking value 2 when both selectors are active, and at most otherwise; so by the same procedure in Appendix D.I, with BóS taking value 1 or 2, we can construct W = T QK WBS)) that produces the desired attention pattern in the post-softmax attention weights.
QK are given, and
We can compose these constructions, negating the two given selectors before combining them with and, to get nor, with WoK QK = -Wa QK W&k QK and BoO taking value - or O, resulting in an implementation of select select k pred_a) nor select (query_b bb, key_b, pred_b).
So far these are fairly natural constructions the boolean operators not and and can be used to construct all other possible boolean operators, so we might expect that indeed all combinations of selectors via boolean operators can be compiled to transformer weights this way.
Alas, it is not so. Unlike the implementation of not, the implementations of and and nor above did not result in a direct attention matrix that produces the correct pattern (potentially shifted by constant) in the attention logits, but rather only in the attention weights after temperature-adjusee softmax, meaning they cannot be composed further to produce arbitrary logical statements.
If we were to try to implement or, the easiest way would be to negate the nor by composing the transformations but the resulting Wor - (-Wōk W is actually the same direct attention matrix we used for and. This produces attention logit I or 2 where the selectors' or is active, and O where it isn't. However, the temperature adjustment with T-1 >> 1 that forces the attention to be near-zero where neither selector is active will then also do the same thing when only one selector is active, so the attention weights will be different between tokens where both selectors are active versus only one selector.
In fact, this obstruction to implementing or can be generalized, as follows.
Lemma F.1. Consider two selectors select (query_A, key_A, pred_A) and select (query_B key_B, pred_B), B), with direct attention matrices W& and W QK- For ease of analysis, let's suppose query_ A, key_A, A, query_ B, and key_ B are stored in separate, orthogonal subspaces Q A. KA. QB, KB-
Now suppose there exists an attention matrix W QK. with row space contained in QA + QB and column space contained in R. + Ry, that, after adjustment by some BOS logit offset Bãos and some temperature T - O, produces attention weights converging to the normalized selector weights for se select ( (query_A, key_A, pred_ A) or select ( (query_B, B, key_B, pred_ B). Then, these se- lectors are not generic- - they satisfy some very limiting constraints about their predicates.
Proof. Let's begin by assuming the second selector, B, is not constant, selecting some tokens and not- selecting other tokens. This implies the existence of basis vectors 9:: 9 E QB and kg. k) K such that 9B "W QK kg = and q &K
QK k/B = 1. Holding these constant, consider some basis vectors qΛ E QA and kA, K4, E kA- Then, for query vector qл + q all tokens with key vector k + k) or K + kB must be selected, which means they must have equal attention logits. Therefore,
War (KA + kbb)= = (9. + q/3 W or (K's + kg). so q Wor K. q% Warks K'4 (qA + 9/)) QK
QK
Now, consider k = KA + kg. k' = K + kg q = qn + L and, for some basis vector q's E QA. let q' = 44 + 9h- We have logit differences q W QK k q @TWak = q WoR (k' k) q'TWokk' Q k' q @ Therefore, among tokens where key_B has vector kg (let's call these kg -toees,,)), the tokens that have highest logit for query vector q' are the same as those for query vector q. However, the selected tokens among the kb -toens are either none of them, or exactly those with the highest logit (which depends on ke)_ Because of the definition of or, kg -tokens are selected exactly if select (query_A, key_A, pred_A) would select them.
Work
Putting the above observations together, it follows that for query_A _A vectors q4 and d'a pred_A will either select no keys for one of them, or will select the same keys for both of them. In other words, pred_ must be rewritable in the form
(query_A) and key_pred_A A
19

---

(key_A), Equivalently, pred_A's matrix has rank 1; we can in short that pred_A is rank- predicate, or that seleciqqeey_A. key-A, pred_A) rank-1 selector.
we suppose our initial assumption to be false, then pred_ is constant, and can thus be just as well rewritten to be predicate query and key_A; A: then, is easy to derive the necessary Wor from QK
select (query_A. key_A, pred_A or pred_B).
We can repeat the argument interchanging the selectors, to conclude that either the operation is trivial (because one predicate constant), or both selectors must rank-1.
The above conclusion may averted in the case that we have priori information that certain values of q^, A 9B. kg cannot co-occur, if some of the input s-ops shared. We leave exploring that, as well as whether or can be implemented in the case of rank-1 predicates, to future work.
A notable special case of the above is the case where query_A and query_B compute the same s-op, and key_A and key_B also compute the same s-op. (They may the same s-op, or redundant copies.) Then simple rewriting is possible, similarly to the or case explained earlier. For example:
simplifsable_selere stuplirised
(tokenn,,, indices, <-) solect (tokens, "a", --) (tokena, indices, 9 <= or
"a")
similar strategy of matching s-ops can be used to circumvent the lemma and straightforwardly implement operators like or, by constructing combined s-ops query_ both and key_both with output types representing all pairs of queries and keys of the two selectors. These s-ops may be computed by the preceding MLP- however, the encodings occupy dimensionality multiplicative in far.
the sizes of the constituent s-op output types, which an impediment to scaling these circuits very
Due to the composability limitations of each approach considered, we did not implement boolean operators acting selectors, apart from simple cases where the query and key s-ops agree.
G More Compiled Models
Here, present few additional RASP programs and the compiled Tracr models.
Figure shows and extended sort program. It works similarly to the cort_unique program in Figure 5, but sorts any sequence of values by sequence of keys and can handle duplicates occurring the keys.
Figure 10 shows the pair_balance balance program, which computes the difference in the fraction of open and closed parenthesis tokens. We can now use as subroutine for the dyck-n program, which checks if sequence of n different types of parentheses is balanced:
Input: pairs
Compute running balance of each type parenthesis
2 balances [pair_ (pair) for pair in pairs]
If balances were negative anywhere -> parentheses not balanced
any any_negative balances [0]
for balance in balances [1:]:
any_negative = any_negative or (balance < 0)
select_all select (1, 1, =-)
10 has_neg aggregate (select_all. any any_negative)
11
12 If all balances are 0 at the end -> closed all parentheses
13 all_zero = balances [0] == 0 for balance in balances [1:]:
20

---

15
all_2ero all_zero and (balance == 0)
17
select_last select (indices, length 1, ==)
18
last_zero aggregate aggregate (select_ltt, all_zero)
19
20
dyck_n (last_zero and not has_neg)
Figure 11 shows the compiled dyck-2 model for pairs ("0", "(P).
21

---

Input: keys, vals, min key, context_lengtt
keys (keys indices min_key) context_
2 smaller select (keys, keys, <=)
target_pos =
sel_sort = select (target_pos, indices, ==) sort aggregate (sel sort, vals)
Input
Attn 1
MLP
Attn 2
MLP
Attn 3
MLP
tokens:
pad
tokens:
bas
tokens:
tokens:
tokens:
tokens:
tokens:
output
target pas:
target_pos:
target_pos:
target_ pos:
target_pos:
target pos:
sart:
sort:
sort:
sort:
sort:
sequence nар: 5.8
sequence nap: 5.6
sequence mар: 5.4
sequence_ned nар: 5.2
sequence pар: 5.0
sequence_nng nар: 4.8
sequence pар: 4.6
4.4
sequence nap: 4.2
sequence_nas nap: 4.0
sequence nap: 3.8
sequence_nas pap: 3.6
sequence pap: 3.4
sequence mар: 3.2
sequence nap: 3.0
sequence pар: 2.8
sequence_nap
sequence pар: 2.4
sequence pар:
sequence mар: 2.0
1.8
sequence nap: 1.6
sequence nар: 1.4
sequence nap: 1.2
sequence nap: 1.0
one
indices:
indices:
indices:
indices:
indices:
bos
3 3 bos 4
3 bos 4 bos 4 3 bos 4'3 4 bos
3 bos
3
Figure 9: Compiled sort program. Attn 1 no-op. MLP adds small multiple of indices to the keys, and the rest of the model essentially implements sort_ port_unique.
22

---

Input: open_token, close_token
bools_open (tokens open_token)
opens frac-prevs (bools_open)
boose (tokens close_token) closes = frac_preva, (bools_close)
pair_balance = opens closes
Input
Attn
MLP
Attn
MLP
tokens: pad
takens: bos
tokens:
takens:
pair_ba balaacee
opens
indices:
indices
indices:
indices:
indices:
closes
bools close
bas
( bos
bas
C'y' ('bos ('' bos C'
Figure 10: RASP program that uses frac_ srac_prevs subroutine to compute the fraction of open and closed parenthesis tokens and computes the difference. The compiled model uses open_tote token ="(" and token """. Note that the compiled model has the same number of layers the single frac_preva model in Figure 2, Attn is still no-op, MLP and Attn compute both calls to frac_prews in parallel, and MLP 2 computes the final result.
23

---

(y ( hnd ( box (1 boo ( hood 3 yood e ( had () yood ? ( bod t yood (1) ( poooo ( ( Food ( ybod ( Thod () (
MUP
Attn
MLF
Att
MLPS
Atti 6
Figure 11: Compiled dyck 2 program for pairs ("O". "(").
MLP
Atn
MLP
Attn
MLF 3
Az2 3
MLF2
Att
bod bod ( () bad ( ( hane of yook (
MLL
Altn
Input
bos
shuffie Syck 8: True
shutf fie f dyck 4: False
siquence ton 8: True
siquence mbn 8: False
sequence mep 10: True
sequence ne 18: False
True
not haz neq 6: False
True
mа 25: Fa
mао 24: True
mо0 24: False
Tа0 12: True
mе 12: Fat
ma 11: True
11: False
tength 15: 0
Lokens:
tokene:
tokens:
tokens:
Tokans:
tokans:
opens 22
opens 28
one
not has nea 6:
TаP 25:
10:
18:
10:
andices:
indices:
indices:
boold is opan 27
boo' i hat lalce.e.. 17
29
be Lance () 16
length
Length
closes
closas
boo open
bool i (lose
any megel
seleccorr
24


Reviews:
[{"review_id": "IfFX6FEf9s", "reviewer": "Reviewer_767m", "summary": "In this submission, the authors present Tracr, a compiler for RASP (a DSL for transformer computations) into transformer weights. The authors introduce their compilation approach, which includes an “assembly” language called Craft, which is used to represent the transformer weights agnostic to explicit implementations. They provide example models produced by Tracr (counting tokens & sorting, and refer to more examples in the repo). Most interestingly, they provide a way to compress these compiled models using SGD to study more involved topics in transformer explainability research, such as superposition.", "questions": "- How do the authors try to approach the mentioned limitations (in their last section) in future work? What is the path from here? \n- Can the authors elaborate more on the choice of compression (page 6, line 224+)? Can you provide more detail how this would affect the assumption that it can serve as a “ground truth”?", "limitations": "See their limitations section and Appendix A2. I feel that the authors have honestly addressed the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The compilation (section 3) is exceptionally well explained \n- Their approach is especially powerful as a didactic tool \n- The approach can serve as a guide for explainability research, and in providing a “ground truth” for explainability research \n- The open source implementation can serve as a starting point for many researchers interested in explainability of transformer models \n- Section 5 on compression is interesting and circumvents some of the limitations (otherwise it would be impossible to do research based on the compiled models on more involved concepts) ", "weaknesses": "- Examples feel toyish; the limited size of programs (see their limitation section) \n- That the compiled models could serve as a “ground truth” has to be taken with a grain of salt: Trained transformers, of course, optimize towards a different objective and solve a different problem.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> How do the authors try to approach the mentioned limitations (in their last section) in future work? What is the path from here? \n\nThere are technical limitations of Tracr that mostly come from design choices we made for simplicity (e.g., Tracr models embed different variables orthogonally, which can be quite inefficient). These limitations can be overcome with more sophisticated features. We plan to develop such features as they are needed for specific applications and we will encourage the research and open-source communities to contribute to Tracr.\n\nThe RASP language has some limitations, such as using only binary attention patterns. It is not fully understood which of these limitations are most severe, and we are excited about recent and ongoing work that studies the expressivity of different models of transformer computations (e.g., for binary attention patterns, Merrill et al. 2022 is relevant). In future versions of Tracr, we might extend RASP to remove some of these limitations.\n\nFinally, there are fundamental limitations to compiling models. Of course, we will never compile a fully-fledged language model with Tracr. Instead, we think compiled models will be useful as an intermediate step between analysing toy models and real learned models.\n\nFor a more detailed discussion of Tracr’s limitations and possible future work, see Appendix A.2.\n\n> Can the authors elaborate more on the choice of compression (page 6, line 224+)? Can you provide more detail how this would affect the assumption that it can serve as a “ground truth”?\n\nWe aimed to find a compression procedure that:\nmakes the compiled models more efficient and realistic;\nmaintains the “ground truth”-property of the compiled models, i.e., that we can identify which part of the model implements which computation;\nis conceptually simple and natural.\n\nRequirements (1) and (3) suggest a gradient-descent-based compression is a natural choice. To achieve (2), we freeze most of the compiled weights and only learn a new embedding matrix. Limiting the compression to a single shared matrix W ensures that the model still has the same structure. In most cases, we still fully understand the compressed models after investigating the learned W (e.g. see Section 5.2), which suggests that the compressed models are still useful as a ground truth. However, sometimes the learned compression can be quite complex, making the model more difficult to use as a ground truth (e.g. see Section 5.3).\n\nWe could try many other possible compression procedures here (and we might explore some in the future). But we think the proposed approach is a good trade-off between the three requirements and a natural starting point for studying compressing the compiled models. There is a fundamental trade-off between having a ground truth understanding of the model and having a very natural-looking model. While we provide a preliminary investigation of this in the paper, there is undoubtedly much room for future work here.\n"}, {"review_id": "HdGt5r2npf", "reviewer": "Reviewer_PkEE", "summary": "This paper presents Tracr, a compiler that can take 'program' specifications and translate them into GPT (decoder only) style transformer models. Tracr is built on the RASP 'programming' language introduced by Weiss et. al., and translates a RASP program into model weights via an intermediate representation termed craft. The entire compilation step occurs first by translating a RASP program into a computational graph where operations in the program correspond to nodes in the graph. The RASP language is equipped with: sequence operations and selectors, and instructions: element-wise and select-and-aggregate operations, which correspond to the MLP layers and attention operation respectively. The nodes in the graph are then translated to MLPs and attention heads which are then assembled into a complete model. The authors show how to use Tracr in two examples: sorting and token counting. Lastly, they use RASP to examine the influence of various model compression schemes on the model's logic.", "questions": "See weaknesses section for a longer discussion of the key questions. \n\n", "limitations": "There is extensive discussion of various limitations in the Appendix and Conclusion section.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "**Originality**\\\nThe key previous work is the RASP paper by Weiss et. al., however, this work significantly expands the scope of that previous paper into a usable compiler for assessing the effectiveness of interpretability methods on decoder-only transformer models. I find the series of choices made in the program translation process to also be intuitive. This paper presents an intriguing new tool will be useful for evaluating interpretability methods.\n\n**Quality/Clarity**\\\nThe paper is quite well written and the figures help demonstrate the point of the work. The use of the is_x example was very helpful in following the compilation steps that tracr implements. This work is of high quality and delivers a useful tool for the purpose for which it sets out.\n\n**Significance**\\\nIt is hard to judge significance of a paper, but I'd hazard a guess that this paper opens up a new line of work  for testing mechanistic interpretability methods, which are a new increasingly popular method for reverse-engineering large-scale models. Personally, I am quite skeptical of most research in mechanistic interpretability as I think it is susceptible to just reading tea leaves, so the tracr compiler could end up taking the place of unit tests for new mechanistic interpretability schemes. Imagine compiling a program to a set of weights and then asking whether some new mechanistic interpretability technique can identify the underlying mechanism. Overall, this paper has interesting ideas that can likely be adapted in future work.\n\n\n\n", "weaknesses": "Overall, I think this work is an important one, and opens several questions that could be follow-up work. None of the weaknesses I discuss below are disqualifying.\n\nI'll start with my major qualms with this work:\n\n**How different are tracr model weights from those that gradient descent learns?**: The model weights that one gets from tracr are going to be quite different from those that you get when you collect data from the RASP function and simply learn a model on that data. And I just don't mean that the exact values are different, but the weights could encode qualitatively different behaviors. It seems to me that based on the rules here, tracr models will have sparser weight matrices than traditional models learned via SGD. Since interpreting a sparse model is easier than a dense one, do the authors see problems with tracr models being quite different from learned ones?\n\n**Details**: As expected, a paper like this is packed with insights, and the details here actually matter. Specifically, I am referring to the heuristics used to 1) decide values that an s-op will take, 2) combining layers. Of course one has to make concrete choices here, but it is a bit unclear to me how these rules should impact the 'quality'/type of function you get out of tracr. Similarly, the last step, a crucial one, is also under explained in the main draft. Specifically, it is unclear what the authors mean by \"factor\", and exactly how they arrive at the specific values for these weight matrices. Some of this information is discussed in the appendix.\n\n**The interpretability assessment pitch**: A main motivation of this work is to test interpretability methods. However, this paper stops short of that goal. It would've been amazing for the authors to apply a mechanistic interpretability to say a dyck example to see whether these methods can recover such logic. Even further, while the tool is an important component of testing an interpretability method, it is unclear to me how you would actually use it to do so. Section 5 attempts to do this, but not in a straightforward manner.\n\n**Minor Issues**\n1) I had to go read Elhage et. al. to understand what a 'residual stream' is. Is this interpretation now standard in the literature? Section B.2 in the appendix is actually important background notation. \n2) Defn of mechanistic interpretability in lines 54-55 is circular. What is a mechanistic explanation? I think the second sentence is actually the real definition?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n> How different are tracr model weights from those that gradient descent learns?\n\nYour description is accurate: Tracr models are significantly sparser than real trained transformers, and, in particular, small Tracr models tend to be easy to interpret. There are straightforward ways to \"obfuscate\" the compiled models, e.g., by rotating their representations. Doing this makes the models significantly harder to understand. However, as you mentioned, Tracr models' weights look qualitatively different than the weights of compiled models. This is what we mean by Tracr models being “unrealistic” in Appendix A.2 when discussing limitations; we will make sure to clarify this in the main paper.\n\nThis observation motivates the second part of our paper on compressing Tracr models using SGD. We find that this compression procedure makes the models significantly more realistic while maintaining the \"ground truth\" computation. There is some trade-off between models being natural and models having a ground truth interpretation (e.g. see Section 5.2). But we primarily see Tracr models filling the role of \"unit tests\" for mechanistic interpretability (Thanks for this metaphor!). So, we think it is acceptable to sacrifice some amount of \"naturalness\" for having the ground truth interpretation.\n\n\n> Specifically, it is unclear what the authors mean by \"factor\", and exactly how they arrive at the specific values for these weight matrices.\n\nWe assume you are referring to the sentence “we then factor the W_QK and W_OV matrices into separate W_q, W_k, W_o, W_v weight matrices” in lines 142 and 143.\n\nIn our implementation, we simply set $W_q = W_{QK}$ and $W_k = I$ as an identity matrix and $W_o = W_{OV}$ and $W_v = I$ as an identity matrix. However, we might use more sophisticated factorisations in the future. We agree this deserves some more details in the main paper.\n\n\n> I had to go read Elhage et. al. to understand what a 'residual stream' is. Is this interpretation now standard in the literature? Section B.2 in the appendix is actually important background notation.\n\nWe agree that Section B.2. contains important background, and we aim to move it to the main paper when updating the paper.\n\n\n> The interpretability assessment pitch: A main motivation of this work is to test interpretability methods. However, this paper stops short of that goal. It would've been amazing for the authors to apply a mechanistic interpretability to say a dyck example to see whether these methods can recover such logic. Even further, while the tool is an important component of testing an interpretability method, it is unclear to me how you would actually use it to do so. Section 5 attempts to do this, but not in a straightforward manner.\n\nWe agree that testing an interpretability method with a Tracr model would make the impact much more apparent. We are addressing this in follow-up work, but we felt it would have expanded the scope of this paper to the point that the contributions would have been harder to communicate clearly.\n\n"}, {"review_id": "ugtDZf0Bhf", "reviewer": "Reviewer_XH5L", "summary": "The paper presents Tracr, which is a compiler from RASP programs (a language designed to showcase a possible computational model used by Transformers) to Transformer architectures and weights. The paper first details the operation of the compiler, then discusses examples of several simple RASP programs and their respective Tracr-generated Transformer models. Finally, the paper experiments with compressing the dimensionality of these models, exploring an emerging hypothesis about superposition of feature representation in Transformers.", "questions": "* How precisely do the compiled Transformers implement the original algorithms?\n* Regarding \"Disallow[ing] arbitrary selector combinations\" in Appendices C and G: does this restriction reduce the set of programs that it is possible to represent with RASP?\n* This is more of a curiosity than a criticism, but in Section 5.1 when projecting out of the compressed space, is there a reason to apply W^T rather than say the psuedoinverse of W?\n", "limitations": "The authors adequately discuss limitations in the appendix.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper addresses a significant gap in prior work, which is that while RASP purports to encode the computational model for Transformers, RASP programs are not directly compilable to Transformer weights. This is an active and interesting area of research, and providing such a compiler is likely to significantly impact work in this domain.\n* Overall, the paper is very clearly written\n* The examples given (Section 4) strongly demonstrate the success of the technique, and broadly help to explain RASP's computational model as well\n* Section 5 is an interesting study on the effects of compression, though (as noted in Weaknesses below) I am not sure about the relevance or significance of these results.\n", "weaknesses": "* The paper does not sufficiently discuss the accuracy of the compiled programs (Section 3). Line 49 says \"Any elementwise operations in RASP can be approximately computed by an MLP layer\", and line 119 says that the MLPs and attention blocks \"approximate arbitrary functions\". Having error in the compiled programs is a perfectly reasonable limitation, but it would be very good to know the fidelity of these approximations.\n* The paper assumes a high level of familiarity with prior work, which could be inlined into the paper.\n  * Lines 59-62: I'm not familiar with Elhage et al. (2021)'s concept of a residual stream, but this seems to be a core concept in the paper. It would help to clarify these concepts and definitions in the paper. Appendix B.2 provides some information here, but this is still a bit insufficient (a \"residual stream\" is never actually defined).\n  * In Section 3, it could also help to include a figure with RASP's syntax to understand the specific functions that must be compiled (seemingly, `select`, `selector_width`, `aggregate`, and the element-wise operations; `selector_width` isn't mentioned until Section 4.2), and the specific operations that are not yet supported by Tracr (e.g., those on Line 187). Again, Appendix B.3 helps, but is still not quite sufficient (e.g., it also does not define `selector_width`).\n* Section 5 reads as an entirely different paper, with its own significant limitations. As the authors note, \"even with a fairly restrictive compression setup, compressed models may not stay faithful to the original RASP program\". This is certainly an interesting finding (and has ramifications on compression techniques broadly). But, given that these compressed Transformers neither follow the original RASP program, nor is it clear that this the behavior observed when compressing non-Tracr Transformers, it's not clear what to take away from these results.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> How precisely do the compiled Transformers implement the original algorithms?\n\nTracr can compile any algorithm to a finite model that implements it exactly, i.e., with zero approximation error. This is because we know the full (discrete) input vocabulary for the model at compile-time. So, while some functions are implemented “approximately” by MLPs, we ensure that for all values that can occur for inputs from the vocabulary, the MLP approximation is exact (this is also briefly discussed in Appendix E.1). We will add a sentence clarifying this in the main body.\n\n\n> Regarding \"Disallow[ing] arbitrary selector combinations\" in Appendices C and G: does this restriction reduce the set of programs that it is possible to represent with RASP?\n\nStrictly, yes, this reduces the set of programs that we can compile. For example, without selector combinations, we cannot compile the sort program by Weiss et al. 2021 because it uses a combined selector to handle duplicates in the input:\n\n```\nselect(keys, keys, <) or (select(keys, keys, ==) and select(indices, indices, <)\n```\n\nHowever, any such program can be refactored into an equivalent RASP program that we can compile. We describe this procedure at the end of Appendix G, which also discusses the technical reasons for disallowing arbitrary selector combinations. In brief: this is primarily because combined selectors a) produce large and inefficient models, and b) would break the correspondence between selectors and attention heads. We therefore leave doing this refactoring to users to avoid surprises in compilation.\n\nIn practice, we have not found that this restriction introduces practical difficulties.\n\nThinking about this question again helped us clarify our thinking on this issue; thank you – we will update the paper and the appendix to reflect that.\n\n\n\n> This is more of a curiosity than a criticism, but in Section 5.1 when projecting out of the compressed space, is there a reason to apply W^T rather than say the psuedoinverse of W?\n\nWe made this choice primarily to be consistent with Elhage et al. 2021. Using the pseudoinverse of W would also be a reasonable choice.\n\n\n> The paper assumes a high level of familiarity with prior work, which could be inlined into the paper.\n\nThanks for flagging this! We aim to address the insufficient discussion of necessary background by moving Appendix B.2 to a dedicated background section and improving it. We will also improve our discussion of RASP syntax to make the paper more self-contained.\n"}, {"review_id": "6X4Arl2yrr", "reviewer": "Reviewer_ACvV", "summary": "This submission builds on the RASP language proposed by Weiss et al. It proposes a way to compile RASP program into real Transformer weights. It also comes with a case study showing the use case of Tracr to study a phenomenon called superposition which is widely known in Mechanistic Interpretability field. It's claimed by the authors that their Tracr can serve as a ground-truth for interpretability research. ", "questions": "+ On Line 209. The choice of $W^\\top$ in Elhage et. al. to recover the compressed feature actually make more sense than in the case of this submission. They are forcing $W$ to be orthonormal so that $WW^\\top=I$ stands, which is believable if it happens since their loss function is encouraging that. In Tracr's case, I cannot see why the residual added into the stream is supposed to be the same as the input.", "limitations": "+ This is indeed a good pedagogical tool for Transformer architecture. However, it's worth highlighting in the paper that LayerNorm or any similar operations, RMS Norm are left out, so to not mislead readers. ", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This piece of construction result is of evident significance to the community, considering Transformer is the dominant architecture these days. It furthers our understandings theoretically how Transformer can be implementing various algorithms inside of it. I hope this can inspire us to discuss what can and what **cannot** be expressed in a specified architecture. ", "weaknesses": "+ On the first use case: it is studying a phenomenon not even well defined in the literature, called superposition. Wish to see more rigorous framing of what is really studied, e.g., compressed sensing. The results in Section 5 isn't impressive to me since it isn't an organic combination of compressing and Tracr. \n\n+ On the second use case: It feels to me an overclaiming that Tracr could serve as a _ground-truth_ for evaluating interpretability methods. This is an understandable imagination but such idealism isn't correct. An interpretability algorithm can discover very different underlining algorithms and if the discovered is different from Tracr compilation; only the Tracr is to be challenged. \n\n+ Summary up, I acknowledge this is an important theoretical results but framing it as interpretability or anything related to real Transformer behaviour sounds far-fetched. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> On the second use case: It feels to me an overclaiming that Tracr could serve as a ground-truth for evaluating interpretability methods. This is an understandable imagination but such idealism isn't correct. An interpretability algorithm can discover very different underlining algorithms and if the discovered is different from Tracr compilation; only the Tracr is to be challenged.\n\nThis might be a misunderstanding of what we mean by ground-truth. We do not mean that if we can compile an algorithm in Tracr, this algorithm is necessarily the correct interpretation for a learned transformer model completing the same task. Say we compile a model using Tracr to sort a sequence of numbers and train a transformer to do the same task. You are correct that the trained model might use a different algorithm from the compiled model. The advantage of the compiled model is that if we run an interpretability method on it, we know exactly the correct interpretation for _that compiled model_ – the RASP program serves as a ground-truth. We can then see how that ground truth is reflected in the interpretability method.\n\n\n> On Line 209. The choice of in Elhage et. al. to recover the compressed feature actually make more sense than in the case of this submission. They are forcing to be orthonormal so that stands, which is believable if it happens since their loss function is encouraging that. In Tracr's case, I cannot see why the residual added into the stream is supposed to be the same as the input.\n\nYou are right; there is no principled reason why the model needs to use the same embedding at each point in the model. We choose the same embedding matrix throughout the model primarily for simplicity, but also because this is how uncompressed Tracr models are structured. As mentioned in response to Reviewer 767m, we found our compression procedure to be a particularly good tradeoff between being simple and producing more efficient and natural models. However, it is worth noting that recent empirical evidence from real transformers suggests that the input embedding matrix is still a meaningful interpretation of the residual stream at other points of the model (e.g., see [1]).\n\n\n> However, it's worth highlighting in the paper that LayerNorm or any similar operations, RMS Norm are left out, so to not mislead readers.\n\nThanks, we will make sure to highlight this point.\n\n**References**\n\n[1] Guy Dar, Mor Geva, Ankit Gupta, and Jonathan Berant. 2023. Analyzing Transformers in Embedding Space. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).\n"}]

Author rebuttal:
We thank all reviewers for their insightful comments. We are glad the reviewers found our paper clear and appreciated the contribution of Tracr to studying the computational models of transformers (Reviewers XH5L, ACvV), to advancing interpretability research (Reviewers 767m, PkEE), and as a didactic tool (Reviewer 767m).

We respond to any open questions in the individual responses to the reviewers, and we hope this addresses any remaining uncertainties about our paper. We will be available during the discussion period if any further questions arise.

Meta-review: The paper gives a method for compiling human-readable programs into decoder-only transformers. The structure of the compiled models is used to design experiments (for example, a study of superposition in transformers) and to evaluate interpretability methods. 

The reviewers were uniformly positive about both the paper's contributions and the presentation. There are some questions about how broadly the methods here are applicable to transformers that implement natural tasks; however, the paper is above the bar in spite of these questions. Given this, I am enthusiastically recommending acceptance. Please make sure to incorporate the reviewers' feedback carefully in the final version.
Decision: Accept (spotlight)