StagFormer: Time Staggering Decoder only Transformers
Authors: Dylan J Cutler
Arun Kandoor
Nishanth Dikkala
Nikunj Saunshi
Xin Wang
Rina Panigrahy
Venue: COLM 2025 (2025)
Keywords: staggered execution
decoder only language models
efficiency
novel architectures
generative models
TL;DR: We propose a novel variant of the Transformer architecture for decoder only language modeling where we break the causal flow of information along the layers of a model by staggering in the time axis.

Abstract: Standard decoding in a Transformer based language model is inherently sequential as we wait for a token’s embedding to pass through all the layers in the network before starting the generation of the next token. In this work, we propose anew architecture StagFormer (Staggered Transformer), which staggered execution along the time axis and thereby enables parallelizing the decoding process along the depth of the model. We achieve this by breaking the dependency of the token representation at time step $i$ in layer $l$ upon the representations of tokens until time step $i$ from layer $l−1$. Instead, we stagger the execution and only allow a dependency on token representations until time step $i−1$. The later sections of the Transformer still get access to the ”rich” representations from the prior section but only from those token positions which are one time step behind. StagFormer allows for different sections of the model to be executed in parallel yielding up to 33% speedup in decoding while being quality neutral. We also explore many natural variants of this idea.  We present how weight-sharing across the different sections being staggered can be more practical in settings with limited memory. We show how one can approximate a recurrent model during inference using such weight-sharing. We explore the efficacy of using a bounded window attention to pass information from one section to another which helps drive further latency gains for some applications. We also explore demonstrate the scalability of the staggering idea over more than 2 sections of the Transformer.

Published as a conference paper at COLM 2025
StagFormer: Time Staggering Decoder-only Transformers
Dylan Cutler
Arun Kandoor
Nishanth Dikkala
Google
Google Research
Google Research
dy.lancutler@google..oom
akandoor@googlan com com
nishanthdegoogle.co com
Nikunj Saunshi
Xin Wang
Rina Panigrahy
Google Research
Google Research
Google Research
nsaunshi@google.com
wanxin@google.com com
rinap@google.c com
Abstract
Decoding in a Transformer based language model is inherently sequential token's embedding needs to pass through all the layers in the network before the generation of the next token can begin. In this work, we propose a new architecture StagFormer (Staggered Transformer), which staggers execution along the sequence axis and thereby enables partial paralleliza- tion of the decoding process along the depth of the model. We achieve this by breaking the dependency of the token representation at time step in layer upon the representations of tokens until time step from layer 1. Instead, we stagger the execution and only allow a dependency on token representations until time step i i- 1. The later sections of the Transformer still get access to the "rich" representations from the prior section but only from those token positions which are one time step behind. StagFormer allows for different sections of the model to be executed in parallel yield- ing speedup decoding while being quality neutral in our simulations. We also explore many natural extensions of this idea. We present how weight-sharing across the different sections being staggered can be more practical in settings with limited memory. We explore the efficacy of using a bounded window attention to pass information from one section to another which helps drive further latency gains for some applications. We also explore the scalability of the staggering idea over more than sections of the Transformer enabling more parallelism. Finally, we show how one can approximate a recurrent model during inference using weight-sharing. This variant can lead to substantial gains in quality for short generations while being neutral in its latency impact.
1 Introduction
The Transformer architecture (Vaswani et al., 2017) has seen tremendous success as the primary backbone for language models (Chowdhery et al., 2022; Hoffmann et al., 2022; Brown et al., 2020). It lends itself particularly well for causal language modeling by allowing efficient, highly parallelized training over large datasets. Moreover, the model can be efficiently partitioned across multiple devices (Pope et al., 2022) enabling model parallelism across machines. However, it is well known that, during inference, decoding from Transformer an inherently sequential task. This task becomes more expensive when trying to decode long sequences due to the cost of attention, which requires computation that scales linearly with respect to sequence length for each additional token.
Consequently, there has been significant body of research which tries to make inference from Transformers more efficient in practice. Speculative decoding, local attention and other efficient attention variants (Tay et al., 2022), KV cache optimizations, blockwise parallel decoding (Stern et al., 2018) etc. are few such works. However, there haven't been many works which try to tackle the sequentiality imposed by the depth of the Transformer. Depth,
1

---

Published as a conference paper at COLM 2025
while known to be essential for the strong performance of Transformers (Raffel et al., 2023; Zhao et al., 2023; Ye et al., 2024), introduces a proportional cost in terms decoding latency. In this work, take look at how can introduce some degree of parallel execution along the depth axis causally trained Transformer language model while decoding.
Noxt lotion
6
T2
Intermedallo
IM Brann
Intermedion Activaions
[1
TI
wnttwcbliggs
1254
[
embodshngs
12:45
t-1
1=1
1
(a) Transformer Timing Diagram
(b) StagFormer Timing Diagram
Figure 1: Depiction of forward pass a standard Transformer compared with that of StagFormer. Note that in StagFrrmer, the data dependency in given time step has been broken for the two stacks T1 and T2.
We introduce StagFormer (Staggered Transformer), novel Transformer variant which breaks the sequential dependency of the upper layers the lower layers by staggering the time dependency of token embeddings passed from the lower layers to the upper layers. In particular, we devise mechanism by which, at time step i, the upper layers of the model use the rich representations of tokens computed by earlier layers only until time step i 1. Note that in a standard Transformer this dependency allowed until time step i. We show how one can train and decode efficiently while matching the standard Transformer's 's quality using our architecture.
We demonstrate the performance of the StagFormer architecture on language modeling on the Pile dataset (Gao et al., 2020). We show that we can get significant latency savings during decode due to parallel execution of different parts of the Transformer stack while being neutral in quality. Finally, we also explore some natural variants and generalizations of the StagFormer architecture and demonstrate their efficacy well for language modeling. We include thorough downstream task evaluation for our trained language models across wide suite of tasks involving summarization, reasoning, and factuality among others.
1.1 Related Work
The Transformer was originally proposed in the seminal work of Vaswani al. (2017). Decoder-only language modeling using the Transformer originally proposed by Radford (2018) and has since become a standard backbone to many frontier language models today,
There has been enormous body of research dedicated towards making Transformer training or inference more efficient (Tay et al., 2022). These involve approaches which focus pre-training such distillation (Xu et al., 2024), layer stacking (Panigrahi et al., 2024), Alternating-updaar (Baykal et al., 2024), Matryoshka Transformer (Kusupati et al., 2022) among others. Quantization (Xiao et al., 2023) has been another widely successful approach
2

---

Published as a conference paper at COLM 2025
Table StagFormer vs Standard Transformer: Pretrained on the Pile dataset for 300B tokens. We see that StagFormer is able to match or outperform the 36 layer baseline model. Even though the time staggering should give some quality loss, this is being offset by the additional cross-attention weights.
Model
File Pplx.
HellaSwag
ARC-E
ARC-C
WinoGrande
SuperGLUE
SQuADv2
GEM-XSum
rouge2
Avg.
Baseline (181.)
1.68 params
4.026
49.8
60.1
31.8
53.4
59.3
31.8
41.0
Baseline (361.) 2,88 params
3,780
33,3
66.7
62.1
45.0
StagTormee p (2x 181. Stacks)
2.98 params
61.3
47.0
at speeding up inference of language models. There have been other approaches specifically focused on improving the decoding speed from language models such as speculative decoding and related works (Leviathan et al., 2023; Sun et al., 2024; Santilli et al., 2023).
There has also been huge body of work focusing making the self-attention more efficient. Some of these works have introduced the idea of introducing form of recur- rence mechanism into models, such Transformer-XI. and State Space Models (SSMs) like Mamba (Dai et al., 2019; Gu et al., 2022; Gu & Dao, 2024). Recurree Transform- ers use cross-attention introduce per-layer recurrence mechanism into Transformer networks (Hutchins et al., 2022).
More closely related to our effort are works such Medusa (Cai al., 2024) which uses parallel heads to decode multiple tokens ahead at once, Staircase Attention (Ju et al., 2022) which uses similar idea of staggering attention window context as we advance deeper into the Transformer stack. However, they mainly explore variant of the idea which allows one to bring in the benefits of RNNs rather than efficiency gains, our main focus here.
Our seared-weight variant of StagFormer, introduced in Section 4.1, is closely related to the idea of looped Transformer, where the hidden activation signals are sent through the layers of the network multiple times (Dehghani et al., 2018; Giannou et al., 2023; Gatmiry et al., 2024). Part the intuition behind looping is that the lower layers a network can reuse the more-information-rich activations from layers later in the same network in the next iteration of the loop to create higher quality representations.
key difference in our method from looping is that it breaks the strict data fata-dependency on each prior loop, allowing for parallel execution of different passes through the network.
2 Staggered Transformers (StagFormer)
In this section we describe Staggered Transformer (StagFormer) architecture. We begin with a brief background on a decoder-only language models based on the standard Transformer, the backbone for most state-of-the-art language models today.
Language Modeling with the Transformer A Transformer of depth L is a sequence-to- sequence model which takes in token sequence of length N and generates an output sequence of length N. The tokens are each first mapped to a d-dimensional representa- tion using an embedding layer. Positional information may also be combined into the embedding at this stage. Denote the token embeddings so obtained by t th....N Then, these representations are progressively modified by applying a sequence of Transformer lay-
ers, Li...., Lp: Rd Rd iteratively: t....N L; ( ( for E {1,..., Each layer Lj consists of two main operations: (a) self-attention which combines information across the different token embeddings and (b) a feed-forward network which modifies each individual token embedding. The two main operations applied along with residual connections and layer normalization. There may additionally be position encoding incorporated into the embedding during self-attention stage well.
3

---

Published as a conference paper at COLM 2025
To use Transformers as decoder-only language models, popular paradigm is that of causal language modeling, Given a train dataset of examples each of which is a sequence of tokens of length N, causal language modeling simultaneously trains to minimize N loss terms each sequence. These loss terms minimize the cross-entropy loss of the model's prediction for token t using the prefix t!...-1. During training, all N of these loss terms can be evaluated in parallel with the use of causal masking. During decoding, the model iteratively generates one new token at a time by passing token t; through the layers sequentially to obtain 111- This means that growing the network depth incurs a linear cost on the time taken to decode the next token during inference. However, there is ample evidence that depth is crucial for good quality models (Devlin et al., 2019; Raffel et al., 2023). There is fundamentally no way to avoid this cost in a Transformer, since every token relies on the completed predictions of every other prior token.
StagFormer StagFormer introduces way to break the sequential dependency of layers within Transformer network and still be able to perform efficient and performant causal language modeling. We first partition our é layers into p sub-networks we call stacks. For ease of exposition we will first focus on the simplest case p = 2. Let L/2. StagFormer allows for execution of the stacks of layers I, and h + 1,.., E in parallel in given time step by staggering the dependency between t, and t 1+11-
In particular, we compute f/,11 function of the original token sequence t and the ht layer representations taken until time step i - 1: t...//111 Crucially we exclude a dependency on t/. This allows the lower half of layers to begin computing the predictions for the next token in the sequence, t while the upper layers in the network are finishing computing the final prediction for position f--
We realize this by passing the original token embedding, t as input to the second half of the layers, Lu+1e-.-,, Le, and by augmenting these layers with cross attention to the final activations of the first half of the network on the prior tokens, t f/=1, when computing the final predictions for the next token after position i. Thus #+1 does not actually depend on the prior layers' representation the token, ti, it is a function of the initial token embedding, fo, and cross-attends to the previous layers' representations of only past tokens, ,
Figure 1 shows a timing diagram of how decoding works in StagFormer. The parallel execution of the two stacks shown more clearly in Figure 2. During training, to faithfully simulate StagFormer's 's decoding, we sequentially pass our token sequence over the two stacks of layers where we allow the second stack to cross-attend to the outputs of the first stack with masking such that at position i we can only cross-attend to the first i 1 outputs from the first stack. This completes a description of how we can train and decode using StagFormer. The full algorithm is given is Algorithm 1.
This idea can be generalized to partitions of the layers by having each new partition stagger additional time-step. We call this technique staggering the Transformer network over stacks. A full description of this generalization is presented in Section 4.3.
Quantifying the Latency Benefits of StagFormer. The main advantage of StagFormer is the potential to save latency during decoding by executing stacks in parallel. This can be realized efficiently on today's hardware accelerators such as TPUs and GPUs. Staggering the dependency on the processed representations of tokens until time step i between the first and second stacks of StagFormer can, in principle, lead to a decrease in quality of the model. However, the additional cross-attention parameters in the second stack help ameliorate this decline.
How much compute savings does StagFormer afford us? We first perform a theoretical analysis to understand the optimal savings can hope to achieve.
Assume our baseline transformer has depth l which is evenly divisible by 2. If the number of FLOPS in the embedding/ uuriiidddnnnnb/iinngggggg layers are each and the number of MLP and
4

---

Published as a conference paper at COLM 2025
1' 3
1 4
T1
T2
1
2
1 4
1
t=1
t=2
Figure 2:
Depiction of the parallel execution of stacks T1 and T2 2-stack StagFormer. In a given time step, both T1 and T2 can run in parallel: T1 producing the intermediate activation to be used in the next time step and T2 producing the output token for the next time step.
Algorithm 1 StagFormer algorithm
Input: t E Rd Token embeddings for positions 1,... in the input sequence.
Output: t E Rd The predicted token embedding for position i+ in the input sequence where l is the total number of Transformer layers in the network.
1: First pass : for each layer Li, Ln where = f/2 compute t/: = Li (E())
Each application of Lj using standard Transformer layer with self-attention and feed- forward layers.
Second pass for each layer L441 L, compute = 11j I..... ( 1)...--1)
2:
Where и 0 when = ht + and otherwise.
Where I, is Transformer layer that has additional cross-attention layer between the self-attention and feed-forward layers that uses t). t!....--1 for KV inputs.
3: Return t

---

Published as a conference paper at COLM 2025
attention flops per transformer layer are m, respectively, then the total FLOPs of forward pass through the baseline are 2e {mm (1m++)) a))
Now consider StagFormer of depth & with 2 stacks. The total number of FLOPs are 2e + 36 + a2/ However, due the staggering we can execute the bottom stack and top stack in parallel. In modern accelerator hardware, when two computation graphs can be executed in parallel, the amount of time it takes to complete both of them depends on how busy we can keep the hardware. That is, how close to the arithmetic intensity of the hardware (which roughly measures how peak FLOPs/ sec the hardware supports) the computation graph executes. This requires minimizing the communication time across chips and memory read/ write time as well. In the ideal scenario when we are able to double the number of chips we have, we can achieve latency equivalent that of model executing 2e + [(m + 2 FLOPs.
In Section 3, we train and evaluate StagFormer for language modeling and observe that depth é StagFormer with 2 stacks outperforms depth é regular Transformer (Table 1). At the same time, we measure the potential latency speedup using setup that simulates Stagformer with 2 stacks on twice the number of chips used by a baseline model. While we faithfully account for every segment of the StagFormer model, we ignore the inter- chip communication cost between the first and second stacks of the Stagformer. This communication cost is expected to be minimal in practice under an optimized hardware setup. We compare the time it takes StagFormer to generate text compared to baseline Transformer model in Figure 3. Overall, we see strong performance gains on tasks such as SQuADv2, Lambada and HellaSwag while being neutral with the baseline on some others such as SuperGLUE.
Figure Simulated Latency Benchmarking for a baseline Transformer (dotted) vs a compa- rable quality StagFormer model (solid). We plot the total time it takes to decode up to 2,048 tokens after 8,192 token prefill. We observe that StagFormer is able to decode 2,048 tokens than baseline even with double the effective batch size.
ToitP dea
Dakk
Thme decode (s)
4
2006
Impact of StagFormer on Memory Usage One limitation of StagFormer is that it increases the memory used by Transformer's KV caches during decoding. This due to the additional cross-attention layers that we add to the latter stack of the model which introduce a cross-attention KV cache in addition to the self-attention KV cache for the latter stack. This increases the total KV cache memory usage by precisely 50%. In Section 4.2 discuss variant of StagFormer where we use local attention in order to reduce the increase KV cache size.
In the next section, we describe some extensions of the StagFormer idea which might be more applicable for specific settings.

---

Published as a conference paper at COLM 2025
3 Experiments
In this section, we describe our pre-training downstream evaluation setup used for the different variants of the StagFormer via causal language modeling on the Pile dataset (Gao et 2020). We begin by outlining our experiment setting.
3.1 Experimental Setting
We performed our experiments using a standard Transformer architecture. The model uses vocabulary size of 256,000. The model adds global positional embeddings to initial token embeddings and applies Rotary Positional Embeddings (RoPE) in the attention layers (Su al., 2023). We compare StagFormer an 18 layer baseline model with 1.6 billion parameters as well as a baseline where we double the number of layers, resulting in 2.8 billion parameter model. We pretrained our model on The Pile dataset with global batch size of 1024 and max sequence length of 1280 (Gao et al., 2020). We trained the model for 250, 000 steps or 327 billion tokens which Gu Dao (2024) demonstrated should be enough tokens for the model to develop few-shot learning capabilities.
We evaluate the model's performance on several few-shot learning tasks (Brown et al., 2020). The evaluation benchmarks include HellaSwag, ARC-E/C WinoGrande, SuperGLUE, SQuADv2, and GEM-XSum (Zellers et al., 2019; Ma et al., 2023; Sakaguchi al., 2019; Wang et al., 2020; Rajpurkar et al., 2018).
3.2 Results
We first present latency benchmarking results on accelerator hardware which demonstrate the gains are able to see during decoding with StagFormer compared to quality matched standard Transformers, The analysis is presented in Table 3.
We compare model with double the number of Transformer layers with a sepirate-weights StagFormes model which uses the same number of layers as the original baseline model in each pass. We chose compare StagFormer to a Transformer with double the number of layers to compare the benefits of using staggered passes with adding more layers to the model. See Table 1 in the Introduction for results.
Step
Layer
2,0092
230.000
Shared Sharee-Waigns SragFormer 2 Pasa 18 Layers
1.8619
290,000
30 myer Booosiine
Sepaeste- StagFormer Presa 18 Layara
1.9102
290,000
Figure 4: Plot of the training perplexity loss for the 18 layer baseline (black), 18 layer shared- weights StagFormer (blue), the 36 layer baseline (red), and separate-weights StagFormer with stacks of 18 layers (yellow). We see that shared-weights StagFormer closes much of the gap between the 18 and 36 layer baselines. We also that separate-weights StagFormer is matches and slightly surpasses the 36 layer baseline.
7

---

Published as a conference paper at COLM 2025
4 Extensions of the StagFormer
In this section, we describe certain natural extensions and variants of the StagFormer architecture.
4.1 Shared-Weights StagFormer
A simple variant of the StagFormer architecture is one where the two stacks share all weight parameters (except the cross-attention parameters). This can useful in scenarios where we are bound tightly on memory requirements. Moreover, it allows us to execute both stacks on the same core using batching of the input data.
One limitation of shared-weigee weights StagFormer is that it increases the KV cache size. A 2 stack wheghd- StagFormer will have triple the KV cache size Transformer model: StagFormer will need KV cache for self-attention for each stack and one additional KV cache for the cross attention. In Section 4.2 we discuss a variant of StagFormer which uses local cross-attention to reduce the size of the KV cache. In Section A we discuss variant of shared-weights StagFormer which reduces the KV cache size by simulating recurrent neural network (RNN).
Experimental Results At the 1-3 billion parameter scale, we compare shared-weights StagFormer baseline model with the same number of layers. The results with Shared- Weights StagFormer are presented in Table 2. We remark that 2 stack shared-weights StagFormer with each stack having 18 layers performs significantly better than 18 layer baseline model which has similar number of parameters. Therefore, StagFormer is effective way of boosting the performance given a fixed parameter budget (similar to Dehghani et al. (2018); Lan al. (2019); Saunshi et al. (2025) and other related works).
Table 2: Performance of shared-weight StagFormer pretraining and inference. We see that shared-weights StagFormer is able to close much of the gap between the 18 and 36 layer baselines and more parameter efficient than the 18 layer baseline model.
CEM-XSum
Pplx.
HellaSwag
ARC-C
WinoGrande
SuperGLUE
SQuADv2
rouge2
Avg.
Baseline (181) 1.68 params
4.026
49.8
60.1
31,8
53.4
59.3
31.8
0.9
41.0
Baseline (361. 2.8 params
3,780
66.7
34.6
60.4
62.1
36.3
1.6
45.0
Flamer p Shared Weights I8L.
1,8B params
3,896
51.3
61.7
31.7
57.7
59.5
16.9
2.1
41.8
4.2 StagFormer with Local Cross-Attention
If we want stronger latency savings, further optimization for StagFormer that is simple to implement to use local attention for the cross-attention between passes (Beltagy et al., 2020). Using local attention also reduces the additional memory need to store in the KV cache. We observe that StagFormer still performs well even when using local cruss-attention with relatively small attention window sizes. StagFormer is also capable of giving non- trivial quality when using an attention window size of 1, which converts the application of the cross-attention in layer L, token E to linear function of t
(recall h = L/2).
Experimental Results We run experiments using StagFormer with local cross-attention with both the separate- and sttared-eightt variants. We present results for experiments with local attention using window sizes 512, 128, and 1. We remark that local can be used successfully with both the separate-weight and shared-weights variants. See Table and 5 in the Appendix for the results of our experiments using local attention with and shared-weight respectively.
8

---

Published as a conference paper at COLM 2025
4.3 StagFormer with More Than Two Stacks
A natural extension of StagFormer idea we had touched upon earlier is to have ht be less than f/2 and to stagger over more than 2 stacks through the network. For instance, we could have h = f/3 and stagger the network over 3 stacks. Let be the number of stacks we stagger the network over, then h = t/p. Intuitively, as we increase the number of stacks p, due to progressive staggering, at time step i stack only gets to see tokens until time step p + but needs to produce activations which help predict token i+ 1. Thus the job becomes more difficult to learn p increases, and the depth of each stack reduces which contributes to eventual degradation in quality.
Our early experiments found that model quality suffers when 2. However, we find that we can recover significantly by imploring simple change for StagFormer: rather than using just the final stack's output for computing the final logits, use linear combination of each stack's 's output with learnable coefficients, #,,..., Ap. Algorithm defines separate-weights StagFormer for when p 2 in the Appendix.
We explored the settings of p 3,4 in this work. We find that as we increase p model quality suffers, but there might be ways extend the approach effectively to even larger values of p which we leave for future work. Our experimental results for this variant are shown on Table in the Appendix.
Experimental Results
We also explore another variant of the shared-weights StagFormer one which allows for recurrent inference pathway in Section A in the Appendix. This variant has very efficient decode step, however the forward pass between training and inference differs leading to slight regression on long generation tasks.
Conclusion
We present the StagFormer architecture as way increase the capacity of Transformer models by allowing lower-level layers to attend to the final activations produced by the same or different networks. With separate-weighh can use higher level representations of prior tokens to run data-independent Transformer layers in parallel to process the current token without sacrificing quality.
StagFormer, we demonstrate that we
Future work and limitations We believe the StagFormer idea offers many interesting avenues for future research. For example, training weights StagFormer only approx- imates recurrent inference, since training requires discrete number of passes. Future work could explore how to extend the StagFormer algorithm that either better approximates or fully realizes recurrent decoding with better quality.
Another limitation is that cross-attention incurs additional computational cost in both time and space with respect to the input length. In the separate-weights StagFormer variant, cross attention increases the KV cache size by 50%. Instead of passing the output from the first stack as an extra input attend to, if we could somehow merge it with the token embeddings then we would remove the additional cross-attention cost and KV cache. Such a design would greatly increase efficiency.
The seirhd- variant of StagFormer also increases the KV cache size. Because of the additional pass with self- and cross-attention in the second pass, sharhd- Stagformer triples the KV cache size relative to traditional Transformer with the same number of parameters.
Future work can explore alternate ways to reuse information-rich higher level activations in lower-level layers allow parallel execution of layers in a way that incurs less computa- tional cost than attention and matches a deeper model's quality.
One material challenge of StagFormer's 's parallel execution of layers is that would require communication cost to copy the result from one network over to the other. This prevents one from realizing the full theoretical latency benefit of running the StagFormer towers in parallel. Furthermore, since most models rely on the single program, multiple data (SPMD)

---

Published as a conference paper at COLM 2025
paradigm (Xu et al., 2021), parallel execution of StagFormer stacks would require storing a copy of the token embeddings and final softmax tables in both cores when executing StagFormer stacks. Further work could explore how to extend this algorithm the SPMD paradigm to help realize greater latency benefits when executing Transformer networks in parallel.
References
Cenk Baykal, Dylan Cutler, Nishanth Dikkala, Nikhil Ghosh, Rina Panigrahy, and Xin Wang. Alternating updates for efficient transformers. Advances in Neural Information Processing Systems, 36, 2024.
Iz Beltagy, Matthew F. Peters, and Arman Cohan. Longformer: The long long-document trans- former, 2020. https://///axi wwwp://///bbb/0.....mm 05150.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL https: http:////aarii.ooaaa///////
14165.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason Lee, Deming Chen, and Tri Dao. Medusa: Simple Ilm inference acceleration framework with multiple decoding heads. Xrvivv preprint ar arXiv:2401.10774 2024.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arviv preprint ar Xr21:2204.023111 2022.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhut- dinov. Transformer-xl: Attentive language models beyond fixed-length context, 2019. wrtp:/iiii..o///aaoaa http://rrxiiv.orr//////9// /b////99///1111///////// 02860.
URL http:::////rxiii.rrrgg//e
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser.
Universal transformers. arXiv preprint ar trXi::1807.033 00019, 2018.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019. URL https://arxiv.
org/abs/181e 148800
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling, 2020. URL .0027.
https://arxiv.org/abbs
Khashayar Gatmiry, Nikunj Saunshi, Sashank J. Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi-step gradient descent for in-context learning? In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, Pp. 15130-15152, PMLR, 21-27 Jul 2024. URL https://proceeding wwweddinnnss...........m press//233/ gatmiry24b.u html.
Angeliki Giannou, Shashank Rajput, Jy yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped transformers programmable computers, 2023, URL wrtp/////////...... 13196.
https://arxxv.org//aa
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024. URL https: wwwaaawwwwwoo/b////1//.. 00752.
10

---

Published as a conference paper at COLM 2025
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, 2022. URL /ttpss///axiiv.. wwwp://wwgg/bbb//// org/abs/ 11.00996.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXio preprint ar xrx:::2003.15556 2022.
DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block- recurrent transformers, 2022. URL https://arxix wwtrrriwv...o/////eee..ee org/abs/e 2203. 228220852..
Da Ju, Stephen Roller, Sainbayar Sukhbaatar, and Jason E Weston. Staircase attention for recurrent processing of sequences. Advances in Neural Information Processing Systems, 35: 13203-13213, 2022.
Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ramanujan, William Howard-Snyder Kaifeng Chen, Sham Kakade, Prateek Jain, et al. Matryoshka representation learning. Advances in Neural Information Processing Systems, 35: 50233-30249, 2022.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. arxiv preprint ar xrXit:1909.11944 11942, 2019.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding, 2023. URL httpp::/ //axiv wwwpaawwgwww.aaa....e..ereaa.aee
17192.
Yuhan Ma, Haiqi Jiang, and Chenyou Fan. Sci-cot: Leveraging large language models for enhanced knowledge distillation in small models for scientific qa, 2023. URL https: 04679.
//arxiv.org/aas/////
Abhishek Panigrahi, Nikunj Saunshi, Kaifeng Lyu, Sobhan Miryoosefi, Sashank Reddi, Satyen Kale, and Sanjiv Kumar. Efficient stagewise pretraining via progressive subnet- works. Xio preprint ar arXto:2402.05913 2024.
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonaahan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently
Anselm Levskaya, Jonathan scaling transformer inference, 2022. URL httttt:////xiiiii .ogg/oggggeo.. abs//21111111111///////////////1 221.10102..
Alec Radford.
Improving language understanding by generative pre-training. 2018.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer, 2023. URL
Pranav Rajpurkar, Robin Jia, and Percy Liang, Know what you don't know: Unanswerable 03822.
questions for squad, 2018.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019. URL https://arxivv org/abs/
Andrea Santilli, Silvio Severino, Emilian Postolache, Valentino Maiorca, Michele Mancusi, Riccardo Marin, and Emanuele Rodolà. Accelerating transformer inference for translation via parallel decoding. arviv preprint ar xrxo:22305.10427 2023.
Nikunj Saunshi, Nishanth Dikkala, Zhiyuan Li, Sanjiv Kumar, and Sashank Reddi. Rea- soning with latent thoughts: On the power of looped transformers. arXio preprint arXit:2502.17416, 2025.
Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models, 2018. URL https: htttp::/aaxxv..or/g/eee 11111111111.1.....
11

---

Published as a conference paper at COLM 2025
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023. URL https://arxxv..rrt 0tss/100409980
Ziteng Sun, Ananda Theertha Suresh, Jae Hun Ro, Ahmad Beirami, Himanshu Jain, and Felix Yu. Spectr: Fast speculative decoding via optimal transport. Advances in Neural Information Processing Systems, 36, 2024.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, 2022. URL https: / /arxiv..v. wwtrriiw.////.o.oo..o...e abs/2009 000.2209900027 06732.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30, 2017. URL https: http:::///ooocooooeeeeeeee newringg...nnuppiee nwurips.cc. cc/
Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. Superglue: A stickier benchmark for general- purpose language understanding systems, 2020. URL https://arxii.org/abb///// 00537.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087-38099, PMLR, 2023.
Xiaohan Xu, Ming Li, Chongyang Tao, Tao Shen, Reynold Cheng, Jinyang Li, Can Xu, Dacheng Tao, and Tianyi Zhou A survey on knowledge distillation of large language models. Xiv preprint ar Xiv::4402..13116 2024.
Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni, Ruoming Pang, Noam Shazeer, Shibo Wang, Tao Wang, Yonghui Wu, and Zhifeng Chen. Gspmd: General and scalable parallelization for ml computation graphs, 2021. URL https: .rgg/ 0t5//2105.046666
Tian Ye, Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu. Physics of language models: Part 2.1, grade-school math and the hidden reasoning process. arviv preprint arXio:2407.20311 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence?, 2019. URL httpp://arxiv.crr wrgg 1905 07838.
Haiteng Zhao, Shuming Ma, Dongdong Zhang, Zhi-Hong Deng, and Furu Wei. Are more layers beneficial to graph transformers? In The Eleventh International Conference Learning Representations, 2023. URL
A A Recurrent Approximation to the Shared-Weights StagFormer
Consider the shared-weights StagFormer with = 2. It suggests the following variant which we will see can have many benefits. We perform the first pass prefill on the input tokens as before storing the final activations. During decoding, rather than having the network process each token twice, we only do a single pass. When doing so, the network cross-attends to the final activations of all prior tokens including the prefill tokens and generated tokens so far.
This method of decoding resembles a recurrent neural network (RNN) where the final activations of prior tokens are the RNN's hidden state and cross-attention serves as gating mechanism while processing the current token. Though one key difference is that traditional RNN's hidden state remains a fixed size, but shared-weights StagFormer grows the size of its "hidden state" of upper layer activations as the sequence length grows.
12

---

Published as a conference paper at COLM 2025
This method reduces the size of the additional KV cache needed for StagFormer. Recall that a pass sheihd-weights StagFormer increases the KV cache size because it needs additional caching for both the self- and cross-attention layers. The recurrent StagFormer variant, the other hand, only needs to add additional KV cache for cross-attention and does not require any additional caching for the self-attention layers.
We show that it is possible to use wheged-weights StagFormer for recurrent decoding using this scheme, even when the model is trained using two separate passes. However, find that the generated text's quality is not as good as when we process decode new tokens the original way, with two networks running in parallel.
2 3
T (without cross-attention)
1
1''
1
1 3
2
t=1
t=
Prefill
Decode
Figure 5: Timing Diagram of Prefill Decode steps for Recurrent Inference with Shared- Weights StagFormer. During prefill, the Transformer is run without cross-attention and during decode it is run with cross-attention.
Table 3:
Performance of Shared-Weight StagFormer pretraining and recurrent inference using Shared- Weight StagFormer. We see that the recurrent variant regresses in sampling tasks like SQuADv2 and GEM-XSum.
CEM-XSum rouge2
Pplx.
HellaSwag
WinoGrande
SuperGLUE
SQuADv2
Avg.
Baseline (18L) 1.68 params
4.026
49.8
60.1
31,8
53.4
59.3
31.8
0.9
41.0
Raseline (361.) 2.8 params
3,780
66.7
34.6
60.4
62.1
36.3
1.6
45.0
StagFormer p Shared-Weights 18L. Two-Netwarks 1.8B params
3,896
54.3
61.7
31.7
57.7
59.5
46.9
2.1
448
StagFormrr
Weights Recurrent 1,88 params
3,896
54.3
61.7
31.7
57.7
59.5
42
0.4
43.9
Recurrent StagFormer with More Than Two Passes One can also increase the number of staggered passes with shared- StagFormer when training to better approximate recurrent decoding. Since the Transformer layer weights are shared between passes, shared- weights StagFormer would process the same input multiple times, cross-attending to prior tokens' final activations from prior passes.
For wharhd- StagFormer, match training during prefill and run all p stacks, and then switch to recurrent inference for decoding. Note that for = 4, some evaluation tasks failed due to memory constraints. We find that increasing p surprisingly has negative
13

---

Published as a conference paper at COLM 2025
impact on model quality. See Table 7 in the Appendix for results using p 2,3,4 for the recurrent variant of shared- StagFormer.
B Additional Details on StagFormer Extensions and Experiments
We begin by presenting the full algorithm for the shared-weights variant of the StagFormer in Algorithm 2.
Algorithm 2 Shared- weights StagFormer algorithm
Input: tg....,4 E Rd Token embeddings for positions i, in the input sequence. Output: tj E IR4 The predicted token embedding for position i i in the input sequence where is the total number of Transformer layers in the network.
i
1: First pass: for each layer Li, L1 compute t/ L; ((-r).
Each application of Lj using standard Transformer layer with self-attention and feed- forward layers.
Second pass for each layer Li,..., L, compute tj L (Landd A.rnd
(f-1" L AL 1).
2:
Where L has an additional cross-attention layer between the self-attention and feed- forward layers to the Transformer layers in the first pass that uses t/m..-1 for KV inputs. Return f..
3:
Next, we present how with small modification to Algorithm 2 we can approximate recurrent model in Algorithm 3.
Algorithm 3 Recurrent Decoding using Shared- StagFormer
Input: t .., Hj E Rd Token embeddings for positions 1, ., i in the input sequence. Output: tr E Rd The predicted token embedding for position + in the input sequence where L is the total number of Transformer layers in the network.
Prefill Use the shared-weights StagFormer algorithm to process the prefix (Algo- rithm 2).
2: Decoding: : for each layer I,,...,, L compute tj = ((-T(())))))
Where I has an additional cross-attention layer between the self-attention and feed- forward layers to the Transformer layers in the first pass that uses E!...! for KV inputs.
The rest of the parameters in L/ are the same those in Lj used for the prefill.
3: Return t
The last algorithm present is the one for extending the StagFormer idea to more than stacks (Algorithm 4).
We also present experimental results for both the separate- and shared- variants StagFormer with local cross-attention.
Finally, present experimental results for using StagFormer with more than 2 passes, or in other words, when > 2.
14

---

Published as a conference paper at COLM 2025
Algorithm Separate- StagFormer p algorithm
Input: fgr...,tt E R²: Token embeddings for positions 1,..., in the input sequence.
Output: tj E R The predicted token embedding for position + in the input sequence where the total number of Transformer layers the network.
1: First pass: for each layer L,, Ln where h p compute jj L, (+=T)
Each application of L, using standard Transformer layer with self-attention and feed- forward layers.
2: Subsequent passes for each E (,...., p} do:
L compute = L1 (t) ( (k 1))
for each layer in (k-
Where и when j It (k - + and otherwise.
Where L is Transformer layer that has an additional cross-attention layer between the
self-attention and feed-forward layers that uses ( for KV inputs.
(.-k-1)
F Return Lax '
k Where each ag is learnable scalar.
Table 4: Performance of separate-weights StagFormer pretraining and eval tasks with local cross-attention. We see that StagForme, is able to maintain its quality when the window size is 512. However, we see when reduce the cross-attention window size further that quality begins to degrade.
Model
Pile Pplx.
HellaSwag
ARC-E
ARC-C
WinoGrande
SuperGL
SQuADvz
GEM-XSum
Avg.
rouge2
Baseline
Layers (36L)
2.8B params
3,780
33.3
66.7
34.6
60.4
62.1
36.3
1.6
45.0
StagFormer
Soparate Weights
Window 512 299 params
3.767
38.6
68.2
36.9
61.8
63.3
41.5
1.9
47.5
StagFormer
Separate-Weights Window 128 params
3.797
SlagFormer
Seperate-Weightts Window 298 params
38818
33.3
30.9
25.3
51.2
45.6
26.6
Table 5: Performance of shirhd- StagFormer on pretraining and eval tasks with local cross-attention. We see that shared-wert thights StagFormer is able to maintain their quality with cross-attention window sizes of 512 and 128. However, we see that when the window size is that quality begins to suffer.
Model
Pile Pplx.
HellaSwag
ARC-E
ARC-C
WinoGrande
SuperGLUE
SQuADv2
GEM-XSum
Avg.
rouge2
1.68 params
4026
49.8
60.1
31.8
53.4
59.3
31.8
0.9
41.0
StagFormer
Shared-Weights Window 512
1.8 params
3,908
35.7
64.9
33.9
59.4
60.1
39.4
1,6
45.0
SlagFormer
Shared- Weights
Window 128
1.88 params
3,929
56.4
64.9
34
59.4
40.3
1.8
StagFormer
Shored- Shored-Weights
Window
1.8 params
3.951
46.8
56.5
29.4
58.5
58
34.8
0.6
40.7
15

---

Published as a conference paper at COLM 2025
Table 6: Performance of StagFormer pretraining and a subset of evaluation tasks when p 2. In this case, we need the output of StagFormer be a linear combination of each stack's output in order to not degrade quality. For 3 we see the model quality does not degrade much except for significant regression SQuADv2.
Model
Train Pplx.
HellaSwag
ARC-E
ARC-C
WinoGrande SuperCLUE
SQaADv2
Avg.
rouge2
Baseline 18I.
166 params
4.026
49.8
60.1
31.8
33.4
89.3
31.8
0.9
41.0
Baseline
2x Layers (361.)
2,8B params
3,780
53.3
66.7
34.6
60.4
62.1
36.3
1.6
45.0
StagFormer
Separate-Weights (3x 12L)
3.0B params
3.766
52.9
52.7
29.1
55.2
60
13.7
37.8
StagFormer P Separate-Weights (4 9I.)
3.0B params
3.797
51.3
59.3
33.1
41.2
Table 7: Performance of Shared-Weights on pretraining and recurrent inference StagFormer subset of evaluation tasks when > 2. For 4 some evaluation tasks crashed due to memory constraints. We find that as we increase p, the model does not better simulate an RNN and quality degrades.
Model
Train Pplx.
HellaSwag
ARC-E
ARC-C
WinoGrande
SuperCLUE
SQuADv2
GEM-XSum
Avg.
rouge2
Baseline 18L
1.68 params
4.026
49.8
60.1
31.8
53.4
59.3
31.8
0.9
41.0
StagFormer p 2
Shared- Weighhs 18L. Recurrent 1.8B params
3,896
5433
61.7
31.7
57.7
39.5
42
43.9
StagFormer
Shared-Weights 1SL. Recurrent 1.8B params
Stag Farmerrer y Shared-Weiihtt 18L Recurrent 1.88 params
3,870
46.6
51.9
0.6
26.0
16


Reviews:
[{"review_id": "RVEwIL8XIj", "reviewer": "Reviewer_UGau", "summary": "This paper proposes a staggered layer formulation for a transformer, specially applicable to causal autoregressive decoders,  that allows to break the causal dependency between chunks of the model, allowing their parallel execution in a decoding phase while introducing little to no overhead in training and prefill phase. Experimental results show competitive performance in quality and better generation speed.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* Novel transformer variant with enhanced decoding parallelization at no quality cost and performance improvements in inference. \n* Theoretically sound motivation with good results in practice\n* Sufficent experimental results to back claims and sufficiently detailed description to allow independent reproduction.", "reasons_to_reject": "* Some portions of the manuscript seem incomplete (ex. \"TODO\" annotations)\n* Lack of considerations for real world deployment of the proposed decoding acceleration (batched inference vs single sequence)\n* Explanation of theoretical inference speedup could be more detailed, without over-reliance on computational experiments to prove a speedup (on TPU, the operations are compiled from jax and it is unclear what computational graph (ex. loop unrolling, etc) is actually happening)", "questions_to_authors": "* The Stagformer models seem to have a 0.1B params extra for each \"P-1\". That is, 2.9B for p=2, 3.0B for p=3. Did you try adding a few extra params to the baseline (ex. +1 transformer block) to compensate for this?\n\n* The benefit of stagformer seems to come mostrly from parallelization of decoding. However, in many real world scenarios, inference can be batched to increase efficiency, and under that setting, the model becomes memory bound. It would be interesting to benchmark the model in such cases (batched inference) and measure the relative performance. There might still be benefits due to reused KV cache, but gains might diminish. Did you consider this?\n\n* A clarification could be added in section 3.2 for the RNN parallelism: the key defining feature of an RNN state is its fixed (or, in some cases, bounded) size. A transformer is usually not referred to as an RNN because the state grows with the sequence. Could you please distinguish this case when drawing the analogy?", "extra_scores": {}}, {"review_id": "9en5qjKUx6", "reviewer": "Reviewer_n1xq", "summary": "This paper introduces StagFormer, a Transformer variant with a direct and simple yet original idea for reducing decoding latency by staggering execution across model depth. This ambitious approach shows promising experimental results in latency reduction, and its \"Recurrent Approximation\" variant offers insightful avenues. However, the significance of these findings is somewhat tempered by relatively thin experimental validation across broader conditions. Additionally, the clarity of its core mechanism is impacted by a confusing Figure 1. While the concept is of good quality and addresses an important problem, the work would benefit from more robust empirical backing and clearer visual explanations to fully realize its potential impact.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "mentioned in Summary", "reasons_to_reject": "1. The figure 1 is confusing: My understanding is that at t=2, the prediction of the fifth token still depends on the output of the fourth token at t=1. If so, there is no latency improvement. figure 2 is much more clear, I think these two figures are duplicate.\n2. The experiments are too thin: I suggest adding some comparisons with speculative sampling.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "185jyHIsww", "reviewer": "Reviewer_TYAW", "summary": "This paper proposes StagFormer, a method to improve the decoding efficiency through introducing some degrees of parallelization. In StagFormer, transformers can be divided into 2 stacks with different layers (lower half and upper half, though weights can be shared). The representations from the lower half can be passed into the upper half by skipping time steps, therefore achieving multi-step decoding. Experiments show that StagFormer can improve efficiency while maintaining the same performance as vanilla transformers. The authors also discuss many different implementations of StagFormer.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The proposed StagFormer can be considered novel, building on existing works. The recurrent formulation is particularly interesting and can inspire future works.\n2. The improvement over vanilla transformers seems pretty well established, given the experiments.", "reasons_to_reject": "1. Lack of comparison. The authors only compare the proposed method with vanilla transformers. Yet there are already many improvements over the vanilla transformers. It could help the readers better contextualize the results with some well-established baselines, such as Medusa and speculative decoding. \n2. Result presentations. There are some issues with the result presentations.\n    - The main contribution of this paper is an efficient decoding method. Yet, the efficiency performance is only presented in Table 2 for 1024 tokens for the separate-weight variant. It will be more helpful if the efficiency can be presented across sequence length? \n    - The authors mention that StagFormer without local attention can increase KV cache. How much more compared to a vanilla performance? Quantifying this trade-off can help better illustrate the performance of StagFormer. \n    - The labels for StagFormer variants seem a bit inconsistent.  In the Appendix, variants are labeled either as shared-weights or separate-weights. In Table 1, I am assuming that the StagFormer is a separate-weight variant, though it is not labeled explicitly.", "questions_to_authors": "1. The authors only considers l/2 case where the upper half and the lower half are equal in weights. I know this facilitates weight sharing. If the two halves are unequal (like a heavy lower half and a lightweight upper half), how will the performance change? Medusa only uses very lightweight decoding heads at the last layers. It seems that equal weights are not always necessary.", "extra_scores": {}}]

Author rebuttal:

Meta-review: StagFormer introduces an interesting staggered‐layer transformer variant that breaks causal dependencies between layer chunks to enable parallel decoding with minimal training or prefill overhead. The paper is technically sound and the proposed architecture is novel.

Reviewers UGau and TYAW both highlight its theoretically sound motivation and demonstrate strong empirical speedups and quality retention compared to vanilla transformers, with UGau praising its novel parallelization and TYAW confirming its well‐established inference improvements and clear recurrent formulation.

The paper also has several areas needing attention. UGau highlights the lack of clarity of writing, and the absence of real‐world deployment analysis—particularly under batched inference—and requests a more detailed theoretical account of the promised speedup and an equitable, parameter‐matched baseline comparison. TYAW recommends adding established decoding acceleration baselines like Medusa and speculative decoding, broadening latency studies across sequence lengths, explicitly labeling StagFormer variants, and quantifying the KV‐cache overhead. Addressing these points will greatly bolster the paper’s clarity and practical relevance.

Overall, I think it's a good contribution to the community of improving the efficiency of the current LLM paradigm.
Decision: Accept