PaperlayerView exact revision on arXiv

Source-linked AI summary

Tracr: Compiled Transformers as a Laboratory for Interpretability

David Lindner, János Kramár, Sebastian Farquhar, Matthew Rahtz, Thomas McGrath, Vladimir Mikulik

arXiv:2301.05062v5cs.LGcs.AIstat.ML

TL;DR

Because learned transformer programs are unknown, it can be difficult to tell whether an interpretation is correct. Tracr compiles human-readable programs into transformers with known structure, enabling experiments and evaluation against known computations. Its compressed sort_unique models approach perfect task accuracy even while their layer-output cosine similarity remains around 0.8.

  • Problem

    Interpretability research lacks ground-truth explanations, making it difficult to determine whether an explanation captures a model’s computation.

    Source: §1 · p. 1

  • Method

    Tracr compiles human-readable RASP programs into transformer weights and compresses compiled models using gradient descent to study superposition.

    Source: §1 · p. 1, §5.1 · p. 7

  • Results

    Around 0.8 average cosine similarity persists across sort_unique model layer outputs despite the compressed model achieving almost perfect task accuracy.

    Source: §5.3 · p. 9

  • Takeaways & Limitations

    Compiled models with known computations can provide test cases for calibrating and evaluating interpretability methods.

    Source: §7 · p. 10

    RASP and Tracr are limited in expressivity, efficiency, and realism, so compiled-model evaluations should be treated as a minimum bar rather than full validation.

    Source: §7 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Tracr compiles human-readable RASP programs into decoder-only transformers with known computational structure. These models support interpretability experiments, including studies of superposition in multi-step algorithms.

1 Introduction

  • Tracr compiles human-readable RASP programs into standard decoder-only transformer models with known computational structure.That structure can serve as ground truth for comparing interpretability tools’ explanations with the models’ mechanisms.Source: §1 · p. 1, Figure 1 · p. 1
  • Tracr offers a response to the lack of ground-truth explanations that makes it difficult to determine whether interpretability methods succeed.The compiler’s models could support future tests of classifier probes, gradient-based attribution, and causal tracing.Source: §1 · p. 1
  • Known-structure models support experiments on superposition in transformers implementing multi-step algorithms.The paper studies models compressed using gradient descent and reports that they drop unnecessary features while representing less important features in superposition.Source: §1 · pp. 1–2
  • Compiled models can also serve as a didactic tool for developing a more concrete understanding of transformer mechanisms.Source: §1 · p. 2

2 Background

The background introduces transformer architecture and RASP, then situates the paper’s superposition study within mechanistic interpretability. It emphasizes the need to understand superposition in larger models performing more complex computations.

2.1 Transformer Architecture

  • Transformers alternate attention and MLP layers with residual connections; attention heads read and write to the residual stream, while MLPs apply nonlinear transformations.The residual stream acts as memory through which earlier layers pass information to later layers.Source: §2.1 · pp. 2–3

2.2 RASP

  • RASP represents sequence computations with s-ops, elementwise operations, and select-aggregate operations that roughly correspond to residual-stream states, MLPs, and attention.Source: §2.2 · p. 3

2.3 Mechanistic Interpretability and Superposition

  • Mechanistic interpretability seeks mechanistic explanations of model behavior, including reverse-engineering how neural networks implement specific behaviors.Source: §2.3 · p. 4
  • Superposition—approximately representing more features than embedding dimensions—remains poorly understood in networks beyond two layers that implement simple algorithms.The paper studies it in larger models because understanding superposition there could advance mechanistic interpretability.Source: §2.3 · p. 4

3 Tracr: A Transformer Compiler for RASP

Tracr compiles RASP programs through computational graphs and intermediate model components into transformer weights. Its design includes constraints and heuristics that shape which programs it supports and how it allocates computations.

3 Tracr: A Transformer Compiler for RASP

  • Tracr restricts RASP programs by disallowing boolean combinations of selectors, requiring annotated categorical or numerical embeddings, and requiring a beginning-of-sequence token.The authors report that programs for all tasks described by Weiss et al. (2021) could be implemented after modifying some programs to use supported features.Source: §3 · p. 4
  • The compiler proceeds through six stages: build a computational graph, infer s-op values, translate s-ops into blocks, assign layers, construct the model, and assemble weights.The RASP-to-Craft-to-weights pipeline is analogous to translating a programming language through assembly into machine code.Source: §3 · p. 4, Figure 3 · p. 4
  • The compiler translates elementwise operations into MLP blocks and select-aggregate operations into attention blocks, using engineered components to approximate functions.Numerical MLP inputs and outputs use piecewise-linear approximations; attention translation supports categorical inputs.Source: §3 · p. 5
  • Tracr assigns computational components to alternating attention and MLP layers using a heuristic based on graph depth and parallelizable blocks.The heuristic returns a correct but sometimes suboptimal layer allocation.Source: §3 · p. 5
  • Tracr constructs a residual space from component input and output spaces, placing each s-op in its own orthogonal subspace before assembling the model weights.Its intermediate Craft representation is independent of a specific transformer implementation, and the authors describe translating compiled models to standard decoder-only transformers without layer norm.Source: §3 · pp. 4–5

4 Exploring Compiled Transformers

The examples show how Tracr compiles counting and sorting programs into transformer components, while additional examples extend its reach to more tasks. The models illustrate computations but are not necessarily realistic.

4 Exploring Compiled Transformers

  • Tracr compiled RASP programs for all tasks described by Weiss et al. (2021), although some programs required modifications to use supported features.The authors characterize the examples as potentially useful for understanding complex computation, while noting the models are not necessarily realistic.Source: §4 · p. 6

4.1 Example 1: Counting Tokens

  • For token counting, an MLP computes the indicator is_x and a causal attention head computes the fraction of previous “x” tokens.The compiled model has a 14-dimensional residual stream, with 12 dimensions used for input embeddings and two for is_x and frac_prevs; its first and last layers are no-ops.Source: §4.1 · p. 6

4.2 Example 2: Sorting

  • For sorting unique tokens, selector_width counts smaller tokens to determine each target position, and a second attention layer moves tokens there.The paper notes that Tracr does not support the composite selector used by Weiss et al. (2021) for duplicates; an alternative adds a small multiple of indices to keys before applying sort_unique.Source: §4.2 · pp. 6–7

4.3 More Examples

  • Tracr’s additional examples include a program for checking balanced parentheses, alongside a library of further programs in the open-source implementation.Source: §4.3 · p. 7

5 Compressing Compiled Transformers

Tracr compression provides a way to study superposition in models running multi-step algorithms, but compressed models can depart from the original computation even when task accuracy remains high.

5 Compressing Compiled Transformers

  • The case studies use frac_prevs and sort_unique compression to explore superposition in multi-step algorithms, while Tracr's orthogonal subspaces for each s-op can make its models sparse and inefficient.The shared projection and frozen weights are a restrictive design choice, though the authors note that other choices could make models more expressive and potentially more realistic.Source: §5 · p. 7, §5.1 · p. 8

5.1 Gradient Descent Based Compression

  • Compression uses one shared projection W across layers, freezes the other weights, and trains W to match model outputs and per-layer outputs in a smaller residual stream.The output loss is cross-entropy for categorical outputs and mean-squared error for numerical outputs; both loss terms receive equal weight.Source: §5.1 · p. 7, Figure 6 · p. 8

5.2 What Does the Compression Learn?

  • The frac_prevs model retains essentially its original performance at d=6 rather than D=14 dimensions, but performance suffers when d is reduced further.Figure 7 compares loss across embedding sizes and shows the learned compression retains task-relevant features unlike PCA.Source: Figure 7 · p. 8, §5.2 · p. 8
  • Known features reveal that tokens:x, is_x, and frac_prevs mostly retain separate dimensions, while unnecessary token features are discarded and sparse indices overlap in the compressed space.is_x and tokens:x also share part of the embedding space because they encode similar information.Source: §5.2 · pp. 8–9

5.3 Do the Compressed Models Still Implement the Same Computation?

  • Some compressed sort_unique models approach 100% accuracy while layer-output cosine similarity stays around 0.8, with a different target_pos encoding that can break fidelity to the original RASP program.The compressed model appears to encode target_pos numerically rather than categorically.Source: §5.3 · p. 9
  • Compressed models can provide a way to investigate unexpected algorithmic representations, although their annotations may cease to be exact ground truth when their computations diverge from the compiled models.The authors suggest studying these discrepancies to learn how neural networks represent computations without reverse-engineering entire models.Source: §5.3 · p. 9

6 Related Work

Related work includes mechanistic reverse engineering, alternative transformer construction approaches, and efforts to create ground truth for interpretability evaluation. The paper positions Tracr as a complementary approach that compiles fully human-readable code into transformers.

6 Related Work

7 Discussion & Conclusion

Tracr compiles human-readable code into transformer weights, providing a testbed for interpretability research. The paper outlines applications and cautions that compiled-model evaluations are a minimum bar, not full validation.

7 Discussion & Conclusion

  • Tracr compiles human-readable code into transformer weights as a testbed for developing and evaluating interpretability tools.The authors also provide an open-source implementation.Source: §7 · p. 10, A Applications and Limitations of Tracr · p. 14
  • Compiled models could support interpretability benchmarks, component replacement, study of isolated phenomena, and stronger intuitions about transformer algorithms.These applications include using known computations to calibrate and evaluate methods.Source: §7 · p. 10
  • Tracr has limitations, and experiments on compiled models should be interpreted carefully because they do not fully represent real transformers.The authors recommend treating evaluations as a minimum bar rather than full validation of a technique.Source: §7 · p. 10
  • The authors see Tracr as a promising approach for studying transformers and evaluating interpretability tools against models whose computations are known.They compare this calibration opportunity to testing a microscope lens on familiar, well-understood shapes.Source: §7 · p. 10

A.1 Applications of compiled models in interpretability research

Compiled models offer controlled settings for studying model phenomena and evaluating interpretability methods. The paper describes applications ranging from superposition experiments and explanation tests to intuition-building about transformer algorithms.

A.1 Applications of compiled models in interpretability research

  • Compiled models let researchers set up controlled experiments on transformer computational structure that may otherwise be intractable.They serve as a laboratory for interpretability research.Source: §A.1 · p. 14

Understanding model phenomena and developing new techniques

  • Compressing compiled models can induce superposition, offering a testbed for studying how computational structure affects which features are stored that way.The authors also describe predicting compression outcomes and reversing induced superposition as possible research directions.Source: §A.1 · p. 14

Test cases for interpretability tools

  • Compiled models provide a basis for testing whether interpretability explanations faithfully describe model computation and for falsifying incorrect explanations.The paper discusses classifier probes as one example of a method whose results can be difficult to interpret.Source: §A.1 · p. 14
  • Libraries of compiled-model test cases could enable quantitative evaluation metrics for interpretability tools.Source: §A.1 · p. 14

Replacing model components

  • Compiling hand-coded components could support evaluating understanding by replacing parts of a trained model with idealized implementations.The authors cite Nanda et al. (2023), who found that component replacement increased downstream performance for modular addition.Source: §A.1 · p. 14

Building intuition for algorithms implementable by transformers

  • Tracr spells out the relationship between a program and its transformer implementation, but real learned models may use a greater variety of algorithms.Source: §A.1 · p. 14

A.2 Limitations of RASP and Tracr

RASP and Tracr trade expressivity, efficiency, and realism for models with known structure. The appendix also describes constraints on RASP selectors and encodings, and explains how to read compiled-model forward-pass figures.

A.2 Limitations of RASP and Tracr

  • RASP targets discrete sequence outputs and binary attention, so it cannot model probabilistic token computations and supports a limited range of algorithms.Numeric attention patterns are discussed as a possible extension.Source: §A.2 · p. 15
  • Tracr reserves orthogonal residual-stream subspaces for variables across all layers, even when a variable is used only in part of the computation.Real models use more compressed representations and likely reuse dimensions for multiple features.Source: §A.2 · p. 15
  • Tracr's hand-coded matrices and computational-basis-aligned features make its models unrealistic and easy to interpret by inspecting residual-stream activations.The authors suggest learning layers in isolation, matrix factorization, or basis rotation as possible remedies.Source: §A.2 · p. 15
  • Compiled models are simpler than learned models and may not generalize; evaluations should be treated as a minimum bar, not full validation.They are more likely to be interpretable and may omit features present in real models.Source: §A.2 · p. 15

B Modifications to RASP

  • RASP selectors are restricted to two input variables because arbitrary compositions lack a straightforward efficient implementation; the restriction did not prevent implementing the tasks described by Weiss et al. (2021).Refactoring composites into atomic selectors can be inefficient and would break the one-to-one correspondence between RASP selectors and compiled attention heads.Source: B Modifications to RASP · pp. 15–16
  • Tracr encodes categorical variables as one-hot vectors and numerical variables as magnitudes in dedicated residual-stream subspaces.Some aggregate operations, such as mean across token positions, are natural for numerical encodings but not categorical ones.Source: B Modifications to RASP · p. 16
  • A mandatory BOS token serves as a default attention position when a RASP selector would otherwise produce an all-zero row, which softmax attention cannot represent.Source: B Modifications to RASP · p. 16

C Reading Model Output Figures

  • Forward-pass figures show residual-stream contents after successive layers as heatmaps, with sequence positions on the x-axis and residual dimensions on the y-axis.Red borders mark entries changed by the corresponding layer; labels distinguish categorical dimensions from numerical ones.Source: C Reading Model Output Figures · pp. 16–17

D Tracr Implementation Details

Tracr translates RASP operations into transformer components, with implementations for MLPs, attention, selector width, and causal attention. Causal evaluation changes RASP semantics and can require adapting programs.

D Tracr Implementation Details

  • Tracr's implementation constructs MLP and attention blocks, implements selector width, and supports causal attention.The full implementation and documentation are in the code repository.Source: D Tracr Implementation Details · p. 17

D.1 MLP and Attention Blocks

  • Categorical Map operations become lookup-table MLPs, while numerical MLPs discretize known inputs and map each bucket to an output to minimize approximation error.SequenceMaps with categorical inputs and outputs first encode input pairs, then apply a lookup table.Source: §D.1 · p. 17
  • Attention blocks derive their attention pattern from RASP selectors, with a BOS rank-one update ensuring BOS is attended to always or when no other token is selected.The attention matrix is scaled by an inverse-temperature parameter, 100 by default, to produce weights very close to 1/#{selected tokens} on selected tokens and 0 elsewhere.Source: §D.1 · pp. 17–18

D.2 Selector Width Primitive

  • Tracr implements selector width with an attention layer that computes x = 1/(1 + w), followed by an MLP that recovers w = 1/x − 1 and cleans the BOS position.Source: §D.2 · p. 18

D.3 Causal Attention

  • Causal evaluation applies a causal mask to each selector and changes some RASP operation semantics, so programs may need adapting.For example, frac_prevs no longer needs a manual causal mask, while the cited length implementation no longer computes sequence length correctly because it needs future tokens.Source: §D.3 · p. 18

E Compression Training Details

Compression models were trained in Jax with AdamW, using a fixed batch size and a linearly decayed learning rate; each run took 1–4 hours on two CPU cores.

E Compression Training Details

F Theoretical Results on Combining Attention Heads

The section analyzes when selectors can be combined through Boolean operators in attention. It shows that general composition is constrained, while matching query and key s-ops offer a workaround with a dimensionality cost.

F Theoretical Results on Combining Attention Heads

G More Compiled Models

The section presents compiled models for sorting and parenthesis tasks. These examples extend sorting to duplicate keys and use a parenthesis-balance computation as a subroutine for checking balanced sequences.

G More Compiled Models

Related papers

Loading 2301.05062v5…