Learning Transformer Programs
Authors: Dan Friedman
Alexander Wettig
Danqi Chen
Venue: NeurIPS 2023 oral (2023)
Keywords: mechanistic interpretability
transformers

Abstract: 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.

Learning Transformer Programs
Dan Friedman Alexander Wettig
Danqi Chen
Department of Computer Science Princeton Language and Intelligence Princeton University
{dfriedman, i dee
Abstract
Recent research in mechanistic interpretability has attempted to reverse-engineer Transformer models by carefully inspecting network weights and activations. How- ever, these approaches require considerable manual effort and still fall short 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 RASP [Weiss al., 2021], programming language that can be compiled into Transformer weights. Instead of compiling human-written programs into Transformers, we design modified Transformer that can be trained using gradient- based optimization and then automatically converted into discrete, human-readable program. We refer these models as Transformer Programs. To validate our approach, we learn Transformer Programs for variety of problems, including an in-context learning task, suite of algorithmic problems (e.g. sorting, recognizing Dyck- languagsg, 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 compa- rable size; and, more importantly, they are easy to interpret. To demonstrate these advantages, we convert Transformers into Python programs and use off- fff-lh code analysis tools to debug model errors and identify the "circuits" used to solve different sub-problems. We hope that Transformer Programs open new path toward the goal of intrinsically interpretable machine learning.
Introduction
Transformers [Vaswani et al., 2017] have become the predominant neural network architecture in chine learning, representing the state-of-the-au in natural language processing (NLP) and increasingly in computer vision, molecular modeling and other domains. Prominently, large Transformer -oaseoemeeerrrdddddddee language models ILLMs; Brown et al., 2020] have demonstrated impressive general-purpose capabil- ities and are now widely deployed 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 audit, debug, and maintain. This black- box quality poses practical challenges and limits the usefulness of these models in high-stakes applications.
As result, considerable body of work has aimed improving our understanding of Transform- ers. In NL.,, 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 [c.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]; J; in any case, they do not provide complete or faithful description of the algorithm that the model implements. More recently,
Our code is available at https://githen com/ wwwp://rinnetttoott/ataaaattteee number of example Transformer Programs.
along with a
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Task Data
1. Train
2. Sample
3. Decompile
# annotate code
Input
Target
baxtad Bosd a Lestedee
seat sld watt
use debugger
b3c
d2c4a2b
unk
relaaxed
discrete
suu peelil 13h -enta
autonatic ana myalysis
Transformer Program
Transformer Program
Python program
Figure We design a modified Transformer that can trained on data and then automatically discretized and converted into numan-readale program. The program is functionally identifical to the Transformer, but casier to understand- for example, using an off-the-shelf Python debugger.
research mechanistic interpretability has attempted to gain algorithmic understanding of these models, with the goal of iivers--engineeniie Transformers into human-interpretablc 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 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, aim to train Transformers that are mechanistically interpretable by design. We take inspiration from RASP [Weiss et al., 2021], 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- traseent-bbsed optimization and then automatically decompile them into human-readable programs.
To this end, we propose method for learning Transformer Programs- Progrums-Transforme Transformers that are constrained to implement interprrteereetable algorithms. First, we introduce a set of constraints that restrict the Transformer to lic within an interpretable subpace of the parameter space, by which mean the subset of parameter values that can be mapped to a family human-readable programs. Second, we develop 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 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
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, 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 "irccuits" used to solve different sub- broblem,, and to debug model errors. Overall, aim demonstrate that Transformer Programs represent promising first step towards building intrinsically interpretable machine learning models.
Background
Transformer architecture. The Transformer [Vaswani et al., 2017] is a neural network archi- tecture for processing sequential data. The input to the Transformer is a sequence of tokens w W....... WN E V in discrete vocabulary V. At the input layer, the model maps the to- kens to a sequence of dddimensiinal embeddings X E Nxd, In the standard Transformer, this initial embedding is defined as the sum of learned token embedding and positional embedding. Each subsequent layer consists of a multi-head attention layer (MHA) followed by a multilayer perceptron layer (MLP): x, x;-1 MLP, + MHA, Following the presentation
"The standard Transformer also includes dayer-normalization layer [Ba al., 2016), which we omit here.

---

of Elhage et al. I, multi-head attention can be written as
H
MHA(x) = Σ softmax(xW) (xWW) x J W
h=1
where H. is the number of heads, dy, is the attention embedding dimension, W W W E Raxd, are referred to as the query, key, and value projections respectively, and O
o E &R xd projects the output value back to the model dimension. By default, 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, is two-layer feedforward network: ML = xWW) where W E Rdddm W2 E Rdm xd and a is a non-linear activation, commonly the ReLU function. The output of the model a sequence of token embeddings. XL E R xd 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 neural network processes information. Informally, Transformer be viewed series of nodes (attention heads and -erd-fr layers) reading and writing information to the residual stream, which refers the sequence of token embeddings . Xt, R computed each layer. Each node "reads" from the residual stream via a projection matrix Win E Rddda, calculates some function fxx;-1 Win); and then "writes" the residual stream via another projection: x; x,-1 + f W W ouns with W out E Rd, 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 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 al., 2021] is 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 Transformer network, by mapping between RASP primitives and Transformer components. Tracr [Lindner et al., 2023] is 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 e KNN; sequence of queries
QM, and a boolean predicate E Q x K. -> (0, The output is an attention matrix A E {0, M with A. predicate (queries, keys,). The select operator is combined with an aggregate operation, which takes input an attention matrix A and a sequence of values E yM and outputs, at each row i, the weighted average of values, with weights given by A,. 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 huran-interppetable programs Transformers, but there is 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 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 simple in-context learning task (described in more detail in Section 4. and is depicted in Figures 2 and 3. In Section 3.3, we show we extend our framework to support 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 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 new variable dedicated address.
Reading from the residual stream. To ensure that each module reads a fixed set of variables, parameterize each projection matrix by one-hot indicator vector over the available variables.
3

---

Classifier weights
tokens
position
attention
attention layer
features & classify
output
actn output
Input:
Output:
Figure We constrain the Transformer to have disentangled residual stream: the token embeddings encode a fixed set of discrete variables in orthogonal subspaces, and each module reads fixed set of variables and writes a new variable to a dedicated address. Here, we depict the residual stream of Tranformer Program that was trained simple in-context learning task (Section 4.1).
Suppose the residual stream encodes m categorical variables, all with cardinality k. resulting in input embeddings x E {0, 1) Nxmk Then each projection matrix W Romk is defined by an indicator [0, Σx= म.
That is, W [ .......... Tm T where L the identity matrix. Concretely, this means that each attention head is associated with three gate vectors, K, Q nV. 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 dedicated, orthogonal subpace, taking inspiration from Tracr. In practice, we accomplish this simply by concatenating the module output to the input-i.e., if layer consists of a single attention head h. then x; [x h(x;-1)]- (This is equivalent adding the output to the input after padding with zeros.) The output each attention head is another
categorical variable with cardinality k the final embedding dimension is then (2 + T. x H) where I, is the number of layers and H the number of heads per layer.
3.2 Transformer Program modules
Next, constrain each module to implement interpretable, rule-based mapping between inputs and outputs. Here, we describe the primary modules in our programs, categorical attention heads; 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 cate, which maps every combination of key and query to a value {0, 1}. We implement simple learnable predicate by associating each attention head with 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, W prodicate {0, 1] with the constraint that each row W Wprddicates sums to one. The self-attention pattern is then defined by a score matrix E {0, 11 NxN with xWQ W predicatee (xWx)T.
Aggregating categorical variables. We adopt 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 A, One-hot (arg max, Sty), where A, denotes the 7th row of the attention matrix. We implement hard attention so that, in the event that there is matching key for query, the model defaults to

---

attention layer
attention layer
Tokens
positions
Eggropete
sttn
Octput
tokens
oggrapete
3130
output
positiions
atin
output
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 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 implement a simple induction head [Olsson al., 2022], mapping each letter token to number that is preceded by the same letter.
Figure 3:
def predicate_ (q_.oss___qpppptiiin x_position):
def predicate_ 2(oken,, atta_1_outpat):
if q.position in (o, 8):
if token in "<pad>", "0", "1", "2", "3", "<g>"
return k_position
if g.position in {1, 9}:
}:
raturn k_position
return attn, 1_output "<pad>"
q.position in (2):
token in {*a"):
return k_position
return attn____uuppn "a"
q.position (3):
token ("b"):
return K_position --
return atta_ abta__1output -- "b"
g.position (4):
token {"c"):
return k_position
return attn___outgun "c"
q-position (s):
token {td"):
k_position
return atin, 1output nd"
g.position in (6):
k_position
attn attn_2_patterm select_closest(
if q.position in (7):
stta Lata__Loutpute, tokens, predicate_2)
return k_position
attn_2_outpee
takens)
attm 1.patern select_ccosest(
positione, positions, predicate_1)
attn attn___ i_outputs Earroaate(attn___ptttee tokens)
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.
attend 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.I.
Optimization. Consider Transformer Program with single categorical attention head, with input variables of cardinality k. This model defined by the indicators π K. π QQ Ty and the predicate matrix W predicate. To learn Transformer Programs, we optimize distribution over these discrete program weights. For each gate ग K, TQ. TV. we define parameter & QQ Φy E R". For row in the predicate matrix W predicate: we define parameters , 2/x E R.,, with ab, defining distribution over the tth row. Referring to these parameters collectively 4, we define distribution over discrete Transformer weights p(0 (b) as the product of categorical distributions, which we optimize using the Gumbel reparameterization [Jang et al., 2017]. Given classification dataset 'D consisting of inputs w and labels y. we minimize the loss
C = B0)p(()+) [Ew.yD [- log w. 0)]
Bw. |o p(y w; 0.)],
#=1
where 0, is a sample from the product of Gumbel-Softnee distributions, with softmax temperature 0, and S is a hyper-parameter denoting the number of samples per training step. Transformer Programs also include a discrete arg max operations hard attention, which we also relax using the Gumbel-Softmax. We anneal the Gumbel temperature over the course of training. As 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, 0" arg maxp p(0) (b), 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 into predicate function, which maps each possible pair of query and key values to value in {0, 1}. The attention operation is then
5

---

implemented by library function, select_closest which selects the closest key satisfying the predicate, or the first token in the sequence there no match.
3.3 Extensions
This framework can be extended to support additional modules, provided that the module can be mapped 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 small vocabulary (i.e., |V] 100), use fixed, one-hot token embeddings. For larger-scale NLP tasks, we learn a factored categorical embedding, representing each token as the product of categorical variables, where 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 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 limited form of numerical attention, which guarantees that the output values are integers with a bounded range. This enables possible imput values. This module loosely corresponds to the selector_ _withh 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 key and query and a numerical variable the value. As above, the module learns binary predicate matrix mapping queries to keys, but aggregates values by taking a weighted sum. That is, given attention scores S E (0, 1] M and value variable var, the output for the 7th token is defined as Σ)n vir]
characterize downstream modules by enumerating the
Feed-forward layers. Finally, implement a simple feed-forward layer, designed to correspond to a lookup-table. Each feed-forward layer reads é input variables, which are designated advance to be either numerical categorical variables, and outputs one categorical variable. For example, assuming the residual stream encodes m categorical variables with cardinality k,, and n numerical vari- ables. The output of a categorical feed-forward layer is defined where Win R [(mk+n)xck is a projection matrix constrained, as above, to read E categorical variables, and f IR tk is an arbitrary transformation. (We implement f as an MLP with one hidden layer). This module can be fully described as a look-up up table, mapping ke possible inputs outputs. An example of learned MLP is illustrated in Figure 9.
One-hot (arg max( fel fx(xVVV Win)).
Experiments
In this section, we learn Transformer Programs for variety of tasks, including simple in-context learning experiment ($4.1), I); the suite of algorithmic tasks introduced by Weiss et al. (2021) ($4.2); and two NL. 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 SC.
4.1 In-context learning
First, by way of illustration, we learn a Transformer Program for toy task (Figure 5), designed to elicit a simple form of in-context learning [Brown et al., 2020]. The input to the model is sequence of alternating letters and numbers that end with a letter. 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 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 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
Deseription
Example
M
Acc.
Reverse
Reverse a string. For each token, the number of occur- rences of that letter in the sequence. For each token, the number of unique tokens with the sume histogram value. Sort the input lexicographical order. The unique input tokens order of fre- quency, using position to break ties. For each position i, is the input up until valid string Dyck-I (T); a valid prefix (P); or invalid (F))
xeverse("aabe" "ebba"
8
3
8
99.79
Histogram
hist("ubbc") "1221"
100.0 4
Double hist.
hist2("abbc")) = "21 22112"
8
3
98.40 4
Sort
sort("chab")> "abbe"
8
99.83
Most-Freg
mont fret__rre"""hhe""" "hac"
8
3
75.69 8
Dyck-1
(yck)""00")="PPP"""""
16
3 8
99.30
Dyck-2
The same us above, but in Dyck-2.
dycx2((()D))== "PPPTPF"
"PPPT""
16
3
4
99.09
Table 1; The RASP tasks, as introduced by Weiss et al. [2021]. We train Transformer Programs on scall-scale instance of each task and report the number of layers (L), attention heads (H), and MLPs (M) used in the best-performing model. kt denotes the variable cardinality, which is fixed at the maximum sequence length for each task.
The learned program achieves perfect accuracy on 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 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, 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_O__ _______________ outputss) as key. the query is a letter, the head attempts to map it to key that has the same letter written to head _adddd____o___ -utttottutttsttt outt is, number preceded by the same letter.
Results.
4.2 RASP tasks
Next, we test whether we can learn Transformer Programs for wider variety of algorithmic tasks. We use the tasks introduced by Weiss et al. [021] 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 2 and setting the maximum sequence length and vocabulary size to for the remaining tasks. As above, we use fixed one-hot token and position embeddings and set the variable cardinality & 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-Histogne (98.4) and Most-Freq (75.69). These results show that, at least on short inputs, Transformer Programs can learn effective solutions variety of algorithmic tasks. On the other hand, in Appendix S 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 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 or 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 4, it increases the likelihood of outputting 4 at positions three through six. Appendix Fig. 11, we use an interactive Python debugger to find 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 variety of non-trivial solutions, composing
7

---

First attentson head:
attn 0 0 outputs
8 If g se carly in the sequence, look for I. IJ q sa tate the sequesce, took for a
<pad>
def predicate, Sarriicat___0(pee8_ opoopiiiionn, token): if position in (0):
Classifier weight
Feature value
return token 13"
if poastion in {1, 23:
if position in (3, 4, 6, 6): return token '4"
-5
if position (7):
return token '0"
</5>
<3>
-15
attn_ select_ (tokens, positions,
predicate_0.00
<pads0
attm_ octn_ttt Segragateeatmm____paaatt tokens)
Output class
Figure 6: Left: This code corresponds an attention head in a Transformer Program that was trained to sort sequences of up eight tokens long, with beginning- and end-of-sequence tokens. At early positions, this attention head checks if there is a the sequence, and at later positions it looks for 4. Right: The classifier weights associated with this feature.
operations at different layers 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 paring Transformer Programs with standard Transformers on RASP tasks with different lengths and vocabulary sizes. Standard Transformers out-perform Transformer Programs across most tasks." 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 50% accuracy gap relative to the standard Transformer in the most difficult setting. We discuss these scaling challenges 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 space constraints, 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,4
test sentences). The data labeled according to IOB2 tagging scheme, with each word being assigned 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 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 MI.P 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 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 Appendix $B.3.
of four entity
Model
D
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
32x4
2
8
81.0
71.8
76.1
Table 2: Named entity recognition the CoNLL 2003 shared task. For each method, 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 the variable cardinality multipled by the number of input variables. The unigram baseline predicts the class most commonly associated with each word in the training set.
'The exceptions are the histogram related tasks. This could be because the built-in numerical attention mechanism provides inductive bias for learning histograms.

---

Favors B-LOC
asts, Copy uard froe pressons taken daf predicate,0, 0.1,pt_..t. k_panision).
position::
29-
17 q.pocition in (2):
k.position
varl: Var1. v30
Classifie weight
e.poxition (9):
var3: Var3. V12
zeturn K.position
17 a.poestion in 643:
attn attn_a_3: Var3.
return k.poxition
9.positiin in [6):
var3: Var3. v10
return k.posstion
attn 0 1: Var3. V20
if i.ponitiin in (S):
k.poestion ==
attn Vare.
atto,_03 0_9_pattern 1amn_0,3,outpppee aggrogate (attn, O.J.pattem, varo, enbeddinn))))
solect, coloct_cleccst/pper clocest/pociticcos positions, prodicate_
var2: Var2. vBo
Favors B-ORG
(b) Code for the attention features.
varl: Var1.v13 V13
Var3(8mm):
attn 0 1: Var3. V16
weight
V00
['Garman'. "talaviaion', , 'Forwign' 'necepappr"
var2: Var2.
-2
VO1 Fankk>' "Johnzon', 'Norrie', 'Service' , c..]
V02
C'cpp '</a>" 'Bank', 'Yori', "Comiiasion
varl: Var1. v02
V03
'In', "Saturday' 'mwtchh' "Az'y
var2: Var2. V16
-8
VO5 L'during' "leader', "quatad" 'manager's "oxee ..d
V04
C*/1, 'up', "no", 'nevsroax', "Attondance"
attn aetn_a_3: 3: Var3.
V06
('Axraz', "TOROKIO', "BALTIHORS' "BOSTO"" L'sasd', "e", eecondd "tald', "b', 'did', 'spokeeran" Auutrraiia,, 'France', 'Spazz", 'England',
'gth', 'other', "eharee'
attn 0 Vare. ve7
VO8
attn 0 1: Var3. V22
VO9
V10
Metherlandd'' 'Pinland', 'countriee' 'Surdich'
B-LOC B D-ORG
(c) The most common words assigned to different values of the Var3 embedding variable.
(a) Feature weights
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 gather information from the context. For example, attention head attn_0 ___ copies one of the embedding variables from the previous position (7b). It promotes the B-LOC label if the value is Var3. VO3, which includes prepositions like "at" and "In" (7c). We represent the embedding variables as Python Enum objects, which facilitates analysis using a debugger.
Figure 7:
Results. The results are in Table 2. We compare the best- perss-erffrmig Transformer Program with 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 predict tags.
Interpretability. We illustrate the resulting program by examining how 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, find that the model copies information from the neighboring words- for example, increasing the likelihood of B-LOC if the word is preceded by preposition like "at" or "In".
5 Related work
Learning programs. Our work has precedent in a variety of existing work on program induction and neuro-symbolic methods [c.g. Reed and De Freitas, 2015, Cai et al., 2017, Andreas et al., 2016, Inala et al., 2020, Cranmer et al., 2020, Kim, 2021]. J. In particular, a long line 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 sequence modeling problems.
Transformers and formal languages. In addition to RASP [Weiss et al., 2021], 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 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., 2018J. Merrill and Sabharwal [2022] show theoretically that any ixxe-pprcisioon Transformer can be expressed as a formula in 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, 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]. motivation we share in this work.
6 Conclusion and discussion
We introduced a method for training Transformers that can be deterministically mapped to discrete, human-readabe 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-analysi. tools debug errors. On the other hand, learning discrete programs introduces considerable modeling and optimization challenges. We conclude by mentioning some limitations 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 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 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- thel-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 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.
Acknowledgeents 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.
References
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE Conference on Computer Vision and Patlern Recognition, pages 39-48, 2016.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint 1607000007 2016.
10

---

Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Computational Linguisties, 48(1):207-219, 2022.
Tolga Bolukbasi, Adam Pearce, Ann Yuan, Andy Coenen, Emily Reif, Fernanda Viégas, and Martin Wattenberg. An interpretability illusion for BERT. arXiv preprini urXiv:2104.07144 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal,
Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems (NeurIPS), 33:1877-1901, 2020.
Jonathon Cai, Richard Shin, and Dawn Song. Making neural programming architectures generalize via recursion. In International Conference on Learning Representations (ICLR), 2017.
Chaofan Chen, Oscar Li, Daniel Tao, Alina Barnett, Cynthia Rudin, and Jonathan K Su. This looks like that: Deep learning for interpretable image recognition. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D Manning. What does BERT look at? An analysis of BERT's attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pages 276-286, 2019.
Ronan Collobert, Jason Weston, Léon Bottou, Michael Karlen, Koray Kavukcuoglu, and Pavel Kuksa. Natural language processing (almost) from scratch. The Journal of Machine Learning Research (JMLR), 12:2493 2537, 2011.
Miles Cranmer, Alvaro Sanchez Gonzalez, Peter Battaglia, Rui Xu, Kyle Cranmer, David Spergel, and Shirley Ho. Discovering symbolic models from deep learning with inductive biases. Advances in Neural Information Processing Systems (NeurIPS), 33:17429-1177
2020.
Andrew Cropper and Sebastijan Dumančić. Inductive logic programming at 30: new introduction. Journal of Artificial Intelligence Research, 74:765-8515 2022.
Jacob Devlin, Ming-Wes Weei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT), pages 4171-4186, 2019.
N Elhage, N Nanda, C Olsson, T Henighan, N Joseph, B Mann, A Askell, Y Bai, A Chen, T Conerly, et al. A mathematical framework for Transformer circuits. Transformer Circuits Thread, 2021.
Nelson Elhage, Tristan Hume, Catherine Olsson, Neel Nanda, Tom Henighan, Scott Johnston, Sheer ElShowk, Nicholas Joseph, Nova DasSarma, Ben Mann, et al. Softmax linear units. Transformer Circuits Thread, 2022.
Atticus Geiger, Zhengxuan Wu, Christopher Potts, Thomas Ieard, and Noah D Goodman. Finding alignments between interpretable causal variables and distributed neural representations. arXiv preprint
2023.
Mor Geva, Avi Caciularu, Kevin Ro Wang. and Yoav Goldberg. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. arXiv preprint 46060 2022.
Angeliki Giannou, Shashank Rajput, Jy yong Sohn. Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped Transformers as programmable computers. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association of Computational Linguistics (TACLL) 8: 1711, 2020.
Yiding Hao, Dana Angluin, and Robert Frank. Formal language recognition by hard attention trans- formers: Perspectives from circuit complexity. Transactions of the Association of Computational Linguistics (TACL), 10:800-8 810, 2022.
11

---

Jeevana Priya Inala, Yichen Yang, James Paulos. Yewen Pu, Osbert Bastani, Vijay Kumar, Martin Rinard, and Armando Solar-Lezama, Neurosymbolic Transformers for multi-agent communication. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pages 13597-13608, 2020.
Henrik Jacobsson. Rule extraction from recurrent neural networks: laxonomy and review. Neural Computation, 17(6): 1223-1263, 2005.
Sarthak Jain and Byron C Wallace. Attention is not explanation. In North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT), pages 3543-3556, 3556, 2019.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations (ICLR), 2017.
Yoon Kim. Sequence-to-sequemn learning with latent neural grammars. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pages 26302-26317, 2021.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 11Xiv:1412.6980, 2014.
Quentin Lhoest, Albert Villanova del Moral, Yacine Jernite, Abhishek Thakur, Patrick von Platen, Suraj Patil, Julien Chaumond, Mariama Drame, Julien Plu, Lewis Tunstall, et al. Datasets: A community library for natural language processing. In Empirical Methods in Natural Language Processing (EMNLP): System Demonstrations, pages 175-184, 2021.
David Lindner, János Kramár, Matthew Rahtz, Thomas McGrath, and Vladimir Mikulik. Tracr: Compiled transformers as a laboratory for interpretabilityy. arXiv preprint arXiv:2301.0566. 2023.
Thomas McGrath, Matthew Rahtz, Janos Kramar, Vladimir Mikulik, and Shane Legg. The Hydra effect: Emergent self-repair in language model computations. arXiv preprint 1xXix:3307.15771 2023.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in GPT. Advances in Neural Information Processing Systems (NeurIPS), 35:17359- 17372, 2022.
William Merrill and Ashish Sabharwal. Transformers implement first-order logic with majority quantifiers. arXiv preprint arXiv:2210.02671 2022.
William Merrill, Ashish Sabharwal, and Noah A Smith. Saturated Transformers are constant-depth threshold circuits. Transactions of the Association of Computational Linguistics (TACL), 10: 843-856, 2022.
Stephen Muggleton and Luc De Raedt. Inductive logie programming: Theory and methods. The Journal of Logic Programming, 19:629-679, 1994.
Hiroki Nakayama. seqeval: A Python framework for sequence labeling evaluation, 2018. URL https: //gi thub. com/ whakkiiwwrks//seeeeeeeo seqeval. Software available from http://github.co//hakii works/seqeval.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In International Conference Learning Representations (ICLR), 2023.
Nostalgebraist, Interpreting GPT: The logit lens. https://www. wttp://wwww.immmmnfroformmmm wtw.ooorrmrooogg/ooo/ooo/ooo org/posts/ http://DDdddaoeveee/nneeerreeeee
gpt- the- 2020. Accessed: 2023-05-15.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11888 11895, 2022.
Bo Pang and Lillian Lee. A sentimental education: Sentiment analysis using subjectivity summariza- tion based on minimum cuts. In Asxociation for Computational Linguistics (ACL), pages 271-278, 2004.
12

---

Bo Pang and Lillian Lee. Seeing stars: Exploiting class relationships for sentiment categorization with respect to rating scales. In Association for Computational Linguistics (ACL), pages 115-124, 2005.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-perffrmance deep learning library. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019.
Ali Payani and Faramarz Fekri. Learning algorithms via neural logic networks. arXiv preprint 2019.
Jeffrey Pennington, Richard Socher, and Christopher Manning. GloVe: Global vectors for word representation. In Empirical Methods in Natural Language Processing (EMNLP), pages 1532-1543 2014.
Felix Petersen, Christian Borgelt, Hilde Kuchne, and Oliver Deussen. Deep differentiable logic gate networks. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations (ICLR), 2022.
Scott Reed and Nando De Freitas. Neural geggramme--intepreeteer In International Conference on Learning Representations (ICLR), 2015.
Cynthia Rudin. Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead. Nature Machine Intelligence, 1(5):206-215, 2019.
Erik Tjong Kim Sang and Fien De Meulder. Introduction to the CoNLL-2003 shared task: Language- independent named entity recognition. In Human Language Technologies: North American Chapter of the Association for Computational Linguistics (HLT-NAACL), pages 142-147, 2003.
Ian Tenney, Dipanjan Das, and Ellie Pavlick. BERT rediscovers the classical NLP pipeline. In Association for Computational Linguistics (ACL), pages 4593-4601, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems (NeurIPS), 30, 2017.
Ellen M. Voothees and Dawn M. Tice. Building a question answering test collection. In Proceedings of the 23rd Anmual International ACM SIGIR Conference on Research and Development in Information Retrieval, page 200-207. Association for Computing Machinery, 2000.
Fulton Wang and Cynthia Rudin. Falling rule lists. In Artificial Intelligence and Statistics (AISTATS). pages 1013-1022. PMLR, 2015.
Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: A circuit for indirect object identification in GPT-2 small. In International Conference on Learning Representations (ICLR), 2023.
Qinglong Wang, Kaixuan Zhang, Alexander G Ororbia II, Xinyu Xing, Xue Liu, and C Lee Giles. An empirical evaluation of rule extraction from recurrent neural networks. Neural Computation, 30(9):2568-259 2018.
Zhuo Wang. Wei Zhang, Ning Liu, and Jianyong Wang. Scalable rule-based representation learning for interpretable classification. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Extracting automata from recurrent neural networks using queries and countercxamples. In International Conference on Machine Learning (ICML), pages 5247-5256. PMLR, 2018
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like Transformers. International Conference Machine Learning (ICML), pages 11080-11090. PMLR, 2021.
13

---

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 (ACI-IJCNIP)) 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.
14

---

A Method details
A.1 Categorical attention
As described in Section 3.2, implement categorical attention by associating each attention head with a boolean predicate matrix, W predicate E {0, 1) where is the variable cardinality, with the constraint that each row W predicate, sums to one. The aelf-attention pattern is then defined by score matrix S {0, 1 N with S xWa q W predicate (x W к) ) To ensure that each query token attends to a single key token, we use hard attention, defining the attention pattern at position i as A, One-hot (arg max, S.,). During training, we define A; Gumbel-SSftmaxx(S5
Defaulting to the beginning of sequence. We implement hard attention so that, in the event that there is no matching key for query, the model defaults to attend to the first token in the sequence. Let S E {o, M_NNNN denote the score matrix for a sequence with length N. We define a modified score matrix S such that, for each row i,
S. = [Sh + (1 maxy Sia) S
ifj=1
otherwise.
Breaking ties. Furthermore, we implement the attention mechanism so that, in the event that there more than one matching key. the model attends to the closest match. Given the score matrix S E (0, 1) JNNN we define the modified score matrix S such that, for each row i, St. = S. b-jj- where b E [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 j = 1, to I/N, when j jl N, with 6g = 1/N to bias the query at position against attending to itself. This is similar to types of relative positional bias that have been proposed for regular Transformer-biee 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 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 S {0, 1 M 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 S E {0, 1] M and value variable var, the output for the 7th
token is defined as Σ varj]] At the input layer, there is a single numerical variable, ones, which is frozen and equal to for all positions; attention heads that read ones as the value variable equivalent the selector_widde 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 standard Transformer, because computes weighted sum rather than weighted average. But numerical attention head can be implemented a standard Transformer by composing an attention head with a feed-forward layer, using the beginning- soae-oo--quuueee of-seueence token to allow the model attend to nothing. This is how the selector__ wiith primitive is implemented in Tracr [Lindner et al., 2023].
Feed-forward layers. In RASP, feed-forward layers are used to implement arbitrary element- wise 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 decompiled. Each foed-forward layer reads input variables, which are designated in advance to be either numerical or categorical variables, and outputs one new categorical variable. We convert fred-forward layers to programs by enumerating all possible inputs and forming a lookup-table. For all experiments, we set f = 2. For categorical attention, given variable cardinality k. there are k² 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 the value
15

---

def nun. wun.predicate.c 1__ token, k_ Rokke))):
def nun_nip_ 0_1 (nun attn_t_ttt_....tttt __1__otppr))) key nut_atta_O__ 1_tta_____uuuuus if key in {0, 1): 4
g_token in ("0"):
return k_token =- "0"
alif q_token in {"1"):
return k_token -- "1"
retura 0
q_token in {"2"):
return k_taken == "2" ellf q, token in ("3");
num_nip
=
nun_nlpp (k0)
return k_token =- "3"
for k0 in aun_atta 0 __1_outptte
clif g_tcken in ("d"):
k_taken "4"
elff q.token in ("5");
return k_token "g"
Feature value
num mlp 8 1 outputs
Classifier weight
elif q_token in {"<a>">:
0-
0
return k_token -- "Spad>"
nun aattn_0_1_pptteee select(
4-
-18
tokena, takens, nun.predicato_o_]
0_i.n_,,,_uuputte aggregate_aunt nun nun_attn_0__ p_t_pattern, ones)
Output class
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, attention head (lef)) calculates the histogram for each position. An MIP (lop right) reads the histogram values and outputs a value of O 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 i-increases the likelihood that the double-histogram value is 1 or 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, is unlikely that another number appears the same number of times. An output of 4 (meaning 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 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 and 9.
Attention. Each attention head is represented by a predicate function, which takes as input key and query and outputs value in {0. 1}. Transformer Programs, all keys and queries are categorical variables with cardinality k. so the predi cate function can be defined by enumerating the possible query values, which we do with 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 illustrated in Figure In the first attention head (lef),, each query position attends to the previous position (with the exception of the first token), we cannot apply any compression. In the second attention head (bottom right), fifteen out of the sixteen possible query values are mapped to single key value, we combine them into single branch.
MIPs. We convert feed-forward modules functions by enumerating all possible inputs and forming a key/value lookup-table. For all experiments, we set each MLP to read & = 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 value of if the input is 0 or 1 and value of 0 otherwise.
16

---

First aftention head: copy previous def predicate_0_0(e 0_0(q__oeiiiion, k_position);
and previous 8 Dubputa 13 if it sees "()" or *()". det mlp_0_0(token, acta_ (ttn.0.0_outpat))
MLP:
g.position in fo, 13):
return K_position 12
key (token, atta_0_0_outpur if key in (()), ")*),
elif g_pesition in (1):
k_position
(*)*, "}*),
allf q_peaition in {2}:
("(", ")"),
return K.position
(*)', ')*),
elif q_position in (3):
(*)*, '))):
return k.position
returo 4 elif key in (C)", '0", (*), *(")):
elif q.position in (4):
return k_position
alif q-pesition {
return k_position ==
elif key in {(*(', ??",
q_position {6):
(*(,, 1)",
return k_position --
(*), '(",
elif q.position in {
(*(", 1)")
return k_position
(?", '("):
alif g.pesition in {8):
return 0
k_peaition
7
elif q.position in {9)
11p_0_0.
1_p_0____uupppppppppppp 1.0_oupttt [[
return k_position
mip_0_0(k0, for k1 zip(tokens, attn_0_0_outpurte
alif q_position in {10}:
return k_position
elif q.position {11):
return k_pasition 10
2nd layer attention: check for *(3" or "()"
elif q.pesition {12}:
def predicate_1 aedicate....(paatt..0 1.posiiii mlp_0_0.000 ouuoput):: if position in (0, 1, 2, 4, 5, 6, T, 8, 9, 10. 11, 12, 13, 14, 15): 13
return K_position -- 11 q.position (1d):
return k_position = 13 elif g.position in (15):
return mlp_0_0_outpor
olif position in (3):
k_position 14
mlp_0_0_output
attn_0_0_ p.t.par select_cloneat (poaition,, pasitions,
attn_ 1_2_______pp_ pattarn select_closeat(
1_2______
predicate_ attn_ atta_0_0_ 0_0_outputs aggregate 0.0_pttter, tokens)
positione, predicate_1.2)
1ita_1_2_ouupurte aggregate(( attn_1,2_patterm mlp.0_0_outputs)
Figure 9: An example of circuit in a program for Dyck-2. For this task, the imputs consist of strings from the vocabulary { (.), {. At each position i, the model must output T if the string up to position i a valid string in 2 P it is the prefix of a valid string; and otherwise. string is valid if every parenthesis is balanced by parenthesis of the same type. This circuit recognizes an invalid pattern, where left parenthesis (C or {) immediately followed by a right parenthesis the wrong type (J or respectively). First, 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 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 represented integers in the corresponding program. To improve readability, we replace these integers with symbolic values where appropriate. At the input layer, 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 (C),, C, }; the first attention head reads tokens as the value variable, so we can automatically determine that attention, ae__oo___________ 0_0_outpute variable takes on values of the same type; finally, the MLP reads tokens and attention. lattiooo__oonttiiiiieeeeee __ouutuuussssttt input variables, so we define the mlp mlp_0_0 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-non mnn-to-many attention predicates Our Transformer Programs use a form of binary categorical attention that associates each query value with single key value. This is enforced by parameterizing the attention head with a predicate matrix W Wpredicate E {0, j (where & the variable cardinality) and constraining W predicate that each row contains only single non-zero entry. This choice results
17

---

in more concise programs, as each attention head can be characterized in terms of one-to-on mapping between key and query values, but 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 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 understand.
Score- attention As 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 oot attention, where each attention head associated with both a binary function, predicate: Q K. - {0, 1}. and a continuous function, score: Q x K. + R. that maps each combination of key and query to a real value. Given a query q and keys kj...., ,.....,u kn. the attention between q and K is defined as I if both predicate(g: k;) = 1 and score(q, k;) = max]== score(g, kz), 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 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(g, k.) i.
would be relatively straightforward to extend Transformer Programs to include 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 [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 attends to token j. our default, one-to-one form of attention, this observation means that j is the closest position such that predicate(y:: ky) = 1. In score-based attention, the observation also implies that there is no position such that score(q, ke) > score (q: ky), 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, the cost of some interpretability tradeoffs. For example, it may no longer be possible to decompile MLPs by enumerating all possible arguments.
B Experiment details
In this section we describe additional implementation details for the experiments in Section 4. Code to reproduce the experiments is available at https://github. wwwpprincetonn-nnenennoeeeeee n1p/ TransformerProgramt
B.1 In-context learning
Data. For our in-context learning experiment (Section 4.1), sample 20,000 sequences of length 10. The first token is a beginning-of-sequenee token, and the remainder of the sequence is formed by randomly sampling many-to-one mapping between letter types (either a, b, c, d) and numbers (0,
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 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 S 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 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 PoToh [Paszke al., 2019] and use the Adam optimizer [Kingma and Ba, 2014].
18

---

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 exception of the Dyck tasks, we sample inputs by sampling tokens uniformly from the vocabulary. For the Dyck tasks, follow Weiss et al. [2021] and sample with bias towards strings with 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 with length |s| I N, and append N - | 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 and-of-sequence token, eos. Following Weiss et al. [2021], we report the token-level accuracy.
instances respectively. With the
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. Simililarly, the ML.s are evenly divided between MI.Ps with two categorical inputs, and ML.s 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.I), we similarly do grid search for number of layers (2, 3) and attention heads (4, 8) but otherwise follow the ayper-parameters described by Weiss 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 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 caiin/vaiidaiion/esstt 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 batch size of 32 and perform 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 Glo Ve 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 Additional results
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 larger input vocabulary. Because we use one-hot encodings for the token and position embeddings, these values determine the cardinality, A of the categorical variables used in the program. We compare Transformer Programs with standard Transformers, for each model performing grid search over the number of layers and
19

---

Reverse
Hitt 2-H62
Sort
Most Freq
120
Sett
[see Awi5
1x--1, Mi:66
20
0
Figure 10: RASP aceuracy after increasing the size of the input vocabularies ((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 test accuracy of the model that performs best 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 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 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. Table 3, we compare results on RASP tasks with only categorical variables. The experiment setting is otherwise the same 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, Transformer Programs still achieve good performance on the other four tasks. These include balanced parenthesis languages (Dyck-1 and Dyc which are most naturally solved by keeping count of unmatched parentheses at each position.
Dataset
M
Acc.
Reverse
3
99.75
Hist
3
8
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
99.43
Dyck-2
3
4
2
99.17
Table 3: Accuracy of Transformer Programs on RASP tasks using only categorical variables, with number of layers (L.), attention heads (H), and MI.Ps (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 lasks. Complete, generated programs are available at https://githun wtt.hu/.ccoo/piinttooonn.tt.eeee Transf oransformerPPogrramm
Code analysis tools. One advantage of representing models as Python programs is that we can use off frr-the-shelt code analysis tools to interpret them. Figure 11 illustrates 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 Appendix A.3. The program length includes the basic library functions used to We use automated Python code formatter, which applies a number of standard style conventions and, in particular, enforces a maximum line length of 88 characters.
'https:/ peattps:/ggithann
20

---

PV
wotty
?
picgram sort
soitay
2 get, Jestures
progame mart EDrt.N O prL_iattuee
predisatte 3.3issiiiin,, Whot::
Reed arkedcing siddle assitions,
9 322
att patters select closesttrakess pastrions, predicate !
atts. putppts aggraqeeejittn 1.pttteer, EinnAL.
wergy -12 -pwgactteentteeeee wwwprreeegggggapgggggreeeoo
Delittions (1)
36
d:f predicate, R.J T.0 pesition?:
37
Store the 2nd toses at the beglinning of sequence enteddingg token ("ap""y
positiin
114
tohen
reture posilion
Resd from arbedcing free later positions. clal positiin D, (11
5tore 6th toses the ass sequence cmbo3ding
121
127
retum taken "cgp"
121
124
LLvvte_ettter sctect witosst_iittttttkkooo. pasitimr, arctiostic 1 arth 1.1 Aitputs aggrepatelstn, 2 potters, MED A.L.ppt
all eiif tooes (18, "2" "2", "3", 14%)
Store
123
OUTPUT
000.0 CONSOLE
TERMINAL
Fil (ng bet Vec.
position
otta 8.1 artouts
tokers
[Taas" 12" 14", 11, 137, 181, 33, "o/s"")
4
46 stts Luitgate aggregarejattnn & patters, telend
att1 5.1 pattere a select clossttpaettton tekens, predicate a
p4r. 2,, 12" 12%
12%, 131
32 attr p
L.iovts
125, (47)
def predicate 2ipositiinn
126
Figure 11: After converting a Transformer Program into Python code, we can analyze using off- the-shelf debugging tools. Here, inspect 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, attention head (left) reading a value from a first-layer attention head (right). Inspecting the code, can that this pair of heads has the effect of propagating an early-position token the later positions and a late-position token to the earlier positions.
find circuit involving two attention heads, with second-laye layeed
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.
What information do the attention heads read? Because each attention head reads a fixed set of named variables, 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 positi as key and query variables and tokens as the value. At higher layers, positi ons 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 positi iisiiiinn ons 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.
characterize how information flows through the programs by examining
C.3 Optimization challenges: case study
What are the challenges to scaling Transformer Programs to 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, 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.
write Transformer Program that would achieve perfect accuracy,
Setting We focus on the Reverse task. This task can solved by the following program, adapted from Weiss et al. [2021], using Transformer with two layers and one attention head per-layer:
Färst layer attention length aggregate (tokee,,,, tokens, lambda q, k: k "</S>"), positions)
Pérst- layer HLP
targets one_ho (length positions)
21

---

Cat. head: query
Cat. head: key
Cat. head: value
Num. head: query
Num.
head: key
Num. head: value
tok
mlp
tok
attmmij
wbtn x
pos attm nlt
tak
wittn
atth
0.21
0.79
0.00
0:00
0.17
0:83
0.02
0.00
0.33
0:11
0.00
0.00
0:33
0.62
0.00
0.00
0.62
0:38
0.00
0.02
1.00
9.00
0.34
0.29
0.01
0:01
0.00
0.47
9.00
0:38
0:21
0.00
0.17
0.46
9:30
0.25
0:12
0.75
For each of the RASP tasks, learn 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 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 previous attention head (attn); or the output of a previous MI. (mlp). We aggregate over RASP programs and compare categorical attention heads (left) and numerical attention heads (right).
Figure 12:
Attention
MLP1
Attention 2
Read Predicate
Read
Read
Predicate
Accuracy
23.2/23.6/24
99.9/99 99/9/99.9/80.8
37.9/40 3
17.1/13. 7/200.2
95. 1t.11/94.1// 1/955.
99. 1/83.9/78..
35.5/44.1/41.
Table 5: Results on the Reverse task ([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 (V/) initialize the other weights randomly (-), and then train the model, reporting the (per-token) accuracy 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 MIPP 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 positions, ==), tokens)
The first attention layer copies the position of the end-of-sequenc token to determine the length of the sequence; the larst-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 - 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 two-layer Trans- former 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 manu- ally initialize: (1) the parameters associated with и,, TQ. and TV, which identify the key, query, and value variables read; and (2) the parameters associated with the predicate matrix W predicate. In the MLP, we manually initialize the parameters associated with the input projection matrix W... 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.
22

---

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 particular, finds the internal MI.P weights needed to calculate
solve the problem successfully. the difference between the two input variables (length and position) and output one-hot encoding of the result. Interestingly, 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 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 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 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 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 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 Appendix B.3. We hold out 10% of each training set to use as 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 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 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.
23

---



20
verl_erbeddinns: 9
crbeddingg:
3
how
var2_enbeddinns:
many
vara_erbeddings:
2
people
10
positions:
attn_ 0_0_outputs:
Token
in
attn 0_1_outputs::
the
outputs:
score
world
attn 3_3_outputs:
nlp _________tttt____
speak
-5
nlp 0 outputs:
french
outputs:
-10
attn outputs:
attn 1 2_outuuts: 15
-15
attn 1.3. oouuptts:
ABBR
DESC
ENTY
HUM
LOC
HUM
outputs:
nlp 1 outputs: 14
Class
ABOR
Bess ENTY may RUy
Loc
(a) Classification scores for each token embedding.
Feature-level scores for the "world" token.
varo...
var2...
7
7
X
astn, 1. Copy vard from early positions. def predicate.1_ enbdddin ponition):
how
4
if var2_anbedding in [0, 4, 5, 6, 10, 12, 183:
are
11
4
0
position
olif anbdding in (1, 7, 8, 9, 11, 14]:
for
2
10
zeturn positicn --
elif cabedding
when 6
14
6
elif var2_cabedding in (16): return position =-
position
there
year
attt p.ta....tattrrn cllectt.coccest
people
positions, varl_enbediinge, prodicate_1_1) atta___i_outpen aggregace( 1tta_1_1_pptterr varo, cnbe_enbe)))nne)
time
6
15
0
date
6 14
6
(c) The code for attn 1_1 outputt feature.
computing the
(d) A subset of the embedding table, filtering to words with var0 embedding values of
Figure 13: Visualizing the predictions for example from the TREC question classification dataset. We classify sequences by pooling the final-layer token embeddings, so we can visualize the classifica- tion 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 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. fatt_____ouuputt copies the var0 cerr=_emmeddinnge variable from first or fourth position-
7). This attention head generally position-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 occur in numerical questions.
24


Reviews:
[{"review_id": "VIkjTu5DYo", "reviewer": "Reviewer_Kcxc", "summary": "This paper proposes a new technique for training Transformer models while ensuring that the learned weights obey certain constraints, designed to facilitate mechanistic interpretability. In particular, the technique works by (1) freezing certain weight matrices throughout training (e.g., to always write the output of an attention layer to a previously unused linear subspace of the residual stream) and (2) introducing Gumbel-softmax sampling at various steps of the computation, and annealing the temperature down during training to ensure that at convergence, those computations perform discrete operations. The resulting Transformers can be automatically converted into symbolic programs that read, write, and perform lookup-table-style computation with a fixed number of categorical and integer variables. The authors train (quite small) Transformers in this way, to solve several tasks, including named entity recognition and sentiment classification on short sequences. The authors also provide some analysis of the programs extracted from these Transformers.", "questions": "- At the end of training, your Gumbel-Softmax's have been annealed to Categorical distributions, but these distributions are not necessarily *peaked* or low-entropy categoricals. Therefore, your extracted program (which takes the argmax of the categorical) may behave quite differently from your trained network (which randomly samples the categorical). Have you investigated (a) the extent to which the learned categoricals are peaked at a single mode, and (b) the difference in task performance / accuracy between the final learned Transformer and the extracted program?\n- Why limit the predicate matrices to contain only one ‘1’ per row? Is it just because this constraint is more amenable to the Gumbel-softmax training procedure?\n- In the description (Appendix A.1) of the categorical attention mechanism, you appear to use an argmax operation that is not differentiable. During training do you actually use Gumbel-softmax? If so, how are the resulting samples used to compute the output of the attention layer? (More generally, it would help if the paper more clearly laid out exactly where in the architecture all the Gumbel-softmax samples are used. Perhaps algorithm boxes in the style of https://arxiv.org/pdf/2207.09238.pdf could be provided for each component of the neural architectures you build?)\n- For the in-context learning experiment, did you try (1) training an ordinary Transformer of the same size on the task? and/or (2) training a larger (i.e., overparameterized) Transformer program on the task? I’d be curious about the results of either experiment, if you did run them. Under the ‘lottery ticket hypothesis,’ overparameterization is an important part of how algorithms are learned by gradient descent — it would be interesting to see what this looks like in a Transformer program. E.g., in a larger model, do we end up with a small interpretable subprogram that correctly completes the task?\n- Can you clarify how the word embeddings are trained? What does it mean that the 100-d GloVe embeddings are “randomly projected” to e.g. 32*4=128-d embeddings? Do you use a randomly generated dense 100x128 matrix? Then how are the embeddings trained? Is there a Gumbel sample for each of the categorical variables?\n- Why not include MLP layers in the NLP tasks?\n- For the RASP tasks that were not successfully solved by learned Transformer programs, do you believe this is because (1) good solutions exist, but are not found by gradient descent, or (2) your subset of RASP is too limited and the tasks cannot be solved with this subset? If (2), what are the prospects for handling a broader subset of RASP?", "limitations": "Limitations are adequately discussed, but I do think it could be useful to say a bit more about the challenge of scaling this approach -- what are the key obstacles & prospects for addressing them?", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This is a neat paper that advances two exciting lines of research:\n\n- RASP and Tracr been helpful as intuition pumps for how trained Transformers *could* work, but it has not been clear whether the (correct) RASP algorithms researchers manually developed could be discovered automatically by gradient descent, and if not, what different properties these learned Transformer algorithms would have. This paper proposes a method for learning Transformer programs that could be used to start to answer these questions.\n- A key challenge in mechanistic interpretability research (e.g. on Transformer circuits) is that Transformers trained by (unconstrained) gradient descent appear to collapse multiple conceptually distinct concepts via superposition. This paper proposes training methods that may help to mitigate superposition by forcing the network to read, write, and manipulate individual ‘variables’ that live in orthogonal subspaces of the residual stream.\n\nAlthough terse in some places (perhaps due to length constraints), the methodology is for the most part clearly described, and the limitations of the approach seem sensible for a first attempt at this kind of constrained training. Another strength is that the framework seems to be extensible—the examples in the paper paint a clear enough picture of how to modularly extend the target symbolic programming language with new features (and then how to extend the neural net architecture and training procedure accordingly).", "weaknesses": "- Although the motivation is interpretability, I was not convinced that the learned programs in the experiments were particularly interpretable. It would have been nice to see more detailed analysis of the programs. Outside the toy in-context learning example, are there any interpretable algorithms that are learned? Can examples on which the learned programs fail be traced back to particular ‘bugs’ in the symbolic programs? Can the symbolic programs be used to engineer adversarial examples on which the Transformer will fail? Do ordinary transformers trained without constraints also fail on these adversarial examples (indicating that they might be using similar flawed algorithms)? How stable are the learned programs across different training runs?\n\n- It is unclear how scalable the presented technique is, to more complex problems, longer sequences, or bigger networks. I also would have appreciated more discussion / investigation of how easy or tricky the optimization was, and why it appears to work more poorly for larger Transformers (e.g., that handle larger sequence lengths). What are the key challenges in scaling up? Do the Gumbel gradient estimates become too noisy, e.g.? Or are there other main challenges?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! Please see our comments in the GR about whether the programs are interpretable, and the main challenges to scaling up this approach.\n\n> More analysis of the programs\n\nThank you for the great suggestions! Regarding debugging: We do not necessarily expect that standard transformers will make the same errors, but we agree it would be interesting to measure the extent to which (a) whether standard Transformers and Transformer Programs learn similar solutions, and (b) the similarity between Transformer Programs with different random seeds, and we will add this analysis to our next draft.\n\n> Are the categorical distributions peaked?\n\nWe did not investigate these questions in detail but we will run this analysis and add it to our final draft. Anecdotally, we found that the accuracy of the argmax model closely approximates the final training accuracy.\n\n> Why limit the predicate matrices to contain only one ‘1’ per row?\n\nThis was a design choice that we made because it leads to programs that are (qualitatively) easier to understand. It is also possible to allow the predicate matrix to have multiple 1’s per row, and using a Gumbel-sigmoid function. In our preliminary experiments, we trained models with this form of predicate matrix and found they obtained similar classification performance but were qualitatively more difficult to understand. We omitted this discussion for lack of space but will add it to the appendix in our final draft.\n\n> Argmax attention / location of Gumbel-softmax operations\n\nThank you for pointing this out. We use the Gumbel-softmax during training, in the place of hard attention (this is mentioned in section 3.2 but not in the appendix–we will correct this in the next draft). Using the notation from Section 2, the output of the attention layer is defined as $AxW_V$, where A is an attention matrix with each row adding up to one, and S is a matrix of attention scores. With argmax attention, $A_i = \\text{One-hot}(\\text{arg}\\max_j S_{i, j})$. During training, we define $A_i = \\text{One-hot}(\\text{Gumbel-softmax}(\\log (S_i + \\epsilon)))$ for some small $\\epsilon$. We will explain this more clearly in our next draft, and include a table laying out the different components of the model.\n\n> Additional in-context learning experiments\n\nWe did try (2), training larger Transformer programs. We found that over-parameterized models are more likely to learn the correct solution (as measured by the percentage of training runs that achieve perfect test accuracy). For example, if we train a model with two heads per layer, we are more likely to find that one pair of heads implements the induction head, perhaps analogously to the lottery ticket hypothesis. One interesting question for future work is whether we can identify the interpretable subprogram automatically, for example by adapting pruning methods. We will add this discussion to our updated draft.\n\n> Can you clarify how the word embeddings are trained?\n\nBefore training, we initialize the word embeddings to $E = E_G  \\Pi$, where $E_G \\in \\mathbb{R}^{|\\mathcal{V|} \\times 100}$ denotes the GloVe embedding matrix, and $\\Pi \\in \\mathbb{R}^{100 \\times 128}$ is a matrix with each entry sampled independently from $\\mathcal{N}(0, 1/128)$. This embedding matrix is then split into four $|\\mathcal{V}| \\times 32$ matrices (one for each variable), and we sample categorical word embeddings by taking a Gumbel sample for each row of each of the four matrices. We will explain these details more clearly in our updated draft.\n\n> Why not include MLP layers in the NLP tasks?\n\nWe used one MLP layer in all of the NLP tasks, but this is not clearly reflected in the paper–thank you for pointing this out. We will update the method description to be more clear.\n\n> Explaining the RASP tasks that were not successfully solved\n\nWe believe that the main issue is (1): good solutions exist but are not found by our optimization procedure. This is because it is possible to manually construct Transformer Programs that solve RASP tasks perfectly, so the performance limitations must be attributed to optimization challenges. One caveat is the “less than” predicate, which appears in RASP programs for  tasks such as sorting. This predicate is difficult to express using our version of categorical attention, where each query value attends to a single key value, but could be expressed using the one-to-many version of attention discussed above. It is straightforward to extend our approach to include this kind of attention. In our preliminary experiments, we found that, even using one-to-many attention, the  Transformer Programs did not learn the “less-than” predicate. We will add these experiments to our updated draft."}, {"review_id": "HcYNixFhF5", "reviewer": "Reviewer_vQnY", "summary": "This paper presents an approach to constraint the parameter space of Transformers so that when trained, the resulting weights can be directly translated to a human-readable program, e.g. in Python. To do so, attention is modified so that the internal model representation is, in practice, composed of a set of distinct variables that the attention processes read/write to. The base approach focuses in attention-only models (no MLP layers), but the authors propose extensions to support MLP layers.\n\n- page 5: line 165-166: \"For row in the predicate matrix\" -> there seems to be some word missing around \"row\".\n", "questions": "- Section 3.2: the explanation for how to extract the part of the program that corresponds to the attention layers is clear. But how about the MLP layers? How do you infer what the MLP layers have learned to do to conver that into a program? [edit: this is explained immediately below, but it's still unclear; so, for EACH MLP in EACH LAYER of the model, you would have to extract a look-up table that will be added to the Python program?]\n- (more a comment than a question): concerning MLPs. I wonder if it would be possible to define a library of commmon operations (e.g. addition/multiplication/etc., or whatever makes sense with other categorical values), and allow non-learnable, pre-defined MLP weights to be used via some gating (e.g., the model choosing to use one of these built-in methods, or a custom-learned weight matrix). to increase the chance of having an interpretable function in the resulting model, rather than having one look-up table per MLP layer.\n", "limitations": "See weaknesses above.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "I found the idea of constraining the parameter space extremely interesting. While the types of programs that can currently be learned seem a bit restricted (e.g., as the number of variables grows, the model/program could potentially grow very large), I think this is a very promising direction, and a greag foundation to build on.\n", "weaknesses": "While it seems clear that programs for simple tasks can be interpreted, it is not clear that many tasks (specially vision/natural language tasks) actually even afford interpretable algorithms. So, perhaps future work needs to focus on more complex primitives for the programming language supported by the models that would simplify the programs, or enable higher-level constructs that are interpretable.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the comments!\n\n> Higher-level constructs for more complex tasks\n\nThank you for the great suggestions! We agree that future work should explore extending the framework to include more complex primitives, or perhaps automatically identifying higher-level programming abstractions in learned programs (e.g. as in [1]). We will mention this direction in our updated draft.\n\nPlease also see the GR for our more general comment about whether complex tasks admit interpretable algorithms.\n\n[1] Ellis et al., 2023. DreamCoder: Growing generalizable, interpretable knowledge with wake–sleep Bayesian program learning.\n\n> Converting MLP layers into a program\n\nThank you for pointing this out – this is correct. For each MLP, for each layer of the Transformer, we extract a lookup table that will be added to the Python program. (However, within a single Transformer layer, each MLP module is converted into a single lookup table, regardless of the number of hidden layers within the MLP itself.) The size of the lookup table is determined by the number and cardinality of the input variables, although in practice we can compress the lookup tables in situations where the model does not use all of the output values. Some example MLP functions are presented in the appendix in Figures 8 and 9. We will update Section 3.2 to explain this more clearly.\n\n> Defining a library of common operations for MLPs\n\nThank you for the suggestion! This is a great idea and it would be simple to incorporate this type of module into this framework. We chose to treat MLPs as lookup tables following Tracr, but we agree the resulting lookup tables are not necessarily easy to interpret, and incorporating common operations like addition would be a good alternative. We will mention this suggestion in our final draft as a possibility for future work.\n"}, {"review_id": "3wMac7Vzkf", "reviewer": "Reviewer_HW9X", "summary": "There has been prior work on compiling programs written in a domain-specific language (RASP) into a transformer that emulates the function. In this paper, the authors target the opposite direction of training transformers and generating RASP programs that are faithful to the model's computation. The authors do this by proposing an architecture that serves as a discretized version of a transformer, allowing translation to a discrete, interpretable python program after training. They find that for a variety of tasks, it is possible to 1) learn a solution in this architecture and 2) compile it into RASP.", "questions": "No questions beyond the one in the weaknesses/limitations.", "limitations": "Are there tasks where you tested this method and either the code was uninterpretable or the discretized transformer couldn't achieve high accuracy, unlike the standard transformer? If not, are there any tasks where you wouldn't expect this to work?\n\nOther than this, the authors have adequately addressed limitations.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This work demonstrates a novel architecture that allows for direct translation to code. The paper does a great job of demonstrating their method for synthetic tasks. Though prior work has noticed sparsity improved interpretability, this work leverages it for full programmatic translation. I hope this can inspire future work in interpretable neural networks.", "weaknesses": "1. Learning solutions to the RASP tasks makes it easier for the model to learn interpretable solutions. The paper would be stronger if it targeted other tasks where the compilation is more surprising.\n2. How well can a discrete transformer solve problems relative to the original transformer? I imagine the discretized model has a lower representation capacity or worse optimization. I think it's important to experimentally/theoretically analyze what is lost as a result of this architectural choice, considering it's not a standard transformer in many key aspects.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! Please also see the GR for our comments about whether the solutions are always interpretable, and about the comparison between Transformer Programs and standard Transformers.\n\n> Learning solutions to the RASP tasks makes it easier for the model to learn interpretable solutions. The paper would be stronger if it targeted other tasks where the compilation is more surprising.\n\nIn addition to the RASP tasks, we trained models on two non-synthetic NLP tasks: named entity recognition (Section 4.3) and text classification (Appendix C.3). In the paper, we show examples from the learned programs to illustrate how they allow us to interpret the solution. For example, in Figure 7, we identify some of the features that the program uses to distinguish location entities from organization entities. These include attention heads that copy lexical information from the preceding and subsequent tokens–e.g., a token is more likely to be the beginning of a location entity if the following word is a word like “Germany”, “Italy”, or “Netherlands.” On a question classification dataset (Appendix Fig. 11), we can see that the model focuses on tokens at the beginning of the input–for example, predicting that the input is a numerical question if it starts with a question word like “when”, or relevant nouns like “year”, “time”, and “date.”\n\nWe chose to experiment on the RASP tasks because these reflect a variety of algorithmic subproblems. We would also note that, even though these tasks can be solved with concise, interpretable programs, it is still non-trivial to automatically learn these solutions given only input-output pairs.\n\n> How well can a discrete transformer solve problems relative to the original transformer?\n\nThe paper and appendix include experimental results comparing the original transformer and our discrete transformers on the NLP tasks: Appendix C.4 (Table 6) shoes that the Transformer Programs are are either competitive with or slightly worse than the standard transformer on four  text classification tasks, and Section 4.3 (Table 2 in the main paper) shows that the Transformer Program actually slightly outperforms a standard transformer in our NER experiments.\n\nOn the other hand, the Transformer Programs under-perform a standard transformer on some RASP tasks, and on longer inputs. We did not include RASP results with a standard transformer, but Weiss et al. [1] report that standard transformers get > 99% accuracy on tasks with a sequence length of up to 100 tokens, while our Transformer Programs perform worse on longer sequences (Appendix C.1, Table 3). (We will include RASP results with a standard transformer in our next draft.) This is an optimization issue rather than an expressivity issue: we can write Transformer Programs that achieve perfect accuracy at arbitrary sequence lengths, but our optimization procedure fails to find these programs in practice. We hope to address these optimization challenges in future work. (In general, Transformer Programs can be as expressive as RASP models, with the caveat that we have restricted the size of the feed-forward layers; like RASP models, the discrete transformers are less expressive than standard transformers due to the limitations of hard attention [2].)\nWe will highlight this discussion in greater detail in our final draft.\n\n[1] Weiss et al., 2021. Thinking Like Transformers.\n\n[2] Hao et al., 2022. Formal language recognition by hard attention transformers: Perspectives from circuit complexity.\n\n> Are there tasks where the code was uninterpretable or the discretized transformer couldn't achieve high accuracy?\n\nAs noted above, the discretized model performs worse than a standard transformer on RASP tasks with longer sequences, and performs slightly worse on some text classification datasets. The Transformer Programs also perform relatively poorly on the Most-Freq task (Table 1), which requires the model to count, sort, and identify unique elements. In general, we might expect Transformer Programs to perform worse on tasks like Most-Freq that require the composition of a series of higher-order functions, because of optimization challenges (i.e. local minima).\n\nRegarding tasks where the code was difficult to understand: In Section 4.2, we highlight some subsets of the _sort_ program with non-obvious interpretations. For example, Figure 6 illustrates a circuit of two attention heads that propagates early-position tokens to later positions and late-position tokens to early positions by first copying them to the beginning- and end-of-sequence tokens. In general, for larger-scale tasks, it might become more difficult to understand the code at the level of individual attention heads. Some interesting directions for future work are  to develop tools for identifying higher-level abstractions in large programs, and for imposing an inductive bias in favor of sparser programs. Please also see our comment in the GR about whether the learned programs are interpretable."}, {"review_id": "LYliWQt63L", "reviewer": "Reviewer_TaJd", "summary": "_Background_: RASP is a tensor processing language which provides a language to hand-write transformers. Tracr presents a framework to compile RASP programs to a transformer architecture automatically. \n\nThe authors of this project are interested in distilling a transformer architecture to an equivalent program in RASP. This is challenging because RASP is not a sound language: the behavior of many trained transformers is equivalent to that of a single RASP program. The authors go around this problem by constraining the transformer architecture in strategic places so that there is a deterministic mapping from a transformer to an equivalent program in RASP. The authors present a model distillation procedure to construct such reduced transformers (called Transformer Programs) automatically. The authors test the Transformer Programs on three tasks: 1) in-context learning, 2) recovering handwritten RASP programs, and 3) simple real-world NLP tasks. They observe mixed results: transformer programs can fully solve (1), they cannot generalize to larger sequences for (2), and they are better than non-sequential baselines on (3) but worse than generic transformers. The authors also show the utility and interpretability of the distilled programs by analyzing the programs using off-the-shelf debuggers.", "questions": "* Constraint 1 “the token embeddings encode … fixed set of variables… in orthogonal subspace. What is a variable here? I’m going to assume variable=token, but it would be good to define this. \n   * Does this mean that encoding 100 variables requires a 100 dimensional vector? Does this assumption violate some of the properties we know about transformers? For instance: nonlinear superposition[1]?\n* 155: “require that each query token attend to a single key token”. This constraint seems very limiting. Do most sequential problems where transformers are used adhere to this constraint?\n* 256 “results degrade on longer sequence tasks.” This seems concerning. Are these results for the trained model or for the extracted program? Ideally, model distillations are more robust to OOD tasks than the model itself[2]. Is there any indication for why this might be the case?  \n. 259 “our solution might be different from human written solutions.” This seems concerning because RASP programs attempt to upper-bound the number of transformer layers that might be needed for computation. Do we have any indication for why this might be happening? Is this maybe because the human written solutions reuse previously computed values?\n\n1. https://transformer-circuits.pub/2022/toy_model/index.html\n2. https://arxiv.org/abs/2006.11287\n", "limitations": "The authors have addressed the limitations of the method adequately to the best of my knowledge.", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "Originality\n* RASP programs are inherently unsound: one RASP program is equivalent to many trained transformers. Solving this problem by constraining the transformer to elicit a bijective mapping is a simple yet elegant idea. \n\nSignificance\n* This paper shows a simple method to distill transformers into interpretable programs. This work opens the door for work in program synthesis, static analysis, and model checking to be used for sequential processing tasks.\n* Many downstream tasks rely on transformers to processing temporal sequences of data into labels. While the transformer architecture presented is constrained, an interpretable transformer can be very useful for tasks where high degree of interpretability and accuracy are required. \n\nClarity and Quality\n* The framework is well motivated, and the writing is clear and precise.", "weaknesses": "__Weaknesses__\n\nI have three main concerns (paraphrased from the questions section below):\n* I’m concerned that the constraints mentioned here violate properties[1] that the mechanistic interpretability community have identified as being intrinsic to the utility of transformers.\n* Poor robustness: Model distillation has been observed to produce models that are more robust to outliers while being slightly less accurate[2]. However, in experiments it seems that the distilled models are neither robust to outliers nor as accurate as baselines, which makes it harder to use this as a drop-in replacement.\n* Lack of Reusability: An interesting property of the RASP programs (and RASP distilled transformers) is that many tasks can be efficiently solved by precomputing useful primitives. The algorithm presented here does not seem to model this property. For instance, in the reverse proposed in RASP, the authors' solution precomputes length of the sequence and then uses length to flip the indices of an identity matrix. The derived transformer also captures this behavior. Does the Transformer Program also capture this behavior? If not, why not? \n\nClarity:\n* 135. subpace -> subspace\n* Related work: I defer to the authors, however, I feel adding a section on other model distillation approaches would be relevant. In the context of model distillation, [3] seems similar to this work in learning a transformer (for multi-agent communication) and then by discretizing it (using MCMC) to obtain an interpretable policy.\n\nOverall, I've given the paper a __borderline rejection__. The problem statement presented has promise to have significant impact, however, the experimental results do not instill confidence in the downstream utility of this method and I believe this paper will benefit from another round of review. I'm more than willing to change my recommendation after the rebuttal period.\n\n====\n__EDIT__: I misinterpreted the paper as a _model distillation_ paper instead of a continuous relaxation one. In light of this, and other clarifications presented by the authors, I'm raising my score to a __weak accept__.  \n\n1. https://transformer-circuits.pub/2022/toy_model/index.html\n2. https://arxiv.org/abs/2006.11287\n3. https://arxiv.org/abs/2101.03238\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! We believe there is some misunderstanding in the main concerns, and we have addressed them below. We kindly request the reviewer to reconsider the evaluation, and we  are more than happy to further clarify if anything is still unclear.\n\n> Do the constraints violate properties (e.g. superposition)?\n\nOur constraints do prevent the model from using feature superposition. As noted by [1], superposition is a major obstacle to interpretability. Our work is one attempt to solve it. It is not clear that superposition is “intrinsic to the utility of transformers” – in [1], the authors suggest that superposition is useful but might not be necessary; in fact, every model can be seen as equivalent  to a “hypothetical disentangled model” that uses higher-dimension embeddings. On the other hand, we agree that superposition is useful, because it enables the use of lower-dimensional representations. One avenue for future work is to extend our framework to enable the use of more compressed representations, while still preserving interpretability.\n\nPlease let us know if there are other properties that raise concerns (aside from superposition).\n\n> Poor robustness / relationship to model distillation\n\nWe believe this concern might be based on a misunderstanding:  our work is not really a model distillation approach, or analogous to the method in [2] (which targets a different domain and architecture–physics modeling with GNNs). Our approach is better thought of as a continuous relaxation of discrete search, where the search space is the set of Transformer networks with discrete weights. As such, we do not have any prior reason to expect that the discrete programs will be more robust. On the contrary, our primary goal is interpretability, and we expect that Transformer Programs might sacrifice accuracy for this goal. Nevertheless, we show that Transformer Programs can achieve competitive accuracies on synthetic RASP tasks and real-world NLP tasks. We also highlight some areas where our method struggles to learn robust programs, due to the challenges of discrete optimization. Replacing large transformers with drop-in interpretable counterparts is an ambitious project, and we see this work as a first step. Addressing these optimization challenges is an avenue for future work.\n\n> Lack of Reusability\n\nOur experiments show that the Transformer Programs do in fact learn to compose operations, computing primitive values at lower layers and using these results at higher layers to compute higher order patterns. For example: the simple in-context-learning program learns an  induction-head pattern  (section 4.1), and the double-histogram program (appendix Figure 9) learns to compute the first-order histogram in the first layer, and consume this value in the second layer. In Appendix C.2 (Figure 10), we show that higher-layer attention heads often read information from intermediate-layer attention and MLP layers, showing that the programs do indeed learn to reuse previously computed values.\n\n> Results on longer sequence tasks\n\nTo clarify, this line is referring to models that are trained on longer sequences, not to OOD generalization. As noted above, our work is not a model distillation approach, and we don’t expect other results about distillation to apply here. The reason the results degrade on longer sequences is related to optimization challenges: our programs tend to find solutions that use relatively shallow features of the input, rather than the most robust, parsimonious solution. As discussed above, addressing these optimization challenges will require significant research advances, which we are leaving for future work.\n\n> Learned solutions vs. human-written solutions\n\nWe did not expect in advance that the Transformer Programs would learn the same solutions as the human-written programs for RASP, which represent just one possible way to solve the problems. In fact, in the RASP paper (Table 1), they find that even standard Transformers do not always learn the same solution as the human-written program. As noted above, many of the programs described in the paper learn to reuse previously computed values. However, there could be other inductive biases that make some solutions more difficult for our model to learn–for example, it might be harder for our model to learn arbitrary feed-forward operations.\n\n> What is a variable here?\n\nWe use the word variable in the same sense as a variable in a computer program: a named container that takes on a particular value when the program is executed. In the example in section 3.1, the program has four variables: _tokens_, _positions_, _attention 1 outputs_, and _attention 2 outputs_. We follow the design of the Tracr compiler and use a one-hot encoding for categorical variables, meaning the embedding size is equal to the number of variables times the number of possible values they can take on. We will define this more precisely in our final draft.\n\n> requiring each query token attend to a single key token\n\nIn this paper, we support two kinds of attention: categorical attention, where each query token attends to a single key token, and numerical attention, where each query token can attend (uniformly) to any number of key tokens. This is a limiting constraint, but we adopt it from RASP/Tracr, and it is important for ensuring that the resulting programs are discrete and easy to interpret. Moreover, existing work has shown that these attention patterns are empirically common in Transformer LMs [4], and, theoretically, are expressive enough to model a large class of formal languages [5].\n\n[4] Merrill et al., 2021. Effects of Parameter Norm Growth During Transformer Training: Inductive Bias from Gradient Descent.\n\n[5] Merrill et al., 2022. Saturated transformers are constant-depth threshold circuits..\n\n> Related work\n\nThank you for the pointer! We will add a discussion and related references in our final draft."}]

Author rebuttal:
We thank the reviewers for their thoughtful comments! The reviewers have made a number of great suggestions that we will incorporate into our next draft, including clarifying some details about the method; providing more discussion about optimization challenges; and including some additional analysis of the learned programs. We respond to each of the reviewers in detail, and to some common issues below. 

> Are the programs interpretable?

Some of the reviewers questioned whether the learned programs are actually interpretable (Kcxc, HW9X), or whether more complex tasks even admit interpretable solutions (vQnY). We provided a number of examples in Section 4 and Appendix C illustrating interpretable subsets of programs for different tasks, including sorting (Figure 5); NER (Figure 7); double histogram (Appendix Figure 8);  Dyck-2 (Appendix Figure 9), and a classification task (Appendix Figure 11). For more complex tasks, the program can be thought of as a collection of individually interpretable feature functions, rather than a single interpretable algorithm. For example, a circuit in the Dyck-2 program checks for invalid bigrams using an attention layer and an MLP. In the NER program (Fig. 7), the model is more likely to predict that a token is the beginning of a location entity if the following word is a word like “Germany”, “Italy”, or “Netherlands.”

On the other hand, we acknowledge that the learned programs can still be complicated and non-intuitive. (For example, Figure 6 illustrates a subset of a _sort_ program whose interpretation is not immediately obvious.) We would emphasize that this method always provides a baseline level of interpretability: even when the program is complicated, 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. Nonetheless, we agree that the programs can still be difficult to interpret, and we think a good 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.

> Scaling, expressivity, and optimization challenges

Several of the reviewers asked if we could provide more details about how the Transformer Programs compare to standard Transformers (HW9X, TaJd); provide more insight into the cases where the Transformer Programs under-perform standard Transformers (Kcxc, HW9X, TaJd); and discuss the key challenges to scaling up this approach. We provided relevant analysis in the paper by experimenting on RASP tasks with longer sequences (Appendix C.1) and on NLP tasks (Section 4.3 and Appendix 3.3), but we add some additional discussion here. In particular, Reviewers HW9X and Kcxc ask whether the main limitation of this method is (a) limited expressivity or (b) optimization challenges.

We think that the main scalability challenges are optimization challenges. In particular, Transformer Programs are expressive enough to represent robust solutions to the majority of RASP tasks, but in practice they do not learn these solutions. This is reflected in the results in Table 3 (in the appendix), where we learned programs for longer sequences. For example, we can construct a Transformer Program that achieves perfect accuracy on the _reverse_ task, by first calculating the sequence length, then calculating the target position indices using a feed-forward layer (_target_index_ = _length_ - _position_ + 1), and then copying the token from the target index. In ablation experiments, we found that Transformer Programs could learn to calculate these intermediate values if they were trained directly, but failed to learn this decomposition when trained end-to-end on the _reverse_ task. We think that the underlying optimization issue is that it is more difficult for these discrete models to escape from local minima. (We will add this discussion to our next draft.)

Addressing these challenges will require more research on discrete optimization. While our approach might not be a drop-in replacement for regular Transformers just yet, we think that our work represents a significant first step towards this goal, and can help lay the groundwork for future research.


Meta-review: This paper introduces a novel method to learn transformer programs, which are executable programs that can be interpreted by a transformer model. The reviewers recognized that the paper presents an interesting and original idea that has potential applications in natural language processing and program synthesis. They also pointed out some issues such as the interpretability and scalability of the generated programs, the optimization challenges, and performance comparing to the baseline Transformer models. The authors provided a rebuttal that addressed most of the issues raised by the reviewers and promised to improve their work in the final version. Based on the reviews, the rebuttal, and the discussion, I recommend an ACCEPT to this paper. I believe this work will inspire further research and applications in this direction and suggest that the authors address the remaining concerns of the reviewers and provide more details and comparisons in the final version.
Decision: Accept (oral)