SlowGuess/NeurIPS_2023_OCR huggingface dataset:
What information do the attention heads read? Because each attention head reads a fixed set of named variables, we can characterize how information flows through the programs by examining which variables are read by each head. In Figure 12, we summarize this information for RASP programs. At the first layer, the majority of categorical attention heads read positions as key and query variables and tokens as the value. At higher layers, positions remains the most common key variable, but the models are more likely to read the outputs of lower-layer attention heads as the value variable. Numerical attention heads are less likely to read positions and more likely to read tokens, attn, and mlp outputs. Both kinds of attention successfully learn to compose modules, with higher-layer modules reading the outputs of modules at lower layers.
C.3 Optimization challenges: case study
What are the challenges to scaling Transformer Programs to more complex problems? As discussed in Section 6, some of the main obstacles are optimization challenges. This is evident from the fact that, for many RASP tasks, we can write a Transformer Program that would achieve perfect accuracy, but our method fails to learn these solutions (as illustrated in Fig. 10). In this section, we conduct a case study to better understand these optimization challenges.
Setting We focus on the Reverse task. This task can be solved by the following program, adapted from Weiss et al. [2021], using a Transformer with two layers and one attention head per-layer:
# First-layer attention
length = aggregate(select(tokens, tokens, lambda q, k: k == "
") positions)
# First-layer MLP
targets = one_hot(length - positions)
Figure 12: For each of the RASP tasks, we learn a Transformer Program with three layers and eight heads per-layer, divided evenly between categorical and numerical attention heads, and summarize the types of variables that are read at different layers. For each layer, we list the key, query, and value variables read by attention heads at that layer, and calculate the proportion of heads that read the tokens variable; positions; ones (for numerical attention); the output of a previous attention head (attn); or the output of a previous MLP (mlp). We aggregate over RASP programs and compare categorical attention heads (left) and numerical attention heads (right).
Attention 1	MLP 1	Attention 2	Accuracy
Read	Predicate	Read	Read	Predicate
-	-	-	-	-	23.2/23.6/24.1
✓	✓	✓	✓	✓	99.9/99.9/80.8
✓	✓	✓	✓	-	37.9/40.3/18.5
✓	✓	✓	-	✓	17.1/13.7/20.2
✓	✓	-	✓	✓	95.1/94.1/95.3
✓	-	✓	✓	✓	99.1/83.9/78.2
-	✓	✓	✓	✓	35.5/44.1/41.8
Table 5: Results on the Reverse task  $(|\mathcal{V}| = N = 16)$  after manually initializing model components to encode a generalizing solution. (See Section C.3.) In each row, we choose a subset of components to manually initialize  $(\checkmark)$ , initialize the other weights randomly (-), and then train the model, reporting the (per-token) accuracy of the resulting program. The components we consider are the projection matrices, which determine which variables each module reads (Read), and the attention predicate matrices (Predicate). All other components (i.e., the internal MLP parameters and the final classifier weights) are always initialized randomly. We run each experiment with three random seeds and report the accuracy from each run. When we manually initialize the attention Read and Predicate weights to encode the generalizing solution, the model successfully learns the remaining components. Performance degrades considerably when even a single attention components is initialized randomly, suggesting that our optimization procedure struggles to find effective attention patterns.
Second-layer attention
output = aggregate(select(targets, positions, ==), tokens)
The first attention layer copies the position of the end-of-sequence token to determine the length of the sequence; the first-layer MLP calculates the difference between length and positions; and the second attention layer uses the MLP output to read, at each position  $i$ , the token at position length  $-i$ . This solution works for sequences of all lengths, but Transformer Programs evidently fail to learn it.
Method To better understand why, we manually initialize some components of a two-layer Transformer Program model to encode this solution and train the model on the Reverse task with a vocabulary size and maximum sequence length of 16. Specifically, in the attention layers, we manually initialize: (1) the parameters associated with  $\pi_K$ ,  $\pi_Q$ , and  $\pi_V$ , which identify the key, query, and value variables to read; and (2) the parameters associated with the predicate matrix  $\mathbf{W}_{\mathrm{predicate}}$ . In the MLP, we manually initialize the parameters associated with the input projection matrix  $\mathbf{W}_{\mathrm{in}}$ , which determines which variables the MLP reads from the residual stream. We do not initialize the other MLP weights or the classifier weights. We initialize different subsets of these parameters and observe whether the model learns the remaining components needed to complete the solution. We train each model for 100 epochs but otherwise use the same hyperparameters described in Appendix B.
Results The results of this experiment are in Table 5. When all of the components are randomly initialized, the model fails to learn the generalizing solution. When we manually initialize the attention layers and the MLP input weights, the model learns the remaining components needed to solve the problem successfully. In particular, it finds the internal MLP weights needed to calculate the difference between the two input variables (length and position) and output a one-hot encoding of the result. Interestingly, if we randomly initialize either even a single component in either one of the attention layers, the model fails to learn the correct solution; but if we randomly initialize only the MLP, the model learns successfully. These results suggest that the attention layers may be more difficult to optimize—for example, they may be more likely to get stuck in local minima.
C.4 Text classification
Model	TREC	MR	Subj	AG
Bag of words	74.8	77.8	92.6	89.6
Standard Transformer	88.6	77.2	92.3	91.7
Transformer Program	85.6	77.9	93.0	90.8
Table 6: Classification accuracy on question classification (TREC); sentiment analysis (MR); subjectivity classification (Subj); and topic classification (AG news). The Bag of words baseline is a multinomial Naive Bayes model trained on unigram features.
Next, we train Transformer Programs for three standard text classification datasets: classifying questions as one of six topics [TREC; Voorhees and Tice, 2000]; classifying movie reviews as positive or negative [MR; Pang and Lee, 2005]; classifying sentences as objective or subjective [Subj; Pang and Lee, 2004]; and classifying sentences from news articles as one of four topics [AG News; Zhang et al., 2015]. We filter the datasets to sentences with at most 64 words and use a vocabulary of the 10,000 most common words, replacing the remaining words with an unknown token, and fixing the variable cardinality at 64. As above, we compare the Transformer Program with a standard Transformer. For both models, we use the Transformer to extract token embeddings, obtain a sentence embedding by averaging the token embeddings, and train a linear classifier on the sentence embedding. We initialize both models with 300-dimensional GloVe embeddings and use grid search to select the model dimension, number of layers, and number of heads, and training hyper-parameters (learning rate and number of training epochs), as described in Appendix B.3. We hold out  $10\%$  of each training set to use as a validation set. For both models, we train with three random seeds, pick the model with the highest accuracy score on the validation set, and report the results on the standard test set.
Accuracy. Table 6 compares the accuracy for Transformer Programs, standard Transformers, and a bag-of-words baseline. The Transformer Programs are competitive with the standard Transformer on three out of four datasets, performing slightly worse on TREC. This could be because the standard Transformer can more effectively leverage pre-trained word embeddings, or because it is easier to regularize.
Interpretability. We illustrate the interpretability of these programs in Figure 13 by inspecting the features for an example from the TREC dataset.
(a) Classification scores for each token embedding.
(b) Feature-level scores for the "world" token.
# attn_1.1: Copy var0 from early positions.  
def predicate_1_1(var2_embedding, position):  
    if var2_embedding in {0, 4, 5, 6, 10, 12, 13}:  
        return position == 1  
    elif var2_embedding in {1, 7, 8, 9, 11, 14}:  
        return position == 4  
    elif var2_embedding in {2, 3}:  
        return position == 5  
    elif var2_embedding in {15}:  
        return position == 2
attn_1_1_pattern = selectclosest( positions, var2_embedding, predicate_1) attn_1_1.Outputs = aggregate( attn_1_1_pattern, var0_embedding)
(c) The code for computing the attn_1_1_outputs feature.
word	var0...	var1...	var2...	var3...
丫	7 ×	丫	丫	丫
how	7	4	5	15
are	7	11	4	0
for	7	2	11	10
when	7	6	14	6
there	7	5	14	15
year	7	6	6	14
people	7	13	15	15
time	7	6	15	0
date	7	6	14	6
(d) A subset of the embedding table, filtering to words with var0_embedding values of 7.
Figure 13: Visualizing the predictions for an example from the TREC question classification dataset. We classify sequences by pooling the final-layer token embeddings, so we can visualize the classification scores for each token embedding (Figure 13a). In this example, the first three tokens ("how", "many", and "people") have the highest scores in favor of the NUM class, but most other tokens have high scores for this class as well. To see why, we can inspect the feature-level scores for individual token embeddings. In Figure 13b, we display the categorical features of the "world" token along with the corresponding classifier weights. For this token, the input features favor the LOC label, but attention features increase the score for NUM. Figure 13c displays the subset of the program that calculates one of these attention features (attn_1_1_output = 7). This attention head generally copies the var0_embedding variables from first or fourth position—positions that can be expected to be informative for question classification. In Figure 13d, we display the most frequent words that have var0_embedding values of 7: they include question words like "how" and "when" and nouns like "year", "time", and "date", which are more likely to occur in numerical questions.