Recursive Transformer: Boosting Reasoning Ability with State Stack
Authors: Kechi Zhang
Ge Li
Jia Li
Huangzhao Zhang
Yihong Dong
Jia Li
Jingjing Xu
Zhi Jin
Venue: NeurIPS 2025 poster ()
Keywords: Foundation Model
Primary area: deep_learning

Abstract: The Transformer architecture has emerged as a landmark advancement within the broad field of artificial intelligence, effectively catalyzing the advent of large language models (LLMs).
However, despite its remarkable capabilities and the substantial progress it has facilitated, the Transformer architecture still has some limitations.
One such intrinsic limitation is its inability to effectively recognize regular expressions or deterministic context-free grammars.
Standard Transformers lack an explicit mechanism for recursion and structured state transitions, which can hinder systematic generalization on nested and hierarchical patterns.
Drawing inspiration from pushdown automata, which efficiently resolve deterministic context-free grammars using stacks, we equip layers with a differentiable stack and propose StackTrans with recursion to address the aforementioned issue within LLMs.
Unlike previous approaches that modify the attention computation, StackTrans explicitly incorporates hidden state stacks between Transformer layers. This design maintains compatibility with existing frameworks like flash-attention.  
Specifically, our design features stack operations -- such as pushing and popping hidden states -- that are differentiable and can be learned in an end-to-end manner.
Our comprehensive evaluation spans benchmarks for both Chomsky hierarchy and large-scale natural languages. 
Across these diverse tasks, StackTrans consistently outperforms standard Transformer models and other baselines. 
We have successfully scaled StackTrans up from 360M to 7B parameters. In particular, our from-scratch pretrained model StackTrans-360M outperforms several larger open-source LLMs with 2–3x more parameters, showcasing its superior efficiency and reasoning capability.

Recursive Transformer: Boosting Reasoning Ability with State Stack
Kechi Zhang 1. Ge Li Li Huangzhao Zhang, Yihong Dongl. Jia Li Jingjing Xu Zhi 12,3-
Key Lab of High Confidence Software Technology (PKU), Ministry of Education
School of Computer Science, Peking University, China SScoool of Computer Science, Wuhan University, China *College of AI, Tsinghua University *ByteDance
{zhangkechi, ,iige, zhi phijjin jin edu. cn, zhijin @ whu.edu.cn
Abstract
The Transformer architecture has emerged landmark advancement within the broad field of artificial intelligence, effectively catalyzing the advent of large language models (LLMs). However, despite its remarkable capabilities and the substantial progress
facilitated, the Transformer architecture still has some limitations. One such intrinsic limitation is its inability to effectively recognize regular expressions or deterministic context-free grammars. Standard Transformers lack an explicit mechanism for recursion and structured state transitions, which can hinder systematic generalization on nested and hierarchical patterns. Draw- ing inspiration from pushdown automata, which efficiently resolve deterministic context-free grammars using stacks, we equip layers with a differentiable stack and propose STACKTRANS with recursion to address the aforementioned issue within LLMs. Unlike previous approaches that modify the attention computation, STACK- TRANS explicitly incorporates hidden state stacks between Transformer layers. This design maintains compatibility with existing frameworks like flash-attention. Specifically, our design features stack operations such as pushing and popping hidden states - that are differentiable and can be learned in an end-to-end manner. Our comprehensive evaluation spans benchmarks for both Chomsky hierarchy and large-scale natural languages. Across these diverse tasks, STACKTRANS consis- tently outperforms standard Transformer models and other baselines. We have successfully scaled STACKTRAM up from 360M to 7B parameters. In particular, our crrmm-ccrach pretrained model STACKTRANS --0M--60 outperforms several larger open-source LLMs with 2-3x more parameters, showcasing its superior efficiency and reasoning capability.
Introduction
In the era of large language models (LLMss GPT-4 2023 the Transformer architecture has emerged as the nearly universal backbone, achieving remarkable success across various domains and beyond [Bi et al. 22224] Bai et al. 2023 Zhang et al.) 2024a However, recent empirical studies [Delétang ctt al. | 2020 have demonstrated that Transformers struggle with tasks involving Chom- sky hierarchy Chomsky 1956), such as regular expressions (REs) and deterministic context-free grammars (DCFs) with hierarchical/recursive structure. formal language theory, deterministic context-free languages a proper subset of context-free free languages. For example, in RE match-
"Ge Li and Zhi Jin are the corresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

52
Transtormer Layer 1+1
FaeePee
Porty
48-
obed
Consoirt
300M
State
Mulb-bead Orfenentaaree Stock
SmoliM3-360M
Queen2 5-0,50
Poweraadrr
SmollM 300M
Models
Vorofornor
OLMo
TeeyLLMMA
LSTM
Trareformer Lager
Pretraining Tokens
(a) Architecture
(b) Formal language evaluation
(c) Natural language evaluation
Figure 1: STACKTRANS architecture and its overall performance (best view in color). (a) The hidden state stack can be integrated between Transformer layers to model the Chomsky hierarchy. (b) Radar chart for formal language tasks, where task difficulty indicated by color (from red to blue to green, with green representing the most difficult tasks). (c) Results on natural language tasks, where the r-axis represents the pretraining corpus size, the y-axis represents the average performance across multiple benchmarks (see Table 2p, and the dot size indicates the model size.
ing task, each training example consists of string within specified length range, paired with a matching-or-not label. While Transformers can perform well within the length boundaries of the training set, they often fail to maintain consistent performance when evaluated on input strings that are longer or shorter than those in the training data. One theoretical explanation is that standard Transformers lack suitable inductive biases Tmases-most most notably, an explicit mechanism for recursion and reusable structured state transitions [Mitchel]] 1980; Battaglia et al. 2018 Sartran et al.) [022] Without inherent assumptions to guide learning, Transformers struggle to effectively capture the underlying grammar in the training set. Consequently, they cannot generalize beyond the training data and perform poorly on inputs of lengths that differ from those seen during training. On the other hand, natural languages are generally considered to belong to class of languages that extends beyond DCF [Gazdar and Pullum 1982] Chomsky 1956 [hieber] 1985). This inherent limitation of Transformers may possibly hinder their application in real-world natural language modeling. Such deficiencies may also have the risk to block LLMs to achieve more advanced forms of intelligence.
In the realm of formal language theory and computational linguistics [Chomsky 1956 Savage 1998] Sipser! 1996]. it is widely recognized that automata augmented with stacks correspond to different lev- els of the Chomsky hierarchy grammars Given the success of stack-equipped automata in handling rather complex grammars, it is a natural progression to incorporate the data structure of stack into the Transformer architecture. Recently, researchers have proposed the stacked attention mechanism DuSell and Chiang 2024 Li al. 20244 and have examined its viability upon multiple relatively small benchmarks, but these approaches primarily restructure attention rather than introducing an explicit, reusable recursive state mechanism.
Drawing inspiration from pushdown automata and prior research, we introduce hidden state stacks into the Transformer architecture, proposing a novel method, STACKTRANS. We use the term "stack" because the update mechanism of our module is fundamentally based on differentiable push and pop operations. The proposed mechanism improves hicrarchical generalization and recursive reasoning in Transformers. Unlike the stacked attention mechanisms [DuSell and Chiang 2024
Li et al. 2024], which replace the standard attention, STACK TRANS incorporates differentiable stacks between the Transformer layers, meanwhile preserving the integrity of the Transformer layers (Figure ((a)) This integration allows us to embed the assumptions of the Chomsky hierarchy into the model, enabling STACKTRANS to inherently model and learn REs and DCFs. Moreover, this design maintains compatibility with existing frameworks like flash-attention [Dao et al. 2022 Dao 2024), enabling seamless integration into efficient L.MM training pipelines. Specifically, the stack stores hidden states generated by the preceding layer and updates through operations such as push
'The Chomsky grammars of type 3 (REs) and type 2 (DCFs) be recognized by pushdown automata (automata equipped with stack) [Chomsky 1956 Chomsky Schüizenberger 1959]. With proper enhance- ment, pushdown automata can even resolve type (context sensitive, CS) and type (unrestricted) grammars - 2- stack automata (automata equipped with two independent stacks) are equivalent to Turing machines, and can recognize type 0 grammars Yau 1969].

---

and pop at each decoding time step. To enable end-to-end training of STACKTRANS, we design soft stack operations, thereby making the hidden state stack differentiable. STACKTRANS also adopts a multi-head stack to improve its representation capability. Additionally, we find that the standard stack reading operation (which only returns the top element in the stack) may cause unstable training. Therefore, we propose the global reading operation through learnable query-over-stack attention, stabilizing the training process and enriching the expressiveness of STACKTRANS.
We conduct comprehensive experiments on multiple benchmarks spanning both formal languages (Delétang al., 2022] and natural languages [Groeneveld et al., 2024). Specifically, in RE and DCF tasks, TTACKTRANS outperforms the standard Transformer by at least 30%, achieving nearly 100% test accuracy on all regular language tasks and demonstrated excellent performance on the majority of deterministic context-free grammar tasks, as shown in Figure [[b] Notably, during natural language evaluation, STACKTRANS also demonstrates substantial improvements on tasks such as common sense reasoning and question answering. Furthermore, we have successfully scaled "TACKTRANS up from 360M to parameters. In particular, our open-sourced STACKTRANS TRANS3330,,,,, which is pretrained on corpus of approximately IT tokens, performs better than or comparably to state-of-the-art LLMs with 2-3x more parameters, as shown in Figure [(c)
The contribution of this paper can be summarized as below:
We introduce STACKTRANS, which incorporates hidden state stacks between Transformer layers, providing explicit mechanism to model recursive structure. This integration enables STACK- TRANS to inherently learn grammars from the Chomsky hierarchy, including REs and DCFs.
We design the hidden state stack to be differentiable, which employs soft push, pop, and no-op operations, a multi-head stack mechanism, and global stack reading operation. These innovations ensure stable and end-to-end training of STACK TRANS.
3
We conduct extensive experiments on multiple formal language benchmarks, demonstrating STACKTRANS's effectiveness in learning Chomsky hierarchy grammars such as REs and DCFs. Furthermore, we have successfully scaled STACKTRANS up from 360M to 7B parameters for general language modeling. Evaluations on large-scale natural language benchmarks show that STACKTRANS-360N outperforms baselines with even 2 - 3x more parameters.
2 Background
The foundational success of LLMs can be attributed to the development of the Transformer archi- tecture Vaaswani et al. 2017 and its numerous variations, which serve as the backbone for LLMs [Radford et al., 2018 GPT-4 2023 GPT-4o 2024). Despite their widespread adoption, the Trans- former architecture has inherent expressivity limitations (Hahn] 2020). J. Recent studies have shown that standard Transformers struggle with recursive and hierarchical patterns across both synthetic and real-world tasks (Joulin and Mikolov] 2015 Grefenstette et al. 2015 Sartran et al.) 2022], and need to equip neural networks with external data structures. These limitations pose a potential risk to the natural language modeling capabilities of Transformers, as suggested by discussions surrounding the classification of natural languages within the Chomsky hierarchy [Gazdar and Pullum 1982, Chomsky 1956 Shieber 1985]. For more detailed related work, please refer to $A Based on the same principle, STACK'T RANS introduces differentiable hidden state stacks in modular and scalable manner. Without altering the attention mechanism, STACKTR. TRANS integrates stack operations into layer-wise hidden state updates. This design maintains compatibility with existing frameworks like flash-attention (Dao et al., 2022; Dao 2024] and supports architectures ranging from 0.36B to 7B parameters. By learning stack operations explicitly, STACKTRANS is capable of addressing broader linguistic and algorithmic challenges.
3 Method
As introduced previously, the standard Transformer architecture (Vaswani et al.) 2017] struggles to learn the Chomsky hierarchy due to its lack of inductive biases. To address this limitation, we propose STACKTRANS, which incorporates hidden state stacks into the Transformer architecture. These stacks introduce the assumptions of the Chomsky hierarchy, enhancing the model's ability to capture hierarchical structures. In STACKTRANS, the hidden state stacks augment the information flow by
3

---

① Updating Differentiable Stack States & Maintaining Stack Mask
spol Awads
ft
stack
mask,-1
U/
push
pop
op
stack,
mask;
query-Ove/-atae attantoo
merge
Reading from Differentiable Stack
Figure 2: Illustration of the multi-head differentiable stack. Our designed differentiable stack includes the updating mechanism based on three action probabilities (i.e., push, pop, and no-op), the stack mask maintenance, well gated global reading mechanism. To improve both memory efficiency and representational flexibility, we also add multi-head and low-rank mechanisms.
routing token-level hidden states through learnable stack operations, such as stack updating and reading. Specifically, the stacks are integrated between standard Transformer layers, where they store hidden states and perform updates and readings via soft operations (see for details). To further improve the expressiveness of STACKTRANS we introduce a multi-head stack mechanism (please refer to $3.2 This design enhances the ability of STACKTRANS low-rank representations. Finally, to ensure robust training and avoid issues such as stack operation collapsing, introduce stack regularization techniques. We also implement stack truncation to facilitate batching and parallel training ($3.3). These innovations and designs collectively enhance the effectiveness and stability of STACKTRANS in learning Chomsky hierarchy grammars.
capture diverse patterns with
3.1 Hidden State Stack
In computational linguistics,, it is well established that the Chomsky hierarchy grammars can be resolved by different classes automata, with pushdown automata being the minimal computational model for DCFs [Chomsky and Schützenberger 1959). Drawing inspiration from this, we incorporate stacks into the Transformer architecture to address its limitations in learning the Chomsky hierarchy. Recall that standard stack is a read or updated (i.e., operations like peck, push, and pop). Given hidden state sequence h,, fy (for now, do not consider any computational dependencies among he '), at each time-step t, "TACKTRANG is designed to either push the current hidden state h2 into the stack or pop the top element from the stack, and then peek the current top element (at this stage, ignore how these operations are determined or executed by the model). This procedure results a new stack ttack-operated sequence, which is permutation of h,, h- Ideally, if the stack operations are correct, it is highly plausible that the model can effectively learn the target Chomsky hierarchy grammar.
storage structure that allows the top element to be
Since the hidden state stack will be incorporated into STACK TRAS, it must be differentiable to en- able end-to-end training. However, standard stack operations such as push and pop are discrete, which disrupts gradient pockppropagttion and hinders the training process. To address this challenge, introduce a soft operation mechanism, following earlier explorations [Grefenstette et al. 2015 Joulin and Mikolov 2013). In this mechanism, the results of the stack operations are continuously interpo- lated based on some trainable parameters. This design not only makes the operations differentiable but also allows the model to learn the stack operations through these trainable parameters.
Soft update The stack time-step can be formalized as a list of vectors, where St, i refers to the i-th element from top to bottom in the current stack St,. We define three candidate operations for STACKTRANS determined by the current hidden state ht (here, St, [1] and h, both belong to Rd meaning they share the same width d): "push" shifts every element down by one position and puts h at the top; "pop" removes the top element and moves every remaining element up by
one position; and "no-op" does not alter St, at all. Instead of discretely selecting of these

---

operations, the soft update mechanism computes a distribution over the candidate operations. This is achieved through a learned linear projection A 3xxd followed by a softmax function:
[a dz P] Softmax(Aly,
(I)
where each scalar in a, represents the probability corresponding to one operation. We then combine the results of each operation based on their respective probabilities to update the stack as follows:
ap he push Steli -1 ap Steli 1
+
dt pup
St,[1]
+ "t maop St. [0],
if =0
Ste+1[i]
+
at
+ a, noop St. [1],
= (2)
+
at
St. +1
+ noop St, [.],
otherwise
where S denotes the size of the stack. The first and the second rows in Equation correspond to the top element [0)) and the bottom element (St++1[S- - 1]) respectively, while the last row pertains to the intermediate elements. It is important to note that zero vector is always maintained at the bottom of the stack, as indicated in the "pop" term of the second row in Equation 2
The soft update mechanism is fully differentiable, enabling end-to-end parameter tuning. Meanwhile, the dynamic of the information flow aligns with that of standard stack. When a.uah dominates in a,, all elements the stack tend to shift downward as more information from h2 flows into the top element; conversely, when the pop operation dominates, the elements in the stack shift upward the information of the top element is mostly removed.
Stack mask To implement proposed hidden state stack using list, the overall available stack size S must be sufficiently large. Assuming S is large enough, the tail of the stack is always padded with zero vectors. This padding ensures that stack operations comply with logical constraints and prevents invalid behaviors. However, this process is inherently discrete. To address this, we propose maintaining a differentiable stack mask for STACKTRANS. Specifically, the i-th element in the mask M, E R$ suggests how likely the corresponding element in the stack is active. The mask is updated with dynamics similar to those described in Equation 2
push a
+ pop [1[1]
+ at d. M,[0], i 0
M
push d push
M, - 1]
+ a 0
+ at noop
M.[i], -i=s-1
Ml -1]
1]
+ &eP M +1
noop + a, Me(i], otherwise
M, serves as an activation controller STACKTRANS- if quah dominates, one more stack element is further activated, otherwise at' dominates in a,, the last activated element is more likely to be deactivated. When accessing the stack, we pad the stack by element-wise production of St, and M,
Global read The standard stack simply peeks and returns the top element during reading. Although peeking is quite straightforward, we notice that may cause unstable training during our initial experiments. Furthermore, limiting access to only the top element restricts gradient flow, reducing learning efficiency and leading to unstable training dynamics. By relaxing this constraint and allowing global read operations, our differential stack achieves greater performance. A detailed discussion is provided in $B.4 Therefore, we propose the global read mechanism, by collecting information over the stack. Global read is achieved through learnable query over-stackk attention:
R = Softmax(W, (St, @ M,)) St,
where @ refers to wise production and W, e R is a trainable query vector. The Softmax term computes the attention score, and the content read from the stack is the weighted sum of all stack elements, where each element is weighted by its corresponding attention score. The final output is residual-like connection hi 9h. h4 + R,, where 9r. is trainable parameter.
3.2 Multi-Head Low-Rank Stack
In the Transformer architecture, the multi-head attention mechanism (Vawwani al. 2017 processes multiple attention patterns in parallel across different representation subspaces. The design enables
5

---

the model to capture diverse relationships within the input sequence. Following a similar design philosophy, we propose the
low-rank stack. Specifically, we doonnproojen the hidden state h4 R& and split it into subspaces (R as: hh h/ h (H) Reshape down he), h
where H denotes the number of stack heads and W down ( is the down-projection matrix. Reshape is standard operation in deep learning that changes the shape of a tensor. In this context, it reshapes the pown-projected vector into independent vectors of dimension dg for the multi-head stacks. Each head corresponds to an independent stack as introduced in $3.1
Given the down-projected hidden state h Rdd for the i-th head, the stack element St." and the mask M (both in xxxx)) along with the final read-out R are computed independently for each head following Equations
After performing soft updates and global reads for all H heads,
concatenate their outputs to obtain the final result: h, 9h h4 + Wap Concat ( (2) ;...; ;. R(")),
where Wap up (d)((Hd)) is the up-rrojection matrix. This multi- head mechanism allows the model to organize stack operations into different patterns in parallel, thereby capturing diverse relationships and dependencies more effectively. Empirically,
find that small H (e.g., 4 or 8) is sufficient to achieve notable performance. For more details on ablation studies and discussions, please refer to {6 On the other hand, the overall computational cost with the low-rank design is much smaller than counterpart of single-head stack with full dimensions due to the reduced dimension of low-rank adaptation in each sub-stack. Refer to for ablations.
3.3 Key Implementation & Training Know-Hows
The modular stack of STACKTRANS enables it to augment the Transformer architecture without altering the Transformer layers themselves. There are some key implementation and training insights.
Stack Overflow In $31 and 3.2 we assume that the stack size S is sufficiently large or infinite. However, due to limitations in computational power and storage resources, S is typically relatively small, making overflow inevitable. In our implementation, we address this by truncating the stack and setting all overflow elements zero, that is, St, [i] 0 if 2 S. This truncation can be seen as a form of "forgetting", where the information carried by the overflow elements is discarded.
Training parallelism Ideally, the stack is supposed to process hidden states according to their temporal or layer dependencies, i.e., it should prioritize hidden states generated from earlier tokens or shallower layers. One feasible sequence fed into the hidden state stack would be [he,.0; ht.,1."" Lt L,, he, ht,
"...) ". layer at token t, and represents the total number of layers. Although such a behavioral pattern is clearly beneficial for learning stack operations, the temporal dependencies conflict with the parallel training of the Transformer layers. To facilitate training parallelism, we implement STACKTRANS by breaking these temporal dependencies. Specifically, STACKTRANS learns stack operations based on the hidden state sequence at token 4; from layer 0 to L ((b,, 00 ... ,h,,,.)). allowing all tokens to be trained parallel. We provide detailed discussion in $B.3
where h4, E Rd denotes the hidden state of the i-th Transformer
Stack regularization During training. STACKTRAN optimizes the standard autoregressive lan- guage modeling loss (CLM) over the token sequence. To prevent the operation probabilities a, from col- lapsing into uniform, we introduce entropy- -based regularization term, defined as Ls. E, H(a)), where H(-) calculates the entropy. The overall loss function combines the language modeling loss and the stack regularization term, L LLM + X Ls., where À is hyperparameter. As regularization term, we give À a small weight in experiments, e.g., 0.001.
4 Evaluation against Formal Languages
Understanding formal languages is fundamental for modeling many aspects of real-world natural language processing tasks. To highlight the motivation behind our stack-enhanced mechanism, first evaluate TRANS on formal language modeling tasks inspired by the Chomsky hierarchy.
Experimental Setup this section, we evaluate STACKTRANS on three groups formal language modeling tasks aligned with the Chomsky hierarchy (Delétang et al., 2022 These tasks assess

---

Table 1: Test accuracy on formal language tasks compared to Transformer Vaswani al., 2017
LSTM Graves and Graves 2012]. StackRNN [Joulin and Mikolov 2013]) and StackAttn Li et al. 2024
Task
LSTM
StackRNN
StackAttn"
Transformer
STACKTRANS
Regular (RE) Tasks
Even Pairs
1.00
1.00
0.49
1.00
Parity Check
1.00
1.00
0.50
1.00
Cycle Navigation
0.89
1.00
0.20
1.00
Deterministic Context-Free (DCF) Tasks
Stack Manipulation
0.66
0.85
0.93
0.53
0.92
Reverse String
0.71
0.80
1.00
0.55
1.00
Modular Arithmetic
0.43
0.42
0.30
0.30
0.60
Context-Sensitiv (CS) Tasks
Missing Duplicate
0.68
0.67
0.53
1.00
Odds First
0.60
0.55
0.51
0.53
Binary Addition
0.56
0.51
0.48
0.48
Binary Multiplication
0.56
0.53
0.50
0.48
Compute Sqrt
0.64
0,63
0.51
0.57
Bucket Sort
0.79
0.75
0.79
0.89
Results listed here are reported by Li et al. [2024).
model's ability to learn underlying compositional rules of formal languages and generalize to input lengths beyond those during training. Please refer to S for details. Following prior work Delétanng et al., 2022], we implement STACKTRANS with relatively limited parameters. Concretely, five Transformer layers with d 64. H is
4 and d, is set to 8. We compare STACKTRAND to some representative baselines, including the standard Transformer [Vaswani et al.)
2017. LSTM [Graves and Graves 2012). StackRNN [Joulin and Mikolov 2013)) and StackAttn
et al.
2024 maintaining identical experimental settings for all models.
Evaluation Results Table 1 shows that STACKTRANS consistently outperforms the standard Transformer, particularly on RE and DCF tasks. For RE tasks, most evaluated models attain near- perfect accuracy. However, Transformers tend falter in the absence of explicit inductive biases. This further underscores the effectiveness of the hidden state mechanism introduced in STACKTRANS. When compared with the state-of-the-ard stakk-augmented approaches, such StackRNN and StackAttn, STACKTRANS either outperforms them or is at least on par with them in nearly all tasks.
The hidden state stack mechanism likely endows STACKTRAN with characteristics akin pushdown automata. This evidenced by its superior performance over all baselines on both RE and DCF tasks, which are known be solvable by pushdown automata. However, pushdown automata theoretically incapable of resolving CS tasks, we observe that all approaches, including STACK AAN,,, perform poorly on CS tasks. Despite this limitation, STACKTRanK demonstrates the ability to handle a subset of CS tasks, which we attribute to specific design enhancements such as the multi-head mechanism and the global reading capability. These features provide STACKTRAN with stronger modeling capacity than traditional pushdown automata, enabling it to capture additional dependencies and complexities beyond the theoretical limits of pushdown automata.
5 Evaluation against General Natural Languages
From the perspective of computational linguistics, natural languages are generally considered to belong to class of languages that includes DCF [Gazdar and Pullum 1982 Chomsky 1956, Shieber 1985]. Therefore, to thoroughly assess STACKTRAAS language modeling tasks. We examine the scalability of STACKTRAN through scaling law studies, analyzing the effect caused by model size and training tokens. @ STACKTRANS with 360 million parameters is pretrained with nearly 980 billion tokens. We evaluate its performance standard benchmarks, comparing models with similar or larger parameter scales. We provide additional empirical observations through ablation studies and deeper analysis (please refer to ())
we conduct further evaluations on general
variety of

---

Lunguage Wodkiigg Love
SuackTamm Standard Transforme
Standard transtommr
2.2
Standard Waanforme"
Meddlling
2.0
Adva raaee
1.5
100 2003
sa0e v00e
Model Size (Millions Parameten)
Datsset (Blliions Thienss
Training Steps
(a) Scaling law on model size
(b) Scaling law on dataset size
(c) Training loss dynamics
Figure 3: Scaling law and training loss dynamics across early steps of STACKTRANS.
Table 2: Evaluation results on natural language tasks.
STACKTRANS
SmolLM
SmolLM2
Qwen2.5
OLMo
TinyLLaMA
Param. (B)
0.36
0.36
0.36
0.5
1.1
1.1
Token # (T)
0.6
4
18
3
HellaSwag
59.5
51.8
54.5
51.2
60.7
55.2
ARC
59.0
50.1
53.0
45.4
44.0
43.4
PIQA
71.7
71.6
71.7
69.9
75.2
72.9
MMLU
36.5
34.4
35.8
33.7
31.9
32.2
CommonsenseQA
37.1
35.3
38.0
31.6
40.3
37.0
TriviaQA
11.2
9.1
16.9
4.3
2.8
9.8
Winogrande
52.8
52.8
52.5
54.1
53.2
55.7
OpenBookQA
37.5
37.2
37.4
37.4
38.0
33.2
GSM8K (5-shot)
33.6
1.6
3.2
33.4
1.8
1.7
Average
44.3
38.2
40.3
40.1
38.7
37.9
We follow the OLMo framework AllenAI 2024 to pretrain STACKTRANS. Our corpora from Dolma (Soldaini al., 2024] and Smoll Allal et al. 2025 which contain high-quality natural language, math, and Python code examples with diverse domains. We carry out data filtration, ultimately obtaining approximately 980 billion tokens. The data filtering followed standard LLM pre-training procedures, including deduplication, removal of low-quality text, and filtering of harmful content. To scale up model parameters, configuration in OLMo, using roughly 80 billion tokens for each variant model training. For scaling up training tokens, we train STACKTRANS with 360M parameters on sampled subset of 200 billion tokens from pretraining corpora.
Experimental Setup
adapt the Dolma vi.6-sample
Scaling Law of STACKTRANS We train STACKTRAN models with range of parameter sizes (360M, 600M, 1.OB, 1.5B, and 7B) under the same training budget in terms of tokens. The language modeling loss is tracked throughout the training process, and the scaling trends are depicted in Figure 3 Our observations find that STACKTRANT exhibits smoother convergence and attains lower final loss compared standard Transformers equivalent size. Notably, even with 360M parameters, STACKTRANS consistently demonstrates smaller loss, which underscores the significant contribution of the hidden state stack mechanism improved generalization capabilities. Overall, STACKTRANS aligns well with the predicted scaling trends and delivers superior performance. To analyze the optimization process, we compare training loss dynamics between STACKTRANT and the standard Transformer in Figure 3(c) Detailed analysis is shown in B.5
from scratch, and the detailed model configuration is shown in $ To assess the downstream capabilities, we evaluate STACKTRANS -60M on a comprehensive suite of widely-used benchmarks, and details are shown in $EI As listed in Table STACKTRANS-360M, outperforms all baseline models, including those with significantly larger parameter sizes. Notably, achieves substantial gains on GSM8K and ARC, highlighting its strength in reasoning tasks that require compositional generalization, recursion, or latent state management. Despite having fewer parameters, STACKTRAM performs competitively on PIQA and CommonsenseQA. further indicating that the Btack-augmente memory module improves
Evaluation against STACKTRANS--30N
We pre-train STACKTRANS-

---

Table 3: Training and validation results of stack ablations (~20B tokens).
Model
Train Loss 4
V2 Loss 4
V2 PPL +
V3 Loss
V3 PPL 4
Transformer
2.411
3.518
34.38
3.195
25.33
STACKTRANS
2.359
3,432
32.89
3.092
24.50
TraueTrans (Suck=-Queue)
2.679
3.679
35.14
3.211
25.97
Push- Push-Only (Pix a."""
2.875
4.032
39.02
3.407
27.13
Single-Head (H
2,493
3,552
33.56
3.130
24.91
Hall-Dimension (H
2.370
3.457
33.05
3.105
24.73
Validation Training Loss
Training Loss
Validatio Lose Triiing
Lons Valer
14
12520
5.412
Valistttion
3,60
3.00
345
1.457 3,448
Optamal
3,455
3,432
3,436
3
3.4
12
140
Vell Lowa: 3.492
Val Lomu:J433 7
Oplind at
240
Vall Loss: 3.82
2.585
28
23%
2,30
242,093
2,350
234
23000
23550
24
2.350
23361
2.70
236
Singlo-Hed
16
Full-Dan 128
32
16
28
(a) Stack head If
(b) Stack Dim de
Stack size
Figure 4: Hyperparameters ablations of STACKTRANS.
representation capacity without compromising generalization. Overall, STACKTRANS-360M achieves an average performance of 44.3 across 9 diverse tasks, exceeding comparable models in the table with only fraction of the parameter size and dataset size. Our proposed STACK TRANS enhances LLM's generalization ability, especially in scenarios with limited computation and parameter budgets.
Discussion
Ablations of Key Designs To assess the impact of the stack design in STACKTRANS. investigate We replace the stack with a queue, adopting first-in- first-out
two alternative configurations. storage structure, which we term QueueTrans Apart from this modification, the queue mask performs similar functions to the stack mask in maintaining valid and activated elements. @ In this extreme setting, we modify the stack in STACKTRANT to "push-only", where mush fixed to 1. This configuration essentially disables pop and no-op operations. The multi-heas stack and the low-rank compression are the other two crucial design elements in STACKTRAN To study their impact, we conduct ablation studies as follows: We disable the multi-head splitting, reverting back single-head stack. 4 We remove the down- and up-projections (low-rank mechanism) from the full stack. In total, we create four variants for ablation studies.
We evaluate all the variants introduced above on the V2 and V3 validation sets Zhu et al. 2024 Ex- perimental details are provided in | The ablation results are presented in Table The QueueTrans variant exhibits notably higher perplexity and lower overall accuracy compared to STACKTRANS particularly on tasks involving hierarchical recursive patterns. This outcome is consistent with our expectation that queue operations are inherently less effective at modeling nested dependencies and grammars. Similarly, the push-only variant performs poorly, with both training and validation losses significantly deteriorating. The absence of pop operations impairs its ability to dynamically manage and retrieve stored information, thereby reducing its overall effectiveness, Both single- and mull-dimensionll variants are consistently outperformed by STACKTRANS. The multi-head mechanism enhances flexibility by enabling parallel decomposition of stack streams, while the low-rank mechanism reduces computational costs without compromising much modeling capacity. The results show that the single-head stack more efficient but performs slightly worse, whereas full-dimension stack is extremely costly for limited performance gain. This precisely demonstrates that our multi-head low-rank design is the optimal trade-off, achieving most of the performance gains at cost close to the bascline Transformer.

---

Ablations of Hyperparametesr In addition to the configuration of the standard Transformer layers. STACKTRANS has three key hyperparameters the number of stack heads H. the dimension of each stack head d,, and the stack size S. We perform grid search over reasonable ranges for these hyperparameters, training STACKTRAN ARS with 20 billion tokens. The curves of training loss and validation loss (evaluated on V2 (Zuu et al. 2024]) are plotted in Figure + It is clear that H is crucial for parallelism, but Figure ((a indicates that performance plateaus once H surpasses a certain threshold. From Figure 4(b) we observe that setting d, within the range from 16 to 64 balances the computational cost and the model's expressiveness effectively, Similarly, Figure [cC shows that increasing S from 24 to 32 has nearly impact performance. Given that larger S leads to higher storage overhead, we make a trade-off and ultimately set S to 24 during our evaluation.
More Detailed Discussion Due the length constraints of the paper, provide more discussions in the appendices. For in-depth investigations of STACKTRANS' training and inference efficiency, please refer to $B.1 In implementation, parallelism of STACKTRANS, as briefly introduced in $3.3 We further discuss why this approximation works in |B.3 We adopt global reading rather than top peeking for STACKTRANS, and we explain the rationale behind this design in $B.4 We provide visualizations of the stack action patterns across different tasks in BB.2 and analyze the training dynamics [B.5]
break the temporal dependencies facilitate training
7 Conclusion
Inspired by pushdown automata, we propose STACKTRANS, novel Transformer variant architecture integrating differentiable hidden state stacks in between Transformer layers. STACKTRANG improves generalization in both formal language tasks and natural language modeling tasks. In particular, our from-scratch pretrained STACKTRANS-BAR 2-3x more parameters, showcasing its superior efficiency and reasoning capability.
outperforms several larger open-source LLMs with
Acknowledgments and Disclosure of Funding
This research supported by the National Key R&D Program under Grant No. 2023YFB4503801, the National Natural Science Foundation of China under Grant No. 62192731, 62192730, 62192733, the Major Program (JD) of Hubei Province (No.2023BAA024)
References
GPT-4. 2023.
--tp:///gg---------------------- OpenAl,
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek 11m: Scaling open-source language models with longtermism. arXiv preprint arXh::2401.02944 2024.
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint arXiv:2309. 16609, 2023.
Huangzhao Zhang, Kechi Zhang, Zhuo Li, Jia Li, Jia Li, Yongmin Li, Yunfei Zhao Yuqi Zhu, Fang Liu, Ge Li, al. Deep learning for code generation: a survey. Science China Information Sciences, 1719)191101, 2024a.
Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, al. Neural networks and the chomsky hierarchy. arXiv preprint arXiv:2207.02098, 2022.
Michael Hahn. Theoretical limitations of attf-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 171, 2020.
Noam Chomsky. Three models for the description of language. IRE Transactions on information theory, 2(3):113-124 1956.
Tom M Mitchell. The need for biases in learning generalizations. 1980.
10

---

Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzales Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arxiv 2018. arXiv preprint arXiv:/806001200 2018.
Laurent Sartran, Samuel Barrett, Adhiguna Kuncoro, Miloš Stanojević, Phil Blunsom, and Chris Dyer. Transformer grammars: Augmenting transformer language models with syntactic inductive biases at scale. Transactions of the Association for Computational Linguistics. 10:1423 -0:3993-1439 2022.
Gerald Gazdar and Geoffrey K Pullum. Generalized phrase structure grammar: a theoretical synopsis. (No Title), 1982.
Stuart M Shieber. Evidence against the context-freeness of natural language. The Formal complexity of natural language, pages 320-334. Springer, 1985.
John E. Savage. Models of computation exploring the power of computing. Addison-Wesley, 1998. ISBN 978-0-201-89539-1.
Michael Sipser. Introduction the theory of computation. ACM Sigact News, 27(1):27-22 1996.
Noam Chomsky and Marcel P Schützenberger. The algebraic theory of context-free languages. In Studies in Logic and the Foundations of Mathematics, volume 26, pages 118-161, Elsevier, 1959.
SS Yau. Computation: Finite and infinite machines (marvin I. minsky), 1969.
Brian DuSell and David Chiang. Stack attention: Improving the ability of transformers to model hierarchical patterns. In ICLR, 2024.
Jiaoda Li, Jennifer C White, Mrinmaya Sachan, and Ryan Cotterell. A transformer with stack attention. arXiv preprint arXiv:2405.045/5 2024.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with 10-awareness. Advances in Neural Information Processing Systems (NeurIPS), 2022.
Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024.
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, al. Olmo: Accelerat- ing the science of language models. arXiv preprint arXiv:2402.00838, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
GPT-4o. https://openal: http:///nnnicom///cdeee/elll-o/o
OpenAl. 2024.
Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack-augmented recurrent nets. Advances in neural information processing systems, 28, 2015.
Edward Grefenstette, Karl Moritz Hermann, Mustafa Suleyman, and Phil Blunsom. Learning to transduce with unbounded memory. Advances neural information processing systems, 28, 2015.
Alex Graves and Alex Graves. Long short-term memory. Supervised sequence labelling with recurrent neural networks, pages 37-45, 2012.
AllenAI. httpe::/ggtaabb.co///lllleaaa///
Github, 2024.
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, et al. Dolma: An open corpus of three trillion tokens for language model pretraining research. arXiv preprint arXiv:2402.001530 2024.

---

Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Gabriel Martín Blázquez, Guilherme Penedo, Lewis Tunstall, Andrés Marafioti, Hynek Kydlíček, Agustín Piqueres Lajarin, Vaibhav Srivastav, et al. Smollm2: When smol goes big -ig-data-centric training of small language model. arXiv preprint FrXXv:2502.00737, 2025.
Defa Zhu, Hongzhi Huang, Zihao Huang, Yutao Zeng. Yunyao Mao, Banggu Wu, Qiyang Min, and Xun Zhou. Hyper-connectonns. CoRR, abs/2409. 19600, 2024.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint 19X447 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1397. 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.088: .08361, 2020.
Ankur Arjun Mali, Alexander Ororbia, and C. Lee Giles. The neural state pushdown automata. CoRR, abs/1909.00 00333,,,33,, 2019, http:::p////rrtoooorr/b/////
org abspg///b////99//
19099033333
John Stogin, Ankur Arjun Mali, and C. Lee Giles. A provably stable neural network turing machine with finite precision and time. Inf. Sci., 658: 120034, 2024. doi: IN0.///.NS..023.. 102..4003444 URL
https://doi. org/ 10. 1016/ iisinnsnn 2023. 120034
Keito Kudo, Yoichi Aoki, Tatsuki Kuribayashi, Shusaku Sone, Masaya Taniguchi, Ana Brassard, Keisuke Sakaguchi, and Kentaro Inui. Think-to-talk or talk-to-think? when Ilms come up with an answer in multi-step reasoning. CoRR, abs/2412.01 113, 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? In ACL (1), pages 4791-4800. Association for Computational Linguisties, 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the A12 reasoning challenge. CoRR, abs/1 803.05557, 2018
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In ICLR. OpenReview.net, 2021.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In AAAI, pages 7432-7439, AAAI Press, 2020.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In AAAI, pages 8732-8740. AAAI Press, 2020.
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsensega: A question answering challenge targeting commonsense knowledge. In NAACL-HLT (1), pages 4149-4158. Association for Computational Linguistics, 2019.
Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. AC), pages 1601-1611. A Asooiiation for Computational Linguistics, 2017.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? A new dataset for open book question answering. In EMNLP, pages 2381-2391. Association for Computational Linguistics, 2018
12

---

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichrro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. CoRR, abs/21 1416.18.,,.... 2021.
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Leandro von Werra, and Thomas Wolf. Smollm blazingly fast and remarkably powerful, 2024.
An Yang, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2. 5 technical report. arXiv preprint arXiv:2412.15115, 2024.
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model, 2024b.
huggingface.
hwt.h/bb..oo/uugggggggg///////// Github, 2024.
A Related Work
Evolvement of Large I Language Models The success of L.LMs is deeply rooted in the Transformer
architecture |Vaswani et al.; 2017] and its subsequent variations, which serve as the backbone for
LLMs such as GPT Radfford et al. 2018 GPT-4 2023 GPT-4o 2024]. DeepSeek [Bi et al. 2024
Liu al. 2024) LLaMA [ouvron al. 2023 and other prominent LLM series [Brown
2020
Scaling laws Kaplan et al. [020 have shown that increasing model parameter size leads to emergent capabilities, enabling LLMS to tackle increasingly complex tasks and exhibit surprising generalization
behaviors. In addition to proprietary closed systems, open-source initiatives like OLMo [Groeneveld
al., (2224 demonstrate the potential of community-scall pretraining using meticulously curated datasets such as Dolma (Soldaini et al. 2024). These efforts highlight how transparent methodologies can accelerate innovation, making powerful I.LMS more accessible for academic research and industrial applications. Despite their widespread adoption, the Transformer architecture has inherent expressivity limitations [Hahn] 2020]. Although Transformers are theoretically Turing complete [Chomsky 1956), they often underperform tasks tied to formal languages within the Chomsky hierarchy. Recent studies [Delétang et al. 2022 DuSell and Chiang 2024] have shown that standard Transformers struggle with recursive and hierarchical patterns across both synthetic and real-world tasks. These limitations underscore the need for fundamental architectural enhancements to better model the Chomsky hierarchy, such as rich linguistic and algorithmic structures.
Building on prior work and drawing inspiration from pushdown automata [DuSell and Chiang 2024 [i et al., 2024, (Sartran et al. 2022). we address these limitations by introducing hidden state stacks into the Transformer architecture. Our proposed method, STACKTRANS, enhances the model capacity to represent hierarchical dependencies and recursive grammars, enabling it to learn Chomsky hierarchy grammars effectively. We believe that fostering transparent and open discussions around the underlying architectural challenges will accelerate the evolution of Transformer-based models and propel the development of large language models to new heights.
Stack Augmentation Equipping neural networks with external data structures, such stacks, has been widely explored to enhance models' ability to recognize hierarchical and context- -fre languages. Although earlier studies primarily focused on recurrent neural networks [Joulin and Mikolov 20155] Grefenstette et al., 2015). J, recent efforts have adapted these thoughts to Transformer-basel models [DuSell and Chiang, 2024 Li et al.) 2024 Sartran al.) 2022 They aim embed stack-like operations into Transformers to address their shortcomings in modeling the Chomsky hierarchy, particularly DCFs. For example, [i et al. (2024) augment standard attention layers with differentiable stacks, enabling soft push/pop operations to model recursive structures. However, this comes with architectural trade-offs, where the stack control tightly coupled with the attention mechanism, leading to increased entanglement and reduced modularity. Mali et al.] [2019] and [Stogin et al. (2024) develop a rigorous theory: they prove orbital stability of continuous stack encodings and show that finite-precision RNNs equipped with such stacks are furing-complee
and can simulate any PDA/TM- providiii tight neuron bounds. Further advances, such as those in DuSell and Chiang 2024], embed stack operations directly within attention heads, providing stronger inductive biases. While these approaches better model Chomsky hierarchy grammars, their validation is largely limited to small-scale models and synthetic datasets. It raises questions about their scalability and
13

---

Table 4: Training and inference efficiency, including time cost and GPU memory usage.
Time Cost
Peak GPU Memory Usage
Model Variant
Training
Inference
Transformer' STACKTRANS
x 1.00
x 1.000
x 1.000
x 11.16
x1.09
x1.12
Single-Head Stack
x 1.03
x1.04
x1.00
Full-Dimensional Stack
x3.78
x3.30
x x1.73
All values represent multiples of results of the baseline Transformer.
Network Depth
Push
0.45
No-op
0.40
two-digit addition
0.35
20
Push
0.30
Pop
No-op
two-digit multiplication
Figure 5: Average probabilities of three operations across the network depth for the two-digit addition task and the two-digit multiplication task.
generalization ability to large-scale tasks and larger Transformer-base models with millions or billions of parameters. In contrast, STACKTRANS introduces differentiable stack. Rather than tempering the attention mechanism, STACKTRANS integrates stacks in between standard Transformer layers. This design decouples stack manipulation from attention computations, enabling seamless integration with pretrained Transformer models while preserving compatibility across architectures of varying sizes (ranging from 0.36B to 7B parameters). By focusing on hidden states, STACKTR ANS maintains flexibility and scalability, addressing broader linguistic and algorithmic challenges without compromising the core principles of the Transformer architecture.
B Discussion (Cont.)
Following
$
we provide some more discussions in this section.
B.1 Training and Inference Efficiency
Considering that the training and inference efficiency is an important factor for model applications. In this section, investigate the training and inference efficiency of STACK compared to the standard Transformer. To keep a fair comparison, all comparisons are conducted on the same hardware setup. We measure both training and inference time over 100 consccutive steps under identical hyperparameter and batch size configurations. To explore the design trade-offss we also compare several STACKTRANG variants, including single-head stack and full-dimennional stack described in ş0 As shown in Table 4 despite incorporating differentiable stack modules, STACKTRANS achieves competitive training and inference efficiency. Concretely, it introduces only marginal overhead (around 10%) compared to the standard Transformer while yielding significant performance improvements. Besides, memory usage increase is moderate, well within the typical consumption range of large-scale L.M.. This suggests that STACKTRANK offers practical and scalable approach, and its stack mechanism can be integrated into existing Transformer without compromising deployment efficiency.
14

---

B.2 Stack Action Patterns across Layers and Tasks
STACKTRANS introduces mayer-wise stack modules that manipulate memory using three soft actions: push, pop, and no-op. Since stack dynamics play a central role in the model's 's expressiveness, we conduct an in-depth analysis of the stack action patterns across layers and downstream tasks. We select two arithmetic tasks from our synthetic benchmark suite: the two-digit addition task and the two-digit multiplication task. Although both tasks involve structured numerical reasoning, multiplication generally requires deeper or more nested intermediate steps than addition. Our STACKTRANS-3OR achieves 100% accuracy on both tasks, indicating the performance of our model on these basic arithmetic questions. For every layer and timestep. compute the average action probabilities of the three operations for each Transformer layer in STACKTRANS and visualize their trends across the network depth, as shown in Figure [
The results reveal consistent action distribution pattern on both tasks: earlier layers predominantly favor push operations, while later layers exhibit an increased use of pop, with no-op remaining rela- tively stable throughout. This trend suggests that STACKTR. ANNSS automatically learns to incrementally store information during lower layers and retrieve it in upper layers. The intuitive behavior mirrors how hierarchical or recursive structures are processed.
Figure 5 further shows that multiplication elicits markedly more push operations in middle layers and deferred pop activity in higher layers, reflecting the deeper computation graph required by the task. In contrast, addition induces flatter push/pop pattern distributed more evenly across layers, consistent with its shallower reasoning structure. These findings confirm that STACKTRAN learns to adapt memory access patterns dynamically according to task complexity, and that its stack behavior both interpretable and task-sensitive.
B.3 Approximations for Training Parallelism
In our implementation, we introduce necessary approximations to maintain training parallelism while preserving model performance, as detailed in $3.3 Temporal dependencies among hidden states are a fundamental aspect of the Transformer's ' processing pipeline. Let h,,; E IR represent the hidden state at layer i for token I, where the ideal processing order for our stack would follow the complete sequence:
(5) with L denoting the total number of layers. However, to enable practical training parallelism, we introduce controlled truncation between h and 1,+1.0. This approximation allows us compute token losses for all elements in a sequence simultancously, which would otherwise be computationally prohibitive.
[bs,00, htg.1...., L ht. .....
The decision break these temporal dependencies is guided by two key considerations, First, the self- attention mechanism inherently captures cross-token relationships, which can partially compensate for the truncation. Second, the stack mechanism introduced our model complements the attention layers by retaining sequential dependencies through external memory operations. Empirical evidence from prior work [Kudo et al. 2024) supports this design, showing that those lntermediiteellyer hidden states for subsequent tokens effectively preserve information from earlier tokens. Overall, this approximation achieves an optimal trade-off between computational efficiency and model performance, enabling scalable and parallelizable training while maintaining our designed stack mechanism.
B.4 Global Reading Capability
Traditional stacks only permit access to the top element. However, in neural network modeling, such a restriction is unnecessary since tensor vectors can be efficiently accessed through operations like matrix multiplication. To enhance the stack's representational ability, our differential stack eliminates this constraint by introducing global reading capabilities and enabling full random access.
Furthermore, we find that enforcing strict top-only access during training leads to unstable and suboptimal model performance. We attribute this to the frequent stack operations neural networks: limiting access to the top element disrupts gradient flow and reduces parameter learning efficiency. By relaxing this constraint and enabling global read operations, our differential stack achieves greater representational ability and improved adaptability across diverse tasks.
15

---

B.5
Training Loss Dynamics
To analyze the optimization process, we conduct comparative analysis of training loss dynamics between STACKTRANS and the standard Transformer. The training curves are presented in Figure One may find out that STACKTRANS exhibits a slightly slower decrease in loss during the very early training stages compared to the standard Transformer. We attribute this behavior to the additional learning complexity introduced by the hidden state stack, where STACKTRAN must learn when and how to carry out stack operations. As training progresses, however, STACKTRANS not only catches up but eventually surpasses the standard Transformer in convergence speed. Once the stack operation distribution stabilizes, STACKTRAN begins to leverage the stack more effectively, leading to a steeper decline in loss and an overall lower convergence plateau. This phenomenon shows that STACKTRANS, while requiring slightly longer warm-up phase, ultimately achieves greater learning efficiency and a superior asymptotic performance ceiling compared to the standard Transformer.
3(c)
C Limitations
While STACKTRANE demonstrates strong performance across variety of tasks, there are several limitations to our work that we aim to address.
Limitations of Model Size and Dataset Size Constrained by computational resources, we limit our final pre-trained model to 360M parameters and use approximately trillion training tokens. Although the results show competitive performance, the scaling law discussed in suggesss that larger models and datasets could further amplify the strengths of STACKTRANS. Particularly, scaling up the number of model parameters and training tokens may enhance its ability tackle more complex tasks. This limitation highlights the importance of access to large-scale computing infrastructure for future research. We hope to leverage the power of the open-source community validate this new architecture.
Necessary Approximation in Design To achieve training parallelism, we introduce controlled approximations in the sequence processing pipeline, as detailed in $3.3) Specifically, the truncation of temporal dependencies between tokens facilitates scalable training but may reduce the model's ability to fully exploit fine-grained sequential patterns. While the self-attention mechanism and stack-based memory mitigate this limitation, the truncation approximation may still pose risks, especially in tasks requiring deep inter-token dependencies. We plan to explore more robust and efficient settings in future work.
D Formal Language Task Details
We follow the experimental settings of Delétang et al. [002). formal language tasks and their complexity level within the Chomsky hierarchy. These tasks assess models' ability to learn underlying compositional rules of formal languages and generalize to input lengths beyond those seen during training. The three task groups are categorized by the Chomsky hierarchy type, including RE (type-3), DCF (type-2), and CS (type-1). 1). The classification adheres to formal automata theory, associating the three tasks with tnite--tate automata, pushdown automata, and linear- bounded automata, respectively. Please refer to Table 5 in for definitions and examples of these tasks. Despite the presence of classification tasks, all tasks are formulated as sequence mapping problems. In this setup, the model takes an input sequence and decodes it into an output sequence. STACKTRANS is trained on sequences with input length uniformly sampled from to 40 tokens. At test time, we evaluate STACKTRANT on sequences with significantly longer lengths up to 500 tokens, thereby measuring its length generalization. Following the same procedure as Delétang et al. 2022]. token-level accuracy is used as the evaluation metric. We repeat each experimental configuration ten times and report the best accuracy achieved.
an overview of the
E General Natural Language Task Details
To assess the downstream capabilities, we evaluate STACKTRANS- on comprehensive suite of widely-used benchmarks [Brown al. 2020) Touvron et al., [023] Groeneveld et al., [024]] including those for common sense reasoning (e.g., HellaSwag, PIQA) Zellers et al. 2019
Clark
16

---

Table 5: Formal language task descriptions and input-output examples.
Task Name
Description
Regular (RE) Tasks
Even Pairs
Check if the count of ab/ba pairs is even.
Parity Check
Check if the count is even.
Cycle Navigation
Navigate movements on modulo-5 cycle.
Deterministic Context- Free (DCF) Tasks
Stack Manipulation
Perform stack operations and return the final state.
Reverse String
Reverse the input string using stack.
Modular Arithmetic
Evaluate nested arithmetic expressions modulo 5.
Solve Equation
Find a variable satisfying modular equation.
(CS) Tasks
Binary Addition
Compute binary addition of two numbers.
Binary Multiplication
Compute binary multiplication.
Compute Sqrt
Compute the integer square root of a binary number.
Bucket Sort
Sort sequence over a fixed alphabet.
Duplicate String
Output the string concatenated with itself.
Missing Duplicate
Find the missing character in duplicated string.
Odds First
Interleave odd and even indices of sequence.
Task Name
Input Example
Output Example
Regular (RE) Tasks
Even Pairs
aabba
True
Parity Check
aaabba
True
Cycle Navigation
011210
2
Deterministic Context-Free (DCF) Tasks
Stack Manipulation
abbaa POP PUSH a POP
abba
Reverse String
aabba
abbaa
Modular Arithmetic
-(1 2) (4-33- (-2))
0
Solve Equation
[= 2) (4 3- (-2)
1
Context-Sensitiee (CS) Tasks
Binary Addition
10010 101
10111
Binary Multiplication
10010 101
1001000
Compute Sqrt
101001
101
Bucket Sort
421302214
011222344
Duplicate String
abaab
abaababaab
Missing Duplicate
ab_aba
Odds First
aaabaa
aaaaba
ct al. 2018 Hendrycks et al., 2021 Bisk et al.
2020:
Sakaguchi et al./ 2020 question answering
(e.g.. OpenBookQA) [Tlmor et 2019 Joshi al.
2017 Mihaylov al. 2018), and math-based
reasoning (GSM8K) Cobbe et al. 2021 We compare STACKTRANS 360M with other open- source models around IB parameters, including SmolLM-3660
[Alal al. 2024 SmoILM2-3660
[Allal et al.) 2023), Qwen2.5-0.. 0.55B (Yang et al.) [024]. OLMo- IB [Groeneveld et al. [222]]. and TinyLLaMA- [Zhang et al. 2024b). We use the lighteval framework [huggingface] [2024]. and for all applicable tasks, we adhere to zero-shot evaluation settings, unless otherwise specified.
"The evaluation protocol strictly follows Allal et al. 2025 and we obtain nearly identical results to those reported in the paper.
17

---

Table 6: Model configuration of STACKTRANS-360M.
Parameter
Value
Vocabulary Size
49152
Number of Attention Heads
15
Number of Hidden Layers
32
Hidden Size
960
Intermediate Size (FFN)
2560
Attention Dropout
0.0
Activation Function
Silu
Number Stack Heads
4
Stack Dimensionality
16
Stack Size
24
Maximum Position Embeddings
4096
RoPE Scaling
None
RoPE 0
100000
Table Overview of V2 and V3 Validation Sets. Each validation set includes diverse text sources to ensure comprehensive evaluation.
Validation Set
Datasets Included
V2 Validation Sets
22-sall--chaan-validtion v2-small-c4 100 1o00dmmains-valdaationn
v2-small-ice-validatio. y2-small-gab--aaiiattn
12-small-m2202_iikivaaadda 22small-manopphrreealiitttiiii
2-smaallnooe_enn-alldaii.nn v2-small-pile--aaidaiinn,
22-small-ptb-vaaidattonn, Y
v2 2-small-wikiteett103---1 100-aliiatt000
V3 Validation Sets
(3-smallec_een-vaiiattioo besssaal-doomma_hookssvallliiiie
v3-small-blmma_commnn commanacmmmmmmmmaa crawl-vvliddation
v3-small-dolma
v3-small-dolma d---mmll-doomaa_rdddddi-lviddddd reddit-validation,
v3-small-dolma aaamallldollmaawi---aaaaaaaaaiid wiki-validation,
D-small-ice-validatiie
v3-small-m2d22 s2ore- validaaai,,,
-3-small-pile-vaaiattonn w3-small-wikiiext 1trpiiiieexx_1133--111111 103-validation
F Model Configuration of STACKTRANS-360M
Table T) shows the detailed model configuration of our STACKTRANS-360M, which is inspired by the
similar setting in Allal et al. (0244 and Allal al. 20255) The stack-related setting is decided by grid search in $
G Validation Dataset Details for General Language Modeling
Following the method of Zhu et al. 2024 we evaluate all variants on the V2 Validation Sets and V3 Validation Sets curated within the OLMO framework. The specific datasets for V2 and V3 validation
Zhu et al. 2024] are shown in Table 7
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claim for this paper enhance LLMs with the stack mechanism. The content of this paper focuses on the design and evaluation of this method, which is the main contribution of this paper. We give full contribution list in I
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations in CC
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiattinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
19

---

Answer: [NA]
Justification:
The theoretical result and proof is mainly from existing related work in A
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide detailed experimental setup for each task in $4
1 and
$
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
20

---

The authors should provide scripts to reproduce all experimental results for the new reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We will provide the key code for dataset processing and training in the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ httpppscccppbllccguuudoooo Coaee///ddeSSbbmmmma ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide all the training and test details for each task in
$
D
and
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[No]
Justification: All the evaluation settings strictly follow the default settings, and details are provided in and $ The randomness in these evaluations is small.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
21

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide all the model size and dataset setting for each task
and $
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
The research conducted in the paper conform with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is no societal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
22

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
The dataset we use follows Apache-2... 2.0 license.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
23

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [Yes]
Justification: This paper is an enhancement method for LLMs. The LLM is the basis of our method.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https:/ /ttprr/////..p....eee. /t//oonfeeennce/22222 for what should or should be described.
25


Reviews:
[{"review_id": "855SUz4myQ", "reviewer": "Reviewer_FFsc", "summary": "The submitted work proposes blending the well-known Transformer (LLM) architecture with an extension for stack automata to overcome parsing limitations of the former. While the attempt is not new, the submission claims that this one is more informed, theoretically more sound and better performing. The latter is proven by much better ratio between accuracy versus number of model parameters and also shows higher coverage of evaluating formal language tasks compared to previous work.", "questions": "* What information exactly is being stored on the stack? I personally failed to understand from Figure 1(a), Figure 2, the text (e.g. Sections 3.1 & 3.2) and also Appendix F. I have a \"feeling\" that it may be aligned with the attention head configuration, but Figure 2 shows the output $h_t$ of the input transformer layer only as abstract tensor and *multiple* up-projecting transforms in parallel. Please describe that mapping in more detail!\n\n* Lines 65-66 claim \"..., achieving nearly 100% test accuracy in most scenarios, as shown in Figure 1(b).\" Looking at Figure 1(b) I count 100% for StackTrans only 5 times out of the 13 test scenarios. That's 38.5%. Can you please shed more light on your claim or correct the phrasing?\n\n* Lines 138-139 mention \"...here, $St_t[i]$ and $h_t$ both belong to $IR^d$, meaning they share the same width $d$: ...\". I suggest to write that more formally, e.g. $St_t[i], h_t \\in IR^d$\n\n* The two definitions of $S_{t+1}[i]$ and $M_{t+1}[i]$ between lines 136-164 can be compressed by defining them as tuple $(S_{t+1}[i],M_{t+1}[i])$ referencing tuples accordingly also on the right side of the equation. That also simplifies the text and eliminates the remark about element-wise product in 171. In fact that tuple notation is motivated by the used in Eq. (4). The gained space can be  used to better explain e.g. the aforementioned mapping between output vectors of the transformer layer and the stack.\n\n* I'm not 100% sure, but think that it is sufficient to define the zero vector only with the mask. Do you share that view? If yes, does that allow for optimizations, e.g. to manage stack overflows or for training?\n\n* Line 180 mentions a unary \"... = $Reshape(W_{down} \\cdot h_t)$\". What's the exact definition of that function?\n\n* The paragraph (lines 212-216) introduces \"\\lambda\" as hyperparameter. What are typical obtained values? It maybe worth to mention by use case or add it to Figure 4 or otherwise move it to an appendix.\n\n* Figure 3(c): The green curve is hard to see and lateron the text refers to \"...smoother convergence...\" (line 264). Please try to improve that. I'm unsure wether it is optimal, but it may help to display the overlap of the two curves in a third color and explain in the subtitle.\n\n* Line 255: \"... comes...\" should be plural: \"... come...\"\n\n* Line 257 mentions \"data filtration\". Can you please provide more details?\n\n* Lines 284-293: for all ablation studies where applicable it may be interesting to compare the computational complexity as well.\n\n* Lines 294-297: in addition to loss / perplexity, how does accuracy compare?\n\n* Table 3 (QueueTrans): it would be interesting to see the effect over size S of both. In light with global reading queues may have a different optimal size setting compared to stacks (with pop).\n\n* Table 3 (last 2 lines): both ablations almost don't hurt, but that in particular isn't being discussed in more detail (e.g. single head may possibly be computationally less costly). I suggest to do add a more detailed discussion.\n\n* Figure 4: are the rightmost points with increasing trend an effect of two many parameters? For those measurements did you apply exactly the same training length or did you account for a possible increase in # of parameters by a longer training with more epochs?", "limitations": "Yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "The big strengths of this submission are the combination of an ambitious idea (differentiable stack), made to work well (e.g. extra complexity to improve the training) and a solid description of the whole effort. The basic description is augmented with additional constrasting (stack vs. queue) / ablation studies and some discussion of the results. I do see more possible refinements to build on top of this work by the community after publishing. The topic is of significant interest to the audience.\n\nThere are some minor weaknesses w.r.t. clarity and structure (a detailed list is in Questions below).", "was_revised": true, "final_justification": "Rebuttal read and acknowledged. Keeping the overall rating unchanged as it reflects the good work already.", "rebuttal": "We sincerely thank Reviewer FFsc for their review and valuable suggestions for revision.\n- On the specific mapping of information stored in the stack: We thank the reviewer for pointing this out. The specific mapping is described in §3.2: The hidden state h_t \\in \\mathbb{R}^d from a Transformer layer is first passed through a down-projection matrix W_{down}, then reshaped and split into H independent heads, i.e., [h_{t}^{(1)}, ..., h_{t}^{(H)}] = \\text{Reshape}(W_{down} \\cdot h_t). Each h_t^{(i)} is then fed into its corresponding independent stack for operations. We will clarify this textual description in the revised version and update the caption for Figure 2 to more clearly illustrate this mapping from h_t to the multi-head inputs.\n- On the \"near 100% accuracy\" claim: Your analysis of Figure 1(b) is correct, and our statement was indeed too broad. We originally intended to highlight that the model achieved 100% accuracy on all RE (Regular Expression) tasks and also demonstrated strong performance on DCF tasks. Out of the 13 formal language tasks, StackTrans did achieve 100% on 5 of them. We will revise the wording to be more precise: \"achieved 100% accuracy on all regular language tasks and demonstrated excellent performance on the majority of deterministic context-free grammar tasks,\" to accurately reflect our experimental results.\n- On the standardization of mathematical symbols: Thank you for the suggestion. We fully agree and will uniformly use the more formal \\mathbb{R}^d to denote vector spaces in the revised version.\n- On using tuple notation to simplify formulas: This is an excellent suggestion. Adopting the tuple notation (S_{t+1}[i], M_{t+1}[i]) will indeed make the formulas more compact and simplify the main text. We will adopt this suggestion in the revision and use the saved space to better elaborate on other key details.\n- On zero vectors and the mask: This is a profound question. In our design, the mask M_t \\in \\mathbb{R}^S is a vector of probabilistic values indicating the \"activeness\" of each position in the stack. The stack itself, St_t, stores the actual hidden state vectors. The final effective stack is obtained through an element-wise product St_t \\otimes M_t. Therefore, both are necessary: M_t controls which positions are valid in a differentiable manner, while the zero vectors in St_t are the actual padding values for invalid positions. While the idea of using only a mask is interesting, our current design ensures the entire process is fully differentiable. We will clarify this point in the revision.\n- On the definition of the Reshape function: We apologize for the lack of clarity. Reshape is a standard operation in deep learning that changes the shape of a tensor. In this context, it reshapes the down-projected vector (with dimension H \\cdot d_s) into H independent vectors of dimension d_s for the multi-head stacks. We will add this explanation in the revised paper.\n- On the hyperparameter \\lambda: Thank you for the reminder. We will add the value of \\lambda used in our experiments. As a regularization term, we give \\lambda a small weight in experiments, e.g., 0.001.\n- On the visibility of Figure 3(c): We agree that the green curve in Figure 3(c) is difficult to discern. In the revised version, we will modify the colors and styles of the chart to enhance its clarity and readability.\n- Grammar correction: Thank you for pointing out the typo. We will change \"comes\" to \"come\" in the revision.\n- On the details of \"data filtering\": On line 257, we mentioned that our corpus comes from Dolma and Smoll. The data filtering followed standard LLM pre-training procedures, including deduplication, removal of low-quality text, and filtering of harmful content. We will add a sentence to briefly explain this in the revision.\n- On the computational complexity of the ablation study: This is a great suggestion. In fact, we have already compared the training and inference efficiency for some ablation variants (single-head and full-dimension) in Table 4 of the appendix. The results show that our multi-head, low-rank design achieves the best performance with a minimal increase in overhead. We will more explicitly reference and discuss these results in the main text in §6.\n- On the accuracy in the ablation study:  In our ablation study (Table 3), we primarily used validation perplexity (PPL) as the core evaluation metric. This is the gold standard for assessing generalization ability in language model research and is highly correlated with downstream task performance. While we did not evaluate every ablation variant on all downstream tasks, the significant differences in PPL already provide strong evidence for the effectiveness of our design.\n- On the size S for QueueTrans: Your point is well-taken; the optimal size for a queue might differ from that of a stack. In our ablation study, to ensure a fair comparison by controlling variables, we used the same stack/queue size S for all variants. Exploring the optimal hyperparameters for different data structures is a valuable direction for future work, which we will consider.\n- On the discussion of single-head/full-dimension ablations: Table 3 and Table 4 show that the single-head stack is more efficient but performs slightly worse, whereas the full-dimension stack is extremely costly  for limited performance gain. This precisely demonstrates that our multi-head, low-rank design is the optimal trade-off, achieving most of the performance gains at a cost close to the baseline Transformer. We will discuss this cost-performance advantage more deeply in §6.\n- On the trend and training length in Figure 4: This is actually because, with a fixed total parameter count, having too many heads leads to a lower dimension per head, which weakens the model's expressive power. All experiments in Figure 4 used the same training dataset setting to ensure a fair comparison. We will clarify this."}, {"review_id": "8oZLJqwBWK", "reviewer": "Reviewer_1yQj", "summary": "This paper introduces a new architecture designed to enhance the capabilities of Large Language Models by integrating concepts from automata theory. The authors identify a key limitation in Transformers: its difficulty in learning and generalizing tasks based on the Chomsky hierarchy, such as regular expressions and deterministic context-free grammars. To address this, the authors incorporate differentiable hidden state stacks between Transformer layers and introduce an inductive bias for hierarchical structures. Comprehensive evaluations demonstrate that the proposed architecture significantly outperforms standard Transformers on Chomsky hierarchy benchmarks.", "questions": "N/A", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n- The paper is well-written, clear, and easy to follow.\n-  The paper addresses a critical weakness of the standard Transformers in the inability to model and generalize on tasks requiring an understanding of hierarchical or recursive structures, such as those in the Chomsky hierarchy. \n- The authors conducted a comprehensive set of experiments. The authors validate their approach across multiple model scales, demonstrating consistent and significant performance gains.\n\nWeaknesses:\n- The primary natural language results are centered on a 360M model. Although the paper demonstrates scaling up to 7B parameters, it could be beneficial to investigate how these architectural benefits translate to the very large-scale models (e.g., 70B+).", "was_revised": true, "final_justification": "After considering the rebuttal and discussions, I maintain my rating of 5. The paper presents a technically solid contribution by integrating differentiable stack mechanisms into Transformers. The authors addressed my key concerns.", "rebuttal": "We sincerely thank the reviewer 1yQj for their high praise and strong support for our work. We are delighted that you found our paper to be \"well-written, clear, and easy to follow,\" and that you recognized the value of our work and the thoroughness of our experiments.\n\nWe agree with your point that validating the advantages of StackTrans on larger-scale models (e.g., 70B+) would be a crucial and exciting next step. As you noted, our current experiments have successfully scaled the model from 360M to 7B parameters, and in this process, we observed consistent and stable performance gains and scaling trends. This provides strong preliminary evidence for our architecture's continued effectiveness on even larger models. Our decision to focus the main natural language evaluations on the 360M model while demonstrating scalability up to 7B was primarily due to computational resource constraints, a point we also mentioned in the limitations section.\nOur work lays a solid foundation for this new architecture and provides compelling evidence of its potential. We hope this will inspire the open-source community and institutions with greater resources to pursue larger-scale validation of this architecture in the future."}, {"review_id": "xsnety5dDD", "reviewer": "Reviewer_nPse", "summary": "This paper describes a new architecture based on the transformer but augmented with Joulin-Mikolov style differentiable stacks. Unlike previous work on stack RNNs and stack attention, these stacks run from \"bottom to top,\" one for each position of the sequence. The new model is tested on both formal languages and natural languages. Especially notable is the fact that the model is scaled up to 7B parameters and 1T tokens.", "questions": "- Is a comparison with DuSell and Chiang (2024) possible?\n- In figure 4, how many layers were used?", "limitations": "The limitations section is relegated to the appendix. I actually don't really agree with the limitations: the model and dataset size are already impressive to me, and the choice to make the stack \"bottom-up\" is different from previous models but not, in my opinion, a limitation.\n\nI don't see any potential negative social impact.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# Strengths\n\nThis is a nice model, with good results. The fact that it could be scaled up to 7B parameters and 1T tokens is extremely impressive.\n\nThere are two features that distinguish this model from previous similar models: (1) the network can attend to the entire stack; (2) the stack runs from \"bottom to top.\" These ideas are strengths, and I am listing them here to make clear that my comments below about them are not criticisms of the ideas themselves.\n\n# Weaknesses\n\n## Model\n\nRegarding (1), \n- If the network can attend to the entire stack, then perhaps it shouldn’t be called a stack. Though, I admit I don't know what to call it instead. \n- I would like to see a more careful explanation of why it’s bad just to read the top stack element. The current explanation is brief, even in the appendix, and has a lot of buzzwords but doesn't really shed much light on the question.\n\nRegarding (2),\n- The decision to make the stack run \"bottom to top\" is presented as an \"implementation know-how\" (lines 202-211) or an \"approximation\" (531). It is not. It is a fundamental modeling decision (and a completely reasonable and original one), and it should be made known much earlier than page 6. It should be in the abstract and on page 1.\n- However, the \"bottom to top\" direction does call into question the motivation of the paper, which is modeling the kinds of dependencies found in context-free languages. There isn't, as far as I can tell, any reason to believe that this model would be better at these kinds of dependencies than a standard transformer.\n- There also isn't any reason to believe that increasing the maximum stack size S would have any effect for S beyond the number of layers in the transformer.\n\n## Evaluation\n\n- In Table 1, please give citations for both StackRNN and StackAttn. They are at line 228, but they should also be in the table caption. The footnote makes it appear as if StackAttn is from DuSell and Chiang (2024), but it is not.\n- Is a comparison with DuSell and Chiang (2024) possible?\n\n## Formal language terminology\n\nIn general the writing about grammars, inductive biases, and especially the Chomsky hierarchy needs to be thoroughly revised.\n\nline 6, 15, 26, 33: The “hierarchy” in “Chomsky hierarchy” does not refer to hierarchical structure, but a sequence of four language classes (regular, context-free, context-sensitive, and recursively enumerable). \n- It wouldn’t be reasonable to expect Transformers to “capture the Chomsky hierarchy.”\n- It’s simply not correct to talk about “Chomsky hierarchies” (plural).\n- There’s no such thing as “the underlying grammar of the Chomsky hierarchy.”\n- I stopped keeping track of all of the misuses of this term. Please search for all occurrences of “Chomsky hierarchy” and fix all of them.\n\nThere seems to be some confusion between context-free languages and deterministic context-free languages.\n\n# Other comments\n\nline 31: Of course Transformers have inductive biases; the question is whether they have the right ones.\n\nline 37: “This inherent limitation” appears to refer to problems with length generalization, but it’s not clear what the relationship of this to context-freeness (line 36) is.\n\nlines 95-100: A summary of new results doesn’t belong in a background section.\n\nline 198: You say that “S is typically relatively small” but don’t say what you actually set S to until line 312.\n\nThere are many incorrect citations. Here are the ones I noticed from the text: Battaglia et al., 1806; Savage and Computation, 1998; allenai, 2024. Glancing at the bibliography, though, I see many more.", "was_revised": true, "final_justification": "Thank your for your responses. I don't have anything to add and will keep my score the same.", "rebuttal": "We would like to sincerely thank the reviewer nPse for their thorough and highly constructive feedback. We will address each of your points below.\n- On the term \"Stack\" and \"Global Read\": Thank you for this insightful conceptual point. We use the term \"stack\" because the core update mechanism of our module is fundamentally based on differentiable push and pop operations. Regarding why reading only the top element is suboptimal, you are right that our explanation in the main text was too brief. As detailed in Appendix B.4, our initial experiments showed that a strict top-only access leads to \"unstable and suboptimal model performance\" because it \"disrupts gradient flow and reduces parameter learning efficiency\". In the revised paper, we will clarify this distinction between the stack-like update mechanism and the global read mechanism. We will also move a more detailed justification for the \"global read\" from the appendix to the main text to make the rationale clearer.\n- On the \"Bottom-to-Top\" Design and its Motivation: We will revise the manuscript to introduce this concept much earlier, including in the abstract and introduction, to properly frame its significance as a core contribution. Regarding its effectiveness for context-free languages, our empirical results in Table 1 provide strong validation. StackTrans achieves scores of 0.92 and 1.00 on the DCF tasks \"Stack Manipulation\" and \"Reverse String,\" respectively, whereas the standard Transformer fails with scores of 0.53 and 0.55. This demonstrates a clear and significant advantage in modeling these dependencies, directly validating our design choice.\n- On Stack Size S vs. Number of Layers L: This is a very keen observation. You are correct that in our current implementation, which is designed for training parallelism, the stack processes hidden states layer-wise for each token. This means the effective sequence length fed to the stack at any given time is the number of model layers, L. We will add a note to the text clarifying this relationship between S and L in our specific implementation. This is also empirically supported by our hyperparameter experiment in Figure 4(c).\n- On Table 1 Results: Thank you for pointing out this ambiguity. In the revision, we will add the citations directly into the table caption (which are already listed in Line 228 in the paper). We will also fix the citation typo in the footnote (the results are reported by [1]). Regarding the baseline in DuSell and Chiang [2024], in our initial experiments, we selected stack-augmented baselines, closely following the setup in [2] to ensure our results were solid. We are now working on the experiments for the baseline you mentioned. Preliminary results already show that our method, StackTrans, outperforms it by a significant margin. The detailed experimental setup and results will be provided in the revised paper.\n\n[1] Jiaoda Li, Jennifer C White, Mrinmaya Sachan, and Ryan Cotterell. A transformer with stack attention.  2024. \n\n[2] Delétang G, Ruoss A, Grau-Moya J, et al. Neural networks and the chomsky hierarchy. 2022.\n\n\n\n- For descriptions for Chomsky hierarchy, We are extremely grateful for these critical corrections. We acknowledge that we were misusing some terms. We will perform a thorough revision of the entire manuscript to fix all instances of these errors. We will replace incorrect uses of \"Chomsky hierarchy\" with more precise phrasing (e.g., \"modeling dependencies found in specific classes of formal languages\"). We will also ensure the distinction between context-free and deterministic context-free languages is handled correctly. This will significantly improve the paper's formal precision.\n- We will also revise those lines to address all points in \"other comments\". \n\nFor questions:\n\n- Comparison with some baseline: As noted above, we will fix the citations in the footnote in Table 1 to properly contextualize our work among other stack-augmented models discussed in the literature. The detailed experimental setup and results will be provided in the revised paper.\n- Layers in Figure 4: The experiments in Figure 4 were conducted using our StackTrans-360M model configuration. We will clarify in the caption that these hyperparameter settings used the 360M model, which has 32 layers."}, {"review_id": "wcUnUKGOSW", "reviewer": "Reviewer_j9mp", "summary": "This paper proposes to solve the limitations of the standard transformers when processing formal languages of the Chomsky hierarchy.\nDue to the equivalence of the pushdown automata (PDA) and the context-free grammar (CFG), the authors design a differentable hidden stack between consecutive transformer layers, yielding the StackTrans model. The setting injects the inductive bias of the way the PDAs work, and enhances the ability of processing grammars of the Chomsky hierarchy.", "questions": "According to the strengths and weaknesses above, some concerns are below:\n\n- Can the learnt stack be shown having effectively captured the grammar?\n- Is there any result of the StackTrans model in different sizes, and are they significant?\n- Is there any performance degeneration of StackTrans on general natural language datasets where grammars of the Chomsky hierarchy are not that important?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths.\n\n- The proposed hidden stack module is more elegant than previous works and can work with other transformer designs with much less invasion in the existing codes. For example, the stacks can be neatly used with Flash Attentions.\n- The StackTrans model is also shown satisfying the scaling law on model sizes and dataset sizes.\n- Experiments have shown the effectiveness of StackTrans on several formal language tasks and some Natural Language Tasks.\n\nWeaknesses.\n\n- To some extent it remains unknown if the hidden stack is indeed a stack. For example, can we restore and check the grammar bottom-up, from the trained dataset? But in StackTrans, the stack is globally read, more like a global memory instead of a stack.\n- The model details are not given. Since there are lots of variants of a transformer layer, could StackTrans work with most of them? If so, perhaps adding the hidden stack into LLMs of different families are needed. Furthermore, can we use the proposed hidden stack as the way of using LoRA?\n- The results of different sizes of StackTrans are not given.\n- Though the CFG grammar had been considered as one way to develop natural language grammars, it's also been widely admitted that there're lots of long-context dependencies in natural languages which cannot be described in CFG. Despite the long debates, at least we hope to see a CFG/PDA-inspired inductive bias for the neural models will not hurt the general performance on common corpus. So perhaps the StackTrans model is also required to be evaluated on general tasks (for example, the GLUE benchmark)", "was_revised": true, "final_justification": "My concerns listed in the weaknesses have been carefully addressed by the authors. I believe the overall quality is sufficient to make a score raising.", "rebuttal": "We would like to sincerely thank the reviewer j9mp for their thorough and highly constructive feedback. We will address questions here in detail one by one.\n- On whether the \"hidden stack is really a stack\" and if \"the learned grammar can be inspected.\"\nThis is a very central question, and we will respond from two perspectives:\n  - On the name and the global read mechanism: We call it a \"stack\" based on its core update mechanism. The information flow is dynamically managed through differentiable push and pop operations. The \"global read\" you observed is a reading mechanism we introduced to solve a key problem in practical training. As described in Appendix §B.4, a traditional \"stack-top only\" read can lead to restricted gradient flow and unstable training. Therefore, our design can be seen as: a module with stack-like update logic, supplemented by a more flexible global read mechanism to optimize the learning process.\n  - On whether we can show that it captures grammar: Yes, we can.\n    - Quantitative Evidence: In the formal language tasks in §4, StackTrans achieved good performance on tasks requiring precise grammar. For example, it scored 1.00 on the \"Reverse String\" task, whereas the standard Transformer only scored 0.55. This directly demonstrates that it has learned the underlying grammatical rules.\n    - Qualitative Evidence: In Appendix §B.2 (Figure 5), we visualized the stack's operational patterns. The results clearly show that the model automatically learns meaningful, task-aligned patterns: it tends to push (store information) in the shallower layers of the network and pop (retrieve information) in the deeper layers. This interpretable behavior could be seen as the evidence that it has learned the underlying \"grammar\" of the task.\n- On model details, compatibility with variants, and comparison with LoRA: \nOne of the core advantages of our design is its modularity and non-intrusive nature.\nAs stated in the paper, StackTrans \"integrates a differentiable stack between Transformer layers while maintaining the integrity of the Transformer layers\" and \"decouples the stack operations from the attention computation.\"This design allows it to be naturally combined with various Transformer backbones. Therefore, we have strong reason to believe that it can also be combined with parameter-efficient fine-tuning (PEFT) methods like LoRA (e.g., by freezing the backbone LLM and fine-tuning only the StackTrans modules). This is a very promising direction for future research.\n- On providing results for different model sizes and their significance: \nThe reviewer may have overlooked the relevant sections, as we have provided detailed results in the paper:\n  - Scaling Law Study: In §5, we systematically studied the model's scaling laws, clearly analyzing the impact of both model size and the amount of training data.\n  - Performance Curves for Different Sizes: Figure 3(a) visually presents the loss curves as the model size increases from 360M to 7B, clearly showing that larger model sizes lead to better performance and outperform the standard Transformer.\n  - Evaluation of a Specific Model Size: All comparisons in Table 2 are centered around our StackTrans-360M model, which is compared against models with larger parameter counts (such as the 1.1B OLMo and TinyLLaMA). The results are significant.\n- On CFG bias and performance on general-purpose natural language tasks: \nYour question—\"whether introducing a CFG bias would harm general-purpose performance\"—is a core question that all methods with structured inductive biases must answer. We address this challenge head-on by evaluating our model on a wide range of general-purpose tasks.\n  - It did not harm performance; it improved it. As shown in Table 2, we tested on 9 diverse general-purpose benchmarks, including commonsense reasoning (e.g. HellaSwag), question answering (e.g. OpenBookQA), and mathematical reasoning (e.g. GSM8K).\n  - The results are significant. The outcomes show that not only did StackTrans not suffer from performance degradation, but it actually achieved significant performance gains. Our 360M model's average score (44.3) far surpasses baseline models that are 2-3 times its size in parameter count. This proves that our stack mechanism acts as an effective supplementary module that enhances the model's structured reasoning capabilities without sacrificing its general language abilities. While GLUE is a good benchmark, we believe the existing 9 benchmarks (for example, MMLU is a popular benchmark that is adopted by many pretraining LLM studies) already provide strong evidence of the model's generality."}]

Author rebuttal:

Meta-review: **Summary:**  Transformers are widely used on tasks like code generation, which involves generating text described by a formal grammar.  Nevertheless, it is well known that they lack the ability even in theory to handle truly recursive structures, and when used for formal languages, they often fail to generalize beyond the lengths and recursion depths that they are trained on.  Inspired by pushdown automata (which can handle context-free grammars), the authors implement a differentiable stack, with push and pop operations, and add it to the transformer.  They test the new architecture on both formal language and natural language modeling tasks. 

Stack architectures have been proposed before, but a key innovation of this work is the fact that the stack runs layer-wise, rather than length-wise.  I.e. each layer can do push/pop operations in parallel for each token position, rather than consuming tokens sequentially.  While this change makes training practical and efficient, it means that the stack is **NOT** being used in the same way that a push-down automata ordinarily operates, which is a clear limitation of the work.

**Metareview:**

The reviewers unanimously vote to accept this paper (score 5).  They reviewers cite clarity of writing, solid implementation, and good scaling behavior (up to 7B params, 1T tokens, and compatibility with flash attention).  Most of the reviews are high quality, drawing attention to weaknesses, errors, and incorrect use of terminology where appropriate, and the authors responded positively to the constructive criticism.  I have looked over the paper myself, and see no reason to disagree with the reviewer consensus.

That being said, I think the fact that the stack runs layerwise, rather than lengthwise, does call into question the analogy with push-down automata.  The resulting transformer still can't handle arbitrary recursion depth, and it's not actually clear what the model is actually using the stack for... it can't use the stack in the same way that an automata would.  The reviewers agree:

Reviewer nPSE writes:  "However, the "bottom to top" direction does call into question the motivation of the paper, which is modeling the kinds of dependencies found in context-free languages. There isn't, as far as I can tell, any reason to believe that this model would be better at these kinds of dependencies than a standard transformer."

Reviewer j9mp writes: "To some extent it remains unknown if the hidden stack is indeed a stack. For example, can we restore and check the grammar bottom-up, from the trained dataset?"
Decision: Accept (poster)