Transformers Learn Shortcuts to Automata

Introduction

Modern deep learning systems demonstrate increasing capabilities of algorithmic reasoning. Particularly in modalities such as natural language, math, and code, neural networks can successfully parse and synthesize sequences containing symbolic information and compositional structure. To exhibit these functionalities, these networks are required to learn and execute the relevant discrete algorithms within their internal representations. A core open question in this domain is that of mechanistic understanding: how do neural networks encode the primitives of algorithmic reasoning?

When considering this question, there is an apparent mismatch between classical sequential models of computation (e.g. Turing machines) and the Transformer, the state-of-the-art architecture for neural algorithmic reasoning. If we are to think of algorithms as sequentially-executed computational rules, why should we use a shallowIn practice, a Transformer’s context length (which can be thousands of tokens) is typically far greater than its depth (which can be as small as $6$). and non-recurrent architecture to represent them?

We study this question through the lens of semiautomata, which compute state sequences $q_{1},\ldots,q_{T}$ from inputs $\sigma_{1},\ldots,\sigma_{T}$ by application of a recurrent transition function $\delta$ (and initial state $q_{0}$):

Semiautomata describe the underlying dynamics of automata, which are simply semiautomata equipped with mappings from states to outputs. With unbounded state spaces, automata can represent all algorithms; however, even bounded automata form a rich class of sequence processing algorithms, containing regular expression parsers and finite-state transducers. In reinforcement learning and control, semiautomata are better known as deterministic Markov models (where $\sigma_{t}$ are actions); thus, in addition to algorithmic reasoning, the results in this work also pertain to Transformer dynamics models.

We perform a theoretical and empirical investigation of whether (and how) non-recurrent Transformers perform the computations of semiautomata. We find that Transformers learn shortcut solutions, which correctly and efficiently simulate the sequential transitions of semiautomata using a shallow parallel circuit, rather than naively iterating the single-step recurrence. Shortcuts arise from hierarchical reparameterizations of a semiautomaton’s global transition dynamics.

Our contributions. Our theoretical results provide structural guarantees for the representability of semiautomata (thus, iterative algorithms) by one pass through a shallow, non-recurrent Transformer. In particular, we show that:

Shortcut solutions, with depth logarithmic in the sequence length, always exist (Theorem 1).

Constant-depth shortcuts exist for solvable semiautomata (Theorem 2). These are understood via the Krohn-Rhodes theorem, a landmark result in semigroup theory. Conversely, there do not exist constant-depth shortcuts for non-solvable semiautomata, unless $\mathsf{TC}^{0}=\mathsf{NC}^{1}$ (Theorem 4).

For a natural class of semiautomata corresponding to path integration in a “gridworld” with boundaries, we show that there are even shorter shortcuts (Theorem 3), beyond those guaranteed by the general structure theorems above.

We accompany these with an extensive set of experimental findings:

Transformers learn shortcuts with standard training (Section 4). Across a wide variety of semiautomaton simulation problems, we find that shallow non-autoregressive Transformers successfully learn shortcut solutions: despite the non-convex optimization problem, gradient-based training works. This suggests that shortcuts are plausible mechanisms for algorithmic reasoning in non-synthetic sequence models, and lies beyond our current theoretical understanding.

Shortcuts are statistically brittle (Section 5). We identify empirical weaknesses of the shortcuts found by Transformers: poor out-of-distribution generalization (including to unseen sequence lengths) and worse performance than RNNs under limited supervision. Toward mitigating these drawbacks and obtaining the best of both worlds, we show that with recency-biased scratchpad training, autoregressive Transformers can easily be guided to learn the iterative RNN-like solutions (chain-of-thought generation).

Emergent reasoning in neural sequence models. Neural sequence models, both recurrent (Wu et al., 2016, Peters et al., 2018, Howard and Ruder, 2018) and non-recurrent (Vaswani et al., 2017, Devlin et al., 2018), have become an era-defining tool for parsing and transducing data with combinatorial structure, such as natural language and code. A nascent frontier is to leverage neural dynamics models, again both recurrent (Hafner et al., 2019) and non-recurrent (Chen et al., 2021a, Janner et al., 2021), for decision making. At the highest level, the present work seeks to understand the mechanisms by which these models perform combinatorial and algorithmic reasoning.

Computational models within neural networks. Despite the preponderance of empirical successes, many mysteries remain, towards understanding the internal mechanisms of neural networks capable of algorithmic reasoning. It is known that self-attention realizes low-complexity circuits (Hahn, 2020, Elhage et al., 2021, Merrill et al., 2021, Edelman et al., 2022), declarative programs (Weiss et al., 2021), and Turing machines (Dehghani et al., 2019, Pérez et al., 2021, Giannou et al., 2023). Interpretable symbolic computations have been extracted from trained models (Clark et al., 2019, Vig, 2019, Tenney et al., 2019, Wang et al., 2022). Our conclusions are closest to the literature on the universal representation on Turing machines (which are automata with unbounded states); however, our work is unique in characterizing the recurrent machines whose execution loops can be efficiently unrolled into a single pass of a shallow Transformer.

Learning elementary algorithms with Transformers. Our work provides a unifying lens on many recent investigations on whether (and how) Transformers represent certain classes of fundamental algorithmic computations. These include bounded-depth Dyck languages (Yao et al., 2021), modular prefix sums (Anil et al., 2022), adders (Nogueira et al., 2021, Nanda and Lieberum, 2022), regular languages (Bhattamishra et al., 2020), and sparse logical predicates (Edelman et al., 2022, Barak et al., 2022), which are all special cases of simulating finite-state automata. Thus, our work provides guarantees of shallow Transformer solutions in all of these settings. Zhang et al. (2022) empirically analyze the behavior and inner workings of Transformers on random-access group operations and note “shortcuts” (which skip over explicit program execution) similar to those we study.

We provide an expanded discussion of related work in Appendix A.5.

Preliminaries

A semiautomaton ${\mathcal{A}}:=(Q,\Sigma,\delta)$ consists of a set of states $Q$, an input alphabet $\Sigma$, and a transition function $\delta:Q\times\Sigma\rightarrow Q$. In this work, $Q$ and $\Sigma$ will always be finite sets. For all positive integers $T$ and a starting state $q_{0}\in Q$, ${\mathcal{A}}$ defines a map from input sequences $(\sigma_{1},\ldots,\sigma_{T})\in\Sigma^{T}$ to state sequences $(q_{1},\ldots,q_{T})\in Q^{T}$: $q_{t}:=\delta(q_{t-1},\sigma_{t})$ for $t=1,\ldots,T$. This is a deterministic Markov model, in the sense that at time $t$, the future states $q_{t+1},\ldots,q_{T}$ only depend on the current state $q_{t}$ and the future inputs $\sigma_{t+1},\ldots,\sigma_{T}$.

We define the task of simulation: given a semiautomaton $\mathcal{A}$, starting state $q_{0}$, and input sequence $(\sigma_{1},\ldots,\sigma_{T})$, output the state trajectory ${\mathcal{A}}_{T,q_{0}}(\sigma_{1},\ldots,\sigma_{T}):=(q_{1},\ldots,q_{T})$. Let $f:\Sigma^{T}\rightarrow Q^{T}$ be a function (which in general can depend on $\mathcal{A},T,q_{0}$). We will say that $f$ simulates $\mathcal{A}_{T,q_{0}}$ if $f(\sigma_{1:T})=\mathcal{A}_{T,q_{0}}(\sigma_{1:T})$ for all input sequences $\sigma_{1:T}$. Finally, for a positive integer $T$, we say that a function class $\mathcal{F}$ of functions from $\Sigma^{T}\rightarrow Q^{T}$ simulates $\mathcal{A}$ at length $T$ if, for each $q_{0}\in Q$, there is a function in $\mathcal{F}$ which simulates $\mathcal{A}_{T,q_{0}}$

Every semiautomaton induces a transformation semigroup ${\mathcal{T}}(\mathcal{A})$ of functions $\rho:Q\rightarrow Q$ under composition, generated by the per-input-symbol state mappings $\delta(\cdot,\sigma):Q\rightarrow Q$. When ${\mathcal{T}}(\mathcal{A})$ contains the identity function, it is called a transformation monoid. When all of the functions are invertible, ${\mathcal{T}}(\mathcal{A})$ is a permutation group. See Figure 1 for some examples which appear both in our theory and experiments; additional background (including a self-contained tutorial on the relevant concepts in finite group and semigroup theory) is provided in Appendix A.2. An elementary example is a parity counter (leftmost in Figure 1): the state is a bit, and the inputs are {“toggle the bit”, “do nothing”}; the transformation semigroup is $C_{2}$, the cyclic group of order 2.

2 Recurrent and non-recurrent neural sequence models

An $L$-layer (or depth-$L$) Transformer is a different sequence-to-sequence network, consisting of alternating self-attention blocks and feedforward MLP blocks

Typical Transformers are shallow, in the sense that $L\ll T$. In practice, this makes their inference and gradient computations highly parallelizable, with the number of sequential computation steps scaling linearly in $L$, while RNNs require a scaling linear in $T$. While there is always a canonical way for RNNs to simulate semiautomata, the answer to the analogous question for shallow Transformers is far less obvious, and forms the main topic of this paper.

Theory: shortcuts abound

To simulate a semiautomaton at length $T$, a $T$-layer Transformer can implement the same sequential solution as an RNN: let the $t$-th layer embed the state transition $q_{t-1}\mapsto q_{t}$. We define shortcuts as solutions which implement the same functionality with a significantly smaller depth.

Let ${\mathcal{A}}$ be a semiautomaton. For every $T\geq 1$, let $f_{T}$ be a sequence-to-sequence neural network which simulates ${\mathcal{A}}$ at length $T$. Then, we call this sequence $\{f_{T}\}_{T\geq 1}$ a shortcut to ${\mathcal{A}}$ if the sequence of network depths $D:=\{D(f_{T})\}_{T\geq 1}$ satisfies $D\leq o(T)$.

By this definition, shortcuts are quite general, and some are less interesting than others. For example, it is always possible to construct a constant-depth neural network which memorizes all $|\Sigma|^{T}$ values of ${\mathcal{A}}_{T,q_{0}}$, but these networks must be exceptionally wide. There are also solutions which emulate transitions in “chunks”, letting each of (say) $\sqrt{T}$ layers perform $\sqrt{T}$ consecutive state transitions; however, without exploiting the structure of the semiautomaton, this would require width $\Omega(|\Sigma|^{\sqrt{T}})$. To rule out these cases and focus on interesting shortcuts for Transformers, we want the other size parameters (attention and MLP width) to be small: say, scaling at most polynomially in $T$, $|Q|$, and $|\Sigma|$. To construct such shortcuts, we need ideas beyond explicit iteration of state transitions.

We begin by noting that polynomial-width shortcuts always exist. This may seem counterintuitive if we restrict ourselves to viewing a Transformer’s intermediate activations as representations of states $q_{t}$, like the RNN solution. Instead, a Transformer can encode and hierarchically compose transformations $\delta(\cdot,\sigma):Q\rightarrow Q$ (see Figure 2a), leading to far shallower solutions:

Transformers can simulate all semiautomata $\mathcal{A}=(Q,\Sigma,\delta)$ at length $T$, with depth $O(\log T)$, embedding dimension $O(|Q|)$, attention width $O(|Q|)$, and MLP width $O(|Q|^{2})$.

This is proven in Appendix C.2, and leverages the ability of a self-attention head to approximate hard attention (i.e. concentrate its mixing weights on a single position). However, self-attention heads can also perform soft attention (i.e. depend on a large number of previous positions), enabling even shallower implementations of certain sequential computations. For example, the parity automaton can be simulated by a single Transformer layer (see Lemma 6): soft attention computes prefix sums in parallel, then the MLP computes “mod 2”. This leads to a significantly more nuanced question: when are there even shallower shortcuts? At first glance, such solutions may seem rare, and specialized to simple cases such as parity.

Our resolution to this question comes from the Krohn-Rhodes decomposition theorem (Krohn and Rhodes, 1965), a landmark result which vastly generalizes the uniqueness of prime integer factorizations, and created the mathematical field of algebraic automata theory (Rhodes et al., 2010). The conclusion is quite unintuitive: allowing for both hard and soft modes of attention, constant-depth shortcuts are surprisingly common!

Transformers can simulate all solvableSee Definition 6. Intuitively, the only obstructions are when the semiautomata contain non-solvable groups such as $S_{5}$, the group of all permutations of $5$ elements. semiautomata $\mathcal{A}=(Q,\Sigma,\delta)$, with depth $O(|Q|^{2}\log|Q|)$, embedding dimension $2^{O(|Q|\log|Q|)}$, attention width $2^{O(|Q|\log|Q|)}$, and MLP width $|Q|^{O(2^{|Q|})}+2^{O(|Q|\log|Q|)}\!\,\cdot\!\,T$.

Much of the appendix is dedicated to providing a user-friendly exposition of the relevant algebraic concepts, culminating in the proof of Theorem 2. We provide a few high-level notes below:

Intuitively (illustrated in Figures 2b and 2c), the Krohn-Rhodes decomposition “factorizes” every solvable semiautomaton into modular counters and memory units, glued together via a feedforward cascade product (Definition 4) whose depth only depends on $|Q|$, not $T$. These two types of “prime” semiautomata can be efficiently simulated by depth-$1$ Transformers.

The decomposition depends on the transformation semigroup $\mathcal{T}(\mathcal{A})$. It is non-constructive (much like how the existence of prime factorizations doesn’t entail a procedure to find them). Computationally, these solutions still have to be found by a search procedure. Remarkably, we find that gradient-based training succeeds empirically, despite the worst-case computational hardness of related problems.

What makes Transformers special (vs. other universal function approximators)? The same underlying semiautomaton-to-circuit constructions could be applied to any universal function approximator (like a vanilla MLP with the same depth). Transformers embed all of the constructions in Theorems 1 and 2 with exceptional efficiency, in terms of the network complexity measures discussed in Section 2. Most importantly, the constructions leverage Transformers’ positional weight sharing, which removes allThe only part of Theorem 2 requiring $T$ non-tied neurons is the implementation of mod-$p$ gates. It disappears entirely if we can add auxiliary MLP neurons with periodic activation functions such as $x\mapsto\sin(x)$. suboptimal $T$ factors from the parameter count. We discuss this further in Appendix A.5.

The proof builds a parallel nearest boundary detector for the two boundary (i.e. leftmost and rightmost) states, and can be found in Appendix C.4. We note that gridworlds are known to be extremal cases for the holonomy decomposition in Krohn-Rhodes theory (Maler (2010) discusses this, calling it the elevator automaton). It would be interesting to generalize our improvement and characterize the class of problems for which self-attention affords $O(1)$ instead of $\text{poly}(|Q|)$-depth solutions.

2 Lower bounds

Can Theorem 2 be improved to handle non-solvable semiautomata? (Equivalently: can Theorem 1 be improved to constant depth?) It turns out that as a consequence of a classic result in circuit complexity (Barrington, 1986), this question is equivalent to the major open question of $\mathsf{TC}^{0}\stackrel{{\scriptstyle?}}{{=}}\mathsf{NC}^{1}$ (thus: conjecturally, no). Unless these complexity classes collapse, Theorems 1 and 2 are optimal. In summary, simulating non-solvable semiautomata with constant depth is provably hard:

Let $\mathcal{A}$ be a non-solvable semiautomaton. Then, for sufficiently large $T$, no $O(\log T)$-precision Transformer with depth independent of $T$ and width polynomial in $T$ can continuously simulate $\mathcal{A}$ at length $T$, unless $\mathsf{TC}^{0}=\mathsf{NC}^{1}$.

This is proven in Appendix C.5. The smallest example of a non-solvable semiautomaton has $|Q|=60$ states, whose transitions generate $A_{5}$ (all of the even permutations).

Finally, we note that although our width bounds might be improvable, an exponential-in-$|Q|$ number of hypotheses (and hence a network with $\text{poly}(|Q|)$ parameters) is unavoidable if one wishes to learn an arbitrary $|Q|$-state semiautomaton from data: there are $|Q|^{|Q|\cdot|\Sigma|}$ of them, which generate $|Q|^{\Omega(|Q|^{2})}$ distinct semigroups (Kleitman et al., 1976). If we wish to study how machine learning models can efficiently identify large algebraic structures, we will need finer-grained inductive biases to specify which semiautomata to prefer, a direction for future work.

Experiments: can SGD find the shortcuts?

The results in Section 3 only provide a precise understanding of representability: they show that shortcut solutions exist within the parameter space of a shallow Transformer. To understand whether Transformers can actually learn these shortcuts from data, we must introduce the additional considerations of generalization and optimization. It is notoriously difficult to derive meaningful analyses which account for all of these factors in deep learning; thus, we do not attempt to do so in this work.Our bounds on the parameter count and weight norms do imply classical generalization bounds for appropriately norm-constrained Transformers (Edelman et al., 2022), but these are too coarse-grained to provide non-vacuous predictions of generalization behavior. Instead, we approach the end-to-end question with an empirical lens: trained on sequences arising from a variety of automata, does a shallow (depth-$L\ll T$) Transformer converge to correct simulators of these automata?

For a selection of 19 semiautomata corresponding to various groups and semigroups (detailed descriptions in Appendix B.1.1), we train shallow Transformer (GPT-2-like (Radford et al., 2019)) models to map randomly sampled sequences $(\sigma_{1},\ldots,\sigma_{T})$ to their corresponding state sequences $(q_{1},\ldots,q_{T})$, and evaluate their accuracy on held-out sequences. We vary the depth $L$ from $1$ to $16$, and use freshly-sampled sequences of length $T=100$. In this setup, the number of sequences encountered during training ($\leq 10^{6}$) is far smaller than the number of distinct input sequences ($|\Sigma|^{100}$). Thus, brute-force memorization cannot solve this task, and generalization is necessary to achieve nontrivial performance.

Strikingly, we obtain positive results ($>\!99\%$ in-distribution accuracyOur primary goal is to understand if gradient-based training can find shortcut solutions at all, rather than whether such training is stable. Accordingly, unless otherwise noted, we report the performance of the best model among 20 replicates. See Appendix B for details and sensitivity analyses.) for every finite-state semiautomaton we considered, including ones which generate the non-solvable groups $A_{5}$ and $S_{5}$. Figure 3a gives a selection of our full results (in Appendix B.1). We find that more complex semiautomata (corresponding to non-abelian groups) require deeper networks to learn, in agreement with our theoretical constructions.

Our theoretical results show that there are logarithmic-depth ($S_{5},A_{5}$) and constant-depth (all the others) solutions which simulate these semiautomata with exactly 100% accuracy. Of course, with such long sequences, black-box evaluation of whether this accuracy is reached in the population distribution is computationally infeasible. However, we note that without periodic activations (or some other mechanism for extrapolating to unseen count values), our theoretical constructions require MLPs to memorize the mod-$n$ function. This will be revisited in the out-of-distribution evaluation experiments in Section 5.2, but there is even a corresponding implication for in-distribution mistakes: if a model never sees outlier counts during training, it is expected to make mistakes on those outliers when they appear in evaluation.

Our theoretical results identify shortcut solutions which follow multiple, mutually incompatible paradigms. In general, we do not attempt a full investigation of mechanistic interpretability of the trained models. In particular, we do not claim that the networks discover implementations are isomorphic to those described in the proofs of Theorem 1, 2, and 3. However, as a preliminary exploration, we visualize some of the attention patterns in Figure 3b within successfully-trained models, finding attention heads which perform flat summations (with uniform attention) and conditional resets, agreeing with the construction in Theorem 3.

Although sufficiently deep networks find the solutions with non-negligible probability, the training dynamics are unstable; Figure 3b,c show example training curves, which exhibit high variance, negative progress, or accuracy that decays with continued training. In the same vein as the “synthetic reasoning tasks” introduced by Zhang et al. (2022), we hope that semiautomaton simulation will be useful as a clean, nontrivial testbed (with multiple difficulty knobs) for debugging and improving training algorithms, and perhaps the neural architectures themselves. More details are deferred to Appendix B.1.1.

Further experiments: more challenging settings

The results from Sections 3 and 4 show that Transformers can learn shortcuts end-to-end, unobstructed by depth, generalization, or optimization. However, the experiments in Section 4 are idealized in several ways; a natural question is whether these findings are robust to various challenges that arise in practice. In this section, we investigate the robustness of the shallow Transformer solutions, compared to those found by RNNs (the “natural” architecture for simulating semiautomata). We consider harder forms of supervision (Section 5.1) and evaluation (Section 5.2); details are deferred to Appendix B.2.

2 Out-of-distribution shortcomings of shortcut solutions

The theoretical construction of modular counters (Lemma 6) suggests a possible failure mode: if attention performs prefix addition and the MLP computes the sum modulo $n$, the MLP could fail on sums unseen during training. This suggests that if the distribution over $\sigma_{1:T}$ shifts between training and testing (but the semiautomaton remains the same), a non-recurrent shortcut solution might map inputs into an intermediate latent variable space (like the sum) which fails to generalize.

More ambitiously, we could try to use these models to extrapolate to longer sequence lengths $T$ than those seen in the training data. Promoting this difficult desideratum of length generalization is an intricate problem in its own right; see Yao et al. (2021), Anil et al. (2022) for more experiments similar to ours. Figure 7 shows the performance on sequences of various lengths. In contrast to LSTM’s perfect performance on all scenarios, Transformer’s accuracy drops sharply as we move to lengths unseen during training. This is not purely due to unseen values of the positional encoding: randomly shifting the positions during training can cover all the positions seen during testing, which helps improve the length generalization performance but cannot make it perfect; we see similar results for removing positional encodings altogether. Finally, we empirically show that the above flaws are circumventable. Using a combination of scratchpad (a.k.a. “chain-of-thought”) (Nye et al., 2021, Wei et al., 2022) and recency bias (Press et al., 2022), we demonstrate that Transformers can be guided towards learning recurrent (depth-$T$) solutions, which generalize out-of-distribution and to longer sequence lengths (Figure 7, yellow curves). Details are deferred to Section B.2.3.

Throughout the deep learning literature, the term shortcut is often used to refer to undesired (i.e., misleading, spurious, or overfitting) statistical properties of learned representations (Geirhos et al., 2020, Robinson et al., 2021). Meanwhile, under our computational ($o(T)$ circuit depth) definition, shortcut solutions are perfectly valid ways to represent recurrent computations. The results in this section establish a connection between these notions: partially-learned computational shortcuts can be statistical shortcuts. Specifically, a non-recurrent architecture can “hallucinate” intermediate variables other than the state (e.g. the “count” variable for the parity automaton), is thus sensitive to the coverage of these variables in the training data. This leads to out-of-distribution generalization failures (e.g. on rare counts) which are not present in recurrent models.

The experiments in this section highlight a statistical penalty for learning recurrent computations with a non-recurrent architecture. However, the computational advantage of a shallow architecture is extremely appealing: maximally leveraging parallel computation, training and inference can be done much faster ($O(\log T)$ or $O(1)$ time, compared to $O(T)$). This highlights a delicate tradeoff between RNNs and Transformers, where neither architecture dominates the other, even when considering this elementary class of algorithmic problems. Attaining the best of both worlds with a practical architecture is an interesting avenue for future work.

Conclusions and future work

We have conducted a theoretical and empirical analysis of how shallow Transformers can learn shortcuts to the recurrent computations of finite-state automata. These shortcuts replace $T$ sequential iterations of a recurrent unit with a single pass through $L\ll T$ parallel self-attention layers. Our theoretical results show that shortcuts are ubiquitous, and characterize extremely shallow ones (with $L$ independent of the context length $T$) using algebraic machinery (Krohn-Rhodes theory). Empirically, we have shown that gradient-based optimization successfully finds these shortcuts. While the solutions found in practice generalize near-perfectly in-distribution (Section 4), they lack out-of-distribution robustness (Section 5). We hope that these results shed new light on the internal reasoning mechanisms of Transformers, as well as the design space for architectures capable of algorithmic reasoning.

Future work. This work is an initial foray into the interplay between neural architectures, algebraic automata theory, and circuit complexity. We list some open questions:

Finer-grained circuit complexity of self-attention: For certain automata of interest (e.g. bounded Dyck language parsers (Yao et al., 2021), and the gridworld automata from Theorem 3), there exist extremely shallow Transformer solutions, with depth independent of both $T$ and $|Q|$. Which other natural classes of automata have this property of “beyond Krohn-Rhodes” representability?

When and why does gradient-based optimization work? The precise understanding of hierarchical representation learning in neural networks is an active frontier of research. The empirical tractability of finding “algebraic” shortcut solutions with gradient descent is an especially striking case, as related problems are known to be $\mathsf{PSPACE}$-hard (kozen1977lower, beaudry1992membership, Cho and Huynh, 1991).

Mechanistic interpretability: Automaton simulation problems generate a rich class of challenging test cases for the reverse engineering of trained networks (see Nanda and Lieberum (2022)), and may yield further insights on the inductive biases of Transformers. In our preliminary attempts, we were only able to interpret a small number of simple models.

Acknowledgements

We are very grateful to Abhishek Shetty for helpful discussions about circuit complexity. We also thank Ashwini Pokle for thoughtful comments and suggestions towards improving clarity and readability.

Reproducibility Statement

Complete proofs of the theoretical results are provided in Appendix C, with a self-contained tutorial of relevant algebraic concepts in Appendix A.2. For the empirical results, all our datasets are derived from synthetic distributions, which are clearly described in Appendix B.1 and B.2. The architectures, implementations (with references to popular base repositories), and hyperparameters (including training procedure) are documented in Appendix B.3. Our open-source data-generating code is available from our project page.

References

Appendix

Below, we list our notational conventions for indices, vectors, matrices, and functions.

For a natural number $n$, $[n]$ denotes the index set $\{1,2,\ldots,n\}$.

We will sometimes index vectors and matrices with named indices (such as $\bot$ for padding tokens) instead of integers, for clarity.

$e_{i}$ denotes the $i$-th elementary (one-hot) unit vector. Likewise as above, we sometimes use non-integer indices (e.g. $e_{\bot}$).

For a function $f:X\times Y\rightarrow Z$ and all $y\in Y$, we will let $f(\cdot,y):X\rightarrow Z$ denote the restriction of $f$ to $y$ (and similarly for other restrictions). This appears in the per-input state transition functions $\delta(\cdot,\sigma):Q\rightarrow Q$, as well as the functions represented by neural networks for a particular choice of weights.

For functions $f,g$, $f\circ g$ denotes composition: $(f\circ g)(x):=f(g(x))$. When we compose neural networks $f:X\times\Theta_{f}\rightarrow Y,g:Y\times\Theta_{g}\rightarrow Z$ with parameter spaces $\Theta_{f},\Theta_{g}$, we will use $f\circ g:X\times(\Theta_{f}\times\Theta_{g})\rightarrow Z$ to indicate the composition $f(g(x;\theta_{g})\theta_{f})$.

A.2 Automata, semigroups, and groups

Recall that a semiautomaton ${\mathcal{A}}=(Q,\Sigma,\delta)$ has a state space $Q$, an input alphabet $\Sigma$, and a transition function $\delta:Q\times\Sigma\to Q$. For any natural number $T$ and a starting state $q_{0}\in Q$, by repeated composition of the transition function $\delta$, one can use ${\mathcal{A}}$ to define a map from a sequence of inputs $(\sigma_{1},\ldots,\sigma_{T})\in\Sigma^{T}$ to a sequence of states $(q_{1},\ldots,q_{T})\in Q^{T}$ via:

Here and below, it is helpful to use a matrix-vector notation to express the computation of semiautomata. For a given semiautomaton we can always identify the state space $Q$ with index set $\{1,\ldots,|Q|\}$ and use a one-hot encoding of states into $\{0,1\}^{|Q|}$. For each input symbol $\sigma\in\Sigma$, we associate a transition matrix $\delta(\cdot,\sigma)\in\{0,1\}^{|Q|\times|Q|}$ with entries $[\delta(\cdot,\sigma)]_{q^{\prime},q}=\mathbf{1}\{\delta(q,\sigma)=q^{\prime}\}$. This implies that for all $q,\sigma$, we have $e_{\delta(q,\sigma)}=\delta(\cdot,\sigma)e_{q}$, so that the computation of the semiautomaton amounts to repeated matrix-vector multiplication.

While semiautomata are remarkably expressive, we discuss a few simple examples throughout this background section to elucidate the key concepts.

Let $Q=\Sigma=\{0,1\}$ and let $\delta(q,0)=q$ and $\delta(q,1)=1-q$. Then, starting with $q_{0}=0$, the state at time $t$, $q_{t}$, is $1$ if the binary sequence $(\sigma_{1},\ldots,\sigma_{t})$ has an odd number of $1$s.

Let $Q=\{1,2\},\Sigma=\{\perp,1,2\}$ and let $\delta$ be given by

As the name suggests, this semiautomaton implements a simple memory operation where the state at time $t$ is the value of the most recent non-$\perp$ input symbol.

Let $S$ be a natural number, $Q=\{0,1,\ldots,S\}$ and $\Sigma=\{L,\perp,R\}$. Then the transition matrices are given by:

This semiautomaton describes the movement of an agent along a line segment where actions $-1$ and $+1$ correspond to decrementing and incrementing the state respectively, except that the decrement input has no effect at state and the increment input has no effect at state $S$.

Note that we have chosen a convention which differs slightly from the main paper (i.e. Figure 1): we enumerate the indices starting from rather than $1$. This is because the proofs are stated more naturally when the boundaries of the gridworld are identified with the indices and $S$.

For a semiautomaton ${\mathcal{A}}=(Q,\Sigma,\delta)$ each input symbol $\sigma\in\Sigma$ defines a function $\delta(\cdot,\sigma):Q\to Q$. These functions can be composed in the standard way, and we use $\delta(\cdot,\sigma_{1:t})$ to denote the $t$-fold function composition. Note that $\delta(q_{0},\sigma_{1:t})$ is precisely the value of the state at time $t$ on input $\sigma_{1:t}$. Thus, the set of all functions that can be obtained by composition of the transition operator, formally

plays a central role in describing the computation of the semiautomaton. This object is a transformation semigroup. We now turn to describing the necessary algebraic background.

Recall that a group $({\mathcal{G}},\cdot)$ is a set ${\mathcal{G}}$ equipped with a binary operation $\cdot:{\mathcal{G}}\times{\mathcal{G}}\to{\mathcal{G}}$ such that

(identity) There exists an identity element $e\in{\mathcal{G}}$ such that $e\cdot g=g\cdot e=g$ for all $g\in{\mathcal{G}}$.

(invertibility) Every element $g\in{\mathcal{G}}$ has an inverse $g^{-1}\in{\mathcal{G}}$ such that $g\cdot g^{-1}=g^{-1}\cdot g=e$

(associativity) The binary operation is associative: $(g_{1}\cdot g_{2})\cdot g_{3}=g_{1}\cdot(g_{2}\cdot g_{3})$.

A monoid is less structured than a group; there must be an identity element and the binary operation must be associative, but invertibility is relaxed. A semigroup is even less structured: the only requirement is that the binary operation is associative.

It is common to let ${\mathcal{G}}$ be a subset of functions from $Q\to Q$ where $Q$ is some ground set and let the binary operation be function composition. In this case, the structure is called a permutation group or transformation monoid/semigroup depending on which subset of the above properties hold. For transformation groups, since every element has an inverse under function composition, it is immediate that every element is some permutation over the ground set.

In fact, taking ${\mathcal{G}}$ to be a subset of functions as above is without loss of generality: by Cayley’s theorem every group is isomorphic (equivalent after renaming elements) to a transformation group on some ground set, and we can take the ground set to have the same number of elements as the original group (for finite groups). Analogously, all semigroups are isomorphic to a transformation semigroup, but the ground set may need one additional element (for the identity); this is Cayley’s theorem for semigroups. It is also clear that every transformation semigroup can be realized by some semiautomaton by trivially having the input symbols correspond to the functions in ${\mathcal{G}}$.More succinctly, inputs can correspond to a generating set of the group, but this is not relevant for our results. Therefore we have lost no structure when passing from finite semiautomata to finite semigroups.

Before discussing the compositional structure of semigroups, we give one more canonical example.

Let $S$ be a natural number let $Q=\{0,1,\ldots,S-1\}$ and let $\Sigma=\{1\}$ have only one element. The dynamics are given by $\delta(q,1)=(q+1)\mod S$. Clearly this semiautomaton implements counting modulo $S$. The underlying group is the cyclic group, denoted $C_{S}$, which is isomorphic to the integers mod $S$ with addition as the binary operation. Note that in this case, the operation is commutative, which makes the group abelian.

The most natural way to compose larger groups from smaller ones is via the direct product. Given two groups $G$ and $H$, we can form a new group with elements $\{(g,h):g\in G,h\in H\}$ with a binary operation that is applied component-wise $(g,h)\cdot(g^{\prime},h^{\prime})=(g\cdot g^{\prime},h\cdot h^{\prime})$ (here, $\cdot$ is overloaded to be the group operation for all three groups). This direct product group is denoted $G\times H$. In the context of permutation groups, say $G$ is a permutation group over ground set $Q_{G}$ and $H$ is over ground set $Q_{H}$. Then $G\times H$ has ground set $Q_{G}\times Q_{H}$ and every function in $G\times H$ factorizes component-wise, i.e., every element in $G\times H$ is identified with a permutation $(q_{G},q_{H})\mapsto(g(q_{G}),h(q_{H}))$ where $g\in G,h\in H$.

Observe that $G\times H$ contains normal subgroups which are isomorphic to both $G$ and $H$. To see this, take $N=\{(e_{G},h):h\in H\}$ where $e_{G}$ is the identity element in $G$. Then since $ge_{G}=e_{G}g$ and since $H$ is closed under its group operation, we have $(g,h)N=N(g,h)$ for all $(g,h)\in G\times H$. A symmetric argument shows that $G$ is also a normal subgroup of the direct product.

Note that we can analogously define direct products in the absence of the group axioms, and thus for monoids and semigroups. This gives a natural construction of the semigroup corresponding to moving around both axes of a 2-dimensional rectangular gridworld, as a concatenation of two non-interacting 1-dimensional gridworlds:

If $G_{S}$ is the transformation semigroup of the 1-d grid world with $S+1$ states, then $G_{S}\times G_{S}$ corresponds to a 2-dimensional gridworld. A semiautomaton that yields this transformation semigroup has state space $Q=\{(i,j):i,j\in\{0,\ldots,S\}\}$ and 5 actions: increment or decrement $i$ or $j$, subject to boundary effects, or do nothing.

The definition of direct product extends straightforwardly to more than two terms $G_{1}\times G_{2}\times\ldots\times G_{n}$; we identify the items with tuples $(g_{1},g_{2},\ldots,g_{n})$.

However, it is possible to compose larger groups so that one of the subgroups is not a normal subgroup. This operation is called a semidirect product, with the group law $(g,h)\cdot(g^{\prime},h^{\prime})=(g\cdot\phi_{h}(g^{\prime}),h\cdot h^{\prime})$ for some $\phi_{h}$ to be defined later. Observe that in the direct product $G\times H$, we have constructed the elements from ordered pairs $(g\in G,h\in H)$, lifting $G$ and $H$ into a shared product space (i.e., the Cartesian product of the underlying sets of $G$ and $H$), defining the group operation as simply applying those of $G$ and $H$ separately.

In fact, there are other ways, to define the group operation in the product space, but a difficulty arises: we need to find other nontrivial multiplication rules on pairs $(g,h)$, and we cannot take for granted that an arbitrary binary operation satisfies the group axioms. We would like to define other operations $(g,h)\cdot(g^{\prime},h^{\prime})$ which output an element of $g$ and an element of $h$. An attempt would be to pick two arbitrary injective homomorphisms $\phi_{G},\phi_{H}$ which embed $G$ and $H$ into a “shared space,” so that elements of $G$ and $H$ can be multiplied together:

However, the middle equality may not hold, because $\phi_{G}(g^{\prime})$ and $\phi_{H}(h)$ are not guaranteed to commute. (Observe that for the special case of $g\mapsto(g,e_{H}),h\mapsto(e_{G},h)$, these two elements always commute, giving rise to the direct product.)

which is of the form $\phi_{G}(\cdot)\cdot\phi_{H}(\cdot)$ since both $G$ and $H$ are themselves closed. This condition is precisely that $G$ is a normal subgroup.

This object is the semidirect product, and it is denoted $G\rtimes H$. Note that the choice of mapping $\phi$ is unspecified in the notation, and, in general, different choices of $\phi$ will yield different structures for the semidirect product.

Finally, when $G=N\rtimes H$, both $N$ and $H$ are subgroups of $G$, but $N$ is also a normal subgroup. To see this, we need to check that $hN=Nh$ for any $h\in H$. This is equivalent to $hnh^{-1}\in N$ for each $h,n$, but we defined the group operation to be $hnh^{-1}=\phi_{h}(n)\in N$, specifically so this would hold. On the other hand, $H$ may not be a normal subgroup, and in this sense the semidirect product is a generalization of the direct product (for which both subgroups are normal). However, when the mapping $\phi$ is trivial, that is $\phi_{h}(n)=n$ then both $N$ and $H$ are normal subgroups, and one can verify that in this case the semidirect product and direct product coincide.

The transformation semigroup for this semiautomaton is $C_{S}\rtimes C_{2}$ where $C_{S}$ is the cyclic group on $S$ elements (cf. Example 4). $C_{2}$ has two elements, the identity $e$ and one element $h$ such that $hh=e$. $C_{S}$ has $S$ elements where each element $g$ is a function that adds some number $k\in\{0,\ldots,S-1\}$ to the input modulo $S$. The inverse $g^{-1}$ is naturally to subtract $k$ to the input, modulo $S$. The homomorphism $\phi$ in the semidirect product is such that $\phi_{e}(g)=g$ and $\phi_{h}(g)=g^{-1}$.

We define one more type of product between groups $N$ and $H$: the wreath product $N\wr H:=(N\times\ldots\times N)\rtimes H$. This is a group containing $|N|^{|H|}\cdot|H|$ elements (rather than $|N|\cdot|H|$, like the direct and semidirect products). Intuitively, it is defined by creating one copy of $N$ per element in $H$ via the direct product, then letting $H$ specify a way to exchange these copies. Formally, $N\wr H$ is the unique group generated by

where we have enumerated the elements of $H$ in arbitrary order, such that each $\pi_{h}:[H]\rightarrow[H]$ is the permutation defined by right multiplication $h^{\prime}\mapsto h^{\prime}h$ (by convention).

A naive way to construct the Rubik’s Cube is to assign labels $\{1,\ldots,54\}$ to the stickers on the cube, and define the Rubik’s Cube group $G$ via the sticker configurations reachable by the $6$ face turns (which each specify a permutation $\delta_{L},\delta_{R},\delta_{U},\delta_{D},\delta_{B},\delta_{F}:\rightarrow$ of the stickers). This establishes $G$ as a subgroup of $S_{54}$. First, notice that the 6 central stickers never move (so this is really improvable to $S_{48}$). Next, notice that the $24=8\times 3$ vertex stickers never switch places with the $24=12\times 2$ edge stickers. The vertex stickers form a subset of the wreath product $C_{3}\wr S_{8}$, while the edge stickers form a subset of the wreath product $C_{2}\wr S_{12}$. In all, this realizes $G$ as a subgroup of a direct product of wreath products:

Among other consequences towards solving the Rubik’s Cube, this gives an improved upper bound on the size of $G$ (which turns out to still be off by a factor of 12, because of nontrivial invariants preserved by the face rotations, a.k.a. unreachable configurations).

When $N$ is a normal subgroup of $G$, the quotient group $G/N$ is defined as $\{gN:g\in G\}$ with binary operation $(gN)(g^{\prime}N)=(gg^{\prime})N$. The fact that $N$ is a normal subgroup implies that this is a well defined group. We can also check that if $G=N\rtimes H$ then the quotient group $G/N$ is isomorphic to $H$, which matches the intuition for multiplication and division.

A $G$ group is simple if it has no non-trivial normal subgroups. Intuitively, a simple group cannot be factorized into components; this generalizes the fact that a prime number admits no non-trivial factorization. When $G$ is not simple then it has a non-trivial normal subgroup, say $N$. We call $N$ proper if $N\neq G$. We call a proper subgroup $N$ maximal if there is no other proper normal subgroup $N^{\prime}\triangleleft G$ such that $N\triangleleft N^{\prime}$. Equivalently, $N$ is a maximal proper normal subgroup if and only if $G/N$ is simple. This is akin to extracting a prime factor from a number, since the quotient group $G/N$ cannot be further factorized. We will revisit this idea of factorization when defining composition series and solvable groups in Section C.3.2.

Finally, we provide some additional terminology related to these different notions of products, which provide a cleaner unifying language in which to state our constructions. Let $N,H$ be arbitrary groups. Which groups $G$ contain a normal subgroup isomorphic $N$, such that the quotient $G/N$ is isomorphic to $H$? Such a group $G$ is said to be an extension of $N$ over $H$. The direct product $G=N\times H$ is known as the trivial extension. A semidirect product $G=N\rtimes H$ is known as a split extension. However, not all extensions are split extensions; the smallest example is the quaternion group $Q_{8}$, the group of unit quaternions $\{\pm 1,\pm i,\pm j,\pm k\}$ under multiplication ($i^{2}=j^{2}=k^{2}=ijk=-1$), which cannot be realized as a semidirect product of smaller groups. In general, it is very hard to derive interesting properties of a group extension based on the properties of $N$ and $H$. Fortunately, there is a characterization of general extensions. The Krasner-Kaloujnine universal embedding theorem (Krasner and Kaloujnine, 1951) states that all extensions $G$ can be found as subgroups of the wreath product $N\wr H$. The proof of Theorem 2 essentially shows how to implement the different kinds of group extensions, given constructions which implement the substructures $N,H$. In the worst case, we will have to implement a wreath product.

A.3 Shallow circuit complexity classes

We provide an extremely abridged selection of relevant concepts in circuit complexity. For a systematic introduction, refer to (Arora and Barak, 2009). In particular, we discuss each circuit complexity class and inclusion below:

$\mathsf{NC}^{0}$ is the class of constant-depth, constant-fan-in, polynomial-sized $\mathsf{AND}$/$\mathsf{OR}$/$\mathsf{NOT}$ circuits. If a constant-depth Transformer only uses the constant-degree sparse selection constructions in (Edelman et al., 2022), it can be viewed as representing functions in this class. However, the representational power of these circuits is extremely limited: they cannot express any function which depend on a number of inputs growing with $T$.

$\mathsf{AC}^{0}$ is the class of constant-depth, unbounded-fan-in, polynomial-sized $\mathsf{AND}$/$\mathsf{OR}$ circuits, allowing $\mathsf{NOT}$ gates only at the inputs. A classic result is that the parity of $T$ bits is not in $\mathsf{AC}^{0}$ (Furst et al., 1984); Hahn (2020) concludes the same for bounded-norm (and thus bounded-Lipschitz-constant) constant-depth Transformers.

$\mathsf{ACC}^{0}$ extends $\mathsf{AC}^{0}$ with an additional type of unbounded-fan-in gate known as $\mathsf{MOD}_{m}$ for arbitrary number $m$, which checks if the sum of the input bits is a multiple of $m$. Theorem 2 comes from the fact that the semigroup word problem (which is essentially identical to semiautomaton simulation) is in this class; see (Barrington and Thérien, 1988).

$\mathsf{TC}^{0}$ extends $\mathsf{AC}^{0}$ with an additional type of unbounded-fan-in gate called $\mathsf{MAJ}$, which computes the majority of an odd number of input bits (a threshold gate). It is straightforward to simulate modular counters using a polynomial number of parallel thresholds (i.e. $\mathsf{ACC}^{0}\subseteq\mathsf{TC}^{0}$). Whether this inclusion is strict (can you simulate a threshold in constant depth with modular counters?) is a salient open problem in circuit complexity. Threshold circuits are a very natural model for objects of interest in machine learning like decision trees and neural networks (Merrill et al., 2021).

$\mathsf{NC}^{1}$ is the class of $O(\log T)$-depth, constant-fan-in, polynomial-sized $\mathsf{AND}$/$\mathsf{OR}$/$\mathsf{NOT}$ circuits. It is an extremely popular and natural complexity class capturing efficiently parallelizable algorithms. It is unknown whether any of the inclusions in the “larger” classes $\mathsf{TC}^{0}\subseteq\mathsf{NC}^{1}\subseteq\mathsf{L}\subseteq\mathsf{P}$ are strict.

A.4 The Transformer architecture

In this section, we define the Transformer function class used in our theoretical results, and discuss remaining discrepancies with the true architecture.

the causally-masked softmax at row $t$ is defined to be $\mathsf{softmax}(z_{1:t})$ on the first $t$ coordinates, and 0 on the rest. To implement the causal masking operation, it is customary to set the entries above the diagonal of the attention score matrix $XW_{Q}W_{K}^{\top}X^{\top}$ to $-\infty$, then obtaining $\mathsf{CausalAttn}(XW_{Q}W_{K}^{\top}X^{\top})$ via a row-wise softmax (letting $e^{-\infty}$ evaluate to ).

In general, for any positive integer $H$, a multi-headed self-attention block consists of a sum of $H$ copies of the above construction, each with its own parameters.

This component is often called soft attention: the softmax performs continuous selection, taking a convex combination of its inputs. In contrast, hard attention refers to attention heads which perform truly sparse selection (putting weight $1$ on the position with the highest score, and on all others).

At the level of granularity of the results in this paper (up to negligible changes in the width and weight norms), this changes very little from the viewpoint of representation. A residual connection can be implemented (or negated) by appending two ReLU activations to a non-residual network:

Similarly, a residual connection can be implemented with one attention head (with internal embedding dimension $k=d$), as long as it is able to select its own position (which will be true in all of our constructions).

For simplicity of presentation, we omit the normalization layers which are usually present after each attention and MLP block. It would be straightforward (but an unnecessary complication) to modify the function approximation gadgets in our constructions to operate with unit-norm embeddings.

Finally, it will greatly simplify the constructions to add padding tokens: to simulate a semiautomaton at length $T$, we will choose to prepend $\tau$ tokens, with explicitly chosen embeddings, which do not depend on the input $\sigma_{1:T}$. Theorem 1 uses $\tau=\Theta(T)$ padding, and Theorem 2 uses $\tau=1$. In both cases, padding is not strictly necessary (the same functionality could be implemented by the MLPs without substantially changing our results), but we find that it leads to the most intuitive and concise constructions.

We define the following quantities associated with a Transformer network, and briefly outline their connection to familiar concepts in circuit complexity:

The dimensions according to the definition of a sequence-to-sequence network: sequence length $T$ and embedding dimension $d$. Up to a factor of bit precision, this corresponds to the number of inputs in a classical Boolean circuit. We will exclusively define architectures where $d$ is independent of $T$.

A bound on its $\infty$-weight norms: the largest absolute value of any trainable parameter. These can be converted into norm-based generalization bounds via the results in Edelman et al. (2022). Note that the results in this paper go beyond the sparse variable creation constructions of bounded-norm attention heads; in general, the norms scale with $T$. The attention heads express meaningful non-sparse functions. Aside from the positive experimental results, we do not directly investigate generalization in this paper.

A.5 Additional discussion of related work

We first provide references for the “reasoning-like” applications of neural networks mentioned in the main paper.

Program synthesis: (Chen et al., 2021b, Schuster et al., 2021, Li et al., 2022).

Mathematical reasoning: (Lample and Charton, 2019, Polu and Sutskever, 2020, Drori et al., 2022).

Neural dynamics models for decision-making: recurrent (Hafner et al., 2019, Ye et al., 2021, Micheli et al., 2022) and non-recurrent (Chen et al., 2021a, Janner et al., 2021).

We provide an expanded discussion of empirical analyses of neural networks trained on synthetic combinatorial tasks.

Pointer Value Retrieval: Zhang et al. (2021a) propose a benchmark of tasks based on pointer value retrieval (PVR) to study the generalization ability (in-distribution as well as distribution shift) of different neural network architectures. Their key idea behind the task is “indirection through a pointer rule”, that is, a specific position of the input acts as a pointer to the relevant position (window) of the input which contains the answer. Using our results, we can implement a certain sub-class of PVR tasks: (1) we use the first attention layer to identify the pointer, and (2) we use the second attention layer to select the window between the pointer value and the width. (2) is doable with $O(1)$ attention heads if we are computing a function that is based on the sum (for example,$\mod n$). Otherwise it would require the window size number of attention heads similar to our grid-world construction.

LEGO: Zhang et al. (2022) propose a task based on solving a simple chain-of-reasoning problem based on group-based equality constraints. They study the ability of transformers to generalize the entire chain of reasoning given only part of the chain while training. A direct comparison to our setting is not clear since this task is not modelled as a sequence-to-sequence task, however it serves as another example of the emergence of “shortcut” solutions: transformers solve certain variables without resolving the chain of reasoning.

Dyck: Several works (Hahn, 2020, Ebrahimi et al., 2020, Newman et al., 2020, Yao et al., 2021) have studied the ability of Transformers to represent Dyck languages, both for generation and closing bracket prediction. The most closely related to our work is Yao et al. (2021), which constructs a clever depth-2 as well as a depth-$D$ solution for bounded-depth $D$ Dyck languages. Bounded-depth Dyck can be captured by our semiautomata formalism and our main construction would recover the depth-$2^{D}$ solution by default. Their depth-2 construction bears semblance to the constructions we use in Theorem 3: they implement a counter in the first layer similar to our $\mod n$ construction, and implement a proximity-based depth matching in the second layer. Our grid-world construction generalizes their construction to a significantly more complex problem. We view our work as a generalization of their results to a wider class of semiautomata.

Parity: Another commonly studied synthetic setup is the task of learning parities. Edelman et al. (2022), Barak et al. (2022) perform a theoretical and empirical study of the ability of Transformers (and other architectures) to learn sparse parities where the support size $k\ll T$. Bhattamishra et al. (2020), Schwarzschild et al. (2021) study the task of computing prefix sum in the binary basis (which is essentially parity of the prefix sum) for Transformers and recurrent models, repsectively. Anil et al. (2022), Wei et al. (2022) study essentially the same problem however they model the task as a natural language task and use pretrained Transformers.

Modular addition: In the pursuit of understanding grokking, Nanda and Lieberum (2022) focus on the task of adding two 5 digit numbers modulo a large prime ($113$ in their setting). They take the viewpoint of mechanistic interpretability and attempt to reverse engineer what the Transformer is learning on this task in the low sample regime. They claim that the trained model learns sinusoidal encodings that we also use in our theoretical constructions. Note that our setting of modular counters performs a $T$-way summation, while their setting involves only a 2-way summation (with carryover). Inspired by their work, we do some preliminary investigation into interpreting the trained Transformer on the grid world (see Figure 10).

Dyck languages are particularly interesting for their completeness property: the Chomsky-Schützenberger representation theorem (Chomsky and Schützenberger, 1959) states that all context-free languages can be (homomorphically) represented by the intersection of a Dyck language and a regular language. For more on this topic, see the discussion in Yao et al. (2021). In the context of regular languages (which in general induce finite-state automata), our findings imply that $O(\log T)$-depth networks can simulate all context-free languages (Theorem 1), and $O(1)$-depth networks can represent some of them. The obstructing regular languages are the ones whose associated syntactic monoids are non-solvable. We further note that the gridworld semigroups are aperiodic and thus simulable by star-free regular expressions (Schützenberger, 1965) and $\mathsf{AC}^{0}$ circuits (Chandra et al., 1983, Barrington and Thérien, 1988). We did not see a way for this to generically entail $O(1)$-depth shortcuts with self-attention. For the relation between the Chomsky hierarchy and various neural networks in practice, Delétang et al. (2022) provide an extensive empirical study for memory-augmented RNNs and Transformers on tasks spanning all 4 levels of the hierarchy, and conclude the Transformers lack the ability to even recognize regular languages. Their results do not contradict with ours, since they measure performance on “inductive inference”, which is similar to our length generalization setup where we also see the failure of Transformer.

There has been much recent interest in quantifying out of distribution generalization of trained models under distribution shifts that maintain some notion of “logical” invariance. Wei et al. (2022), Anil et al. (2022) empirically investigate the ability of pre-trained Transformers to generalize to longer sequence length for parity-like problems modelled as language tasks. Xu et al. (2020) study size generalization in graph neural networks where they train on small graphs and evaluate on larger sized graphs with similar structural properties. Schwarzschild et al. (2021), Bansal et al. (2022) focus on length and algorithmic generalization for recurrent models where they train on simple/easy instances of the underlying problem and evaluate on harder/complex instances using the power of recurrence to simulate extra computational steps, inspired by the ideas of Neural Turing Machines (Graves et al., 2014) and Adaptive Computation Time (Graves, 2016). We view our results as complementing those of Yao et al. (2021), Anil et al. (2022) for a richer class of problems. Our use of scratchpad is inspired by Nye et al. (2021), Wei et al. (2022), Anil et al. (2022).

Our work is not the first to notice that Transformer architectures make brittle predictions out-of-distribution. Indeed, even the seminal paper introducing the architecture (Vaswani et al., 2017) notes that length generalization is promoted by a subtle hyperparameter choice (namely, the positional encoding scheme). Furthermore, there have been several attempts to reconcile this gap by modifying Transformers to behave more like RNNs; (Dehghani et al., 2019, Nye et al., 2021, Wei et al., 2022, Anil et al., 2022, Hutchins et al., 2022). Kasai et al. (2021) consider training a non-recurrent Transformer, and finetuning it into an RNN. All of these works have some element of natural language experiments: either the task is end-to-end language modeling, or the synthetic reasoning task is framed as a natural language problem, for a pretrain-finetune pipeline. We view our work as strengthening the foundations of these lines of inquiry. Theoretically, we provide structural guarantees for how shallow non-recurrent models can (perhaps deceptively) fit recurrent dynamics over long sequences. Empirically, we perform a pure (no confounds arising from the influence of a natural langauge corpus) analogue of the experiments seeking to help neural networks follow long chains of reasoning.

As mentioned briefly towards the end of Section 5, the setting of indirectly-supervised semiautomata matches that of autoregressive generative modeling (a.k.a. next-token prediction), if the continuations of the sequence depend on the state of a latent semiautomaton. This is the case in (for example) generating Dyck languages (Yao et al., 2021), where the possible continuations are $\{$all possible open brackets, if the stack $q_{t}$ is not full$\}\cup\{$close bracket which pairs with the top of the stack $q_{t}\}$. We note that when an autoregressive model is used for sequence generation via a token-by-token inference procedure, this amounts to a special case of scratchpad inference (with a naive $1$-step training procedure): the constant-depth network is used as a single iteration of a recurrent network, whose state is the completed prefix of the current generated sequence. Non-autoregressive natural language generation and transduction are an exciting area of research (Gu et al., 2017); for a recent survey, see Xiao et al. (2022). Our results are relevant to this line of work, suggesting that there may not be an expressivity barrier to expressing deep recurrent linguistic primitives, but there may be issues with out-of-distribution robustness.

Pérez et al. (2021) show that an infinite-precision Transformer achieves Turing completeness, as a single forward pass through a 3-layer decoder can simulate one transition step of a Turing machine. Giannou et al. (2023) exhibit a 13-layer Transformer whose weights are hard-coded to a universal Turing machine, and can be looped to perform any computation. These works show that one pass through a Transformer can implement a single computational step of a Turing machine. In contrast, our results show how a shallow Transformer can sometimes execute a computational loop over the entire context in a single non-recurrent pass. This requires a significantly more refined analysis, which depends on the global algebraic structure induced by the automata in question.

Another area where tools from abstract algebra are used to reason about neural networks is geometric deep learning, a research program which seeks to understand how to specify inductive biases stemming from algebraic invariances. For a recent survey, see Bronstein et al. (2021). In contrast, this work studies the ability of a fixed architecture to learn a wide variety of algebraic operations, in the absence of special priors (but a large amount of data). There are certainly possible connections (e.g. “how do you bias an architecture to perform operations in a known group, when there is limited data?”) to explore in future work.

Our theoretical results can be interpreted as a depth separation result: contingent on $\mathsf{TC}^{0}\neq\mathsf{NC}^{1}$, it takes strictly more layers to simulate non-solvable semiautomata, compared to their solvable counterparts. In a similar spirit, there have been several works establishing depth separation for feed-forward neural networks (mostly using ReLU activations) (Telgarsky, 2016, Eldan and Shamir, 2016, Daniely, 2017, Lee et al., 2017, Safran et al., 2019). These results are usually constructive in nature, that is, they show the existence of functions that can be represented by depth $L$ but would require exponential width for depth $L-1$ (or $\sqrt{L}$, depending on the result).

More elementary neural architectures, such as MLPs, have the ability to represent arbitrary functions, given sufficiently many neurons (Hornik et al., 1989, Cybenko, 1989). The $\mathsf{ACC}^{0}$ circuit construction described in (Barrington and Thérien, 1988) can be implemented by any such architecture, not just the Transformer– there is a naive black-box way to “compile” each gate in the circuit into a network with the same depth as the $\mathsf{ACC}^{0}$ circuit. A natural question is: why, then, should we prefer Transformers? The primary advantage of Transformers comes from the position-wise weight sharing of the attention layers and the casual structure from causal attention maps. Unlike MLPs, the shared parameters in Transformers allow for significant reduction in the total parameter count for representing the two main operations across positions: modular counters, and resets. In particular, these functions can be represented with $O(1)$ trainable parameters in the attention and MLP weight matrices (i.e. independent of $T$), as opposed to the $\Theta(T)$ parameters in a vanilla MLP, where position-wise parameter sharing is not available. In a sense, the Transformer architecture is naturally suited for implementing this construction.

Appendix B Experiments

This section contains a full description and discussion of the in-distribution simulation experiments from Section 4.

This defines a supervised learning problem over sequences.

Note that without intermediate states in the input these problems exhibit long-range dependencies: for example, in the parity semiautomaton (and for any semiautomaton whose transformation semigroup is a group), every $q_{t}$ depends on every preceding input $\{\sigma_{t^{\prime}}:t^{\prime}<t\}$. Indeed, this is why previous studies have used group operations as a benchmark for reasoning (Anil et al., 2022, Zhang et al., 2022).

We proceed to enumerate the semiautomata considered in these simulation experiments.

Cyclic groups $C_{2},C_{3},\ldots,C_{8}$. For each cyclic group $C_{n}$ (realized as $Q:=\{0,1,\ldots,n-1\}$ under mod-$n$ addition), we choose the generator set $\Sigma$ to be the full set of group elements $\{0,\ldots,n-1\}$. An alternative could be to let $\Sigma$ be a minimalIn the sense that it induces a non-trivial learning problem on this group.If we only pick the generator $\{1\}$, the output sequence is deterministic, and there is no learning problem. set $\{0,1\}$, which we do not use in the experiments.

Direct products of cyclic groups $C_{2}\times C_{2},C_{2}\times C_{2}\times C_{2}$, realized as concatenated copies of the component semiautomata. Note that $C_{6}$ (which is isomorphic to $C_{2}\times C_{3}$), included in the above set, is another example.

Dihedral groups $D_{6},D_{8}$. Our realization of $D_{2n}$ chooses $Q=\{0,1,\ldots,n-1\}\times\{0,1\}$ and $\Sigma=\{(1,0),(0,1)\}$. Since these groups are non-abelian, it is already not so straightforward (compared to parity) to see why constant-depth shortcuts should exist.

Permutation groups $A_{4},S_{4},A_{5},S_{5}$. We choose $Q$ to be the set of $n!$ permutations for $S_{n}$ (symmetric group), and $Q$ to be the set of $\frac{n!}{2}$ even permutations for $A_{n}$ (alternating group on $n$ elements). The generator set for $S_{n}$ consists of the minimal generators, a transposition and an $n$-cycle, as well as 6 other permutations. These other permutations are chosen following the ordering given by the $\mathsf{sympy.combinatorics}$ package. They are not necessary for covering the state space (since the minimal set of 2 permutations already suffice to cover $Q$), but can help speed up the mixing of the states. For $A_{n}$, we choose the $3$-cycles of the form $(12i)$ for $i\in\{3,4,\cdots,n\}$. Note that $A_{4},S_{4}$ are solvable (leading to constant-depth shortcuts), while $A_{5},S_{5}$ are not. Also, note that to learn a constant-depth shortcut for $A_{4}$, a model needs to discover the wondrous fact that $A_{4}$ has a nontrivial normal subgroup, that of its double transpositions.

The quaternion group $Q_{8}$. This is the smallest example of a non-abelian solvable group which is not realizable as a semidirect product of smaller groups, thus requiring the full wreath product construction (Lemma 10) in our theory.

We focus on the online learning setting for all experiments in this paper: at training iteration $i$, draw a fresh minibatch of samples from ${\mathcal{D}}_{\mathcal{A}}$, compute the network’s loss and gradients on this minibatch, and update the model’s weights using a standard first-order optimizer (we use AdamW (Loshchilov and Hutter, 2017)). This is to mitigate the orthogonal challenge of overfitting; note that the purpose of these experiments is to determine whether standard gradient-based training finds shortcut solutions in these combinatorial settings (in a reasonable amount of time), not how efficiently. We do not investigate how to improve sample efficiency in this paper. The results in the paper are based on sinusoidal positional encodings (Vaswani et al., 2017) unless otherwise specified.

We report our main results with sequence length $T=100$, which is large enough to rule out memorization: for this choice of $T$, the inputs come from a uniform distribution over $|\Sigma|^{100}>10^{30}$ sequences, rendering it overwhelmingly unlikely for a sample to appear twice between training and evaluation. We observed positive results in most of the settings for larger $T$, but training became prohibitively unstable and computationally expensive; mitigating this is an interesting direction for future empirically-focused studies.

We seek to investigate the sufficient depth for learning to simulate each semiautomaton. Thus, for each problem setting, we vary the number of layers $L$ in the Transformer between $1$ and $16$. Note that we do not attempt in this work to distinguish between depths $O(\log T)$ and $O(1)$, nor do we attempt to tackle the problem of exhaustively enumerating and characterizing the shortcut solutions for any particular semiautomaton.

The minimum number of layers required to achieve 99%+ performance reflects our beliefs on the difficulty of the task: a high-level trend is that the semigroups which don’t contain groups (which only require memory lookups) are the easiest to learn, and among the groups, the larger non-abelian groups require more layers to learn, with the non-solvable group $S_{5}$ requiring the largest depth.However, we stress that these experiments do not control for the fact that larger groups have richer supervision (for example, $A_{5}$ has more informative labels than $A_{4}$), possibly accounting for the counterintuitive result that the latter requires more layers, despite being a subgroup of the former. Between the non-abelian groups, the difficulty of learning $Q_{8}$ compared to $D_{8}$ (which has the same cardinality) agrees with our theoretical characterizations of the respective constant-depth shortcuts for these groups: $D_{8}$ can be written as a semidirect product of smaller groups, while $Q_{8}$ cannot, so our theoretical construction of a constant-depth shortcut must embed $Q_{8}$ in a larger structure (i.e. the wreath product).

Throughout these experiments, we observe the following forms of training instability: high variance in training curves (based on initialization and random seeds for the gradient-based optimization algorithm), and negative progress (i.e. non-monotonic loss curves), even for training runs which eventually converge successfully. This is evident in Figure 3(b),(c) and in the significant difference between the maximum accuracies in Figure 8 and the median in Figure 9.

To stabilize training, we experiment with dropout and exponential moving average (EMA)We use the EMA implementation from https://github.com/fadel/pytorch_ema.. The effectiveness of dropout varies across datasets; for example, we find using a dropout of 0.1 (the best among $\{0,0.1,0.2,0.3\}$) to be helpful for Dihedral and Quaternion, while such dropout hurts the training of Dyck and Gridworld. We find EMA to be generally useful, and fix the decay parameter $\gamma=0.9$ in the experiments since the performance of the EMA model does not seem to be sensitive to the choice of $\gamma\in\{0.85,0.9,0.95\}$. Further, increasing the patience of the learning rate scheduler can be helpful.

B.2 Section 5: Failures of shortcuts in more challenging settings

Our theoretical and main empirical findings have shown that not only do shallow non-recurrent networks subsume deeper finite-state recurrent models in theory, these shallow solutions can also be found empirically via standard gradient-based training. However, experiments in Section 4 and Appendix B are in an idealized setting, with full state supervision during training and in-distribution evaluation at test time. This section studies more challenging settings where these assumptions are relaxed. We consider training under indirect (Section B.2.1) or incomplete (Section B.2.2) state supervision, and evaluation on sequences that is out-of-distribution (Section B.2.3) or of longer lengths (Section B.2.4).

Permutations with single-element observations: We take the permutation group $S_{5}$ with $Q$ is the set of $5!$ operations. The observation function $\varphi:Q\rightarrow\{1,2,3,4,5\}$ returns the first value of the permutation. For example, $\varphi((2,1,4,3,5)=2$. We use a set of 5 generators for the experiments.

We consider two distributions on the input sequences: (1) the input is always a sequence of the form $abababa\cdots$ (i.e. the process is never in the absorbing state), which is the setup in Bhattamishra et al. (2020); and (2) the input is of the form $abababa\cdots$ with probability 0.5, and is some randomly drawn string of $a,b$ otherwise. Note that case (1) can be solved purely based on the positional encoding, since the label is 1 when the position is a multiple of 4 and 0 otherwise, while case (2) is more difficult since the model needs to take into account the input tokens.

We train GPT-2-like models on sequences of length 40. We use 16 layers for $S_{5}$ and 8 layers for other tasks, with embedding dimension $d=512$ and $H=8$ attention heads. As shown in Figure 11, the model is able to achieve near-perfect in-distribution accuracies for all tasks. An interesting side finding is that the choice of positional encoding turns out to be important for both cases of $(abab)^{*}$: learning is challenging for linear encoding (i.e. $p_{i}\propto i$) but is easy when using sinusoidal positional encoding, which is likely because the sinusoidal encoding naturally matches the periodicity in $(abab)^{*}$. In all other experiments, we use sinusoidal positional encodings unless otherwise noted.

Another challenge of limited supervision is that the observation sequence may be incomplete, that is, we may not be able to get supervision on the states at every time step. We consider the task of learning length 100 sequences, where the state at each position is revealed with some probability $p_{\mathsf{reveal}}\in(0,1]$.

The previous subsections show positive results on learning shallow non-recurrent shortcuts with limited supervision during training, either in the form of indirect observations or incomplete observation sequences. In this section, we study challenges at test time, and evaluate Transformers on their out-of-distribution generalization performance. For this and the next subsection, the models are trained in the standard way with full state supervision. The training sequences are of length 40, where each position has an equal probability of being 0 or 1, i.e. $\mathop{\bf Pr\/}[\sigma=1]=0.5$. At test time, the sequences of the same length as training, but the Bernoulli parameter $\mathop{\bf Pr\/}[\sigma=1]$ varies in the range $\{0.05,0.1,0.15,\ldots,0.9,0.95\}$.

Figure 13 (left) shows the accuracy as $\mathop{\bf Pr\/}[\sigma=1]$ varies. The performance of the Transformer degrades sharply as the test distribution changes away from training, failing at out-of-distribution generalization. Given the theoretical construction of modular counters (Lemma 6), our hypothesis is that Transformer may be learning a shortcut solution that computes the parity by counting the number of 1s, and that counts less frequently seen during training will cause the model to fail. The experimental results agree with the hypothesis: as $\mathop{\bf Pr\/}[\sigma=1]$ deviates from 0.5, it is less likely for the value of the count (which concentrates around $T\times\mathop{\bf Pr\/}[\sigma=1]$) to be seen during training, hence the performance degrades. In contrast, an LSTM recurrent network maintains perfect accuracy when evaluated on all values of $\mathop{\bf Pr\/}[\sigma=1]$.

We further test this hypothesis by checking how the accuracy changes as we vary the count (i.e. the number of 1s) in the input sequence. As shown in Figure 13 (right), Transformer’s performance degrades as the count moves away from the expected number during training, agreeing with the hypothesis. It might appear strange that GPT fails at a lower count more than a higher count. However, this may be because the shortcut learns a correlation between the count and the position: during training, a lower count is more likely to appear early in an input sequence, as opposed to the testing scenario where a lower count is equally likely to appear at a later part of an sequence. This is further supported by the observation that training the model with randomly shifted positions significantly improves the performance at lower counts.

We investigate one established mitigation for the out-of-distribution brittleness of non-recurrent Transformers: scratchpad training and inference. Given a sequence of inputs $(\sigma_{1},\ldots,\sigma_{T})$ and states $(q_{1},\ldots,q_{T})$, in the standard (non-recurrent) sequence-to-sequence learning pipeline, the network receives $\sigma_{1:T}$ as input, and outputs the sequence of predictions for $q_{t}$. In scratchpad training (Nye et al., 2021, Wei et al., 2022), we instead feed the network an interleaved sequence of inputs and states $(\sigma_{1},q_{1},\sigma_{2},q_{2},\sigma_{3},q_{3},\ldots,q_{T-1},\sigma_{T})$ (with an appropriately expanded token vocabulary), and define the network’s state predictions to be those at the appropriately aligned positions: $(\hat{q}_{1},\bot,\hat{q}_{2},\bot,\ldots,\bot,\hat{q}_{T})$ (where $\bot$ denotes a position where the prediction is ignored by the loss function). During inference, we iteratively fill in the state predictions. This removes the need for the network to learn long-range dependencies in a single non-recurrent pass, by splitting it into $T$ sequential state prediction problems which can depend on previous predicted state $\hat{q}_{t-1}$; one can think of this as a way to guide a shallow Transformer to learn the recurrent solution (i.e. explicit depth-$\Theta(T)$ iteration of the state transition function), rather than a shortcut.

We note that introducing the scratchpad itself is not sufficient to remove the parallel solution, since the model can simply ignore the scratchpad positions and find the same parallel shortcut as before. The good news is that we can couple scratchpad with an explicit recency bias in the attention mechanism (Press et al., 2022) which biases the model towards putting more attention weights on closer input. Intuitively, if the model is only allowed to put attention on the current input token and the current scratchpad (which is simply the current state), then the model is forced to be recurrent; recency bias can be considered as a soft relaxation of the same idea. Combining scratchpad and recency bias, we are able to train a Transformer to learn the recurrent solution, which is resilient to distribution shift; see Figure 13 (left). Notice that this mitigation completely foregoes the computational advantage of a shallow shortcut; we leave it to future work to obtain shortcuts which are resilient to distribution shift. Towards this, the constructions used in the proof of Theorem 1 may be helpful. Finally as a side note, even though the state transitions are Markov, the dependency in the input sequence can still be long range, so we do not expect recency bias to help without scratchpad, since in this case the output can depend uniformly on each input positions (e.g. consider parity).

Figure 14 shows the performance on sequences of various lengths. In contrast to LSTM’s perfect performance on all scenarios, Transformer’s accuracy drops sharply as we move to lengths unseen during training. This is not purely due to unseen values of the positional encoding: randomly shifting the positions during training can cover all the positions seen during testing, which helps improve the length generalization performance but cannot make it perfect; we see similar results for removing positional encodings altogether. However, similar to the OOD setup in the previous subsection, we empirically show that the above flaws are circumventable. Using a combination of scratchpad (a.k.a. “chain-of-thought”) (Nye et al., 2021, Wei et al., 2022) and recency bias (Press et al., 2022), we demonstrate that Transformers can be guided towards learning recurrent (depth-$T$) solutions, which generalize out-of-distribution and to longer sequence lengths (Figure 14, yellow curves). The results also confirm that the inclusion of recency bias is necessary: without it, scratchpad training shows no improvement on length generalization.

Figure 14 also shows some interesting findings related to positional encoding, which is believed to be a key component for Transformers and a topic with active research (Ke et al., 2020, Chu et al., 2021). While this work does not aim to improve positional encoding, some of our results may be of interest for future research.

Training with shifted positions: In general, unseen positions appear to be a major contributor to Transformer’s failure of length generalization. This is evidenced by the comparison between Transformer trained with absolute positional encoding, and Transformers trained with random shifts added to the positional encoding: for each batch, we sample a random positive integer in and add it to the position indices before calculating the positional encoding; this random integer is the same for each batch and varies across batches. Figure 14 shows that adding such random shifts gives a significant boost to Transformer’s length generalization performance, for both Dyck and $C_{2}$. This suggests that a main challenge to length generalization is the distribution shifts due to positions unseen during training, and finding better positional encoding could be a potential remedy for poor length generalization.

As a side note, we also find that removing positional encoding altogether helps improve generalization for both parity and Dyck. For the former, removing positional encodings makes sense since parity is a symmetric function where the ordering of the arguments does not matter, Empirically, we are able to achieve non-trivial accuracy (even when evaluated at the sequence level) without positional encoding, whereas Bhattamishra et al. (2020) reports 0 accuracy. The discrepancy may be due to different model size: Bhattamishra et al. considers Transformers with up to 4 layers, 4 heads and dimension up to 32, whereas for the parity experiments we consider Transformers with 8 layers, 8 heads, and dimension 512. though the positive result for Dyck is less clearly understood. Note that removing positional encoding does not mean having no position information, since the use of the causal mask implicitly encodes the position, which is also noted in Bhattamishra et al. (2020) and concurrent work by Haviv et al. (2022). Understanding this phenomenon is tangential to the current work and is left to future work.

B.3 Additional details

For GPT-2 models, we fix the embedding dimension and MLP width to 512 and the number of heads to 8 in all experiments in Section 4, and vary the number of layers from 1 to 16. For LSTM, we fix the embedding dimension to 64, the hidden dimension to 128, and the number of layers to 1. We use the AdamW optimizer (Loshchilov and Hutter, 2017), with learning rate in {3e-5, 1e-4, 3e-4} for GPT-2 or {1e-3, 3e-3} for LSTM, weight decay 1e-4 for GPT-2 or 1e-9 for LSTM, and batch size 16 for GPT-2 or 64 for LSTM. As detailed in Section B.1.1, the models are trained in an online fashion with freshly drawn samples in each batch. The number of freshly drawn samples ranges from 600k to 5000k for different datasets, which is much fewer than the number of possible strings of length 100.

Our experiments are implemented with PyTorch (Paszke et al., 2019). The Transformers architectures are taken from the HuggingFace Transformers library (Wolf et al., 2019), using the GPT-2 configuration as a base. The LSTM architecture is the default one provided by the PyTorch library.

The experiments were performed on an internal cluster with NVIDIA Tesla P40, P100, V100, and A100 GPUs. For the experiments in Section 4, each training run took up to $10$ hours on a single GPU, for a total of $\approx 10^{4}$ GPU hours. The remaining experiments in Section 5 amount to less than $1\%$ of this expenditure.

Appendix C Proofs

We first recall the notions of simulation introduced in Section 2:

A function can simulate an automaton for particular choices of $T,q_{0}$. For a semiautomaton ${\mathcal{A}}=(Q,\Sigma,\delta)$, a function $f:\Sigma^{T}\rightarrow Q^{T}$ simulates ${\mathcal{A}}_{T,q_{0}}$ if $f(\sigma_{1:T})={\mathcal{A}}_{T,q_{0}}(\sigma_{1:T})$ for all input sequences $\sigma_{1:T}$. Here, the right-hand side denotes the sequence of states $q_{1:T}$ induced by the input sequence $\sigma_{1:T}$ under the transitions $\delta$ starting from state $q_{0}$.

A function class can simulate multiple functions associated with a semiautomaton. For a semiautomaton ${\mathcal{A}}=(Q,\Sigma,\delta)$ and a positive integer $T$, a function class ${\mathcal{F}}$ (a set of functions $f:\Sigma^{T}\rightarrow Q^{T}$) simulates ${\mathcal{A}}$ at length $T$ if, for every $q_{0}\in Q$, there is function $f_{q_{0}}\in{\mathcal{F}}$ which simulates ${\mathcal{A}}_{T,q_{0}}$.

When $W$ is a linear threshold function $z\mapsto\operatorname*{arg\,max}_{q}[Wz]_{q}$, this corresponds to a standard classification head. However, our constructions may leverage other encodings of discrete objects.

We provide some simple function approximation results below.

The inner dimension is $d^{\prime}=4|{\mathcal{X}}|$, and the weights satisfy

For each $x_{0}\in{\mathcal{X}}$, we construct an indicator $\psi_{x_{0}}(x)$ for $x_{0}$, out of 4 ReLU units. Letting $\Delta^{\prime}:=\Delta/4$, the construction is

The second layer simply sums these indicators, weighted by each $f(x_{0})$. ∎

Letting ${\mathcal{X}}_{i}$ denote the set of unique values in coordinate $i$, the inner MLP dimensions are as follows:

When we apply Lemmas 1 and 2 in recursive constructions, and $B_{x}/\Delta\geq 1$, we will opt to use the bound $\left\|b_{1}\right\|_{\infty}\leq 6B_{x}/\Delta$, to reduce the clutter of propagating the $2$ term without resorting to asymptotic notation.

The inner dimension is $d^{\prime}=2$, and the weights satisfy

We construct the threshold using 2 ReLU units. The construction is

We record some useful lemmas pertaining to approximating hard coordinate selection with soft attention. The following is a simplified version of Lemma B.7 from (Edelman et al., 2022) (which generalizes this to multi-index selection):

Let $t^{*}:=\operatorname*{arg\,max}_{t}z_{t}$. Suppose that for all $t^{\prime}\neq t^{*}$, $z_{t^{\prime}}\leq z_{t^{*}}-\gamma$. Then,

Without loss of generality, $\max z=\gamma$ (since the softmax function is invariant under shifting all inputs by the same value), so that all other coordinates are non-positive. Also, assume $T\geq 2$ (the $T=1$ case is trivial). We have

Thus, the 1-norm of the difference is bounded by

We note the following elementary fact about 2-dimensional circular embeddings.

Consider $p_{1},\ldots,p_{T}$, the $T$ equally-spaced points on the $2$-dimensional circle:

C.2 Proof of Theorem 1: Logarithmic-depth shortcuts via parallel prefix sum

In this section, we give the full statement and proof of the universal existence of logarithmic-depth shortcuts.

Let $\mathcal{A}=(Q,\Sigma,\delta)$ be a semiautomaton, $q_{0}\in Q$, and $T\geq 1$. Then, there is a depth-$\lceil\log_{2}T\rceil$ Transformer which continuously simulates ${\mathcal{A}}_{T,q_{0}}$, with embedding dimension $2|Q|+2$, MLP width $|Q|^{2}+|Q|$, and $\infty$-weight norms at most $\max\{4|Q|+2,10T\sqrt{\log|Q|+\log T}\}$. It has $H=2$ heads with embedding dimension $|Q|$ implying $2|Q|+2$ attention width, and a $3$-layer MLP.

The basic idea is that all prefix compositions $\delta(\cdot,\sigma_{t})\circ\ldots\circ\delta(\cdot,\sigma_{1})$ can be evaluated in logarithmic depth using a binary tree whose leaves are the per-input transition functions $\delta(\cdot,\sigma):Q\rightarrow Q$. The attention heads select the pairs of functions that need to be composed, while the feedforward networks implement function composition. The network will manipulate functions in terms of their transition maps: for example, the encoding of $f:=(1\mapsto 1,2\mapsto 1,3\mapsto 2)$ is

We will produce a construction for the case where $T$ is a power of 2; general $T$ can be handled via padding. To simplify the construction, we also introduce $T$ padding positions $-(T-1),\ldots,0$ at the beginning; while this greatly simplifies the positional selection construction, this padding construction could be replaced with a slightly more complicated MLP. Also, in this construction, we do not need to use residual connections; the parallel prefix sum algorithm we use can be executed “in place”, saving a logarithmic factor in the width. We do assume access to the 2 positional embeddings at each layer; in the absence of residual connections, the identity function restricted to these 2 dimensions can be implemented by the MLP and attention heads.

Let $L=\log_{2}T$ be the depth of the binary tree. We choose $d:=2|Q|+2$. Instead of indexing the dimensions by $[d]$, we give them names:

Left function encoding dimensions $(q,\mathsf{L})$ for each $q\in Q$.

Right function encoding dimensions $(q,\mathsf{R})$ for each $q\in Q$.

Positional encoding dimensions $\mathsf{P}_{1},\mathsf{P}_{2}$.

Without loss of generality, let $Q=[|Q|]=\{1,\ldots,Q\}$ (selecting an arbitrary enumeration of the state space). Also, assume $|Q|\geq 2$ (if not, add a dummy state). We choose $E(\sigma_{t}):=\sum_{q\in Q}\delta(q,\sigma_{t})\cdot e_{(q,\mathsf{R})}$, mapping each input symbol to the “transition map” of its transitions. At the padding positions $-(T-1),\ldots,0$, we will encode the “go to $q_{0}$” function: $\sum_{q\in Q}q_{0}\cdot e_{(q,\mathsf{R})}$.

We first introduce the construction for function composition with a 3-layer ReLU MLP, which will be used by all layers. It gives an exponential improvement over the generic universal function approximation gadget from Lemma 2.

The intermediate dimensions are $d_{1}=|Q|^{2}+|Q|$ and $d_{2}=|Q|^{2}$, and weight norms are bounded by $4|Q|+2$.

We also add $Q$ more weights which let the inputs pass through along the $e_{(q,\mathsf{R})}$ directions (add $Q$ more rows $e_{(q,\mathsf{R})}^{\top}$ to $W_{1},W_{2}^{\prime}$, calling these indices $\bullet q$ for all $q\in Q$; set biases to 0), for a total of $4|Q|^{2}+|Q|$ hidden units and $|Q|^{2}+|Q|$ output dimensions of $W_{2}^{\prime}$.

The rest of the construction uses a standard parallel algorithm for computing all prefix function compositions: at layer $l\in[L]$, compose the function at position $t$ with the function at position $t-2^{l-1}$. This is a standard algorithm for computing all prefix compositions of associative binary operations with a logarithmic-depth circuit (Hillis and Steele Jr., 1986). We choose the position embeddings to enable implementing these “look-backs” with rotation matrices. For each $t\in\{-T+1,\ldots,0,1,\ldots,T\}$, we use the circle embeddings

Let $\theta:=-\frac{\pi 2^{l-1}}{T}$, $\gamma:=100T^{2}(\log|Q|+\log T)$.

Select $W_{Q}^{[\mathsf{L}]}=W_{Q}^{[\mathsf{R}]}=W_{K}^{[\mathsf{R}]}:=\sqrt{\gamma}\cdot(e_{\mathsf{P}_{1}}e_{1}^{\top}+e_{\mathsf{P}_{2}}e_{2}^{\top})$.

Select $W_{K}^{[\mathsf{L}]}:=\sqrt{\gamma}\cdot(e_{\mathsf{P}_{1}}e_{1}^{\top}+e_{\mathsf{P}_{2}}e_{2}^{\top})\rho_{\theta}$, where $\rho_{\theta}$ is the rotation matrix

Select $W_{V}^{[\mathsf{L}]}=W_{V}^{[\mathsf{R}]}:=\sum_{q\in Q}e_{(q,\mathsf{L})}e_{q}^{\top}$.