SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In the main paper and Appendix G, we show figures of a forward pass in a compiled model. We found that these figures can be confusing to read at first, especially as the compiled models get bigger. This section serves as a reference for how to interpret these figures.
As an example, let us walk through the figure for the frac_prevs model from Figure 2:
The figure has 5 panels, each of which shows the content of the residual stream after the corresponding layer in the model. This allows us to follow what the model does step-by-step. The residual stream has size [sequence length x dimensionality], therefore we visualize it as a 2-dimensional heatmap. In this example, we have a size of  $5 \times 14$  including the BOS token position.
We show a forward pass for a specific input sequence [bos, x, a, c, x]. On the x-axis of each panel we label the token positions with the corresponding input token. The y-axis of the plot contains the dimensions of the residual stream. Thanks to our knowledge of the program the model implements, we can label each dimension according to what it encodes. Dimensions starting with 'tokens' contain the (categorical) input embeddings. Dimensions starting with 'indices' contain the (categorical) position embeddings. Labels that contain a ':' are dimensions that correspond to a categorical ('one-hot') encoding and the value after the ':' is the value encoded in this dimension.
Labels without a ':' mean that this dimension encodes a numerical value. In each of the four panels we show the full residual stream content as a heatmap. Entries that were changed by the layer corresponding to the panel are highlighted with a red border.
Let's go through the plot step by step and map it to the code we used to compile this model:
In the leftmost panel we see the residual stream after the input embedding layer. It contains the categorical encoding of the tokens and the categorical encoding of the indices. For example at the token position of the 'a' token, the dimension tokens: a contains 1, and the dimension indices: 1 contains 1. The auxiliary dimension one contains 1 at every token position.
The first attention layer is a no-op, so the residual stream afterwards (shown in the second panel) is the same as before. No entry is highlighted.
MLP 1 computes the first line of the rasp program is_x = (tokens == "x"). It writes the result into a numerical dimension in the residual stream labelled with is_x. For this concrete sequence the layer writes a 1 into both token positions that contain the token 'x' in the input.
Attn 2 computes the select-aggregate operations in lines 2 and 3 of the RASP program. It computes the fraction of previous 'x' tokens, and writes the result into a single dimension labelled frac_prevs. It writes values between 0 and 1 in all token positions except for the BOS token position. For this example the result will be [1, 1/2, 1/3, 1/2].
The final MLP 2 layer is a no-op again and does not change anything in the residual stream. The output unembedding layer will then read the result from the frac_pre vs dimension.
D Tracr Implementation Details
This section highlights a few more implementation details of Tracr. We describe how we construct MLP and attention blocks, how we implement the selector width primitive, and how we extend RASP and Tracr to use causal attention. For the full implementation and documentation, refer to the code repository at https://github.com/google-deepmind/tracr.
D.1 MLP and Attention Blocks
For MLP layers, we distinguish between Map operations with a single input and output and SequenceMap operations with two inputs and one output. We can recursively represent functions with more than two inputs using SequenceMaps.
We translate Maps with categorical inputs and outputs to MLPs that act as a lookup table. SequenceMaps with categorical inputs and outputs become MLPs where the first layer maps to an encoding of all pairs of inputs and the second layer acts as a lookup table.
For numerical inputs and outputs, we explicitly construct MLP layers as universal function approximators. In these MLPs, the first layer discretises the input, and the second layer maps each discrete bucket to a corresponding output value. We know which input/output values can occur, so we can choose the discretisation around these known input values to minimise the approximation error.
We now turn our attention to the attention blocks, which we construct from RASP selectors.
We first construct the  $\tilde{W}_{QK}$  matrix to implement the desired attention pattern in the attention logits. We will refer to this as the direct attention matrix. This matrix has low rank, with its row space being the part of the residual stream where the query s-op is stored, and the column space being where the key s-op is stored. We adjust the direct attention matrix by adding a rank-one update  $W_{BOS} = \beta_{BOS}x_{\mathrm{one}}x_{\mathrm{tokens:bos}}^{\intercal}$  with  $\beta_{BOS} = 1$  or  $\beta_{BOS} = \frac{1}{2}$ , to ensure that the BOS token is attended to either always, or whenever no other token is. ( $x_{\mathrm{one}}$  and  $x_{\mathrm{tokens:bos}}$  here are unit vectors for the special embedding dimensions introduced in Section 4.) We then scale up the matrix by an inverse-temperature parameter  $T^{-1}$  (100 by default), getting  $W_{QK} = T^{-1}(\tilde{W}_{QK} + W_{BOS})$ . As a result,
the attention weights  $A_{ij} = \mathrm{softmax}\left(\mathbf{q}_i^\intercal W_{QK}\vec{\mathbf{k}}\right)_j = \exp (\mathbf{q}_i^\intercal W_{QK}\mathbf{k}_j) / \sum_{j'}\exp (\mathbf{q}_i^\intercal W_{QK}\mathbf{k}_{j'})$  are very close to 1/#{selected tokens} on selected tokens and 0 elsewhere.
The  $W_{OV}$  matrix maps the value input to the corresponding output dimensions. Attention layers only support categorical key and query inputs. The value inputs can be numerical or categorical. We can only use categorical values if the head never attends to more than one token.
D.2 Selector Width Primitive
RASP provides the selector width primitive, which counts the number of 1s in each row of a selector. It provides an alternative to aggregate for processing selectors.
Weiss et al. (2021) provide a selector width implementation in pure RASP, making it not necessarily a language primitive. However, the most efficient implementation uses the BOS token, which exists Tracr but is not exposed to the RASP program.
Therefore, Tracr translates selector width directly into an efficient implementation in Craft consisting of an attention layer and an MLP layer. The attention layer implements an attention pattern that matches the selector to compute the width of. It uses the BOS token as value input, resulting in the attention head computing  $x = 1 / (1 + w)$  where  $w$  is the desired selector width output. The next MLP layer then computes  $w = 1 / x - 1$  and cleans the BOS token position.
D.3 Causal Attention
Most transformer models used in practice use causal attention, i.e., they apply a mask to the attention patterns that allows the model to attend only to previous tokens. This allows training the models autoregressively. However, RASP assumes non-causal (i.e. bidirectional) attention by default. While all models discussed in the main paper use non-causal attention, Tracr also supports causal attention.
To enable this, we extend RASP to support causal attention via a flag set during evaluation. To evaluate a RASP program in the causal evaluation mode, we apply a causal mask to the output of each selector. Causal evaluation changes the semantics of some RASP operations, and, in general, it is necessary to adapt RASP programs to function with causal attention. For example, the frac_prevs program no longer needs to compute a causal mask manually. However, for example, the length implementation by Weiss et al. (2021) no longer correctly computes the length of a sequence because it requires attending to future tokens.
Similarly, Tracr has a flag to enable causal compilation. Most of the compilation process does not change, and we only need to ensure to compile selectors to causal attention heads.
E Compression Training Details
We implemented the compression described in Section 5 in Jax on top of the Haiku transformer implementation that comes with Tracr. We train  $W$  using the AdamW optimizer (implemented in Optax) with a weight decay factor of 0.1, and parameters  $\beta_{1} = 0.9$ ,  $\beta_{2} = 0.99$ . We train for  $3 \times 10^{5}$  steps with a batch size of 256. We decay the learning rate linearly from  $10^{-3}$  to  $10^{-6}$  over the first half of training. Each compression run requires between 1 and 4 hours of run time on two CPU cores (depending on the size of the model to compress).
F Theoretical Results on Combining Attention Heads
The RASP language permits combining arbitrary selectors elementwise using boolean operators, such as and, or, and not. It is not immediately obvious what operators can be implemented given the way we encode selectors as attention matrices  $W_{QK}$ , as described in Appendix D.1.
First, let's consider not operator for a selector select(query, key, pred) with given direct attention matrix  $\tilde{W}_{QK}$ . One way to implement not select(query, key, pred) is to note that it's equivalent to select(query, key, not pred). Another is to use a transformed direct attention matrix  $\tilde{W}_{QK}^{\mathrm{not}} = -\tilde{W}_{QK}$ , alongside a  $\beta_{BOS}^{\mathrm{not}}$  that's 0 or  $-\frac{1}{2}$ .
Next, let's consider the and operator on two selectors select(query_a, key_a, pred_a) and select(query_b, key_b, pred_b) whose direct attention matrices  $\tilde{W}_{QK}^{A}, W_{QK}^{B}$  are given, and produce 0-1 attention logits. We can observe that taking  $\tilde{W}_{QK}^{\text{and}} = \tilde{W}_{QK}^{A} + \tilde{W}_{QK}^{B}$  results in attention logits taking value 2 when both selectors are active, and at most 1 otherwise; so by the same procedure in Appendix D.1, with  $\beta_{BOS}^{\text{and}}$  taking value  $\frac{3}{2}$  or 2, we can construct  $W_{QK}^{\text{and}} = T^{-1}(\tilde{W}_{QK}^{\text{and}} + W_{BOS}^{\text{and}})$  that produces the desired attention pattern in the post-softmax attention weights.
We can compose these constructions, negating the two given selectors before combining them with and, to get nor, with  $\tilde{W}_{QK}^{\mathrm{nor}} = -\tilde{W}_{QK}^{A} - \tilde{W}_{QK}^{B}$  and  $\beta_{BOS}^{\mathrm{and}}$  taking value  $-\frac{1}{2}$  or 0, resulting in an implementation of select(query_a, key_a, pred_a) nor select(query_b, key_b, pred_b).
So far these are fairly natural constructions - the boolean operators not and and can be used to construct all other possible boolean operators, so we might expect that indeed all combinations of selectors via boolean operators can be compiled to transformer weights this way.
Alas, it is not so. Unlike the implementation of not, the implementations of and and nor above did not result in a direct attention matrix that produces the correct pattern (potentially shifted by a constant) in the attention logits, but rather only in the attention weights after temperature-adjusted softmax, meaning they cannot be composed further to produce arbitrary logical statements.
If we were to try to implement or, the easiest way would be to negate the nor by composing the transformations - but the resulting  $\tilde{W}_{QK}^{\mathrm{or}} = -(-\tilde{W}_{QK}^{A} - \tilde{W}_{QK}^{B})$  is actually the same direct attention matrix we used for and. This produces attention logit 1 or 2 where the selectors' or is active, and 0 where it isn't. However, the temperature adjustment with  $T^{-1}\gg 1$  that forces the attention to be near-zero where neither selector is active will then also do the same thing when only one selector is active, so the attention weights will be different between tokens where both selectors are active versus only one selector.
In fact, this obstruction to implementing or can be generalized, as follows.
Lemma F.1. Consider two selectors select (query_A, key_A, pred_A) and select (query_B, key_B, pred_B), with direct attention matrices  $\tilde{W}_{QK}^{A}$  and  $\tilde{W}_{QK}^{B}$ . For ease of analysis, let's suppose query_A, key_A, query_B, and key_B are stored in separate, orthogonal subspaces  $Q_{A}$ ,  $K_{A}$ ,  $Q_{B}$ ,  $K_{B}$ .
Now suppose there exists an attention matrix  $\tilde{W}_{QK}^{\mathrm{or}}$ , with row space contained in  $Q_{A} + Q_{B}$  and column space contained in  $R_{A} + R_{B}$ , that, after adjustment by some BOS logit offset  $\beta_{BOS}^{\mathrm{or}}$  and some temperature  $T \to 0$ , produces attention weights converging to the normalized selector weights for select(query_A, key_A, pred_A) or select(query_B, key_B, pred_B). Then, these selectors are not generic - they satisfy some very limiting constraints about their predicates.
Proof. Let's begin by assuming the second selector,  $B$ , is not constant, selecting some tokens and not-selecting other tokens. This implies the existence of basis vectors  $\mathbf{q}_B^0$ ,  $\mathbf{q}_B^1 \in Q_B$  and  $\mathbf{k}_B^0$ ,  $\mathbf{k}_B^1 \in K_B$  such that  $\mathbf{q}_B^0 \top \tilde{W}_{QK}^B \mathbf{k}_B^0 = 0$  and  $\mathbf{q}_B^1 \top \tilde{W}_{QK}^B \mathbf{k}_B^1 = 1$ . Holding these constant, consider some basis vectors  $\mathbf{q}_A \in Q_A$  and  $\mathbf{k}_A$ ,  $\mathbf{k}_A' \in K_A$ . Then, for query vector  $\mathbf{q}_A + \mathbf{q}_B^1$ , all tokens with key vector  $\mathbf{k}_A + \mathbf{k}_B^1$  or  $\mathbf{k}_A' + \mathbf{k}_B^1$  must be selected, which means they must have equal attention logits. Therefore,  $(\mathbf{q}_A + \mathbf{q}_B^1) \top \tilde{W}_{QK}^{\mathrm{or}}(\mathbf{k}_A + \mathbf{k}_B^1) = (\mathbf{q}_A + \mathbf{q}_B^1) \top \tilde{W}_{QK}^{\mathrm{or}}(\mathbf{k}_A' + \mathbf{k}_B^1)$ , so  $\mathbf{q}_A^T \tilde{W}_{QK}^{\mathrm{or}} \mathbf{k}_A = \mathbf{q}_A^T \tilde{W}_{QK}^{\mathrm{or}} \mathbf{k}_A'$ .
Now, consider  $\mathbf{k} = \mathbf{k}_A + \mathbf{k}_B^0$ ,  $\mathbf{k}' = \mathbf{k}_A' + \mathbf{k}_B^0$ ,  $\mathbf{q} = \mathbf{q}_A + \mathbf{q}_B^0$ , and, for some basis vector  $\mathbf{q}_A' \in Q_A$ , let  $\mathbf{q}' = \mathbf{q}_A' + \mathbf{q}_B^0$ . We have logit differences  $\mathbf{q}^{\intercal}\tilde{W}_{QK}^{\mathrm{or}}\mathbf{k}' - \mathbf{q}^{\intercal}\tilde{W}_{QK}^{\mathrm{or}}\mathbf{k} = \mathbf{q}_B^0\mathbf{T}\tilde{W}_{QK}^{\mathrm{or}}(\mathbf{k}' - \mathbf{k}) = \mathbf{q}'\tilde{W}_{QK}^{\mathrm{or}}\mathbf{k}' - \mathbf{q}'\tilde{W}_{QK}^{\mathrm{or}}\mathbf{k}$ . Therefore, among tokens where key\_B has vector  $\mathbf{k}_B^0$  (let's call these  $\mathbf{k}_B^0$ -tokens), the tokens that have highest logit for query vector  $\mathbf{q}'$  are the same as those for query vector  $\mathbf{q}$ . However, the selected tokens among the  $\mathbf{k}_B^0$ -tokens are either none of them, or exactly those with the highest logit (which depends on key\_A). Because of the definition of or,  $\mathbf{k}_B^0$ -tokens are selected exactly if select(query\_A, key\_A, pred\_A) would select them.
Putting the above observations together, it follows that for query_A vectors  $\mathbf{q}_A$  and  $\mathbf{q}_A'$ , pred_A will either select no keys for one of them, or will select the same keys for both of them. In other words, pred_A must be rewritten in the form query_pred_A(query_A) and key_pred_A
(key_A). Equivalently, pred_A's matrix has rank 1; we can say in short that pred_A is a rank-1 predicate, or that select(query_A, key_A, pred_A) is a rank-1 selector.
If we suppose our initial assumption to be false, then pred_B is constant, and can thus be just as well rewritten to be a predicate of query_A and key_A; then, it is easy to derive the necessary  $\tilde{W}_{QK}^{\mathrm{or}}$  from select(query_A, key_A, pred_A or pred_B).
We can repeat the argument interchanging the selectors, to conclude that either the operation is trivial (because one predicate is constant), or both selectors must be rank-1.
The above conclusion may be averted in the case that we have a priori information that certain values of  $\mathbf{q}_A, \mathbf{k}_A, \mathbf{q}_B, \mathbf{k}_B$  cannot co-occur, or if some of the input s-ops are shared. We leave exploring that, as well as whether or can be implemented in the case of rank-1 predicates, to future work.
A notable special case of the above is the case where query_A and query_B compute the same s-op, and key_A and key_B also compute the same s-op. (They may be the same s-op, or redundant copies.) Then simple rewriting is possible, similarly to the or case explained earlier. For example:
simplifiable_selector = select(tokens, indices, <=) or select(tokens, "a", ==)  
simplified_selector = select(tokens, indices, q <= k or q == "a")
A similar strategy of matching s-ops can be used to circumvent the lemma and straightforwardly implement operators like or, by constructing combined s-ops queryboth and keyboth with output types representing all pairs of queries and keys of the two selectors. These s-ops may be computed by the preceding MLP - however, the encodings occupy dimensionality multiplicative in the sizes of the constituent s-op output types, which is an impediment to scaling these circuits very far.
Due to the composability limitations of each approach considered, we did not implement boolean operators acting on selectors, apart from simple cases where the query and key s-ops agree.
G More Compiled Models
Here, we present a few additional RASP programs and the compiled Tracr models.
Figure 9 shows and extended sort program. It works similarly to the sort_unique program in Figure 5, but sorts any sequence of values by a sequence of keys and can handle duplicates occurring in the keys.
Figure 10 shows the pair_balance program, which computes the difference in the fraction of open and closed parenthesis tokens. We can now use it as a subroutine for the dyck-n program, which checks if a sequence of  $n$  different types of parentheses is balanced:
Input: pairs
1 # Compute running balance of each type of parenthesis  
2 balances = [pair_balance(pair) for pair in pairs]  
3  
4 # If balances were negative anywhere -> parentheses not balanced  
5 any_negative = balances[0] < 0  
6 for balance in balances[1:] :  
7 any_negative = any_negative or (balance < 0)  
8  
9 select_all = select(1, 1, ==)  
10 has_neg = aggregate(select_all, any_negative)  
11  
12 # If all balances are 0 at the end -> closed all parentheses  
13 all_zero = balances[0] == 0  
14 for balance in balances[1]:
all_zero = all_zero and (balance == 0)  
select_last = select(indices, length - 1, ==)  
last_zero = aggregate(select_last, all_zero)  
dyck_n = (last_zero and not has_neg)
Figure 11 shows the compiled dyck-2 model for pairs = {")}, {")}.
15   
16   
17   
18   
19   
20
Input: keys, vals, min_key, context_length
1 keys = (keys + indices + min_key) / context_length  
2 smaller = select_keys, keys, <=)  
3 target_pos = selector_width(smaller)  
4 sel_sort = select(target_pos, indices, ==)  
5 sort = aggregate(sel_sort, vals)
Figure 9: Compiled sort program. Attn 1 is a no-op, MLP 1 adds a small multiple of indices to the keys, and the rest of the model essentially implements sort_unique.
Input: open_token, close_token
1 bools_open = (tokens == open_token)  
2 opens = frac_prevs(bools_open)  
3 bools_close = (tokens == close_token)  
4 closes = frac_prevs(bools_close)  
5 pair_balance = opens - closes
Figure 10: RASP program that uses frac_prevs as a subroutine to compute the fraction of open and closed parenthesis tokens and computes the difference. The compiled model uses open_token = "(") and close_token = ")". Note that the compiled model has the same number of layers as the single frac_prevs model in Figure 2. Attn 1 is still a no-op, MLP 1 and Attn 2 compute both calls to frac_prevs in parallel, and MLP 2 computes the final result.
Figure 11: Compiled dyck-2 program for pairs = ("()", {"}").