ALTA: Compiler-Based Analysis of Transformers
Authors: Peter Shaw
James Cohan
Jacob Eisenstein
Kenton Lee
Jonathan Berant
Kristina Toutanova
Venue: Accepted by TMLR ()

Abstract: We propose a new programming language called ALTA and a compiler that can map ALTA programs to Transformer weights. ALTA is inspired by RASP, a language proposed by Weiss et al. (2021), and Tracr (Lindner et al., 2023), a compiler from RASP programs to Transformer weights. ALTA complements and extends this prior work, offering the ability to express loops and to compile programs to Universal Transformers, among other advantages. ALTA allows us to constructively show how Transformers can represent length-invariant algorithms for computing parity and addition, as well as a solution to the SCAN benchmark of compositional generalization tasks, without requiring intermediate scratchpad decoding steps. We also propose tools to analyze cases where the expressibility of an algorithm is established, but end-to-end training on a given training set fails to induce behavior consistent with the desired algorithm. To this end, we explore training from ALTA execution traces as a more fine-grained supervision signal. This enables additional experiments and theoretical analyses relating the learnability of various algorithms to data availability and modeling decisions, such as positional encodings. We make the ALTA framework --- language specification, symbolic interpreter, and weight compiler --- available to the community to enable further applications and insights.

Published in Transactions on Machine Learning Research (2/2025)
ALTA: Compiler-Based Analysis of Transformers
Peter Shaw?, James Cohan", Jacob Eisenstein', Kenton Lee', Jonathan Berant", Kristina Toutanova'
"Google DeepMind, Google
Reviewed on OpenReview: htips: // openreview. net/forumt id- d7-51ll99i8
Abstract
We propose new programming language called ALTA and compiler that can map ALTA programs to Transformer weights. ALTA is inspired by RASP, a language proposed by Weiss et al. (2021), and Tracr (Lindner et al.) 2023). a compiler from RASP programs to Translormer weights. ALTA complements and exlends this prior work, offering the ability to express loops and to compile programs to Universal Transformers, among other advantages. ALTA allows us constructively show how Transformers can represent length- length-invariant algorithms for computing parity and addition, well as solution to the SCAN benchmark of compositional generalization tasks, without requiring intermediate scratchpad decoding steps. We also propose tools to analyze cases where the expressibility of an algorithm is es- tablished, but end-to-end training on a given training set fails to induce behavior consistent with the desired algorithm. To this end, we explore training from ALTA execution traces as a more fine-grained supervision signal. This enables additional experiments and thcoretical analyses relating the learnability of various algorithms to data availability and modeling decisions, such positional encodings. We make the ALTA framework - language speci- fication, symbolic interpreter, and weight compiler - available to the community to enable further applications and insights
1 Introduction
There has been significant discussion and debate about the degree to which Transformers can perform compositional generalization and "System 2" reasoning, prompted by negative results various evaluations for certain classes of Transformers (e.g.. Dziri el al. 2023 Qiu et al. 2023 Shaw et al. 2021 Wu et al., 2024 Delétang et al. 2022 Mitchell et al. 2023 Valmeekam et al. (023). Do such negative results reflect some mutable aspect of how such models were trained, or more fundamental architectural limitations? To better understand the conceptual limitations of Transformers, would be useful to have an interpretable framework for understanding whether and how Trannformers can represent and learn solutions various tasks of interest. Such framework could also potentially help elucidate path towards improving these capabilities.
We present a new framework for compiling interpretable, symbolic programs to Transformer model weights. The framework is based on a new programming language called ALTA, A Language for Transformer Analysis. It includes an interpreter for symbolically executing ALTA programs, and a compiler for converting ALTA programs Transformer model weights. ALTA inspired by prior work that introduced a programming language for Transformers called RASP (Weiss et al., (2021), and prior work that built compiler from RASP programs model weights called Tracr (Lindner et al.) 2023). ALTA complements and extends this prior work, with two key conceptual differences.
First, ALTA supports dynamic control flow operations such as loops. While Zhou et al. 2023b showed how RASP programs can be executed within the context of an auto-regressive decoder to implement some forms
Code available

---

Published in Transactions on Machine Learning Research (2/2025)
ALTA Program
>_
ALTA Compiler
Transformer Weights
Figure 1: Overview of ALTA. We propose new programiiig language called ALTA, and a "compiler" that can map ALTA programs to Transformer weights. ALTA is inspired by RASP, language proposed by Weiss et al. (2021) and Tracr (Lindner et al. 2023), a compiler from RASP programs to Transformer weights. ALTA complements and extends this prior work, offering the ability to express loops and to compile programs to Universal Transformers, among other advantages.
of loops by leveraging scratchpads (Nye et al., 2021 Wei et al, 2022). ALTA can implicitly support such operations without relying on intermediate decoding steps. This is useful to study because such additional decoding sleps can be computationally inefficient and are non-differentiable, typically necessitaling addi- tional supervision. ALTA accomplishes this by compiling to Transformers with layer-wise weight sharing. From one perspective, Transformers with weight sharing are simply a special case of standard Transformers, However, they have also been shown to have an inductive bias that improves performance on composi- tional tasks (Csordás et al.) 2002 Ontanon et al. 2022 Yang al. 2024), which warrants further study. ALTA also supports a conditional computation mechanism, enabling compilation of programs to Universal Transformers (Dehghani et. al. 2019), thereby enabling new constructive expressivily results for this class of models.
Second, ALTA represents the computation of the MLP sub-layer as sparse set of transition rules. We show that this enables compilation of complex programs to reasonably sized Transformers. In contrast, Tracr compiles functions over multiple variables expressed in RASP to dense lookup tables encoded in the MLP parameters, which can suffer from combinatorial explosion in the number of possible variable combinations. The ALTA compiler can leverage the sparsity expressed the set of transition rules to reduce the number of MLP hidden dimensions required, in some cases by many orders of magnitude. Additionally, representing the MLP sub-layer computation as set of sparse transition rules supports new insights into the generalization potential of MLP layers, and therefore of Transformers, which we explore both theoretically and empirically,
We highlight two primary applications this framework. First, we show new constructive expressivity results for Transformers and Universal Transformers. This includes showing how Transformers can implement fongth-nvarrian algorithms for computing parity and addition. We also demonstrate shift-reduce parsing algorithm that solves the SCAN (Lake & Baroni 2018) benchmark of compositional generalization tasks. Second, provide tools to analyze cases where the expressibility of algorithm is established, but end-to- end training on given training set fails to induce behavior consistent with the desired algorithm. Specifically, we propose to use intermediate supervision from ALTA execution traces over a given training set as a learning signal. In some cases, we show this additional supervision is sufficient to learn the dcsired algorithm, but in other cases failures can highlight limitations of the underlying architecture due to, e.g., the type of positional encoding used. To complement this empirical assessment, we also introduce the analytical notion of whether a program is minimal with respect to a training set, based on whether certain components of a program could be removed without affecting the training set predictions. We prove that if program is not minimal with respect to training set, then the compiled model will contain parameters that can be freely changed without affecting any predictions on the training set, i.c. some parameters are under-specified by the training set, even with intermediate supervision. We demonstrate cases where this analysis predicts that test set performance would be under-specified by a given training set, and show agreement with empirical results from training with intermediate supervision. We hope these tools can help provide insights to bridge the gap between expressibility and lcarnability of algorithms Transformers.
2

---

Published in Transactions on Machine Learning Research (2/2025)
In summary, our main contribations are follows:
We propose a new programming language called ALTA, and a compiler that can map ALTA programs to Transformer weights. ALTA supports loops and the ability to compile programs to Universal Transformers, among other advantages. (Section
We use ALTA to demonstrate new constructive expressivity results for Universal Transformers, showing how they can express length invariant algorithms for computing parity and addition, well solution to SCAN. (Sections 3
We show how the ALTA framework can be used to analyze the learnability various algorithuss with respect to a training set, both theoretically and empirically via trace supervision. (Sections
We make the ALTA framework available to the community to support further applications and insights.
2
Proposed Framework
Here we give an overview of how ALTA programs are specified, their computational model, and how they can compiled to Transformers. More details on the ALTA program API is Appendix A.1 and compilation details and examples are in Appendix A.2
2.1 Overview
We give an example of an ALTA program in Figure 2 components: a set of variables, a set of attention heads, and "MLP function". We explain each of these below in the context. of how they affect the execution of an ALTA program. The computational model of ALTA program aligns closely with the computational model of Transformer (Vaswani et al. 2017). In this paper we focus on encoder-only Transformers for simplicity. However, we note that ALTA programs can alternatively be executed in the context of a decoder-only Transformer. This involves adding a causal attention mask and outer auto-regressive decoding loop, but does not otherwise affect the definition and compilation of ALTA programs (see Appendix A..).. We also focus on Transformers with layer-wise weight sharing, i.c. where all attention and MLP parameters are shared across layers. We also support Universal Transformers which have input-dependent number of layers.
An ALTA program specification includes three key
Notably, not all types of computation expressible by Transformers can be represented in ALTA, i.e. the range of the ALTA compiler is a relatively small subspace of all possible parameter values. For example, ALTA has limited support for numeric computation and does not support modeling of probabilistic output distributions. However, ALTA provides broad support for implementing various types of deterministic algorithms.
The ALTA framework includes an interpreter, which symbolically executes a program, and a compiler which compiles programs Transformer weights. The input an ALTA program / E P sequence of integers inputs (E X') and the output is sequence of integers of equal length (E )). The interpreter implements function I P X + D. The interpreter is useful for development and understanding the compulational model in an abstract way. The compiler implements a function C' such that 0 = C(P) where T(x, 0) s I(P,x) for all x E X' and where T denotes the output of a Transformer encoder. The equality holds up to the limits of mumerical approximation for well formed programs.
2.2 Variables
Similarly to Lindner et al. (2023). we adopt the residual stream view of Transformers as proposed by Elhage ett al. (2021). In this view, the attention and MLP sub-layers within the Transformer read and write to the residual stream, which represented by the activations between these sub-layers. While Transformers represent the residual stream for each element as a vector, our interpreter represents the residual stream for each element symbolically mapping of variables to values. The residual stream of the interpreter for the
3

---

Published in Transactions on Machine Learning Research (2/2025)
vars =
# Initialize parity with input.
"parity": var (range-2,, input_init_ti snit_fn--lmberr x: x),
# Whether parity has been updated.
"done": var var (range=2, tot fn-laddd x: 0)
Position of current element.
"idx": var (reagge-NUM_PP
position, init_ x: x)
Index of preceding element.
"idx_left", var (range-MUM_PPS, position, init fn=lambda max (0, 1).
attention, heads -
# Valuea of 'parity' and 'done' for preceding element. "parity_lefte ("idx_leftt) "idx", "parity") (done_left?y okv ("idx_left", "idx", "done*)
def fin_fn(z):
if not 2 =["done"] and "]]
# Update parity based on parity of preceding elenent.
z[ [['parity"] = =l"parity_left") =["parity"]
2[ "onee
return progran_spaa (vars=vars, heads=attention seattteentio_hhhdd, ifn fn-ffn_fff output= parity" halt__pec-halt_ee apec ("done", 1).
input_range=2, position.range:NUX.
Example ALTA Program. The parily program computes whether given binary sequence contains an even or odd number of "1" tokens. For an input of length N, the parity variable of the final input element will equal the parity of the overall sequence after N layers, and computation will halt. The program specification contains all the necessary information compile the program to a Transformer.
Figure 2:
parity program of Figure is shown in Figure 3 The set of variables and their possible values are specified by the program. There are three kinds of variables ALTA: categorical variables have bounded integer values, numerical variables have real-valued values, and set variables have sets of bounded integer values. Variables representing the output of attention heads can also take null undefined value (see 2.3 We establish bijective mapping between variable assignments and activation vectors. Each possible value of categorical variable is assigned standard basis vector in the activation space, i.e., a one-hot encoding. Set variables are similarly represented, but with a multi-hot encoding. The scalar value of numerical variable is directly represented in a single dimension. This mapping can be seen as establishing an approximate isomorphism with respect to the sub-layer operations of the interpreter and those of compiled Transformer.
2.3 Execution
In this section, we explain the execution of ALTA program in the interpreter, and summarize how each operation is encoded in compiled Transformer, with more details in Appendix [A] We denote the value of variable foo for element at sub-layer & as (i,too) Let a denote the vector of values for foo across all elements at sub-layer k.
&
(s,to))
Initialization Given an input sequence (I1, "2, " xx)) we initialize every variable at every position.
Specifically, for variable foo, we initialize (1.fo)) as function of I,, a function of the positional index i, or as a constant, based on how the initialization for foo specified the program. This operation is encoded the parameters of the Transformer's embedding tables for input and position values. The number of possible input values and possible positional indexes must be specified the compiler. Alternatively, positional embeddings can be omitted if no variable is initialized a function of position.
Encoder Loop We then proceed itcratively exccute the self-attention sub-layer and the MLP sub- layer, which share parameters across all layers. A dynamic halting criterion, as proposed for the Univeral Transformer (Dehghani et al. 2019), optionally be specified by the program, which consists of specify-
4

---

Published in Transactions on Machine Learning Research (2/2025)
parity
parity_left left 2
done
MLP
MLP
done done_left
1dk
idx idx_1eff left
parity
parity_left
left
Attention
Attention
MLP
MLP
dome_leff left
idx idf_left
parity
parity_left
done
MLP
MLP
foft_left
1dx
1dx left
Figure 3: Visualization of the interpreter's symbolic residual stream for the parity program shown in Figure for the input sequence [1, 0, I]. The computed output sequence after the second layer is 11,1,0] which corresponds to the parily of the inpul sequence up to the current position, with the final value containing the parity of the entire imput sequence. Output values are outlined in red, and values that have changed are highlighted in blue.
ing which variable and corresponding value indicate that computation has completed for a given element. Alternatively, a maximum number of layers can be specified as an argument to the interpreter or when running compiled Transformer. Similarly to Tracr (Lindner al. 2023),, our compiled Transformers do not include layer normalization operations, which simplifies compilation. Otherwise, the self-attention and MLP sub-layers align with those of standard Transformer, and are described below.
Self- Aulf-Attention For each attention head, the interpreter computes a selection matrix, containing "weight" for every pair of inputs, and uses this to aggregate information across positions. A simplifying assumption of ALTA, similarly to RASP, is that this matrix binary. Attention heads in ALTA are specified by pointers to query, key, value, and oulput variables. Each head must have a unique output variable. The query variable must be categorical set variable, the key variable must be categorical, and the value and output variables must both be either categorical numerical. For cach attention sub-layer k every attention head updates the value of some outpul variable:
A+1 ((oout))
2/ aggregate(selectt (uuuur) =(, ): 2t =(vval))))
where query, key, value, and out are the variable names specified by the given attention head. The definition select is similar to that used by RASP. The main difference that we do not allow specifying custom binary predicate as an argument to select, which simplifies compilation The select operation returns square selection matrix, St,j, where St [=((xay)) qujry if query is categorical and Si,
I (ekk)) E ak ((i.qe if query is set-valued. The aggregate operation returns new sequence of values z k+1 (oout)) Each value (iiout))
.+1 is determined by aggregating over the set of selected values, {= (= |Sa, = 1}
specified by the selection matrix row St.- When value is numeric, aggregate is defined the same as in RASP, outputting the average of the selected values, and will be undefined if no value selected. When value categorical, the output will be undefined if there is not exactly one value selected. The interpreter will raise an exception if any undefined variable is used as input to any operation in subsequent layers, as the encoding of an undefined variable is not well specified
compiled model.
While this may to restrict the expressivity of the query, key, and value projections in the Transformer, that the MIP sub-layer prior to the attention sub-layer can the query, key, and value variables arbitrary values. By using variables, it still possible specify arbitrary binary selection matrices.
5

---

Published in Transactions on Machine Learning Research (2/2025)
These operations can be encoded in the parameters of the query, key, value, and output projections for given attention head. A scalar hyperparameter controls the degree to which the softmax operation approxinates generating a binary selection matrix. Each attention sub-layer is followed by a residual connection. We also support the option of using relative position representations (Shaw et al. 2018 by specifying a mask of relative positions that each attention head can allend to, which is applied to the selection matrix following the select operation. This binary mask is compiled to relative position biases using the parameterization of Raffel et al. (2020).
MLP Function The MLP sub-layer implements a mapping from set of variable assignments lo new set of variable assignments, which is applied at every element. For compilation and analysis purposes, ALTA programs internally represent this operation as a set of transition rules, which can be interpreted as logical implications with conjunctive antecedents. For example, here is one of the transition rules for the example program in Figure
consequent
antecedent
1+1 -aprtty
k { dona) = O A &
dd = 1 A 2 (.parsty_ = 1 A ((-parsty) = 0
When the antecedent of a rule is satisfied by the MLP input (i.e., all of the conditions hold with respect to the variable assignments at the MLP input), then the consequent determines the value of some output variable for the next sub-layer. By construction, for a given output variable, there should never be more than one rule satisfied by the MLP input. If no rule is satisfied, then the value of that output variable is unchanged from the MLP input. We also include transition rules that ensure that every atlention oulput variable is set to a null value so that it can be updated by the next attention sub-layer without conflicting with the residual connection. No attention output variable can otherwise be the output variable of any transition rule.
The set of transition rules can be specified in two ways when defining ALTA programs. First, shown in Figure one can simply write a Python function with the signature shown. The set of transition rules can then be determined by executing this function for every possible set of variable assignments. In cases where this is not feasible, or where it is desirable to have more control over the set of transition rules, we offer an alternative AP1 for specifying the set of transition rules more directly (see $A.1). In either case, the representation of numerical variable values in the antecedent of a transition rule is based on the set of discrete buckets specified for the given variable. Leveraging the sparsity represented in set transition rules rather than compiling a lookup table consisting of all variable combinations can significantly reduce the number of MLP dimensions required? [
The set of transition rules is represented in the parameters of the MLP layers. We generate 4-layer MLP with clipped RcLU activations. The first 2 layers are only responsible for converting numerical and set variables into a one-hot representation, representing the possible values of these variables. For numerical variables, these correspond to a specified set of discrete buckets. Note that if program contains only categorical variables, these 2 layers could be omitted. The final 2 layers of the MLP are based on the set of transition rules. The parameters of these layers are compiled such that the hidden activations are a binary vector where each value corresponds to whether particular transition rule was satisfied by the ML. input. Each row of the first matrix is a function of the antecedent of a particular rule, and each column of the second matrix is a function of the consequent of a particular rulc. Each MLP sub-layer is followed by a residual connection.
Output Fach ALTA program specifies an output variable, which must be categorical. If execution termi- nates after & sub-layers, and the output variable is output, then the program returns zt (otppuz) Selecting the subset of dimensions associated with the output variable is encoded in the parameters of the oulput projection. The Transformer then computes a softmax over this one-hot vector, and outputs the argmax.
For example, cousider a string of length N represented by set of categorical variables, 21, 22, ... aN, each with K possible values. We want determine if in some vocabulary consisting of V strings. A naive lookup table approach requires KN hidden dimensions, but this function can be represented with only V transition rules.
6

---

Published in Transactions on Machine Learning Research (2/2025)
3 Expressibility and Learnability
While there are many potential applications for ALTA, we focus on two applications in this paper: new constructive expressivity demonstrations, and analysis of whether such algorithms are learnable given a particular training set, with varying amounts of supervision. We give an overview of these applications and our proposed analytical tools here, with results Section
3.1 Expressibility
There has been considerable interest in establishing the theoretical expressivity of various classes of Trans- formers (Pérez et al., 2021: Chiang et al.) 2023 Yun et al., 2020 Fong et al., 2023 Merrill & Sabharwal 2024) However, such theoretical works often do not demonstrate specific constructions of how Transformers can express algorithms of interest with limited resources, which can require manually specifying weight matrices. Tools like RASP, Tracr, and ALTA make it easier to construct such constructive demonstrations. For examn ple, RASP and its extensions have supported many new expressivity results for Transformers (Zhou et al.
2023b Angluin et al. 2023 Kazemnejad et al. 2024 Yang & Chiang 2024 Strobl et al. 2024 Friedman
et al. 2024).
Much of the recent work on the expressivity of Transformers has focused on Transformer decoders that can execute an input-dependent number of intermediate decoding steps before producing final output (Pérez
ct al. decoding steps have been empirically successful at improving sequential reasoning capabilities (Nye et al. 2021 Wei et al.) (022), this typically requires additional supervision during training (Pfau et al. 2024). Additionally, leveraging such intermediate decoding steps may not be the most computationally efficient way to represent certain algorithms. Therefore, it is interesting to study whether and how Transformers can represent various algorithms without relying on such intermediate decoding steps, as an alternative or complementary approach. We use ALTA to provide new constructive expressibility results for Universal Transformer encoders for various resource bounds, detailed in Section 4 When considering finite limit on the number of layers, such results also hold as special case of standard Transformer encoders.
2021 Feng et al., 2023 Merrill & Sabharwal, 2024 Zhou et al. 2023b). While such intermediate
3.2 Learnability
In many cases, we can establish that algorithm is expressible by given class of Transformers, but training model from this class on input and output examples of particular algorithm can fail to induce model that generalizes outside of the training set. It can be difficult to diagnose the reason for such failures, and to determine what to change regarding the architecture, training objective, or training data to improve generalization. We provide two tools to help bridge the gap between expressibility and learnability, which we discuss next.
Trace Supervision We propose to use intermediate supervision from ALTA execution traces over a given training set as learning signal. Given program P and set of model inputs X, we run P for every input in X, and extract the variable assignments at the input and output of every sub-layer, for every position. These traces can be used to derive trace supervision, which encourages the behavior of the model to align with that of the program being used to provide supervision. In our experiments, for simplicity, we focus on training the MLP parameters to reconstruct the desired output vector for each input vector, and compile the remaining parameters. In some cases, we show this additional supervision is sufficient to learn the desired algorithm, but in other cases failures can highlight limitations of the underlying architecture due to, e.g., the type of positional encoding used. We report results on the parity task in 4 and details of the training procedure in Appendix C.1
Theoretical Analysis To complement the empirical results from trace supervision, we also take a step towards analytically characterizing the conditions under which any particular ALTA program can be learned from a given training set. To this end, we introduce criteria for determining whether program minimal with respect to training set, which depends on whether certain components of a program could be removed without affecting the training set predictions.

---

Published in Transactions on Machine Learning Research (2/2025)
Here we give an overview of our theoretical analysis, which is detailed in Appendix B We focus the set of ML.P parameters, and consider a setting similar to that of training with trace supervision, where we derive a set of MLP inputs, D, corresponding to running given program over some set of model inputs. The MLP parameters are specified by the set of transition rules, R, in the given program. A rulc set is minimal with respect to a set of MLP inputs, D, if it is not possible to remove any rule or any constraint from any rule without changing the program output on some input in D. Similarly to training with trace supervison, we consider a reconstruction loss over D that quantifies how well the MLP outputs correspond to the outputs specified by R. Our results are with respect to training objective that combines this reconstruction loss with a regularizer on the parameter weights. We show that:
Theorem 1 (Informal). If the rule set R is not minimal with respect to D, then the compiled MLP parameters are not strict coordinate-wise local optimum of the regularized reconstruction loss.
Theorem 2 (Informal). If the rule set R is minimal with respect to D, then the compiled MLP parameters are strict coordinale-wist local oplimum of the regularized reconstruction loss.
Therefore, when the minimality conditions are met, we can conclude that there exists a local optimum point with respect to the training objective that corresponds to the behavior of the given program, although we are not guaranteed to find this point during training. See Theorems I and in Appendix B for the formal statements and proofs.
We can apply this theory to analyze a particular program, training set, and test set. We can determine a minimal version of the program with respect to the training set by removing aspects of the program that do not change any predictions then assess whether this minimal version generalizes, i.e,, has the same behavior as the original program a Lest set. This analytical (est does not require training models or selecting hyperparamelers, and provides an interpretable assessment of the potentially underspecified aspects of program. Our theoretical analysis suggests that a Transformer trained with trace supervision from a program on a training set is more likely to generalize if the minimal version of the program generalizes, and we evaluate this in Section
the training set (Appendix C.2 describes the exact procedure used). We
Notably, this notion of minimality can also help generalize some of the ideas introduced by Zhou et al.) (2023b) with respect to RASP. As Zhou et al. (2023b) were interested in understanding length generalization, they proposed some restrictions on RASP programs that would otherwise be "difficult to learn". First, they proposed restrictions on programs containing certain operations over positional indices. Second, they excluded certain programs from consideration on an intuitive basis, such as a program for solving the parity task using sum and modulo operations. In both cases, such programs would not be minimal with respect to some finite length training set, according to our proposed criteria, as we show in Section
4
Experiments and Analysis
We detail experiments and analysis on several tasks, with further details and results in Appendix D
4.1
Parity
The parity task requires the model to compute whether a binary sequence contains an even or odd number of ones. The ability of transformers to learn parity has been studied extensively (Hahn 2020), particularly the degree to which they exhibit, length generalization (Bhattamishra et al. 2020 Chiang & Cholak 2022;
Ruoss et al. 2023 Delétang ct al. 2022). Empirically successful solutions have relied on scratchpads (Anil et al. 2022 Zhou al. 2022). Zhoo al. (023b) used a variant of RASP to investigate why Transformers struggle with length generalization parity and why scratchpads help.
We study three ALTA programs for computing parity detailed in Appendix D.1 First, the Sequential (Absolule) program computes parily by iterating through each position (one per layer), flipping a parily bit every time a one is encountered. While this program uses absolute positions to enable propagating information to neighboring tokens, we also consider a Sequential (Relative) version that uses relative positions instead. Finally, the Sun + Modulo program computes parity in a single layer, using an attention head to compute the total number of ones, and then computing a mod 2 operation in the following ML.P sub-layer.
8

---

Published in Transactions on Machine Learning Research (2/2025)
Length 20
Ones 20
1.0
Accuracy
Accuracy
Compiled (Original)
Compiled (Minimal)
0.5
0.5
Trace Supervision
0.0
Sequential (Relative)
Sequential (Absalute)
Sum Modulo
Sequential (Relative)
Sequential (Absolute)
Sum Modulo
Figure 4: Generalization by length and number of ones for different parity programs in various settings. The training set consists of inputs of up length 20 (and therefore up 20 "ones") while the test set consists of inputs with lengths 21 to 40. While the original compiled models have perfect accuracy, the minimal versions of the programs exhibit different generalization behaviors and closely mirror the behaviors of the trace supervision experiments.
Accuracy Length
Accuracy by Number of Ones
1.0
1.0
Accuracy
Accuracy
Standard Supervision
Sum + Modulo
0.5
Trace Supervision
Sequential (Relative)
Trace Supervision
0.0
0.0
Length 20
Length
Ones s 20
Ones 20
Figure 5: Accuracy by length and number of ones for Transformers trained with trace supervision vs. standard, end-to-end supervision. Transformers trained with standard supervision exhibit behavior similar those trained with trace supervision from the Sum Modulo program, ie., they exhibit no generalization to numbers of ones greater than those seen during training, but they do exhibit some length generalization.
For each program, Figure f compares the generalization behavior of Transformers compiled from the original version, compiled from the minimal version, and trained from trace supervision. The standard compiled models achieve 100% accuracy, exactly emulating the behavior of the symbolic program, as expected. The minimal versions of these programs exhibit different generalization behavior. The minimal Sequential (Abso- lute) program does not generalize to examples longer than those seen during training, because the embeddings for positions not seen during training not specified. The minimal Sum + Modulo program does not gen- eralize to examples with more ones than those seen during training, as it does not contain transition rules related to the numerical values corresponding to larger numbers of ones than those seen during training. However, it can handle examples longer than the training examples the numbers of oncs were seen during training. Only the minimal Sequential (Relalive) program generalizes supervision results exhibit the same generalization behavior as the minimal programs, as predicted in Sec- tion 3 The Sequential (Relative) program is also notable because it provides constructive demonstration of how a Universal Transformer encoder can express length-invarianl solution to parity with finite set of parameters, without relying intermediate decoding steps T
input lengths. Notably, the trace
We also evaluated several different Transformer variants using standard, end-to-end supervision, with results shown in Figure Theoretically, a transformer trained with weight sharing and at least as many layers as the longest. examples the train set could learn the minimal Sequential (Relative) algorithm and generalize to examples of length up to the number of layers. However, in practice, all variants exhibit behavior similar to that of the minimal Sum Modulo program, i.e., they exhibit some degree of length generalization but do not generalize to examples with more ones than were seen during training. See D.1 for additional results and experiment details.
WWill Chiang & Cholak (2022) previously demonstrated how Transformer express solution to parity for arbitrary lengths, their approach requires encoding activations and parameters with degree of numerical precision that scales with the maximum input length. The Sequertial (Relative) program does have this limitation, but does require more layers of computation. Zhou al, (2023b) also provide length-invariant construction, but their approach requires intermediate decoding steps. See Appendix
details.

---

Published in Transactions on Machine Learning Research (2/2025)
Simplicity Bias Prior work has attempted to understand the implicit inductive bias of Transformers in relation to a simplicity bias given some measure of complexity (Zhou et al. 2023b Abbe et al. 2023;
Bhattamishra et al. 2023 Tsoy & Konstantinoy 2024). For instance, inspired by the Minimum Description Length (MDL) principle Rissanen 1978 Grunwald 2004), Zhou et al. (2023b) hypothesized that Trans- the simplest RASP program that fits the training set, if one exists. Properties of the set of transition rules in an ALTA program can potentially provide new measures of Transformer complexity. For example, we can empirically compare the degree to which Transformers with layerwise weight sharing have an implicit bias towards learning "simpler" programs according to the number of transition rules. While the minimal Sequential (Relative) program contains 7 rules and the minimal Sum + Modulo program contains 31 rules, our results show that end-to-end model behavior is more consistent with the Sum + Modulo program. This indicates that Transformers, in this context, do not have an inherent simplicity bias that aligns with the number transition rules expressed in ALTA. In this particular context, the end-to-end behavior is more consistent with the algorithm that requires the fewest. layers to execule. The results for SCAN in Section 4.4 shosw a similar finding. This suggests potentially interesting direction to explore for future work.
formers are biased towards learning behavior corresponding
4.2 Addition
Another common benchmark for evaluating Transformer generalization is multi-digit addition. This task has been studied extensively (Nogueira. al. 2021
Liu et al. 2022), particularly with respect length generalization (Zhou et al. 2024 Shen et al. 2023 Zhou et al. 2023b Lee et al. 2024 Kazemnejad et al,
2024 Ruoss et al. 2023). While new positional encodings such as FIRE (Li et al. 2024) can improve
performance, Transformers still struggle with length generalization on this task, unless provided with carcfully
constructed supervision over intermediate decoding steps (Zhou et al. 2023b 2024
In Appendix D.2 we detail an ALTA program with dynamic halting that can add two positive integers of unbounded size. This program compiles to a Universal Transformer encoder with relative position repre- sentations (as mentioned in Section 2 we use the parameterization of Raffel et al. (2020)). The number of layers required to compute the sum is N + 2, where N is the number of digits in the larger of the two inputs. Notably, the minimal version of this program with respect to training set that includes only inputs with < 3 digits can generalize to unbounded input lengths.
We verified the correctness of our compiled model by considering samples of pairs of integer inputs with 1-50 digits. The compiled model achieved 100% accuracy across all instances, emulating the behavior of the symbolic program.
4.3 SUBLEQ
SUBLEQ is a single instruction language that has been shown to be to infinite memory (Mavaddat & Parhamil 1988). Giannou et al. (2023) previously showed how a Looped Transformer can implement an interpreter for a variant of SUBLEQ. In Appendix D.3 we demonstrate an
when given access
ALTA program for implementing an interpreter for less restrictive version of SUBLEQ in Universal Transformer encoder.
4.4 SCAN
The SCAN (Lake Baroni 2018) suite of compositional genoralization tasks requires mapping natural lan- guage commands (e-g., "jummp twice") to action sequences (e.g,, JUMP JUMP). Certain train and test splits
have been shown to be challenging for Transformer- based models (Keysers et al. 2020 Furrer et al. 2020
Qîu et al. 2022b Kazemnejad et al. 2024). Empirically successful solutions have involved symbolic decom-
positions of some form (Shaw et al.
2021 Chen et al. 2020
Herzig & Berant 2021 Qiu et al. 2022a Zhou
et al. 2023a
In Appendix D.4 we demonstrate an ALTA program that solves the SCAN task. The program achieves
100% accuracy across all examples in the SCAN dataset. We also verified that the compiled model emulates the behavior of the program, similarly achieving 100% accuracy. First, the program executes shift-reduce
10

---

Published in Transactions on Machine Learning Research (2/2025)
parse of the imput sequence, representing the parse tree. Second, the ALTA program decodes the output sequence by traversing the parse tree. The program represents the necessary variable-length data structures (a stack, parse tree, and buffer) using a variable number of input tokens. Compiled models require fewer than 2, 000 MLP hidden dimensions despite there being more than 10e0 possible variable combinations in the program. This highlights the importance of sparsity, which ALTA enables by representing the MLP computation as set of transition rules.
Notably, the minimal version of our program with respect to the training set generalizes to the test set, achieving 100% accuracy for all of the most challenging length- based and Maximum Compound Divergence (MCD) Keysers et al., 2020) splits. Our ALTA program for SCAN thus gives a constructive demonstration of how Transformers can represent algorithms exhibiting systematic generalization: a finite set of transition rules and attention operations can be recombined in novel ways to process novel inputs.
length-based
Because our ALTA program decomposes the task into many small steps, it can require up to 512 layers handle the longest SCAN examples. To undersland whether Transformers trained with standard, end-to-end supervision can also benefit from such large number of layers, we trained several Transformer variants, varying the total number of layers up to maximun of 256 encoder layers and 256 decoder layers. However, all generalize poorly on the test splits, and increasing the number of layers does not improve generalization. Consistent. with the parity results, end-to-end I training does not learn a sequential algorithm that generalizes well, despite demonstrating that such an algorithm expressible for a similar class of Transformers. See Appendix D.4 for additional results, experiment details, and discussion.
5 Discussion
Limitations ALTA has many of the same limitations as RASP and Tracr with respect to the potential differences between compiled models and those learned in practice, as discussed by Lindner et al. (2023). In particular, the framework provides limited support for numerical computations and modeling probabilistic output distributions. The properties of compiled models may not reflect those of models learned in practice. However, ALTA can still be a useful tool these limitations are kept mind when interpreting results.
Opportunities In this paper, we focused on analysis related to expressibility and learnability. However, there are many potential applications of the ALTA framework that could be interesting for future work. For example, we discuss more flexible alternatives for learning from trace supervision Appendix [C.]] ALTA could also potentially help develop test cases for interpretability tools. This was one of the primary motiva- tions for Tracr, which has been applied to help design interpretability benchmarks (Thurhherr & Scheurer 2024) and more interpretable Transformers (Friedman et al. 2023). Additionally, compiled components can potentially be integrated within learned models, such as circuits for arithmetic (Nanda et al., 2023) or induction heads (Akyürok et al., 2024). Finally, ALTA could potentially be extended to compare the relative expressivity of Transformers with self-attention compared with other architectures such as Linear Transformers (Katharopoulos et al., (020)) or State-Space Models (Gu et al., 2022). More broadly, we hope the ALTA framework can help the community better understand how Transformers can represent and learn various algorithms, and inspire new methods and insights.
Acknowledgments
We would like to thank Anian Ruoss, David Lindner, Adam Fisch, Chris Dyer, and the anonymous reviewers for helpful feedback and discussions.
References
Emmanuel Abbe, Samy Bengio, Aryo Lotfi, and Kevin Rizk. Generalization on the unseen, logic reasoning and degree curriculum. In International Conference on Machine Learning, pp. 31 60. PMLR, 2023.
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Arhitectures and algorithms. Xiv preprint ar arXiv:2401. 12973, 2024.
11

---

Published in Transactions on Machine Learning Research (2/2025)
Dana Angluin, David Chiang, and Andy Yang. Masked hard-attention transformers and boolean rasp recognize exactly the star-free languages. Xiv preprint arXiv:2310. 13897, 2023.
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. Advances in Neural Information Processing Systems, 11:35546-38555 2022.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint (XXt::2009.11266 2020.
Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions. In The 61st Annual Meeting Of The Association For Computational Linguistics, 2023.
Xinyun Chen, Chen Liang, Adams Wei Yu, Dawn Song, and Denny Zhou. Compositional generalization via neural-symbolic stack machines. Advances in Neural Information Processing Systems, 33:1690-1701, 2020.
David Chiang and Peter Cholak. Overcoming a theoretical limitation of self-atlention. arv preprinl orXiv:2200.12172, 2022.
David Chiang, Peter Cholak, and Anand Pillay. Tighter bounds on the expressivity of transformer encoders. In International Conference on Machine Learning, pp. 5544-5562. PMLR, 2023.
Noam Chomsky, Syntactic Structures. Mouton, Berlin, Germany, 1957.
Róbert Csordás, Kazuki Tric, and Jucrgen Schmidhuber. The devil is in the detail: Simple tricks improve systematic generalization of transformers. In Proceedings of the 2021 Conference Empirical Methods in Natural Language. Processing, pp. 619-619444 0-4,,, 2021.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal trans- formers. In International Conference on Learning Representations, 2019.
Grégoire Delétang, Anian Ruoss, Jordi Gran-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, et al. Neural networks and the Chomssy hierarchy, arXiv preprint XrXi::2207. 02099,, 2022.
P Kingma Diederik. Adam: A method for stochastic optimization. (No Title), 2014.
Nouha Dairi, Ximing Lu, Melanie Selar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, et al. Faith and fate: Limits of transformers compositionality. Advances in Newral Information Processing Systems, 36, 2023.
Javid Ebrahimi, Dhruv Gelda, and Wei Zhang. How can sclf-attention networks recognize dyck-n languages? 2020.
In Findings of the Associalion for Compulational Linguislies: EMNLP 2020, pp.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer cir- cuits. Thansformer Circuits Thread, 2021. fedexxhtmll
pwb/2021/ffamewrr//// ttttpp:://rrassforeeet httpa//cmoooooootcoooooo//////ar
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: A theoretical perspective. In Thirth-seventh Conference on Neural Information Processing Systems, 2023. URL https://ppenreviee htt.aanneeeeooooooddd/taaaoaa
Dan Friedman, Alexander Wettig, and Danqi Chen. Learning transformer programs. Advances in Neural Information Processing Systems, 36, 2023.
Dan Friedman, Abhishek Panigrahi, and Danqi Chen. Representing rule-based chatbots with transformers. arXiv preprint or Xio:24407. 10949, 2024.
12

---

Published in Transactions on Machine Learning Research (2/2025)
Daniel Furrer, Mare van Zee, Nathan Scales, and Nathanael Schärli. Compositional generalization in semantic parsing: Pre-training vs. specialized architectures. arXiv preprinl arXi::2007.089700, 2020.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopon- los. Looped transformers programmable computers. In International Conference Machine Learning, pp. 18898-114433 PMLR, 2023.
Peter Grunwald. A tutorial introduction to the minimum description length principle. arXiv preprint math/0406077 2004.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modcling long sequences with structured state spaces. In International Conference on Learning Represenlations, 2022. URL htttps///penneeiiww.nt
forum?id=uYLFozIv1..
Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the
Association for Computational Linguistics, 8:156-171, 2020.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint wxXb:1606.08844. 2016.
Jonathan Herzig and Jonathan Berant. Span-based semantic parsing for compositional generalization. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 908 921, 2021.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pp. 5156-5165. PMLR, 2020.
Amirhosscin Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural In Information Processing Systems, 36, 2024.
Daniel Keysers, Nathanael Schärli, Nathan Scales, Hylke Buisman, Daniel Furrer, Sergii Kashubin, Nikola Momchev, Danila Sinopalnikov, Lukasz Stafiniak, Tibor Tihon, Dmitry Tsarkov, Xiao Wang, Marc van Zee, and Olivier Bousquet. Measuring compositional generalization: A comprehensive method on realistic data. In International Conference on Learning Representations, 2020. URL https::/ net/ Syggelld=
Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 66 66-71, 2018.
Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sxquence-to-sequeece recurrent networks. Tn International conference on machine learning, pp. 2873-2882. PMLR, 2018
Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. In The Twelfth International Conference on Learning Representations, 2024. URL
Shanda Li, Chong You, Guru Guruganesh, Joshua Ainslic, Santiago Ontanon, Manzil Zaheor, Sumit Sanghai, Yiming Yang, Sanjiv Kumar, and Srinadh Bhojanapalli. Functional interpolation for relative positions improves long context transformers. In The Twelfth International Conference Learning Representations, 2024. URL
David Lindner, János Kramár, Sebastian Farquhar, Matthew Rahtz, Tom McGrath, and Vladimir Miku- lik. Tracr: Compiled transformers as a laboratory for interpretability. Advances in Neural Information Processing Systems, 36, 2023.
13

---

Published in Transactions on Machine Learning Research (2/2025)
Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35:34651 34663, 2022.
Farhad Mavaddat and Behrooz Parhami. Urisc: the ultimate reduced instruction set computer. International Journal of Electrical Engineering Education, 25(4):327 334, 1988.
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. In The Twelfth International Conference on Learning Representations, 2024.
Mclanic Mitchell, Alessandro B Palmarini, and Arseny Moskvichov. Comparing humans, gpt-4, and gpt-4v on abstraction and reasouing lasks. ar preprint wrXie:2311.09217 2023.
Richard Montague. Universal grammar. Theoria, 36(3):373-395, 1970.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. arXiv preprint ar (rX::23011.05211, 2023.
Michael Nielsen. A visual proof that neural nets can compute any function.
URL:
com/chap4. hlml, 2016.
Rodrigo Nogucira, Zhiying Jiang, and Jimmy Lin. Investigating the limitations of transformers with simple arithmetic lasks. arviv preprinl ar Xiv:22102. 13019, 2021.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXix:2112.001 00110, 2021.
Santiago Ontanon, Joshua Ainslie, Zachary Fisher, and Vaclav Cvicek. Making transformers solve composi- tional tasks. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 3591 3607, 2022.
Jorge Pérez, Pablo Barceló, and Javier Marinkovie. Attention is turing-complete. Journal of Machine Learning Research, 22(75):1-33 35, 2021.
Jackson Petty, Sjoerd Steenkiste, Ishita Dasgupta, Fei Sha, Dan Garrette, and Tal Linzen. The impact of depth on compositional generalization in transformer language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 7232-7245, 2024.
Jacob Pfau, William Merrill, and Samuel R. Bowman. Let's think dot by dot: Hidden computation in transformer language models. arXiv preprint 157558884,, 2024.
Linlu Qiu, Peter Shaw, Panupong Pasupat, Pawel Nowak, Tal Linzen, Fei Sha, and Kristina Toutanova. Improving compositional generalization with latent structure and data augmentation. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 4341 4362, 2022a.
Linlu Qiu, Peter Shaw, Pamupong Pasupat, Tianze Shi, Jonathan Herzig, Emily Pitler, Fei Sha, and Kristina Toutanova. Evaluating the impact of model scale for compositional generalization in semantic parsing. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 9157-9179, 2022b.
Linlu Qiu, Liwei Jiang, Ximing Lu, Melanie Sclar, Valentina Pyatkin, Chandra Bhagavatula, Bailin Wang, Yoon Kim, Yejin Choi, Nouha Dziri, al. Phenomenal yet puzzling: Testing inductive reasoning capa- bilities of language models with hypothesis refinement. arXie preprint arXiv:2810.08533 2023.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lec, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 (140):1 1-77 2020.
14

---

Published in Transactions on Machine Learning Research (2/2025)
Jorma Rissanen. Modeling by shortest data description. Automatica, 14(5):463 471, 1978.
David E Rumelhart, Geoffre E Hinton, and Ronald J Williams. Learning internal representations by error propagation, parallel distributed processing, explorations in the microstructure of cognition, ed. de rumelhart and j. meclelland. vol. 1986. Biometrika, 71(599-607);6, 1986.
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mchdi Bennani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pp. 1889-1903, 2023.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume (Short Papers), pp. 464-468, 2018.
Peter Shaw, Ming-Was Wei Chang, Panupong Pasupat, and Kristina Toutanova. Compositional generalization and natural language variation: Can semantic parsing approach handle both? In Proceedings of the 59th Annual Meeling of the Associalion for Compulalional Linguisties and the 11th International Joinl Conference on Natural Language Processing (Volume 1: Long Papers), pp. 922-938, 2021.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost.
International Conference on Machine Learning, pp. 4596 4604. PMLR, 2018.
Ruoqi Shen, Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, Yuanzhi Li, and Yi Zhang. Positional description matters for transformers arithmetic. arXiv preprint orXtt:2811.14787, 2023.
Lena Strobl, Dana Angluin, David Chiang, Jonathan Rawski, and Ashish Sabharwal. Transformers transducers. arXiv preprint Xii:2444.0204/ 2024.
Hannes Thurnherr and Jérémy Scheurer. Tracrbench: Generating interpretability testbeds with large lan- guage models. In ICML 2024 Workshop on Mechanistic Interpretability, 2024.
Paul Tseng. Convergence of a block coordinate descent method for nondifferentiable minimization. Journal of optimization theory and applications, 109:475-499 2001.
Nikita Tsoy and Nikola Konstantinov. Simplicity bias of two-layer networks beyond linearly separable data. arXiv preprint arXie:2405.1123 7299, 2024.
Karthik Valmeekam, Matthew Marquez, Sarath Sreedharan, and Subbarao Kambhampati. On the planning abilitics of large language models-a critical investigation. Advances in Neurol Information Processing Systems, 11775993-76003 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought, prompting elicits reasoning large language models. Advances in neural in information processing systems, 35:24824 24837, 2022.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference Machine Learning, pp. 11080 11090. PMLR, 2021.
Zhaofeng Wu, Linlu Qiu, Alexis Ross, Ekin Akyürek, Boyuan Chen, Bailin Wang, Najoung Kim, Jacob An- dreas, and Yoon Kim. Reasoning or reciting? exploring the capabilities and limitations of language models through counterfactual tasks. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 1819-1862, 2024.
15

---

Published in Transactions on Machine Learning Research (2/2025)
Andy Yang and David Chiang. Counting like transformers: Compiling temporal counting logic into softmax transformers. arXiv preprinl or Xrx:::204..00 04.4499, 2024.
Liu Yang, Kangwook Lee, Robert D Nowak, and Dimitris Papailiopoulos. Looped transformers are better at learning learning algorithms. In The Twelfth International Conference Learning Representations, 2024.
Shunyu Yao, Binghui Peng, Christos Papadimitrion, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. In Proceedings of the 59th Annual Meeting of the Association for Compulational Linguislics and the 11th International Joinl Conference Natural Language Processing (Volume 1: Long Papers), pp. 3770-3788 2021.
Chulhee Yun, Srinadh Bhojanapalli,, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transform- ers universal approximators sequencc-to-ssoquenc functions? In International Conference on Learning Representations, 2020. URL
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuurmans, Claire Cui, Olivier Bousquet, Quoe V Le, and Ed H. Chi. Least-to-most prompting enables complex reasoning large language models. The Fleventh International Conference on Learning Representations, 2023a. URL
Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. Teach- ing algorithmic reasoning via in-context learning. arXiv preprint arXtw:2211.09066, 2022.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. arXiv preprint Xrx::22300. 16028, 2023b.
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Transformers achieve length generalization but not robustly. orXiv preprint ar Xiv:24002 09371, 2024.
16

---

Published in Transactions on Machine Learning Research (2/2025)
A Framework Details
In this scction we provide additional details about the ALTA Framcwork, as introduced in 2 First, we detail the program API in A.1 which is referenced by the ALTA programs in this paper. Second, we provide more details and examples of how programs are compiled to Transformer weights in A..2
A.1 Program API Details
Here we detail the functions of the ALTA API used to build the programs shown in this paper. We refer the reader to our open-source implementation for further details.
Variables and Attention Heads The module contains several methods for defining variable specifica- tions:
var defines a categorical variable, the most common variable type. The cardinality must be specified.
numerical_ var
var defines numerical variable. set of discrete buckets must be specified, and values rounded to the closest bucket in the MLP sub-layer of compiled models.
set_var var defines a set variable. A set of sets of possible values must specified.
For each of these functions, it is also necessary to specify how the variable is initialized.
There are two methods for defining attention heads:
qkv defines attention head, with arguments that specify the query, key, and value variables. Optionally, a set of relative positions can also be passed, as well as an explicit specification for the output variable.
relative_ v is shorthand function for defining an attention head that attends a specific relative position. The query and key are implicitly set to single-valued categorical variable.
The output variable specification can be optionally specified explicitly, or is otherwise inferred from the type of the value variable. The output variable for each head is also included in the overall of program variables.
MLP Functions There are two ways specify the MLP function, mentioned in 2 For simplicity, the programs listed in this paper specify the MLP function as Python function, which is passed dictionary-like object for accessing and updating variable values. Alternatively, the set of transition rules can be specified directly, allowing more control to manage the scope of variables included in the antecedent of each rule, and avoid the combinatorial explosion of possible variable values for more complex programs. Figure 6 gives an example of specifying the transition rules for the parity program of Figure 2 using the MLPBuilder class. The class has two methods. The get method returns generator over possible variable values. The class automatically tracks which variables and variable values are in scope. The set method generates transition rule, generating the antecedent of the rule automatically based the current scope. In almost all cases, this results in a more compact set of transition rules than specifying the MLP function as a Python function. All of the results this paper related to analyzing the minimal versions of programs or computing the number of MLP hidden dimensions in compiled models are based on versions programs where the transition rule set has been specified directly.
17

---

Published in Transactions on Machine Learning Research (2/2025)
def cat_transition_run rules (variablen, sttention_heads).
x MLPBuilder (variables, attention_heads)
done get ("done"e))";:
if done 1:
for done_left in x.get (. d' if done_left =* 1:
set (.set("donee ("done", 1)
for parity_lefs left,, parity in x-get ("parity_left", "parity"): ("paritp",,,"", parity_left parity)
return x.rules
Figure 6: Function for directly specifying the transition rules for the parity program of Figure
A.2 Compiler Details
In this section we will give an example of the compilation process introduced in $2 using the parity program of Figure as an example.
Notation
Let of denote the activation vector in compiled model for element i at sub-layer
Encoding Variable Assignments Consider a set of variable assignments:
parity = 1
parity_leff left 0
idx = 2
For brevity, we only include a subset of program variables in this example. This set of assignments is represented as the following vector in a compiled model:
parity
parity =
parity_left = 0
parity_ _eftt = 1
idx 0
idx 1
idx 2
Initialization The imput embedding is computed as: 2f W + W where 2; is the input token ID at position i, and W and W represent the embodding matrices for token and positional embeddings, respectively. For brevity, we only consider up positional indices for this example. We also omit some variables in the matrices below, and transpose them for clarity of the row and column labels. Note that
18

---

Published in Transactions on Machine Learning Research (2/2025)
variables representing attention outputs, e.g. parity_left, are initialized to null value.
0
21
0
parity =
parity = 0
1
parity =
parity =
0
0
parity_lest _left =
parity_left saatt_left
0
parity_left left
parity_led
(WX)T=
(W!)T=
0
idx
1
idx 0
0
idx
0
idx 1
0
idx =
1
idx 2
Self-attention
The self-attention operation sums set of attention heads:
h+1
of
where of the output of head h parameterized by matrices wh. W WY, and W
of = W Σ Oaj
W"
chu
0gj
Exele ela
45 (W WW
For simplicity, we ignore the constant scalar term commonly applied to the dot product. Simlarly to Tracr (Lindner el al. 2023), we also adopt the parameterization of Elhage et al. (2021) for the alten- tion output, which can be shown to be equivalent to that of the original Transformer, but allows for a clearer exposition.
The parity program of Figure has two attention heads. Here describe the WP WK WV and wo matrices corresponding to the attention head with query idx idx_teft, key idx, value parity, and output aaittylefft (ommitted cells are zeros):
idx 0
X
0
idx
Y
0
idx
2
idx
(W9)T=
idx 2
(WK)T
0
X
idx
A
0
1dx Left = 0
0
0
0
0
A
0
idx left =
0
0
0
idx_left
left = 1
idx left = 2
0
idx_left = 2
0
parity =
parity =
0
1
parity
parity =
(WV)T=
0
0
parity_left
Wo
0
parity_left leiitylleft.
0
0
parity_left left
1
parity_left left
where À is a hyperparameter that controls the degree which the selection matrix approximates a binary- valued matrix, set to 100 by default.
19

---

Published in Transactions on Machine Learning Research (2/2025)
MLP Sub-layer Our approach to compiling MLP parameters takes loose inspiration from Nielsen 2016 which provides helpful visual exposition of how MLPs can encode arbitrary functions. The MLP function of the parity program can be represented with 7 transition rules, using the notation of Section 2
R1 k+1
( (pparity) = 1 == (=,don))
= A
(odone dddet = A 2 /.parityleffs = A < ((.parity) = 0
&
R2 (sparity)
X+1
=
=
z
(=,doze)
A _eft] A parity left)) = A 2 (-pariby)
R3 (,done)
done)
0 done left) =
k+1 R4 (-paritty_aas left)
{: par perfty_left)
Rg
1++ ( (apaityyleft)
( parrity_eeft) 1
k+1 Rg (done_leffy
2f ddddd_neeee _eft) 0
1+1 R7 ddddn_nertt
2 = doone_left) 1
Appendix B provides a more formal description of the MLP sub-layer and the compilation process. Here we provide an example of the compiled parameters W1, b2, and W for the parity program (the second bias term, b2, is always set to vector of zeros). Note that use clipped ReLU function the non-linearity As described in Section A there are also two initial MLP layers that are responsible for converting numerical and set variables into one-hot representations, but those are not necessary for the parity program we are considering which uses only categorical variables. We also omit the cells corresponding to idx and idx__t idx_left in the matrices and vector below (the corresponding cells are all zeros).
Each column of the transposed first matrix corresponds to rule, with 1 in cells corresponding to variable values in the andecedent of the rule. Each scalar in the bias vector has value 1 1 N, where N is the number of conjuncts in the antecedent of given rule. This ensures that the output of the following non-linearity is when all conjuncts are satisfied, and 0 otherwise. Each column of the second matrix is similarly associated with a rule, but the values are determined by the consequent of the rule. There is a 1 in the row corresponding to the new value of the output variable (or 0, if the output value is null), and a -1 corresponding to the value the output variable in the antecedent of the rule.
Rı
Ra
Ra
Ra
Rs
Rs
R7
1
0
0
0
0
0
0
parity
1
0
0
0
0
0
parity
0
0
0
1
0
0
0
parity_lett =
I
HII
1
0
0
1
0
0
parity_left left =
(WI)T
1
1
0
0
0
0
done =0
0
0
0
0
0
0
done =
0
0
0
0
1
0
done Line_lefttt
1
1
1
0
0
0
1
done left
R1
R2
Rs
R4
Rs
Rr
(b')T
|-3
-3
-1
0
0
0
0]
R
Ra
Ra
Ra
Rs
Rs
0
0
0
0
0
parity
1
1
0
0
0
0
0
parity =
0
0
0
1
0
0
0
parity parity_left left 0
O
0
0
0
-
1
0
0
parity_lest left = 1
W
1
0
-1
0
0
0001
00
0
done =
1101
0
done =
1
0
0
1
done done_left,
0
0
-1
done_left
20

---

Published in Transactions on Machine Learning Research (2/2025)
MLP Example Consider set of variable assignments:
parity 1
parity_ left le 1
done 0
done left
Let denote the vector encoding of these assignments
parity
parity
parity_left
parity_left
done = 0
done =1
done left done_left 0 done_left
Here the hidden activations in the MLP layer, which is binary vector corresponding to which rules satisfied by the input:
0 Ri
Ra
Rs
h(2) = oW>
R4
Rs
Ra
Rr
Note that is clipped ReLU non-linearity applied element-wise. Here is the output of the MLP sub-layer before and after the residual connection:
parity 0
parity
1
parity 1
0
parity 1
0
parity_left _lett = 0
0
parity_ left 0
parity_leff = 1
0
parity_ _eft = 1
w2h(2)
done 0
w2h((2) +
done =
done
done
done_left
done_left 0
left
done done_left 1
Which corresponds to the expected output assignments:
parity- 0
parity_ left Ø
done
done left
where the attention output variables are set null values they are ready be updated the next self-attention layer.
21

---

Published in Transactions on Machine Learning Research (2/2025)
A.3 Decoder-only Extensions
While this paper focuses using ALTA to compile programs for encoder-only Transformers, the framework can be extended to decoder-only models without significant modifications. Similarly, RASP was originally used to study encoder-only Transformers (Weiss et al. 2021), but was later extended to study decoder-only Transformers (Zhou et al., 2023b). The set of parameters necessary to specify an encoder-only and a decoder- only Transformer are the same, so the ALTA language specification and compiler do not require any changes. Only the interpreter and the Transformer implementation require minor modifications. First, it would be necessary to apply a causal attention mask. In the interpreter, this mask can be applied after determining the selection matrix. Second, rather than computing a single forward pass, the model should be run in the context of an auto-regressive decoding loop, i.c. where the previously decoded output is iteratively appended to the input sequence. For computational efficiency, caching can be used to avoid redundant computation for previously computed timesteps. We leave analysis of decoder-only Transformers to future work.
Theoretical Analysis of Minimal Rule Sets and MLP Parameters
If a program is minimal with respect to some training set, can we draw any conclusions about whether the transformer implementation of that program will be learned on that same training set? We now explore this question from the perspective of the MLP parameters in the trace supervision setting, Specifically, we identify conditions under which the compiled MLP parameters are strict coordinate-wise local optimum of regularized reconstruction loss. Extensions to consider all parameters of the compiled transformer, end-to-end training objectives, or non-coordinate-we wise local optima are left to future work.
B.1 MLP Specification
Here we introduce notation that differs from that in other sections, but is more appropriate for the goals of this section. Let us focus on two components of an ALTA program P for this analysis: a set of possible variable assignments, V, and set of rules, R. The rules implicitly define an MLP function, fR V + V, which is a sub-component of the overall program. For simplicity, we consider only programs with categorical variables for this analysis, and exclude transition rules related to attention outputs.
Variable Assignments An assignment V E V is tuple of Ny elements, with V (Vo, Vi, VNy-1); where Vi E 10, D; 1} and D; is the dimensionality of variable V. The number and dimensionality of variables is given by the program specification.
Transition Rules Let R = R2, RNa), where R2 is a transition rule. As described informally in Scction 2.3 transition rules consist of the following components:
The antecedent of RG, denoted An, consists of a set of Nn, constraints (k, where each k e refers to variable indez, and each (o..... .....,,, 1} refers to a variable value. A transition rule is satisfied by an assignment V V if and only if V2 for all (k.v) in the rule's antecedent.
The consequent of kn, includes an the oulpul variable index, denoted kR and a new value for this variable, denoted ''2, Per construction, for every R2 we require that the output variable appears in the antecedent of the rule. Let un. denote to the value associated with the output variable in the antecedent.
The rule R2 can therefore be interpreted as updating the value of the variable kn, from Un to v'R, if R, is satisfied.
Finally, we also require that for any assignment E V, no more than one rule satisfied for a given output variable. In other words, if two rules share the same output variable index, then they should never both be salisfied for any assignment E V.
22

---

Published in Transactions on Machine Learning Research (2/2025)
Logical MLP Function We can define the logical MLP function fa V V in terms of R. Let V' = =Vn(V). Then:
r'R.: Vi.
if 3R, E R. that is satisfied by V
V/2
otherwise.
B.2 Compiling MLP Parameters
As described in Section 2.3 and Appendix A.2 in our compiled models, we represent variable assignments as vectors, and compile the MLP function into the parameters of an MIP. Here we introduce notation and a more detailed description of the compilation procedure to support our theoretical analysis.
Assignment. Embeddings In compiled models, variable assignments are represented Nz- dimensional vectors. Let Z RNx represent the space of such vectors, which include the inputs and outputs of MLP layers.
We can define mapping Gy : V + 2. First, we define bijective function m(k, v) which for a given variable index E f0, 1,..., No 1) and corresponding variable value e E 10. D&), returns an index in {0, I,................ Nz 13. Let (zo, 21," ...
=N,-1) = -VV... Then:
Dk V v A v) =t].
(1)
other words, for a given assignment V, the vector ev(V) will have a 1 for dimensions associated with variable values in V, and 0's elscwhere.
MLP Function The compiled MLP function fo : Z - Z is defined by parameters E θ, where 0 = (W 6., W and W e NaxNz 6² E RNR, and W e NeNNN, The bias term in the second layer is set to all zeros in practice, and omitted here for simplicity.
For E 2, we can the define the MLP function:
9(0, 2) W 2 +b,
(2)
4(0,2) = ((((02))
(3)
f(0,z) 2²h((0, +
(4)
where (2) = max((0,iinL,,) min(1, clipped ReLU activation function applied deement-wise to the vector
For notational convenience, let f((0 2) = (f(0,2)), i.c denote the ith element of the output of f, and
similarly for g, h. Following equations 2 the values of these individual scalars are:
(:(0,2) Σ W Wl +
(5)
h₂(0, 2) (0, 2))
(6)
f((0, 2) = 2, + Σ (2,,)
(7)
As noted by Pêrcz et 2021 the elipped ReLU function be implemented terms of the standard ReLU function, (2)) m0 ma Therefore, there is equivalent formulation that additional layers but standard ReLU activations. However, we clipped ReLU for simplicity.
23

---

Published in Transactions on Machine Learning Research (2/2025)
MLP Parameters We can define the compiled parameters 1 E è such that ev(fa(V)) fo(ev(V)) for all E V. Let 0 = (W,, 01, W2).
W = [ (k, E An, s.t. s.t. m(k, =
(8)
8 =1 NR.
(9)
- 1, if m(kn, vn,)
W3,
1,
if m(RR, d'ny)) =
(10)
0,
otherwise
This construction ensures that for any assignment V E V, we have:
1, R is satisfied by V
0, NR, constraints of R2 are satisfied by V
(;(O)ev(V))
(11)
1, NR, constraints of R2 are satisfied by V
ev()) = [Ry is satisfied by V].
(12)
In other words, each element in the hidden layer activations corresponds to whether a particular rule was satisfied by the variable assignment represented in the input. Each column of W encodes an "update" the assignments related particular rule. This construction also ensures that:
e(0))v(V)) ev(fR(V)).
(13)
An example demonstrating equations 11
12
and
13 is given in Appendix A.2
B.3 Minimal Rule Set
In Section 3 we introduced the notion of a minimal program. Here we focus on the minimality conditions relating only to the rule set and the corresponding MLP layer. We define minimal rule set, encompassing the conditions necessary to obtain guarantees about a reconstruction loss landscape around the parameters of the compiled MLP.
Definition 1. (minimal rule set) Given a dataset of N variable assignments D = (V1, V2, Ve.... VN), set of rules R is a minimal rule set over D if:
It is not possible to remove any individual rule in R and not change the output of SR. for some ezample Vn D.
It is not possible to remove any individual constraint of the antecedent of any rule in R. without changing the output of fR for some eromple Vn D.
B.4 Reconstruction Loss
We consider a setting related to that of training models with trace supervision, discussed in $C.1 For program P and set of model inputs X, let D = (Vi, V2,..., VN) be the set of variable assignments at the input to the MLP for every position and layer when we run P over X. We use the shorthand:
z" = ev(V))
(14)
to denote the vector encoding of assignment Vn.
Now let us define reconstruction loss, LR, over individual predictions. For model parameters 0 and a given variable assignment V, with corresponding vector encoding ey(V,), the reconstruction loss quantifies how well the predicted variable encodings 20,,"") match the variable encodings specified by R, (v.))(V..)
24

---

Published in Transactions on Machine Learning Research (2/2025)
Per equation 13 this vector can alternatively be written in terms of compiled parameters I as f(0, 2"). For simplicity, we consider an In loss over the predicted and expected encodings, which simplifies the analysis of the local minimum point. We write LR as:
LR(0,z") = -(0,==")(((===)))
(15)
Note that for all
I=((0==))= O.
(16)
B.5 Coordinate wise Local Minimum
As it is difficult to analyze whether a point is a strict local minimum, we will instead analyze whether a point a coordinate-wise local minimum. Prior work has defined the notion of a coordinate- wise minimum, which is defined in terms of axis-aligned movements of parameter values (e.g., Tseng 2001).
Definition 2. (coordinale-wise local minimeum) For a multivuriate function f Xi X2 x ... x X, + R, the value X (Q1, .2...., Rx) is strict coordinate-une local minimum iff for each YE 0, such that [ 2,, ),)) <
8; .,.). for all (-),
We will be ovaluating loss function at a scmi-differentiabbl point (due the non-lincarity σ), so it is useful to note the following proposition:
Proposition 1. Consider a multivariate function and point X. If J a continuous function that is semi- cemi-differentioble at X with both left and right partial derivatives, and for every the right derivative with respect to X, evoluated al & is posilive, (x) > O, and the left dericalive with respect to X; evaluated al 8 ax. (X) 0, then X is strict coordinate-wise local minimam as defined in Definition
X is megalive,
This follows from the definition of left and right derivatives, and the requirement that f is continuous.
B.6 Main theorems
Suppose we learn 0 by optimizing a regularized sum of reconstruction losses over a dataset,
=() + Σ [n((Oeev(V)))
(17)
= =aI((0) + Σ Ln(0,²"),
with > penalizing the sum of L2 norms || + || + l1. We consider the L2 norm the regularization term for simplicity.
We want to show that when the weights are generated by compiling a set of rules that is minimal with respect to D, then a coordinato-wise local optimum of L for a 1.
Theorem 1 (Formal). If is the compilation of rule set R that is not minimal for D, then is not strict coordinate-wic wise local optimum of L(0,D,,),
Theorem 2 (Formal). If Ô is the compilation of rule set R that is minimal for D, then 0 is strict coordinale-wise local oplämum of L(0, D, for a <
Remark. Both theorems can also be proved for squared L2 regularizer with coefficient E (0, (2K)), where K = max(1, max; NN. 1) is the largest parameter magnitude in 0 and 2K is therefore the largest element of the gradient of the squared Ly regularizer evaluated at 0.
B.7 Proof of Theorem i
Proof. There two ways in which ruleset might violate minimality with respect to D.
25

---

Published in Transactions on Machine Learning Research (2/2025)
We can remove rule R, without affecting the output of fR any example. All rules affect the output when their antecedent conditions are met, so we can infer that the conditions of Rg are never met. In this case, there is guaranteed to be a parameter W of # 0, setting the value of f((0, = when the conditions of R3 are met. Tf Rj is never active for any z", then W2 does not affect the reconstrution loss LR-
There is condition of R3 that any example. This condition represented as nonzero element in W (and also b,, which is not necessary for the proof). By construction this parameter does not affect the reconstruction loss.
remove without affecting the output of fR on
In both cases, the gradient with respect to the nonzero parameter (W2, and w is set only by the regularizer. The right and left derivatives have the same sign, violating the conditions of strict coordinate- wise local optimality.
B.8 Proof of Theorem 2
To prove Theorem we will analyze the left and right derivatives of the loss function (eq. 17 with respect to each parameter in the compiled parameters 0, and show that the left derivative is positive and the right derivative is negative, satisfying the conditions of Proposition 1 The left and right derivatives of the loss function with respect to w, evaluated at 0, can be writlen as follows.
Φ= dw
T, D, = Ow (alw=) + a
Σ IT((0,z"))
(18)
Ow
n
0+ L(8, D, a) = dw
I+ (a|w() + g+
Σ LnR( 2m)
(19)
dw
Dw
We analyze the terms of these derivatives related to the regularizer and the reconstruction loss separately. First, consider the regularizer term.
(al (/)
a2
W
(20)
Om
a2 >>>>
0+ (a/w))
02
(21)
Dw
2 0
Next, let us consider the term related to the reconstruction loss. Let Oc be set of parameters equal to 0, but with e added to parameter w, e.g. 0 = (Wr BJ + 6 eg W22), where ej is a standard basis vector with a at position j and |e| < I. The left derivative can be written as follows,
θ- ELR(O
lim --0-
Σ L=(0 -ER00=((
(22)
Ow
= lim --0
(23)
using [nn per eq. 16 The right derivative be obtained similarly.
There are several ways we can show that the conditions of Proposition 1 are met for a given parameter at 8u ( D, a.) is positive.
i.e. that the left derivative 8_ an (, D, is negative and the right derivative
Lemma 1. If there exists some such that En Ln( xn)) 2 lel for all e (-d, d), then the conditions of Proposition are satisfied for w.
26

---

Published in Transactions on Machine Learning Research (2/2025)
Proof. The left and right derivatives of the reconstruction loss term are:
lim Σ LR(O =²)) - lim -lel c-0
(24)
c00
72
1 Σ LR( >²)² x0")))
lim
lim -le|
(25)
c+0
(26)
Because a 1, for any value of we have:
0- -11+
-
(27)
dw
0+
L( D, a)
21-a>0.
(28)
dw
Therefore the conditions of Proposition
1 are satisfied.
Lemma 2. If w 0, then the conditions of Proposition 1 are satisfied for w.
Proof. Because the reconstruction loss is zero 1 and non- negateettvve everywhere, its left and right derivatives are guaranteed to be s O and 2 0, respectively. Therefore:
0- L(,, (.)) 20-
(29)
Ou
θ- dw
2 0+a>0 2
(30)
Therefore the conditions of Proposition
1 are satisfied.
Lemma 3. If w > and there crists some d > 0 such that Σn La(0s 2n) 2 le| for all E (-d,0) then the
conditions of Proposition 7 salisfied for
Proof. This case therefore combines the reasoning of Lemmas (for the left. derivative) and 2 (for the right derivative).
θ- L -1+
(31)
Dw
Q+ (.),
(32)
du
Therefore the conditions of Proposition I are satisfied.
Remark. Compared to Lemma 1 this Lemma combines stronger condition the sign of with looser
condition the reconstruction loss, which in this case must grow with the perturbation only for 0.
We prove lemmas separately for each group of parameters, W 81, and W2 showing that the conditions of Lemma or are met for every parameter in
Lemma 4. Each W2, satisfies the conditions of Propostion
27

---

Published in Transactions on Machine Learning Research (2/2025)
Proof.
LR(0 :"")
Σ 10 (inz -V- (0 (l)(
(33)
Σ
2/ +
Σ +
(34)
2/
Σ
Wy.
=|e (0,00,2))
(35)
Lines 33 and 34 plug in the definitions of In (eq. 15 and Ji (eq. respectively, and line algebra.
35 follows from
Recall that 20 e {0,1} for all (),),, per eq. 12 For all where h,(0.2") 0, the value of En(00 0 for all i, but by minimality there must be some 2n for which h, 2n) 1 or else rule j could be removed. Therefore:
EiR(dig, x") - hg(0,2") 2 lel.
(36)
Therefore the conditions of Lomma are met.
Lemma 5. Each 0} satisfies the conditions of Propostion
Proof. We can derive the following by substituting the definitions of LR (eq. 15 and fa (eq.
7 hy (eq.
6
and gj (eq. 5]. and simplifying the resulting expression, for 1:
Ln(Oj 2")) Σ 15(0 z²) - f:(0, 20)
(37)
Σ W xex
+CE (0,
(38)
le| [o)(0,=") +
+
(0, -
Σ W
(39)
(40)
By construction of W22 for all there exists [W?,] 1, i.e. there is a non-zero entry in every column of W/2
(see eq. 10h. Therefore, E; [W2 2 1, and:
In =n) 2 x (0,(0,=") + (), By construction of 0, for all n, (,(,2²)) possible cases:
(41)
... (per eq. 11
Therefore can enumerate the
[e/, 0,(0==²) 00
LR(0ig,²')
le]. =
(42)
0, otherwise.
By the minimality criteria, for all j there exists some n where g,, 2n) = 0, where N - conditions of rule j are satisfied (see eq. 11]. If not, some condition could be removed, violating minimality. Also, for all j there exists some n where (,, 2²) 1, where all N conditions of rule j are satisfied. If not, the rule could be removed, violating minimality. Therefore:
Σ Ln(06 20"()),2") lc/.
(43)
Therefore, the conditions of Lemma
1 are met.
28

---

Published in Transactions on Machine Learning Research (2/2025)
Lemma 6. Each W),k satisfies the conditions of Propostion 1
Proof. The analysis is similar to that of öj. and we can reuse the result of eq. 41
LR(Oin W1 ,²²)
f=(0,=")]
(44)
Σ l2 W., xex (0,)) + (0,
(45)
|2gl x
W. x E x 19, 2n) TEE (0,
(46)
| x |e| x Ig,(0,z") teE (0, 1)]]
(47)
Recall that 2.k E 10 (eq. 8 and 2 E (0. 1} (eq. There are two cases consider.
Case 1: W 1. By minimality, for all k there exists some such that 2े 1 and g, (0,2"") 1, i.e. where rule R2 is satisfied by 2", and thus necessarily the constraint represented by W is satisfied by the input. However, minimality does not guarantee that there exists some such that 2े and g, 2") 0. Note that e < O, then |=[] x |c| [9, 2*) + E (0, 1)] = |e| 2 = and g,(0, =x) = 1, which the case for some n. Therefore:
+ x |e| x [os( +6E (0, 2 [t]. <<0 0
Σ T
InllO
(48)
0, 2 0
n
Since W 0, the conditions of Lemma are met.
Case 2: W1 0. The conditions of Lemma are met.
C Learnability Analysis Details
In this section, we provide additional discussion and details related to trace supervision ($ C.1 and the procedure for determining the minimal version of a program with respect to a training set
C.1 Procedure for Training with Trace Supervision
Given a program P and a set of model inputs X, we can run P for every input in X, and extract the variable assignments at the input and output of every sub-layer, for every position. In experiments, for simplicity, we focus on training the MLP parameters. After using the ALTA interpreter to collect pairs of variable assignments at the input and output of the MLP layers, we map these assignments to vectors, using the mapping described in 2 Finally, we can train a MLP layer based on these pairs of input and output vectors. Specifically, we use an L2 loss to encourage the MLP to produce the output vector given the input. We then use the ALTA compiler to compile parameters other than those used for the MLP layer. By combining the learned MLP parameters with those provided by the compiler, we have a full set of Transformer parameters. See DD.1 for the hyperparameters and training details for the parity task.
For future work, it may be possible to provide supervision from execution traces in way that makes fewer assumptions about how variables are encoded in the residual stream, i.c. by encouraging the residual stream to encode such values in a way that allows them to be accurately predicted from the residual stream with a linear classifier. This would enable training models with intermediate supervision without any compiled parameters, but is out of scope for this work. It would then be interesting to assess whether there is any potential for transfer learning from training with such supervision cases where a ground truth program is known to cases where ground truth program is unknown not feasible to express.
29

---

Published in Transactions on Machine Learning Research (2/2025)
C.2 Procedure for Determining Minimal Versions of Programs
Given a program P and of model inputs X, we can determine the minimal version of P with respect to X. For simplicity, we focus on the set of transition rules and the input embedding operations for this analysis. First, as motivated by our theoretical analysis in Appendix are never satisfied when running P on the inputs in X. While the formal definition of a minimal rule set also puts conditions on the constraints of satisfied rules, we assume that there are no unnecessary constraints to simplify our analysis. Second, we analyze the set of input IDs and positions seen when executing P over X. We restrict the variable initialization functions of the minimal program to output default values for variables outside of this minimal of observed token IDs and positions.
remove any transition rules which
Program Details and Additional Results
In this section provide program details and additional results. Some program listings omit the definitions of various constants and helper functions for brevily, and we refer the reader to the open-source code for the complete program specifications for all programs discussed in this paper.
D.1 Parity
Programs Here we provide the ALTA code for the parity programs that we study. The code for the Sequential (Relative) program in Figure g and the code for the Sum Modulo program is in Figure 7 The code for the Sequential (Absolute) program given in Figure 2
vars
"parity": var (range==))
"start": aumeric_var (input_ init_ _nt_tiaaaaa_aaaeeeeeee fn-lambda
float (x START),
values= (0, 1)),
"start_or or_eeer; var
var (range=2,
x: in (1, START))
"query": var (range=2, default=1),
attention heads
"x": (* query", *start_ or or one", "start",
ottput_ppccemueric..re Sappee/mannicccccaaaaaassseeee (values=BUCXETSS)
def (fn_fn(z)
aun_ones round (1 / 2["x"]) z 2["parity"] int (nun_an ones % !- 0)
return p
variables-vars, heads ion ion_attetttinondddtsesee ffn ffa_fn-ffa fa-ffn_ffffff output_name=
parity", pparaiiyyy="""
input_ranges poaition_rango-llon
Figure 7: Program for computing parity using sum and modulo operation.
Model Interface Each program expects the input to be binary sequence. The Sum Modulo and Sequential (Relative) programs additionally expect START token to be prepended to the sequence. For all programs, the final token in the output sequence contains the parity of the input sequences.
The compiled model for the Sequential (Relative) program has an MLP width (i.e., number of transition rules) of 7, and attention heads. The number transition rules for the minimal version of the Sequential (Absolute) program depends on the maximum input length in the training set, and the number of transition rules for the minimal version of the Sum + Modulo program depends on the maximum number of ones in the training set..
Model Sizes
30

---

Published in Transactions on Machine Learning Research (2/2025)
variables
"parity": var (range=2, aaput_init_dtte fn-la__f x: START elae x)
"done": var (range-2, input_init__falamn x: START),
attention heads : (_partts -1),
"done_loft"; "oolative""donee -1),
def ffn_fn(z):
if 2 [["done"] !- and z["done_ 1eft"] 1:
2 p["parity"] = al"parity_left") - 2["parity"] 2["done"]
return progran_spec(
variablee=variiallss heads=attention_ ffta_fn-ffa_th antput_name="partie input_range=3, position_range-llone. halt-halt_spet ("done", halt_value=1j
Figure Program for computing parity sequentially using relative positions.
Alternative Programs Thore is anothor algorithm for parity proposed by Chiang & Cholak (002)), which was inspired by algorithms for MLPs from Rumelhart et al. (1986). Similarly to the Sum + Modulo program, this algorithm first computes sum operation using an attention head, but they avoid explicitly computing the modulo operation by instead using attention heads that separately attend to even and odd elements. This enables the algorithm to be implemented with fixed set of parameters that is invariant to the maximum input length, aside from the positional encodings. However, we did not explore this algorithm because it is not possible to implement in way where the minimal ALTA program is invariant to the maximum input length considered. This is because ALTA requires discretization of numerical variables in order to perform numerical computations, a current limitation of the framework. Implementing the computations required would therefore require a number of transition rules that scales with the maximum input length considered, similarly how the Sum + Modulo program requires number of buckets that scales with the maximum number of ones considered. Regardless, Chiang & Cholak (2022) showed this algorithm is difficult to learn in practice, and it also requires specialized positional encodings. Furthermore, their algorithm requires encoding parameters and activations with degrec of numerical precision that scales with the maximum input length. In contrast, the Sequential (Relative) program compiles to a Universal Transformer that only needs to encode binary variables in the residual stream, and consists of only transition rules, i.e., requires only 7 hidden MLP dimensions. However, the construction of Chiang & Cholak (2022) requires only 2 layers, where as the Sequential (Relative) program requires N layers, where N is the maximum input length.
Train and Test Sets The train and test sets are the same for all experiments (including both trace and end-to-end supervision). The train set consists of examples between lengths 0 and 20, and the test set contains examples between lengths 0 and 40. The sets include roughly an equal number of examples per number of ones.
Trace Supervision Details Hyperparameters for training with trace supervision listed in Table All are standard hyperparameters for training neural networks except 'Noise Std Dev. We added small amount of Gaussian noise to the neural network input at training time to make it robust to numeric imprecision at inference time.
The Sequential (Absolute) and Sum + Modulo experiments used four hidden layers instead of the standard two, as we explored various hyporparameters to ensure the MIP had the capacity to fit the training set. Similarly, we used Adam (Diederik 2014) for both experiments instead of Adafactor (Shazeer & Stern 2018 it helped fit the training set..
31

---

Published in Transactions on Machine Learning Research (2/2025)
Table 1: Trace supervision hyperparameters.
Program
Hyperparameter
Sequential (Relative) Sequential (Absolute)
Sum + Modulo
Hidden Layers
4
Hidden Layer Size
4,096
4,096
Batch Size
256
256
256
Steps
50,000
50,000
400,000
Learning Rate
1e-2
1c-4
1e-4
Activation Fn
ReLU
ReLU
ReLU
Optimization Fn
Adafactor
Adam
Adam
Noise Std Dev
0.1
0.1
0.1
The Sum Modulo program used for trace supervision differs slightly from the program in Figure 7 num_ones
ones is stored as an intermediate categorical variable, as we found it casier for the MLP to learn to map a categorical variable to the correct parity output than a numeric variable.
End-to-end Training Details We trained transformers with various configurations using standard super- vision varying the number of layers, whether weight sharing is used, and the type of positional encoding, Constant in all standard supervision experiments are the following hyperparametens: embeddings with di- mension 512, hidden layer sizes of 2048, attention heads with dimension 64, an Adafactor optimization function, GeLU (Hendrycks & Gimpel 2016 activation functions, learning rate of 5e-4, 50,000 steps, and byte vocabulary.
Table 2: Length generalization accuracy for Transformers trained with intermediate supervision on parity.
Accuracy
Program
Length < 20
Length >> 20
Ones >> 20
Soquential (Relative)
100%
100%
100%
Sequential (Absolute)
100%
52%
51%
Sum Modulo
100%
78%
50%
Table 3: Longth gencralization accuracy for Transformers trained with standard supervision on parity.
Accuracy
Layers
Weight Sharing
Positional Encoding
Length s 20
Length > 20
Ones > 20
8
No
Relative
100%
60%
51%
Yes
Relative
100%
65%
52%
40
No
Relative
100%
55%
51%
40
Yes
Relative
100%
54%
48%
No
Absolute
100%
49%
50%
Yes
Absolute
100%
50%
52%
40
No
Absolute
100%
50%
51%
40 Yes
Absolute
100%
52%
51%
Results Table 2 presents the accuracy for each intermediate supervision experiment and Table presents the accuracy for each standard supervision configuration on different slices of the data. comppres the intermediate and standard supervision results, showing that standard supervision exhibits behavior similar to the Sum + Modulo program.
32

---

Published in Transactions on Machine Learning Research (2/2025)
Relative Positon Embeddingg
Absolute Position Embedddings
1.0
1.0
0.8
0.8
Layers
Accuracy
Accuracy
Layers
Wright Sharing
Layers
6 layers
Weight Sharing
0.2
0.2
0.0
0.0
10
15
20
25 Length
30
35
10
15
20
255
30
35
40
Length
Figure Accuracy by length for Transformers trained with standard supervision. There some length generalization with relative position embeddings (Ieft), but none with absolute (right).
Figures
9
10 and break down the standard supervision results in more detail.
Figure shows that there is no length generalization when using absolute positional embeddings, while with relative positional embeddings there is some length generalization which decreases as length increases.
Accuracy by Number of Ones
1.0
Layers
Accuracy
0.66
Layers Weight Sharing
Layers
Layers Weight Sharing
0.4
Number of Ones
Figure 10: Accuracy by number of ones for Transformers trained with standard supervision using relative position emboddings. There is no generalization beyond ones, which is the maximum number seen during training.
Figure shows that the length generalization we observe with relative positional embeddings due entirely longer examples sometimes containing the same number of ones as shorter examples in the training distribution, consistent with results in Anil et al. (2022). After 20 ones, accuracy starts oscillating between 0 and 1. The oscillations are due to the model predicting either 0 for all examples once examples contain more ones than were seen during training, which is shown in Figure
33

---

Published in Transactions on Machine Learning Research (2/2025)
Average Prediction by Number
Average Prediction
0.8
40 Layers
0.6
40 Layers Weight Sharing
8 Layers
--
Layers Weight Sharing
0.2
0.0
10 15
20
25
30
35
40
Number
Figure 11: Average prodiction by number of ones for Transformers trained with standard supervision using relative position embeddings. Once examples have slightly more than 20 ones, the maximum number seen during training, the models resort to predicting either 0 or 1 for all examples.
D.2 Addition
While the addition task has been previously studied in the context of decoder-only Transformers, it also straightforward to compute addition in an encoder-only Transformer. We give an example input and output sequence pair in Table
Table 4: Example input and output sequences for addition task.
Input Sequence
START
3
4
+
5
6
END
Output Sequence
PAD
PAD
PAD
PAD
9
0
PAD
The algorithm initializes pointer to the first digit of each input, and pointer to output buffer. The algorithm then iterates over the digits, adding the current digits, adding the value to the output buffer, and updating the carry value. When the last digit has been processed, the final carry value is added to the output buffer.
Our program for computing addition is given in Figure 13
The minimal compiled MLP width (i.e., number transition rules) is 884, and the program uses attention heads. The program requires N layers, where N the number of digits the larger of the two inputs.
34

---

Published in Transactions on Machine Learning Research (2/2025)
def init(x): Initializes pointers.
if xi"token_right" END_TOKEN: x["ptr_b"]
x["ptr_out"] = if x[""toke_rright]] == ADD_TOKEN: x["ptr_a"]
def iterate(x):
Execute one step of addition. rav_sum x["value_carry") if x("ptr_out");
"'"ptr_atokke
x["value_out"] rav_sum X
x s["value_carry) = raw_eum // 10 # Nove all pointers to the left.
a Attention heads attending to the right will be undefined.
if x["token"] != END_TOKEN:
x x["ptr_out"] x["ptr_out_right x["ptr_a"] x["ptr_b"] x("ptr_b.right"]
def finalize(x):
Finalize output by adding the final carry to the output. 1f x["ptr_out");
x["value_out"] = x["value_carry"
x x["step"] STEP DONE
def (x)
1f x["step
INIT:
init(x) x["step"] STEP ITERATE
elif x["step"] STEP ITERATE:
if aa = START_TONE TOKEN:
x["step"] = STEP FINALIIE
else:
(x)
elif st"step"]
STEP TTEP_FINALIZE:
finalize(x)
Figure 12: MLP function for addition program.
35

---

output to, Published in Transactions on Machine Learning Research (2/2025)
variables
"token": pb. var (INPUT_ RANGE, ieput__nittff=llame
x),
# This variable tracks the current processing step.
"step": pb. var (NUM_STEPS),
# These are pointers to which digit is currently being processed.
# They are '1" at the position of the current digit to process, and 10" # otherwise.
"ptr_a": pb. (2),
"ptr_b": pb. (2),
This pointer and '0" othervise.
'1' the position write
"ptr_out"" var
# This tracks the "carry" value form the previous iteration. "value_carry": pb. var (10),
a This tracks the final output value for given digit. "value_out": var (10), # Static variables used as attention query. "one": pb. raage=_,,_ default-1),
attention
For these relative attention heads, alvays want attend the # position inmediately the right.
"token_right"; "eelative""toker 1),
"ptr_a_right"; v_relative ("ptr_a", 1),
"ptr_b.right"; v_relative (_"ttrt 1), "ptr_out_right" a For these attention heads, de vant to attend to the positions associated # with the current pointers.
(uptr 1),
"ptr_a_token"; (kv("one", "ptr_a", "token"), "ptr_b_token": qkve", "ptr_b", "token"),
return ( variables=variabless headsrattention
tin_f==fin_ta in,
output output_names name- "value_out _ott", input_rangee Sapppt__anng--NUT_RAAAA RANGEE position_rangenitne halt_spec_pb. halt halt_ spec("
Figure 13: Program for adding two positive integers. The MLP function is defined in Figure 12
36

---

Published in Transactions on Machine Learning Research (2/2025)
D.3 SUBLEQ
Our ALTA program for implementing SUBLEQ interpreter is given in Figure 15 The input tokens define the set of memory registers, and program execution starts at position O. SUBLEQ stands for SUBtract and branch Less-than or EQual to zero. Commands in SUBLEQ are specified by three memory addresses A, B, and C. Executing a command consists of subtracting the value at memory address A from the value at address B, and writing the result to address B. If the result less than or equal zero, then the program jumps to the command address C, or will halt if C' < |Giannou al. (023) showed how a Looped Transformer implement interpreter for a restricted form of SUBLEQ that did not allow self-modifying code, i.c., memory registers specifying SUBLEQ instructions could not. be modified during program execution. Our program implements a SUBLEQ interpreter without such restrictions, i.c. without differentiating between program and memory registers. Additionally, our implementation executes 1 SUBLEQ instruction every 3 layers, as opposed to every 9 layers for the implementation of Giannou et al. (2023). However, our approach requires ( MLP hidden dimensions (i.e., transition rules), where N the number of possible memory values. The construction of Giannou al. (2023) requires only (loggN) MLP hidden dimensions.
Our compiled models use attention heads. Each token in the input and output of our model corresponds to single register, with the input and output sequences encoding the initial and final values of the registers, respectively. We verified the correctness of our program and compiled models using unit tests involving simple SUBLEQ programs.
37

---

Published in Transactions on Machine Learning Research (2/2025)
def encode (value):
Encodes register value as positive integer.
return value MIN_VALUE
def decode (value):
Decodes register value from positive integer. return value MIN_ VALUE
def .update_position(z, position_a):
2 ='"position_a"
position_a 2 [""osiition_"" position_a 1
2 =["position_c"] position_a
def fin_fa(z):
if z["state"] STATE_1:
if decode (x[">"])
(x""de < 0:
# 'a' or 'b' not valid register positions. z["state"] STATE_DONE return =["state"] STATE_2
elif z["state"] == STATE_2: 2:
Compute mem [b] mem [a].
men_b decode (2 [immem_b"]] j[mjum"" int (nem <= 0)
(2
Update menory value at position 'b'. if 2["position"] == =["b"]:
=["mem"] = encode(men_b
z =["state"] STATE_3
elif z["state"] == STATE_3:
Determine next instruction. if []"jump"]:
Junp instruction 'c', if decode decode(x("""] (z["<"]) < 0:
# Break if 'c' is negative.
2 =["state"] = STATE_DONE
elae:
popdatee_posttion/e =["state"] = STATE_1
2["c"])
else:
Proceed to next instruction.
_update_positionn =["position_a""
2['state] STATE_1
Figure 14: MLP function for interpreting SUBLEQ instructions.
38

---

Published in Transactions on Machine Learning Research (2/2025)
men ren_range MAAX_VLUUE VALUE)
variables
Value of register.
"mem" (men_range, imput_init_fo--lmbbll x),
Position of register.
"pos": var (men_range, josttion_intt_fmmnccee)
Position of current instruction.
"pos_a": var var
default=encode (0)),
"poa_b": var var
default=encode (1))
"pos_c": var (arem_rrange, default=encode (2))
Program state.
"state": (NUM_STATES),
Whether jump next instruction.
"jump": (2)
attention_heatt heads =
# Valuea of registers at 'poa_a', 'poa_b', and "p
"a": qkv("pos_a", "pos", "mem").
"b": qkv("pos_b", "pos", "mem")
"c"; ("v "pos", "mem" Value of registers 'a' and bb'. "mem_a": ("a", "pos", "mem"), "mem_b" qkv("b", "pos", "mem"),
return progran_spect
variableesvariabbles
beddpatteentioo.baadds fin_fa=fin_fa.
output_namer men"" input_range rappttrraggemee_rrnng,, poo_range-MUN.POSTTTOO nalt-halt_spec ("state", halt balt.,valu--777EEE.008)) sepeaall_peeeesttte_e_
Figure 15: Program for interpreting SUBLEQ instructions. The MLP function is specified in Figure 14
39

---

Published in Transactions on Machine Learning Research (2/2025)
Table 5: Example of shift-reduce parsing for SCAN, for the input "jump twice". Our ALTA program
represents the state of the stack and parse tree using a variable number of input tokens.
Action
Stack
Parse Tree
Input Buffer
Initialize
(
jump, twice) >
Shift
jump }
()
{ (wice) )
Reduce
{ NN
( [NT - jump)
{ (twice) )
Shift.
NT, twice) )
( + jump
() ()
Reduce
{ )
{ NT + jump, NT - NT twice )
D.4 SCAN
Background The SCAN suite of compositional generalization tasks (Lake & Baroni 2018 require map- ping natural language commands (e.g., "jump twice and look left") to action sequences (e.g., JUMP JUMP LTURN LOOK). The suite is inspired by the linguistic notion of systematic compositionality, i.c., the ability to recombine a finite set of elements in novel ways (Chomsky 1957 Montague 1970). Certain splits have been shown to be challenging for Transformer-best based models (Keysers et al. [2200] Furrer et al. 2020 Qiu ct al. 2022b Kazemnejad et al. 2024), especially the length split and the Maximum Compound Divergence (MCD) splits proposed by Furrer et al. (020)). Notably, no Transformer- based model has been shown to reliably solve these tasks, without relying on some symbolic decomposition of the task (Zhou et. al. 2023a)
or training data augmented by a symbolic system (Qiu et al. 2022a). Other successful solutions have also involved symbolic parsing of some form (Shaw et al. 2021 Chen et al.. 2020 Herzig & Berant 2021). While prior work has studied the expressivity of various classes of Transformers with respect to formal languages, it
has primarily focused recognizing and generating Dyck languages (Yao al. 2021 Bhattamishra et al.,
2020 Hahn 2020 Weiss et al., 2021 Ebrahimi et al. 2020), and thus has not previously shown constructive
demonstration of how a Transformer can solve the SCAN task.
Program Our approach follows Shaw el. al. (2021) in formalizing the SCAN task translation given a quasi-synchronous context-free grammar. Notably, the SCAN grammar is unambiguous and can be parsed in linear time. First, the program executes shift-reduce parse of the input sequence, representing the parse tree. Second, the ALTA program decodes the output sequence by traversing the parse tree. The program represents the necessary variable-length data structures (a stack, parse tree, and buffer) using a variable number of imput tokens. We include additional "memory" tokens in the imput to ensure there are a sufficient number of tokens to represent these structures. We give an example of shift-reduce parsing for SCAN in Table 5
Similarly to approach to the addition task, while SCAN has commonly been studied using encoder- decoder or decoder-only Transformers, we can also represent the task using an encoder-only Transformer. We ensure there is sufficient number of "memory" tokens so that the entire output sequences can be represented in the encoder output, even if the output sequence exceeds the length of the input sequence.
In Figure 17 we show a program for parsing SCAN inputs using a shift-reduce parser. This program is a fragment of the overall SCAN program, which also decodes the output from the parsed representation, but is too long to include in this paper. We refer the reader
open-source code for the full SCAN program.
Program Size The minimal version of the SCAN program compiles to a Transformer encoder with fewer than 2,000 MLP dimensions and 13 attention heads. The number of layers required scales with the input length and number of tree traversal operations to decode the output from the parse tree. All examples in the dataset can be parsed with fewer than 512 layers. However, it is likely that more layer-efficient implementations exist.
End-to-end Training Details We trained mncodrrddecoeer Transformers end-to-end on the SCAN length and MCD splits. We varied the number of layers, whether weight sharing was used, and the type of positional
40

---

Published in Transactions on Machine Learning Research (2/2025)
def shift_ shift_stackkpointters (z, stack_pointer -fffattte_offe:;
seu_stack_pointer_0 at """stack_pointer_
stack,poiater_
2 =["stack_pointer_ new_stack_pointer_00
"'"stack_poitter_1" tow_stack.pointer_.0
""sstackpointer__
nev_stack.pointer__
2 ["" "'ttck_pointer_S"
new_stack.polnter_0
def reduce(z, matched_rule).
# Pop RHS elements and add LES nonterminal to stack.
if z["position"] == ssl"stackppointrr.
rule_len (mate_ee((atthhee_rule))))
=["eymbol_id"] rule_ hili__hhe_dmmatcceddd.eee
eiaft__racc..ootterrte
rule_len rule_een((atccee__rr)))
Add rule parse tree.
if z["position"] z("tree_pointer"] Use 1-indexing to reserve 0 for rule. 2["rule_id"] rule_id rule z [""sroe_pointer"] +=
def s2)
# Shift the next token to the stack. if 2["position"] al"stack_poiater_0 z["symbol_id"]
tbift__ttac_ppinterrr [["input_pointer"] +=
1)
def fin_fn(z):
if not z d]"done"
# Check if top-3 stack symbols (and 1 lookahead token) match any rule. matched_rule
baybe_match_rule
z [''inptt_oointer_ee "oree___ook__") 2 2""stack_syabol_1""
z [""stack_symbol_22 2["stack_symbol,
if matched_rule is not None:
(z, matched_rule)
else:
Check if parsing is complete.
if 2 [["done"]
== EOS_ID:
else:
sz)
Figure 16: MLP function for parsing SCAN.
encodings used in the encoder. When using weight sharing, we trained with up to 256 encoder layers and 256 decoder layers (as our ALTA program requires at most 512 total layers). Without weight sharing, we only trained with up to 64 encoder layers and 64 decoder layers due to memory constraints.
Constant in all experiments were the following hyperparameters: emboddings with dimension 128, hidden layer sizes of 512, 8 attention heads with dimension 128, an Adafactor optimization function, GeLU activation functions, a learning rate of 5e-4, 100,000 steps, and SentencePiece vocabulary (Kudo & Richardson 2018
Results Figure 18 plots test accuracy when using weight sharing and relative positional encodings. Accu- racy does not improve as the number of layers increascs. (If anything, there is a slight inverse relationship between number of layers and test accuracy.) This is the case all experiments, regardless of the type of positional encoding used in the encoder and whether weight sharing is used. (See Table 6 for all results.)
The fact that incrcasing the number of layers docs not improve gencralization indicates that when trained with standard supervision, Transformers are unable to take advanlage of extra layers to learn a function consistent with the sequential ALTA program that generalizes perfectly. Notably, all end-to-end experiments fit the training set, even with just two layers, so there exists some other algorithm that fits the training set that requires at most two layers. We speculate that in all cases, the end-to-end supervised Transformers unable to learn the sequential algorithm with perfect generalization because of bias towards learning
41

---

Published in Transactions on Machine Learning Research (2/2025)
variables {
"token": (NUM, INPUT, TOKENS, iaput, intt_iiiit..naambdee
"position": (OUK__POSITIOOS,,
itttsioonnitt_tnllllele x: Whether parsing is complete. "done": var (2),
Pointer to the next stack position, and then the elements # the stack.
"stack_pointer_O"" var (EUX_NNNITIION POSITIONS default=STACK (FFaULt=STAACKOFFFEET)
default=STACK_0PE OFFSET),
"stack_pointer_ var (arONK_POSITTONN
1).
"etack_pointer" 2": var (NUX, (WWDX_OSSTIONNS
defaul:-STAcK OFFSET 2)
"stack.pointer_9" var (ar(KK_POOSITON the "tree_pointer"; (SUM_FOSITIONS, default=TREE_0FFSEE))
default=STACCK 3),
Pointer
Pointer the next input token to process.
"saput_pointer"; var (NUM_POSITIONS, default=1NNUT_OFFSEE)
# Stores index of associated parsing rule
"rule_id": var (NUM_RULES),
# Stores аymbol ID associated with stack element. "symbol_id"; var (NUM (HUM_SYMBOLS),
heads
Get token input pointer.
"iuput_pointer_ taatter_toke__i""" qkv qtv("iautt.pointe" "position", "token"), # Get top symbols on stack.
"stack_symbor _1: ddt"stack_pointer,
"position", "symbol_ide
"stack_eymbol 2": qkv ("stack,pointer,22 "position", sdymbol_id"). "stack_symbol_33" qkv ("stack,pointer_3", "position", "symbol_id"),
return progran_spect
veriables-varrislles,,
heads-heads, ffa_fn-ffa_fa.
srrput_name="rul___ie
www.rrraae-NUK__PPP_TTOOO TOKENS. position_ranae Baapaao_aaanne-U_PPPPOTTTAA
Figure 17: Program for parsing SCAN, The MLP function is defined in Figure 16
algorithms that require fewer layers to execute. These results are consistent with our results on Parity, which Transformers trained with end-to-end supervision do not learn the sequential algorithm with perfect generalization, instead seeming to mimic algorithm that requires just one layer (see @D.1).
Similarly, prior work has generally not shown a consistent benefit from increasing model size or number of layers on SCAN or other similar compositional generalization tasks when fine-tuning. Furrer et al. (2020) evaluated models of various sizes on SCAN and did not find that increasing size improves generalization. Ontanon al. (2022) showed some improvement on the SCAN length split when incrcasing the number of layers with weight sharing, but only evaluated up to six layers. Qiu et al. (2022b) and Petty et al. (2024) systematically evaluated the impact of model size and depth, respectively, on compositional generalization (though not on SCAN) and found that when fine-tuning, scaling curves are flat quickly saturate. However, these papers evaluated models with at most 118 layers, and only Ontanon et al. (2022) evaluated models with varying numbers of layers using weight sharing.
42

---

Published in Transactions on Machine Learning Research (2/2025)
Table 6: Test accuracy for all Transformers trained with standard, end-to-end supervision on SCAN. Re-
gardless of the configuration and split, increasing the number of layers does not improve test accuracy.
Number of Layers
Split
Weight Sharing
Encoder Positional Encoding
2
6
16
32
64
128
256
Length
Yes
Relative
11%
10%
9%
7%
7%
7%
5%
Length
Yos
Absolute
8%
9%
7%
5%
5%
6%
4%
Length
No
Relative
9%
9%
11%
10%
10%
Length
No
Absolute
9%
6%
7%
6%
4%
MCD1
Yes
Relative
5%
3%
2%
4%
3%
2%
MCD1
Yes
Absolute
4%
1%
3%
2%
2%
2%
1%
MCD1
No
Relative
4%
4%
3%
2%
2%
MCD1
No
Absolute
5%
2%
1
1%
2%
MCD2
Yes
Relative
9%
13%
9%
5%
4%
7%
3%
MCD2
Yes
Absolute
5%
6%
2%
2%
2%
2%
2%
MCD2
No
Relative
13%
7%
4%
5%
4%
MCD2
No
Absolute
4%
4%
2%
2%
1%
MCD3
Yes
Relative
12%
12%
6%
2%
3%
MCD3
Yes
Absolute
3%
4%
3%
3%
3%
2%
2%
MCD3
No
Relative
8%
8%
7%
3%
2%
MCD3
No
Absolute
4%
5%
2%
3%
1%
Accuracy by Number Layers With Weight Sharing
Length
MCD1
MCD2
MCD3
0.25
0.25
0.25
0.25
Accuracy
0.15
0.15
0.15
0.15
0.30
0.10
0.10
0.10
0.06
0.05
0.05
0.05
2
16
32
128
255
2
16
32
64
128
256
16
32
64
128
256
16
32
64
128
256
Number of Layers
Figure 18: Test. accuracy by number of layers SCAN splits when trained using standard, end-to-end supervision with weight sharing and relative positional encodings. Test accuracy does not increase as the number of layers increases.
43


Reviews:
[{"review_id": "riWe7mXD4p", "reviewer": "Reviewer_rgVx", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The atuhor conducts experiments to try to support their claims.\n\nWeakness:\n* **The paper writing is not clear**\n   * It seems that the task of this work is not very common. Therefore, after reading the paper several times, I still find it difficult to know what the task is.\n      * What is the input and what is the output?\n      * And moreover, could you please suggest what the task belongs to? Is it a sequence classification task, as the author uses an encoder-only transformer?\n* The work is based on RASP. Therefore, which one is better for RASP and this work?\n* **For Experiment Part**\n   * How the dataset is constructed?\n   * It seems that the author does not clearly present the dataset tasks. For example, what is Sum + Modulo? Could you give an example?\n   * For Section 4.2, it seems that the dataset is Addition.\n      * Traditionally, the addition task could be regarded as a next-token prediction task. And we usually solve it via a decoder-only Transformer.\n      * Therefore, how does ALTA solve such ta ask?\n\nAs I have several above concerns, I have to ask the author to give further explanation.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work propose a new programming language called ALTA and a compiler that can map ALTA programs to Transformer weights. Also, this work complements and extends this prior work, offering the ability to express loops and to compile programs to Universal Transformers, among other advantages.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please check the weakness.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "NxnKupe173", "reviewer": "Reviewer_SgJx", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\nThis paper is well-organized, and relatively easy to follow given the complexity of the proposed approach. The experiments on the parity problem provide some evidence of the effectiveness of the proposed framework.\n\nThis work presents a theoretical analysis on both the expressivity of Transformers, and the learnability of algorithms. Even if an algorithm can be expressed with a Transformer class, training such model and making it generalize beyond its training set can be challenging. This work shows that training on execution traces from an ALTA program can improve performance for some tasks.\n\nThis work introduces the analytical notion of *minimal* programs with respect to a training set: a minimal program will be compiled to a Transformer which uses all of its parameters fully to represent the training set. A non-minimal program will be compiled to a set of weights that can be changed without affecting the prediction power of the training set, ie: some parameters become under-specified.\n\n**Weaknesses**\n\nOne weakness of this work, as recognized in the paper and as with similar approaches, is the relatively small variety of representable Transformers by the ALTA framework. For instance, the selection matrix of the self-attention layers can only be binary. There is limited support for numeric computations and no support for probabilistic output distributions.\n\nAnother weakness is the small amount of experimental results on different tasks. The paper shows that the ALTA framework allows to generalize to various difficulty levels of the parity problem, but does not show the same thing for the addition, SUBLEQ, and SCAN tasks. Instead, the ALTA programs are presented in respective Appendices, but it is not clear if the compiled Transformer, for each of these task achieves a strong performance across all difficulty levels. Does having an ALTA program automatically guarantee that the problem is solved (ie: 100% accuracy regardless of the input difficulty)? Some clarifications around this and the deterministic nature of the approach would greatly improve the paper. \n\nIn particular, in the case of the SCAN benchmark the paper claims multiple times that it “solves” the SCAN benchmark, however, when looking at the results in Figure 18 and Table 5, we see that the best accuracy achieved by training a model is 13%. Does the trained model use ALTA traces? it would be nice to show this experiment, otherwise the motivation for training a model is not clear. What are the accuracies of the ALTA compiled transformer? If it is 100% by design of the ALTA program, it must be clearly stated. Until this is clarified, I cannot say that *the claims made in the submission are supported by accurate, convincing, and clear evidence*.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work presents a new programming language (ALTA) and its compiler that compiles symbolic programs to Transformer weights. Based on previous work (RASP language), ALTA additionally allows loops to be represented with layer-wise Transformer weight sharing. This mechanism enables mapping programs to Universal Transformers and the implementation of length-invariant algorithms. In addition, ALTA is more efficient than previous approaches as it can compile more complex programs to reasonably sized Transformers. This is done by representing MLP layers with sparse transition rules.\n\nOverall this work presents new theoretical and empirical insights on the representation power and learnability of Transformers.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- [*critical to secure my recommendation*] Clarify the SCAN results: mention the accuracies of the ALTA compiled transformer, if it is not 100%, then rephrase all occurrences saying that you “solve” the SCAN benchmark (see weakness paragraph 3)\n\n- [*to strengthen the work*] Accuracy numbers on Addition (and SUBLEQ if that is possible) (see weakness paragraph 2)", "broader_impact_concerns": "no concern", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "2YVUAIQHng", "reviewer": "Reviewer_XoJj", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- ALTA's sparse representation enables compilation of complex programs into reasonably-sized Transformers by leveraging the sparsity of transition rules to reduce required MLP hidden dimensions, sometimes by several orders of magnitude. This is in contrast to RASP's dense lookup table approach which can suffer from combinatorial explosion in variable combinations.\n\n- The representation of MLP sublayer computations as sparse transition rules provides valuable insights into the generalization potential of MLP layers and Transformers as a whole, offering a new lens to understand their behavior.\n\n- The paper constructively demonstrates how Transformers can solve the SCAN task through a shift-reduce parsing algorithm, proving that Transformers can represent systematic generalization through a finite set of transition rules and attention operations.\n\n- The set of transition rules in ALTA programs provides a new metric for measuring Transformer complexity, enabling studies of whether Transformers have an inherent simplicity bias by comparing the number of transition rules across different programs.\n\n- The work establishes a solid theoretical framework for analyzing minimality and learnability, effectively bridging the gap between expressivity and practical implementation.\n\nWeaknesses:\n\n- While the paper presents attention as a mechanism for computing binary selection matrices and aggregating information across positions, recent advances in MetaFormer [R1] architectures suggest that simpler operations like pooling could potentially replace attention. Similarly, State Space Models and Linear Transformers have demonstrated success using recurrent structures instead of traditional attention. The paper would benefit from exploring whether ALTA's framework could extend to these alternative architectures, potentially offering insights into the essential computational requirements for tasks like SCAN and parity.\n\n- Although the paper mentions that ALTA programs can be executed in decoder-only Transformers by adding causal attention masks and outer auto-regressive decoding loops, this aspect receives limited discussion. A more detailed analysis of decoder-only implementations would be valuable, particularly given the growing prominence of decoder-only models in contemporary research.\n\n[R1]: Yu, Weihao, et al. \"Metaformer is actually what you need for vision.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces ALTA, a new programming language and compiler that maps programs to Transformer weights, extending previous work (RASP and Tracr) with support for loops without intermediate decoding steps and compilation to Universal Transformers. The authors demonstrate novel expressivity results by showing how Transformers can represent length-invariant algorithms for parity and addition, as well as a solution to the SCAN benchmark, all without requiring intermediate decoding. To analyze the gap between expressivity and learnability, they introduce tools including trace supervision for analyzing when algorithms can be learned from training data, and a theoretical framework for analyzing program minimality. Through empirical analysis, they show that Transformers trained end-to-end tend to learn simpler algorithms requiring fewer layers rather than more complex algorithms that generalize better, and demonstrate that increasing model depth does not necessarily improve compositional generalization. The authors make ALTA available as open source, enabling further research into Transformer capabilities and limitations. In general, this paper seems make a good contribution to this field and doing a very thorough experiment.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to my weakness part.", "broader_impact_concerns": "I do not have any concerns about the broader impacts for this paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper advances understanding of Transformer architectures through ALTA. Compared to previous languages such as RASP, ALTA supports programs with loops, enables compiling programs to Universal Transformers and can support compilation of more complex programs by leveraging sparsity in the MLPs. Practically, the authors demonstrate how transformers can represent length-invariant algorithms and solve the SCAN benchmark, which is a non trivial result. ALTA allows to compile symbolic programs into Transformer weights with loop support and sparse transition rules, which offers an important perspective on Transformer expressivity, generalization and learnability.
Decision: Accept (poster)