Title: Learning Transformer Programs
Venue: NeurIPS 2023
Accepted: yes
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

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. 1

1 Introduction

Transformers [Vaswani et al., 2017] have become the predominant neural network architecture in machine learning, representing the state-of-the-art in natural language processing (NLP) and increasingly in computer vision, molecular modeling and other domains. Prominently, large Transformer-based language models [LLMs; Brown et al., 2020] have demonstrated impressive general-purpose capabilities and are now widely deployed as 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 to audit, debug, and maintain. This black-box quality poses practical challenges and limits the usefulness of these models in high-stakes applications.

As a result, a considerable body of work has aimed at improving our understanding of Transformers. In NLP, 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 [e.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]; in any case, they do not provide a complete or faithful description of the algorithm that the model implements. More recently,

Input

Target

a 1 b 2 b 2 a

1

b 3 c 4 a 3 c

4

d 2 c 4 a 2 b

unk

relaxed

1. Train

u

🔎

use debugger

automatic

analysis

discrete

Transformer Program

2. Sample

Python program

3. Decompile

🤖

research on mechanistic interpretability has attempted to gain an algorithmic understanding of these models, with the goal of reverse-engineering Transformers into human-interpretable 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 can 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, we aim to train Transformers that are mechanistically interpretable by design. We take inspiration from RASP [Weiss et al., 2021], a 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-based optimization and then automatically decompile them into human-readable programs.

To this end, we propose a method for learning Transformer Programs -Transformers that are constrained to implement human-interpretable algorithms. First, we introduce a set of constraints that restrict the Transformer to lie within an interpretable subpace of the parameter space, by which we mean the subset of parameter values that can be mapped to a family of human-readable programs. Second, we develop a 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 to 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 1.

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, we 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'circuits' used to solve different sub-problems, and to debug model errors. Overall, we aim to demonstrate that Transformer Programs represent a promising first step towards building intrinsically interpretable machine learning models.

2 Background

Transformer architecture. The Transformer [Vaswani et al., 2017] is a neural network architecture for processing sequential data. The input to the Transformer is a sequence of tokens w = w 1 , . . . , w N ∈ V in a discrete vocabulary V . At the input layer, the model maps the tokens to a sequence of d -dimensional embeddings x 0 ∈ R N × d . In the standard Transformer, this initial embedding is defined as the sum of a learned token embedding and a positional embedding. Each subsequent layer i consists of a multi-head attention layer (MHA) followed by a multilayer perceptron layer (MLP): x i = x i -1 +MLP i ( x i -1 +MHA i ( x i -1 )) . 2 Following the presentation

of Elhage et al. [2021], multi-head attention can be written as

where H is the number of heads, d h is the attention embedding dimension, W h Q , W h K , W h V ∈ R d × d h are referred to as the query , key , and value projections respectively, and W h O ∈ R d h × d projects the output value back to the model dimension. By default, we 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, it is a two-layer feedforward network: MLP( x ) = σ ( xW 1 ) W 2 , where W 1 ∈ R d × d m , W 2 ∈ R d m × d , and σ is a non-linear activation, commonly the ReLU function. The output of the model is a sequence of token embeddings, x L ∈ R N × d . 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 a neural network processes information. Informally, a Transformer can be viewed as a series of nodes (attention heads and feed-forward layers) reading and writing information to the residual stream , which refers to the sequence of token embeddings x 0 , . . . , x L ∈ R N × d computed at each layer. Each node 'reads' from the residual stream via a projection matrix W in ∈ R d × d h ; calculates some function f ( x i -1 W in ) ; and then 'writes' to the residual stream via another projection: x i = x i -1 + f ( x i -1 W in ) W out, with W out ∈ R d h × d . 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 to 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 et al., 2021] is a 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 a Transformer network, by mapping between RASP primitives and Transformer components. Tracr [Lindner et al., 2023] is a 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 ∈ K N ; a sequence of queries ∈ Q M ; and a boolean predicate ∈ Q × K → { 0 , 1 } . The output is an attention matrix A ∈ { 0 , 1 } M × N with A i,j = predicate ( queries i , keys j ) . The select operator is combined with an aggregate operation, which takes as input an attention matrix A and a sequence of values ∈ V N and outputs, at each row i , the weighted average of values , with weights given by A i . 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 human-interpretable programs to Transformers, but there is no 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 a 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 a simple in-context learning task (described in more detail in Section 4.1) and is depicted in Figures 2 and 3. In Section 3.3, we show we extend our framework to support a 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 a 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 a new variable to a dedicated address.

Reading from the residual stream. To ensure that each module reads a fixed set of variables, we parameterize each projection matrix by a one-hot indicator vector over the available variables.

tokens

position

attn 1

output

attn 2

output

q

k

v

out

attention

layer 1

attention

layer 2

Extract

features

&  classify

/gid00053

Classifier

weights

q

k

v

out

Input:

0 1 2 3 <unk>

Output:

<bos> b 1 a 2 a 2 b

pad

?12m

bos

ab℃dθ123456789θ123456789θ123456789

Suppose the residual stream encodes m categorical variables, all with cardinality k , resulting in input embeddings x ∈ { 0 , 1 } N × mk . Then each projection matrix W ∈ R mk × k is defined by an indicator π ∈ { 0 , 1 } m : ∑ m i =1 π i = 1 . That is, W = [ π 1 I k ; . . . ; π m I k ] ⊤ , where I k is the k × k identity matrix. Concretely, this means that each attention head is associated with three gate vectors, π K , π Q , π V , 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 to a 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 i consists of a single attention head h , then x i = [ x i -1 ; h ( x i -1 )] . (This is equivalent to adding the output to the input after padding with zeros.) The output of each attention head is another categorical variable with cardinality k , so the final embedding dimension is then (2 + L × H ) × k , where L is the number of layers and H is the number of heads per layer.

3.2 Transformer Program modules

Next, we constrain each module to implement an interpretable, rule-based mapping between inputs and outputs. Here, we describe the primary modules in our programs, categorical attention heads; we 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 , which maps every combination of key and query to a value in { 0 , 1 } . We implement a simple learnable predicate by associating each attention head with a 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 predicate ∈ { 0 , 1 } k × k , with the constraint that each row W predicate ,i sums to one. The self-attention pattern is then defined by a score matrix S ∈ { 0 , 1 } N × N , with S = xW Q W predicate ( xW K ) ⊤ .

Aggregating categorical variables. We adopt a 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 i = One -hot ( arg max j S i,j ) , where A i denotes the i th row of the attention matrix. We implement hard attention so that, in the event that there is no matching key for a query, the model defaults to

q

k

v

select(q, k, predicate)

positions

positions

tokens

aggregate

attention layer 1

attn 1

output

q

k

v

select(q, k, predicate)

tokens

attn 1

output

tokens

aggregate

attention layer 2

attn 2

output

attend to a 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.1.

Optimization. Consider a Transformer Program with a single categorical attention head, with m input variables of cardinality k . This model is defined by the indicators π K , π Q , π V and the predicate matrix W predicate . To learn Transformer Programs, we optimize a distribution over these discrete program weights. For each gate π K , π Q , π V , we define a parameter ϕ K , ϕ Q , ϕ V ∈ R m . For row in the predicate matrix W predicate, we define parameters ψ 1 , . . . , ψ k ∈ R k , with ψ i defining a distribution over the i th row. Referring to these parameters collectively as Φ , we define a distribution over discrete Transformer weights p ( θ | Φ ) as the product of categorical distributions, which we optimize using the Gumbel reparameterization [Jang et al., 2017]. Given a classification dataset D consisting of inputs w and labels y , we minimize the loss

where θ s is a sample from the product of Gumbel-Softmax 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 in 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, θ ∗ = arg max θ p ( θ | Φ ) , 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 h into a predicate function, which maps each possible pair of query and key values to a value in { 0 , 1 } . The attention operation is then

implemented by a library function, select_closest , which selects the closest key satisfying the predicate, or the first token in the sequence if there is no match.

3.3 Extensions

This framework can be extended to support additional modules, provided that the module can be mapped to a 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 a small vocabulary (i.e., |V| < 100 ), we use fixed, one-hot token embeddings. For larger-scale NLP tasks, we learn a factored categorical embedding, representing each token as the product of m categorical variables, where m is another hyperparameter. That is, each word embedding is the concatenation of mk -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 a 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 a limited form of numerical attention, which guarantees that the output values are integers with a bounded range. This enables us to characterize downstream modules by enumerating the possible input values. This module loosely corresponds to the selector_width 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 as key and query and a numerical variable as the value. As above, the module learns a binary predicate matrix mapping queries to keys, but aggregates values by taking a weighted sum. That is, given attention scores S ∈ { 0 , 1 } M × N and value variable var , the output for the i th token is defined as ∑ N j =1 S i,j var [ j ] .

Feed-forward layers. Finally, we implement a simple feed-forward layer, designed to correspond to a lookup-table. Each feed-forward layer reads ℓ input variables, which are designated in advance to be either numerical or categorical variables, and outputs one new categorical variable. For example, assuming the residual stream encodes m categorical variables with cardinality k , and n numerical variables. The output of a categorical feed-forward layer is defined as z = One -hot(arg max( f θ ( xW in )) , where W in ∈ R ( mk + n ) × ℓk is a projection matrix constrained, as above, to read ℓ categorical variables, and f : R ℓk → k is an arbitrary transformation. (We implement f as an MLP with one hidden layer). This module can be fully described as a look-up table, mapping k ℓ possible inputs to k outputs. An example of a learned MLP is illustrated in Figure 9.

4 Experiments

In this section, we learn Transformer Programs for a variety of tasks, including a simple in-context learning experiment (§4.1); the suite of algorithmic tasks introduced by Weiss et al. [2021] (§4.2); and two NLP tasks: named entity recognition and text classification (§4.3). Additional implementation details for all experiments are provided in Appendix §B, and we provide further ablations and analysis of the generated code in Appendix §C.

4.1 In-context learning

First, by way of illustration, we learn a Transformer Program for a toy task (Figure 5), designed to elicit a simple form of in-context learning [Brown et al., 2020]. The input to the model is a sequence of alternating letters and numbers that end with a letter. If the letter has appeared already in the sequence, the model should output the number that followed it. If the letter has not appeared before, the model should output an unk token. We train an attention-only Transformer

with two layers and one attention head per layer, with a vocabulary of four numbers and four letters, and trained on sequences of up to 10 tokens long. The token and position variables are fixed, one-hot indicators, meaning that the input embeddings encode two variables and the final embeddings encode four: tokens , positions , and one output variable for each attention head. The cardinality of each variable is k = 10 . For this task, we use a causal attention mask.

Results. The learned program achieves perfect accuracy on a held-out test set, and we can observe that the model learns to compose attention heads according to the induction heads pattern identified by Elhage et al. [2021]. The residual stream is illustrated in Figure 2, and the attention heads in Figure 3. The first attention head learns to read the position variable for both key and query. At each position, it attends to the key at the previous position and writes the value of the tokens variable as the value. The second attention head reads tokens as the query and the output of the previous head ( head_0_0_outputs ) as key. If the query is a letter, the head attempts to map it to a key that has the same letter written to head_0_0_outputs -that is, a number preceded by the same letter.

4.2 RASP tasks

Next, we test whether we can learn Transformer Programs for a wider variety of algorithmic tasks. We use the tasks introduced by Weiss et al. [2021] to illustrate the RASP language, which are listed in Table 1. We train on small-scale instances of each task, setting the maximum sequence length to 16 for Dyck-1 and Dyck-2, and setting the maximum sequence length and vocabulary size to 8 for the remaining tasks. As above, we use fixed one-hot token and position embeddings and set the variable cardinality k to be equal to the maximum sequence length. For this setting, we introduce numerical attention and MLPs. We equip each model with an equal number of categorical and numerical attention heads, and categorical and numerical MLPs, fixing the number of MLP input variables to two, and perform a grid-search over the number of layers, attention heads, and MLPs per-layer.

Results. The results of this experiment are in Table 1. On five out of seven tasks, our method is able to find a program that gets more than 99% accuracy. The exceptions are Double-Histogram (98.4) and Most-Freq (75.69). These results show that, at least on short inputs, Transformer Programs can learn effective solutions to a variety of algorithmic tasks. On the other hand, in Appendix §C.1, we find that the results degrade for longer inputs, suggesting that the model may not learn robust solutions to all tasks. Additional, we can observe that the learned models often use more layers and attention heads than we might expect. For example, Weiss et al. [2021] present RASP solutions for Sort and Reverse requiring only two layers and one attention head, while our best models use three layers and eight attention heads (four categorical and four numerical). These observations indicate that the solutions found by our method might be quite different from human-written solutions.

What kinds of programs does the model learn? To get a better sense for how the model solves these tasks, we examine a program that the model learns for the Sort task. We analyze a model with two layers and four attention heads per layer, which achieves greater than 99% accuracy, and compile it into a Python program of approximately 300 lines, excluding the classifier weights. We find that the program uses a variety of heuristics for sorting these sequences. For example, Figure 6 depicts the code for one of the first-layer attention heads. This head looks for a 1 or a 4 in the sequence. If it finds a 1 , it increases the likelihood of outputting 1 at positions one or two, and if it finds a 4 , it increases the likelihood of outputting 4 at positions three through six. In Appendix Fig.11, we use an interactive Python debugger to find a more sophisticated circuit, composing attention heads at two layers. We provide more examples in Appendix §A.2 and analysis of the learned programs in Appendix §C.2. In general, we find that the model learns a variety of non-trivial solutions, composing

<pad>0

1

2

3

4 
 


Output class

<pad>

0

1

2

3

4







Feature value

attn_0_0_outputs

15

10

5

0

5

Classifier weight

operations at different layers to compute higher-order patterns. These solutions include both brittle heuristics and sensible strategies for the different tasks.

Comparison to standard Transformers. In Appendix Fig. 10, we provide additional results comparing Transformer Programs with standard Transformers on RASP tasks with different lengths and vocabulary sizes. Standard Transformers out-perform Transformer Programs across most tasks. 3 Interestingly, we observe that both models show similar trends, with model performance degrading for tasks with longer sequences and larger vocabularies. Nevertheless, the Transformer Programs deteriorate more dramatically, with over a 50% accuracy gap relative to the standard Transformer in the most difficult setting. We discuss these scaling challenges in more detail in Section 6.

4.3 NLP tasks: named entity recognition and text classification

Now, we turn to standard NLP tasks including named entity recognition [Sang and De Meulder, 2003] and text classification [Pang and Lee, 2005, 2004, Zhang et al., 2015]. (Due to space constraints, we present text classification results in Appendix §C.4.) We first experiment with the English portion of the CoNLL-2003 Named Entity Recognition task [Sang and De Meulder, 2003]. This dataset consists of sentences from Reuters news articles (14,041 train/3,250 validation/3,453 test sentences). The data is labeled according to an IOB2 tagging scheme, with each word being assigned to one of four entity types ( PER , ORG , LOC , MISC ) or marked as not an entity ( O ). We filter the dataset to sentences with at most 32 words and use a vocabulary of the 10,000 most common words, replacing the remaining words with an unknown token. For these experiments, we use only categorical attention heads, with one categorical MLP per-layer. We fix the variable cardinality at 32 and perform a grid search over number of embedding variables, layers, and heads. The embedding parameters are initialized with 300-dimensional GloVe embeddings [Pennington et al., 2014]. We compare the Transformer Program with a standard Transformer, also initialized with GloVe embeddings with grid search of model dimension and number of layers and heads. More details are provided in Appendix §B.3.

(a) Feature weights.

(c) The most common words assigned to different values of the Var3 embedding variable.

Results. The results are in Table 2. We compare the best-performing Transformer Program with a standard Transformer and with a unigram baseline, which predicts the tag most frequently assigned to each word in the training data. The Transformer Program achieves reasonable performance, on par with the standard Transformer. In particular, the Transformer Program surpasses the unigram baseline, demonstrating that the method learns to make use of contextual information to predict tags.

Interpretability. We illustrate the resulting program by examining how it distinguishes between two commonly confused classes, location entities ( B-LOC ) and organizations ( B-ORG ). To see how the model makes this distinction in general, we extract the linear classifier weights and identify the features with the largest difference between the two classes (Figure 7). The highest ranked features include word-level features (that is, components of the word embeddings); position information; and features computed by the attention heads. Working backwards through the program, we find that the model copies information from the neighboring words-for example, increasing the likelihood of B-LOC if the word is preceded by a preposition like 'at' or 'In'.

5 Related work

Learning programs. Our work has precedent in a variety of existing work on program induction and neuro-symbolic methods [e.g. Reed and De Freitas, 2015, Cai et al., 2017, Andreas et al., 2016, Inala et al., 2020, Cranmer et al., 2020, Kim, 2021]. In particular, a long line of work on Inductive Logic Programming [Muggleton and De Raedt, 1994, Cropper and Dumanˇ ci´ c, 2022] has sought to learn symbolic logical programs from data, and a number of recent works have used neural networks to search for discrete logical expressions using differentiable reparameterizations [Payani and Fekri, 2019, Wang et al., 2021, Petersen et al., 2022]. We differ from these methods in targeting programs for the Transformer and focusing on sequence modeling problems.

Transformers and formal languages. In addition to RASP [Weiss et al., 2021], a body of research has explored the connection between Transformers and formal languages. Much of this has aimed to

formally characterize the expressivity of Transformers with hard attention [Hahn, 2020, Merrill et al., 2022, Hao et al., 2022]. Giannou et al. [2023] show how Transformers can act as programmable computers by designing a Transformer that can execute programs written in a single-instruction language. Another line of work has attempted to extract deterministic automata, or rules, from neural networks, largely focusing on recurrent neural networks [Jacobsson, 2005, Wang et al., 2018, Weiss et al., 2018]. Merrill and Sabharwal [2022] show theoretically that any fixed-precision Transformer can be expressed as a formula in a type of first-order logic. In contrast to this work, our goal is to design a more interpretable Transformer, rather than extract rules from an existing network.

Interpretable machine learning models. Some prior work has introduced architecture changes aimed at making Transformers more interpretable, including sparse attention [Zhang et al., 2021] and changes to the activation function [Elhage et al., 2022]. These methods make some components of the Transformer qualitatively easier to understand; in contrast, our method results in a model that can be fully described by a discrete program. More generally, a growing body of work has sought to develop intrinsically interpretable machine learning models [e.g. Wang and Rudin, 2015, Chen et al., 2019, Rudin, 2019], a motivation we share in this work.

6 Conclusion and discussion

We introduced a method for training Transformers that can be deterministically mapped to discrete, human-readable programs. We showed that our method is capable of learning effective solutions for a several synthetic and real-world tasks, and that the resulting programs are easy to interpret. In particular, converting Transformers into programs offers practical benefits-for example, we can use standard code-analysis tools to debug errors. On the other hand, learning discrete programs introduces considerable modeling and optimization challenges. We conclude by mentioning some limitations of our method, which we believe represent promising avenues for future work.

Scaling challenges. While we were able to learn a number of effective Transformer Programs, we found that they can struggle to learn on longer inputs, and to learn parsimonious solutions. Are these errors because the framework is not expressive enough, or are they due to optimization issues? Our results suggest that the main issue is optimization. This is evident from the fact that, for many tasks, we can write a Transformer Program that would achieve perfect accuracy, but our method fails to learn these solutions. We conduct a more detailed case study of some optimization failures in Appendix C.3. Therefore, future work is needed to develop better methods for discrete optimization. Additionally, our experiments use a restricted set of modules that may be less expressive than standard Transformers. For example, hard-attention Transformers are less expressive than standard Transformers [Hahn, 2020, Hao et al., 2022]. These restrictions can be relaxed in part by introducing new modules within this framework. We discuss some extensions in Appendix A.4, but leave this direction to future work.

Are the programs interpretable? We have argued that our discrete Transformer Programs are more interpretable than standard Transformers, but one might question whether programs are actually easier to understand, especially as we learn larger programs for more complex tasks. On one hand, these programs are interpretable in the sense that we can trace how information flows between different positions and model components, and we can inspect the program using off-the-shelf code analysis tools, like debuggers. For more complex tasks, the program can be understood as a collection of individually interpretable feature functions, rather than a single interpretable algorithm (as illustrated in Figure 7). Nonetheless, the learned programs can still be complicated and non-intuitive. Another avenue for future work is to explore methods for automatically analyzing the resulting programs, and for imposing an inductive bias in favor of more interpretable programs.

Acknowledgments We thank the members of the Princeton NLP group for their helpful feedback and discussion. This research is funded by the National Science Foundation (IIS-2211779) and a Sloan Research Fellowship.

References

Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 39-48, 2016.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances. Computational Linguistics , 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 preprint arXiv:2104.07143 , 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-17442, 2020.

Andrew Cropper and Sebastijan Dumanˇ ci´ c. Inductive logic programming at 30: A new introduction. Journal of Artificial Intelligence Research , 74:765-850, 2022.

Jacob Devlin, Ming-Wei 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 Icard, and Noah D Goodman. Finding alignments between interpretable causal variables and distributed neural representations. arXiv preprint arXiv:2303.02536 , 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 arXiv:2203.14680 , 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 (TACL) , 8:156-171, 2020.

Yiding Hao, Dana Angluin, and Robert Frank. Formal language recognition by hard attention transformers: Perspectives from circuit complexity. Transactions of the Association of Computational Linguistics (TACL) , 10:800-810, 2022.

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: A taxonomy 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, 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-sequence 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: A method for stochastic optimization. arXiv preprint arXiv: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 interpretability. arXiv preprint arXiv:2301.05062 , 2023.

Thomas McGrath, Matthew Rahtz, Janos Kramar, Vladimir Mikulik, and Shane Legg. The Hydra effect: Emergent self-repair in language model computations. arXiv preprint arXiv:2307.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:1735917372, 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 logic 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: //github.com/chakki-works/seqeval . Software available from https://github.com/chakkiworks/seqeval.

Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In International Conference on Learning Representations (ICLR) , 2023.

Nostalgebraist. Interpreting GPT: The logit lens. https://www.alignmentforum.org/posts/ AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens , 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.11895 , 2022.

Bo Pang and Lillian Lee. A sentimental education: Sentiment analysis using subjectivity summarization based on minimum cuts. In Association for Computational Linguistics (ACL) , pages 271-278, 2004.

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-performance 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 arXiv:1904.01554 , 2019.

Jeffrey Pennington, Richard Socher, and Christopher D Manning. GloVe: Global vectors for word representation. In Empirical Methods in Natural Language Processing (EMNLP) , pages

1532-1543, 2014.

Felix Petersen, Christian Borgelt, Hilde Kuehne, 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 programmer-interpreters. 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: Languageindependent 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, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems (NeurIPS) , 30, 2017.

Ellen M. Voorhees and Dawn M. Tice. Building a question answering test collection. In Proceedings of the 23rd Annual 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. Interpretability 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-2591, 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 counterexamples. In International Conference on Machine Learning (ICML) , pages 5247-5256. PMLR, 2018.

Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like Transformers. In International Conference on Machine Learning (ICML) , pages 11080-11090. PMLR, 2021.

Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. In Association for Computational Linguistics and International Joint Conference on Natural Language Processing (ACL-IJCNLP) , pages 3770-3785, 2021.

Biao Zhang, Ivan Titov, and Rico Sennrich. Sparse attention with linear units. In Empirical Methods in Natural Language Processing (EMNLP) , pages 6507-6520, 2021.

Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. Advances in Neural Information Processing Systems (NIPS) , 28, 2015.

A Method details

A.1 Categorical attention

As described in Section 3.2, we implement categorical attention by associating each attention head with a boolean predicate matrix, W predicate ∈ { 0 , 1 } k × k , where k is the variable cardinality, with the constraint that each row W predicate ,i sums to one. The self-attention pattern is then defined by a score matrix S ∈ { 0 , 1 } N × N , with S = xW Q W predicate ( xW K ) ⊤ . 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 i = One -hot ( arg max j S i,j ) . During training, we define A i = Gumbel -Softmax( S i ) .

Defaulting to the beginning of sequence. We implement hard attention so that, in the event that there is no matching key for a query, the model defaults to attend to the first token in the sequence. Let S ∈ { 0 , 1 } N × N denote the score matrix for a sequence with length N . We define a modified score matrix ¯ S such that, for each row i ,

Breaking ties. Furthermore, we implement the attention mechanism so that, in the event that there is more than one matching key, the model attends to the closest match. Given the score matrix S ∈ { 0 , 1 } N × N , we define the modified score matrix ¯ S such that, for each row i , ¯ S i,j = S i,j × b i -j , where b i -j ∈ [0 , 1] is a bias associated with the offset between position i and j . For most experiments, we fix the bias to decrease from 1, when | i -j | = 1 , to 1 /N , when | i -j | = N , with b 0 = 1 /N to bias the query at position i against attending to itself. This is similar to types of relative positional bias that have been proposed for regular Transformer-based language models [Press et al., 2022].

A.2 Additional modules

Numerical attention. We implement limited support for numerical variables, designed to ensure that all variables within the program are integers within a bounded range, which allows us to discretize the program by enumerating all possible inputs. First, we include numerical attention heads, which read categorical variables as key and query, and numerical variables as value. The numerical attention head computes attention scores S ∈ { 0 , 1 } M × N using the select operator. Unlike in categorical attention, each query can attend to more than one key, and queries can attend to nothing if there is no matching key. Given attention scores S ∈ { 0 , 1 } M × N and value variable var , the output for the i th token is defined as ∑ N j =1 S i,j var [ j ] . At the input layer, there is a single numerical variable, ones which is frozen and equal to 1 for all positions; attention heads that read ones as the value variable are equivalent to the selector_width primitive in RASP. At higher layers, numerical attention heads can read the output of lower-layer attention heads as values. Figure 8 depicts a numerical attention head from a program for the Double Histogram task.

,

Numerical attention does not directly correspond to attention in a standard Transformer, because it computes a weighted sum rather than a weighted average. But a numerical attention head can be implemented in a standard Transformer by composing an attention head with a feed-forward layer, using the beginning-of-sequence token to allow the model to attend to nothing. This is how the selector_width primitive is implemented in Tracr [Lindner et al., 2023].

Feed-forward layers. In RASP, feed-forward layers are used to implement arbitrary elementwise operations, and they play an important role in many human-written RASP programs. In our Transformer Programs, we restrict the capacity of the feed-forward layers to ensure that they can be decompiled. Each feed-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 feed-forward layers to programs by enumerating all possible inputs and forming a lookup-table. For all experiments, we set ℓ = 2 . For categorical attention, given variable cardinality k , there are k 2 possible inputs. An example of a feed-forward layer is given in Figure 9, which depicts a circuit in a program for the Dyck-2 task.

For numerical attention, we can bound the input range by determining the cardinality of the numerical attention heads. For each numerical attention head, the minimum output value is equal to 0 and the maximum output value is equal to the maximum input length multiplied by the cardinality of the input variable (that is, the value variable). For example, if an attention head reads ones as the value

Figure 8: An example of a numerical attention head and MLP in a program for the Double Histogram task. For this task, the model must output, for each position, the number of unique tokens with the same histogram value. In this example, an attention head ( left ) calculates the histogram for each position. An MLP ( top right ) reads the histogram values and outputs a value of 0 if the histogram value is greater than one, and 4 otherwise. Inspecting the corresponding classifier weights ( bottom right ), we see that an output value of 0 -meaning a histogram count greater than 1 -increases the likelihood that the double-histogram value is 1 or 2 , and decreases the likelihood of larger values. Because the input length is limited to 8, this reflects the fact that if one number appears many times, it is unlikely that another number appears the same number of times. An output of 4 (meaning a histogram count of 1 ) increases the likelihood that the double-histogram is greater than 1 . Note that we configure all MLPs to read two input variables, but some MLPs learn to read the same variable for both inputs, as in this example. This allows us to compress the corresponding function.

variable, the maximum output value is equal to the maximum input length. If an attention head reads the output of a first-layer attention head as value, the maximum output value is equal to the square of the maximum input length.

A.3 Extracting programs

In this section, we provide more details about our procedure for converting our trained models into Python programs. While there are many possible ways to express the discretized model as Python code, we choose a mapping that facilitates analysis with a standard Python debugger. We use several simple strategies to improve the readability of the code, by annotating variable types, compressing statements, and removing unreachable branches. Illustrative programs are depicted in Figures 8 and 9.

Attention. Each attention head is represented by a predicate function, which takes as input a key and query and outputs a value in { 0 , 1 } . In Transformer Programs, all keys and queries are categorical variables with cardinality k , so the predicate function can be defined by enumerating the possible query values, which we do with a series of if statements. Additionally, if multiple query values are mapped to the same key, we condense the predicate by combining these queries into a single branch. This is illustrated in Figure 9. In the first attention head ( left ), each query position attends to the previous position (with the exception of the first token), so we cannot apply any compression. In the second attention head ( bottom right ), fifteen out of the sixteen possible query values are mapped to a single key value, so we combine them into a single branch.

MLPs. Weconvert feed-forward modules to functions by enumerating all possible inputs and forming a key/value lookup-table. For all experiments, we set each MLP to read ℓ = 2 input variables. In some cases, the MLP learns to read the same variable for both input variables, allowing us to reduce the length of the function. We further reduce the length of these functions by identifying the MLP output value associated with greatest number of keys, and returning this value as the default value without explicitly evaluating the corresponding condition. These two forms of compression are illustrated in Figure 8 ( top right ). This MLP reads a single numerical input variable and outputs a value of 4 if the input is 0 or 1 and a value of 0 otherwise.

Annotating variable types. In our Transformer model, all categorical variables are encoded using one-hot embeddings, which are represented as integers in the corresponding program. To improve readability, we replace these integers with symbolic values where appropriate. At the input layer, we represent the values of the tokens variable as strings rather than integer indices. At subsequent layers, we determine the appropriate type by following the computation graph. For example, in Figure 9, the tokens variable takes on values in { ( , ) , { , } } ; the first attention head reads tokens as the value variable, so we can automatically determine that attention_0_0_outputs variable takes on values of the same type; finally, the MLP reads tokens and attention_0_0_outputs as input variables, so we define the mlp_0_0 function in terms of the token value type.

A.4 Possible extensions

In our experiments, we used a limited set of model components that were designed to ensure that the resulting programs are interpretable. However, these do not cover all of the primitive components included in RASP. In this section, we discuss some of these differences; how these limitations can be addressed by introducing additional modules; and the relationship between (a) introducing more expressive modules and (b) the complexity of the corresponding program.

One-to-many attention predicates Our Transformer Programs use a form of binary categorical attention that associates each query value with a single key value. This is enforced by parameterizing the attention head with a predicate matrix W predicate ∈ { 0 , 1 } k × k (where k is the variable cardinality) and constraining W predicate so that each row contains only a single non-zero entry. This choice results

in more concise programs, as each attention head can be characterized in terms of a one-to-one mapping between key and query values, but it also makes it more difficult to express certain common attention patterns. For example, it is difficult to implement the 'less than' predicate, by which a query with a value of n attends to all keys with values < n . This predicate is used in a number of RASP solutions [Weiss et al., 2021] to implement a sorting routine. It is straightforward to allow the predicate matrix to have multiple non-zero entries per row. These weights could be learned by replacing the row-wise Gumbel softmax operation described in §3.2 with an entry-wise Gumbelsigmoid operation. In preliminary experiments, we trained models with this form of predicate matrix and found they obtained similar performance but were qualitatively more difficult to understand.

Score-based attention As in Tracr and RASP, we require that attention heads with categorical values always attend to exactly one position (hard attention). We enforce this constraint by deterministically breaking ties in favor of the position nearest to the query position. Weiss et al. [2021] discuss an extension to RASP ( select_best ) that provides for a more general form of one-hot attention, where each attention head is associated with both a binary function, predicate : Q×K→{ 0 , 1 } , and a continuous function, score : Q×K → R , that maps each combination of key and query to a real value. Given a query q and keys k 1 , . . . , k n , the attention between q and k i is defined as 1 if both predicate( q, k i ) = 1 and score( q, k i ) = max n j =1 score( q, k i ) , and 0 otherwise. This type of operation has a number of applications. For example, Yao et al. [2021] describe a construction for modeling sequences of balanced parentheses with an autoregressive Transformer, in which the model attends to the most recent opening bracket at the same nesting depth as the current position, by defining score( q, k i ) = i .

It would be relatively straightforward to extend Transformer Programs to include score-based attention modules, but this would also result in longer and more complex programs, because each query must now be defined in terms of an ordering of all |K| possible keys, instead of a single key. To see why this makes the program more difficult to understand, consider the information gained by observing that some token i attends to token j . In our default, one-to-one form of attention, this observation means that j is the closest position such that predicate( q i , k j ) = 1 . In score-based attention, the observation also implies that there is no position ℓ such that score( q i , k ℓ ) > score( q i , k j ) , conveying additional information about the rest of the sequence.

Learnable numerical variables In our experiments, we constrain the numerical variables to be either constant and equal to one (at the input layer), or the output of a numerical attention head. In RASP, numerical variables can take on any scalar value. Transformer Programs could be extended to include learnable numerical input variables, at the cost of some interpretability tradeoffs. For example, it may no longer be possible to decompile MLPs by enumerating all possible arguments.

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.com/princeton-nlp/ TransformerPrograms .

B.1 In-context learning

Data. For our in-context learning experiment (Section 4.1), we sample 20,000 sequences of length 10. The first token is a beginning-of-sequence token, and the remainder of the sequence is formed by randomly sampling a many-to-one mapping between letter types (either a , b , c , d ) and numbers ( 0 , 1 , 2 , 3 ), and then alternating letters and numbers until reaching the target length. The model is trained to predict the number following each letter, or a special unk token if the letter has not appeared earlier in the sequence. We use an autoregressive mask so that at each position i , the model can attend only to keys with positions j ≤ i .

Training. We train each model for 250 epochs with a batch size of 512, a learning rate of 0.05, and annealing the Gumbel temperature geometrically from 3.0 to 0.01, decreasing the temperature at each training step. We take one Gumbel sample per step. The model has two layers with one categorical attention head per layer, with a variable cardinality of 10. We train five models with different random initializations and pick the one with the lowest test loss. We implement all models in PyTorch [Paszke et al., 2019] and use the Adam optimizer [Kingma and Ba, 2014].

B.2 RASP tasks

Data. For each RASP task, we sample 20,000 inputs without replacement and partition them into train, validation, and test sets containing 16,000/2,000/2,000 instances respectively. With the exception of the Dyck tasks, we sample inputs by sampling tokens uniformly from the vocabulary. For the Dyck tasks, we follow Weiss et al. [2021] and sample with a bias towards strings with a valid prefix. Specifically, given maximum length of N , with probability 0.5, we sample N tokens uniformly from the vocabulary; otherwise, we sample a valid Dyck string s with length | s | ≤ N , and append N -| s | randomly sampled tokens to the end to obtain a string with length N . For all tasks, we obtain 20,000 samples without replacement by sampling strings uniformly until the set of unique strings has the intended size. We prepend all inputs with a beginning of sequence token bos . For the sort and reverse tasks, we also append an end-of-sequence token, eos . Following Weiss et al. [2021], we report the token-level accuracy.

Training. As above, we train the model for 250 epochs with a batch size of 512, and a learning rate of 0.05. We anneal the Gumbel temperature geometrically from 3.0 to 0.01, decreasing the temperature at each training step, and taking one Gumbel sample per step. These hyperparameters were chosen after initial experiments on the RASP tasks. We do a grid search for number of layers (2, 3), number of attention heads (4, 8), and number of MLPs per layer (2, 4). The attention heads are evenly divided between categorical and numerical heads. Simililarly, the MLPs are evenly divided between MLPs with two categorical inputs, and MLPs with two numerical inputs. We train models with five random initializations, pick the model with the best accuracy on a validation set, and report the accuracy on a held-out test set. For experiments with standard Transformers (described in Appendix C.1), we similarly do a grid search for number of layers (2, 3) and attention heads (4, 8) but otherwise follow the hyper-parameters described by Weiss et al. [2021]: the hidden dimension is 256, the learning rate is 0.0003, the batch size is 50, and we train for up to 100 epochs, picking the model checkpoint from the epoch with the highest validation accuracy. Each model takes between five and fifteen minutes to train on an Nvidia RTX 2080 GPU, depending on the number of layers.

B.3 Named entity recognition

Data. For our named entity recognition experiments (Section 4.3), we train on data from the CoNLL-2003 shared task [Sang and De Meulder, 2003], using the distribution from HuggingFace Datasets [Lhoest et al., 2021]. This data is pre-tokenized and we filter the dataset to sentences with a maximum length of 30 tokens and add special beginning- and end-of-sequence tokens. Following Collobert et al. [2011], we preprocess the data by replacing all contiguous sequences of numbers with a special number symbol, so, for example, the string '19.99' is replaced with '#.#'. We use the standard train/validation/test split and evaluate the results using a Python implementation of the standard CoNLL evaluation script [Nakayama, 2018].

Training. For both the standard Transformer and Transformer Programs, we use a batch size of 32 and perform a grid search over the number of layers (1, 2, or 3) and the number of attention heads (4, 8). For the standard Transformer, we search for model hidden dimension (64, 128, 256) and train for up to 100 epochs, taking a checkpoint after each epoch and picking the checkpoint with the highest performance on the validation set. For the Transformer Program, we fix the variable cardinality to 32 and search for the number of embedding variables (2, 4, 8) and training epochs (30, 50, 100). As above, we anneal the temperature geometrically from 3.0 to 0.01, and we report results after discretizing the model at the end of training. For both models, we initialize the word embeddings using 300-dimensional GloVe embeddings For both models, we train with three random seeds, pick the model with the highest F1 score on the validation set, and report the results on the test set.

C 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 a larger input vocabulary. Because we use one-hot encodings for the token and position embeddings, these values determine the cardinality, k , of the categorical variables used in the program. We compare Transformer Programs with standard Transformers, for each model performing a grid search over the number of layers and

Program

Standard

0

20

40

60

80

100

Reverse

Program

Standard

Hist

Program

Standard

2-Hist

Program

Standard

Sort

Program

Standard

Setting

|

|=8,

N

=8

|

|=8,

N

=16

|

|=16,

N

=16

attention heads and report the test accuracy of the model that performs best on the validation set. In Figure 10, we increase the vocabulary size and maximum sequence length from eight to sixteen. All other experiment details are the same as in Appendix B.2. Performance degrades on longer sequences, underscoring some of the optimization challenges in learning Transformer Programs for larger scales.

No numerical variables. In our main experiments, we train Transformer Programs with an equal number of categorical attention heads and numerical attention heads, and categorical MLPs and numerical MLPs. In Table 3, we compare results on RASP tasks with only categorical variables. The experiment setting is otherwise the same as in Appendix B.2, but all attention heads and MLPs are constrained to read only categorical variables. Not surprisingly, performance degrades on three tasks that are primarily based on counting: Histograms, Double Histograms, and Most Frequent. However, the Transformer Programs still achieve good performance on the other four tasks. These include the balanced parenthesis languages (Dyck-1 and Dyck-2), which are most naturally solved by keeping a count of unmatched parentheses at each position.

C.2 Analyzing the generated code

Here, we provide some additional analysis of the generated code, focusing on the RASP tasks. Complete, generated programs are available at https://github.com/princeton-nlp/ TransformerPrograms .

Code analysis tools. One advantage of representing models as Python programs is that we can use off-the-shelf code analysis tools to interpret them. Figure 11 illustrates one example, showing how we can use a Python debugger to understand a non-trivial circuit in a Sort program.

Program length. Table 4 shows the number of lines in our best-performing program for each RASP task, before and after applying the compression strategies described in Appendix A.3. The program length includes the basic library functions used to We use an automated Python code formatter, 4 which applies a number of standard style conventions and, in particular, enforces a maximum line length of 88 characters.

2sort.py

×

G

sort.py

×

G

programs >sort > sort.py > get_features

programs>sort> sort.py> get_features

113

114

112

def predicate_1_1(position,token):

#Read from the eos embedding at middle positions.

125

124

attn_1_1_outputs = aggregate(attn_1_1_pattern,attn_0_1_outputs)

attn_1_1_pattern = select_closest(tokens,positions, predicate_1_1)

115

if position in {0，2，3，4}:

sort.py ~/project/transformer-programs/programs/sort - Definitions (1)

116

return token == "
"

36

def predicate_0_1(token,position):

117

# Look for 0s from position 1.

37

#Store the 2nd token at the beginning of sequence embedding.

118

elif position in {1}:

38

if token in {"
"}：

119

return token == "o"

39

return position ==2

120

#Read from the bos embedding from later positions.

40

#Store the 6th token at the endof sequence embedding

121

elif position in {5，6，7}:

41

#(or
iflen<6).

122

return token == "
"

42

elif token in {"
"}:

123

43

return position ==6

124

attn_1_1_pattern = select_closest(tokens,positions,predicate_1_1)

44

#Store the 1st token at all other embeddings.

125

attn_1_1_outputs=aggregate(attn_1_1_pattern,attn_0_1_outputs)

45

elif token in{"0"，"1"，"2"，"3"，"4"}：

OUTPUT

DEBUG CONSOLE

TERMINAL

Filter (e.g. text, !exc...

46

return position == 1

tokens

47

48

attn_0_1_pattern = select_closest(positions,tokens,predicate_0_1

>['
',

attn_0_1_outputs

'2'

'4'

1'

1'

‘0'，'3'，'
']

49

，+]<

attn_1_1_outputs

2'，

'2'，

'2'，

'2'，

'3']

51

50

################################################

>['3'，'2′，

'3'

'3'，

'3'，

'4'

‘4'，

'4']

52

def predicate_0_2(position,token):

126

>

What information do the attention heads read? Because each attention head reads a fixed set of named variables, we can characterize how information flows through the programs by examining which variables are read by each head. In Figure 12, we summarize this information for RASP programs. At the first layer, the majority of categorical attention heads read positions as key and query variables and tokens as the value. At higher layers, positions remains the most common key variable, but the models are more likely to read the outputs of lower-layer attention heads as the value variable. Numerical attention heads are less likely to read positions and more likely to read tokens , attn , and mlp outputs. Both kinds of attention successfully learn to compose modules, with higher-layer modules reading the outputs of modules at lower layers.

C.3 Optimization challenges: case study

What are the challenges to scaling Transformer Programs to more complex problems? As discussed in Section 6, some of the main obstacles are optimization challenges. This is evident from the fact that, for many RASP tasks, we can write a Transformer Program that would achieve perfect accuracy, but our method fails to learn these solutions (as illustrated in Fig. 10). In this section, we conduct a case study to better understand these optimization challenges.

Setting We focus on the Reverse task. This task can be solved by the following program, adapted from Weiss et al. [2021], using a Transformer with two layers and one attention head per-layer:

tok

pos

attn mlp

0

1

2

Layer

0.21

0.79

0.00

0.00

0.04

0.54

0.12

0.29

0.04

0.58

0.17

0.21

Cat. head: query

tok

pos

attn mlp

0.17

0.83

0.00

0.00

0.04

0.83

0.12

0.00

0.29

0.54

0.08

0.08

Cat. head: key

tok

pos

attn mlp

0.83

0.17

0.00

0.00

0.42

0.00

0.38

0.21

0.25

0.00

0.58

0.17

Cat.

head: value

tok

pos

attn mlp

0.38

0.62

0.00

0.00

0.00

0.17

0.38

0.46

0.04

0.12

0.58

0.25

Num. head: query

tok

pos

attn mlp

0.62

0.38

0.00

0.00

0.38

0.25

0.12

0.25

0.29

0.00

0.12

0.58

Num. head: key

ones

attn

1.00

0.00

0.17

0.83

0.21

0.79

Num. head: value

0

.

25

0

.

50

0

.

75

1

.

00

0

.

00

output = aggregate(select(targets, positions, ==), tokens)

The first attention layer copies the position of the end-of-sequence token to determine the length of the sequence; the first-layer MLP calculates the difference between length and positions ; and the second attention layer uses the MLP output to read, at each position i , the token at position length -i . This solution works for sequences of all lengths, but Transformer Programs evidently fail to learn it.

Method To better understand why, we manually initialize some components of a two-layer Transformer Program model to encode this solution and train the model on the Reverse task with a vocabulary size and maximum sequence length of 16. Specifically, in the attention layers, we manually initialize: (1) the parameters associated with π K , π Q , and π V , which identify the key, query, and value variables to 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 in, which determines which variables the MLP reads from the residual stream. We do not initialize the other MLP weights or the classifier weights. We initialize different subsets of these parameters and observe whether the model learns the remaining components needed to complete the solution. We train each model for 100 epochs but otherwise use the same hyperparameters described in Appendix B.

Results The results of this experiment are in Table 5. When all of the components are randomly initialized, the model fails to learn the generalizing solution. When we manually initialize the attention layers and the MLP input weights, the model learns the remaining components needed to solve the problem successfully. In particular, it finds the internal MLP weights needed to calculate the difference between the two input variables ( length and position ) and output a one-hot encoding of the result. Interestingly, if we randomly initialize either even a single component in either one of the attention layers, the model fails to learn the correct solution; but if we randomly initialize only the MLP, the model learns successfully. These results suggest that the attention layers may be more difficult to optimize-for example, they may be more likely to get stuck in local minima.

C.4 Text classification

Next, we train Transformer Programs for three standard text classification datasets: classifying questions as one of six topics [TREC; Voorhees and Tice, 2000]; classifying movie reviews as positive or negative [MR; Pang and Lee, 2005]; classifying sentences as objective or subjective [Subj; Pang and Lee, 2004]; and classifying sentences from news articles as one of four topics [AG News; Zhang et al., 2015]. We filter the datasets to sentences with at most 64 words and use a vocabulary of the 10,000 most common words, replacing the remaining words with an unknown token, and fixing the variable cardinality at 64. As above, we compare the Transformer Program with a standard Transformer. For both models, we use the Transformer to extract token embeddings, obtain a sentence embedding by averaging the token embeddings, and train a linear classifier on the sentence embedding. We initialize both models with 300-dimensional GloVe embeddings and use grid search to select the model dimension, number of layers, and number of heads, and training hyper-parameters (learning rate and number of training epochs), as described in Appendix B.3. We hold out 10% of each training set to use as a validation set. For both models, we train with three random seeds, pick the model with the highest accuracy score on the validation set, and report the results on the standard test set.

Accuracy. Table 6 compares the accuracy for Transformer Programs, standard Transformers, and a bag-of-words baseline. The Transformer Programs are competitive with the standard Transformer on three out of four datasets, performing slightly worse on TREC. This could be because the standard Transformer can more effectively leverage pre-trained word embeddings, or because it is easier to regularize.

Interpretability. We illustrate the interpretability of these programs in Figure 13 by inspecting the features for an example from the TREC dataset.

>

2

w

1

ny

.e

1

.n

5

le

d

0

ak

ch

？

9

3

6

12

15

2

6

7

1

12

2

6

2

6

15

2

12

14

RCAMCM

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\n_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.\n\nSTRENGTHS:\nOriginality\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.\n\nWEAKNESSES:\n__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\nQUESTIONS:\n* 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\nLIMITATIONS:\nThe authors have addressed the limitations of the method adequately to the best of my knowledge.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThere 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\n1. 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\nQUESTIONS:\nNo questions beyond the one in the weaknesses/limitations.\n\nLIMITATIONS:\nAre 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nI 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\nWEAKNESSES:\nWhile 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\nQUESTIONS:\n- 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\nLIMITATIONS:\nSee weaknesses above.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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).\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nLimitations 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?", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}