Transformer Working Memory Enables Regular Language Reasoning And Natural Language Length Extrapolation
Authors: Ta-Chung Chi
Ting-Han Fan
Alexander Rudnicky
Peter Ramadge
Venue: EMNLP 2023 Findings (2023)
Keywords: Transformer
Algorithmic reasoning
Length extrapolation
Working memory
TL;DR: Transformer Working Memory Enables Regular Language Reasoning And Natural Language Length Extrapolation

Abstract: Unlike recurrent models, conventional wisdom has it that Transformers cannot perfectly model regular languages.
Inspired by the notion of working memory, we propose a new Transformer variant named RegularGPT. With its novel combination of Weight-Sharing, Adaptive-Depth, and Sliding-Dilated-Attention, RegularGPT constructs working memory along the depth dimension, thereby enabling efficient and successful modeling of regular languages such as PARITY.
We further test RegularGPT on the task of natural language length extrapolation and surprisingly find that it rediscovers the local windowed attention effect deemed necessary in prior work for length extrapolation.

Transformer Working Memory Enables Regular Language Reasoning And Natural Language Length Extrapolation
Ta-Chung Chi
Ting-Han Fan
Cameegie Mellon University
Princeton University
aachungc@andred cmu. edu
tinghanf@prrnncetonn. edu
Alexander I. Rudnicky
Peter J. Ramadge
Carnegie Mellon University
Princeton University
airOcs. cmu. edu
ramadge@princetoon edu
Abstract
et al., 2020; Armeni et al., 2022), and Transform- ers' (Vaswani al., 2017) salient tokens height- ened by self-attention.
Unlike recurrent models, conventional wis- dom has that Transformers cannot per- fectly model regular languages. Inspired by the notion of working memory, we pro- pose new Transformer variant named Regu- larGPT. With its novel combination of Weight- Sharing, Adaptive-Depth, and Sliding-Dilated- Attention, RegularGPT constructs working memory along the depth dimension, thereby enabling efficient and successful modeling of regular languages such as PARITY. We fur- ther test RegularGPT on the task of natural language length extrapolation and surprisingly find that it rediscovers the local windowed at- tention effect deemed necessary in prior work for length extrapolation.
In reality, self-attention often attends broadly (Clark et al., 2019), violating the limited amount of information notion of WM. Our hypothesis is that such violation to blame for Transformers' failure on algorithmic reasoning of regular languages (Deletang et al., 2023; Liu et al., 2023) such as PARITY, seemingly simple task that checks if the number of is in bit string is even. Surprisingly, Transformer can only count the number of is correctly when the sequence length held fixed training sequence length Tir, and fails miserably when the testing sequence length extrapolates to Tex > Tir (Hahn, 2020; Bhattamishra et al., 2020; Chiang and Cholak, 2022; Deletang et al., 2023; Liu et al., 2023). In contrast, an RNN can extrapolate perfectly.
1 Introduction
It is long believed that Working Memory (WM), a term coined in 1960s to liken human minds to computers, plays important role in humans' rea- soning ability and the guidance of decision-making behavior (Baddeley and Hitch, 1974; Baddeley, 1992; Ericsson and Kintsch, 1995; Cowan, 1998; Miyake et al., 1999; Oberauer, 2002; Diamond, 2013; Adams et al., 2018). While no single defini- tion encompasses all applications of WM (Adams et al., 2018), the following one should be shared by all of the theories of interest:
The goal of this work is therefore to improve Transformers' WM by limiting the amount of ac- cessible information at time. Existing attempts that use combination of scratchpad and recency biases (Wei et al., 2022; Nye et al., 2022; Anil et al., 2022; Liu et al., 2023) is not optimal it completely foregoes the parallelization property of a Transformer, making it as computationally ineffi- cient as an RNN.
This begs the question: Does there exist more efficient Transformer working memory design? The answer is affirmative thanks to the proposed Reg- ularGPT, which boils down to the three design choices: Weight-Shhring, Adaptive-Depth, and Sliding- Dilated-Attenin
Working memory is a system of compo- nents that holds a limited amount of in- formation temporarily in a heightened state of availability for use in ongoing processing. Adams et al. (2018)
Each of them has been proposed previously but it is the unique combina- tion that sparks the successful and efficient learning of regular languages. We will further demonstrate its: 1) similar recursive parallel structure as lin- ear RNN (Orvieto et al., 2023), resulting log Ter or log Tex layers, and 2) generalizability by show- ing strong performance on the task of Transformer
WM is instantiated in the two major driving forces of sequence modeling: Recurrent neu- ral networks' (RNN) (Elman, 1990; Jordan, 1997; Hochreiter and Schmidhuber, 1997) short term memory modulated by their recurrent nature and gate design (Rae and Razavi, 2020a; Nematzadeh

---

natural language length extrapolation (Press et al., 2022: Chi al., 2022, 2023).
Previous work sheds light on this by showing that a Transformer might settle on the naive-summation approach (Anil et al., 2022; Deletang et al., 2023; Liu et al., 2023). Concretely, it sums up all the bits and outputs the summation modulo 2. This ap- proach fails since unseen summations will be pro- duced when the model takes sequences of length Tex T as input or PS,)) deviates from 0.5.
In this work, we use [N] to denote the list of The Trans-
non-negative integers [o....,,N' N -1. former model used in this work is always causal. It takes in an input sequence of T < Tir units (can be tokens or bits) [E[TT]. passes them through fixed amount of L transformer layers, and finally computes the distribution over the vocabulary V via the prediction head W..
To the best of our knowledge, the existing rem- edy (Liu et al., 2023; Anil et al., 2022) is to use scratchpad (Wei et al., 2022; Nye et al., 2022) along with recency biases (Press et al., 2022) to enforce the correct learning: They create scratchpad that interleaves the sequence of input bits and interme- diate answers (01, 02; Q2, " oT; (r), where 9i = solve -.. o.). The model is trained to pre- dict all the S.ery. Recency biases play the role of limiting a Transformer's ' receptive field to only a few most recent and q at every timestep i. This is to prevent self-attention from ignoring @ and giving the same naive-summation solution.
2 Background
2.1 Regular Language and Algorithmic Reasoning
The Chomsky hierarchy (Chomsky, 1956b) clas- sifies formal languages into different hierarchies based on their increasing complexity. Each hierar- chy represents family of formal languages that can be solved by the corresponding automaton. At the lowest level resides the family of regular lan- guages, which can be expressed using finite state automaton (FSA), a computational model compris- ing a set of states and transitions connecting them. Our primary objective is to enhance the algorith- mic reasoning of the Transformer model on regular languages by testing its language transduction ca- pability under the extrapolation setting. Concretely, the model is trained only to predict desired outputs on a set of short length-T sequences with T < Ter- Still, it must also predict the correct outputs for longer testing sequences of length Tex >> Tir. It is worth noting that we evaluate our model via language transduction following recent work (Dele- tang et al., 2023; Liu et al., 2023), instead of the conventional language recognition protocol. Both settings are equally hard as they are underpinned by the same finite state semiautomaton. Interested readers may refer to Deletang et al. (2023) for fur- ther details regarding the two evaluation protocols. We also reveal the connection between RegularGPT and finite state semiautomaton later in $7.
Scratchpad and recency biases jointly create the notion of WM along the temporal dimension simi- lar to RNNs, thereby enabling successful extrapo- lation on regular languages. Nevertheless, we note that this fix is inefficient during inference since all the intermediate answers qi have to be generated sequentially before reaching the final answer ÇT- A desirable fix should only take in the input bits (01, 02 .T.) and directly generate the final an- swer 4r. In other words, our goal is to find an efficient WM design for a Transformer.
2.3 A Desirable Fix for PARITY (Figure 1)
An alternative solution to the PARITY problem is based on the spirit of divide-ande -nd-uuq where we first divide the sequence into T/C chunks with each chunk of length C' T, and we compose the final answer by recursively merging the chunk outputs. This approach does not suffer from the unseen summation issue as the model was trained to handle fixed amount of C' bits at a time in its WM (chunk). It then recursively applies the alrcady-seen results compose the final solution when it encounters longer sequences during infer- ence. More importantly, it is more efficient than the scratchpad and recency biases approach since it only requires logg T layers of parallel computa- tions instead of 27 steps of sequential decoding.
2.2 Failure Mode and An Inefficient Fix
The PARITY task involves a length T bit string 01 oT where each bit o is randomly sampled from Bernoulli distribution with P(a; = 1) = contains an even or odd number of 1s.
0.5. The goal is to determine whether the sequence
has been observed that a Transformer is inca- pable of performing length extrapolation on PAR- ITY, but what could be its potential failure mode?

---

Layer = Layer
Layer
thickness
0
FFN
FFN
FFN
1
0
1 0 0
1
0
0
Figure 1: This is the seride-and-conqqee approach that solves the PARITY problem. The lightly shaded blue cells represent Min in eq. (I). The darkened blue cells represent the routing path to solve the result for the last bit specifically. As we can see, this approach requires at most log, T layers to obtain the result for length input efficient approach compared to the combination of scratchpad and recency biases.
sequence, rendering it
3 Proposed Architecture of RegularGPT
et al., 2022, 2023). M will then be added to the original self-attention matrix, creating the pro- posed Sliding- Cieding-Dilated IillatedAAttettion effect. The out- put of SA (1) will be transformed by the positional- independent FFN!)) to produce (e(T--
We present our modifications to the vanilla Trans- former below. Only the related operations will be expanded, and we follow all the other details of GPT2 (Radford et al., 2019).
(1)
The case of C = 2 used as possible construc- tion of Theorem in Liu et al. (2023). However, their focus is not on length extrapolation, hence lacking the below two proposed modifications.
3.1 Miding--DiateddAAtenniii
A Transformer layer layer / consists of self- attention operation denoted SA and feed- forward network denoted as FFN Originally, (A computes the inter-token relationships across all T units. Instead, we set the chunk size to C and produce T/C non-overlapping chunks;! Only the units within the same chunk inter-attend with each other. In practice, this can be achieved by attention mask (0) E RTx at layer l. (i) shares the same shape as the self-attention matrix (see Figure 1) and is defined as:
3.2 Adaptive-Depth and Weight-Sharing
Since our Bliding-Dlltted Attention limits the num- ber of accessible tokens at time, we need an adap- tive depth I. = loge T that the final output can utilize every single piece of input information. However, when Ter > Tir, the depth during infer- ence will be higher than that during training. The simplest way to solve this challenge without further parameter updating is to perform Weight-Sharing across layers. To account for the possible perfor- mance loss due to Weight-Sharing, we first thicken the model by K times, resulting in total number of K L layers. Next, share the weights across the K L layers in the following way for [K]:
(8) mn
Iim w)/Ce, inf,
C² m-n c [C]
(1)
otherwise
Note that M is a lower triangular matrix due to the causal nature of our model. ri's with i [C] learnable relative positional scalars. To be precise, each attention head has different set of Iearnable biases Ty's Here, we drop the dependency on the head for notational simplicity.
(A--++)) -SA (R) for l E [Z]
FFN (FFN()) for E [L]
can be equivalently interpreted as stacking more SA and FFN components within every Transformer layer, and the same thickened layer reused L times. This layer thickening design is only used the natural language modeling experiments in $6.
The use of ry's is similar to the positional scalars of T5 (Rae and Razavi, 2020a) except that we do not use the log-binning strategy over m n. It is to facilitate the extraction of global information instead of enforcing the windowed-attention ef- fect (Raffel al., 2020; Press et al., 2022; Chi
3.3 Where the WM Notion?
Instead of instantiating WM along the temporal dimension as the combination of scratchpad and
Whenevvrr 7' is not divisible hy C, we pad the input se- quence such that its length is a multiple of C.

---

(A' + Av + 1/4 )1* (A + =/ Av, + va)
range of attention in lower layers of a Transformer is also corroborated in Rae and Razavi (2020b). where they find such design does not deteriorate the performance.
1 (AV,- + v2)) + Av,+ + V4
4 (AVs + v6) + Av; Vg
Adaptive-Depth and Weight-Sharing AL- BERT (Lan et al., 2020) and Universal Trans- former (Dehghani et al., 2019) share the parameters across layers. The weight sharing design makes them compatible with the idea of Adaptive Computation Time (Graves et al., 2014) and Dynamic Halting (Dehghani et al., 2019; Elbayad et al., 2020), which allocate different compu- tational budget depending on the complexity of tasks (Simoulin and Crabbé, 2021; Csordás et al., 2022). However, they lack the special Beading-DilatedAAtteee
Av + 1/2
Avg + 1/4
Avs + V/ Au, + Da
P1
V2 F/3
1/4
F
g
Figure 2: This is the parallel scan algorithm that can accelerate a linear RNN. In this example, we visualize the routing path for computing æs. Blocks at the same layer can be computed in parallel on GPUs.
recency biases, RegularGPT limits the amount of information along the depth dimension. As we have seen, the idea of breaking 7 units into several chunks limits the amount of accessible informa- tion at each layer, thereby enabling the WM no- tion. A similar argument was made by Yogatama et al. (2021) in sense that they categorized Long- former (Beltagy et al., 2020), a transformer variant with local attention pattern, as a model of working memory. Finally, thanks to modern accelerators such as GPU, all chunks at a layer can be processed concurrently, and this further makes RegularGPT more favorable over the scratchpad and recency biases approach.
design that is necessary for ruling out naive solutions.
Linear RNN Given xo 0 E R and the in- put vectors 01 ... UT, linear RNN (Orvieto et al., 2023) for & E [T] can be written as:
Jt Ark 1 + Buk
A Bakk
1=0
Σ A Ukー 1=0
where we set j Buk j. The operation can be accelerated by the parallel scan algorithm that per- mits efficient cumulative sum (Ladner and Fischer, 1980; Blelloch, 1990; Lakshmivarahan and Dhall, 1994; Martin and Cundy, 2018; Liu et al., 2023; Smith et al.. 2023). As we can see in Figure 2, the routing path specified by the parallel scan algo- rithm is the same as our Diitte---Dlatt-Attttte illustrated in Figure 1.
3.4 Complexity Analysis
The sparse attention pattern of RegularGPT sug- gests it might enjoy the same speedup provided by sparsified Transformers. The complexity of our model is O(TCK logo T') where TC is the com- plexity of each self-attention module and K logo T is the total number of layers. On the other hand, the vanilla Transformer follows O To il- lustrate the possible speedup, if T = 512 and K log 512 < 5 when K < 128 1og 512 R 3.11L. Namely, long K < 3L, our model is likely to be more efficient than vanilla Transformer.
5 Regular Language Experiments
C = 128, then 512 128
512/.
5.1 Task Descriptions
We focus on the four tasks in section 1) (Deletang et al., 2023) of Table I as they will also be used in our analysis in $7. For tasks in section 2), please refer to Bhattamishra et al. (2020) for details.
4 Connection to Prior Work
This special atten- tion pattern dates back to pre- Transformer era such as Wavenet (van den Oord et al., 2016) with dilated convolution. It can also be viewed as a special form of Longformer attention pattern with system- atic dilation (Beltagy et al., 2020).² 2 Limiting the
Sliding-Dilated-Atteeee
Even Pairs A model needs to predict whether the total count of "ab" and "ba" pairs is even. In the example of "aabba", there is one "ab" and one "ba". resulting in a total count of 2, which even. This task is equivalent to checking whether the first and last characters in a string are identical.
"The original Longformer also adopts dilated attention on a few heads at higher layers hut without the systematic pattern
used in this work.

---

RegularGPT
Task
RNN
Transformer
= 2
C=
1) Deletang et al.
Even Pairs
100.0/ 100.0
99.7 733.2
100.0/
89.3
100.0 19.6
Modular Arithmetic
100.0 100.0
21.9
96.4/
82.6
21.2/ / 20.5
Parity Check
100.0
98.9
52.3
100.0
100.0
100.0 / 88.7
Cycle Navigation
100.0 100.0
21.7 20.6
100.0
100.0
100.0/ 78.6
2) Bhattamishra et al.
D2
100.0 / 100.0
100.0 / 80.1
100.0
100.0
99.8 96.5
D3
100.0 100.0
10..0/77.8
100.0/
99.7
98.6/ 93.0
D4
100.0 / 100.0
82.6
100.0/
98.7
97.7 /9.6
D12
100.0 100.0
/00.0/80.3
100.0/
99.8
94.1 90.4
Tomita 3
100.0 100.0
/94.4
100.0/
99.7
100.0 99.9
Tomita 4
100.0 100.0
70.00
100.0/
99.8
100.0 99.3
Tomita 5
100.0 / 100.0
74.5/74.5
100.0/
99.8
98.2 84.1
Tomita 6
100.0 100.0
/0.0/5001
100.0 98.5 100.0 65.7
Table 1: Length generalization results Regular Languages (Max/Avg). All models in the first section (Dele- tang et al.) are trained on sequences of length 40. 'he reported numbers are the average of length extrapolation results from to 500. Each result is an average over seeds. All models in the second section (Bhattamishra al.) are trained sequences of length 50. The reported numbers are the average of length extrapolation results from 51 to 100. Each result an average over 3 seeds. Please refer to Appendix A for the detailed hyperparameters.
Modular Arithmetic Given sequence of num- bers in {0, 1, 2, 3, 4) and operations in (+,-, model needs compute the result modulo For example, x 1 + 2 evaluates y 4.
reports the maximum parameter settings and random seeds because their goal to know if model can extrapolate all. We additionally report the average scores since we want to know the model can consistently obtain good performance. The baseline models we com- pare against are an RNN and vanilla Transformer with Transformer XL. style relative positional em- bedding (Dai et al., 2019). Table 1 shows that RegularGPT with C 2 acheives similar perfor- mance an RNN and substantially outperforms vanilla Transformer.
across different hyper-
Parity Check A model needs to compute whether the number of bs in given binary string is even. For example, the sequence aaabba contains bs, which is even.
Cycle Navigation Given sequence of move- ments on cycle of length 5, model needs to compute the end position. The possible movements are STAY, INCREASE, DECREASE encoded as (0, 2). The agent always starts at position For example, 010211 means the agent stops at position 2 = 0 1 + 1+
5.3 The Effect of Chunk Size C
We vary the chunk size C of RegularGPT to its impact on the performance. The motivation for using a larger is to reduce the number of layers (i.e., L = logg decreases C) and in- crease the degree of parallelization. However, Table 1, larger C' seems to pose challenge to RegularGPT on the Modular Arithmetic task. Mod- ular Arithmetic is hard task with far more states and complicated state transitions. Increasing C' is likely increase the task difficulty by composing more state transitions once. We will have an in-depth discussion of the theoretical reasons in $7.
5.2 Language Transduction and Extrapolation
First, want to know if endowing Transformer with the notion of WM really improves its length extrapolation capability on regular languages. We test RegularGPT and all the baselines on two sets of regular languages from prior work (Deletang et al., 2023; Bhattamishra et al., 2020).3 3 Prior work often
*Our implementation is based on the codebase of Deletang et al. (2023) at: https: / thub. com/ deepmind neural l_nettorrss swoors__chomee ky_ mirrrchyy We addi- tionally implement the regular languages in the second section
of Table 1.

---

Given that RegularGPT has been battle-tested on the main experiment of regular languages, we now shift gear to benchmark its performance in the nat- ural language scenario. Given a model trained on sequences of length Tir, we test it on much longer sequences of length Tex > Tor during inference, and the goal is to observe similar perplexities. To optimize efficiency, we employ random selec- tion process to extract 1,000 chunks, each with Tex tokens from the testing set. Subsequently, we calculate the average perplexity of the last tokens within these chunks to ensure each of them has Ter - tokens as the context, thereby avoiding the issue of early token curse (Press et al., 2022; Chi et al., 2023). We compare our model against the existing methods that are known to demonstrate the ability of length extrapolation including T5 (Raffel et al., 2020), ALiBi (Press et al., 2022), and KER- PLE (Chi et al., 2022).4 4 To counteract the loss of expressive power due to weight sharing, we thicken each layer of RegularGPT K as detailed in $3.
Probability Plaa
Settings
0.1
0.3
0.5
0.7
0.9
1) Same Length
In Table 3, we first observe exploding perplex- ities for C' = 32 after Tex 2 2048. RegularGPT might only learn to model [log32 512] = 2 lay- ers during training, hence it fails to recursively model more than 32 = 1024 tokens during infer- ence. This is validated by C 64 since this time it is able extrapolate until 6-1 G-googe 512] 4096.
RegularGPT
100
100
100
100
100
RNN
100
100
100
100
100
Transformer
98.4
99.8
99.6
97.8
77.2
2) Extrapolation
RegularGPT
100
100
100
100
100
RNN
100
100
100
100
100
Transformer
50.1
49.7
50.3
49.9
50.0
While the above argument seems to suggest large C, setting C 256 also deteriorates the perfor- mance. This might be due to the limited number of chunks (512/256 - 2) and r,'s i (in Eq. (1)) ob- served at the second layer, making the learning of ry's harder. Overall, C' is hyperparameter that needs to be carefully decided for RegularGPT on natural languages. We also observe that 128/12 per- forms better than 128/6, implying RegularGPT's performance could be improved by stacking more layers to counteract the performance loss due to Weight- Shariigg
Table 2: We alter the probability P(o; = 1) used to sample Is of PARITY. The same length setting is 40. The extrapolation setting is from 41 to 500. Each entry is an average over seeds.
5.4 Robust to Probability Changes
Other than the length extrapolation experiment, we alter the probability of sampling 1s of PARITY, i.c., set P(o;) * 0.5. The results in Table 2 show that RegularGPT is robust to different sampling probabilities. indicating its successful modeling of the underlying regular language grammar. In contrast, a vanilla Transformer model struggles to achieve good performance even for the same length setting, again validating the fact that it only finds the naive-summation solution as discussed in $2.2.
It is worth noting that 128/12 performs relatively well and is close to previous methods designed specifically for the task of natural language extrapo- lation. We will analyze its inner workings in depth in Figure 4 and $7, in which we find that Regu- larGPT learns the similar local receptive field prior work, which is likely the key to its successful natural language extrapolation performance.
6 Natural Language Experiments
Discussion and Analysis
7.1 Regular Language and Finite State Semiautomaton
Regular language is the type of formal language recognized by an FSA (Chomsky, 1956a), which is 5-tuple (Q, Σ, 6, go, F), where Q is a finite non-empty set of states, Σ is finite non-empty set of symbols, qu E Q is an initial state, 8: Q x Σ + Q is a transition function; F C Q is a set of final states. However, some of our tasks are better modeled by a finite-state transducer (FST) as discussed in $2.1. To underpin both FSA and FST, consider semiautomation A = (Q, E, 8) (i.e., an FSA without qo and F) and establish its connection to a Transformer model.
Let aeb be the sequence from position a (in- clusive) to b (exclusive) out of a length input sequence (i.e., 0 < a < b < T). We define A(oas) Q + Q as the (b a)-step state transi- tion relation after receiving Ta:h.
'We use the ma codebase: https: www.c////apatyyymneee
and the Open- WebText2 datasot: https: / inggggngtacc.coo datasets/the, hhtassts/thhe_pll_pppepeeeee pie_o__eeeeeeettettt........
A(aws) = 5(-1ab 1) ...... 8(.1au),

---

RegularGPT
(C/C))
Tex
KERPLE
T5
ALiBi
32/6
64/ 6
128 6
128 / 12
256/6
512
24.71
24.50
24.53
32.06
30.17
28.80
26.37
27.90
1024
24.42
24.38
24.90
32.03
30.30
28.94
26.91
34.38
2048
24.21
25.01
25.08
791.74
30.56
29.14
27.08
34.85
4096
24.53
28.91
25.08
812.00
30.80
29.25
27.28
35.11
8192
24.74
39.08
25.08
818.49
1175.91
29.41
27.39
35.42
Table 3: Natural language extrapolation results on OpenWeeText2. The training length is 512. The numbers are averaged over three random seeds. Please refer to Appendix B for the detailed hyperparameters.
where s()) s(-) (((()) denotes function composition. With abuse of notation, define Ag(Dat) Q as the state after receiving aabb if starting at eegee Q
Now, we can relate Lemma I to the FFN lay- in RegularGPT. Following $3, when chuck size C = and thickness K = 1, the output vector of depends on input sequence 01-22-1+1:++ Also, 0 is computed from o and of (-1) depend on input sequences 0, +1-2+++ and 01-2+11t11 respectively. This observation im- plies that of" likely models the transition func- tion A(0,-2-14410 which we denote as of A(a, ((0--2+1-1:441) We will verify this assumption in $7.3.
1) which
Ag(Das) 8(-100-1) .... ... 8(-1oa).
7.2 Modeling Transition Composition
We want to show that the layers of RegularGPT with chunk size C' = 2 can model the composition of two transition functions:
(1)
If og -(0--2+111+41 is true, Lemma implies that RegularGPT's FFN models the tran- sition function composition. This is immediate by setting o (1-1) Flat(A(o; 21:0:---------1 1)))
A(Da.)) A(ora)) fori e [a+11, ..,..,.,b) b).
This way, the regular language problem can be solved recursively using the construction outlined in $3 and Figure 1. To formalize the statement, we first observe that A(o A(o.s)), and A(o can be represented in R QI):
of (1-1) ~ Flat(AAo, ((1((((1-++1=-1) and recognizing the fact that function composition is matrix prod- uet under the representation of Eq. (2).
The next step is explain the use of self- attention layers in RegularGPT. Although Lemma has established a composition, it is unclear how the transitions are concatenated in the first place (i.c., [Flat( A), Flat(B))). With a two-hcad self- attention and the learnable relative positional scalars, it is possible to adjust them so that the attention output contains the concatenated informa-
OneHotiq (Ago (aab))
OneHot (An (Ta:b)) E RIoP
A(aao)=
OneHotiqi (ogg)
(2)
where OneHotiq is a one-hot vector of length [QI with the i-th index being 1.
The next step is to mix A(Fas) and A(ass) to- gether and get A(a.s)) We show in Lemma I that 2-layer ReLU network can learn (and transformer layer) the composition. The proof of Lemma deferred to Appendix C.
tion [Flat ()), Flat(B)).
Recall in Eq. (I), each head has different set of scalars ry's. 's One concrete construction for concate- nation is setting rn = 0 and the remaining - 0o for the first head; = 0 and the remaining -oo for the second head. In other words, each head is only responsible for capturing one state transition. After the multi head self-attention operation, obtain the concatenation of two state transitions.
Lemma (Approximation for Binary Matrix Prod- uct). Let A, B E {0, 1|ax be binary matrices of dimension n x n. Then, there exists two-layer ReLU network such that
Finally, when the prediction head reads out the answer, the operation is equivalent to a map- ping from A(aorr) E x! to Ag (oor) =
(Flat Flai(()) Fla((A)),
where Flat(X)ss (i-1)a15 = X1,5 for i,j E [n] is the operation that flattens a matrix into a vector.
(I) A(our) o qo E IR Q Since we assume that 0f

---

15
15
10
10
5-
PCA2
PCA2
0-
-15
-10
-5
10
15
-1
-10
-5
0 PCA1
10
15
PCA1
(a) PARITY.
(b) Cycle Navigation.
Figure 3: Clustering of FFN output vectors across all layers via PCA the task of PARITY and Cycle Navigation.
models A the transduction readout is per- formed by linear map oT-1 Wo (1) Our tree-structured construction also guarantees that the final answer could be derived using log2 T layers.
extract the output of of all layers perform dimen- sion reduction using PCA, and plot the dimension- reduced points on a 2D plane. Ideally, we want to see limited number of clusters across all layers, indicating the model learns to capture the state tran- sition function. As has 2 clusters and Cycle Navigation has 5 clus- ters. The clear clustering effect demonstrates Reg- ularGPT's correct learning of state transition func- tions. This is in contrast to the naive-summation ap- proach learned by a vanilla Transformer as shown in Figure B.4 of Deletang et al. (2023).
7.3 Verification of Transition Modeling
see in Figure 3, PARITY
To verify whether our model learns the dynamics semiautomaton, we perform a clustering exper- iment to demystify the FFN output representations on the tasks of PARITY and Cycle Navigation. The two tasks are chosen as we can easily derive their state transition functions. For example, there are only two state transitions in PARITY:
7.4 Receptive Field Analysis
1
0
0 1
1
0
We resort to the gradient analysis tool (Chi et al., 2023) to inspect the receptive field of RegularGPT on regular and natural languages. It computes a cu- mulative sum of the gradient norms starting from the most recent token to the earliest one. A large magnitude of slope cent token has high dependency that position. Idcally, we would like to see the receptive field covering the whole input sequence for the case of regular languages because every single bit in the in- put sequence important for the final results. This is equivalent to a slanted line going from the lower right to the upper left, which is validated in Fig- ure 4a. As for natural language, we discover some- thing interesting in Figure 4b in that RegularGPT settles the local windowed-attention pattern those enforced manually in prior work (Press al., 2022; Chi al., 2022, 2023). This suggests the task natural language modeling mostly needs only local context to achieve good performance,
and five state transitions in Cycle Navigation
OneHots ((0 + k) mod 5)
OneHots (1 + mod 5)
position means the most re-
OneHot;((2 + mod 5)
, for l E [O.....4] 1).
OneHot; + k mod 5)
OneHots + ( mod 5)
0
0
1
e.g., = gives
0
0
1
0
0
0
0
0
0
Given a testing input sequence of length 500 that is much longer than the training length 40, we
Cycle Navigation ($5.1) has one-step transitions (i... A(( ) = 3). Composing these transitions yields 5 different multi-stop state transitions (i.c., [A(da))] = 5 5).

---

1.0
1.0
0.8
0.8
0.6
0.6
0.4
0.2-
0
100
200
300
400
500
500
1000
1500
2000
Position
Position
(a) Regular Language PARITY
(b) Natural Language OpenWebText2
Figure 4: Receptive field of RegularGPT via the cumulative gradient analysis tool (Chi et al., 2023).
which aligns with the common belief.
acknowledges the support from the Boeing Com- pany (2019-STU-PA-259), The fourth author ac- knowledges the support from NSF. (MRI Award: 1919452).
8 Conclusion
This paper introduces RegularGPT, a novel variant of the Transformer architecture inspired by the no- tion working memory that can effectively model regular languages with high efficiency. Theoretical explanations and accompanying clustering visual- izations are presented to illustrate how RegularGPT captures the essence of regular languages. More- over, RegularGPT is evaluated on the task of nat- ural language length extrapolation, revealing its intriguing rediscovery of the local windowed atten- tion effect previously observed in related research. Notably, RegularGPT establishes profound connec- tions with various existing architectures, thereby laying the groundwork for the development of fu- ture Transformer models that facilitate efficient al- gorithmic reasoning and length extrapolation.
References
Eryn Adams. Anh T Nguyen, and Nelson Cowan. 2018. Theories of working memory: Differences in definition, degree of modularity, role of attention, and purpose. Language, speech, and hearing ser- vices in schools, 4903):340-35.
Cem Anil, Yuhuai Wu, Anders Johan Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Venkatesh Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. 2022. Exploring length generalization in large language models. In Advances in Neural Information Processing Sys- tems.
Kristijan Armeni, Christopher Honey, and Tal Linzen. 2022. Characterizing verbatim short-term mem- ory neural language models. In Proceedings of the 26th Conference on Computational Natural Language Learning (CoNLL), pages 405-424, Abu Dhabi, United Arab Emirates (Hybrid). Association for Computational Linguistics.
Limitations
Currently we set the chunk size C of RegularGPT to constant. Can we make the chunk size more flexible? flexible and data-driven C might fur- ther boost its performance on natural languages as they often demonstrate diverse patterns unlike reg- ular languages underpinned by simple grammars. This might also improve the performance of Regu- larGPT when C' 128.
Alan Baddeley. 1992. 255(5044):5565559
Working memory. Science,
Alan D Baddeley and Graham Hitch. 1974. Working memory. In Psychology of learning and motivation, volume 8, pages 47-89. Elsevier.
1z Beltagy. Matthew E. Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. arXiv:2004.05150.
Acknowledgment
We thank the anonymous reviewers for their insight- ful feedback and suggestions. We thank Prince- ton Research Computing for the technical support on the Della and Adroit clusters. The first author
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. 2020. On the Ability and Limitations of Transform- ers to Recognize Formal Languages. In Proceed- ings of the 2020 Conference on Empirical Methods

---

Natural Language Processing (EMNLP), pages 7096-7116, Online. Association for Computational Linguistics.
Gregoire Deletang, Anian Ruoss, Jordi Grau-Moya. Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Ve- ness, and Pedro A Ortega. 2023. Neural networks and the chomsky hierarchy. In International Confer- ence on Learning Representations.
Guy E Blelloch. 1990. Prefix sums and their applica- tions. School of Computer Science, Carnegie Mel- lon University Pitshurgh, PA, USA.
Adele Diamond. 2013. Executive functions. Annual review of psychology, 64: 135-168.
Ta-Chung Chi, Ting-Han Fan, Peter Ramadge, and Alexander Rudnicky. 2022. KERPLE: Kemeliied relative positional embedding for length extrapola- tion. In Advances in Neural Information Processing Systems.
Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. 2020. Depth-adaptive transformer. Interna- tional Conference on Learning Representations.
Ta-Chung Chi, Ting-Han Fan, Alexander I. Rudnicky, and Peter Ramadge. 2023. Dissecting transformer length extrapolation via the lens of receptive field analysis. In Proceedings of the 6lst Annual Meet- ing of the Association for Computational Linguistics (Volume 1: Long Papers), Toronto, Canada. Associ- ation for Computational Linguistics.
Jeffrey L Elman. 1990. Finding structure in time. Cog- nitive science, 14(2)179-211.
K Anders Ericsson and Walter Kintsch. 1995. Long- term working memory. Psychological review, 102(2):211.
Alex Graves, Greg Wayne, and Ivo Danihelka. 2014. Neural turing machines. arXiv preprint arXiv:1410.540I.
David Chiang and Peter Cholak. 2022. Overcoming a theoretical limitation of self-attention. Proceed- ings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Pu- pers). pages 7654-7664, Dublin, Ireland. Associa- tion for Computational Linguistics.
Michael Hahn. 2020. Theoretical limitations of self- attention in neural sequence models. Transactions of the Association for Computational Linguistics. 8:156-171.
N. Chomsky. 1956a. Three models for the description of language. IRE Transactions Information The- ory, 203):13--24.
Sepp Hochreiter and Jürgen Schmidhuber. 1997. Long short-term memory. Neural computation, 1735
Noam Chomsky, 1956b. Three models for the descrip- tion of language. IRE Transactions on information theory, 2(3):1 13-124..
Michael I Jordan. 1997. Serial order: parallel dis- tributed processing approach. In Advances in psy- chology, volume 121, pages 471-495. Elsevier.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. 2019. What does BERT look at? an analysis of BERT's attention. In Pro- ceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pages 276-286, Florence, Italy. Association for Computational Linguistics.
Richard E Ladner and Michael Fischer. 1980. Parallel prefix computation. Journal of the ACM (JACM), 27(4):831-838
Sivaramakrishnan Lakshmivarahan and Sudarshan K Dhall. 1994. Parallel computing using the prefix problem. Oxford University Press.
Nelson Cowan. 1998. Attention and memory: An inte- grated framework. Oxford University Press.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel. Piyush Sharma, and Radu Soricut. 2020. Albert: A lite bert for self-supervised learning of language representations. In International Con- ference on Learning Representations.
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. 2022. The neural data router: Adaptive control flow in transformers improves systematic generalization. International Conference on Learning Represen- lations.
Bingbin Liu, Jordan T. Ash. Surbhi Goel, Akshay Kr- ishnamurthy, and Cyril Zhang. 2023. Transformers learn shortcuts to automata. In International Confer- ence on Learning Representations.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Car- bonell. Quoe Le, and Ruslan Salakhutdinov. 2019. Transformer-XLe Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computa- tional Linguistics, pages 2978-2988, Florence, Italy. Association for Computational Linguistics.
Eric Martin and Chris Cundy. 2018. Parallelizing lin- ear recurrent neural nets over sequence length. In International Conference on Learning Representa- tions.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. 2019, Univer- sal transformers. International Conference on Learning Representations.
Akira Miyake, Priti Shah, et al. 1999. Models of work- ing memory. Cambridge: Cambridge University Press.

---

Aida Nematzadeh, Sebastian Ruder, and Dani Yo- gatama. 2020. On memory human and artificial language processing systems. In Proceedings of ICLR Workshop Bridging Al and Cognitive Sci- ence.
Aäron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alexander Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. 2016. Wavenet: A generative model for raw audio. In Arxiv.
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 Augustus Odena. 2022. Show your work: Scratchpads for intermedi- ate computation with language models.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, F. ukasz Kaiser, and Illia Polosukhin, 2017. Attention is all you need. In Advances in Neural Information Pro- cessing Systems, volume 30. Curran Associates, Inc.
Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoe Le, and Denny Zhou. 2022. Chain of thought prompting elicits reasoning in large language models. In Ad- vances in Neural Information Processing Systems.
Klaus Oberauer. 2002. Access to information in work- ing memory: exploring the focus of attention. Jour- nal of Experimental Psychology: Learning, Memory, and Cognition, 28(3):111
Antonio Orvieto, Samuel L. Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. 2023. Resurrecting recurrent neu- ral networks for long sequences. arXiv preprint arXiv:2303.06349.
Dani Yogatama, Cyprien Masson d' A'Atumme, and Lingpeng Kong. 2021. Adaptive Semiparametric Language Models. Transactions of the Association for Computational Linguistics, 9:362 373.
Ofir Press, Noah Smith, and Mike Lewis. 2022. Train short, test long: Attention with linear biases enables imput length extrapolation. International Confer- ence on Learning Representations.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et 2019, Lan- guage models are unsupervised multitask learners. OpenAl blog. 1(8)9
Jack Rae and Ali Razavi. 2020a. Do transformers need deep long-range memory? In Proceedings of the 58th Annual Meeting of the Association for Compu- tational Linguistics. pages 7524-7529.
Jack Rac and Ali Razavi. 2020b. Do transformers need deep long-range memory? In Proceedings of the 58th Annual Meeting of the Association for Compu- tational Linguistics, pages 7524-7529, Online. As- sociation for Computational Linguistics.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020. Exploring the limits of transfer learning with unified text-to-text trans- former. The Journal of Machine Learning Research, 21(1):5485-555 5551.
Antoine Simoulin and Benoit Crabbé. 2021. How many layers and why? An analysis the model depth in transformers. In Proceedings of the 59th Annual Meeting of the Association for Computa- tional Linguistics and the 11th International Joint Conference on Natural Language Processing: Stu- dent Research Workshop, pages 221-228, Online. Association for Computational Linguistics.
Jimmy T.H. Smith, Andrew Warrington, and Scott Lin- derman. 2023. Simplified state space layers for se- quence modeling. In The Eleventh International Conference on Learning Representations.

---

A Hyperparameters for the Regular Language Experiments
To sum these n products into n² results, the second layer's binary weight matrix W(2) {0, 1] jn'xnn xn² constructed as:
We report the hyperpamaters used in the regular language experiments (Table 1) in Table 4.
1 n
On
On
On On
B Hyperparameters for the Natural Language Experiments
W I,2 & Ln
On
In
On
On
On
1n
We report the hyperpamaters used in the natural language experiments (Table 3) in Table 5.
E (0, 1
C Proof of Lemma
where 1 is an n² n² identity matrix, @ is the Kronecker product, O₂ is an n-dimensional column vector of all zeros, and 1n an n-dimensional column vector of all ones. We arrive at a two- layer ReLU network that perfectly approximates the multiplication of two binary matrices:
Lemma (Approximation for Binary Matrix Prod-
uct). Let A, B {0, 1jnx= be binary matrices of dimension n x n. Then, there exists two-layer ReLU network such that
(Fla( A), (B))) Flat(AB),
fmip [Flat(A), (B))
where Flat(X X) (5-1m++ Xi.j for i.j [..... n) is the operation that flattens a matrix into a vector.
=ReLUU (IFla
Flat(AB)
Proof. Observe that ReLU operation can per- fectly approximate the multiplication of two binary scalars:
D Illustration of Lemma 1
ReLU((a + b 1) = a.ab.. b. for a, E (0,
D.1 Illustration of the Binary Weight
The binary matrix product AB is composed of n binary scalar products of the form:
Matrices
We illustrate W()) and W(2) of Lemma as fol- lows:
Aix Bkj = Iii--) (ntt))++= 1)++j
import numpy as np
for i.j. n1.
def get_Wll WI(n): n2 nen
where x (Fltt A), Flat( B)] is the concatenated flattened input. Our goal is construct two neural network layers. The first layer computes all n² binary scalar products. The second layer sums these products into the form of matrix product; i.e., Ei=1 Au Bxj-
np. zeros ((2*n*n, ns*3), dtype=int) for in range for j in range(n):
for k in range(n):
i =
Wl[n2+ken+j [onn++jntk]] =
return WI
The first layer's binary weight matrix W(1) E {0, 102n")2x" xn² is constructed as:
def vet_W2(n):
eye = np. eye (nan, dtype=int)
For 2 E [l,..., 2n² i,j,k E [I,...,n].
ones = np. ones ((n,)) 1), dtype=int) w2 kron (eye. ones) return W2
(1)
1.(1-1)n+40) 1)-
{1 if (i 1)n + k +j
get W1(2) gives:
otherwise.
[[]
0]
[0
0
0]
[0
0
0
0
01
Then, the first layer computes all n binary scalar products follows:
[0
0
0]
ReI.U (Il.
1. (e- (4-1)244 D-D
[0
0]
[0
0]
= Bks for jari,j, E ([1, [1,.., n).
[0

---

# Layers
loge
Hidden Size
256
Attention Heads
8
Train Seq. Len.
40 or 50
# Trainable Params.
4.3 M
Optimizer Adam (Ir le-4, 3e-4, 5e-4)
Batch Size
128
Train Steps
100,000 float32
Precision
Dataset
Regular Languages
Table 4: Hyperparameters for the regular language experiments.
Layers
K logc T 768
Hidden Size
Attention Heads
12
Train Seg. Len.
512
# Trainable Params. 81M (K 6) 123M (K = 12)
Optimizer
Adam (Ir 6e-4)
Batch Size
32
Train Steps
50,000
Precision
b6loat16
Dataset
OpenWebText
Table 5: Hyperparameters for the natural language experiments.
get_ W2(2) gives:
Then, Lemma verified as follows:
ReLU (>W(1) I (2
[[]
0
0
0]
[1
0
0
0]
ReLU ([1 (11201120) 1) (2
[0
1
0
0]
=001000100
[0
1
0
0]
[]
0
1
0]
-0
01
0
1
01
-Flat (le
1 1 1
Flat (AB).
10
0
0
1]
[0
0
0
111
Here is the Python code for the above example:
up. array ([[1 .0] [[1 .0]]). reshape (-1)
B up. array (ITO .11 .1. .011). reshape (-1)
up. concatenate ([A, reshapee)
D.2
An Illustrative Example for n 2
WI get (2)_
W2 get_W2(2)
Suppose the input matrices are:
flat AB np maximum @ WI -1,0) @ W2
0
= 0 1
0
0
The concatenated flattened input becomes:
z = (A), Flat B)] = [1 011011011000 0j.


Reviews:
[{"review_id": "xH57ZldtI4", "reviewer": "Reviewer_wKf5", "summary": "", "questions": "The authors do not compare RegularGPT to other state-of-the-art models for modeling regular languages and natural language length extrapolation. This would be helpful to see how RegularGPT compares to other approaches.\n\nThe authors' claim that RegularGPT is able to learn long-range dependencies effectively is based on its performance on the PARITY task. However, it is not clear whether RegularGPT would be able to learn long-range dependencies as effectively on more complex tasks.\n\nThe authors' use of weight sharing may make RegularGPT less adaptable to different input sequences. This is because weight sharing can lead to a loss of precision, which may make it difficult for RegularGPT to learn the nuances of different input sequences.\n\nThe authors' use of adaptive depth may make RegularGPT less efficient for long input sequences. This is because adaptive depth requires the model to dynamically adjust the number of layers, which can be computationally expensive.\n\nThe paper does not discuss the scalability of RegularGPT. It is unclear how the model would perform on very long input sequences or on large datasets.\n\nThe paper does not provide a detailed analysis of the training process for RegularGPT. This makes it difficult to understand how the model was trained and how it might be improved.\n", "limitations": "", "rating": 2, "confidence": 2, "soundness": 2, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful feedback and suggestions!\n\n**The authors do not compare RegularGPT to other state-of-the-art models for modeling regular languages and natural language length extrapolation. This would be helpful to see how RegularGPT compares to other approaches.**\n- Table 1 and 3 in the original submission compares RegularGPT to other approaches. We also include additional baseline methods suggested by reviewer cBxx. Please check our response to reviewer cBxx below.\n\n**The authors' claim that RegularGPT is able to learn long-range dependencies effectively is based on its performance on the PARITY task. However, it is not clear whether RegularGPT would be able to learn long-range dependencies as effectively on more complex tasks.**\n- We experimented with the twelve tasks in Table 1, and many of which are substantially more difficult than PARITY. RegularGPT is able to achieve perfect accuracy on them.\n\n**The authors' use of weight sharing may make RegularGPT less adaptable to different input sequences. This is because weight sharing can lead to a loss of precision, which may make it difficult for RegularGPT to learn the nuances of different input sequences.**\n- In the original submission, we thickened the number of layers ($K$ detailed in sec. 3.2 in the paper) for the task of natural language to address this problem. We also reported the effect of different $K$ in sec. 6 and Table 3.\n\n**The authors' use of adaptive depth may make RegularGPT less efficient for long input sequences. This is because adaptive depth requires the model to dynamically adjust the number of layers, which can be computationally expensive.**\n- The number of layers increases only logarithmically w.r.t the input sequence length. Concretely, chunk size of 64 can readily process sequences of length 4096 using only $log_{64} 4096=2$ layers. We consider this acceptable. \n\n**The paper does not discuss the scalability of RegularGPT. It is unclear how the model would perform on very long input sequences or on large datasets.**\n- The size of the natural language dataset used to pretrain the model is about 66GB; we believe that this is large enough to provide meaningful results. In the original submsssion, we also presented comprehensive length extrapolation results of RegularGPT on five different sequence lengths up to 8192 tokens.\n\n**The paper does not provide a detailed analysis of the training process for RegularGPT. This makes it difficult to understand how the model was trained and how it might be improved.**\n- We provide a theoretical analysis of function compositions learned by RegularGPT in sec. 7. As well, we present more detailed training setups and hyperparameters in the appendix.", "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": "The paper proposes a novel approach to modeling regular languages that is based on the notion of working memory. This approach is able to learn long-range dependencies effectively, which is important for modeling regular languages. The paper evaluates RegularGPT on a variety of tasks and they show that RegularGPT outperforms state-of-the-art models on all of these tasks.", "reasons_to_reject": "The paper only evaluates RegularGPT on a small number of tasks. It would be interesting to see how RegularGPT performs on a wider range of tasks. This would help to better understand the strengths and weaknesses of the model.", "questions_to_authors": "The authors do not compare RegularGPT to other state-of-the-art models for modeling regular languages and natural language length extrapolation. This would be helpful to see how RegularGPT compares to other approaches.\n\nThe authors' claim that RegularGPT is able to learn long-range dependencies effectively is based on its performance on the PARITY task. However, it is not clear whether RegularGPT would be able to learn long-range dependencies as effectively on more complex tasks.\n\nThe authors' use of weight sharing may make RegularGPT less adaptable to different input sequences. This is because weight sharing can lead to a loss of precision, which may make it difficult for RegularGPT to learn the nuances of different input sequences.\n\nThe authors' use of adaptive depth may make RegularGPT less efficient for long input sequences. This is because adaptive depth requires the model to dynamically adjust the number of layers, which can be computationally expensive.\n\nThe paper does not discuss the scalability of RegularGPT. It is unclear how the model would perform on very long input sequences or on large datasets.\n\nThe paper does not provide a detailed analysis of the training process for RegularGPT. This makes it difficult to understand how the model was trained and how it might be improved.\n", "paper_topic_and_main_contributions": "In this paper, the authors proposed a novel approach to modeling regular languages that is based on the notion of working memory. This approach is able to learn long-range dependencies effectively, which is important for modeling regular languages. The authors evaluate RegularGPT on a variety of tasks, including regular language modeling, and natural language length extrapolation. They show that RegularGPT outperforms state-of-the-art models on all of these tasks.", "excitement": 2, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OAMB2biX9Z", "reviewer": "Reviewer_fhAH", "summary": "", "questions": "Why the choice to only evaluate perplexity? There are plenty of downstream NLP tasks that this method could be applied to, which could be more interesting than perplexity.\n", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful feedback and suggestions!\n\n**The results on natural language are quite limited and seem worse than previous work on length extrapolation. Perhaps more work is needed on this?**\n- Yes, the reviewer is correct that RegularGPT does not outperform previous state-of-the-art (KERPLE) on length extrapolation. As we discussed in the limitation section, we believe this is due to the chunking technique used. Further work is needed to close the performance gap.\n\n**It would be really interesting to see the performance of a model of this kind on actual language tasks, particularly reasoning, which might have some connection to the regular language tasks that this transformer seems to do well on.**\n- We agree with the reviewer. In fact, as simple as PARITY seems, it already encodes the multiple negation problem of logical deduction. It might also be useful for sentiment analysis tasks where multiple positive/negative sentiments exist in a text snippet. This is great future work.\n\n**The \"rediscovery\" of the \"local window attention\" amounts to saying that language modeling mostly needs local context, which seems like a really unsurprising thing to discover.**\n- We apologize for the confusion. We are trying to highlight the fact that RegularGPT does not use the explicit attention bias used in ALiBi and KERPLE but nevertheless learns to focus on local context. Not using explicit attention bias makes RegularGPT more flexible as shown in Fig 4, where it focuses on global information to solve PARITY; Neither ALiBi nor KERPLE are able to solve it due to their strong explicit attention bias.\n\n**Why the choice to only evaluate perplexity? There are plenty of downstream NLP tasks that this method could be applied to, which could be more interesting than perplexity.**\n- We do not experiment with other downstream tasks since our main focus of this paper is on the tasks of regular language. We additionally include more regular language experiments in our response to reviewer AFjC below.", "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": "* The paper is very well written and wonderful to read. It could inspire researchers to experiment with interesting transformer variations -- something we don’t see enough at NLP conferences.\n* The experimental results in Table 1 and 2 seem quite impressive, given how conceptually simple the proposed transformer modification is.\n", "reasons_to_reject": "My main concern is how light the experiments are on natural language tasks.\n\nMore in detail:\n* The results on natural language are quite limited and seem worse than previous work on length extrapolation. Perhaps more work is needed on this?\n* It would be really interesting to see the performance of a model of this kind on actual language tasks, particularly reasoning, which might have some connection to the regular language tasks that this transformer seems to do well on.\n* The \"rediscovery\" of the \"local window attention\" amounts to saying that language modeling mostly needs local context, which seems like a really unsurprising thing to discover.\n", "questions_to_authors": "Why the choice to only evaluate perplexity? There are plenty of downstream NLP tasks that this method could be applied to, which could be more interesting than perplexity.\n", "paper_topic_and_main_contributions": "The paper proposes RegularGPT, a variant of Transformer with weight-sharding, adaptive depth, sliding-dilated attention, for the purpose of modeling regular languages. RegularGPT can learn regular languages such as “PARITY”. In natural language, RegularGPT rediscovers local windowed attention.\n\nEssentially the parameters are tied between layers, like in ALBERT, and additionally self attention is heavily masked similarly to a progressively dilated convolution, and finally the total number of layers is dynamic, so that the attention can dilate fully over the input length.\n\nThe method seems to achieve very good results on artificial tasks, but the results on natural language might require more work. There are likely very interesting problem to tackle here, in making this efficient and scalable to large NLP datasets.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "0OTz39z5Xn", "reviewer": "Reviewer_cBxx", "summary": "", "questions": "1. Since the sequence length varies along the depth, is $O(TCK\\log_C T)$ an upper bound?\n2. Can we reconstruct the outputs and use masks instead of chunks to limit the information along the depth, i.e.,  the current token is only informed by the previous two? For instance, 1) 1st-depth:  input [0,1,unk1,0,1,unk2],  output:[0,1,pred1st_1,0,1,pred1st_2], 2) 2nd-depth: input:  [pred2nd_1,pred2nd_2], output:[final_pred].", "limitations": "", "rating": 3, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful feedback and suggestions!\n\n**The authors should assemble and implement a combination of existing chunking methods and Weight-Sharing methods for a proper baseline, given: a. The method relies on sequence chunking. It is widely explored in long-sequence modeling. b. Weight-Sharing, Adaptive-Depth, or their combination are also widely explored in sequence modeling.**\n- We followed the reviewer's suggestion and examined additional baseline methods.\n- Chunking Only: We follow the idea from Y. Liu et al. suggested by the reviewer. We set the chunk size to be $C=2,4,6,8$ to mimic the page size in the referenced paper. The best results on the four tasks (Even Pairs, Modular Arithmetic, PARITY, Cycle Navigation) in part 1) are 52.2, 22.1, 51.9, and 20.9.\n- Weight-Sharing Only: This is reminiscent of the ALBERT model, which is a highly regularized Transformer. The best results on the four tasks in part 1) are 89.2, 23.1, 50.8, and 20.1.\n- Weight-Sharing + Adaptive-Depth: Following reviewer's suggestion, we use the code of lazy Transformer released by X. Ai et al. and test it on the four tasks in Table 1 part 1). The best results of lazy Transformer are 94.1, 20.7, 50.8, and 20.2, which are similar to the vanilla Transformer model.\n- Overall, we can see that RegularGPT is the only model that performs well across the tasks of regular language, indicating its flexibility to resemble a FSA when the input is regular language.\n\n\n**Long segments might be split into two chunks, which means RegularGPT is sensitive to chunking methods. We can see an example for PARITY. C=2 performs well and is an optimal chunking method (i.e., static size = 2) because RegularGPT can model the composition of two transition functions. Therefore, I think the author should experiment with the different chunking methods (e.g., static, dynamic, overlapping, etc.) to observe the performance and understand RegularGPT in depth.**\n- We follow reviewer's suggestion and explore the overlapping chunking method. For regular languages, we set the overlap between two neighboring chunks to be 1 token when the chunk size is 2. We set the overlap to be 1 or 2 when the chunk size is 3. We find the results to be similar to the current results reported in Table 1 (part 1).\n- Interestingly, we observe different behavior when we move to the natural language domain. We pick the best-performing model parameters ($C/K=128/12$) and set the overlap between two neighboring chunks to  $128/2=64$. The resulting perplexities are shown in the table below.\n\n\n    | Length | RegularGPT (Original) | RegularGPT (Overlapping) | \n    | -------- | -------- | -------- |\n    |  512 | 26.37 | 24.98 |\n    | 1024 | 26.91 | 27.01 |\n    | 2048 | 27.08 | 28.25 |\n    | 4096 | 27.28 | 28.27 |\n    | 8192 | 27.39 | 28.31 |\n    \n    As we can see, the in-distribution (training=extrapolation length=512) performance improves while the extrapolation performance degrades slightly. We believe this is related to the local attention bias. Our Fig.4 shows that RegularGPT demonstrates strong local attention bias, which is helpful for length extrapolation suggested by prior work. However, it becomes harder for the model to select more recent chunks to focus on when there is an overlap, thereby weakening the local attention bias effect and leading to worse extrapolation performance.\n    \n- Overall, if a user only cares about the extrapolation performance on regular but not natural language, the overlapping chunking method could be used. We thank the reviewer for this suggestion.\n    \n- Finally, we agree with the reviewer that dynamic chunking seems to be an appealing approach. However, the primary goal of this paper is to establish a model that can solve the tasks of regular language. More advanced chunking methods that can further boost the performance on natural language are left as future work.\n\n**Since the sequence length varies along the depth, is $O(TCK\\log_C T)$an upper bound?**\n- Yes. It is computed as $\\frac{T}{C}\\cdot C^2\\cdot K\\cdot \\log_C T$, where $\\frac{T}{C}$ is the number of chunks. $C^2$ is the self-attention complexity of a chunk. $K$ is the thickness. $\\log_C T$ is the adaptive depth.\n\n**Can we reconstruct the outputs and use masks instead of chunks to limit the information along the depth, i.e., the current token is only informed by the previous two? For instance, 1) 1st-depth: input [0,1,unk1,0,1,unk2], output:[0,1,pred1st_1,0,1,pred1st_2], 2) 2nd-depth: input: [pred2nd_1,pred2nd_2], output:[final_pred].**\n- Yes, this is possible. In fact, this is the preliminary version of RegularGPT, for which we identified two potential drawbacks: 1) the input sequence becomes longer due to the additionally inserted tokens, and 2) barring the use of highly customized CUDA kernels, the attention mask method still computes the full attention matrix, and doing so incurs redundant calculation and will lower the training speed.", "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": "1. The theoretical explanation is thorough and interesting.\n2. The method is simple.", "reasons_to_reject": "1. The authors should assemble and implement a combination of existing chunking methods and Weight-Sharing methods for a proper baseline, given:\n \n   a.  The method relies on sequence chunking. It is widely explored in long-sequence modeling.\n\n   b.  Weight-Sharing, Adaptive-Depth, or their combination are also widely explored in sequence modeling.\n\n2. Long segments might be split into two chunks, which means RegularGPT is sensitive to chunking methods. We can see an example for PARITY. C=2 performs well and is an optimal chunking method (i.e., static size = 2) because RegularGPT can model the composition of two transition functions. Therefore, I think the author should experiment with the different chunking methods (e.g., static, dynamic, overlapping, etc.) to observe the performance and understand RegularGPT in depth.\n\nI am happy to discuss these concerns.\n\n\n--------------------------After Rebuttal--------------------\n\nI have no major concerns.", "questions_to_authors": "1. Since the sequence length varies along the depth, is $O(TCK\\log_C T)$ an upper bound?\n2. Can we reconstruct the outputs and use masks instead of chunks to limit the information along the depth, i.e.,  the current token is only informed by the previous two? For instance, 1) 1st-depth:  input [0,1,unk1,0,1,unk2],  output:[0,1,pred1st_1,0,1,pred1st_2], 2) 2nd-depth: input:  [pred2nd_1,pred2nd_2], output:[final_pred].", "paper_topic_and_main_contributions": "This paper studies a limitation of Transformer (the lack of a recurrent mechanism) for regular languages. The authors present RegularGPT, a combination of Sliding-Dilated-Attention, Weight-Sharing, and Adaptive-Depth.   RegularGPT split the input sequence into chunks for Sliding-Dilated-Attention and ties Transformer layers across depth with Adaptive-Depth operation.\n\nExperiments show that RegularGPT can work for regular languages and nature languages. The highlight of this paper is that the authors give a theoretical explanation of how RegularGPT works for regular languages.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {"Missing_References": "Weight-Sharing and Adaptive-Depth in sequence modeling:\n\nX. Ai and B. Fang, “Leveraging Relaxed Equilibrium by Lazy Transition for Sequence Modeling,” in Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics, 2022, vol. 1, pp. 2904–2924.\n\n\nChunking in long-sequence tasks:\n\nY. Liu et al., “Leveraging Locality in Abstractive Text Summarization,” in Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, EMNLP 2022, 2022, pp. 6081–6093."}}, {"review_id": "10bWwpVRVu", "reviewer": "Reviewer_AFjC", "summary": "", "questions": "A. The main experimental paradigm in this paper is to train on short sequences and evaluate on longer sequences. Although this is a common setup, I think it begs the question of why these longer sequences should be considered \"correct.\" After all, the training data could be consistent with any number of rules that extrapolate to longer sequences (not necessarily resembling a FSA), and the held-out sequences are technically out-of-distribution because of the difference in length. So these experiments aren't really testing how well the model learns the training data, but its inductive bias for learning regular languages. So, is the goal of the paper to present a new architecture that is expressive enough to learn regular languages such as PARITY, or one that has an inductive bias that favors interpreting the training data as a regular language? Maybe a bit of both?\n\nB. 133: In what way does this approach fail? Doesn't taking the sum modulo 2 produce the correct result for all lengths?\n\nC. 170: Why does the fact that the number of bits C is fixed during training help?\n\nD. Eq 1: What is $\\ell$? Does $C^l$ indicate exponentiation? If so, what's the reason for doing that?\n\nE. 313: Are you selecting the best run based on validation performance? Or directly on test performance?\n\nF. Table 1: How do training and evaluation work here? For the tasks under 2), are you directly training a classifier and using negative samples during training?\n\nG. Table 1: Why is RegularGPT so bad at modular arithmetic? Isn't Cycle Navigation another modular arithmetic problem? Could this have to do with the fact that the layer depth is only $\\log_C T$?\n\nH. Are there any tasks you consider that are expected to require more than $\\log_C T$ layers?\n\nI. 358: Are you really only using the perplexity of the last token? Meaning you're using $p(w_n \\mid w_{1:n})$ to represent $\\exp(-\\log p(w) / T)$?\n\nJ. Table 3: Seeing as you've thickened the layers of RegularGPT, how do these model compare in terms of number of parameters?\n\nK. Sec 7.3: This shows that the model learns the expected number of states. What about verifying the transitions?\n\nL. Fig 3: How did you determine the labels for these points? Is this something you added by hand?\n\nM. Fig 4: What are the units on the y axis?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful feedback and suggestions!\n\n**Comment 1a: The set of regular language tasks seems a bit limited. Something else you could try is to randomly generate regular languages and show performance as a function of language difficulty, as in Valvoda et al. (2022).**\n \n- We follow reviewer's suggestion and test RegularGPT on randomly generated regular languages. Valvoda et al. (2022) experiments only with the in-distribution setting (training = testing length = 50 max); we report the extrapolation results where we double the testing length from 50 to 100. Please find the results in the following table. We measure the averaged per-token accuracy of the output tokens. Every number is the average of 10 randomly sampled FSTs. \n    \n    | # of states | Transformer | RNN | RegularGPT |\n    | -------- | -------- | -------- | -------- |\n    | 10 | **100.0** | **100.0** | **100.0** |\n    | 20 | **99.1** | 98.9 | 98.0 |\n    | 30 | 74.9 | 72.3 | **86.4** |\n    | 40 | **98.3** | 94.0 | 97.7 |\n    | 50 | 97.8 | 96.0 | **98.1** |\n    | 60 | **98.0** | 88.5 | 96.6 |\n    | 70 | **87.4** | 51.3 | 83.6 |\n    | 80 | **97.4** | 87.8 | **97.4** |\n    | 90 | 74.1 | 38.0 | **81.8** |\n    | 100| 86.0 | 73.1 | **94.6** |\n    \n    We can see that RegularGPT and Transformer perform similarly when the number of states is limited. However, RegularGPT performs particularly well when the number of states=90 and 100 possibly due to its better regular language modeling ability.\n\n**Comment 1b: Sec 7.1: Moreover, a weak point of Deletang et al. (2023) is that the relationship between their tasks (which are functions, i.e. sets of pairs of strings) and the Chomsky hierarchy (which contains languages, i.e. sets of strings) is not clearly defined. For the four tasks that you showed in Table 1, it's actually not necessary to express them as transduction tasks; concatenating the input and output still results in a regular language, so they can be reformulated as pure language recognition tasks.**\n- We agree with the reviewer that the four languages presented in Deletang et al. (2023) can be equivalently represented as language recognition tasks. We therefore experimented with the concatenation + negative sampling approach and find the same results as the current numbers in Table 1 (part 1). \n\n**Comment 2: I think the main text of the paper also needs to provide more details about the tasks, training, and evaluation used in Table 1.**\n- Due to limited space, we choose to only provide references to Deletang et al. (2023) and Bhattamishra et al. (2020). The training details and hyperparameters are reported in the appendix section (Table 4 and 5). The evaluation details are discussed more in the responses (Question F) below. If the paper is accepted, we will follow reviewer's suggestion and use the additional page to include more details.\n\n**Comment 3: If I understand correctly, the natural language modeling experiments in Table 3 only measure the perplexity of the last token. This is a poor evaluation setup, because it provides no information how well the model has learned all but the last token. I couldn't find anything in the cited papers that suggests doing this. Why not use the perplexity of the full sequence? Moreover, the results given in Table 3 are not a strength of the paper.**\n- The reviewer understands our setup correctly. However, we disagree that this is a poor evaluation setup. This is the sliding window inference setup used in Press et al. (2022) Appendix B.1 with S=1 (the last token). Since then, follow up length extrapolation research including Chi et al. (2023) and Chen et al. (2023) both adopt the same evaluation strategy. Measuring the perplexity of the last token allows us to accurately adjust the amount of previous tokens. If we measure the averaged perplexity of all tokens, the extrapolation results might be too optimistic due to the early token curse effect detailed in Appendix B of Press et al. (2022). Nevertheless, we include the averaged perplexity over all tokens in the table below:\n\n    | Length | KERPLE | T5 | ALiBi | RegularGPT |\n    | -------- | -------- | -------- | -------- | -------- |\n    |  512 | 26.87 | 26.64 | 26.84 | 28.54 |\n    | 1024 | 25.22 | 25.05 | 25.36 | 27.43 |\n    | 2048 | 24.57 | 24.67 | 25.16 | 27.19 |\n    | 4096 | 24.14 | 25.37 | 24.90 | 27.25 |\n    | 8192 | 23.87 | 28.56 | 24.62 | 27.25 |\n    \n    As we can see, the trend is the same as Table 3 of the original submission.\n- We acknowledge that this is not a strength of RegularGPT. We include the results since we want to provide readers a more comprehensive understanding of the pros and cons of our model. As we discussed in the limitation section: The chunking method limits the information exchange between tokens, possibly hindering the model performance on natural language. We also want to emphasize that the goal of RegularGPT is to model regular languages, on which KERPLE, T5, and ALiBi all fail due to their strong local attention bias.\n\n    Press et al. (2022): Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation.\n    \n    Chi et al. (2023): Dissecting Transformer Length Extrapolation via The Lens of Receptive Field Analysis.\n    \n    Chen et al. (2023): Extending Context Window of Large Language Models via Positional Interpolation.\n    \n**Comment 4: Sec 7.3 provides an analysis of the learned FSA states, but an analysis of the learned transition function is conspicuously absent.**\n- We apologize for the possible confusion. However, Figure 3 in sec. 7.3 verifies the number of learned transition functions. The analysis is on transition functions, not states. We explained this in sec. 7.3: PARITY has two transition functions while Cycle Navigation has five. Fig. 3 shows that PARITY and Cycle have two and five clusters, respectively. Thus, the number of transition functions is verified. These two tasks are chosen so that we can easily derive the number of state transition functions, with the state transition functions detailed in sec. 7.3. We also perform the same clustering analysis on more complicated tasks such as Modular Arithmetic and observe the same clustering effect with many more clusters.\n\n**Comment 5: The decision to set the number of layers to $\\log_C T$ seems naive and restrictive. Maybe you could characterize the class of problems this is expected to work well on. Does Sec 7 show that it works for arbitrary regular languages? If so, it would be a good idea to state this explicitly. And if so, the poor performance on Modular Arithmetic cannot be explained by not having enough depth.**\n- $T$ is a variable representing the input sequence length. In our regular language experiments; it is $T_{tr}=40$ during training and $T_{ex}=500$ during testing . $\\log_C T$ is chosen carefully to ensure the coverage of all input tokens, which is necessary for the regular language tasks.\n- Our chunking and merging strategy makes the function composition pattern a tree of height $\\log_C T$. This fact along with sec. 7 show that $\\log_C T$ works for arbitrary regular languages. We will state this explicitly in the next revision.\n- Modular arithmetic is substantially more difficult than cycle navigation. Given a sequence of numbers in {0, 1, 2, 3, 4} and operations in {+, −, x}, it aims to compute the result modulo 5. It has far more states than that of cycle navigation. For example, consider the input 1 + 2 x 3 - 4, the model cannot naively add 1 with 2 because 2 will be later multiplied with 3, which has higher precedence than addition.\n\n**Question A1 - After all, the training data could be consistent with any number of rules that extrapolate to longer sequences (not necessarily resembling a FSA), and the held-out sequences are technically out-of-distribution because of the difference in length.**\n- We acknowledge that there are many generative processes (any number of rules...) that are consistent on the same short sequences but generate different out-of-distribution longer sequences. We put special focus on FSA since it is one of the most common and useful generative processes.\n\n**Question A2 - So these experiments aren't really testing how well the model learns the training data, but its inductive bias for learning regular languages.**\n- To know how well a model learns the in-distribution training data, we need to check the performance when the testing sequences are as long as the training ones i.e. no extrapolation. In this case, Transformer, RNN, and RegularGPT all achieve perfect accuracy on the tasks of regular language presented in Table 1. As for natural language, the length=512 row in Table 3 shows that RegularGPT can achieve close to state-of-the-art performance. Overall, all methods examined in the paper can learn the in-distribution training data well. It is just RegularGPT which enjoys the additional advantage of out-of-distribution length generalization, demonstrated by the length extrapolation experiments on both natural and regular languages included in the paper.\n\n**Question A3: So, is the goal of the paper to present a new architecture that is expressive enough to learn regular languages such as PARITY, or one that has an inductive bias that favors interpreting the training data as a regular language? Maybe a bit of both?**\n- RegularGPT does not interpret the input data as regular language. We argue that RegularGPT is flexible enough to model different training data: When RegularGPT is fed with regular language, it has the ability to discover the correct state transition compositions that resemble an FSA/FST. On the other hand, it behaves similarly to a length extrapolatable Transformer when the input is natural language. This is verified in Fig. 4, when RegularGPT is fed with PARITY, it evenly puts gradient importance across positions; When it is fed with OpenWebText2 data, it exhibits a strong local attention bias.\n\n\n**Question B - Why does the shortcut solution fail?**\n- As a concrete example, let us consider the maximum summation of PARITY when the sequence length=40, which is 40. However, during the extrapolation stage where a sequence could be of length=500, the maximum summation becomes 500. A model (e.g., Transformer) that learns a naive summation strategy does not know how to process the substantially larger summation during the extrapolation stage, leading to a lower task accuracy.\n\n- This phenomenon was also discussed in Liu et al. (2023) sec. 4, where they theoretically identified such shortcut solutions in Transformer and verified it from the attention pattern. In contrast, a correct FSA/FST or transformer (e.g., regularGPT) that learns the correct function compositions will not produce unseen summation when the length extrapolates since it performs the correct state transition right after it observes a new input.\n\n  Liu et al. (2023): Transformers Learn Shortcuts to Automata\n\n**Question C - Why does the fact that the number of bits C is fixed during training help?**\n- Because the model only needs to: 1) learn the state transition function that takes $C$ bits as input, and 2) learn to compose a fixed amount of $C$ state transition functions of the same size.\n\n**Question D - The purpose of $C^l$**\n- Yes, $l$ is the layer index, and $C^l$ indicates exponentiation. It creates the self-attention matrix mask (the lightly shaded blue cells) illustrated in Fig. 1 and its caption. When $C=2$, it computes $2^0=1$, $2^1=2$, and $2^2=4$, which are the strides of diagonal lines when $l=0,1,2$. This attention pattern computes the overall representation as explained in the caption.\n\n**Question E - How are the best runs selected?**\n- We select the best runs based on the validation set performance.\n\n**Question F - How do training and evaluation work? For the tasks under 2), are you directly training a classifier and using negative samples during training?**\n- For part 1), we follow Deletang et al. (2023) to train the models using sequences of maximum length=40. Take PARITY for example, we sample binary sequences of length ranging from 1 to 40 and train the model to predict if there is even or odd number of 1's. During the evaluation/extrapolation stage, the model will be fed with binary sequences with length ranging from 41 to 500, and it needs to again predict if there is even or odd number of 1's. For part 2), the reviewer is correct in that we use negative samples and train the model to perform classification.\n\n**Question G - Why is RegularGPT so bad at modular arithmetic? Isn't Cycle Navigation another modular arithmetic problem? Could this have to do with the fact that the layer depth is only $\\log_C T$ ?**\n- RegularGPT actually performs much better than the Transformer baseline when $C=2$ (96.4 vs 21.9). When $C=3$, it needs to model the state transition function of three inputs at a time, which is more complex. Also as we discussed in the response to Comment 5 above, operator precendence makes modular arithmetic much more difficult than a simple 5-state cycle navigation task.\n\n**Question H - Are there any tasks you consider that are expected to require more than $\\log_C T$ layers?**\n- As we discussed above in the response to Comment 5, there is no need for more than $\\log_C T$ layers as $\\log_C T$ will change w.r.t the input sequence length $T$. The theoretical analysis in sec. 7 also proves that this is sufficient to model regular languages.\n\n**Question I - Are you really only using the perplexity of the last token?**\n- Answered above in our response to Comment 3.\n\n**Question J - Seeing as you've thickened the layers of RegularGPT, how do these model compare in terms of number of parameters?**\n- When the thickness is 6, which is the best-performing RegularGPT variant, it has the same amount of learnable parameters (123.69M) as the full Transformer baseline model. This is because we reuse the Transformer layers across depth (weight-sharing).\n\n**Question K - This shows that the model learns the expected number of states. What about verifying the transitions?**\n- Answered above in our response to Comment 4.\n\n**Question L - How did you determine the labels for these points? Is this something you added by hand?**\n- We apologize for the confusion. They are dummy labels assigned to highlight the clusters with different colors. We will remove or explain them in the caption if the reviewer thinks they are misleading.\n\n**Question M -  What are the units on the y axis?**\n- The y axis represents the cumulative normalized gradient. When a model is predicting the next output, we derive the gradient norm for each of the previous input tokens. A larger gradient norm means a token affects the prediction more, so we can interpret an input token's gradient norm as its weight/importance. We then normalize all the gradient norms so that their sum is capped at 1.0. Please refer to sec. 5 of Chi et al. (2023) for more details.\n- Similar to a cumulative distribution function (no unit), the cumulation starts from the right most position and sums up the normalized gradient all the way to the first input token. A higher slope magnitude means the model puts more weight on a token. Take Figure 4(a) for example, when RegularGPT is predicting the answer of PARITY, each input token has roughly the same weight since the cumulative sum of them is a relatively staight line going from the lower right to the upper left. This is expected since we need every single input token to determine the solution of PARITY.\n\n    Chi et al. (2023): Dissecting Transformer Length Extrapolation via The Lens of Receptive Field Analysis.\n    \n**Typos Grammar Style And Presentation Improvements:**\n- We thank the reviewer for catching these errors. We will fix them in the next revision.", "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": "The essential idea behind the model is interesting: enable generalization on regular languages by having the transformer simulate $O(\\log n)$ steps of divide-and-conquer computation. Sec 7 does a pretty good job of justifying why the authors expect this to help for regular languages.\n\nTable 1 shows that RegularGPT is clearly better than the baseline transformer on a set of regular languages, remedying a problem identified by Deletang et al. (2023).\n\nFig 3 is an interesting result.", "reasons_to_reject": "My complaints are mainly about the experimental methodology, rather than the proposed method.\n\n~~The set of regular language tasks seems a bit limited. Something else you could try is to randomly generate regular languages and show performance as a function of language difficulty, as in [Valvoda et al. (2022)](https://aclanthology.org/2022.coling-1.525/). Sec 7.1: Moreover, a weak point of Deletang et al. (2023) is that the relationship between their tasks (which are functions, i.e. sets of pairs of strings) and the Chomsky hierarchy (which contains languages, i.e. sets of strings) is not clearly defined. For the four tasks that you showed in Table 1, it's actually not necessary to express them as transduction tasks; concatenating the input and output still results in a regular language, so they can be reformulated as pure language recognition tasks.~~\n\n**Edit:** Addressed in the rebuttal.\n\nI think the main text of the paper also needs to provide more details about the tasks, training, and evaluation used in Table 1.\n\n~~If I understand correctly, the natural language modeling experiments in Table 3 only measure the perplexity of the last token. This is a poor evaluation setup, because it provides no information how well the model has learned all but the last token. I couldn't find anything in the cited papers that suggests doing this. Why not use the perplexity of the full sequence? Moreover, the results given in Table 3 are not a strength of the paper.~~\n\n**Edit:** Addressed in the rebuttal.\n\n~~Sec 7.3 provides an analysis of the learned FSA states, but an analysis of the learned transition function is conspicuously absent.~~\n\n**Edit:** Addressed in the rebuttal.\n\n~~The decision to set the number of layers to $\\log_C T$ seems naive and restrictive. Maybe you could characterize the class of problems this is expected to work well on. Does Sec 7 show that it works for arbitrary regular languages? If so, it would be a good idea to state this explicitly. And if so, the poor performance on Modular Arithmetic cannot be explained by not having enough depth.~~\n\n**Edit:** Addressed in the rebuttal.\n\nDespite my criticisms, I think this is a promising method that merits further study, and addressing these problems would go a long way to making the paper better.", "questions_to_authors": "A. The main experimental paradigm in this paper is to train on short sequences and evaluate on longer sequences. Although this is a common setup, I think it begs the question of why these longer sequences should be considered \"correct.\" After all, the training data could be consistent with any number of rules that extrapolate to longer sequences (not necessarily resembling a FSA), and the held-out sequences are technically out-of-distribution because of the difference in length. So these experiments aren't really testing how well the model learns the training data, but its inductive bias for learning regular languages. So, is the goal of the paper to present a new architecture that is expressive enough to learn regular languages such as PARITY, or one that has an inductive bias that favors interpreting the training data as a regular language? Maybe a bit of both?\n\nB. 133: In what way does this approach fail? Doesn't taking the sum modulo 2 produce the correct result for all lengths?\n\nC. 170: Why does the fact that the number of bits C is fixed during training help?\n\nD. Eq 1: What is $\\ell$? Does $C^l$ indicate exponentiation? If so, what's the reason for doing that?\n\nE. 313: Are you selecting the best run based on validation performance? Or directly on test performance?\n\nF. Table 1: How do training and evaluation work here? For the tasks under 2), are you directly training a classifier and using negative samples during training?\n\nG. Table 1: Why is RegularGPT so bad at modular arithmetic? Isn't Cycle Navigation another modular arithmetic problem? Could this have to do with the fact that the layer depth is only $\\log_C T$?\n\nH. Are there any tasks you consider that are expected to require more than $\\log_C T$ layers?\n\nI. 358: Are you really only using the perplexity of the last token? Meaning you're using $p(w_n \\mid w_{1:n})$ to represent $\\exp(-\\log p(w) / T)$?\n\nJ. Table 3: Seeing as you've thickened the layers of RegularGPT, how do these model compare in terms of number of parameters?\n\nK. Sec 7.3: This shows that the model learns the expected number of states. What about verifying the transitions?\n\nL. Fig 3: How did you determine the labels for these points? Is this something you added by hand?\n\nM. Fig 4: What are the units on the y axis?", "paper_topic_and_main_contributions": "This paper proposes a modified transformer architecture designed to generalize well on regular languages such as PARITY (detecting whether a binary string has an odd number of 1s), as recent work has shown that transformers struggle on such simple tasks. The authors do this by using weight sharing (like Universal Transformers), adaptively scaling the layer depth by $O(\\log n)$, and masking attention so that it only attends to non-overlapping fixed-size chunks of the input. The authors show how this can recognize regular languages, such as PARITY, using a divide-and-conquer approach that works in $O(\\log n)$ steps. The resulting model is called RegularGPT. The experiments are particularly concerned with how well the model generalizes to longer lengths.\n\nThe main experimental results are:\n1. On synthetic regular language experiments from Deletang et al. (2023) and Bhattamishra et al. (2020), RegularGPT outperforms a baseline transformer in every case, but fails to outperform an RNN on Modular Arithmetic.\n2. When trained as a language model on natural language and tested on longer sequences, RegularGPT has mixed results compared to three baseline transformer architectures. Oddly, these experiments only consider the perplexity of the last token.\n3. They show that the hidden states of the transformer cluster into what are presumably FSA states in Fig 3.\n4. They provide some gradient analysis to see how far back the transformer looks.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "A. 059: enable -> improve?\n\nB. 061: attempt -> attempts\n\nC. 078: What is $\\log T_{tr,ex}$?\n\nD. 080: performances -> performance\n\nE. Table 1: The main text of the paper should explain these tasks in more detail.\n\nF. 412: q is not a function, so it can't be composed -- the notation needs to be cleaned up a little here."}}]

Author rebuttal:

Meta-review: Pros:
- Several of the results are strong.
- The method is simple and intuitive, a divide-and-conquer algorithm integrated into a Transformer.
- The approach is well motivated for the specific problem addressed, that of handling regular languages.
- The discussion is in depth and enjoyable to read.

Cons:
- The motivation of handling regular languages is somewhat niche.
- It is not clear why one would use a Transformer variant rather than an RNN or other existing model with a stronger inductive bias for regular languages.
Decision: Accept-Findings