Tracr: Compiled Transformers as a Laboratory for Interpretability

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.

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. + 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. + 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.

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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

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.

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.

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. * Overall, the paper is very clearly written * The examples given (Section 4) strongly demonstrate the success of the technique, and broadly help to explain RASP's computational model as well * 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.

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. * The paper assumes a high level of familiarity with prior work, which could be inlined into the paper. * 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). * 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`). * 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.

Questions

* How precisely do the compiled Transformers implement the original algorithms? * 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? * 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?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors adequately discuss limitations in the appendix.

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.

Strengths

**Originality**\ The 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. **Quality/Clarity**\ The 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. **Significance**\ It 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.

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. I'll start with my major qualms with this work: **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? **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. **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. **Minor Issues** 1) 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. 2) 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?

Questions

See weaknesses section for a longer discussion of the key questions.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

There is extensive discussion of various limitations in the Appendix and Conclusion section.

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.

Strengths

- The compilation (section 3) is exceptionally well explained
 - Their approach is especially powerful as a didactic tool
 - The approach can serve as a guide for explainability research, and in providing a “ground truth” for explainability research
 - The open source implementation can serve as a starting point for many researchers interested in explainability of transformer models
 - 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)
 - 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.

Questions

- How do the authors try to approach the mentioned limitations (in their last section) in future work? What is the path from here?
 - 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”?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

See their limitations section and Appendix A2. I feel that the authors have honestly addressed the limitations.

Concerns addressed

Hello, I have read the rebuttal, and I am satisfied with the response. I think the tracr library is a nice contribution, and should lead to important insights on the effectiveness of mechanistic interpretability methods.

Thank the authors for the response

I appreciate the reply. Here is a little concern I still wish to hear your thoughts on: you mentioned "uncompressed Tracr models are structured", then how do you think of the possibility that there is a unbridgeable gap between real Transformer behavior and the Tracr (post-compression) compiled Transformer. If we apply MI techniques to Tracr Transformer and successfully rediscover the Tracr algorithm, does it mean anything that this MI technique can generalize to real Transformer? I.e., should we therefore trust any conclusion this MI get on real Transformer? To be clear, by no means I'm against the acceptation or any reward of this paper. Just curious about how you think about this.