SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dan Friedman Alexander Wettig Danqi Chen
Department of Computer Science & Princeton Language and Intelligence  
Princeton University
{dfriedman,awettig,danqic}@cs.princeton.edu
Recent research in mechanistic interpretability has attempted to reverse-engineer Transformer models by carefully inspecting network weights and activations. However, these approaches require considerable manual effort and still fall short of providing complete, faithful descriptions of the underlying algorithms. In this work, we introduce a procedure for training Transformers that are mechanistically interpretable by design. We build on RASP [Weiss et al., 2021], a programming language that can be compiled into Transformer weights. Instead of compiling human-written programs into Transformers, we design a modified Transformer that can be trained using gradient-based optimization and then automatically converted into a discrete, human-readable program. We refer to these models as Transformer Programs. To validate our approach, we learn Transformer Programs for a variety of problems, including an in-context learning task, a suite of algorithmic problems (e.g. sorting, recognizing Dyck-languages), and NLP tasks including named entity recognition and text classification. The Transformer Programs can automatically find reasonable solutions, performing on par with standard Transformers of comparable size; and, more importantly, they are easy to interpret. To demonstrate these advantages, we convert Transformers into Python programs and use off-the-shelf code analysis tools to debug model errors and identify the "circuits" used to solve different sub-problems. We hope that Transformer Programs open a new path toward the goal of intrinsically interpretable machine learning.<sup>1</sup>
Transformers [Vaswani et al., 2017] have become the predominant neural network architecture in machine learning, representing the state-of-the-art in natural language processing (NLP) and increasingly in computer vision, molecular modeling and other domains. Prominently, large Transformer-based language models [LLMs; Brown et al., 2020] have demonstrated impressive general-purpose capabilities and are now widely deployed as components in user-facing AI applications such as chatbots, search engines, and code assistants. However, these systems are fundamentally limited by a lack of interpretability, which makes them difficult to audit, debug, and maintain. This black-box quality poses practical challenges and limits the usefulness of these models in high-stakes applications.
As a result, a considerable body of work has aimed at improving our understanding of Transformers. In NLP, much of this work has focused on pre-trained Transformer language models such as BERT [Devlin et al., 2019], using a variety of post-hoc methods, including analyzing attention patterns [Clark et al., 2019] and probing hidden representations [e.g., Tenney et al., 2019, Belinkov, 2022]. These post-hoc approaches can provide partial insight into model behavior but have also been shown to be misleading [Jain and Wallace, 2019, Bolukbasi et al., 2021]; in any case, they do not provide a complete or faithful description of the algorithm that the model implements. More recently,
<sup>1</sup>Our code is available at https://github.com/princeton-nlp/TransformerPrograms, along with a number of example Transformer Programs.
Figure 1: We design a modified Transformer that can be trained on data and then automatically discretized and converted into a human-readable program. The program is functionally identical to the Transformer, but easier to understand—for example, using an off-the-shelf Python debugger.
research on mechanistic interpretability has attempted to gain an algorithmic understanding of these models, with the goal of reverse-engineering Transformers into human-interpretable components. This line of work includes efforts to characterize "circuits" of attention heads [Elhage et al., 2021, Olsson et al., 2022, Wang et al., 2023, Nanda et al., 2023]; align network representations with symbolic causal models [Geiger et al., 2023]; interpret feed-forward layers [Nostalgebraist, 2020, Geva et al., 2022]; and localize and edit factual information [Meng et al., 2022]. However, such methods require extensive manual effort and can be impeded by the inherent complexity of the underlying models [e.g. McGrath et al., 2023].
In this work, instead of attempting to explain black-box models, we aim to train Transformers that are mechanistically interpretable by design. We take inspiration from RASP [Weiss et al., 2021], a programming language for characterizing Transformer operations, and Tracr [Lindner et al., 2023], a compiler for converting RASP programs into Transformer networks. RASP provides a conceptual framework for mapping between Transformer components and human-readable code. However, whereas prior work has compiled human-written programs into Transformers, our goal is to train Transformers using gradient-based optimization and then automatically decompile them into human-readable programs.
To this end, we propose a method for learning Transformer Programs—Transformers that are constrained to implement human-interpretable algorithms. First, we introduce a set of constraints that restrict the Transformer to lie within an interpretable subspace of the parameter space, by which we mean the subset of parameter values that can be mapped to a family of human-readable programs. Second, we develop a continuous relaxation scheme for learning these programs. Transformer Programs can be deterministically mapped to a program in a human-readable programming language, like Python. The Python program is functionally identical to the Transformer, but considerably easier to interpret—for example, we can debug errors by setting breakpoints using the standard Python debugger. Our overall approach is illustrated in Figure 1.
We validate our approach by learning Transformer Programs for a variety of problems, including an in-context learning task; the set of algorithmic problems introduced by Weiss et al. [2021]; and NLP benchmarks for named entity recognition and text classification. While discrete optimization introduces some challenges, we show that the Transformer Programs achieve reasonable performance relative to standard Transformers of comparable size, and, unlike standard Transformers, they are easy to interpret. We illustrate this point by converting Transformers into Python programs and using off-the-shelf code analysis tools to extract the "circuits" used to solve different sub-problems, and to debug model errors. Overall, we aim to demonstrate that Transformer Programs represent a promising first step towards building intrinsically interpretable machine learning models.
Transformer architecture. The Transformer [Vaswani et al., 2017] is a neural network architecture for processing sequential data. The input to the Transformer is a sequence of tokens  $\mathbf{w} = w_{1},\ldots ,w_{N}\in \mathcal{V}$  in a discrete vocabulary  $\mathcal{V}$ . At the input layer, the model maps the tokens to a sequence of  $d$ -dimensional embeddings  $\mathbf{x}_0\in \mathbb{R}^{N\times d}$ . In the standard Transformer, this initial embedding is defined as the sum of a learned token embedding and a positional embedding. Each subsequent layer  $i$  consists of a multi-head attention layer (MHA) followed by a multilayer perceptron layer (MLP):  $\mathbf{x}_i = \mathbf{x}_{i - 1} + \mathrm{MLP}_i(\mathbf{x}_{i - 1} + \mathrm{MHA}_i(\mathbf{x}_{i - 1}))$ . Following the presentation
2The standard Transformer also includes a layer-normalization layer [Ba et al., 2016], which we omit here.
of Elhage et al. [2021], multi-head attention can be written as
$$
\mathrm {M H A} (\mathbf {x}) = \sum_ {h = 1} ^ {H} \mathrm {s o f t m a x} (\mathbf {x} \mathbf {W} _ {Q} ^ {h} (\mathbf {x} \mathbf {W} _ {K} ^ {h}) ^ {\top}) \mathbf {x} \mathbf {W} _ {V} ^ {h} \mathbf {W} _ {O} ^ {h},
$$
where  $H$  is the number of heads,  $d_h$  is the attention embedding dimension,  $\mathbf{W}_Q^h$ ,  $\mathbf{W}_K^h$ ,  $\mathbf{W}_V^h \in \mathbb{R}^{d \times d_h}$  are referred to as the query, key, and value projections respectively, and  $\mathbf{W}_O^h \in \mathbb{R}^{d_h \times d}$  projects the output value back to the model dimension. By default, we assume that attention is bi-directional, corresponding to a Transformer encoder, but we also support causal masking. The MLP layer operates at each position independently; in the standard Transformer, it is a two-layer feedforward network:  $\mathrm{MLP}(\mathbf{x}) = \sigma(\mathbf{x} \mathbf{W}_1) \mathbf{W}_2$ , where  $\mathbf{W}_1 \in \mathbb{R}^{d \times d_m}$ ,  $\mathbf{W}_2 \in \mathbb{R}^{d_m \times d}$ , and  $\sigma$  is a non-linear activation, commonly the ReLU function. The output of the model is a sequence of token embeddings,  $\mathbf{x}_L \in \mathbb{R}^{N \times d}$ . The model is typically trained end-to-end on a prediction task by applying a linear classifier to the final-layer token embeddings.
Transformer circuits. Transformer circuits [Elhage et al., 2021] are an abstraction for characterizing how a neural network processes information. Informally, a Transformer can be viewed as a series of nodes (attention heads and feed-forward layers) reading and writing information to the residual stream, which refers to the sequence of token embeddings  $\mathbf{x}_0,\dots ,\mathbf{x}_L\in \mathbb{R}^{N\times d}$  computed at each layer. Each node "reads" from the residual stream via a projection matrix  $\mathbf{W}_{\mathrm{in}}\in \mathbb{R}^{d\times d_h}$ ; calculates some function  $f(\mathbf{x}_{i - 1}\mathbf{W}_{\mathrm{in}})$ ; and then "writes" to the residual stream via another projection:  $\mathbf{x}_i = \mathbf{x}_{i - 1} + f(\mathbf{x}_{i - 1}\mathbf{W}_{\mathrm{in}})\mathbf{W}_{\mathrm{out}}$ , with  $\mathbf{W}_{\mathrm{out}}\in \mathbb{R}^{d_h\times d}$ . Mechanistic interpretability seeks to find the circuits within a Transformer that implement individual behaviors, but this goal is often difficult in practice; for example, different attention heads may write to the same subspace of the residual stream, making it difficult to disentangle how information flows from one component to another.
Programming languages for Transformer. RASP [Weiss et al., 2021] is a programming language consisting of a set of function primitives for defining operations on sequences. These are designed to ensure that any program written in RASP can be implemented by a Transformer network, by mapping between RASP primitives and Transformer components. Tracr [Lindner et al., 2023] is a compiler for RASP that implements this mapping in practice. The central operation is select, which corresponds to attention in the Transformer. select takes as input a sequence of keys  $\in \mathcal{K}^N$ ; a sequence of queries  $\in \mathcal{Q}^M$ ; and a boolean predicate  $\in \mathcal{Q} \times \mathcal{K} \rightarrow \{0,1\}$ . The output is an attention matrix  $\mathbf{A} \in \{0,1\}^{M \times N}$  with  $\mathbf{A}_{i,j} = \text{predicate}(\text{queries}_i, \text{keys}_j)$ . The select operator is combined with an aggregate operation, which takes as input an attention matrix  $\mathbf{A}$  and a sequence of values  $\in \mathcal{V}^N$  and outputs, at each row  $i$ , the weighted average of values, with weights given by  $\mathbf{A}_i$ . RASP also supports arbitrary element-wise operations, which correspond to the feed-forward layers in a Transformer. RASP and Tracr provide a framework for mapping from human-interpretable programs to Transformers, but there is no clear way to map an arbitrary Transformer to a RASP program.
3 Learning Transformer Programs
Our main idea is to define constraints on the Transformer weights that guarantee that there is a simple, deterministic mapping between Transformer components and RASP-like programming primitives. We refer to the set of Transformers satisfying these constraints as Transformer Programs. To provide an overall picture our approach, we illustrate our method with a minimal Transformer Program, containing only categorical attention layers. This model was trained to perform a simple in-context learning task (described in more detail in Section 4.1) and is depicted in Figures 2 and 3. In Section 3.3, we show we extend our framework to support a broader set of operations.
3.1 Disentangled residual stream
Our first constraint is to ensure that the model maintains a disentangled residual stream. This means that the token embeddings encode the values of a fixed set of named variables, with each variable encoded in an orthogonal subspace. This constraint is illustrated in Figure 2. In this example, the initial input embedding encodes the values of two categorical variables, corresponding to the token and position embeddings. Each attention layer then reads a fixed set of variables from the residual stream and writes a new variable to a dedicated address.
Reading from the residual stream. To ensure that each module reads a fixed set of variables, we parameterize each projection matrix by a one-hot indicator vector over the available variables.
Figure 2: We constrain the Transformer to have a disentangled residual stream: the token embeddings encode a fixed set of discrete variables in orthogonal subspaces, and each module reads a fixed set of variables and writes a new variable to a dedicated address. Here, we depict the residual stream of a Transformer Program that was trained on a simple in-context learning task (Section 4.1).
Suppose the residual stream encodes  $m$  categorical variables, all with cardinality  $k$ , resulting in input embeddings  $\mathbf{x} \in \{0,1\}^{N \times mk}$ . Then each projection matrix  $\mathbf{W} \in \mathbb{R}^{mk \times k}$  is defined by an indicator  $\pi \in \{0,1\}^m : \sum_{i=1}^m \pi_i = 1$ . That is,  $\mathbf{W} = [\pi_1 \mathbf{I}_k; \dots; \pi_m \mathbf{I}_k]^\top$ , where  $\mathbf{I}_k$  is the  $k \times k$  identity matrix. Concretely, this means that each attention head is associated with three gate vectors,  $\pi_K, \pi_Q, \pi_V$ , defining the key, query, and value variables, respectively.
Writing to the residual stream. To ensure that the residual stream remains disentangled, we constrain each module to write its output to a dedicated, orthogonal subspace, taking inspiration from Tracr. In practice, we accomplish this simply by concatenating the module output to the input—i.e., if layer  $i$  consists of a single attention head  $h$ , then  $\mathbf{x}_i = [\mathbf{x}_{i-1}; h(\mathbf{x}_{i-1})]$ . (This is equivalent to adding the output to the input after padding with zeros.) The output of each attention head is another categorical variable with cardinality  $k$ , so the final embedding dimension is then  $(2 + L \times H) \times k$ , where  $L$  is the number of layers and  $H$  is the number of heads per layer.
3.2 Transformer Program modules
Next, we constrain each module to implement an interpretable, rule-based mapping between inputs and outputs. Here, we describe the primary modules in our programs, categorical attention heads; we extend this framework to additional modules in Section 3.3. Categorical attention heads can be decomposed into two operations, corresponding to the select and aggregate operations in RASP: first determining which queries attend to which keys, and then aggregating the corresponding values.
Computing the attention pattern. First, each attention head reads one key and one query variable, and then determines which queries should attend to which keys. In RASP, this operation is defined via a boolean predicate, which maps every combination of key and query to a value in  $\{0,1\}$ . We implement a simple learnable predicate by associating each attention head with a one-to-one mapping between query values and key values (Figure 3). Assuming that all variables have cardinality  $k$ , each attention head is associated with a binary predicate matrix,  $\mathbf{W}_{\mathrm{predicate}} \in \{0,1\}^{k \times k}$ , with the constraint that each row  $\mathbf{W}_{\mathrm{predicate},i}$  sums to one. The self-attention pattern is then defined by a score matrix  $\mathbf{S} \in \{0,1\}^{N \times N}$ , with  $\mathbf{S} = \mathbf{x}\mathbf{W}_Q\mathbf{W}_{\mathrm{predicate}}(\mathbf{x}\mathbf{W}_K)^\top$ .
Aggregating categorical variables. We adopt a constraint from Tracr and require that each query token attends to a single key token; this constraint is necessary to ensure that the output variable will also be categorical. We enforce this by using hard attention, defining the output of head  $h$  as  $\mathbf{A}_i = \mathrm{One - hot}\left(\arg \max_j\mathbf{S}_{i,j}\right)$ , where  $\mathbf{A}_i$  denotes the  $i^{th}$  row of the attention matrix. We implement hard attention so that, in the event that there is no matching key for a query, the model defaults to
Figure 3: We further constrain each module to implement an interpretable, rule-based mapping between input and output variables. This figure depicts the categorical attention module, which defines a one-to-one correspondence between key and query variables. In this example, the first attention layer outputs the value of the tokens variable from the preceding position. The second attention layer uses this output to implement a simple induction head [Olsson et al., 2022], mapping each letter token to a number that is preceded by the same letter.
Figure 4: After training, we convert the model into an equivalent Python program, modeled on RASP. This code corresponds to the attention heads illustrated in Figure 3.
def predicate_1(q_position, k_position):
    if q_position in {0, 8}:
        return k_position == 7
    if q_position in {1, 9}:
        return k_position == 0
    if q_position in {2}:
        return k_position == 1
    if q_position in {3}:
        return k_position == 2
    if q_position in {4}:
        return k_position == 3
    if q_position in {5}:
        return k_position == 4
    if q_position in {6}:
        return k_position == 5
    if q_position in {7}:
        return k_position == 6
def predicate_2(token, attn_1_output):
    if token in {
        "<pad>", "0", "1", "2", "3", "<s"]
    }:
        return attn_1_output == "<pad">
    if token in {"a":
        return attn_1_output == "a"
    if token in {"b":
        return attn_1_output == "b"
    if token in {"c":
        return attn_1_output == "c"
    if token in {"d":
        return attn_1_output == "d"
    attn_2_pattern = selectclosest(
        attn_1_outputs, tokens, predicate_2)
    attn_2.outputs = aggregate(attn_2_pattern, tokens)
attn_1_pattern = selectclosest( positions,positions,predicate_1) attn_1_outputs = aggregate(attn_1_pattern, tokens)
attend to a beginning of sequence token; in the event that there is more than one matching key, the model attends to the closest match. More details are provided in the Appendix A.1.
**Optimization.** Consider a Transformer Program with a single categorical attention head, with  $m$  input variables of cardinality  $k$ . This model is defined by the indicators  $\pi_K, \pi_Q, \pi_V$  and the predicate matrix  $\mathbf{W}_{\text{predicate}}$ . To learn Transformer Programs, we optimize a distribution over these discrete program weights. For each gate  $\pi_K, \pi_Q, \pi_V$ , we define a parameter  $\phi_K, \phi_Q, \phi_V \in \mathbb{R}^m$ . For row in the predicate matrix  $\mathbf{W}_{\text{predicate}}$ , we define parameters  $\psi_1, \ldots, \psi_k \in \mathbb{R}^k$ , with  $\psi_i$  defining a distribution over the  $i^{th}$  row. Referring to these parameters collectively as  $\Phi$ , we define a distribution over discrete Transformer weights  $p(\theta \mid \Phi)$  as the product of categorical distributions, which we optimize using the Gumbel reparameterization [Jang et al., 2017]. Given a classification dataset  $\mathcal{D}$  consisting of inputs  $w$  and labels  $y$ , we minimize the loss
$$
\mathcal {L} = \mathbb {E} _ {\boldsymbol {\theta} \sim p (\boldsymbol {\theta} | \boldsymbol {\Phi})} \left[ \mathbb {E} _ {\boldsymbol {w}, y \sim \mathcal {D}} \left[ - \log p (y \mid \boldsymbol {w}; \boldsymbol {\theta}) \right] \right] \approx \frac {1}{S} \sum_ {s = 1} ^ {S} \mathbb {E} _ {\boldsymbol {w}, y \sim \mathcal {D}} \left[ - \log p (y \mid \boldsymbol {w}; \boldsymbol {\theta} _ {s}) \right],
$$
where  $\theta_{s}$  is a sample from the product of Gumbel-Softmax distributions, with softmax temperature  $\tau > 0$ , and  $S$  is a hyper-parameter denoting the number of samples per training step. Transformer Programs also include a discrete arg max operations in hard attention, which we also relax using the Gumbel-Softmax. We anneal the Gumbel temperature over the course of training. As  $\tau \to 0$ , the samples are identical to one-hot samples of the corresponding categorical distribution.
Extracting programs. After training, we obtain a discrete model by taking the maximum likelihood weights,  $\theta^{*} = \arg \max_{\theta}p(\theta \mid \Phi)$ , and applying the arg max operation for hard attention. We then implement a deterministic mapping to convert the model to a Python program, inspired by RASP (Figure 4). Specifically, we convert each attention head  $h$  into a predicate function, which maps each possible pair of query and key values to a value in  $\{0,1\}$ . The attention operation is then
implemented by a library function, selectclosest, which selects the closest key satisfying the predicate, or the first token in the sequence if there is no match.
This framework can be extended to support additional modules, provided that the module can be mapped to a program and optimized effectively. We introduce three extensions here, providing more details and examples in Appendix A.2.
Learning word embeddings. For synthetic tasks with a small vocabulary (i.e.,  $|\mathcal{V}| < 100$ ), we use fixed, one-hot token embeddings. For larger-scale NLP tasks, we learn a factored categorical embedding, representing each token as the product of  $m$  categorical variables, where  $m$  is another hyperparameter. That is, each word embedding is the concatenation of  $m$ $k$ -dimensional one-hot embeddings. After training, these variables can be understood by inspecting the set of words taking on each value, which we illustrate in Figure 7.
Aggregating numerical variables. We also equip the model with a form of numerical attention, to facilitate operations like counting. We augment the input embeddings with a single scalar variable, which is fixed to be equal to one. In order to ensure that the resulting program is discrete, we implement a limited form of numerical attention, which guarantees that the output values are integers with a bounded range. This enables us to characterize downstream modules by enumerating the possible input values. This module loosely corresponds to the selector_width primitive from RASP. In Tracr and RASP, selector_width is implemented using one attention layer and one feed-forward layer. We "hard-code" it in the form of an attention-like head that reads categorical variables as key and query and a numerical variable as the value. As above, the module learns a binary predicate matrix mapping queries to keys, but aggregates values by taking a weighted sum. That is, given attention scores  $\mathbf{S} \in \{0,1\}^{M \times N}$  and value variable var, the output for the  $i^{th}$  token is defined as  $\sum_{j=1}^{N} \mathbf{S}_{i,j} \text{var}[j]$ .