SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Feed-forward layers. Finally, we implement a simple feed-forward layer, designed to correspond to a lookup-table. Each feed-forward layer reads  $\ell$  input variables, which are designated in advance to be either numerical or categorical variables, and outputs one new categorical variable. For example, assuming the residual stream encodes  $m$  categorical variables with cardinality  $k$ , and  $n$  numerical variables. The output of a categorical feed-forward layer is defined as  $\mathbf{z} = \mathrm{One - hot}(\arg \max (f_{\theta}(\mathbf{xW}_{\mathrm{in}}))$ , where  $\mathbf{W}_{\mathrm{in}}\in \mathbb{R}^{(mk + n)\times \ell k}$  is a projection matrix constrained, as above, to read  $\ell$  categorical variables, and  $f:\mathbb{R}^{\ell k\to k}$  is an arbitrary transformation. (We implement  $f$  as an MLP with one hidden layer). This module can be fully described as a look-up table, mapping  $k^{\ell}$  possible inputs to  $k$  outputs. An example of a learned MLP is illustrated in Figure 9.
In this section, we learn Transformer Programs for a variety of tasks, including a simple in-context learning experiment (§4.1); the suite of algorithmic tasks introduced by Weiss et al. [2021] (§4.2); and two NLP tasks: named entity recognition and text classification (§4.3). Additional implementation details for all experiments are provided in Appendix §B, and we provide further ablations and analysis of the generated code in Appendix §C.
4.1 In-context learning
First, by way of illustration, we learn a Transformer Program for a toy task (Figure 5), designed to elicit a simple form of in-context learning [Brown et al., 2020]. The input to the model is a sequence of alternating letters and numbers that end with a letter. If the letter has appeared already in the sequence, the model should output the number that followed it. If the letter has not appeared before, the model should output an unk token. We train an attention-only Transformer
with two layers and one attention head per layer, with a vocabulary of four numbers and four letters, and trained on sequences of up to 10 tokens long. The token and position variables are fixed, one-hot indicators, meaning that the input embeddings encode two variables and the final embeddings encode four: tokens, positions, and one output variable for each attention head. The cardinality of each variable is  $k = 10$ . For this task, we use a causal attention mask.
Input: a1b2b2a Target: 1
Input: b3c4a3c Target: 4
Input: d2c4a2b Target: ?
Figure 5: Sample inputs and targets for a simple in-context learning task.
Dataset	Description	Example	k	L	H	M	Acc.
Reverse	Reverse a string.	reverse("abbc") = "cbba"	8	3	8	2	99.79
Histogram	For each token, the number of occurrences of that letter in the sequence.	hist("abbc") = "1221"	8	1	4	2	100.0
Double hist.	For each token, the number of unique tokens with the same histogram value.	hist2("abbc") = "2112"	8	3	4	2	98.40
Sort	Sort the input in lexicographical order.	sort("cbab") = "abbc"	8	3	8	4	99.83
Most-Freq	The unique input tokens in order of frequency, using position to break ties.	most_freq("abbc") = "bac"	8	3	8	4	75.69
Dyck-1	For each position i, is the input up until i a valid string in Dyck-1 (T); a valid prefix (P); or invalid (F).	dyck1("()") = "PTPTF"	16	3	8	2	99.30
Dyck-2	The same as above, but in Dyck-2.	dyck2("({})({}) = "PPPTPF"	16	3	4	4	99.09
Table 1: The RASP tasks, as introduced by Weiss et al. [2021]. We train Transformer Programs on small-scale instance of each task and report the number of layers (L), attention heads (H), and MLPs (M) used in the best-performing model.  $k$  denotes the variable cardinality, which is fixed at the maximum sequence length for each task.
Results. The learned program achieves perfect accuracy on a held-out test set, and we can observe that the model learns to compose attention heads according to the induction heads pattern identified by Elhage et al. [2021]. The residual stream is illustrated in Figure 2, and the attention heads in Figure 3. The first attention head learns to read the position variable for both key and query. At each position, it attends to the key at the previous position and writes the value of the tokens variable as the value. The second attention head reads tokens as the query and the output of the previous head (head_0_0_outputs) as key. If the query is a letter, the head attempts to map it to a key that has the same letter written to head_0_0 Outputs—that is, a number preceded by the same letter.
4.2 RASP tasks
Next, we test whether we can learn Transformer Programs for a wider variety of algorithmic tasks. We use the tasks introduced by Weiss et al. [2021] to illustrate the RASP language, which are listed in Table 1. We train on small-scale instances of each task, setting the maximum sequence length to 16 for Dyck-1 and Dyck-2, and setting the maximum sequence length and vocabulary size to 8 for the remaining tasks. As above, we use fixed one-hot token and position embeddings and set the variable cardinality  $k$  to be equal to the maximum sequence length. For this setting, we introduce numerical attention and MLPs. We equip each model with an equal number of categorical and numerical attention heads, and categorical and numerical MLPs, fixing the number of MLP input variables to two, and perform a grid-search over the number of layers, attention heads, and MLPs per-layer.
Results. The results of this experiment are in Table 1. On five out of seven tasks, our method is able to find a program that gets more than  $99\%$  accuracy. The exceptions are Double-Histogram (98.4) and Most-Freq (75.69). These results show that, at least on short inputs, Transformer Programs can learn effective solutions to a variety of algorithmic tasks. On the other hand, in Appendix §C.1, we find that the results degrade for longer inputs, suggesting that the model may not learn robust solutions to all tasks. Additional, we can observe that the learned models often use more layers and attention heads than we might expect. For example, Weiss et al. [2021] present RASP solutions for Sort and Reverse requiring only two layers and one attention head, while our best models use three layers and eight attention heads (four categorical and four numerical). These observations indicate that the solutions found by our method might be quite different from human-written solutions.
What kinds of programs does the model learn? To get a better sense for how the model solves these tasks, we examine a program that the model learns for the Sort task. We analyze a model with two layers and four attention heads per layer, which achieves greater than  $99\%$  accuracy, and compile it into a Python program of approximately 300 lines, excluding the classifier weights. We find that the program uses a variety of heuristics for sorting these sequences. For example, Figure 6 depicts the code for one of the first-layer attention heads. This head looks for a 1 or a 4 in the sequence. If it finds a 1, it increases the likelihood of outputting 1 at positions one or two, and if it finds a 4, it increases the likelihood of outputting 4 at positions three through six. In Appendix Fig.11, we use an interactive Python debugger to find a more sophisticated circuit, composing attention heads at two layers. We provide more examples in Appendix §A.2 and analysis of the learned programs in Appendix §C.2. In general, we find that the model learns a variety of non-trivial solutions, composing
# First attention head:
# If q is early in the sequence, look for a 1.
# If q is late in the sequence, look for a 4.
def predicate_0_0(position, token):
    if position in {0} :
        return token == "3"
    if position in {1, 2} :
        return token == "1"
    if position in {3, 4, 5, 6} :
        return token == "4"
    if position in {7} :
        return token == "0"
attn_0_0_pattern = selectclosest(tokens, positions, predicate_0_0)
attn_0_0_outputs = aggregate(attn_0_0_pattern, tokens)
Figure 6: Left: This code corresponds to an attention head in a Transformer Program that was trained to sort sequences of up to eight tokens long, with beginning- and end-of-sequence tokens. At early positions, this attention head checks if there is a 1 in the sequence, and at later positions it looks for a 4. Right: The classifier weights associated with this feature.
operations at different layers to compute higher-order patterns. These solutions include both brittle heuristics and sensible strategies for the different tasks.
Comparison to standard Transformers. In Appendix Fig. 10, we provide additional results comparing Transformer Programs with standard Transformers on RASP tasks with different lengths and vocabulary sizes. Standard Transformers out-perform Transformer Programs across most tasks. $^3$  Interestingly, we observe that both models show similar trends, with model performance degrading for tasks with longer sequences and larger vocabularies. Nevertheless, the Transformer Programs deteriorate more dramatically, with over a  $50\%$  accuracy gap relative to the standard Transformer in the most difficult setting. We discuss these scaling challenges in more detail in Section 6.
4.3 NLP tasks: named entity recognition and text classification
Now, we turn to standard NLP tasks including named entity recognition [Sang and De Meulder, 2003] and text classification [Pang and Lee, 2005, 2004, Zhang et al., 2015]. (Due to space constraints, we present text classification results in Appendix §C.4.) We first experiment with the English portion of the CoNLL-2003 Named Entity Recognition task [Sang and De Meulder, 2003]. This dataset consists of sentences from Reuters news articles (14,041 train/3,250 validation/3,453 test sentences). The data is labeled according to an IOB2 tagging scheme, with each word being assigned to one of four entity types (PER, ORG, LOC, MISC) or marked as not an entity (0). We filter the dataset to sentences with at most 32 words and use a vocabulary of the 10,000 most common words, replacing the remaining words with an unknown token. For these experiments, we use only categorical attention heads, with one categorical MLP per-layer. We fix the variable cardinality at 32 and perform a grid search over number of embedding variables, layers, and heads. The embedding parameters are initialized with 300-dimensional GloVe embeddings [Pennington et al., 2014]. We compare the Transformer Program with a standard Transformer, also initialized with GloVe embeddings with grid search of model dimension and number of layers and heads. More details are provided in Appendix §B.3.
Model	D	L	H	Precision	Recall	F1
Unigram baseline	-	-	-	82.8	55.1	66.2
Standard Transformer	128	3	4	80.4	62.7	70.5
Transformer Program	32×4	2	8	81.0	71.8	76.1
Table 2: Named entity recognition on the CoNLL 2003 shared task. For each method, we report the embedding dimension (D), number of layers (L), number of heads (H) of the best-performing model, and precision, recall and F1. For the Transformer Program, the embedding dimension is equal to the variable cardinality multiplied by the number of input variables. The unigram baseline predicts the class most commonly associated with each word in the training set.
3The exceptions are the histogram-related tasks. This could be because the built-in numerical attention mechanism provides an inductive bias for learning histograms.
(b) Code for the attention features.
(a) Feature weights.
(c) The most common words assigned to different values of the Var3 embedding variable.
Figure 7: We examine how the program distinguishes location entities (B-LOC) from organization entities (B-ORG) by examining the feature weights with the largest gap between the two classes (7a). Many of the top features are components of the word embeddings, but the model also learns to use the attention heads to gather information from the context. For example, attention head attn_0_3 copies one of the embedding variables from the previous position (7b). It promotes the B-LOC label if the value is Var3.V03, which includes prepositions like "at" and "In" (7c). We represent the embedding variables as Python Enum objects, which facilitates analysis using a debugger.
Results. The results are in Table 2. We compare the best-performing Transformer Program with a standard Transformer and with a unigram baseline, which predicts the tag most frequently assigned to each word in the training data. The Transformer Program achieves reasonable performance, on par with the standard Transformer. In particular, the Transformer Program surpasses the unigram baseline, demonstrating that the method learns to make use of contextual information to predict tags.
Interpretability. We illustrate the resulting program by examining how it distinguishes between two commonly confused classes, location entities (B-LOC) and organizations (B-ORG). To see how the model makes this distinction in general, we extract the linear classifier weights and identify the features with the largest difference between the two classes (Figure 7). The highest ranked features include word-level features (that is, components of the word embeddings); position information; and features computed by the attention heads. Working backwards through the program, we find that the model copies information from the neighboring words—for example, increasing the likelihood of B-LOC if the word is preceded by a preposition like "at" or "In".
Learning programs. Our work has precedent in a variety of existing work on program induction and neuro-symbolic methods [e.g. Reed and De Freitas, 2015, Cai et al., 2017, Andreas et al., 2016, Inala et al., 2020, Cranmer et al., 2020, Kim, 2021]. In particular, a long line of work on Inductive Logic Programming [Muggleton and De Raedt, 1994, Cropper and Dumančić, 2022] has sought to learn symbolic logical programs from data, and a number of recent works have used neural networks to search for discrete logical expressions using differentiable reparameterizations [Payani and Fekri, 2019, Wang et al., 2021, Petersen et al., 2022]. We differ from these methods in targeting programs for the Transformer and focusing on sequence modeling problems.
Transformers and formal languages. In addition to RASP [Weiss et al., 2021], a body of research has explored the connection between Transformers and formal languages. Much of this has aimed to
formally characterize the expressivity of Transformers with hard attention [Hahn, 2020, Merrill et al., 2022, Hao et al., 2022]. Giannou et al. [2023] show how Transformers can act as programmable computers by designing a Transformer that can execute programs written in a single-instruction language. Another line of work has attempted to extract deterministic automata, or rules, from neural networks, largely focusing on recurrent neural networks [Jacobsson, 2005, Wang et al., 2018, Weiss et al., 2018]. Merrill and Sabharwal [2022] show theoretically that any fixed-precision Transformer can be expressed as a formula in a type of first-order logic. In contrast to this work, our goal is to design a more interpretable Transformer, rather than extract rules from an existing network.
Interpretable machine learning models. Some prior work has introduced architecture changes aimed at making Transformers more interpretable, including sparse attention [Zhang et al., 2021] and changes to the activation function [Elhage et al., 2022]. These methods make some components of the Transformer qualitatively easier to understand; in contrast, our method results in a model that can be fully described by a discrete program. More generally, a growing body of work has sought to develop intrinsically interpretable machine learning models [e.g. Wang and Rudin, 2015, Chen et al., 2019, Rudin, 2019], a motivation we share in this work.
We introduced a method for training Transformers that can be deterministically mapped to discrete, human-readable programs. We showed that our method is capable of learning effective solutions for a several synthetic and real-world tasks, and that the resulting programs are easy to interpret. In particular, converting Transformers into programs offers practical benefits—for example, we can use standard code-analysis tools to debug errors. On the other hand, learning discrete programs introduces considerable modeling and optimization challenges. We conclude by mentioning some limitations of our method, which we believe represent promising avenues for future work.
Scaling challenges. While we were able to learn a number of effective Transformer Programs, we found that they can struggle to learn on longer inputs, and to learn parsimonious solutions. Are these errors because the framework is not expressive enough, or are they due to optimization issues? Our results suggest that the main issue is optimization. This is evident from the fact that, for many tasks, we can write a Transformer Program that would achieve perfect accuracy, but our method fails to learn these solutions. We conduct a more detailed case study of some optimization failures in Appendix C.3. Therefore, future work is needed to develop better methods for discrete optimization. Additionally, our experiments use a restricted set of modules that may be less expressive than standard Transformers. For example, hard-attention Transformers are less expressive than standard Transformers [Hahn, 2020, Hao et al., 2022]. These restrictions can be relaxed in part by introducing new modules within this framework. We discuss some extensions in Appendix A.4, but leave this direction to future work.
Are the programs interpretable? We have argued that our discrete Transformer Programs are more interpretable than standard Transformers, but one might question whether programs are actually easier to understand, especially as we learn larger programs for more complex tasks. On one hand, these programs are interpretable in the sense that we can trace how information flows between different positions and model components, and we can inspect the program using off-the-shelf code analysis tools, like debuggers. For more complex tasks, the program can be understood as a collection of individually interpretable feature functions, rather than a single interpretable algorithm (as illustrated in Figure 7). Nonetheless, the learned programs can still be complicated and non-intuitive. Another avenue for future work is to explore methods for automatically analyzing the resulting programs, and for imposing an inductive bias in favor of more interpretable programs.
Acknowledgments We thank the members of the Princeton NLP group for their helpful feedback and discussion. This research is funded by the National Science Foundation (IIS-2211779) and a Sloan Research Fellowship.
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 39-48, 2016.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. In Association for Computational Linguistics and International Joint Conference on Natural Language Processing (ACL-IJCNLP), pages 3770-3785, 2021.
Biao Zhang, Ivan Titov, and Rico Sennrich. Sparse attention with linear units. In Empirical Methods in Natural Language Processing (EMNLP), pages 6507-6520, 2021.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. Advances in Neural Information Processing Systems (NIPS), 28, 2015.
A Method details
A.1 Categorical attention
As described in Section 3.2, we implement categorical attention by associating each attention head with a boolean predicate matrix,  $\mathbf{W}_{\mathrm{predicate}} \in \{0,1\}^{k \times k}$ , where  $k$  is the variable cardinality, 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$ . To ensure that each query token attends to a single key token, we use hard attention, defining the attention pattern at position  $i$  as  $\mathbf{A}_i = \text{One-hot}(\arg \max_j \mathbf{S}_{i,j})$ . During training, we define  $\mathbf{A}_i = \text{Gumbel-Softmax}(\mathbf{S}_i)$ .
Defaulting to the beginning of sequence. We implement hard attention so that, in the event that there is no matching key for a query, the model defaults to attend to the first token in the sequence. Let  $\mathbf{S} \in \{0,1\}^{N \times N}$  denote the score matrix for a sequence with length  $N$ . We define a modified score matrix  $\overline{\mathbf{S}}$  such that, for each row  $i$ ,
$$
\bar {\mathbf {S}} _ {i, j} = \left\{ \begin{array}{l l} \mathbf {S} _ {i, j} + (1 - \max  _ {j} \mathbf {S} _ {i, j}) & \text {i f} j = 1 \\ \mathbf {S} _ {i, j} & \text {o t h e r w i s e .} \end{array} \right.
$$
Breaking ties. Furthermore, we implement the attention mechanism so that, in the event that there is more than one matching key, the model attends to the closest match. Given the score matrix  $\mathbf{S} \in \{0,1\}^{N \times N}$ , we define the modified score matrix  $\bar{\mathbf{S}}$  such that, for each row  $i$ ,  $\bar{\mathbf{S}}_{i,j} = \mathbf{S}_{i,j} \times b_{i-j}$  where  $b_{i-j} \in [0,1]$  is a bias associated with the offset between position  $i$  and  $j$ . For most experiments, we fix the bias to decrease from 1, when  $|i-j| = 1$ , to  $1/N$ , when  $|i-j| = N$ , with  $b_0 = 1/N$  to bias the query at position  $i$  against attending to itself. This is similar to types of relative positional bias that have been proposed for regular Transformer-based language models [Press et al., 2022].
A.2 Additional modules
Numerical attention. We implement limited support for numerical variables, designed to ensure that all variables within the program are integers within a bounded range, which allows us to discretize the program by enumerating all possible inputs. First, we include numerical attention heads, which read categorical variables as key and query, and numerical variables as value. The numerical attention head computes attention scores  $\mathbf{S} \in \{0,1\}^{M \times N}$  using the select operator. Unlike in categorical attention, each query can attend to more than one key, and queries can attend to nothing if there is no matching key. 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]$ . At the input layer, there is a single numerical variable, ones, which is frozen and equal to 1 for all positions; attention heads that read ones as the value variable are equivalent to the selector_width primitive in RASP. At higher layers, numerical attention heads can read the output of lower-layer attention heads as values. Figure 8 depicts a numerical attention head from a program for the Double Histogram task.
Numerical attention does not directly correspond to attention in a standard Transformer, because it computes a weighted sum rather than a weighted average. But a numerical attention head can be implemented in a standard Transformer by composing an attention head with a feed-forward layer, using the beginning-of-sequence token to allow the model to attend to nothing. This is how the selector_width primitive is implemented in Tracr [Lindner et al., 2023].