SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Feed-forward layers. In RASP, feed-forward layers are used to implement arbitrary elementwise operations, and they play an important role in many human-written RASP programs. In our Transformer Programs, we restrict the capacity of the feed-forward layers to ensure that they can be decomposed. 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. We convert feed-forward layers to programs by enumerating all possible inputs and forming a lookup-table. For all experiments, we set  $\ell = 2$ . For categorical attention, given variable cardinality  $k$ , there are  $k^2$  possible inputs. An example of a feed-forward layer is given in Figure 9, which depicts a circuit in a program for the Dyck-2 task.
For numerical attention, we can bound the input range by determining the cardinality of the numerical attention heads. For each numerical attention head, the minimum output value is equal to 0 and the maximum output value is equal to the maximum input length multiplied by the cardinality of the input variable (that is, the value variable). For example, if an attention head reads ones as the value
def num=predicate_0_1(q_token, k_token):
    if q_token in {"0"}:
        return k_token == "0"
    elif q_token in {"1"}
        return k_token == "1"
    elif q_token in {"2"}
        return k_token == "2"
    elif q_token in {"3"}
        return k_token == "3"
    elif q_token in {"4"}
        return k_token == "4"
    elif q_token in {"5"}
        return k_token == "5"
    elif q_token in ["
"]
        return k_token == "<pad>"
def num_mlp_0_1(num_attn_0_1_output
key = num_attn_0_1_output
if key in \{0, 1\}:
return 4
return 0
num_mlp_0_1_outputs = [ 
    num_mlp_0_1(kO)
    for k0 in num_attn_0_1_outputs]
Figure 8: An example of a numerical attention head and MLP in a program for the Double Histogram task. For this task, the model must output, for each position, the number of unique tokens with the same histogram value. In this example, an attention head (left) calculates the histogram for each position. An MLP (top right) reads the histogram values and outputs a value of 0 if the histogram value is greater than one, and 4 otherwise. Inspecting the corresponding classifier weights (bottom right), we see that an output value of 0—meaning a histogram count greater than 1—increases the likelihood that the double-histogram value is 1 or 2, and decreases the likelihood of larger values. Because the input length is limited to 8, this reflects the fact that if one number appears many times, it is unlikely that another number appears the same number of times. An output of 4 (meaning a histogram count of 1) increases the likelihood that the double-histogram is greater than 1. Note that we configure all MLPs to read two input variables, but some MLPs learn to read the same variable for both inputs, as in this example. This allows us to compress the corresponding function.
variable, the maximum output value is equal to the maximum input length. If an attention head reads the output of a first-layer attention head as value, the maximum output value is equal to the square of the maximum input length.
A.3 Extracting programs
In this section, we provide more details about our procedure for converting our trained models into Python programs. While there are many possible ways to express the discretized model as Python code, we choose a mapping that facilitates analysis with a standard Python debugger. We use several simple strategies to improve the readability of the code, by annotating variable types, compressing statements, and removing unreachable branches. Illustrative programs are depicted in Figures 8 and 9.
Attention. Each attention head is represented by a predicate function, which takes as input a key and query and outputs a value in  $\{0,1\}$ . In Transformer Programs, all keys and queries are categorical variables with cardinality  $k$ , so the predicate function can be defined by enumerating the possible query values, which we do with a series of if statements. Additionally, if multiple query values are mapped to the same key, we condense the predicate by combining these queries into a single branch. This is illustrated in Figure 9. In the first attention head (left), each query position attends to the previous position (with the exception of the first token), so we cannot apply any compression. In the second attention head (bottom right), fifteen out of the sixteen possible query values are mapped to a single key value, so we combine them into a single branch.
MLPs. We convert feed-forward modules to functions by enumerating all possible inputs and forming a key/value lookup-table. For all experiments, we set each MLP to read  $\ell = 2$  input variables. In some cases, the MLP learns to read the same variable for both input variables, allowing us to reduce the length of the function. We further reduce the length of these functions by identifying the MLP output value associated with greatest number of keys, and returning this value as the default value without explicitly evaluating the corresponding condition. These two forms of compression are illustrated in Figure 8 (top right). This MLP reads a single numerical input variable and outputs a value of 4 if the input is 0 or 1 and a value of 0 otherwise.
def predicate_0_0(q_position, k_position):
    if q_position in {0, 13):
        return k_position == 12
    elif q_position in {1):
        return k_position == 0
    elif q_position in {2):
        return k_position == 1
    elif q_position in {3):
        return k_position == 2
    elif q_position in {4):
        return k_position == 3
    elif q_position in {5):
        return k_position == 4
    elif q_position in {6):
        return k_position == 5
    elif q_position in {7):
        return k_position == 6
    elif q_position in {8):
        return k_position == 7
    elif q_position in {9):
        return k_position == 8
    elif q_position in {10):
        return k_position == 9
    elif q_position in {11):
        return k_position == 10
    elif q_position in {12):
        return k_position == 11
    elif q_position in {14):
        return k_position == 13
    elif q_position in {15):
        return k_position == 14
attn_0_0_pattern = selectclosest positions, positions, predicate_0_0)
attn_0_0_outputs = aggregate(attn_0_0_pattern, tokens)
MLP: reads current token and previous token
# Outputs 13 if it sees "({}" or "[")."
def mlp_0_0(token, attn_0_0_output):
    key = (token, attn_0_0_output)
    if key in {"(")", "}", (""), [""], [""]
    ("{", ["}", ["]", [""], [""]
    ("{", ["}")": return 4
    elif key in {"("", {""}, {""},
    [""], [""]
    return 13
    elif key in {"("", {""),
    [""], [""]
    [""]
    [""]
    [""]
    return 0
    return 7
mlp_0_0_outputs = [
        mlp_0_0(k0, k1) for k0, k1 in
        zip(tokens, attn_0_0.Outputs)
    ]
# 2nd layer attention: check for "({}" or "[")
def predicate_1_2(position, mlp_0_0_output):
    if position in {0, 1, 2, 4, 5, 6, 7, 8, 9,
        10, 11, 12, 13, 14, 15}: return mlp_0_0_output == 13
    elif position in {3}: return mlp_0_0_output == 4
attn_1_2_pattern = selectclosest( mlp_0_0.Outputs, positions, predicate_1_2)
attn_1_2.Outputs = aggregate( attn_1_2_pattern, mlp_0_0.Outputs)
Figure 9: An example of a circuit in a program for Dyck-2. For this task, the inputs consist of strings from the vocabulary  $\{(),\{\},\{\}\}$ . At each position  $i$ , the model must output T if the string up to position  $i$  is a valid string in Dyck-2; P if it is the prefix of a valid string; and F otherwise. A string is valid if every parenthesis is balanced by a parenthesis of the same type. This circuit recognizes an invalid pattern, where a left parenthesis ( $\langle$  or  $\{\}$ ) is immediately followed by a right parenthesis of the wrong type ( $\}$  or  $\}$ , respectively). First, an attention head (left) copies the tokens variable from the previous position. Second, an MLP (top right) reads the output of this attention head, along with the tokens variable, and classifies the input into one of four categories—in particular, returning 13 if it sees the pattern ( $\}$  or  $\{\}$ ). In the second layer, another attention head (bottom right) looks for 13, propagating this information to later positions.
Annotating variable types. In our Transformer model, all categorical variables are encoded using one-hot embeddings, which are represented as integers in the corresponding program. To improve readability, we replace these integers with symbolic values where appropriate. At the input layer, we represent the values of the tokens variable as strings rather than integer indices. At subsequent layers, we determine the appropriate type by following the computation graph. For example, in Figure 9, the tokens variable takes on values in  $\{(.,), \{.,\}\}$ ; the first attention head reads tokens as the value variable, so we can automatically determine that attention_0_0_outputs variable takes on values of the same type; finally, the MLP reads tokens and attention_0_0_outputs as input variables, so we define the mlp_0_0 function in terms of the token value type.
A.4 Possible extensions
In our experiments, we used a limited set of model components that were designed to ensure that the resulting programs are interpretable. However, these do not cover all of the primitive components included in RASP. In this section, we discuss some of these differences; how these limitations can be addressed by introducing additional modules; and the relationship between (a) introducing more expressive modules and (b) the complexity of the corresponding program.
One-to-many attention predicates Our Transformer Programs use a form of binary categorical attention that associates each query value with a single key value. This is enforced by parameterizing the attention head with a predicate matrix  $\mathbf{W}_{\mathrm{predicate}} \in \{0,1\}^{k \times k}$  (where  $k$  is the variable cardinality) and constraining  $\mathbf{W}_{\mathrm{predicate}}$  so that each row contains only a single non-zero entry. This choice results
in more concise programs, as each attention head can be characterized in terms of a one-to-one mapping between key and query values, but it also makes it more difficult to express certain common attention patterns. For example, it is difficult to implement the "less than" predicate, by which a query with a value of  $n$  attends to all keys with values  $< n$ . This predicate is used in a number of RASP solutions [Weiss et al., 2021] to implement a sorting routine. It is straightforward to allow the predicate matrix to have multiple non-zero entries per row. These weights could be learned by replacing the row-wise Gumbel softmax operation described in §3.2 with an entry-wise Gumbel-sigmoid operation. In preliminary experiments, we trained models with this form of predicate matrix and found they obtained similar performance but were qualitatively more difficult to understand.
Score-based attention As in Tracr and RASP, we require that attention heads with categorical values always attend to exactly one position (hard attention). We enforce this constraint by deterministically breaking ties in favor of the position nearest to the query position. Weiss et al. [2021] discuss an extension to RASP (select_best) that provides for a more general form of one-hot attention, where each attention head is associated with both a binary function, predicate:  $\mathcal{Q} \times \mathcal{K} \rightarrow \{0, 1\}$ , and a continuous function, score:  $\mathcal{Q} \times \mathcal{K} \rightarrow \mathbb{R}$ , that maps each combination of key and query to a real value. Given a query  $q$  and keys  $k_1, \ldots, k_n$ , the attention between  $q$  and  $k_i$  is defined as 1 if both predicate  $(q, k_i) = 1$  and score  $(q, k_i) = \max_{j=1}^{n} \text{score}(q, k_i)$ , and 0 otherwise. This type of operation has a number of applications. For example, Yao et al. [2021] describe a construction for modeling sequences of balanced parentheses with an autoregressive Transformer, in which the model attends to the most recent opening bracket at the same nesting depth as the current position, by defining score  $(q, k_i) = i$ .
It would be relatively straightforward to extend Transformer Programs to include score-based attention modules, but this would also result in longer and more complex programs, because each query must now be defined in terms of an ordering of all  $|\mathcal{K}|$  possible keys, instead of a single key. To see why this makes the program more difficult to understand, consider the information gained by observing that some token  $i$  attends to token  $j$ . In our default, one-to-one form of attention, this observation means that  $j$  is the closest position such that predicate  $(q_i, k_j) = 1$ . In score-based attention, the observation also implies that there is no position  $\ell$  such that score  $(q_i, k_\ell) > \text{score}(q_i, k_j)$ , conveying additional information about the rest of the sequence.
Learnable numerical variables In our experiments, we constrain the numerical variables to be either constant and equal to one (at the input layer), or the output of a numerical attention head. In RASP, numerical variables can take on any scalar value. Transformer Programs could be extended to include learnable numerical input variables, at the cost of some interpretability tradeoffs. For example, it may no longer be possible to decompile MLPs by enumerating all possible arguments.
In this section we describe additional implementation details for the experiments in Section 4. Code to reproduce the experiments is available at https://github.com/princeton-nlp/TransformerPrograms.
B.1 In-context learning
Data. For our in-context learning experiment (Section 4.1), we sample 20,000 sequences of length 10. The first token is a beginning-of-sequence token, and the remainder of the sequence is formed by randomly sampling a many-to-one mapping between letter types (either a, b, c, d) and numbers (0, 1, 2, 3), and then alternating letters and numbers until reaching the target length. The model is trained to predict the number following each letter, or a special unk token if the letter has not appeared earlier in the sequence. We use an autoregressive mask so that at each position  $i$ , the model can attend only to keys with positions  $j \leq i$ .
Training. We train each model for 250 epochs with a batch size of 512, a learning rate of 0.05, and annealing the Gumbel temperature geometrically from 3.0 to 0.01, decreasing the temperature at each training step. We take one Gumbel sample per step. The model has two layers with one categorical attention head per layer, with a variable cardinality of 10. We train five models with different random initializations and pick the one with the lowest test loss. We implement all models in PyTorch [Paszke et al., 2019] and use the Adam optimizer [Kingma and Ba, 2014].
B.2 RASP tasks
Data. For each RASP task, we sample 20,000 inputs without replacement and partition them into train, validation, and test sets containing 16,000/2,000/2,000 instances respectively. With the exception of the Dyck tasks, we sample inputs by sampling tokens uniformly from the vocabulary. For the Dyck tasks, we follow Weiss et al. [2021] and sample with a bias towards strings with a valid prefix. Specifically, given maximum length of  $N$ , with probability 0.5, we sample  $N$  tokens uniformly from the vocabulary; otherwise, we sample a valid Dyck string  $s$  with length  $|s| \leq N$ , and append  $N - |s|$  randomly sampled tokens to the end to obtain a string with length  $N$ . For all tasks, we obtain 20,000 samples without replacement by sampling strings uniformly until the set of unique strings has the intended size. We prepend all inputs with a beginning of sequence token bos. For the sort and reverse tasks, we also append an end-of-sequence token, eos. Following Weiss et al. [2021], we report the token-level accuracy.
Training. As above, we train the model for 250 epochs with a batch size of 512, and a learning rate of 0.05. We anneal the Gumbel temperature geometrically from 3.0 to 0.01, decreasing the temperature at each training step, and taking one Gumbel sample per step. These hyperparameters were chosen after initial experiments on the RASP tasks. We do a grid search for number of layers (2, 3), number of attention heads (4, 8), and number of MLPs per layer (2, 4). The attention heads are evenly divided between categorical and numerical heads. Similarly, the MLPs are evenly divided between MLPs with two categorical inputs, and MLPs with two numerical inputs. We train models with five random initializations, pick the model with the best accuracy on a validation set, and report the accuracy on a held-out test set. For experiments with standard Transformers (described in Appendix C.1), we similarly do a grid search for number of layers (2, 3) and attention heads (4, 8) but otherwise follow the hyper-parameters described by Weiss et al. [2021]: the hidden dimension is 256, the learning rate is 0.0003, the batch size is 50, and we train for up to 100 epochs, picking the model checkpoint from the epoch with the highest validation accuracy. Each model takes between five and fifteen minutes to train on an Nvidia RTX 2080 GPU, depending on the number of layers.
B.3 Named entity recognition
Data. For our named entity recognition experiments (Section 4.3), we train on data from the CoNLL-2003 shared task [Sang and De Meulder, 2003], using the distribution from HuggingFace Datasets [Lhoest et al., 2021]. This data is pre-tokenized and we filter the dataset to sentences with a maximum length of 30 tokens and add special beginning- and end-of-sequence tokens. Following Collobert et al. [2011], we preprocess the data by replacing all contiguous sequences of numbers with a special number symbol, so, for example, the string "19.99" is replaced with "#.#" We use the standard train/validation/test split and evaluate the results using a Python implementation of the standard CoNLL evaluation script [Nakayama, 2018].
Training. For both the standard Transformer and Transformer Programs, we use a batch size of 32 and perform a grid search over the number of layers (1, 2, or 3) and the number of attention heads (4, 8). For the standard Transformer, we search for model hidden dimension (64, 128, 256) and train for up to 100 epochs, taking a checkpoint after each epoch and picking the checkpoint with the highest performance on the validation set. For the Transformer Program, we fix the variable cardinality to 32 and search for the number of embedding variables (2, 4, 8) and training epochs (30, 50, 100). As above, we anneal the temperature geometrically from 3.0 to 0.01, and we report results after discretizing the model at the end of training. For both models, we initialize the word embeddings using 300-dimensional GloVe embeddings. For both models, we train with three random seeds, pick the model with the highest F1 score on the validation set, and report the results on the test set.
C.1 Additional RASP results
In this section, we provide additional code examples and ablations on the RASP tasks.
Larger-scale versions of the tasks. In Figure 10, we show the accuracy of Transformer Programs trained on RASP tasks with a longer maximum sequence length and a larger input vocabulary. Because we use one-hot encodings for the token and position embeddings, these values determine the cardinality,  $k$ , of the categorical variables used in the program. We compare Transformer Programs with standard Transformers, for each model performing a grid search over the number of layers and
Figure 10: RASP accuracy after increasing the size of the input vocabularies  $(|\mathcal{V}|)$  and maximum sequence length  $(N)$ , comparing Transformer Programs with standard Transformers. For each model, we perform a hyperparameter search over the number of layers and attention heads and report the test accuracy of the model that performs best on the validation set. On most tasks, performance degrades moderately when trained on longer sequences, and degrades more when we also increase the vocabulary size, and this trend generally holds for both standard Transformers and Transformer Programs. However, the Transformer Programs deteriorate more dramatically, with over a  $50\%$  accuracy gap relative to the standard Transformer in the most difficult setting.
attention heads and report the test accuracy of the model that performs best on the validation set. In Figure 10, we increase the vocabulary size and maximum sequence length from eight to sixteen. All other experiment details are the same as in Appendix B.2. Performance degrades on longer sequences, underscoring some of the optimization challenges in learning Transformer Programs for larger scales.
No numerical variables. In our main experiments, we train Transformer Programs with an equal number of categorical attention heads and numerical attention heads, and categorical MLPs and numerical MLPs. In Table 3, we compare results on RASP tasks with only categorical variables. The experiment setting is otherwise the same as in Appendix B.2, but all attention heads and MLPs are constrained to read only categorical variables. Not surprisingly, performance degrades on three tasks that are primarily based on counting: Histograms, Double Histograms, and Most Frequent. However, the Transformer Programs still achieve good performance on the other four tasks. These include the balanced parenthesis languages (Dyck-1 and Dyck-2), which are most naturally solved by keeping a count of unmatched parentheses at each position.
Dataset	L	H	M	Acc.
Reverse	3	8	4	99.75
Hist	3	8	4	83.20
2-Hist	2	8	2	55.89
Sort	3	2	2	99.99
Most Freq	3	8	4	64.13
Dyck-1	3	8	4	99.43
Dyck-2	3	4	2	99.17
Table 3: Accuracy of Transformer Programs on RASP tasks using only categorical variables, with the number of layers (L), attention heads (H), and MLPs (M) used in the best-performing model.
C.2 Analyzing the generated code
Here, we provide some additional analysis of the generated code, focusing on the RASP tasks. Complete, generated programs are available at https://github.com/princeton-nlp/TransformerPrograms.
Code analysis tools. One advantage of representing models as Python programs is that we can use off-the-shelf code analysis tools to interpret them. Figure 11 illustrates one example, showing how we can use a Python debugger to understand a non-trivial circuit in a Sort program.
Program length. Table 4 shows the number of lines in our best-performing program for each RASP task, before and after applying the compression strategies described in Appendix A.3. The program length includes the basic library functions used to We use an automated Python codeformatter, which applies a number of standard style conventions and, in particular, enforces a maximum line length of 88 characters.
4https://github.com/psf/black
Figure 11: After converting a Transformer Program into Python code, we can analyze it using off-the-shelf debugging tools. Here, we inspect a subset of the Sort program using the built-in debugger in Visual Studio Code. We can step through the program for a test example, set break-points, and leave comments. In this example, we find a circuit involving two attention heads, with a second-layer attention head (left) reading a value from a first-layer attention head (right). Inspecting the code, we can see that this pair of heads has the effect of propagating an early-position token to the later positions and a late-position token to the earlier positions.
Dataset	Full	Pruned
Reverse	1893	713
Hist	324	160
2-Hist	1309	423
Sort	1503	635
Most Freq	1880	666
Dyck-1	9975	892
Dyck-2	5406	733
Table 4: The number of lines in best programs for each RASP task before and after applying a set of simple pruning strategies based on static analysis of the code.