Transformers Learn Shortcuts to Automata
Authors: Bingbin Liu
Jordan T. Ash
Surbhi Goel
Akshay Krishnamurthy
Cyril Zhang
Venue: ICLR 2023 notable top 5% (2023)
Keywords: Transformer
self-attention
group theory
semigroup theory
algebraic automata theory
shortcut learning
theory of deep learning

Abstract: Algorithmic reasoning requires capabilities which are most naturally understood through recurrent models of computation, like the Turing machine. However, Transformer models, while lacking recurrence, are able to perform such reasoning using far fewer layers than the number of reasoning steps. This raises the question: what solutions are these shallow and non-recurrent models finding? We investigate this question in the setting of learning automata, discrete dynamical systems naturally suited to recurrent modeling and expressing algorithmic tasks. Our theoretical results completely characterize shortcut solutions, whereby a shallow Transformer with only $o(T)$ layers can exactly replicate the computation of an automaton on an input sequence of length $T$. By representing automata using the algebraic structure of their underlying transformation semigroups, we obtain $O(\log T)$-depth simulators for all automata and $O(1)$-depth simulators for all automata whose associated groups are solvable. Empirically, we perform synthetic experiments by training Transformers to simulate a wide variety of automata, and show that shortcut solutions can be learned via standard training. We further investigate the brittleness of these solutions and propose potential mitigations.

Published as a conference paper at ICLR 2023
TRANSFORMERS LEARN SHORTCUTS TO AUTOMATA
Binghin Liu'
Jordan T. Ash
Surhhi Goel Akshay Krishnamurthy
Cyril Zhang?
CCarnegie Mellon University Microsofft Research NYC Uniiverityy of Pennsylvania
ABSTRACT
Algorithmic reasoning requires capabilities which are most naturally understood through recurrent models of computation, like the Turing machine. However, Transformer models, while lacking recurrence, are able to perform such reason- ing using far fewer layers than the number of reasoning steps. This raises the question: what solutions are these shallow and non-recurrent models finding? We investigate this question in the setting of learning automata, discrete dynamical systems naturally suited to recurrent modeling and expressing algorithmic tasks. Our theoretical results completely characterize shortcut solutions, whereby a shal- low Transformer with only o(T') layers can exactly replicate the computation of an automaton on an input sequence of length T. By representing automata using the algebraic structure of their underlying transformation semigroups, we obtain Olog depphh simulators for automata and (11) simulators for all au- tomata whose associated groups are solvable. Empirically, we perform synthetic experiments by training Transformers to simulate a wide variety of automata, and show that shortcut solutions can be learned via standard training. We further in- vestigate the brittleness of these solutions and propose potential mitigations.
INTRODUCTION
Modern deep learning pipelines demonstrate an increasing capability to perform combinatorial rea- soning: pretrained on large, diverse distributions of natural language, math, and code, they nascently solving tasks which seem to require a rigid "understanding" syntax, entailment, and state inference. How do these neural networks represent the primitives of logic and the algorithms they execute internally?
When considering this question, there is immediate mismatch between classical sequential mod- els of computation (e.g., Turing machines) and the Transformer architecture, which has delivered many of the recent breakthroughs in reasoning domains. If we think of algorithm a of sequentially-exxcuuee computational rules, why would we use a shallow! non-recurrent network?
We study this question through the lens of finite semiautomata, which compute state sequences ......, Qr from inputs oI...., by application of a transition function 8 (and initial state go):
= 5(gt-1, 02)-
Semiautomata are the underlying structures governing the computations realizable by automata (such as regular expression parsers finite-state transducers), which are simply semiautomata equipped with mappings from states to output. Thus, one natural motivation for studying them comes from the question of whether Transformers can subsume the structures found in classical NLP pipelines. Another motivation comes from the perspective of reinforcement learning and con- trol, where Transformers are beginning be used as world models: semiautomata specify determin- istic discrete-state dynamical systems.
We perform a theoretical and empirical investigation of whether (and how) nen-recurrent Trans- formers learn semiautomata. We characterize and analyze how shallow Transformers find shortcut
The majority of this work completed while B. Liu was intern Microsoft Research NYC.
This work was completed while Goel was at Microsoft Research NYC. 'Compared to the number of symbols it can process. For example, DistilBERT (Sanh et al., 2019) can handle thousands of tokens with 6 sequential layers.

---

Published as a conference paper at ICLR 2023
parity counter
memory unit
1D gridworkd
@= [even.
Q = (1,2.3.4
@= (11..44
Q = [54 stickers)
I= (0 = lo
= [L.R)
== -=1-...1.3
Σ [6 face rotations]
Figure I: Various examples of semiautomata. From left to right: mod-2 counter, a 2-state memory unit, Grid4, 2-dimensional gridworld constructible via direct product Grid, x Grid4, and Rubik's s Cube, whose transformation semigroup is very large non-abelian group.
solutions, which correctly and efficiently simulate the transition dynamics of semiautomata with far fewer sequential computations than required for iteratively inferring each state qt.
Our contributions. Our theoretical results provide structural guarantees for the representability of semiautomata by shallow, non-recurrent Transformers. 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). There do not exist constant-depth shortcuts for non-solvable semiautomata, unless TC = NC² (Theorem 4). For a natural class of semiautomata corresponding to path integration in a "gridworld" with boundaries, we show that there even shorter shortcuts (Theorem 3), beyond those guaranteed by the general structure theorems above.
We accompany these theoretical findings with an extensive set of experiments:
End-to-end learnability of shoriculs via SGD (Section 4). The theory shows that shortcut solu- tions exist; is the non-convexity of the optimization problem an obstruction to learning them in practice? For variety of semiautomaton simulation problems, find empirically that there is no such obstruction. Shallow non-recurrent Transformers are able to learn shorteuts which generalize near-perfectly in-distribution.
More challenging settings (Section 5). We compare recurreutentt and recurrent models in the presence of additional considerations: out-of-distribution generalization (including to unseen se- quence lengths) and limited supervision. This reveals the brittleness of non-recurrent models, in line with prior "spurious representation" notions of shorteuts in deep learning. Toward mitigat- ing these drawbacks and obtaining the best of both worlds, we show that with recency-biased scratchpad training, Transformers can be guided to learn the robust recurrent solutions.
1.1 RELATED WORK
Emergent reasoning in neural sequence models. Neural sequence models, both recurrent (Wu et al., 2016; Peters et al., 2018; Howard & Ruder, 2018) and non-recurrent (Vaswani et al., 2017; Devlin et al., 2018), have ushered in an era of broadly-applicable and (with pretraining) sample- efficient natural language understanding. Building on this, large-scale non-recurrent Transformer models have demonstrated capabilities in program synthesis, mathematical reasoning, and in-context multi-task adaptation. A nascent frontier to leverage neural dynamics models, again both recur- rent (Hafner et al., 2019) and non-recurrent (Chen et al., 2021a; Janner et al., 2021), for decision making. the highest level, the present work seeks to idealize and understand the mechanisms behind which deep learning solves tasks requiring combinatorial and algorithmic reasoning.
Computational models neural networks. light of the above, it is empirically evident that neural networks are successfully learning circuits which generalize some combinatorial tasks. Many efforts the theory and empirical science of deep learning are dedicated towards the rigor- bounded-complexine
analysis of this phenomenon. Various perspectives map self-attention 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). The research program of BERTology (Clark et al., 2019; Vig, 2019; Tenney et al., 2019) interprets trained models in terms of known linguistic and symbolic primitives.
The most relevant theoretical work to ours is (Barrington & Thérien, 1988). which acts as a "Rosetta Stone" between classical circuit complexity and semigroup theory. The core technical ideas for
2

---

Published as a conference paper at ICLR 2023
Theorems 1 (NC² prefix sum), 2 KKroh-RRhodee and 4 (Barrington) are inspired by the results and discussions therein. In the language of circuit complexity, our work establishes that shallow, non- recurrent Transformers can efficiently represent all of the constructions involved in the (simple) NC² and (significantly more complex) ACC solutions sequential multiplication in semigroups. On the other hand, the shorter shortcut from Theorem 3 carefully leverages self-attention to improve upon these results; we were unable to find an analogous refinement the circuit complexity literature.
Synthetic combinatorial tasks. Our problem setting of simulating finite-state semiautomata unifies the settings of several recent investigations of whether (and how) Transformers learn bounded-depth Dyck languages (Yao et al., 2021), parities (Anil et al., 2022), adders (Nogueira et al., 2021; Nanda & Lieberum, 2022), regular languages (Bhattamishra et al., 2020), and sparse logical predicates (Edelman et al., 2022; Barak al., 2022). Zhang et al. (2022) empirically analyze the behavior and inner workings of Transformers on random-accest 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
2.1 SEMIAUTOMATA AND THEIR ALGEBRAIC STRUCTURE
A semiautomaton A: (Q, م 5) consists of a set of states Q. an input alphabet 2, and a transition function 6 : Q x Σ + Q. In this work, Q and E will always be finite sets. For all positive integers and a starting state qo E Q. A defines a map from input sequences (01..... σr) E ΣT to state sequences (91, (9...... 92) Q²: q: 8(9(-1,0)) for t = 1, ... T. This is a deterministic Markov model, in the sense that at time I, the future states 1t+h.-.. QT only depend on the current state gr and the future inputs 0zt1y-.., ØT.
We define the task of simulation: given a semiautomaton A, starting state go, and input sequence (oI..... ... JT), output the state trajectory AT, nn (oI...., .......,00. TT) == (ql..... .T). Let f E. + Q" be a function (which in general can depend on A, qo). We will say that f simulates Ar, if forrr) = Ar. (F1-)) for all input sequences dT.r Finally, for positive integer T, we say that function class of functions from E" - Q" is said to simulate A at length T if, for each qu E Q, there is function in F which simulates (A,T. go).
Every semiautomaton induces a transformation semigroup T(A) of functions p p:Q Q under composition, generated by the per-input-symbol state mappings ō(-,o) Q + Q. When T(A) contains the identity function, invertible, T(A) is permutation group. See Figure 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 but interesting example is a parity counter (Figure 1, left): the state is a bit, and the inputs are ["toggle the bit", "do nothing"), ); the transformation semigroup is C2, the cyclic group of order 2. Parity has been studied in previous synthetic experiments (Zhang et al., 2022; Anil et al., 2022).
called a transformation monoid. When all of the functions are
2.2 RECURRENT AND NON- DONRRECURENN NEURAL SEQUENCE MODELS
A sequence-to- -eeuuueceeeeeeeeeeeeeee neural network of length T' and dimension d is a function fon Txd x
θ + g'''d, with trainable parameters 0 E O. Equipped with an encoding layer E : E + B& and
decoding layer W Rd + Q (applied position-wise), the function (W fm o E) : Σx + Q has the same input and output types as Ar. This work will investigate when the functions defined by neural networks can simulate semiautomata.
A recurrent neural network (RNN) is a sequence- srquence-to-sequenee neural network defined by iterated composition of recurrent unit g : Rdd Rd + IR.. For a given initial hidden state ho E Rdd. and input sequence .1,.... Wr E R.d it produces an output hidden state sequence
h, g( 0;)),
= I.....
Thus, for any fixed 8, RNN defines semiautomaton with infinitely many states and inputs: Q = Σ = ,dd Thus, as long as g can represent 8. RNNs can simulate all semiautomata. In this sense, the computational models of RNNs and semiautomata naturally coincide.
3

---

Published as a conference paper at ICLR 2023
An L-layer Transformer is another sequence toreenccett-eeeenee network, consisting of alternating self- attention blocks and feedforward MLP blocks
fur = (id + fml)) o (id f (id + o o... (id + fatin) (id P).
Briefly, an attention layer performs & f--normalized mixing operations across positions t, while a constant-layer MLP block performs position-wise function approximation (with mixing between positions); id denotes the identity function (residual connections), and encodes the position t.? We use fairly standard positional encodings in both theory and experiments. Importantly, the standard Transformer is convolutional (in that the weights in sattn and fmip are shared across positions 1), but is non-recurrent: parameters are not shared across blocks.
All architectures have a notion computational depth D (succinctly, depth) when processing inputs of length T. which is the longest path in the computational graph. For RNNs, this is ((T)) while an L-layer Transformer (with constant-layer MLPs) has depth (L). For Transformers, since they co- incide up to constant factors, we use depth and number of layers interchangeably. We will also track the layers L, embedding dimension d, attention width (the largest number of parallel attention head outputs), and MI.P width (the largest number of parallel hidden activations in the MI.P blocks)."
3 THEORY: SHORTCUTS ABOUND
A T-layer Transformer trivially simulate semiautomaton at length T sequentially: like an RNN, the (-th layer can implement (an embedding of) the state transition q:-1 + Qt- Yer, Trans- formers succeed in practice with long contexts (2 10a) and fewer layers (as few 6). A natural theoretical question is that of representability: can Transformers efficiently simulate semiautomata with parallel shortcut solutions. whose depths are much smaller than the sequence length T?
Definition (Shorteut solution). Let A be a semiautomaton. Suppose that for every 2 1, there neural network fr which simulates A at length T. We call this sequence is (fTTT21 a shortcut solution to the problem of simulating A if its depth D satisfies D s o(T).
By this definition, shortcuts are quite general and some are less interesting than others. For exam- ple, it is always possible to construct constant-depth neural network which memorizes all |Σ|T values of Ar. but these networks must exceptionally wide. We could also "fast-forward" state simulation, letting each of (say) VT layers simulate /T consecutive state transitions, but, without exploiting the structure of the semiautomaton, this would require width !(2vT). To rule out these cases and focus on interesting shortcuts for Transformers, we want the other size parameters (at- tention and MLP width) to be small: say, scaling polynomially with T, To construct such shortcuts, we need ideas beyond explicit iteration of state transitions.
dependent only
[Q), [
3.1
SEMIAUTOMATA ADMIT SHALLOW PARALLEL SHORTCUTS
We begin by noting that polynomial-width shortcuts always exist. This may be counterintuitive we restrict ourselves to viewing a network's intermediate activations as representations of states q-- When we instead view them as encoding state transformations 5(., .) Q - Q and their composi- tions, a divide-and andee-and-con and-conquer construction is evident (see Figure 2a), detailed in Appendix C.2:
Theorem 1 (Simulation is parallelizable; informal). Transformers can simulate all semiautomata A = (Q, E, 5) at length T, with depth O(log T), embedding dimension O(Q)), attention width O(QQ),, and MLP width ((
If we assume that an attention head can only select a constant number of indices, Theorem is unim- provable: the receptive field of a sublogarthmiicdeply Transformer is not large enough. However, is known in theory and practice that soft- attenttttnon heads are capable of attending broadly, repre- senting certain non-sparse dependencies (Clark et al., 2019; Yao et al., 2021). Thus, we can ask more challenging question: can the dense operations of attention enable even shallower shortcuts?
We omit layer normalization. This discrepancy is superficial; see the discussion Appendix A.4.
Full statements and proofs also track oc-weight norms (the largest absolute value of any parameter) and bit precision of each floating-poit computation. We defer precise definitions and discussion Appendix A.4.
4

---

Published as a conference paper at ICLR 2023
& 4..
sunto modp
(c)
(d)
Figure 2: Intuitions for the theoretical constructions. (a) Divide-and-conquee function composi-
tion yields logarithmic-depth shortcuts (Theorem 1). (b) The two "atoms" of the constant-depth Krohn-Rhodes decomposition (Theorem 2) of solvable semiautomaton: modular addition and quentially resettable memory. (c) Information flow of the cascade product, which is used to glue these atoms together, and easily implemented with residual connections. (d) An even shorter short- cut solution for gridworld simulation (Theorem 3; see Appendix C.4).
The key to resolving this question comes from Krohn-Rhodes theory, which gives us tools to reason about the structure of arbitrary semiautomata and their transformation semigroups. A landmark result (Krohn & Rhodes, 1965), a vast generalization of the uniqueness of prime factorizations for integers, shows that simulate any semiautomaton, objects: simple groups, and a memory unil (Figure lb). When the Krohn-Rhodes decomposition contains non-abelian groups (we call such semiautomaton solvable"), there exist constant-depth circuits for simulation, which we manifest as neural networks.
only need handle two types of elementary
turns out that positional weight sharing (a.k.a. "width-1 convolutions"). non-recurrence, and self- attention are particularly well-suited for efficiently representing the Krohn-Rhode decomposition of semiautomaton: uniform-sun sum attention heads perform abelian group operations, proximity- selection heads implement memory units, and the rest of the architecture (MLPs and residual con- nections) implements the cascade product (Definition 4) which combines these atomic operations. Overall, we conclude:
Theorem (Transformer Krohn-Rhodes; informal). Transformers can simulate all solvable semi- automata A (Q, Σ, 8), with depth O[IQ (2 log [Q)), embedding dimension 2011Q log (Q1), attention width 20(1 log QQ), and MLP width 12))( + 2011) log (Q1) T.
quite counterintuitive' that as T 00, no additional depth needed for such large class of problems. We provide background and details (including the definition and implementation of this notion of semigroup product) in Appendices A.2 and C.3. In Figure 2b and 2c, we illustrate the three key ingredients: efficient implementations of the two atoms (modular counting and memory lookups), and the procedure for gluing them together (building a transformation cascade).
What does each layer do? The construction in Theorem recursively composes functions, as opposed to the naive solution of directly emulating states. Theorem 2 takes very different approach:: it relies on the holonomy decomposition variant of the Krohn-Rhodes theorem (Eilenberg, 1974). Rather than simulating J. or composing functions, the computational paths correspond to a [Q|- level tree of nested coarsenings of the semiautomaton's dynamics: "which subset of states could be in right now?" Within each level of this tree, the network must implement (generally non- commutative) group operations. This can be done with O([Q) log (Q) layers, by leveraging the Holden- decompositions and the universal embedding theorem (Krasner & Kaloujnine, 1951).
Can we get even shallower shortcuts? Finally, we show that on a natural class of problems, the computational model of self-attention leads to further fine-grained improvements over the guarantees of Krohn-Rhodes theory. Motivated by the application of Transformers in modeling environment
'ee Definition 6. Among the solvable groups are the dihedral groups D2n, the permutation groups Sa, An for S 4, the quatermion group Qs. all groups of order < 120 except As, and all groups of odd order. Perhaps surprisingly, the only place where width of 7 is used is implement mod-n gate. This dependence can removed entirely if we allow for periodic activation functions such as x + sin(z).
From the back cover of Rhodes et al. (2010): the underlying theorem launched theory which "reveals deep and unexpected connections between algebra (semigroups) and areas of science and engineering".
5

---

Published as a conference paper at ICLR 2023
dynamics, consider the semiautomaton Grid, corresponding to "gridworld": n states line, with input symbols "move left if possible" and "move right if possible" (see Figure 1, middle). We show that self-attention enables extremely concise solution, with depth independent of both T and [Q| = nr:
Theorem (Depth-2 shorteut for gridworld: informal). For all positive integers n. T, Transformers can simulate Grid, at length T, with depth 2,' embedding dimension (1)), attention width O(n), and MLP width (T))
The proof builds concise parallel nearest boundary detector, and can be found in Appendix C.4. We note that this particular setting known to be an extremal case for the holonomy construction in Krohn-Rhodes be interesting to generalize our improvement and characterize the class of problems for which self- attention affords (1) instead of poly(IQ))-deplly solutions.
theory (Maler (2010) discusses this, calling it the elevator automaton). It would
Aren't neural networks universal function approximators? Sufficiently wide neural networks with sufficiently expressive nonlinearities can fit arbitrary functions (Hornik et al., 1989; Cybenko, 1989). However, if we constrain complexity measures such as depth and width, one cannot hope to apply universality directly. It is true that one take the discrete circuit constructions in (Bar- rington & Thérien, 1988), "comppll" every gate to a constant-depth network, and recover shortcut solutions with o(T) depth and poly (I') width. However, our constructions go far beyond black-box reductions- the roles of self-attention and positional parameter sharing allow for such efficient con- structions that no parameter count depends on T (except the MLP width, which is removable with a periodic activation function). Furthermore, the constructions are so simple and natural that they are corroborated by the preliminary "reverse engineering" investigation in Section 4.
3.2 LOWER BOUNDS
Can Theorem 2 be improved to handle non-solvable semiautomata? (Equivalently: can Theorem be improved to constant depth?) It turns out that as consequence of classic result in circuit complexity (Barrington, 1986), this question is equivalent to the major open question of TC - NC! (thus: conjecturally, no). Unless these complexity classes collapse, Theorems I and 2 are optimal. In summary, simulating non-solvable semiautomata with constant depth is provably hard:
Theorem 4 (Transformer Barrington). Let A be non-solvable semiautomaton. Then, for suffi- ciently large T, no Ollog T) prrecison Transformer with depth independent of TT and width polyno- mial in T can continuously simulate A at length T, unless TC = NC!.
This is proven in Appendix C.5. The smallest example of non-solvable semiautomaton is the one on [Q| = 5 states, whose transitions generate Ag (all of the even permutations).
Finally, we note that although our width bounds might be improvable, an exponential-in-jQQ number of hypotheses (and hence network with poly(IQ)) parameters) is unavoidable if one wishes to learn an arbitrary /Q|-state semiautomaton from data: there are |Q| (Q--E) of them, which generate (Q)(((()) distinct semigroups (Kleitman et al., 1976). If we wish study how machine learning models can efficiently identify large algebraic structures, we will need finer- merrggraieed inductive biases to specify which semiautomata prefer, direction for future work.
EXPERIMENTS: CAN SGD FIND THE SHORTCUTS?
Our theorems are limited to representability: concise shallow solutions exist, but whether gradient- based local search (i.e., standard training) finds them is another matter entirely. For example, em- bedded within the problem of learning simulate the 2-state parity semiautomaton is well-known non-convex optimization problem (Daniely & Malach, 2020; Edelman et al., 2022; Nichani et al., 2022). In general, even detecting whether T(A) contains cycle is PSPACE- (Cho & Huynh, 1991). Theoretically understanding how the training dynamics of deep learning transcend the worst- case hardness of non-convex optimization a major frontier of research, that we do not attempt to
Thii requires max-pooling we do not use max-pooling, we can instead use an ML. with width 20(x) and depth ((1), or width O(n) and depth O(log n).
As with Theorem 2, the width can be reduced ((z) if employ periodic activation functions.

---

Published as a conference paper at ICLR 2023
5
15
Dyck
29.3
130
102
200
100
100
120
100
C2
17.0
29.5
9
100
100
15.5
100
993
100
100
66.6
41.6
66.7
100
150
100
15.0
27:0
09.3
100
99.0
80.6
Fx
25.4
27.2
47.4
15.2
100
100
100
120
100
100
1st layer, uniform attention
4th layer, 0 detector
4th layer, S detector
Co Accurape
S. Accuracy 02
31.6
49-2
59.
60.4
73.
15.3
100
100
100
100
225
25.2
529
46.7
12.2
28.8
200
120
100
13.6
26.0
03
(a) Accuracy across tasks (rows) and net- work depths (columns).
steps
Training steps
(b) Attention heatmaps (Grids); unstable training (C2 and St).
Figure 3: Overview of the empirical results in Section 4, on in-distribution learnability of shortcuts by standard Transformer training. (a) Truncated table of results (in-distribution accuracy); rows specify semiautomaton simulation problems, and columns specify network depth. (b) Attention heads implement a nearest boundary detector (top); training is highly unstable (bottom).
address here. Instead, we approach the question of optimization through an empirical lens. Our primary goal is to understand if gradient-based training 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; the median performance is provided in Appendix B.
For a selection of 19 semiautomata corresponding to various groups and semigroups, we train shallow Transformers to output their state sequences given random inputs. Specifically, apply GPT-2-like models (Radford et al., 2019) with 1-16 layers on freshly-sampled sequences of length T = 100. Strikingly, obtain positive results (> 99% in-distribution accuracy) for all of them, including ones which generate the non-solvable groups A5 and Sy." 10 Figure 3a gives a selection of our full results (in Appendix B. We find that more complex semiautomata (corresponding non- abelian groups) require deeper networks to learn, in agreement with our theoretical constructions.
Which shallow solutions are learned? Our theoretical results identify shorteut solutions which follow multiple, mutually incompatible paradigms. In general, we do not attempt full investigation of mechanistic interpretability of the trained models. As preliminary evidence, we visualize some of the attention patterns in Figure 3b (top) within successfully-trains trained models, finding attention heads which perform flat summations (with uniform attention) and conditional resets.
Optimization quirks. Although sufficiently deep networks find the solutions with non-negligible probability, the training dynamies are unstable; Figure 3b (bottom) shows some training curves, ex- hibiting 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 semi- automaton simulation will be useful as clean, nontrivial testbed (with multiple difficulty knobs) for debugging and improving training algorithms, and perhaps the neural architectures themselves.
5 FURTHER EXPERIMENTS: MORE CHALLENGING SETTINGS
For wide family of algebraic structures, we have proven that the function class of shallow non- recurrent networks subsumes deeper finite-state recurrent models. Furthermore, the experiments in Section have shown that despite the non-convexity of the optimization problem, standard training works: Transformers can learn shorteuts to semiautomaton simulation, end-to-end. While encourag- ing, the experiments in Section 4 are idealized in several ways, and it is natural to ask Transformers perform similarly in more challenging semiautomaton simulation scenarios. Towards answering this
Using fraphly-sampled data ensures that the model cannot achieve gond performance by brute-force memo- rization in a number of training steps we could ever exccute computationally (for sufficiently large T such as T= 100), since there are an exponential number of sequences.
"Explanatiooss on why certain groups harder learn are provided in Appendix B.
7

---

Published as a conference paper at ICLR 2023
Task
Dycka.8
Gridg
Ss
C4
Dg
(abab)*
Observation
stack top
boundary
T1:t
1 mod
location
accept
Accuracy
100.0
99.7
98.0
99.7
99.8
100.0
(a) Accuracies with indirect supervision. LSTM gets 100% all tasks.
GPT
LSTM
Scratchpad-reeeancy
distribution
Ss Accuracy
0.9
0.8
0.0
0.7
0.7
1/32
1/16
1/8
1/4
1/2
0.0
0.2
0.4
0.6
0.8
1.0
20
40
60
80
100
Preveal
Pr[0=1]
Sequence Length
(b) Varying Preveal (log spacing).
(c) C2 (parity): accuracy at different Pr(a = 1] and T.
Figure 4: Overview of the empirical results in Section 5. (a) Learning in the latent-state setting, with various observation maps (. (b) Learning from incomplete state sequences: final accuracy position-wise probability of a hidden token. (c) OOD generalization: Transformers fail to generalize to different distributions and lengths.
question, in this section, we consider some challenges that may arise in practice and an associated set of experimental results; further details are deferred to Appendix B.2.
5.1 INCOMPLETE AND INDIRECT SUPERVISION
Automata are partially observable semiautomata. Consider the case of partial observabil- ity. For any semiautomaton A (Q, E, and (generally non-invertible) observation function p Q - Q, we can define the problem of predicting q. p(a)). If can only obtain observa- tions Gt (i.e., the state is latent), this fully captures the problem of learning a finite-state automaton from data. The results in this paper have shown that this equivalent to the fully-observable case in terms of representation. However, the learning problem can be much harder; indeed, this may account for Bhattamishra al. (2020)'s 's negative results on learning regular languages with constant- depth Transformers. Note that this also captures autoregressive next-token prediction tasks induced by distributions (e.g., generating Dyck languages (Yao et al., 2021)) where the sequence's contin- uations depend on a latent semiautomaton's 's state (e.g., the current stack for Dyck). Despite these potential challenges, we find that Transformers are able to find a solution with good in-distribution performance for all partially observable settings consider; see Figure 4(a).
Learning from incomplete state sequences. Next, we consider the setting which identical to that described in Section 4, but each state 9: is randomly revealed from the training data with some probability 0 S Preveal I 1. As with partial observability, this does not affect representation issues, but can make learningoppimmization much harder. Figure 4b shows the accuracy of Ss for models trained on length 100 sequences for various Preveal. It can be seen that Transformers may be unable to find good solutions when the labels become sparser, whereas LSTM's performance stays robust across all choices of Preveal.
5.2 OUT-OF-DISTRIBUTIIN SHORTCOMINGS 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 ø shifts between training and testing (but the semiautomaton remains the same), a non-recurrent shortcut solution might map inputs into intermediate latent variable space (like the sum) which fails to generalize.
Indeed, we observe that with the same models which obtain the positive in-distribution results in Section 4, accuracy degrades as distribution shift increases; see Figure 4(c) (left), where the per- formance drops as the probability of seeing input o = 1 deviates from the training distribution (Pr(o 1] 0.5). From the viewpoint of mechanistic interpretation, this is further (but not abso- lutely conclusive) evidence that with standard training. Transformers learn implementations similar to those predicted by the theory. We provide details and further empirical evidence in Section B.2.3.
8

---

Published as a conference paper at ICLR 2023
More ambitiously, could try use these models to extrapolate 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; Yao et al. (2021); Anil et al. (2022) for more experiments similar to ours. Figure 4(c) (right) shows the performance on sequences of various lengths, where Transformer's s accuracy drops sharply as move to lengths unseen during training. In contrast, LSTM performs perfectly in both out-of-distribution scenarios. Details are deferred to Section B.2.4,
Shortcuts as "unintendedd solutions. Throughout the deep learning literature, the term shortcut is often used in a statistical sense to connote undesired (i.e., misleading, spurious, or overfitting) learned representations (Geirhos et al.. 2020; Robinson al., 2021). The experiments in this sec- tion show why our circuit-depth shortcuts are statistical shortcuts. Specifically, we have identified a problem with learning relaxations to sequential state simulation: the models may "hallucinate" statistically suboptimal latent variables. The positive results in Sections 3 and 4 suggest that this may only be robustly diagnosable via -ut-of-distiobution evaluation.
Finally, we empirically show that this flaw is circumventable. Using combination of scratchpad (a.k.a. -fthuu-of- (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-7) solutions that generalize out-of--istributie and to longer sequence lengths (Figure 4(c), yellow curves).
Computationall-satittii tradeofl's. The experiments in this section highlight a statistical price for learning shortcuts to semiautomaton simulation. On the other hand, the shallowness of these shorteuts is computationally appealing: leveraging parallel computation, they enjoy much lower latency (log 7) or ((1), compared to O(T)), in both training and inference. Whether the best of both worlds is attainable is an interesting avenue for future work.
6 CONCLUSIONS AND FUTURE WORK
We have conducted theoretical and empirical analysis of how shallow Transformers can learn shortcul solutions to the problem of simulating the transitions of semiautomata (and thus, the al- gebraic structures which underlie regular expressions, finite-state transducers, and deterministic MDPs), Using tools from semigroup theory and circuit complexity, we have constructed explicit logarithmic-depth and constant-depth shortcuts for semiautomaton simulation. Experimentally, we have shown that gradient-based based optimization finds shorteut solutions which generalize near near-perfeer perfeectyyly in-distribution (Section 4), but are brittle out-of- of-distribuutio (Section 5). We hope that these results shed new light on the power and limitations of applying shallow non-recurrent models, even when the dynamics we wish to represent are deep and recurrent.
Beyond Transformers? The theory and experiments in this work are specialized to the Transformer architecture, to provide a concrete and timely setting. However, we note that the underlying themes (continuous arithmetic circuits; parameter sharing across input positions and/or iterated function compositions; local vs. global computational units) are not specialized to any particular neural architecture" nor the field of deep learning at all. The question of "which architectures are even more natural for learning discrete automata, while being optimizable by gradient-basee search?" is extremely open-ended. We believe that the themes of sufficient depth and recurrent vs. non-recurrent function composition are relevant to the study of other (and future) deep learning methods.
Future topies. In terms of theory, we have only scratched the surface of the possible interplay be- tween neural architectures and classical ideas from the complexity theories of circuits and automata. One salient direction is to generalize the shorter shortcut constructions Theorem 3. Also, we have made no attempt treat stochastic environments, which would fully capture probabilistic Markov models and MDPs. Section alludes to a landscape of algorithm design challenges in the presence of distribution shift and limited supervision. The latter (i.e., latent state inference) is known to lead to worst-case computational hardness (Papadimitriou & Tsitsiklis, 1987), but yields powerful em- pirical tools when tractable. Towards fully understanding and leveraging the cireumstances which allow learning algorithms to decode and simulate ge. there is much work to be done.
IIn fact, the divide-and sovide-and-conee construction of Theorem is almost recurrent with log(?') depth, and re- sembles WaveNet-like hierarchical pooling (van den Oord et al., 2016; Larsson et al., 2016), more than Transformers.

---

Published as a conference paper at ICLR 2023
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 read- ability.
REPRODUCIBILITY STATEMENT
Complete proofs of the theoretical results are provided in Appendix C, with self-contained tutorial of relevant group-theoretic concepts in Appendix A.2. For the empirical results, all our datasets derived from synthetic distributions, which clearly described in Appendix B.1 and B.2. The architectures, implementations (with references to popular base repositories), and hyperparameters (including training procedure) documented in Appendix B.3. We intend release code as open source prior to publication.
REFERENCES
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Am- brose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. arXiv:2207.04901, 2022.
Sanjeev Arora and Boaz Barak. Computational complexity: modern approach. Cambridge Uni- versity Press, 2009.
Arpit Bansal, Avi Schwarzschild, Eitan Borgnia, Zeyad Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. End-to-end algorithm synthesis with recurrent networks: Logical extrapola- tion without overthinking. arXiv: 2xXi==22002.05826, 2022.
Boaz Barak, Benjamin I. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: SGD learns parities near the computational limit. arXiv:2207.08799, 2022.
David A. Mix Barrington. Bounded-width polynomial-size branching programs recognize exactly those languages in NC!. In Symposium on the Theory of Computing, 1986.
David Mix Barrington and Denis Thérien. Finite monoids and the fine structure of NC². Journal
of the ACM, 1988.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transform- ers to recognize formal languages. In Conference on Empirical Methods in Natural Language Processing, 2020.
Michael Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv:2104.13478, 2021.
Ashok K Chandra, Steven Fortune, and Richard Lipton. Unbounded fan-in eircuits and associative functions. In Symposium on Theory of Computing, 1983.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, 2021a.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philipp Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fo- tios Chantzis. Elizabeth Barnes, Ariel Herbert- Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leiki, Josh Achiam, Vedant Misra, Evan Morikawa, Alec
10

---

Published as a conference paper at ICLR 2023
Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob Me- Grew, Dario Amodei, Sam McCandlish, llya Sutskever, and Wojciech Zaremba. Evaluating large language models trained code. arXiv:2107.03374, 2021b.
Sang Cho and Dung T Huynh. Finite-automaton aperiodicity is PSPACE-complete. Theoretical
Computer Science, 1991.
Noam Chomsky and Marcel P Schützenberger. The algebraic theory of context-free languages. In Studies in Logic and the Foundations of Mathematics, 1959.
Xiangxiang Chu, Zhi Tian, Bo Zhang, Xinlong Wang, Xiaolin Wei, Huaxia Xia, and Chunhua Shen. Conditional positional encodings for vision transformers. arXiv preprint arXiv:2102. 10882, 2021.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? An analysis of BERT's s attention. In ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, 2019.
George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of Con- trol, Signals and Systems, 1989.
Amit Daniely. Depth separation for neural networks. In Conference on Learning Theory, pp. 690- 696. PMLR, 2017.
Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Infor- mation Processing Systems, 2020.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2019.
Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Marcus Hutter, Shane Legg, and Pedro A Ortega. Neural networks and the chomsky hierarchy. arXiv preprint arXiv:2207.0000 000099, 2022.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv:1810.048005 2018.
Iddo Drori, Sarah Zhang, Reece Shuttleworth, Leonard Tang, Albert Lu, Elizabeth Ke, Kevin Liu, Linda Chen, Sunny Tran, Newman Cheng, Roman Wang, Nikhil Singh, Taylor Patti, Jayson Lynch, Avi Shporer, Nakul Verma, Eugene Wu, and Gilbert Strang. A neural network solves, explains, and generates university math problems by program synthesis and few-shot learning at human level. Proceedings of the National Academy of Sciences, 2022.
Javid Ebrahimi, Dhruv Gelda, and Wei Zhang. How self-attention networks recognize Dyck-n languages? In Findings of the Association for Computational Linguistics: EMNLP, 2020.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference Machine Learning, 2022.
Attila Egri- Nagy and Chrystopher L Nehaniv. Computational holonomy decomposition of transfor- mation semigroups. wrXiv:1508.0634. 2015.
Samuel Eilenberg. Automata, languages, and machines. Academic Press, 1974.
Ronen Eldan and Ohad Shamir. The power of depth for feedforward neural networks. In Conference learning theory, pp. 907-940. PMLR, 2016.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Gan- guli, Zac Hatfield-Doddt Danny Hernandez. Andy Jones. Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. pub/ pub/2021//rameworr framework / index. html.
URL https: /ttaaraansss.. trans
Merrick Furst, James B. Saxe, and Michael Sipser. Parity, circuits, and the polynomial-time hierar- chy. Mathematical Systems Theory, 1984.
11

---

Published as a conference paper at ICLR 2023
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2020.
Surbhi Goel, Varun Kanade, Adam Klivans, and Justin Thaler. Reliably learning the ReLU in polynomial time. In Conference on Learning Theory, 2017.
Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint 1rXiv:1603.08983 2016.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401, 2014.
Jiatao Gu, James Bradbury, Caiming Xiong, Victor O.K. Li, and Richard Socher. Non-autoregressive neural machine translation. arXiv:1711.02281, 2017.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv:1912.01603 2019.
Michael Hahn. Theoretical limitations of self-attention in neural sequence models. the Association for Computational Linguistics, 2020.
Transactions of
Adi Haviv, Ori Ram, Ofir Press, Peter Izsak, and Omer Levy. Transformer language models without positional encodings still learn positional information. arXiv:2203.16634 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In IEEE Conference on Computer Vision and Pattern Recognition, 2016.
Christoph Hertrich, Amitabh Basu, Marco Di Summa, and Martin Skutella. Towards lower bounds the depth of ReLU neural networks, In Advances in Neural Information Processing Systems, 2021.
W Daniel Hillis and Guy L. Steele Jr. Data parallel algorithms. Communications of the ACM, 1986.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks uni- versal approximators. Neural networks, 1989.
Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. 18Xiv: 060666 /4, 2018
DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block- recurrent transformers. arXiv:2203.07852, 2022.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In Advances in Neural Information Processing Systems, 2021.
Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A Smith. Finetuning pretrained transformers into rnns.
arXiv:2103.1302 2021.
Guolin Ke, Di He, and Tie-Yan Yan Liu. Rethinking positional encoding in language pre-training. arXiv 2020.
preprint ar arXiv:2006.155
Daniel Kleitman, Bruce R Rothschild, and Joel H Spencer. The number of semigroups of order n. Proceedings of the American Mathematical Society, 1976.
László Kovács and Cheryl Praeger. Finite permutation groups with large abelian quotients. Pacific Journal of Mathematics, 1989.
Marc Krasner and Léo Kaloujnine, Produit complet des groupes de permutations et probleme d' d'extension de groupes II. Acta Scientiarum Mathematicarum, 1951.
Kenneth Krohn and John Rhodes. Algebraic theory of machines, I: Prime decomposition theorem for finite semigroups and machines. Transactions of the American Mathematical Society, 1965.
12

---

Published as a conference paper at ICLR 2023
Guillaume Lample and François Charton. arXivv1912.01412. 2019.
Deep
learning for symbolic mathematics.
Gustav Larsson, Michael Maire, and Gregory Shakhnarovich. FractalNet: Ultra-deep neural net- works without residuals. arXiv: 1605 (66400768,, 2016.
Holden Lee, Rong Ge, Tengyu Ma, Andrej Risteski, and Sanjeev Arora. On the ability of neural nets to express distributions. In Conference on Learning Theory, pp. 1271-1296. PMLR, 2017.
Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, Thomas Hubert, Peter Choy, Cyprien de Masson d' d'Autume, Igor Babuschkin, Xinyun Chen, Po-Sen Huang, Johannes Welbl, Sven Gowal, Alexey Cherepanov, James Molloy, Daniel Mankowitz, Esme Sutherland Robson, Push- meet Kohli, Nando de Freitas, Koray Kavukcuoglu, and Oriol Vinyals. Competition-level code generation with alphacode. arXiv:2203.078144 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv:171..5101,, 2017.
Oded Maler. On the Krohn- Bhonn-Rhodes cascaded decomposition theorem. In Time Verification. 2010.
Oded Maler and Amir Pnueli. On the cascaded decomposition of automata, its complexity and its application to logic (Draft). 1994.
Carlo Mereghetti and Beatrice Palano. Threshold circuits for iterated matrix product and powering. RAIRO- TheRo-The Informatics and Applications, 2000.
William Merrill, Yoav Goldberg, Roy Schwartz, and Noah A. Smith. On the power of saturated Transformers: A view from circuit complexity. 17X1v:210..16113, 2021.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. arXiv:2209. 00588, 2022.
Anirbit Mukherjee and Amitabh Basu, Lower bounds over boolean inputs for deep neural networks with ReLU gates. arXiv: 1at::1711.000007 2017.
Neel Nanda and Tom Lieberum. Alignment
mechanistic
interpretability anal- URL
ysis of grokking. / //www. www.aaaiggmennttruu. a-mechani hyg.
Forum, 2022.
https:
Benjamin Newman, John Hewitt, Percy Liang, and Christopher D. Manning. The EOS decision and length extrapolation. In BlackboxNiPP Workshop on Analyzing and Interpreting Neural Networks for NLP, 2020.
Eshaan Nichani, Yu Bai, and Jason Lee. Identifying good directions to escape the NTK regime and efficiently learn low-degree plus sparse polynomials. arXiv:2206.03688 2022.
Rodrigo Nogueira, Zhiying Jiang, and Jimmy Lin. Investigating the limitations of transformers with simple arithmetic tasks. urXiv:2102.1 13019, 2021.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Au- gustus Odena. Show your work: Scratchpads for intermediate computation with language models. arXiv:2112.00114.200 2021.
Christos Papadimitriou and John N Tsitsiklis. The complexity of Markov decision processes. Mathematics of Operations Research, 1987.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas K"opf, Ed- ward Yang, Zach Vito. Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems, 2019.
13

---

Published as a conference paper at ICLR 2023
Matthew E. Peters, Mark Neumann, Mohit lyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representations. arXiv:1802.05336, 2018.
Stanislas Polu and Ilya Sutskever. Generative language modeling for automated theorem proving. arXiv:2009.03399 2020.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl blog. 2019.
John H. Reif and Stephen R. Tate. On threshold circuits and polynomial computation. SIAM Journal Computing, 1992.
John Rhodes, Chrystopher Nehaniv, and Morris Hirsch. Applications of automata theory and algebra: via the mathematical theory of complexity to biology, physics, psychology, philosophy, and games. World Scientific, 2010.
Joshua Robinson, Li Sun, Ke Yu, Kayhan Batmanghelich, Stefanie Jegelka, and Suvrit Sra. Can con- trastive learning avoid shortcut solutions? Advances in Neural Information Processing Systems, 2021.
Itay Safran, Ronen Eldan, and Ohad Shamir. Depth separations in neural networks: what is actually being separated? In Conference Learning Theeory pp. 1664-26666 PMLR, 2019.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistiBBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arX1v:1910001100, 2019.
Tal Schuster, Ashwin Kalyan, Alex Polozov, and Adam Kalai. Programming puzzles. In Advances in Neural Information Processing Systems Track on Datasets and Benchmarks, 2021.
Marcel Paul Schützenberger. On finite monoids having only trivial subgroups. Information and Control, 1965.
Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin, Micah Goldblum, and Tom Goldstein. Can you learn an algorithm? generalizing from easy to hard problems with recurrent networks. In Advances in Neural Information Processing Systems, 2021.
Hava T Siegelmann and Eduardo D Sontag. On the computational power of neural nets. In Confer- Learning Theory, 1992.
Matus Telgarsky. Benefits of depth in neural networks. In Conference on learning theory, pp- 1517-1539. PMLR, 2016.
Ian Tenney, Dipanjan Das, and Ellie Pavlick. BERT rediscovers the classical NLP pipeline. 19X5::100.005550. 2019.
Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. WaveNet: A generative model for raw audio. 1wXv::1609003499 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neurul Informa- fion Processing Systems, 2017.
Ashish Vaswani, Prajit Ramachandran, Aravind Srinivas, Niki Parmar, Blake A. Hechtman, and Jonathon Shlens. Scaling local self-attention for parameter efficient visual backbones. In IEEE Conference Computer Vision and Patten Recognition, 2021.
Jesse Vig. Visualizing attention in transformer-base based language representation models.
arXiv: 1904. 02679, 2019.
14

---

Published as a conference paper at ICLR 2023
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian lchter, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv:2201.11903, 2022.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like Transformers. In International Confer- ence on Machine Learning, 2021.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowiez, Joe Davison, Sam Shleifer, Patrick Von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface's transformers: State- of-the-art natural language processing.
Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoe V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, Jeff Klingner, Apurva Shah, Melvin John- son, Xiaobing Liu, Eukasz Kaiser, Stephan Gouws, Yoshikiyo Kato, Taku Kudo, Hideto Kazawa, Keith Stevens, George Kurian, Nishant Patil, Wei Wang, Cliff Young, Jason Smith, Jason Riesa, Alex Rudnick, Oriol Vinyals. Greg Corrado, Macduff Hughes, and Jeffrey Dean. Google's neural machine translation system: Bridging the gap between human and machine translation. arXiv: 08 44, 2016.
Yisheng Xiao, Lijun Wu, Junliang Guo, Juntao Li, Min Zhang, Tao Qin, and Tic-yan Liu. A survey non-autoregressive generation for neural machine translation and beyond. arXhv:2204.09269, 2022.
Keyulu Xu, Mozhi Zhang, Jingling Li, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. arXiv:2009.11848, 2020.
Shunyu Yao, Binghui Peng, Christos H. Papadimitriou, and Karthik Narasimhan. Self-attention net- works can process bounded hierarchical languages. Association for Computational Linguistics, 2021.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 2021.
H Paul Zeiger. Cascade synthesis of finite-state machines. Information and Control, 1967.
Chiyuan Zhang, Maithra Raghu, Jon Kleinberg, and Samy Bengio. Pointer value retrieval: A new benchmark for understanding the limits of neural network generalization. arXiv:2107. 12580, 2021a.
Linjun Zhang, Zhun Deng, Kenji Kawaguchi, Amirata Ghorbani. and James Zou. How does mixup help with robustness and generalization? In International Conference on Learning Representa- fions, 2021b.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling Transformers with LEGO: a synthetic reasoning task. arXiv:2206.04300, 2022.
Karl-Heinz Zimmermann. On
theory for semiautomata. arXiv:2010. 2222200....26,..00 2020.
15

---

Published as a conference paper at ICLR 2023
Appendix
Table of Contents
A Additional background and notation
17
A.1 Notation
17
A.2 Automata, semigroups, and groups
17
A.3 Shallow circuit complexity classes
21
The Transformer architecture
22
A.5 Additional discussion related work
24
B Experiments
26
B.I Section 4: SGD finds the shortcuts, under ideal supervision
26
B.2 Section 5: Failures of shortcuts in more challenging settings
31
B.3
Additional details
37
Proofs
C.I Useful definitions and lemmas
38
C.2 Proof of Theorem 1: Logarithmic-depp depth shorteuts via parallel prefix
40
C33 Proof of Theorem 2: Constant- shorteuts via Krohn Rhodes decomposition
42
C44 Proof of Theorem 3: Even shorter shortcuts for gridworld
56
C.5 Proof of Theorem 4: Depth lower bound for non-solvable semiautomata
64
16

---

Published as a conference paper at ICLR 2023
A
ADDITIONAL BACKGROUND AND NOTATION
A.1
NOTATION
Below, we list our notational conventions for indices, vectors, matrices, and functions.
For natural number [xx] denotes the index set
For vector v R and E [n]. es denotes the i-th entry. When an expression, we use [v]; for clarity. Vectors can be instantiated by square brackets (like [L 2 3| R²). They can be indexed by slices: Vaab denotes [va Ca+1 1b). We adhere to the convention that all vectors are column vectors.
For a matrix M E Rmxx"ii [m].j E [n]: M, (or (M) denotes the (i. entry. My and M., denote the i-th row and j-th column, respectively. Importantly, we note the convention that this "slice" notation converts all vectors into column vectors.
When the first dimension of matrix M E RTxd is to be interpreted as sequence length, we will implicitly convert sequence of vectors Ui,..., UT Rd into a matrix (v),..., UT) E IR''x whose rows the vectors vt. This an arbitrary choice (compared to concatenating columns and obtaining matrix in R selected to adhere to previously standardized notation for the Transformer.
We will sometimes index vectors and matrices with named indices (such as for padding to- kens) instead of integers, for elarity.
e, denotes the i-th elementary (one-hot) unit vector. Likewise as above, we sometimes use non- integer indices (e.g. e.1).
For vectors u, Rd. (u, v) TTo both denote the inner product.
For function X Y + Z and all y E Y, we will let f(. X + Z denote the restriction of f to y (and similarly for other restrictions). This appears in the per-input state > Q, as well as the functions represented by neural networks for particular choice of weights.
transition functions ((),o)=
For functions f.g. Jog denotes composition: (fog)(o) == (g(()) When compose neural networks f x , + Y.g x 9 - 2 with parameter spaces Oy Oy, we will use fog: X (θ, x 6,)
indicate the composition f(g(;; 0,)0,))
()+: : + R denotes the ReLU (a.k.a. positive part) function: (x) max (2, 0). function compositions, we use to denote the entry-wise ReLU (e.g. f foo og.
A.2 AUTOMATA, SEMIGROUPS, AND GROUPS
Recall that semiautomaton A (Q, Σ, has state space Q, an input alphabet Σ, and transition function 5 Q Σ + Q. For any natural number T and a starting state 9o E Q, by repeated composition of the transition function 5, one can use to define map from sequence of inputs (01,...,TT) E ΣT to sequence of states (q1,...,47) (T) E Q" via:
4 5( Vt E [T].
Here and below, it is helpful automata. For a given semiautomaton we can always identify the state space Q with index set (1...., [Q}} and use a one-hot encoding of states into (0,1)91, Q1, For each input symbol a e Σ, we associate a transition matrix d(.,a) E {0, 13 [Ql with entries [o(,,a))y'' = 1(ő(g, o) = d'). This implies that for all g, a, we have 6g(g.)) 8(,, o)., so that the computation of the semiau- tomaton amounts to repeated matrix-vector multiplication.
matrix-vector notation express the computation of semi-
While semiautomata remarkably expressive, discuss a few simple examples throughout this background section to elucidate the key concepts.
Example (Parity). Let Q = Σ = {0, 1) and let 8(q. = 9 and 8(g. 1) = 1 Then, starting with 4o 0, the state at fime I, qr is 1 the binary sequence (Ø), (a)..... -.. ., 71) has odd number of 1s.
17

---

Published as a conference paper at ICLR 2023
Example 2 (Flip-flop). Ler Q -= and let d be given by
12x2, 8(,,1)= = (0 0) 6(.,2) = ( 1)
As the name suggests, this semiautomaton implements simple memory operation where the state at time t is the value of the most recent non-1 input symbol.
Example (ID gridworld). Let S be a natural number, Q = Then the transition matrices are given by:
{0, ,,...,,77 S} and Σ = 1, R).
!
Isx8
6(, 1) = =1+115++1 8(,,L)= =
R) =
1sxS
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 O and the increment input has no effect at state
Note that we have chosen a convention which differs slightly from the main paper (i.e. Figure I): we enumerate the indices starting from 0 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 semiautomaton A (Q, 8) each input symbol E Σ defines function 6(,, Q Q. These functions can be composed in the standard way, and we use 8(, F1:) to denote the t-fold function composition. Note that 5g0, dia) is precisely the value of the state time t on input 01:- Thus, the set of all functions that can be obtained by composition of the transition operator, formally
T(A)= (0(;, : E N, E E
plays a central role in describing the computation of the semiautomaton. This object transfor- mation semigroup. We now turn to describing the necessary algebraic background.
Recall that a group (G,-) a set G equipped with a binary operation : x G 9 such that
(identity) There exists an identity element E G such that e.g -g = g e = g for g E g.
(invertibilitty Every element g E G has inverse g 11 E 9 such that g g g- 1, g
(associativity) The binary operation is associative: (g) 92) ga = 91 (g2 ga).
A monoid is less structured than a group; there must be an identity element and the binary opera- tion 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 G be subset of functions from Q + 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 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.
which subset of the above properties hold.
In fact, taking 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) 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 G. 12 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.
More succinetly, inputs can correspond a generating set of the group, hut this is not relevant for results.
18

---

Published as a conference paper at ICLR 2023
Example 4 (Cyclic group). Let S be natural number let Q {0, 1.......,.. 1} and let E {1} have only one element. The dynamics are given by 8(q, 1) = (q + 1) mod S. Clearty this semiautomaton implements counting modulo S. The underlying group is the eyclic group, denoted Cs, 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.
Let us now turn to the compositional structure of groups and semigroups. Since it is without loss of generality to consider transformation (semi)groups, we always take the binary operation to be function composition. A subgroup H of a group G' is a subset of the elements of G that is also a group, denoted as H S G. In particular it must be closed under the binary operation. N is a normal subgroup, denoted N 4 G. if in addition to being a subgroup, it satisfies that {on : E N} = {ng- : n E N}. (These sets are known as the left and right cosets of N in G and denoted gN and Ng respectively.) 13 Normal subgroups can also arise as the kernel of mapping from G to a subgroup H of G. Let Φ :GG - H be a mapping that preserves the group operation (i.e., a group homomorphism) and let ker(Ø) == {g : ((g) id]. Then ker(ő) is normal subgroup of G. We will see below that normal subgroups provide a weak form of commutativity, that allows us to construct more complex groups out of simpler ones.
Direct products. 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 new group with elements {(g h) : g e G, h E H) with a binary operation that is applied component component-wise wise (g. h) (g', h²) = (g g', h²) (here, is overloaded to be the group operation for all three groups). This direct product group is denoted G x H. In the context of permutation groups, say G is a permutation group over ground set QG and H is over ground set QH. Then G x H has ground set Qc x QH and every function in G x H. factorizes component- i.e., every element in G x H is identified with permutation (9G. 4H) + (g(ga), h(qH)) where g E G. E H.
Observe that G' x H contains normal subgroups which are isomorphic to both G and H. To see this, take N = {ee, h) : h E H} where EG is the identity element in G. Then since geg = egg and since H is closed under its group operation, we have (g. N)N N(g. h) for all (g, h) E G x H. A symmetric argument shows that G is also a normal subgroup of the direct product.
Note that 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 non-interacting 1 --dimensional gridworlds:
analogously define direct products in the absence of the group axioms, and thus
Example 5 (2D gridworld). Gf is the transformation semigroup of the 1-d grid world with S + states, then Gs x Gs corresponds to 2-dimensional gridworld. A semiautomaton that yields this transformation semigroup has state space Q = {ii.j) : :i,j E S}} and 5 actions: increment or decrement i or j. subject to boundary effects, do nothing.
The definition of direct product extends straightforwardly to more than two terms G1 x identify the items with tuples (g), -.. 9.).
Semidirect products. However, it is possible to compose larger groups so that one of the sub- groups is not a normal subgroup. This operation is called a semidirect product. with the group law (g, (g) ,h) = (g @n (g'), h-hh) for some 6) to be defined later. Observe that in the direct product G x H. we have constructed the elements from ordered pairs (g E G, E 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)- (g', h!') which output an element of g and an element of h. An attempt would be to pick two arbitrary injective homomorphisms de. Φh which embed G and H into a "shared space," so that elements of G and H can be multiplied together:
h) (g', h²) ;= ( 6u (h) de(g') (H').)
'n equivalent definition of normal group is a subgroup N such that g ng N, Vg G,n
19

---

Published as a conference paper at ICLR 2023
However, need ensure that this group operation is closed. Since all elements of the group of the form (g. h) where g E G and h H. we must find pair (g.h) that yields the right hand side of the above display when embedded into the shared space via ΦG. Φu- For this, the most natural choice is 9 = g' and h = h- h', and thus we must check that:
volg g') @n h- 1/ = @c(g) dolg') on (h) Φn (1))
= Φn (h) )olg') Φn (h') = (g. h) (g',h)
However, the middle equality may not hold, because Φc and ΦE (h) are not guaranteed to com- mute. (Observe that for the special case of 9 + (g,en), + (ea, h), these two elements always commute, giving rise to the direct product.)
Fliding $a, Φn and simply using g: h as elements of the shared space, a sufficient condition for this hold is that hggh- G. since then, for some I E G.
(g. h) (g. h²) ghg' (h 'hhh' g(hg' h ' gghh', h)
which is of the form ( nn (-) since both G and H are themselves closed. This condition is precisely that G is normal subgroup.
There is a degree of freedom here: for each pair h and g', we can choose which element of G' is given by hg'h hg'h-". When we make this choice we must ensure all of the group axioms are preserved, e.g. when h eи we should always have engeen = g'. Suppose we make this choice and define ds. g + high1 1 E G (this & homomorphism from H + Aut(G), where (-) denotes the automorphism group, the group of bijections G that preserve the group axioms, under composition). Then, these ordered pairs do indeed form a group. but the group operation is
hg (g', gon(g'))il
This object is the semidirect product, and it is denoted G x H. Note that the choice of mapping is unspecified in the notation, and, in general, different choices of & will yield different structures for the semidirect product.
Finally, when G = N x H, both N and H are subgroups of G, but N is also a normal subgroup. To see this, we need check that hN Nh for any h H. This is equivalent for each h. n. but we defined the group operation to be hnh -1 on (n) E N, specifically so this would hold. On the other hand, H may not be normal subgroup. and this sense the semidirect product is a generalization of the direct product (for which both subgroups are normal). However, when the mapping á is trivial that dh (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.
Anh-1 E N
Example 6 (Dihedral group). Consider semiautomaton with Q = ( and input alphabet Σ (advance, reverse]. The transitions are given by:
S - 1 x +-1,+1}
8((s, b), advance) (s + b mod S,b)
5((s, b), reverse) = -b)
The trunsformation semigroup for this semiautomaton Cs X C2 where Os is the cyclic group on elements (cf. Example 4). C2 has two elements, the identity e and one element h such that hh e. Og has S elements where each element g is a function that adds some number E {O...., . S 1}
to the input modulo S. The inverse 9 is naturally to subtract k to the input, modulo S. The homomorphism & in the semidirect product is such that (e(g) = and on (r)()) g
Wreath products. We define one more type of product between groups N and H: the wreath product N H (N x... x N) x H. This is group containing [N| H elements (rather than [N] - H like the direct and semidirect products). Intuitively, is defined by creating one copy of N per element in H via the direct product, then letting II specify way to exchange these copies. Formally, N i H is the unique group generated by
(g....., MiHl h)
Vg. E N, H,
where
(0) (o1....;9 H1 en) (o/..... - MMM: en) == (0) -.g...,99H4 g(m)en) Vgu E N,
20

---

Published as a conference paper at ICLR 2023
and
(01.....9 HI CH) (e,,...,,y, h) ==
CH) Vg: E E H, (A.1)
where we have enumerated the elements of H in arbitrary order, such that each Th : [H] - [ is the permutation defined by right multiplication h'' + hhh (by convention).
To write this explicitly as a semidirect product N H == (N x... x N) X H. the homomorphism into the direct product's automorphism group & H - Aut(N x ... x N) is given by A. for each h E H. Φ is the automorphism defined by permuting the indices between the terms in the direct product, according to the permutation induced by right multiplication by h.
Example 7 (Rubik's ' Cube). A naive way to construct the Rubik's Cube is to assign labels {l....,54 54] to the stickers on the cube, and define the Rubik's Cube group G via the sticker config- urations reachable by the 6 face turns (which each specify a permutation ot, 8n, dU, dp. 5B, [54] - [54] the stickers). This establishes G as a subgroup of S54- First, notice that the 6 central stickers never move (so this is really improvable to Saa). Next, notice that the 24 = 8 x 3 vertex stickers never switch places with the 24 12 2 edge stickers. The vertex stickers form subset of the wreath product Cs Sg, while the edge stickers form a subset of the wreath product C2 i Sy2. In all, this realizes G as subgroup of a direct product of wreath products:
G S (C3 1 Ss) x (C2 2 S12).
Among other consequences towards solving the Rubik's 's Cube, this gives an improved upper bound on the size of G (which turns out to still be off by factor of 12, because of nontrivial invariants preserved by the face rotations, a.k.a. unreachable configurations).
Quotients, simple groups, and maximal subgroups. When N is normal subgroup of G, the quotient group G/N is defined as {gN: : E G} with binary operation (g N) = (gg') (gg')N. The fact that N is normal subgroup implies that this is a well defined group. We can also check that if G = N x И 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 coiiitivial normal subgroup, say N. We call N proper N * G. We call a proper subgroup N maximal there is no other proper normal subgroup N' G such that N < N'. Equivalently, N is a maximal proper normal subgroup if and only if G/N simple. This is akin to extracting a prime factor from 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.
Group extensions. Finally, we provide some additional terminology related to these different no- tions 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 is said to be an extension of N over H. The direct product G = N x H is known the trivial extension. A semidirect product G = N H is known as a split extension. However, not all extensions are split extensions; the smallest example the quaternion group Qs. the group of unit quaternions =j. =i, under multiplication (i2 o²= =2² 122 ijk- = - 1), which cannot be realized as a semidirect product of smaller groups In general, 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 -aloujnne universal embedding theorem (Krasner& & Kaloujnine, 1951) states that all extensions G can be found as subgroups of the wreath product N 2 H. The proof of Theorem 2 essentially shows how to imple- ment the different kinds of group extensions, given constructions which implement the substructures N, H. In the worst case, we will have to implement wreath product.
A.3 SHALLOW CIRCUIT COMPLEXITY CLASSES
We provide an extremely abridged selection of relevant concepts in eircuit complexity. For a system- atic introduction, refer to (Arora & Barak, 2009). In particular, we discuss each circuit complexity class and inclusion below:
NC c AC c ACC C TC C NC!
21

---

Published as a conference paper at ICLR 2023
NC" is the class of constant-depth, constant-fan-le polynomial-sized AND/OR/NOT circuits. a constant-depth Transformer only uses the constant-degree sparse selection constructions in (Edelman et al., 2022), it can be viewed representing functions in this class. However, the representational power of these circuits is extremely limited: they cannot express any function number of inputs growing with 7.
which depend
AC" is the class of constant-depth, unbounded-fan-in polynomial-sized AND/OR circuits, al- lowing NOT gates only at the inputs. A classic result is that the parity of T bits is not in AC (Furst et al., 1984); Hahn (2020) concludes the same for bounded-norm (and thus bounded- Lipschitz-consstant) constant-depth Transformers.
ACC extends AC with an additional type unbounded fan-in gate known as MOD, for any prime number which checks if the sum of the input bits is a multiple of Theorem comes from the fact that the semigroup word problem (which is essentially identical semiautomaton simulation) is in this class; (Barrington & Thérien, 1988).
TC extends AC with additional type of unbounded-ffa--in gate called MAJ, which computes the majority of an odd number of input bits (a threshold gate). is straightforward to simulate modular counters using polynomial number of parallel thresholds (i.e. ACC" C TC)). Whether this inclusion is strict (can you simulate 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).
NC¹ is the class of Olog dee-thh constant-fan-in, polynomial-sized AND/OR/NOT circuits. It is extremely popular and natural complexity class capturing efficiently parallelizable algo- rithms. It is unknown whether any of the inclusions in the "larger" classes TC C NC SLSP 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.
An L-layer Transformer is sequence-to-sequee sequuence network für RFxx Oir Rxxd consisting of alternating self- relf-attention blocks and feedforward blocks
fet = Tup p(L) o (((--)) o... (()
The parameter space Oi is the Cartesian product of those of the individual blocks (without recurrent weight sharing across layers, by default). We define these two types of blocks below.
Attention. single-headed (H 1) self-attention block is a sruuenc-to-sequemee network fatu RTxd x Oatta + R T xd parameterized by 8 (Wq. WK. Wv, We). With an inner embedding dimension k the shapes of these matrices are as follows: WQ. WK. Wy.WT E Rdxk Each head, indexed by t E [T], computes pairwise query-key alignment scores (W TT, W WT-re)) for each position t' [T], normalizes them with a T T-dimensional causally-masker softmax (forcing weights for positions & 1' to 0), and uses these attention weights a E RT to mix value embed- dings: Ere(Ty O4r W ... This mixture is mapped back into Rd by multiplying by Wc, to form the t-th row of the output matrix. In single equation:
(X; Wo. WK. Wv, Wc) := CausalAttn(XWog) X XXW,
where CausalAttn Rxx R'x applies a row-wise causally-masked T T-ddimensional softmax function. The standard softmax function softmax(2) RT. + RT is defined by
e
[softmax z)], Everyy e
the massdd softmax row is defined to be softmax(214) the first t coordinates, and O on the rest. To implement the causal masking operation, it is customary to set the en- tries above the diagonal of the attention score matrix WWoW to -os, then obtaining GausalAttnXXWoWe XT) via row-wise softmax (letting -c evaluate to 0).
22

---

Published as a conference paper at ICLR 2023
In general, for any positive integer H, multi-headed self-attention block consists of sum of H copies of the above construction, each with its own parameters.
This component is often called seft attention: the softmax performs continuous selection, taking convex combination its inputs. In contrast, hard attention refers to attention heads which perform truly sparse selection (putting weight on the position with the highest score, and 0 all others).
Feedforward MLP. An -laye position-wise feedforward MLP block is a sequence-to-ssequenc network fmip Rxxx x mip Rx'xd, parameterized by Omlp (Wi, by....... WT, by.). For choice of activation function : IR + R (which is always ReLU in our theoretical constructions, for simplicity), fulp applies the same nonlinear map (x + W + by.) oo... ........ oo (z + Wix+ fix++b)) each row of the input matrix X Rxxd (with the same parameters per position E); here, is applied pointwise.
Finally, an extra term P E Txd added to the first layer's imput, the matrix of position encodings.
Residual connections. It is typical to add residual connections which bypass each block. That is, letting id denote the identity function Txx the network (with position encodings) becomes
far (id + ((l))) (id + filta (id / o ... (id + flll)) + (id + P).
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 non-residual network:
x (z)+ (-x)+-
Similarly, residual connection be implemented with one attention head (with internal embed- ding dimension = d), as long as is able to select its own position (which will be true in all of constructions).
In some of our constructions, we choose to use residual connections (sometimes restricted to certain dimensions); it will be very natural to view the embedding space Rd "workspace", where residual connections ensure that downstream layers can access the input (and position) embeddings, well outputs of all earlier layers. We will specify whether use residual connections in each construction, to make the proofs as clear as possible. When we do so, we do not add the extra weights explicitly fatta and fuip.
Layer normalization. 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: sim- ulate a semiautomaton at length T, we will choose to prepend tokens, with explicitly chosen embeddings, which do not depend on the input 01:1- Theorem 1 uses = e padding, and Theorem uses 7 = 1. In both cases, padding is not strictly necessary (the same functionality could be implemented by the MIPs without substantially changing our results), but we find that it leads to the most intuitive and concise constructions.
Padding tokens.
Complexity measures. 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 T' and embedding dimension d. Up to factor of bit precision, this corresponds to the number of inputs a classical Boolean eircuit. We will exclusively define architectures where d is independent of T.
network: sequence length
Its depth L, the number of repeated fulp o fasta blocks. When each of these modules contains a constant number of sequential computations, this coincides with the usual notion of circuit depth, up to a constant factor. This is true in practice and our theoretical treatment (the attention and MI. have constant number of layers).
23

---

Published as a conference paper at ICLR 2023
The other shape parameters from the definition of the architecture: number of heads (per layer and position) H44 and internal embedding dimension k. When filp is I'-layer MI.,, it has MLP intermediate widths d1....,dk d We will exclusively think of L' as a small constant, so that the number of sequential matrix multiplications in the entire network is within a constant factor of L.
Its attention width Wattn is defined to the maximum of {d, Hk), and its MLP width Wmip is defined the maximum of ..i.....,,---13 Taking upper bound we will use to summarize the number per-position trainable embeddings in our constructions. To map this to the usual notion of circuit size, note that the computations repeated position-wwr wise. Thus, Transformers induce a computational graph with Ó(T L w) gates and O(T I, x² wires. The position-wise parameter-sharing induces special notion of circuit uniformity.
max Wattn, Wmlp) as coarse
bound on its oo-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 this paper go beyond the sparse variable creation constructions of bounded-nowe attention heads; in general, the norms scale with T. The attention heads ex- press meaningful non-sparse functions. Aside from the positive experimental results, we do not directly investigate generalization in this paper.
The bit precision (length of finite- finite-precisio precision truncation of real numbers in a computational graph implementing fir), which lets us implement approximate real-valued computations as Boolean (or discrete arithmetic) circuits. With infinite-precision real numbers, there are pathological constructions for RNNs (Siegelmann & Sontag, 1992) and Transformers (Merrill et al., 2021) which give single parameters of neural networks infinite representational power. Throughout this work, our circuits will work with O(log T) bit precision, which can represent real numbers (as integers [x: 2e] in their binary representation, for some choice of c O(log 7)) with magnitude up to (ppoy(TT)) with OLLppoly(T)) approximation error. Since this is far from the focus of our results, we will clide details for the remainder of this paper, returning these considerations only to make Theorem more concrete. All of our constructions are robust up this noise level: this is because the internal weight norms and activations bounded by quantity at most polynomial in T", and the function approximation construction in Lemmas and 2 can tolerate (Tly(TT) perturbations using poly(T) weight norms.
A.5 ADDITIONAL DISCUSSION OF RELATED WORK
Relevant applications. We first provide references for the "reasoning-like" applicatioon 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 & Charton, 2019; Polu & Sutskever, 2020; Drori et al., 2022).
Neural dynamics models for decision-making: recurrent (Hafner al., 2019; Ye et al., 2021; Micheli et al., 2022) and non-recurrent (Chen et al., 2021a; Janner et al., 2021).
Synthetic combinatorial experiments (and relations). 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, 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 function that is based on the sum (for example, mod n). Otherwise it would require the window size number of attention heads similar to grid-world construction.
"There will be notational collision hetween h, H denoting attention heads, and h H denoting an element in group. We keep the overloaded notation for clarity, and this will certainly he unambiguous.
24

---

Published as a conference paper at ICLR 2023
LEGO: Zhang al. (2022) propose task based on solving a simple Mfaai--of-resonning 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-sequune task, however it serves another example of the emergence of "shorteut" 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 Bounded-depth Dyck can be captured by our semiautomata formalism and our main construction would recover the depth-2" 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 construction, and implement proximity-based depth matching in the second layer. Our grid-world construction generalizes their construction to a significantly more complex problem. We view work as a generalization of their results to a wider class of semiautomata.
depth- solution for bounded-depth D Dyck languages.
Parity: Another commonly studied synthetic setup is the task of learning parities. Edelman 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 & 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 & Lieberum (2022) focus on the task of adding two 5 digit numbers modulo 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 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 7).
Formal languages and neural networks. Dyck languages are particularly interesting for their completeness property: the Shoomsky-Schitzenberre representation theorem (Chomsky & 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 (log deppph networks can simulate all context-free lan- guages (Theorem 1), and dlepphh networks can represent some of them. The obstructing reg- ular 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 AC circuits (Chandra al., 1983; Barrington & Thérien, 1988). We did not see a way for this to generically entail d1ee-th 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-augment 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 perfor- mance on "inductive inference", which is similar to our length generalization setup where we also see the failure of Transformer.
RNNs and Transformers on
Different axes of generalization: length, size, and algorithmic. 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) empiri- cally investigate the ability of pre-trained Transformers to generalize to longer sequence length for parity-like 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. ); Bansal et al. (2022) focus on length and algorith- mic generalization for recurrent models where they train simple/easy instances of the underlying
25

---

Published as a conference paper at ICLR 2023
problem and evaluate harder/complex instances using the power of recurrence to simulate tra computational steps, inspired by the ideas of Neural Turing Machines (Graves et al., 2014) and Adaptive Computation Time (Graves, 2016). We view our results 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).
Recurrent Transformers. Our work is not the first to notice that Transformer architectures make brittle predictions autoo--distribution. Indeed, even the seminal paper introducing the architecture (Vaswani et al., 2017) notes that length generalization is promoted by subtle hyperparameter choice (namely, the positional encoding scheme). Furthermore, there have been several attempts to recon- cile 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 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.
Recurrent non-recurrent sequence transduction. As mentioned briefly towards the end of Section 5, the setting of indirectly-uupervised semiautomata matches that of autoregressive gener- ative 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 al., 2021), where the possible continuations are {all possible open brackets, if the stack qi is not full} U (close bracket which pairs with the top of the stack q: J. We note that when an autoregressive model is used for sequence generation via a baken-b-tooken 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 recurrent network, whose state is the completed prefix of the current generated sequence. Non-autorregeessive 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 expressivity barrier to expressing deep recurrent linguistic primitives, but there may be issues with out-of-distribution robustness.
Algebraic structures in deep learning. Another area where tools from abstract algebra are used reason about neural networks is geometric deep learning, research program which seeks to understand how to specify inductive biases stemming from algebraic invariances. For a recent sur- vey, see Bronstein al. (2021). In contrast, this work studies the ability of fixed architecture to learn 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.
Theoretical role of depth. Our theoretical results can be interpreted as a depth separation result: contingent on TC" # NC². it takes strictly more layers simulate non-solvable semiautomata, compared to their solvable counterparts. In a similar spirit, there have been several works establish- ing depth separation for feed-forward neural networks (mostly using ReLU activations) (Telgarsky, 2016; Eldan & 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 I, but would require exponential-winn for depth I, (or VT., depending on the result).
B EXPERIMENTS
B.1 SECTION 4: SGD FINDS THE SHORTCUTS, UNDER IDEAL SUPERVISION
This section contains a full description and discussion of the in-distribution simulation experiments from Section 4.
26

---

Published as a conference paper at ICLR 2023
B.1.1 SHALLOW TRANSFORMERS SIMULATE SMALL GROUPS AND SEMIGROUPS
The main experiments in this paper investigate whether gradient-basee training of Transformers finds low-depth solutions to the problem of simulating semiautomata. In these experiments, we con- sider wide variety of semiautomata A, corresponding to various groups and semigroups, and con- struct a distribution DA over imput sequences (01....07 oT) and their corresponding state sequences (q),..., . = AT. (II:T). ) In each setting, the a are chosen uniformly at random from the set of valid tokens in Σ. 15 Given this distribution D.A, and a teneennetto-squeeee
neural network (with token embedding and linear classification head) which maps Σ" to token predictions Y R 1g"'x(Q)
(such that Y,,- = (g. = 1.)). we establish the task of minimizing the cross-entropy loss
L( = Σ (/
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 semi- group is a group), every ge depends on every preceding input {or t' f}. Indeed, this is why previous studies have used group operations as a benchmark for reasoning (Anil et al., 2022; Zhang al., 2022).
Settings.
We proceed to enumerate the semiautomata considered in these simulation experiments.
Cyclic groups Cz Ca...., Cs. For each cyclic group C, (realized Q f0, 1....... 1} under mod-n addition), we choose the generator set Σ to be the full set of group elements [0....., 1}. An alternative could be to let E be minimal set (0, 1}, which we do not use in the experiments.
Direct products of cyclic groups C2 x C2 C2 C2 x Cz. realized as concatenated copies of the component semiautomata. Note that Cs (which is isomorphic to C2 x C3), included in the above set, another example.
Dihedral groups D6, Dg. Our realization of Dan chooses Q {0, 1...., 1} x {0, and
== {(1,0), 0), (0, })}. 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 A1, S1, As, Sz. We choose Q to be the set of n! permutations for Sn (sym- a! metric group), and Q to be the set of 7 even permutations for An (alternating group elements). The generator set for S, consists of the minimal generators, a transposition and an n-cycle, well as 6 other permutations. (12i) for i E {3, ... .n. Note that A4, S1 are solvable (leading to constant-depth short- cuts), while Aş, S5 are not. Also, note that to learn constant-depth shorteut for As, model needs to discover the wondrous fact that A4 has a nontrivial normal subgroup, that of its double transpositions.
For An, choose the 3-cycles of the form
The quaternion group Q8. This is the smallest example of 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.
The Dyck language Dyck,,, (correctly nested brackets of & types, with depth at most n). We take 4, k 2 in the experiments. To realize Dyck,, as semiautomaton simulation problem, the state Q is the state of the stack which implements Dyck language recognition (there thus Ei ks distinct states); Σ is the set of 2 opening and closing brackets. The
Taak for instance the Dyck language, if the current stack is empty, then o. chosen uniformly from the choices of open parentheses but not the closing parentheses. This is in accordance with Yao et al. (2021). In the sense that it induces non-trivial learning problem this group. If we only pick the generator (1}, the output sequence is deterministic, and there no learning problem.
Thesee other permutations are chosen following the ordering given by the sympy hompy.combinatorics package. They are not necessary for covering the state space (since the minimal set of 2 permutations already suffice cover Q). but can help speed up the mixing of the states.
18 the experiments we use (I: + 1)" classes (i.e. each of the n positions can take (k + 1) possible values), En 2 of which reachable.
27

---

Published as a conference paper at ICLR 2023
distribution in inputs is slightly different, since there is a notion of "illegal" inputs: if the stack is empty, then the set of feasible inputs contain all the opening brackets; if the stack is full (i.e. reaching depth n), then the only feasible input is the closing bracket for the opening bracket at the top of the stack.
Gridworld semiautomata Grid,, Gridg, where Q = -0, ... 1....................
n - 1 (for n = 4 or 9) and == (+1). 19 For this special case, we have a constant-depth solution as stated in Theorem 3.
Training. We focus on the online learning setting for all experiments in this paper: at training iteration i, draw a fresh minibatch of samples from DA. compute the network's 's loss and gradients on this minibatch, and update the model's weights using a standard first-order optimizer (we use AdamW (Loshchilov & 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 experimei
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 sinusoidal positional encodings (Vaswami al., 201 7) unless otherwise specified.
Sequence length. 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 Σ 100 1030 sequences, rendering it overwhelmingly unlikely for 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-focuse studies.
Depth. 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 7)) and (1), nor do we attempt to tackle the problem of exhaustively enumerating and characterizing the shortcut solutions for any particular semiautomaton.
Results. For each task and number of layers, we report the highest (Figure 5) and median (Figure
6) accuracies over 20 runs. The accuracy is calculated at token level (i.e. $ Eie(r) 1 = (.), as opposed to the sequence-level accuracy (i.e. 1 [ = Q1:T) as reported in Bhattamishra et al. (2020). We evaluate in in-distribution accuracy on independent (unseen) samples of DA. which contain 2048 sequences of length T = 100. 20 As shown in Figure 5, Transformers, trained with standard gradient-based methods, are able to find solutions which generalize well (in-distribution) on all of the tasks. Performance tends to improve as the number of layers increases (there is small amount of non-monotonicity in some settings due to training instability); the sufficient depth to achieve high accuracy varies depending on the problem setting, as discussed below.
Trends in sufficient depth. The minimum number of layers required to achieve 99%+ perfor- mance 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 learn, with the non-solvable group Ss requiring the largest depth."! Between the non-abelian groups, the difficulty of learning Qs com- pared to Ds (which has the same cardinality) agrees with our theoretical characterizations of the respective constant-depth shortcuts for these groups: Ds can be written as a semidirect product of smaller groups, while Qs cannot, so our theoretical construction of a constant-depth shortcut must embed Qs in larger structure (i.e. the wreath product).
Improving training stability. Throughout these experiments, we observe the following forms of training instability: high variance in training curves (based on initialization and random seeds for
19 for T., for R. We omit the no-op in the experiment which does not change the difficulty of the task. T0Tis size is sufficient for evaluating the model performance: for example, for Co (i.e parity), evaluating of this size gives standard deviation of 0.031% in the accuracy.
model on 10 evaluation
However we stress that these experiments do not control for the fact that larger groups have richer supervision (for example, As has more informative labels than A4). possibly accounting for the counterintuitive result that the latter requires more layers, despite being subgroup of the former.
28

---

Published as a conference paper at ICLR 2023
1
3
4
5
7
8
10
11
12
13
14
15
16
Dyck
99.3
100
100
100
100
100
100
100
100
100
100
100
100
100
100
100
Grid,
100
100
100
100
100
1.00
100
100
C2
77.6
99.8
99.9
100
100
99.5
100
99.7
100
100
100
100
100
100
100
1000
54.66
94.6
96.7
99.4
100
100
99.8
100
99.9
100
100
100
1.00
100
99.8
100
C
89.0
99.1
99.9
100
100
100
100
100
100
100
100
1000
1.00
100
100
100
Ce
59.8
96.7
75.5
99.9
99.8
99.9
99.9
100
100
100
99.8
99.9
1000
99.8
99.9
99.9
90.9
95.0
99.9
99.9
100
100
100
100
100
99.8
100
100
100
100
C3
90.5
98.8
99.9
100
100
99.9
100
100
99.9
99.9
100
100
1000
100
100
100
C3
65.0
77.9
99.9
97.9
100
99.8
98.2
99.9
100
100
91.9
95.9
91.7
90.6
87.5
80.6
D2
25.4
47.4
75.2
100
100
100
100
100
100
100
100
100
100
100
Qa
31.6
49.2
59.6
60.4
73.5
99.3
100
100
100
100
100
100
100
100
100
100
A4
25.0
35.4
49.1
59.3
62.6
82.3
90.9
98.0
98.0
99.1
99.8
100
99.7
100
100
100
12.5
23.1
32.5
46.7
71.2
98.8
100
100
100
100
100
100
100
100
100
27.1
72.5
S,
7.9
11.8
14.6
19.7
26.0
28.4
32.8
51.8
86.3
94.8
90.2
97.2
99.3
99.1
99.9
99.9
Figure 5: complete version Figure 3, for various tasks (rows) and numbers of network layers (columns). Reported performance is the maximum test accuracy over 20 runs.
29

---

Published as a conference paper at ICLR 2023
1
3
4
5
7
8
10
11
12
13
14
15
16
Dyck
98.8
100
100
100
100
100
100
100
100
100
100
100
100
100
100
100
Grid,
100
100
100
100
1.00
100
C2
56.4
83.0
79.9
80.9
89.1
85.2
84.8
84.9
88.8
94.5
98.3
86.4
90.4
88.7
94.6
99.3
40.3
69.1
78.2
85.0
84.0
84.9
87.9
96.2
99.4
89.2
82.5
99.3
87.4
98.0
89.6
92.0
C
75.6
62.7
99.0
99.5
99.8
99.9
99.8
99.5
99.8
99.8
99.7
99.8
99.8
99.9
99.9
99.7
Ce
45.8
49.0
53.0
59.6
75.5
77.0
95.6
91.2
83.4
59.6
98.4
72.9
89.7
94.5
99.8
87.5
51.0
99.7
99.7
99.6
99.4
99.7
99.7
99.6
99.6
99.6
99.7
99.6
99.8
99.7
C3
62.6
73.1
76.4
73.4
74.9
79.8
84.1
82.4
77.0
70.6
69.0
71.9
70.6
76.9
68.3
59.3
c3
50.0
61.4
60.6
60.7
72.4
63.2
63.8
66.4
69.8
59.0
63.4
54.6
59.5
53.0
44.7
48.4
D2
24.8
26.8
40.8
57.2
81.3
91.6
100
99.6
100
93.0
96.2
100
97.7
99.6
99.3
Qa
29.0
45.8
38.5
42.7
57.4
79.5
84.7
89.2
95.9
98.1
98.8
97.8
99.8
98.3
98.8
99.4
A4
19.7
30.4
41.0
45.4
44.7
52.8
60.0
68.3
72.8
74.1
91.4
82.6
88.2
97.9
99.0
98.55
10.5
26.6
30.5
40.6
63.9
77.2
99.4
99.3
100
100
100
100
100
99.9
100
22.9
S,
7.1
11.0
13.1
16.5
20.9
24.3
29.4
37.6
40.1
59.0
60.4
91.3
91.2
94.6
98.00
99.1
Figure 6:
The median accuracy for various tasks (rows) and numbers network layers (columns). Reported performance is the median test accuracy over 20 runs.
30

---

Published as a conference paper at ICLR 2023
the maadient-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 ((b),cc) and in the significant difference between the maximum accuracies in Figure 5 and the median in Figure 6.
To stabilize training, experiment with dropout and exponential moving average (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.33) 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 y 0.9 in the experiments since the performance of the EMA model does not seem to be sensitive to the choice of y E {0.85, 0.9, 0.95). Further, increasing the patience of the learning rate scheduler can be helpful.
1.1.2 VISUALIZING AND INTERPRETING ATTENTION HEADS
Although defer fine-grained mechanistic interpretability study ("which group/semigroup fac- torizations did these shallow Transformers discover, if any?") to future work, we provide some preliminary visualizations of attention heatmaps which strongly corroborate their theoretical coun- terparts. In particular, consider the gridworld setup in Theorem 3. The theoretical construction consists of two steps: the first attention layer calculates the prefix sum of the actions (i.e. the sum of (E E {0, + and the second attention layer identifies the last time the process is at a boundary state (i.e.
S) where the process can be "reset" (i.e. the model can ignore the history before the boundary state and only needs to calculate the sum of subsequent actions).
We have seen in Figure 3 that the network indeed learns to 1) compute the prefix sum, as evidenced by the uniform attention in the first layer, and 2) detect boundary states, as highlighted by large attention scores in the last layer. Figure 7 provides more examples of attention patterns, which are taken from the last layer of a 4-layer GPT-2 model on two randomly selected Gridg sequences. We highlight the locations where the process is at a boundary state (white strips for state or gray strips for state S 8), which align well with the highly activated positions of the attention heads, showing that the model learns to locate the closest boundary states. Moreover, when processing tokens appearing later in the sequence than these highly activated positions, no attention weight is put on tokens before these positions. This suggests that these highly activated locations reset the state so that the model does not need to look further back past them.
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 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 relaxed. We consider training under indirect (Section B.2.1) or incomplete (Section B.2.2) state supervision, and evaluation sequences that is dusti--disti (Section B.2.3) of longer lengths (Section B.2.4).
B.2.1 CHALLENGES FROM INDIRECT SUPERVISION
One type of limited supervision is that the observations may not provide full information of the underlying state. To model this, consider the case where instead of observing the state directly, we get function of the state, denoted v(g), where & Q + Q is non-injective (i.e. IQ] < IQ). In each of the experiments involving partially-observable semiautomata, we specify the underlying semiautomaton, as well as the observation function yp.
Dyck language with stack top observations: For Dyck,, the state Q is the state of the stack which takes Σn= K2 values. We take &p to be the function that takes in a stack and returns the element at the top of the stack, which is either one of the & open brackets if the stack if non-
We use the EMA implementation from http://ggihuhh.co//aaeee///oopop ema.
31

---

Published as a conference paper at ICLR 2023
Figure 7: Visualization of the entire set of attention heads on two randomly selected length- 128 sequences for models trained on Gridg. The lower triangles visualize the attention patterns, while the upper triangles are ignored by the attention head because of the causal mask. We use the upper triangles to visualize the positions position of state 0 and gray strips mark state = 8. Example heads that clearly detect state 0 and are highlighted with blue and red frames, respectively. Note that in many cases, the white/gray strips align with the locations of high attention scores (the bright yellow patterns). This suggests that the model indeed learns to identify the boundary states and that the construction in Theorem agrees with solutions found in practice.
state and state S 8 in the output: white strips mark the
32

---

Published as a conference paper at ICLR 2023
Task
1yck
Gridg
S5
C4
Dg
(abab)", (1)
(abab)", (2)
Observation
stack top
boundary
(1:
1o mod
location
accept
accept
Accuracy
100.0
100.0
99.6
99.9
100.0
100.0
100.0
Figure 8: Accuracies with indirect supervision. extending results in Figure 4(a). The numbers are the maximum over 25 runs. As a reference, L.TTM gets 100% on all tasks.
empty, a special token 1. indicating an empty stack, i.e. Q (1,2..... ., We consider
L). Gridworld with boundary observations: We consider the case where the underlying semiau-
k = 8 (as opposed to k = 2 in Section 4) to make the prediction task more challenging.
tomaton is Gridy with Q {0, 1. .8). The observation function ip Q - {0, outputs whether the current state is of two boundary states, i.e. at state 0 or state 5 =
Permutations with single-element observations: We take the permutation group S3 with Q is the set of 5! operations. The observation function vp of the permutation. For example, (2( 1, 4,5,) = 2. We use a set of 5 generators for the experiments.
{1, 3, 4, 4.5) returns the first value
Cyclic group with "O mod 4" observations: We take C1 as the underlying group with Q 1,,1,2.3), The observation function computes whether the current state state 0, i.e. (p(g) = =]
Dihedral group with rotation component only: Recall that D2n = Cn x C2. We take n 4 with Q = 1,, x {0, 1}, and let the observation function & output only the the first component (i.e - (0,
(abab)": We consider one semiautomaton which is not featured Section the one which recognizes the regular expression (abab) ", which is also studied in Bhattamishra et al. (2020). The underlying semiautomaton has 5 states: 4 states are in a cyclic fashion when seeing repeated patterns of abab, and a fifth absorbing "failure" state is entered if any other pattern is example, the input sequence ababahaaabab corresponds to states 01 012301244444, where the 5th "a" leads the absorbing state. The observation function yp computes whether the current state is the "accepting" state (i.e. state 3), with Q = {0, 1}. For example, the output of yp for the input sequence ababababa is 000100010, and the output for the input sequence ababbabab is 000111111, i.c. the sequence enters the absorbing state at position 5 and never recovers.
For
We consider two distributions on the input sequences: the form abababa ... (i.e. the process never in the absorbing state), which is the setup in Bhattamishra et al. (2020); and (2) the input is of the form abababa ... with probability 0.5, and is some randomly drawn string of a, b otherwise. Note that case (1) be solved purely based on the positional encoding, since the label is when the position is a multiple of and otherwise, while case (2) is more difficult since the model needs take into account the input tokens.
(I) the input is always a sequence of
Results. We train GPT-2-like models on sequences of length 40. We use 16 layers for S; and 8 layers for other tasks, with embedding dimension d = 512 and H 8 attention heads. As shown in Figure 8, the model is able to achieve near-perfect in 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. P4 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.
B.2.2 CHALLENGES FROM INCOMPLETE SUPERVISION
Another challenge of limited supervision is that the observation sequence may be incomplete, that is, may not be able to get supervision the states every time step. We consider the task of learning length 100 sequences, where the state at each position is revealed with some probability Preveal E (0, 1).
33

---

Published as a conference paper at ICLR 2023
1.0
1.0
C2 Accuracy
Ss Accuracy
0.8
0.5
GPT
0.6
LSTM
0.0
1/32
1/16
1/8
1/4
1/2
1
1/32
1/16
1/8
1/4 Preveal
1/2
Preveal
Figure 9: Learning from incomplete state sequences, extending results from Figure 4(b): accuracy nosition-wise probability of hidden token (i.e. Preveal), for GPT and LSTM. While LSTM able to maintain a perfect accuracy across different values of Preveal- GPT's performance may degrade labels get sparser. The mean and standard deviation are taken over 25 runs.
1.0
C2 Accuracy
0.9
GPT
GPT
0.8
LSTM
0.8
GPT Shifted positions
Scratchpa++recemt
0.7
0.7
0.05
0.25
0.50
0.75
0.95
6
14
22
30
38
Pr[o=1] 1]
Number of 1s
Figure 10: OOD generalization performance on C2: (Left) Accuracy sequences of the same length as training, with varying Pr(a = = 0.5. GPT fails at OOD generalization, whereas recurrent solutions implemented by LSTM and Scratchpad with recency bias is robust to different distributions. (Right) Accuracy on sequences of the same length as training, with a varying number of 1s in each sequence. GPT has worse performance counts less frequently The lines show the mean accuracy (with shadows showing standard error) over 25 replicates.
during training.
Results. Figure 9 shows the accuracy against Preveal. for S₂ and C2 (i.e. parity). Transformer training pipeline is worse than LSTM at tolerating incomplete supervision: while Transformer is able to maintain the performance across Preveal for C2, the performance degrades significantly lower Preveat for S5. We leave improving the robustness to sparse supervision to future work.
B.2.3 OUT-OF-DISTRIBUTION GENERALIZATION
The previous subsections show positive results on learning shallow non-recurrent shorteuts with limited supervision during training, either in the form of indirect observations or incomplete obser- vation sequences. this section, we study challenges at test time, and evaluate Transformers on their nut-of-distributte are trained in the standard way with full state supervision. The training sequences are of length 40, where each position has equal probability of being 0 or 1, i.e. Pr(a 1 0.5. At test time, the sequences of the same length as training, but the Bernoulli parameter Pr(a = 1] varies in the range (0.05, 0.1, 0.15, 0.9, 0.95).
generalization performance. For this and the next subsection, the models
Negative results for vanilla Transformers. Figure 10 (left) shows the accuracy as Pr/a = 1] varies. The performance 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 that Transformer may learning 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 Pr(o = 1] deviates from 0.5, it is less likely for the value of the count (which concentrates around T x Pr(o = 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 Pr(o = 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 10 (right), Transformer's performance
34

---

Published as a conference paper at ICLR 2023
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 shorteut 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 lower count is equally likely to appear at 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.
Guiding the Transformer to learn the recurrent solution. We investigate one established mit- igation for the out-of-distribution brittleness of non-recurrent Transformers: scratchpad train- ing and inference. Given a sequence of inputs (01,..., .T) and states (f1...., ,T), in the standard (non-recurrent) sequence-to-sequees input, and outputs the sequence of predictions for qt. In scratchpad training (Nye et al., 2021; Wei et al., 2022), we instead feed the network an interleaved sequence of inputs and states (T1: 91, 02 92, T3. 93...., "...,, (with an appropriately expanded token vocabulary), and define the network's state predictions to be those at the appropriately aligned positions: (q). 92, L....,,,,, (r) (T) (where 1 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 q:-1; one can think of this as a way to guide a shallow Transformer to learn the recurrent solution (i.c. explicit (T) iteration of the state transition function), rather than a shorteut.
learning pipeline, the network receives our as
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 shorteut as before. The good news is that we can couple scratchpad with an explicit recency bias in the attention mech- anism (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 cur- rent 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 distri- bution shift; see Figure 10 (left). Notice that this mitigation completely foregoes the computational advantage of 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 help without scratchpad, since in this case the output can depend uniformly on each input positions (c.g. consider parity).
B.2.4 LENGTH GENERALIZATION
Settings for length generalization. Our final setup is length generalization, where the model is evaluated sequences of lengths unseen during training. 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 and more discussions on length generalization in A.5. In the following, we check the length generalization performance on Dyck,, and C2 (with Pr(o = 1 = 0.5), where the model is trained on sequences of length 40 and tested on sequences of length {8, 24, 24, ...,, 120, 128).
Results. Figure 11 shows the performance on sequences of various lengths. In contrast to LSTM's perfect performance on all scenarios, Transformer's 's accuracy drops sharply as we move to lengths unseen during training. This 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-thoure 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 -ut-of-distribution and to longer sequence lengths (Figure 11,
35

---

Published as a conference paper at ICLR 2023
GPT
GPT Shifted positions
GPT
No positional encoding
LSTM
Scratchpad
Scratchpad+reeeeccy
In distribution
1.0
Dyck Accuracy
C Accuracy
0.9
0.9
0.8
0.8
0.7
0.7
20
40
60
80
100
120
20
40
60
80
100
120
Sequence Length
Sequence Length
(a) Mean accuracy over 25 replicates, with shadows showing standard error.
1.0
1.0
Dyck Accuracy
C. Accuracy
0.9
0.9
0.8
0.8
0.7
0.7
20
40
60
80
100
120
40
60
80 100
120
Sequence Length
Sequence Length
(b) Max accuracy over 25 replicates.
Figure 11: Length generalization on Dyck and C2: Transformer fails to generalize, but adding Scratchpad (Nye et al., 2021) and recency bias (Press al.. 2022) serves remedy. For "GPT- Shifted positions", the positions a sequence are shifted by a random number. For "GPT-No positional encoding", position encodings are provided but the causal mask is still present.
COS (T)
cos (27)
linear (T)
linear (27)
1.0
1.0
(abab) Accuracy
Grids Accuracy
0.9
0.9
0.8
0.8
0.7
0.7
1e4
2e4
3e4
4e4
1e4
2e4
3e4
Training steps
Training steps
Figure 12: Choice of the positional encoding: while having similar or even superior in-distribution performance (on sequences of length T' 40), sinusoidal positional encoding may suffer larger generalization gap than linear positional encoding when testing on length 27. The lines show the mean accuracy over 25 replicates.
yellow curves). The results also confirm that the inclusion of recency bias necessary: without it, scratchpad training shows no improvement on length generalization.
Impact of positional encoding. Figure 1l also shows some interesting findings related to posi- tional encoding. which is believed to be a key component for Transformers and a topic with active research (Ke al., 2020; Chu al., 2021). While this work does not aim improve positional encoding, some of our results may be of interest for future research.
Sinusoidal linear encoding: We find that the conventional sinusoidal encoding (which is the default for results in this paper) seems to generalize worse to unseen length than linear encoding (where Pe i), despite having comparable bettrr in-distribution performance. Figure 12 shows
36

---

Published as a conference paper at ICLR 2023
examples Gridg and partially observed (abab)', where compare the accuracy freshly drawn samples of the same length as the training sequences (i.e. in-distribution), or of twice the training length. For Gridg. both positional encodings achieve comparable accuracy, however the sinusoidal encoding performs significantly worse when tested on sequences doubled length. For partially observed (abab)* where the label is whether the current string is multiple of abab, the sinusoidal encoding has a clear ad advantage over the linear encoding on in-distribution performance. However, when tested on sequences of lengths twice as those during training, the performance gap between the two positional encodings shrinks significantly.
Training with shifted positions: In general, unseen positions appear be a major contributor to Transformer's s failure of length generalization. This is evidenced by the comparison between Trans- former trained with absolute positional encoding. and Transformers trained with random shifts added to the positional encoding: for each batch, we sample random positive integer in [0,400] and add to the position indices before calculating the positional encoding: this random integer is the same for each batch and varies across batches. Figure 11 shows that adding such random shifts gives significant boost to Transformer's length generalization performance, for both Dyck and C2. This suggests that main challenge to length generalization is the distribution shifts due positions un- seen 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 general- ization for both parity and Dyek. For the former, removing positional encodings makes sense since parity is a symmetric function where the ordering of the arguments does not matter, 23 though the positive result for Dyck less clearly understood. Note that removing positional encoding does mean having no position information, since the use of the causal mask implicitly encodes the posi- tion, 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
Hyperparameters. 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 and vary the number of layers from I to 16. For LSTM, we fix the embedding dimension to 64, the hidden dimension to 128, and the number of layers to We use the AdamW optimizer (Loshchilov Hutter, 2017), with learning rate in (3e-5, le-4, 3e-4] for GPT-2 or fle-3, 3e-3} for LSTM, weight decay Ic-4 for GPT-2 or le-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 5000k for different datasets, which is much fewer than the number of possible strings of length 100.
Implementation details. 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 base. The LSTM architecture is the default one provided by the Py Torch library.
Computational resources. The experiments were performed 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 single GPU, for a total of 10a GPU hours. The remaining experiments in Section 5 amount to less than 1 1% of this expenditure.
Empirically, we are able to achieve non- non-trivial accuracy (even when evaluated at 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 layers, 8 heads, and dimension 512.
37

---

Published as a conference paper at ICLR 2023
C PROOFS
C.I USEFUL DEFINITIONS AND LEMMAS
Formal definitions of simulation. We first recall the notions of simulation introduced in Sec- tion
A function simulate an automaton for particular choices of T, q0- For semiautomaton
A (Q. 8), function f Σ - - simulates Ar. if f((1:) Ag (o1r)) for all input sequences dI: Here, the right-hand side denotes the sequence of states qı.T induced by the input sequence di.7 under the transitions ò starting from state yo-
A function class can simulate multiple functions associated with semiautomaton. For semi- automaton
(Q, 2, and a positive integer T, a function class F (a set of functions Σ - Q²) simulates A at length T if, for every go E Q, there is function Sno E F which simulates Ar.q0-
Our proofs rely composing "gadgets" which simulate various substructures of the transformation semigroup ((A). Thus, it will be useful establish third notion of simulation, which works for functions in the embedding space Rd rather than the symbol spaces Q, E. For clarity, we give this notion a different name (continuous simulation):
For semiautomaton A = (Q, E, 6), function f Rd + F continuously simulates Ar.go if there exist functions E Σ + Rdd W imf Q such that W E simulates Ar.
When W a linear threshold function 2 ~ arg max, (W 222g this corresponds to a standard classi- fication head. However, our constructions may leverage other encodings of discrete objects.
Function approximation. We provide some simple function approximation results below.
Lemma (ID discrete function interpolation with an MLP). Let X be a finite subset of R. such that |x| S Bu for all E X, and Ix x'1 2 for all x & x' E X. Let f X - Rd be such that a 2-layer ReLU network for which
[[ S for all x E X. Then, there is a
fmip + E; mlp)) = f(x)
Va E X,
141 S A/4.
The inner dimension is d' = 4/20, and the weights satisfy
4
4B
|| S A
+2, |1 I By, by = 0.
Proof. For each Zo E X, we construct an indicator vizo (x) for no. out ReLU units. Letting A = A/4, the construction is
(
2A')
E (ro
A')
Mlroo (x)
A'
A
(Fg + A) r
(Io 2A')
A
The second layer simply sums these indicators, weighted by each f(x0).
Lemma 2 (General discrete function interpolation with an MLP). Let X be finite subset of Tedh, such that ||>||| < for all x X, and || x'll A for all x # x' E X. Let X IRdou. be such that [f(x) I By for all x E X. Then, there is a 3-layer ReLU network for which
Smip(x + E; Ouup) = f(x) Yx E X,
181 < A/4.
Letting X, denote the set of unique values in coordinate i, the inner MLP dimensions are as follows:
dj
Σ 144
d2 = 12)
iE iE(di
The weights satisfy
4
4B,
Wa 1
|1al <
+2,
||
1, llbell S dmn
W I By,
by = 0.
A.
A
38

---

Published as a conference paper at ICLR 2023
Proof. The first layer uses the same construction that in Lemma creating indicators for each I E A for each For each æ E X, the second layer has an activation which sums the indicators from each 24, with bias dinn (thus creating indicators for each æ). The third layer outputs f(x) for each indicator.
When apply Lemmas and 2 in recursive constructions, and B²/A 2 1, will opt to the bound | I 6B,/A, to reduce the clutter of propagating the term without resorting to asymptotic notation.
We also introduce simpler version Lemma 1 for the special case of the threshold function f()) = 1 | 0j:
Lemma 3 (Threshold with an MLP). Let X' be a subset of R. and |a| A for all X. Then, there is 2-layer ReLU network for which
fulp(I + mipp) 1[x
Wr K A/4.
The inner dimension is d' and the weights satisfy
1
| $
|1111 - 1/2,
| $1,
by 0.
Proof. We construct the threshold using 2 RRL.U units. The construction is
x + A
x
A
((x)
2A
2A
Selection via soft attention. We record some useful lemmas pertaining to approximating hard coordinate selection with soft attention. The following is simplified version of Lemma B.7 from (Edelman et al., 2022) (which generalizes this to multi-index selection):
Lemma (Softmax approximates hard max). Let & E IRT.. Let softmax(2) RRT - RRT denote the T-dimmensoona softmax function:
e
[softmax(2)) (2)]c Every e#
Let : arg max, Zt. Suppose that for all +1", 1", zv S 21 - 7. Then,
||oftmax(2)) et- || S 2T 2T. e"
Proof. Without loss of generality, max = 7 (since the softmax function is invariant under shifting all inputs by the same value), so that all other coordinates are non- non-positive. Also, assume 7' > (the T case is trivial). We have
E
c e9 +
T-1 2
T-1
[softmax(2)]e- 67 Eepo- el
+7
E
and for +
[softmax(2)]e
et + Enjrr e7'
Thus, the 1-norm of the difference bounded by
T-1
+ (T-1) ([-1)
e7
67 e7
claimed.
39

---

Published as a conference paper at ICLR 2023
We note the following elementary fact about 2-dimensional circular em- Proposition 5 (Circular embeddings). Consider Pı,..., Pr, the T equally-spaced points on the 2-dimensional circle:
Positional embeddings. beddings.
1ml
1
[pe]2 cos
[pe]2 == sin
T
T
Then, for any t t t',
2m²²
19.7 T2
(w)) 5
<
T2
C.2 PROOF OF THEOREM 1: LOCARITIMMC-DEETTI SHORTCUTS VIA PARALLEL PREFIX SUM
In this section, we give the full statement and proof the universal existence of logarithmic-depth shortcuts.
Theorem 1 (Simulation is parallelizable). Let A = (Q) be semiautomaton, Qo Q, and
T 2 1. Then, there is a depth-[loge T Transformer which continuously simulates Ar. with embedding dimension 2 Q + 2, MLP width |Q1 + [Q] and oo-weight norms at most m +
2, 10T /lo [Q] + log T). It has H = 2 heads with embedding dimension [Q] implying 2 + atlention width, and 3-layer MLP.
Proof. The basic idea is that all prefix compositions 0(- ... o(: can be evaluated in logarithmic depth using binary tree whose leaves are the per-input transition functions 0() Q Q. The attention heads select the pairs of functions that need 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 == - + 1,2 1,3
1,,
Σ (E(1.2.)
f(a) eg 1
Small nuances. We will produce 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),.... 0 at the beginning; while this greatly simplifies the positional selection construction, this padding construction could be replaced with 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 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 these dimensions can implemented by the MIP and attention heads.
Let L = log, T be the depth of the binary tree. We choose d = 2/Q1 + 2. Instead of indexing the dimensions by [d], we give them names:
Left function encoding dimensions (g, L) for each e Q.
Right function encoding dimensions (q, R) for each y Q.
Positional encoding dimensions P1, P2.
Without loss of generality, let Q [Q] {1..... Q) (selecting an arbitrary enumeration of the state space). Also, assume [Q1 2 (if not, add dummy state). We choose F(a.) == Egeq 8(g. o,)) R(..R), mapping each input symbol to the "transition map" of its transitions. At the padding positions -(T - .). we will encode the "go to qo" function: Egeq qu (().))-
Function composition gadget. 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.
40

---

Published as a conference paper at ICLR 2023
P R with fixed parameters W1. b1, W2, b2, W3 whose dimensions and weights only depend on Q, such that for all f.g: Q - Q Pmlp outputs the transition map of g given the concatenated transition maps of f and g. That is, for all [Q|2(Q] choices of f.g:
Lemma 5. There exists 3-layer ReLU MLP @ulp
Σ gla) e(g.L) + E(()) c(g.R)
-E (f e(g.R)
BEQ
9EQ
BEQ
The intermediate dimensions are d₁ = |Q| + [Q| and d2 = IQ1 and weight norms are bounded by 4|Q| + 2.
Proof. The first layer uses Lemma I to create 1Q12 indicators: one recognize each value along the (gg,) direction. Let index these by q. o Q. Then, this gives W1 E Rdx44Q"" by R4Q2 W32 E RIQ such that
[((= W22)oook + Wi= b1)) 1(e - d'].
We also add Q more weights which let the inputs pass through along the ((q.R) directions (add Q more rows ((g,R) to Wi, W2. calling these indices oy for all E Q; set biases to 0), for a total of Q| [Q] hidden units and [Q| + |Q] output dimensions of W.-
The second layer implements multiplication between the indicators and function values. The outputs are again indexed by g, q' Q. We define W (101 and b'g RIQ to be such that
[W2 (..44) = QI 1((4,()) = (7. 7'), [Wilag)) = 1 = d). [4 1Q) 44.d. E Q.
Overall, so far we have
[aa o (= + W + b5) o o (z + wi bi))( = ,(d') 1 TT g/).
The third layer W3 E TR 1Rdx simply converts these activations back into an transition map:
W3
Σ
((
V'EQ
Finally, we note the weight norms:
[ I 4 Q),,
||b|ll. S 4|Q11 2,
|W S <4Q)),
|b = =Q],, || = 1.
Recursive parallel scan. The rest of the construction uses a standard parallel algorithm for com- puting all prefix function compositions: layer E [L]. compose the function at position i with the function at position t - 2/-1 This is standard algorithm for computing all prefix compositions of associative binary operations with logarithmic-depth circuit (Hillis & Stecle Jr., 1986). We choose the position embeddings to enable implementing these "look-backs" with rotation matrices. For the circle embeddings
each f. 11T+1,.... 1,1, 1..., T),
.t
Pt,p, == cos
Pt.., = sin
T
T
T
In detail, for each 1 151 L:
Let 0 m2! y : 100T (log [Q + log T).
Let H : 2, = [Q]. Recall that [Q] 2 2. We will index the heads by superscripts [L], [R].
Select w/H w w/Rl = eT + e))
41

---

Published as a conference paper at ICLR 2023
Select W = VT ee + ep,e2 po. where pa the rotation matrix
cos(O)
sin(0)
sin cos(O)
the e1, e2 basis.
Select WHl w = =x ((a.))
Select W EyERR eg wil = (g
At layer 1, let al4, ( E R2 denote the attention mixture weights of the two heads. With this choice of y, Lemma 4 and Proposition 5, for each t [T]. we are guaranteed that ||a - ee/1, and ||a44] -2-22==/11 ||1 are both most [Q]-T Thus, by Hölder's inequality (noting that this mixture is over 7' vectors of oo-norm at most (Q), this attention layer's output is 0.1-close in the co-norm to the concatenated transition maps of the functions at positions & and - 2 allowing us to invoke the perturbation- trouut function approximation guarantee of Lemma with A = For the MLP, use the function composition gadget.
0.1
Thus, at the final layer, the (g. R) dimensions at position t contains the transition map of the prefix composition
(q r qo)-
suffices choose be + R), for an arbitrary q to read out the sequence of states as scalar outputs in [[] To output one-hot encoding, an additional MLP (appended to the end of the final layer) would be required.
KROHN-RHODES C.3 PROOF OF THEOREM 2: CONSTANT-DEPTI SHORTCUTS VIA KROHN- RHODES DECOMPOSITION
We begin with the full statement the theorem:
Theorem 2 (Transformer (oohn-Rhodes) Let = = (Q,E,)) be solvable semiautomaton (see Defimition 6), go E Q, and Z 1. Then, there depth-O(1Q) log (Q)) Transformer which continuously simulates AT with embedding dimension ((219) MLP width (Q) 81) +
O Q [TA)) T), attention width 21 ((())) heads, and weight norms bounded by 6 Q T log + max([[Q),
We will begin by presenting self-contained constructions for the atoms in the Krohn-Rhodes decomposition: a modular counter and memory unit. In Appendix C.3.2, we will introduce neces- sary background from Krohn-Rhodes theory (including the definition of a solvable semiautomaton). In Appendix C.3.3 and C.3.4, we will complete the proof of Theorem 2.
C.3.1 BASE CASES: MODULAR COUNTING AND MEMORY
Base case 1: modular addition. We will start with construction of a tiny network which lets
us simulate any semiautomaton whose transformation semigroup is cyclic group. Later on, will use copies of this unit to handle all solvable groups The construction simply uses attention to perform flat prefix sum, and an MLP to compute the modular sum.
Definition 2 (Modular counter semiautomaton). For any positive integer define the mod-n mod- ular counter semiautomaton A (Q,E,
Q
1},
{0. {O..... 1}.
5(g, (q + σ) mod n. N
Lemma 6 (Simulating a modular counter). Let A (Q,E, be the mod-n modular counter semiautomaton. Let qu E Q and T I. Then, there is a depth-1 Transformer which continuously simulates Ar.go with embedding dimension width 4nT, and xe-weight norms at most 4nI + 2. has H head with embedding dimension
and 2-yyer ReLU MLP.
42

---

Published as a conference paper at ICLR 2023
Proof. The intuition is simply that the lower triangular matrix causal mask can implement simula- tion in this cyclic group by performing unweighted prefix sums. The only subtlety is that selecting We Wx- = 0 does not quite give prefix sums: the attention mixture weights at position t are E rreg Ce', while would like the normalizing factor be uniform across positions (1/T rather than 1/t). It is possible to undo this normalization using the MLP; however, a particulaly simple solution is to use an additional padding input 1 and I-dimensional position embeddings to "absorb" a fraction of the attention proportional to - 1/T.
We proceed to formalize this construction, beginning with the input embedding and attention block:
Select d 3, A 11 == 1. Intuitively, the dimensions implement (input/output, padding, position] "channels".
Select input symbol embeddings B(o) e1 Fdd for each e Σ.
Include an extra position I, with embedding E(L) e2 and position encoding P1.: : O.
Think of this as padding at position 0; it is not masked out by the causal attention mask at any position t 1.
For / E [T], select Pt.: == "r,3, where "/ = logg(2 () 1) is such that c 27
Select Wq == C3, Wx == 62, Wv 61, W =
We do not need residual connections.
In the output of this attention module, for any input sequence 01.7, the 1 channel of the output at position t is then
1
Σ
8 = 27
4E[T]
where 2; (O,...,n 1} is such that o(-,0z) g The MLP simply needs memorize the function
8 + (T's mod n) e1.
We invoke Lemma 1, with A 27 B2 n-1 2h2 By The number of possible values of S (the cardinality of X in Lemma 1) is at most nT.
Base case 2: memory lookups. turns out that simulate semigroups instead of groups, the only additional ingredient is a memory unit, semiautomaton for which there are "read" and "write" operations. The minimal example of this is flip-lop (Example 2), semiautomaton which can sequentially remember and retrieve single bit E {0, 1). and whose transformation semigroup is the flip-flop monoid. will be convenient to generalize this object to Q states:
Definition (Memory semiautomaton). For given state set Q, define the memory semiautomaton
A = E,);
==QU(L)
8(g 0) = 4q E E
8(g.)) EgEQ.
Lemma 7 (Simulating a memory semiautomaton). Let A (Q,E, be the memory semiautoma- ton. Let go E Q and T 2 1. Then, there depth-1 Transformer which continuously simulates AT no with embedding dimension 4, width 4(Q), and so-weight norms at most 2T log([Q[T)) It has H head with embedding dimension
and a 2-layer ReLU MLP.
Proof. We start in state qo E Q. Our goal is to identify the closest non-L token and output the corresponding state. The attention construction is:
Select d = 4, = 2, 1.
43

---

Published as a conference paper at ICLR 2023
Select input symbol encodings
E(o) = (1/o -|40 + 1 + L|o:)e1 1 Llez + 64 E ER"
where the first coordinate denotes the action that sets the state the second coordinate denotes whether the input is the no-op action 1, and the fourth coordinate is padding.
We use positional encoding Pt:: := (t/T) e3.
We = (-21 ea] E RR WK == Icez cea) E RX2 for O(TTog((11T)) as explained below, Wv = fei 0] E R4x2 and W [ei 0 R
The unnormalized attention score computed for position i attending to is cli/T - - -)). Note that the max attention value is achieved the closest reset action: the unnormalized scored non-negative if and only if oj * L, and j/T increases with j ensuring that the closest position is chosen.
Denote this max position as Jmax. In the setting of hard attention, the output for the ins token after attention module is j(a)))))) e1. In particular, this value is qo and only if o, 1, Vj S i, i.e. the semiautomaton never leaves the starting state. Otherwise, the value is the value of the nearest non-L state (including the current state).
By Lemma 4 (with E RR," that is, ||or &jmaalli |ll S 2T e e-//T. This implies, that the output of the attention layer, Ej E(a)
c/T), can approximate hard-attention by soft-attention weights
s 2T e- Then, the MLP can simply round the first coordinate, and we can invoke Lemma 1 with A - 8 exp(-e/T) 1/2 (for = (loog(6|Q[TT)) BN Q), By = [Q] to get weight norm bound log((o(()))TT S TT) and width 4/Q).
C.3.2 PRIME DECOMPOSITIONS OF GROUPS AND SEMIGROUPS
The key idea behind the proof of Theorem 2 is that all semigroups (and thus, all transformation semigroups of semiautomata) admit a "prime factorization" into elementary components, which turn out to be simple groups and copies of the flip-flop monoid, which have both been discussed in Appendix A.2. This is somewhat counterintuitive: the only constraint on the algebraic structure of semigroup associativity (and indeed, there are many more semigroups than groups), but all of these structures can be built using these two types of "atoms". These components, as well as the cascade product under which this notion of "factorization" is defined, are naturally and efficiently implementable by constant-depth self-attention networks.
The special case of groups. We begin by discussing the analogous decomposition for groups, which generalizes the fact that integers have unique prime factorizations. Let G be a finite group, and let
G Hn Hn-1 D H1 > Ho =
be a composition series: each H; is a maximal proper normal subgroup of H1-1: denotes the trivial group with element. Then the quotient group H,+1/ H; is called composition factor. The fordan-Hölde theorem tells us that one can think about the set of composition factors as an invariant of G.
Theorem 6 length n, and the sequences of compositions factors H,41/ H; are equivalent under permutation and isomorphism.
Any two composition series of G are equivalent: they have the same
When each H;-1/ / H; is abelian, G is called solvable group. It turns out that each H H, is a simple group. so the composition factors of solvable groups can only be cyclic groups prime order (because every finitely generated abelian group is a direct product of cyclic groups, and, of these, only those of prime order are simple). The smallest non-solvable group is As. realizable as the
"TTchhncaall I. does not reset the state. We will see that when go is selected, it must be that the semiau-
tomaton is always in state 9o-
44

---

Published as a conference paper at ICLR 2023
group of even permutations of 5 elements. As a part of Theorem 2, we will use the composition series to iteratively build neural networks which simulates solvable group operations, requiring intricate constructions to do this with depth independent of the sequence length T.
Adding memory to handle semigroups. Now, we move on to semigroups. When not all of the input symbols to a semiautomaton induce permutations, we no longer have the group axiom of invertibility (also, if there is no explicit identity symbol, we are not guaranteed to have the monoid axiom of an identity element either). Intuitively, this would seem to induce a much larger family of algebraic structures; an analogy, which is formalizable by representation theory, is that are now considering a collection of general matrices under multiplication, instead of only invertible ones. The non-invertible transitions collapse the rank of the transformations, reducing the set of reachable transformations whenever they are included in an input sequence.
A landmark result of Krohn & Rhodes (1965) tames the seemingly vast and unorderly universe of general finite semigroups. It extends the Jordan-Hölder theorem to the case of semigroups, for a more sophisticated notion of decomposition. Since that work, many variations have arisen, in terms of its precise statement, construction of the decomposition, and proof of correctness. Out of these, an important development is the holonomy decomposition method (Zeiger, 1967: Eilenberg, 1974), which forms the basis of our results. We extract the definitions and theorems from Maler & Pnueli (1994), whose exposition emphasizes explicitly tracking the construction of the semiautomaton. We also refer to Maler (2010); Egri-Nagy & Nehaniv (2015); Zimmermann (2020) recent expositions, containing historical context.
Definition 4 (Cascade product: cf. (Maler, 2010), Definition 11). Let n be a positive integer. For each i [m]. let A(a) (Q(), E00 5(0) be semiautomaton. For i {2, (2,...,yth, .,.,,,,, let 260) Q x... x Q(i-1) x Σ او6 denote dependency function. This object (fA(0)) }: (6 is called transformation cascade, and defines cascade product semiaulomaton A = (Q
x
Q(m), E(), 5) by "feedforward simulation" under the dependency function. We define 5 by the i-th component of its output (which call g(Si); Q(1)) x Q(") x E(); (()):
(( (1) ...... (1)) .) == (g (0))
where
if Fi=1
(g ( (1) o)
otherwise
The corresponding transformation semigroup ((A) is known as a cascade product semigroup.
Intuitively, the cascade specifies a way to compose semiautomata hierarchically: the first layer i maps input sequences to its state sequence, and each internal layer receives input which depends on the states of all of the preceding layers. Algebraically, the cascade product semigroup a sub- semigroup of the larger wreath product of semigroups (the straightforward analogue of the wreath product of groups, discussed in Section A.2). Although this is useful from an algebraic point of view, we will not use this perspective; the cascade product is a smaller substructure of the wreath product which is sufficient for semiautomaton simulation.
Finally, it will be convenient to define permutation--rsel semiautomata, which are useful interme- diate step in the Rhoode decomposition. To obtain our final result, we will further break these semiautomata down into flip-flops and simple groups.
Definition 5 (Permutatioon-rsse semiautomaton; ef. (Maler & Pnueli, 1994), Definition 12). A semi- automaton A = (Q, Σ, 2) is a permutation- reset semiautomaton if, for each a E Σ, the transition function 8(, : Q + Q is either a bijection (i.e. a permutation over the states of Q) or constant (i.e. maps every state to some g) Associated with each permutation-reset semiautomaton is its permutation group, generated by only the bijections.
Now we can state the Krohn-Rhoder theorem, which decomposes every finite semiautomaton into transformation cascade.
Theorem 7 (Krohn-Roe -hoodes,, Let A = (Q, 5) be a semiautomaton. Then, there exists a trans- formation cascade ..,, ,,,))...,,6!! defining a cascade product semiautomaton A', such that:
45

---

Published as a conference paper at ICLR 2023
(i) The input symbol space of A (and thus, that of A') E, the same as that of A.
(i)) Letting Q(9) denote the state space of A there exists a function W Q(1) x... ((m)) Q such that W с Ar simulates Ar. for all T 2 1, 90 E Q. For each E [n], the transformation semigroup T(A is permutation-rese! semiautomaton with at most [Q] states, whose permutation group is (possibly trivial) subgroup of T ((A) (Maler & Prueli (1994), Theorem 4).
(iii) The number of semiautomata the cascade n S 2 Furthermore, the cascade has at most [Q) levels: the indices can be partitioned into at most I, S [Q] contiguous subsets N(I) (1..... .ni,, ,2),, {mi + +,....,, + t N(L) = {n nt + t....., such that o) only depends input indices from previous partitions (Maler & Pnueli (1994), Claim & Covollary 12).
With this decomposition,
able define solvable semiautomaton.
Definition 6 (Solvable semiautomaton). A (Q, Σ, be semiautomaton. We call A solvable all the permutation groups associated with all of the permutation-rest automata from Theorem solvable groups.
The remainder this section will build our construction from the bottom up:
Appendix C.3.3 will build up from the base case of cyclic groups (Lemma 6), using increasingly sophisticated notions of group products, culminating in recursive construction which simulates all stages of the Hölden composition series. The crucial step is construction for simulat- ing the semidirect product of groups, given networks which simulate the individual components; this allows us to handle the solvable non-abelian groups.
Appendix C.3.4 will build networks which simulate permutation- reseet semiautomata. A new base case arises: the memory unit (Lemma 7), semiautomaton whose transformation semigroup a generalization of the flip-flop monoid. Combining the constructions for solvable groups and memory units, we obtain simulators for solvable permutation-reset semiautomata. Finally, the cascade product guaranteed by Krohn-Rhoder (Theorem 7) glues all of these pieces together, giving the final result.
C.3.3 SIMULATING SOLVABLE GROUPS
We begin by handling groups Now, we are ready to specify the recursive constructions which "glue" these components together to form solvable groups. We will proceed in a "bottom-up" order:
(i) Define canonical semiautomaton A corresponding to each group G (Definition 7), such that if network can simulate _AG, (G, it can simulate any other semiautomaton whose trans- formation semigroup T(AA) is G. This lets talk about simulating groups, rather than particular semiautomata. We will show how to turn simulators for groups N and H into simulators for extensions N by H, for increasingly sophisticated extensions, until all cases have been captured.
(ii) Show how to build the trivial extension; given networks which simulate the groups N and H, simulate the direct product G z N H, by simply running the individual simulators in parallel (Lemma 8). Combined with Lemma 6, this immediately allows us to simulate arbitrary abelian groups with depth 1, since every abelian group isomorphic direct product of cyclic groups
(iii) Show how to build a split extension: given networks which simulate a normal subgroup N and quotient H, construct network which simulates any semidirect product G = N H (Lemma 9). This is the first place where we will require a sequential cascade layers. will allow us handle certain families of non-abelian groups (including S3, D2n, A4, S.)
(iv) Show how to build arbitrary extensions (any G which contains N as normal subgroup, and for which the quotient group G/N is isomorphic to H), using the wreath product (Lemma 10), which contains all of the group extensions. The wreath product is itself the semidirect product between a way direct product and H, so this can be done in constant number of layers, by the above. This finally lets implement any step of
46

---

Published as a conference paper at ICLR 2023
composition series. In particular, using eyclic groups as a simulable base case, this shows that we can simulate all solvable groups.
Step (i). will be convenient to associate with each group G a canonical "complete" semiautoma- for the class of all semiautomata A for which T(A) G. It is simply the one whose input symbol space E is every transformation reachable by some sequence of inputs (i.e. every element of ((A).. (For a semigroup. we would also want adjoin the identity element if it is missing, however, we will only find it useful to define this for groups.)
Definition 7 (Canonical group semiautomaton; simulating a group). Let G' be finite group. Then, we define the canonical group semiautonaton for G as the semiautomaton (Q, Σ, 8) defined by:
Q= G, the of elements G. Note that if (for example) G Sn are setting the state space to be the set of n! permutations, not the ground set [n].
Σ G. That is, include all functions in the input symbol space.
= h- g, for all Vg E Q, E. (In algebraic terms, we are embedding the G into its left regular representation, a.k.a. left multiplication action.) Thus, if we take qo to be the identity element, the sequence of states q1, .2....,.7 corresponds to q: 010t-1....00
,(g, k)
When we simulate the canonical group semiautomaton, we will always choose qo to be the identity element 6G-
A sequence-to-sequence network said to continuously simulate G length T if it continuously simulates the canonical group semiautomaton of G at length T.
Notation for composable implementations. Let us furthermore formalize an implementation of group simulation. For any finite group G, T 2 1, go E G, let A = (Q Σ 8) be the canonical semiautomaton for G. Then, we summarize family of concrete implementations of networks which continuously simulate of Ar.ea- We write sim (G,T) + (E : G - Rd fu Rxxd
RTxd W IRd G), where the shape parameters of the output can depend on G,T.
To reduce notational clutter, we will access the shape attributes of implementation via "object- oriented" notation, defining
GG, (depth, dim, heads, headDim, mipWidth, normBound)
to respectively denote the complexity-paramteerne parameterizing quantities
(L, d, k. max(d,), B},
defined in Appendix A.4. Also, will let sim (),(E),,,,))) W) respectively denote the encoding layer E, network parameters Om. and decoding layer W.
Canonical encodings of group elements. We also enforce that throughout our constructions of networks which simulate groups, we will maintain that all networks and their submodules manipu- late encodings via integer vectors in a consecutive range {O...., n - 1}. Furthermore, the identity element will always map to the zero vector. We will keep track of the dimensionality of these vec- tors sim(G. rempii S d, and their maximum entries sim(G. T reesiie 1. All encoders E and decoders W will map all group elements to and from this kind of representation, and we will choose W E-1 In all, the networks will keep DeDDim-dimensiona "workspace" of integer vectors, with entries bounded by repSize- 1. When combining groups via the various products constructions, we will combine the components' individual workspaces to create a larger workspace for the product group's elements.
We make some additional remarks on implementations:
Note that the canonical semiautomaton "forgets" about the semiautomaton abstraction, and never assumes that G is a permutation group on the original state space Q of the semiautomaton we would like simulate. Indeed, when N a G are permutation groups on Q. there is no natural
47

---

Published as a conference paper at ICLR 2023
(fg Sgg (8gg
simulate all
Figure 13: Recursive construction for simulating direct product of groups G( ... x G(*), Any number of groups can be simulated in parallel without increasing the depth.
permutation group on Q associated with the quotient H = G/N; it turns out that will consider simulators for N and H!25
To return to solving the simulation problem for some semiautomaton A (Q, Σ, whose transformation semigroup isomorphic to G (at length 7 and initial state 2o), let G + Sq denote this isomorphism. We use AG T.ea (f1:7) as the network, with an encoding layer F a 1 and decoding layer (i + T(oo)) p W, which can be memorized by an MLP of width O(G)) via Lemma 2.
The modular counter semiautomaton, for which we constructed simulator in Lemma 6, the canonical group semiautomaton for the corresponding cyclic group C.. Calling this construction simc., we can easily verify that it satisfies the canonical simulator's 's conditions, and:
simc, .eeth 1.
simc. =
hhadss = 1.
simc..-eadDinn 1.
simc, mipWidthh -
simc, nomc..normBound < 4|G -T T
-T.
rimc, 1.
simc.ieepSize |G|.
Step (ii). As a precursor to the more sophisticated products, we formalize the obvious fact that two non-interacting parallel semiautomata can be simulated without increasing the depth. First, we define the direct product semiautomaton:
Definition 8 (Direct product of semiautomata). Let A= =Q, (Q', be two semi- automata. Then, A' (Q x Q', [u x U' x 8') denotes the natural direct product semiautomaton. Its states are ordered pairs (q E Q. / E Q'). Its input symbols are defined simi- larly, adjoining identity inputs (so that 8(4,e, e) = q, ('','e)) g'). The transitions 8' are defined such that
(6 x (')(((())))))))))) (', = (5(g () 8),8'(',)0)))
Note that under this definition, we have T(A x A') = ((A) x ((A'). . In particular, for two groups G, H. we have G H T(A9) ((4 = ((QxH)) G H.
Lemma 8 (Direct product via parallel simulation). Lei G GO),..., ((n) be a collection of finite groups, and les T 2 1. Suppose each group admits a simulation sim; == sim (im(G() (G)))) Then, there is a simulation of the direct product group simx = sim(G(1) x .... G(a), T), whose sizes satisfy:
Thhere nothing general preventing quotient groups from being extremely large groups which are not realizable as smaller permutation groups For concrete examples, see (Kovács & Praeger, 1989). When we ultimately specialize to simulating the composition series of solvable groups, the largest groups we will handle will be the eyclic groups of prime order, so will in the end be guaranteed that the groups want simulate realizable with [Q] states, but not directly canonically.
48

---

Published as a conference paper at ICLR 2023
simulate
(g.4k
(Rig Maad
simulate H
mix
unmix
Figure 14: Recursive construction for simulating the semidirect product N H. The quotient group H is simulated first; these outputs are used to "re-map" the inputs into the simulator for N.
himx.depte max; (sim; (sim,.depth)
simx. >, (sim,.dmm)
simx. E, heads.
simx.headDim = max; (smm
simx.mlpWidth mm/smmmm
noox.normouund maaxxsim,nnomoonnd;;
sipx.mepDim = reim
simx max; repizz)
First, we pad all of the individual sim; with layers implementing identity (add residual con- nections, and set attention Wy and all MLP weight matrices to 0), so that all of them have depth (sx
Proof.
Then, the intuition is to construct the direct product semiautomaton by concatenating the "workspaces" of each G In other words, we set the canonical encoding simx.Ei of to be the concatenation of each sim;'s encodings.
The direct product simply lets each sim, take inputs and outputs in its individual workspace. To enable this, we need enough parallel dimensions. We set an embedding space of dimension (sim, } (and similarly within the heads and MLPs), partitioning the coordinates such that in the product construction, each sim; and sim,.0 only reads and writes to its own dimensions.
This clearly simulates the direct product group. Figure 13 provides a sketch of this construction.
Note that the direct product construction already allows us to simulate all finite abelian groups in constant depth, since each such group isomorphic to the direct product collection of abelian groups of prime power order.
Step (iii). Now, as a harder (and conceptually crucial) case, we show how to simulate a group which is a semidirect product of two groups we already know how simulate. This encompasses the direct product as special case, but can now handle some non-abelian groups which admit such decompositions (like the dihedral group D2n). The catch that we will have to simulate these groups using sequential cascade of the individual simulators. This is the key lemma which lets us simulate non-abelian groups:
Lemma 9 (Semidirect product via 4-stage cascade). Let G be a finite group which is isomorphic to semidirect product: G = N H. where N is normal subgroup of G. Let T 2 1. Suppose N, H admit simulations sim := Sim,N, T)), sim == H,, Then, there is a simulation of G. simy sim T), whose sizes satisfy:
sim depth sim depth + sim depth + 2.
simy dim sim .iim + dim.
simx-heads heads = max maxx(simy.hen heads, sim heads).
simy headDDi = max(sim,, max(simy.headDim. sim headDim).
mlpw.dh max 4|Gl).
49

---

Published as a conference paper at ICLR 2023
normoondd
max(sim 6 ( bsta +
sim:n.repDim).
rimx-repDim simy repDim rimp-r
ripsxee max(simy reppsi,,,rrpSSize simg.
Proof. The intuition is as follows, using the dihedral group D2n _tt Cn C2 as an example:
For simplicity, let us think of the "reversible car on circular world" semiautomaton, whose transformation semigroup is Dan- Its state consists of direction E {+1,-1}, and position {0, .,,,,,,,,, n - 1]. It has two types of inputs: "advance by i" (increment the position by in the current direction, modulo n), and "reverse" (flip the sign of the direction). Our simulation task is to track the car's state sequence, given a sequence of inputs (in constant depth, of course).
It is intuitively clear that we can (and should) compute the sequence corresponding to "direction time t", which is equivalent to simulating the parity semiautomaton.
We will convert the "advance" moves via "basis transformation": whenever the current direc- tion is 11 an "advance by i" should be converted into -i. Then, we have reduced the problem to the prefix sum.
Algorithm. This intuition essentially shows how to implement arbitrary semidirect products; we derive the basis change from db Before implementing it with Transformer operations, we formal- ize this "basis transformation". Recall that by the definition of semidirect product, the elements of N x H can be written pairs (g E N, h e H), equipped with a homomorphism ø + Aut(N) which specifies multiplication rule:
(g. h) (g', h²) (gd)(g'),hh))
Let us write down the properties of
is homomorphism. That is, on (ow( = on o Φn as permutations on N. The output of that homomorphism, &b. is also a homomorphism. That is, (g(gg') (g((g)
on (g').
Let us roll the definition of the semidirect product, given a sequence of inputs (ge. h4):
(g2, h2) (g1. h1) (g2 Th, (91), hghh),
(ga, ha) (g2, h2) (gi, = (g3 tn (92 Φng (91) hzh2h1),
(gs, h4) ... (g1, h1) (91 Φn, (gs- Φhg (0r )0, (g1))) hahahohi). In general, by induction, letting (gse, hege) denote (gt, h4) ... (g1, h2), we have
g5t = gr Tn, (9t-1) @here (91-2) ... th.... (g2) wbi...ig (91).
Applying &hg on both sides, we notice that
Tgg (95e) - (9:) Onge. (9t-1) .. (92) ong (g1).
Thus, it suffices compute each hgg h ... h1, map each øe + $7g (9t). compute the prefix products these "coordinates", then invert the mapping to get back gsc.-
Implementation. Like before, we partition the embedding dimension in our construction simx into blocks, one for each component simulator. Let us index the dimensions by the dy := simyy. indices in the "N channel" and analogously for the dæ diymensiionn "H channel". We choose the canonical encoding to map elements to their individual channels:
E(g. h) simy. E(g) (in the N channel) + simg (in the H channel).
We proceed specify the construction layer-by-layer. Let TQN,H3 denote sim sim (ommNN.y-eeeee d.w.t.hpphh.
50

---

Published as a conference paper at ICLR 2023
Layers 1 through LH: quotient group simulation. As suggested by the intuitive sketch, begin with Ly Transformer layers, which are just a copy of simy .8. reading and writing in the H channel, with parallel residual layer in the N channel. So far, after these Ln layers, the output each position & is an integer vector, whose H channel contains lig,, and whose channel contains simy. F(..)
Layer LH + 1: basis change. Now, let us add one more "mixing" Transformer layer, whose attention block identity""; we only need a 3-layer MLP block, which represents the function
G H) + (g)
To do this, we invoke Lemma 2 (choosing the output to be in the same representation as that used by simy. F, in the N channel), with
A = 1, din simy semyyrepDim + rem
B₁ = max repSere repsie) By = simy
giving us construction with
di W 4( HH),
d2 < | [H],
1/ < |6 I 6 msi repSize, sim simy repSize,,
W 1
162ll < sim repDim sim ||W simserepSize.
We also add residual connections in the // channel. In summary, after this layer, the output at each position / an integer vector, whose H channel contains fig,, and whose channel contains (g))
Layers Lu+ +1 through Lu L, + normal group simulation. The next Lx layers are a copy of simw. 0, with residual connections in the H channel. After these layers, the output each position t is an integer vector, whose H channel contains hgg. and whose N channel contains (
Layer LH + IN + 2: undoing the basis change. Now, we add one more Transformer layer, whose attention block identity; we will again use a 3-layer MLP block to represent the inverse of mapping function
(g E N, E H1) + (n(g).
This uses Lemma 2, with exactly the same bounds.
At the end of this final "unmixing" layer, the output each position & an integer vector, whose H channel contains lig, and whose N channel contains geti thus, this is valid simulation of the semidirect product.
This construction is sketched in Figure 14.
Step (iv). Note that H = G/N does not imply that G is semidirect product of N and H. Thus, al- though simulating semidirect products allows us to handle some families of non-abelian groups, this does not yet allow to handle general solvable groups (i.e. general steps of a composition series, even with cyclic quotient group). The smallest example is the non-abelian quaternion group Qs- the group of unit quaternions under multiplication, which cannot be realized as semidirect product of subgroups. Instead, we need to appeal to the Krasner -aluujnnn universal embedding theorem (Krasner & Kaloujnine, 1951): a characterization of all of the groups G which extensions of N by H, as subgroups of the wreath product
H.
Lemma 10 (Wreath product via direct and semidirect products). Let G be a funite group which isomorphic
wreath product: G e N H. Let T > 1. Suppose N, H admit simulations simy == sim T), sim = H,mHHTT Then, there is simulation of G, sim, := (G,T) In the case where rimm.e 1, the sizes satisfy:
Even when an attention block simply implements identity, we choose to include it rather than combining the preceding and subsequent MI.Ps into single MLP. This is ensure that if we compose number of Transformer layers that depends on (Q), the depth of each MI.P P is bounded by absolute constant.
51

---

Published as a conference paper at ICLR 2023
simulate
(identical)
1,0
0000
simulate H
mix
unmix
Figure 15: Recursive construction for simulating the wreath product N H. One independent copy of N is instantiated for each element of H. while the simulator for !/ permutes them.
.imphh simy .depth +
+ 2.
sim H sim sim
himds = max [ simy.heads, sim heads).
sim,. =
Dim))
mimdth = max [| -simy sim mipWidthh N
noruuno < max (sim ) normBond 6 [H]--
rim | repDim
rip,,e = max max/sim, (simy. repSize, repsize.
Proof. Even though the wreath product's algebraic structure can be very complex, the construction just requires implement its relatively simple description. Applying Lemma 8, have net- work sim which simulates N x... x N. Then, we simply apply Lemma 9, using sim to "re-map" inputs simx for the normal subgroup. This construction sketched in Figure 15.
Concise implementation of reindexing. We can make one interesting improvement application of Lemmas 8 and 9: the structure of the mixing function o, which specifies the semidirect product, is extremely regular. Very fortunately, the structure of allows to avoid any dependence on the size of the wreath product group ([NV) H. (H)) ) in the size measures of the implementation. A general automorphism on N x ... N is specified by its |N| values. However, in this case, is just permutation, specified by how each the H channels should switch places. Thus, much like the function composition gadget in Theorem 1, one from Lemma 2.
generic
construct simpler MLP than the generic
Specifically, would like to approximate the function o: (N N (N x x N), which simply applies Ty to the indices:
)n(gi") g (1)) (g((a())) (
(F((111))))
In the component neural networks' representation space, we need the MLP to implement
(") sim.., E(g sim E(())))(())))))))))))))))))
sim E(g. (()
recalling that the elements of g, h represented by integer vectors with oo-norm at most sim rrp.Bonddddd. Notice that when the representation [ is single integer, restricting to any particular coordinate in the representation of element g. this is the same composition prob- lem function transition maps solved by Lemma 5 in the proof of Theorem 1, which uses its left inputs permute its right inputs (modulo converting the representations from {0...., 1) to
52

---

Published as a conference paper at ICLR 2023
{1...., [H], which we can do by shifting the indicators the input and final-layer output weights). Thus, [N] simmy.dii .dim parallel copies of the 3-layer function composition MLP suffice, yielding
d1 = 4 N + [H| |Nd < ||
d2 = |
|b S 1 [ < 4/H1, -6 [H), II = 1. [ill S 4|H,,,
When the information about group elements in H is encoded by multiple integers, it is straightfor- ward to extend this construction, by replacing the one-dimensional indicator with the multidimen- sional indicator from Lemma 2. We will skip the details of this case, since our final results are only about solvable groups; when we want to simulate general group extension, will always come from the composition series, so that H is always cyclic group of prime order.
Thus, for general group extensions G, we can construct sim, the wreath product simulator for N H, and combine the individual simulators. Note that we can throw away the excess group elements from the simulator: only include in sim; E, sim,. W the group elements which correspond to the subgroup isomorphic to G. Then, no part of this construction needs to maintain a width or matrix entry scaling with [N H..
Putting all of this together, we state an intermediate theorem, which our most general result for groups:
Theorem 8 (Simulation of solvable groups). Let G be solvable group which is isomorphic to a permutation group on n elements. Let T 2 1. Then, there is a Transformer network sim == sim(G,T) T') which simulates G at length T,, for which we have the following size bounds:
sim.depth 5 3 3ogg |G|.
sim.dim s 2
sim.heads s 2|
sim.headDim = 1.
sim.mlpWidth v 20]
nim.normBound < 6n7.
sim.repDim C
sim.repSize < n.
Proof. Let
G H He-1
Ho =
denote the composition series. Then, because G is solvable, all of the quotient groups K, H,+1/ H; are abelian, thus cyclic groups of prime order. Since G is assumed to be subgroup of Sn, none of these primes can be greater than Thus, every quotient group K; in the chain satis- fies 2 S K S n. Also, note that the length of the composition series E is at most logg (Gl) (since each inclusion halves the size of the group).
We start with a simulation of Hi, which must be a cyclic group, and build the sequence of group extensions recursively until we obtain G. In the worst case (in the sense that the implementation size bounds from Lemma 10 are maximized), each step in the composition series must be manifested by a wreath products with K Cz. Recall that we have:
simx.depth = 1.
simx.dim 3.
heads 1.
hemd. =
simx.mlpWidth 4nT.
sing.normBound < 6n..
simg.repDim
rimx.repSize < n.
At each step 0, o,... E 1, Lemma 10, with #: K2, N = Hg implies:
53

---

Published as a conference paper at ICLR 2023
i depth dimpt + (1 more layer to simulate the cyclic group K,, and from the wreath product's mixing operations). (noting that all of the components can reuse the same
Him < |K | ,. .dmm+1
and positional encoding dimensions).
sim < | siimsr, heads
headdiin < m ,,
milp..h - max | - 4nT, 5nT, 5| [H:]--
sim normmoodd max (6nT, 6 [(])-
reppin = | si, reppim + 1.
,i -
Iterating these recursive inequalities < Loga T times gives us the desired bounds. Note that using Lagrange's theorem (II, [ [G), well the fact that for positive integers mi,-.. mg 2 2, we have bound on the series of prefix products: E. my S 2 M4
C.3.4 SIMULATING SEMIGROUPS
Now, using this construction and the results developed in the previous section for groups, we com- plete the construction for semigroups:
We combine the memory gate construction (Lemma 7) and any network simulating group to implement the corresponding permutation-reset semiautomaton (Definition 5), the elements of the cascade in Theorem 7.
To finish, we implement the cascade product (Definition 4) of these permutation-reset semiau- tomata, guaranteed to exist by Theorem 7. This gives the full result.
First, we summarize findings of Lemma 7, naming this neural network simyr in our "object- oriented" notation. Note that since we are no longer simulating canonical group semiautomata past this point, repDim, repSize are no longer well-defined.
.dmt. sim 4.
hhmas =
simy.headDim = 2.
simaymlpWWddth = 4/Q1-
simy.normBound s 27 lo)
Lemma 11 (Simulating permutation-reset semiautomaton). Let A (Q,E, be a permutation-
reset semiautomaton (see Definition 5), and let denote its permutation group. Let T 1,40 E Q. Let simg sim(G.T) be Transformer network which continuously simulates G at length T. Then, there is Transformer network sim'g which continuously simulates Ar. with size bounds:
dimth simg-depte + sim deeth + S 3 loga GG + 2. sim simg.dim simg.repDim + sima,din S [G| + [Q] +4
heads simg- heads sim +
head = simg + simas headDDm + rimc. & [Q] + 3.
sim'g-mlpWidth mlppihh [
|G| < 20nI' + 4QQ +
sim < nommsimmnnnnoBound simy.normBound, 6[Q]] < 6Q Tlog
Proof. Without loss of generality, will let [[Q1]-
We split the embedding space in our construction into two channels: the simg dimensions used by G. and channel consisting of 4 additional dimensions, be used by copy of the memory
54

---

Published as a conference paper at ICLR 2023
semiautomaton, whose symbol set is Q. Let call these the G and M channels. For the reset symbols, let (a) denote the 4-dimensional encoding of a from the memory semiautomaton.
Since we defined G' to be isomorphic to the permutation group associated with A. there is bijection G + So between group elements and permutations on Q. We choose the embedding E as follows:
E(a) = Sim.E((#
(((--(((((------))000(0((0011 ()) (G channel) + Eas (1) (M channel), Sim.F(ee) (G channel) + (gx) (M channel),
bijections 6(, resets 5(,,0) 9a
Let LG denote simg.depth.
Layers through Lę: group simulation. The first Lg layers are chosen to be a copy of simg.6 in the G channel, and only residual connections in the M channel. At the end of this, given any inputs σ:. which map via d-1 to g. (letting the group operation be identity when o, reset symbol), the outputs in the G channel will be de := rimii--eDDittiimnssionnn encodings of the prefix group products 9<t = 9t9c-1 ... 91- Now, letting r(t) denote the most recent reset (r s t such that o, is reset token), we notice that the state want can be derived from this sequence:
= 2 (9: (r(t) (r(t))aorc = ( (sr(t) (arrc))
(C.2)
Here, if there have been no resets up to time t, we define r(t) to be O. We treat go like reset symbol at the beginning of the sequence. Also, note that our canonical group semiautomaton simulator always uses Go = eg as its initial state.
Layer Le + 1: memory lookup and copy. To implement the above, at layer LG + 1, we put copy of the memory semiautomaton in channel M, setting its initial state to qo- We will modify this construction slightly, extending Wy with the identity matrix the de group element encoding dimensions of channel G. Intuitively, when the memory unit "fetches" the last non-l token, we would like it to copy the corresponding 95t- Note that simg-repSize, the co-norm bound the group element encodings, is at most [Q] by Theorem 8, so we do not need to modify the Wo, WK norms increase the attention head's precision. The final modification is that append to We an identity matrix copying the dç embedding dimensions to a new dç dimensional channel, which we will call the (for "invert") channel (set to 0 in the embedding all preceding layers). Then, by Lemma 7, at the end of this layer, at each position f, the M channel will contain (orco in dimension 1, and (er(t) in channel Finally, in channel G, we use only residual connections, preserving Ige in channel G.
Layer LG + 2: applying ((gh-1 pointwise. This finally allows us to execute Equation C.2 at each position f. We use one more Transformer layer, with attention block implementing identity. The MLP memorizes the function (g, h q) + ((gh~ (q) e1 (the coordinate is selected arbitrarily). with the concatenated (dinv 2 simg.repSize + 1--dimensional encodings the (G. I, M) channels, whose activations have co-norms bounded by [Q]. We invoke Lemma 2, with parameters
= 1, din = dine; B2 = IQ), By IQ).
giving us construction with
di < 4dmm 4dmm-(1 + 1), d2 < |G
Will S 4, || S 6(Q). || 1 | S dinv;
[Wall $ |Q]--
From this output, W simply decodes the correct q from dimension 1.
Lemma 12 (Implementing the transformation cascade). Let A (Q, be a semiautomaton, ()2)) ....., (())) be the transformation cascade (Definition 4)
and let T > 1. Let which simulates A, as guaranteed by Theorem 7. For each i, let sim, be a Transformer network which continuously simulates the permutation-resst semiautomaton A at length T. Then, there is a Transformer network sim which simulates A at length T.. Its size bounds are:
A d.epth |Q] (max, (max;(sim, depthh + 1) 1 < 3|Q log QQI 7
sim,4.dim Σ dim + 2 (AT( Q + 4) +1
55

---

Published as a conference paper at ICLR 2023
simx.heads Em= himas I 2 ([T(A)) +
simx.headDim- = m
< Q +3.
sim.4.lpWWitth Σ - mlm |Q]2 |>] s ((0
[T(A) + |Q| |
noruunormBound max (sim, (f
max | ,
max 2 (TA() + + (a)
Proof. At this point, most of the work has been done for
We create separate channel for each component permutation-resst semiautomaton A This requires total Σ sim, embedding dimensions. In addition these channels, we keep one dimension (with residual connections throughout the network) to represent the input Ø, Let ex denotes the unit vector along this coordinate. Choosing an arbitrary enumeration to identify Σ with [ we select the embeddings to be E(r) a. exe
The L layers of A are divided into |Q] subnetworks (which are just Transformer networks), which we will concatenate sequentially at the end. Let these subnetworks be indexed by E' {I...., L}. Each subnetwork starts with a parallel simulation (as in the direct product construction of Lemma 8, padding with layers implementing identity if their depths do not match), combining all of the sim,.0 the ě-th level of the Krohn-Rhodes decomposition, as defined by Theorem 7. Each sim, 0 chosen to operate in its own channel i. We add residual connections on all of the input/output dimensions in each channel. Then, at the end of each subnetwork except the final one (1 l < I-1), we append one more Transformer layer with identity attention block, whose MLP implements the "wiring" specified by ())) from the next level of the decomposition.
Namely, we invoke Lemma 2 with A = 1, B = max ( [QQ), [Σ]), By = [E], giving us for each pre-final-layer an MLP which represents the function
(sim, (Q
( E(o)) E(p(g ,...., () o));
where the inputs are stored in the respective < and Σ channels, and the output is written to the i channel. Here, the number of input dimensions is
din Σ sim + < 2 (IT(A) + QQI + 4) + 1 (
Since the state encodings for each predecessor semiautomaton 1' < i are IQ| (olundddded inte- ger vectors and Σ has been assumed be bbu-ddunded positive integer, suffices to use d = 4d,, max Σ}}. The second hidden layer's width d2 is the number of possible inputs |.) which is bounded by [Q)2/91 |Σ| > d1. The weights satisfy
| I 4, || S 6 max max {10),
|| $1, 1, ||22]|.0 S din, ||Wall> S |Σ], by 0.
Between in the same layer, these routing constructions need to be executed parallel, so this incurs another multiplicative factor in the width, bounded conservatively by 2101.
The final construction concatenates these blocks, so that at the output of the last layer, every channel i contains representation of its corresponding component's semiautomaton Qi- The W guaranteed
by Theorem 7 suffices for the overall choice of W.
C.4 PROOF OF THEOREM 3: EVEN SHORTER SHORTCUTS FOR GRIDWORLD
Recall the gridworld semiautomaton in Example where the state (Q (0, 1,.., to the adjacent state based upon seeing input token L R (modulo boundary effects), stay unmoved upon seeing I. More formally, the transition function is defined as:
, S} either move
5(g. L) max(g 1,0)
8(g. R) = min(g + 1,5)
8(g. 1) q.
In this section, we will show how to implement gridworld simulation using only Transformer layers. Here restate the theorem in full generality:
56

---

Published as a conference paper at ICLR 2023
Theorem 3 (Even shallower shortcuts for gridworld). For each positive integer T, Transformers can simulate the (S+1)-stare gridworld semiautomaton with attention layers, where the MLP has either (i) depth O(log 5), width O(T + S), or (ii) depth ((1), width O(T) + 2015), The weight norms bounded by poly(T).
The depth in (i) can be reduced to ((S) if we allow max pooling. and the dependence on T in the width can be removed with sinusoidal activation. We discuss this in detail after the proof along with generalization to the k-dimensional gridworld case.
Note that, in order to find the current state, need to only know the most recent time which the semiautomata was at boundary. It not immediately obvious how to compute the most recent boundary, if one is not allowed
the trivial sequential simulation algorithm. Our key insight is that this boundary detector can be computed without needing to parse the entire sequence, using the most recent S 1 distinct values of the prefix sums in the sequence.
This algorithm is especially well-suited to the Transformer architecture since: (i) the prefix sum can be computed using one attention layer in Lemma and (ii) the identification of distinct values can be implemented by sparse value-de malue-dependent lookup similar to the memory lookup in Lemma 7 with the help of the self attenntoon context-deeendde and (iii) the positional weight sharing and causal masking enable all of these computations to be performed in parallel. Overall, Theorem consists of concise implementation which executes all of these posst-reeet-bounddee detectors in parallel.
retrieval, as opposed a static lookup),
In what follows, we first describe the algorithm (Algorithm 1) for computing the state of the semiau- tomata using the S + 1 distinct prefix sum values, and give proof of its correctness. Subsequently, we formalize the Transformer construction that implements the algorithm. A consolidated list of notations used in the algorithm well the proofs is provided in Table for the reader.
C.4.1 THE ALGORITHM SOLVING ID GRIDWORLD
To convey the essence of the full construction, we first provide pseudocode (rather than Transformer weights) for computing the final state qr (rather than the entire state sequence).
{L, R. 1) T E 1,1,1,0}, i.e. L 11, R + 1, and 1 O. Let 70) denote the sequence of mapped actions, and let O be the initial state. The algorithm (Algorithm 1) has two steps: first, we identify the last time the agent is a boundary (wall) and the type of the boundary (i.e. state 0 or state 5). The final state is then simply the sum of all actions in the sub- sequence, shifted by the last boundary, which is easily computable with I attention layer (Lemma 6). Our key insight is that we can identify the boundary using O(S) attention heads in two attention layers, and therefore do not require a recursive computation from the start state (with depth T).
We map actions
To show the correctness of Algorithm 1, it suffices to show that the boundary state is detected correctly, since after that there is no more boundaries and the only step remaining is to calculate the sum of the actions. Let tuniq: min: max be as defined in Algorithm 1. Then:
Lemma 13. If fain > fmax then state at min is O, otherwise state tuax S.
Proof. The proof follows from two observations:
1. minftmins {max) tuniq.
2. Suppose tuin funig (the argument is symmetric for fuax tunig). Then Atmax S.
First note that E {=1} which implies that the prefix sums increment or decrement by at each index. Therefore, 2tmas - = S + 1. This also applies that between (and including) tmax and bmins there must be indices such that they traverse the S + 1 distinct values. Since we take the shortest suffix satisfying this, bunig 2 min /min ). This proves Observation 1.
Assume mii Luniq- We can break the analysis into the following 2 cases:
(a) The S + 1 distinct values correspond to S + distinct states (covering both boundaries). This implies that the minimum and maximum out of these distinct prefix sums must correspond to the boundaries, that is, Itmax S and It..n. 0.
57

---

Published as a conference paper at ICLR 2023
Notation
Definitions
σ
An imput token; a E R, R, J..
A mapped input token; = = = R, and ē = 0 =1.
Used in Algorithm 1:
The state at position t E [2]; ge E 10, ,
Z
Prefix sums all T positions.
tniiq
The most recent position for which the prefix sums tm contain S + unique values.
maxx bmin
The positions corresponding to the max/min prefix sum among positions lunigu... T.
tfiall
The position of the last boundary state, defined as ffmnal max mmax tmin).
Used in the Transformer construction:
Yt
The positional encoding for position E [T], defined as "/1 == log(2T - f).
Wattn [t]
(1)
The output of the first layer attention position E [T'] defined as (1) [t] 27 Eielg 8i.
w mlp
(1)
[t]
The output of the first layer MLP at position / E [7]. defined as =mpp [t] = | [1]
Yes
1, cos(2 [)), sin (1) [)))).
The position which achieves the max attention score for the tt head at time / (T) (t is omitted for notational convenience). for s [0, 1, .2S].
Tiaxx
x aattn [t]
The output of the second layer attention at position E [T].
defined as
Pattn
h Jmaas , -" (,5))
Jaxx
x (2)
"mip [t]
The output of the second layer MLP at position / E [7]. which gives the state at t.
Table 1: Notations for the proof of Theorem 3.
(b) The S + 1 distinct values correspond fewer than S + distinct states. This implies that only one of the two boundaries visited in the sequence starting from luniq- In order to get S + distinct values, it must be that this boundary wall is hit, i.e., the sequence tries make move that the boundary blocks. If the sequence does not hit a boundary, then at every time the same state revisited, the prefix sum must be the same, and values. Since tmin = fuiig we claim that the visited boundary must be S. Suppose this is not true, then the boundary visited is This implies that Immi O. Since the sequence does not hit S, at any position it is at state 0 before hitting the wall, the value will be 2tmm. Thus, when the sequence first hits the wall (say index T), then 2, 2twin which is not possible by definition of mmin- Thus, the boundary must be S.
will not be able to get S S+ distinct
Given the above, our algorithm identifies the boundary correctly and then can just use the prefix sum to evaluate the current state.
C.4.2 TRANSFORMER CONSTRUCTION FOR ALGORITHM
In this section will show how simulate Algorithm uusing 2-layer Transformer with 28 attention heads.
In construction, the first attention layer will compute the prefix sums. This can be mapped to a cyclic group from Lemma 6, however for completeness, we will restate the main construction. The MLP in the first layer will map this prefix sum to circular embedding
Proof of Theorem 3.
58

---

Published as a conference paper at ICLR 2023
Algorithm 1: 1D gridworld: computing the final state
Data: Ø E {+1]², S E Z
Result:
Final state yr E 10. .S}
// Pad tokens so there are at least S + distinct values
-,.
S-1
Calculate the prefix sum for each index
x + prefix sum(i))
(i.e. 2₁ f Σラ= σT) Find substring containing 4-1 unique values
liiig < max [set(zr:)]
= S +
// Find positions of the last max and min values
min - max 24 min+2(unig 2
max < max [ 2z m 2r}
// Identify the type of boundary
Emin > fmax then
boundary + 0
else
boundary + S
end
// The final state is the sum of the substring after the last boundary
fuall f max(tmiu mii maax]
yr 2T
Ainn + boundary
Case 1: traversal
Case 2: hit one wall
(right boundary)
Tiail
Firool
(left boundary)
Fusiq
L L
L
R
Figure 16: Illustrated examples for 4-state gridworld (Q = 1, 2, 3)). The algorithms compute the prefix sums z (as if there were boundaries; shown green dots); intuitively, it might seem like they have no direct relationship with the state sequence ze (gray dots). However, defining lunig as the start of the shortest suffix containing S + 1 distinct prefix sums, and fint as the most recent minimum or maximum point within this suffix, our case analysis shows that Itanal is always boundary, resulting parallel simulation algorithm.
(see Proposition 5). The second layer attention will the circular embedding structure to find S + closest distinct values to the current value z, (suppose we are considering position / E [7]) by identifying the positions for closest values in the set (z+ - S, 2, S 1,..., , 2t zz + 1, 1, + S}. i.c. S closest distinct values smaller than 2,, and S values larger than 2,. This closest
21 2.... distinct value construction can be viewed as a position dependent flip-flop monoid construction, where we need to identify the closest position with particular action. Note that this set of values would contain the distinet S + values needed by the Algorithm 1, hence the second layer MI.P can implement the state computation using these values.
Input representation:
We select input symbol embedding E(o) = T E Rd where ō is if o L
the action corresponding a, that is,
1
if R otherwise
We will use positional encoding
59

---

Published as a conference paper at ICLR 2023
Pt.: == YeC,, where "t log -) is such that e 27
1. We will include extra position L,
with embedding F((1) e2 and position encoding PL.: 0. Think of this padding at position O; it is not masked by the causal attention mask.
Prefix sum (Layer attention):
The attention construction for the first layer, in full detail:
We select d = 4, k = 1, H. H =
Select Wq == c3, WK == e2. :v
== E1, W == E4.
With this attention module, the 4th channel of the output at position is (1) [t] which is the prefix sum scaled down by 1/21..
27 Eiety 8i.
Circular embedding (MLP D):
(i The first MLP maps Faten [ + (a , sin (aati M }]
where cos, sin are calculated up to Olog 7') precision using the construction in Lemma I with width 4(2T + 1). 27 and weight norms at most 8T. Together with the input using the skip connection (for
%) we get mipp [+] == aetm [t), Ye. 1, cos(2 (1) [T], sin(x (1)
aten [t]m
as the embedding be input
to the second attention layer.
Finding closest S + 1 values (Layer 2 attention): Our goal is to find the shortest subsequence distinct values for attn that is, we want to find the max < t - S such that Kzattn [+] = S + 1. We will do this by using 28 + heads such that Vs E {0, 1, ... ,225, the attention score for the Sen. head on position i E [7] satisfies
(looking back from the current position t) that contains
(1) a (ww)))) x =mppt [t], (w
c loool22 (27 i),
if atta [+] =0=2 [t]
(1)
(1)
S-
27
e clog(2f - i)- 87
otherwise.
where (16 2
That is, for any i.j s.t.
aatt [
Fata
(1)
[
27
(matched) and
(1) [i) * Fasen [t] +
(unmatched), the difference in the unnormalized attention weights is lower
2T
0
bounded by Ci -a >
This can be achieved by letting w(())
1
167
Patn)
5 where Pate) the rotation matrix of angle ((s) =
(e=S)x
such that (Wr()))
T
mlp [t]
21
(o, (o--c.0, cos ((zattn [t] + {كي π.) sin ( (mal [t]
+
8-8
w(s) w-(s) are simply the x
27
identity matrix, and w=(*) e + CzcJ.
Let Jmax denote the position that achieves the max attention score for the th head, then the output of the th head is [x: [[ 3022 0, 0]. We can ignore the last three coordinates (which are 0) as well as 2 [[(]] since we will only need the difference (1( [t] - ==(() [i(e] which is s-S by definition. We then concatenate the outputs from all (28 + 1) heads in a (2S + 1)-dimensional (,0))) as the input the second layer MI.P.
22""
vector vector...2 - = (0) 7,a)
Intuitively, each head in the second attention layer is trying to identify the set of positions for which the prefix sums match particular value specified by the head. Each head selects the last matching position if such positions exist, and selects f. if not. The following observation will be helpful for
The width is if we allow sinusoidal activation instead of relu; see discussion after proof. We assume hard attention here for case of exposition of the proof; soft attention can be handled with Lemma 4 and Lemma I previous constructions. This requires norm poly(T) max.
60

---

Published as a conference paper at ICLR 2023
(2) our subsequent MLP construction: the values of coordinates of ssttn increase both sides of the 5th heads; that is, atm satisfies the following:
Lemma 14. There exist a < E 10,1,..., 25} such that
Fatu Fatiala [a] > x Jatia [a + +1] > ... >..>> x (2) Pattlss < Fatiuls IS + 1] < .... (2 [b]
(2) and all 8 E (0,1, ,,,,, 25} \ fa + l,...,b} we have x Jasen [s] log(2I' t).
Proof. Note that x attm [+] = log(2T - 7,0 which makes the ordering inverse of the position. Observe that the unmatched indices correspond to values that have not been reached. This can only happen for values on either the leftmost or the rightmost coordinates, since the prefix sums continuous on integers. Now let's prove that the value will be decreasing moving away from index in both directions. Suppose this was not true, there indeed was 8 S such that (12) [s] 2 " Is - 1] (s 2 S case is identical), then it implies that the closest index that achieved relative value s further away from than S 8 + 1. However, since the moves can update the prefix sum by magnitude at most I, then to get to relative value 0 from relative value S 8 + 1, we would need to have crossed relative value S - s. This implies that there another position closer to with this value, contradicting our assumption. This proves the result.
(2) Computing state (MLP 2): To compute state from the positional information given by Fattn need to do the following computations:
Step I: consider S + windows of size-(S + 1), each containing the st to (s + S)th heads for
E {0, 1,..,, S), and identify the window that contains positions closest to the end (this would correspond to the closest S + distinct values to (1) [t]);
Step 2: identify the boundary state in the selected window by comparing the indices of the endpoints of the window;
Step 3: output the final state based on the position of the boundary states and its value relative to the current position t.
We will show two constructions for implementing this, one of which will use 0(1) depth and 20(5) width, and the other will use Oloog S) depth and O(S) width. The trade-off essentially lies in how a min function is implemented and can be resolved we allow min-pooling layer, which we will discuss after the proof.
depph construction: The idea is that we can first use O(1) layers construct "features" that contain all the information needed determine the state, then 3-layer MLP with 20(5) width can compute the state as function of these features by Lemma 2. The features we need are the following (the labels underneath are to be consistent with Figure 17, left):
{ 1 [1 Tattn (s + $1]}>
(2)
(C.3)
{ [1[ [1[ 1] > x (2)
(2)
Fattn +S]1
[1
[s + S i]]
(C.4)
71
>R
{ [1(2(e) (2) [s] = loog(2 t)]125
(C.5)
Here the feature in C.3 compares the end points of the S + 1 windows, the two features in C.4 compare the window with its adjacent windows on each side, and the last feature in C.5 will be used to climinate the irrelevant window. Features in C.3 and C.4 can each be computed threshold function (at 0) on the difference between the two elements to be compared, which can be implemented using layers by Lemma
61

---

Published as a conference paper at ICLR 2023
(a) First, show that to compute Step I we only need to compare between adja- cent windows whose + 1 heads are all matched, which can be computed using the features above. Consider any window starting at E [0, 10,1, .]. On ei- ther side, if this window is closer to than its adjacent window on this side, then it is closer to the end of the boundary than all the windows on this same side, which would imply that we can ignore these non-adjacent windows. To prove this, let's consider the left side (the right side analogous) and we want to show: For any E [S], if max(" (2) [s), " S]] < max maxxfattt (s 1]. [s + S - I]}, then max(= (2) [s], x [s + 5}} < max max [s i), (2) [s S :]} for > 1: the if con-
dition gives us x attn [s 1] > Jattn [s + 5], and we know from Lemma 14 that
(2) is i| > Fatta [s 1] Fatin
(2)
(2)
aatnn
(2) attn
8 + S] attm [s + S - 1] Fatin [s S -]].
Combining these inequalities together concludes the proof.
(b) Given the optimal window, we can use feature C.3 for the relevant window to identify the boundary, since the closer-to-t index gives us the last boundary state (see Algorithm I for why this suffices).
(c) Now that we have identified the boundary state (suppose it is at position i), the final state
can be computed as the last boundary state (0 or S) plus the difference 2 = [t] - x [t] The difference is built in to the ordering of the heads, hence we have all the information to compute the final state and we done.
Therefore, we can compute this function using 4S + 3 features each taking value in {0, 1} and the output having S 5 1 values. These features themselves can be constructed using Lemma 3 with A = 1/4T since the indices are separated by at least this gap. For the indicator index, we can compose two such constructions similar to Lemma 1. This gives the first layer of MLP with width ((S) and norms O(T). After this, the rest of the function can be constructed using a 3-layer ReLU network with width 20(5) and norms bounded by O(5) using Lemma 2.
2. Oloo -d-pphhh construction: An alternative solution to the above is to pay O(log(S) depth, but reduce the width be OSS). We will borrow features in equation C.3-C.5, but construct the MLP explicitly rather than calling Lemma 2 as black box: the width and depth trade-off essentially correspond to two ways of implementing the min of S numbers. We describe the corresponding MLP by components (Fig 17, right):
(a) Ignore the unmatched heads; preprocessing step for the cleanness of the proof, we use I I-layer MIP to map (2) [s] for heads where .i(2) [s] = log(27 - t) (i.e. Jmax = t) to 1og(27) such that these unmatched heads can be ignored in the following steps. This can be done by multiplying log(21") to the threshold function given by Lemma 3 (with = } where the network has I hidden layer with width 25 + 1 and x-norm 47'. Note that this map also changes 2 [S] (i.e. the center head that corresponds to position t) but this will not affect the correctness of the proof.
(b) Compute the function
in Figure 17 (right), which computes fi b) (max(a, b), Ia > 6]) for = x 2i (s], 6 = (2) [s + S]}. Vs E (0.1,...,55) The first coordinate max(a, b) can be implemented using I hidden layer with width 1, and |a > b) is the same as feature C.3 and can be implemented with 1 hidden layer by Lemma 3. There are S + choices of 8, hence the overall width is O(S). For notational convenience, let's denote (s) == f1 (= [s], x (2) [s + 51) (with slight abuse of notation).
(c) Find the min value of fa (s) [1]. denoted as f1,min == min, fi [s)[[]]: This can be achieved using 1 min-pooling layer. we allow ReLU only, then this can be implemented with pairwise comparison using network with [log S + 1] depth, 35 width and and constant weight norm." 29
(d) Compute the function ((() in Figure 17 (right): f!() takes two inputs: 1) the second output of f1, which we denote as B, == [1() and 2) M, == 1 [fi.iin I (s)[[]]. which
The log depth is conjectured be unimprovable; see discussion after the proof
62

---

Published as a conference paper at ICLR 2023
state
state
SLmM
2 width
layers
2
hidden layers
S hidden layers
7
R
hiddon layer
hidden layer
X1
X2
X3
X4
X5
X
X0
21
22
X3
X4
X5
X6
Figure 17: Illustration of the two constructions for second-layer MLP in Theorem C.4, with S = replace x1()) with Left: O(1)-depth solution where the first
For ease of readability, block compares the comparison and equality features (see equation C.3- C.S), and the second block computes the state from these features. Right: (log deept solution where first block implements fi (a. b) (max(a, b). 1(a > b]]. second block does min min-pooling operation, the third block implements f2 which takes fi via a residual connection and output of min-pool to compute the final state.
indicates whether the 8th window is the closest-to-l window not and can be computed using a 1-hidden-layer network with width 2. As in the previous construction. the difference fi,min [1] is built-in in the ordering of the head and hence does not need to be passed in explicitly. Then by Lemma 2, a 2-hidden-layer network of width O(1) can take B, M, as input and compute M, [BS - 8) + (1 (2S 1- 1-s]]. The overall width is O(S)
for 1 choices of such f(())
(e) Finally, the state is computed as E, 12 (B,, M,), which can be implemented with layer of width 1.
Improving the construction to remove T width and log(S) depth. Using standard architectural tools, such as max-pooling, we can improve our construction to get Od1)-depth and ((S)- -wiith for the MLP.
Avoiding width T in the MLP using periodic activations. As in the modular addition (Lemma 6) construction, we can use sin activations in the MLP to directly compute the circular embeddings that are used as input to the second attention layer. This would require only two hidden nodes in the MI.. Note that we do not need precision greater than Ollog 7') for these activations since embedding values only close 1/poly(T))
Avoiding log(S) depth the MLP 2 using max-pooling. The O S)) in MLP 2 incurred by calculating the min of S numbers and is conjectured to be necessary for ReLU networks (Goel et al., 2017; Mukherjee & Basu, 2017: Hertrich et al., 2021). However, the depth can be reduced to if we allow max-pooling layers, which are commonly used in both theory and practice (Zhang et al., 2021b; He et al., 2016; Vaswani et al., 2021).
Remark: Yao (2021) use layer-norm to compute cos and embedding with non-uniform angles. This could potentially alleviate the width T' concern; leave this exploration to future work.
Extending beyond dimension. Since a 2-dimensional gridworld is just the direct product of 1-dimensional gridworlds (by the construction in Lemma 8), implement both dimensions in
63

---

Published as a conference paper at ICLR 2023
parallel by concatenating the network for each dimension. This be done by doubling the dimen- sions, parallel attention heads, and parallel hidden units the MLP. The attention head parameters for each dimension be chosen to only focus the relevant dimension and similarly the MLP can zero out dependence on the other dimension. We can extend this to higher dimension with a multiplicative increase in the size of the parameters.
C.5 PROOF THEOREM DEPTH LOWER BOUND FOR NON-SOLVABLE SEMIAUTOMATA
Theorem 4 (Transformer Barrington). Let A be non-solvable semiautomaton. Then, for suffi- ciently large T, no fixed-precision Transformer with depth independent of and width polynomial in T can simulate at length T, unless TC NC².
Proof. This follows straightforwardly from the fact that simulating A length T is NC complett under NC reductions: given any Olog ddephhh bounded-fan-in AND /O// OR/NOT eircuit C, and a depth-D circuit C' which simulates a semiautomaton whose transformation monoid contains non-solvable subgroup. there procedure which generates a (epth-O(D) circuit to simulate C; see (Barrington & Thérien, 1988). This in turn comes from the construction used in Barrington's theorem (Barrington, 1986), which characterizes NC² as exactly the set of languages recognizable by bounded-wied branching programs. For a closely related reference which follows almost exactly the same argument, see (Mereghetti & Palano, 2000).
Thus, it suffices to show that constant-deply Transformer is in TC The details of manipulating floating-poo point numbers with discrete circuits are peripheral to the main results in this paper, so we provide a brief proof sketch. A similar argument is used by Merrill et al. (2021) to establish that hard-attention Transformers), with Ollog T)
"saturated" Transformers multi index analogue bit precision, can be represented with a TC circuit. We outlnee a proof (which applies to the formal setting considered by Merrill et al. (2021)) for the notion of Transformers defined in this paper.
With (log T) bits of precision, all n-way (including unary) arithmetic operations mapping TR" can be represented with constant-depth, -oly(T)-width AC circuit, as long as does not de- certainly possible, it suffices consider the circuit which
pend on T. Although improvements memorizes the i-th bit of the output, which has width 2 log (poly(T)). Thus, the position- wise non non-interacting matrix operations (multiplication by X + WQ etc., the feedforward MLP layers, and the encoding and decoding layers) can be simulated with poly(T)) width.
The only subtlety arises when there is T-wwyy summation over Oloog T'i numbers, which occur in the softmax and attention mixture layers. For this operation, we can use the construction from (Reif & Tate, 1992), which can even add poly numbers TC
64


Reviews:
[{"review_id": "vxsIZu3Q1d", "reviewer": "Reviewer_cJyA", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This is a theoretical work about how transformer models can perform symbolic tasks.  The class of tasks considered is semiautomata (deterministic finite state machines).  The focus is on how a shallow network can compute sequences of states of length L much greater than the network depth D.  It is shown that this is indeed possible, with D \\sim log L at worst, and for a subclass of automata, constant D.\nExperiments are done to validate these predictions and to give evidence that these networks are learnable from data.\n", "summary_of_the_review": "A novel approach to understanding transformers which leads to significant results.", "strength_and_weaknesses": "Strengths: very interesting question and approach, solidly grounded in theoretical computer science, clearly written, strong results.\n\nWeaknesses: the scope and complexity of the ideas make them hard to fully convey in a paper satisfying the length constraints, even with the supplementary material.  In particular a lengthier comparison with the related works (and those in section A.5) would be of great interest.  We hope the authors will write a longer treatment elsewhere.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "tpuOJ1WICz", "reviewer": "Reviewer_zU3W", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper is a contribution to the understanding the capabilities of transformers to model languages. The first part of the paper provides results on the capabilities of (shallow) transformers to simulate computations of semiautomata at length T. The second part of the paper is an experimental study on whether learning allow to find such shallow transformers", "summary_of_the_review": "I am clearly not expert enough to give a firm assessment of the paper. I tend towards acceptance because the theoretical results seem meaningful and sound to me. I would like to have more details to assert the significance of the experimental results.", "strength_and_weaknesses": "I am not expert in Computational Theory and Complexity Theory.\n\nPros\n\n* The theoretical contribution is dense but clear\n* The theorems seem sound\n* Experimental results complement theoretical results\n\nCons\n\n* Could be published at a conference but should be submitted at a review\n* Experimental protocols should be made more precise\n* I am not so convinced by the lessons drawn from the experiments", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "7_XkO12h6w", "reviewer": "Reviewer_Czkt", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the relation between Transformer depth and its ability to solve algorithmic reasoning tasks. The authors formulate reasoning processes in terms of automata, and particularly focus on the question of why and if shallow Transformers are able to simulate automata with much longer horizons. A number of results are presented of the following flavor: Transformers of depth $o(T)$ can be used to simulate automata of length T so long as the width (of e.g., the MLP blocks in the transformer) is big enough, and enough attention heads are used. The width and number of heads needed is quantified in terms of the size of the state space. The shortcuts referred to in the title of the paper refers to the fact that these results show that Transformers can learn to simulate long processes without directly learning each step of the process explicitly. ", "summary_of_the_review": "In all, I am strongly in favor of this paper’s acceptance. I cannot pretend to have fully understood the depths of the ideas in this work - for instance I have not yet had the chance to study in detail how the attention mechanism is able to compete multiple automata steps in parallel (this is probably a good moment to make the disclaimer that I haven’t properly verified any of the sizably appendix to this paper, so all results are taken on face value). However, I have read and understand more than enough to conclude that a this paper makes a contribution that is easily sufficient to merit acceptance at ICLR. I will be interested to see if any points are raised by other reviewers that I wasn’t aware of, but I doubt anything will change my overall view. Congratulations to the authors on an immaculate piece of scholarship, I hope you continue to develop this line of work further. ", "strength_and_weaknesses": "This is a highly unusual ICLR paper. This is not to say that ICLR is an inappropriate venue for this work, since this is clearly not the case: the subject matter is of direct and high importance to the machine learning community. This work is unusual because of the breadth and depth of ideas and connections to different parts of the theoretical computer science literature, as well as very topical considerations to the ICLR community such as out-of-distribution generalization, spurious correlations, and neural algorithmic reasoning. For instance, the results presented in this paper are formulated in terms of the Transformer’s ability to simulate automata—-a now classical concept in AI (albeit being essentially conceptually the same as recurrent models). The analysis tools are also unusual, with many algebraic arguments used, and results from classical complexity theory are drawn upon. \n\nA few brief strengths:\n\n- A complete set of theoretical results, including general results for any arbitrary automata, sharper bounds for certain (solvable) automata, even sharper bounds for a particular class of automata (grid words), and lower bound results.\n- Experimental results validating a number of these findings on a number of algorithmic reasoning tasks, taking pains to use the experiments as a chance to bridge some of the inevitable gaps between the theory and practice (such as checking whether gradient-based training actually finds the shortcuts the theory identifies).\n- Taking care to at least discuss a number of questions a reader might have whilst reading this work, such as the connections to universal approximation, \n\nNo weaknesses come to mind, and I do not wish to work to artificially come up with problems. That’s not to say that the work is perfect, of course it isn’t, but I have no major concerns about the validity or significance of the results. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-5%