StagFormer: A Staggered Transformer for Decoding Layers in Parallel

• Limited exploration of p>2 setting: Since this is the first paper introducing this idea, we wanted to focus more on simpler settings where we found the idea to give strong benefits. We tried extending the idea naturally to p>2 and found diminishing returns. Indeed, this is to be expected as one can imagine that for a very large p, the staggering basically forces a small Transformer network to predict p tokens into the future which becomes information theoretically impossible as p increases. (For instance, it might involve answering a question before the question is even asked). We do expect the performance benefits to drop off after some small value of p. We believe the exploration of techniques to avoid this dropoff is out of scope of the current paper.
• Mitigating communication cost associated with parallel execution: This is an important issue and is a natural by-product of the using the StagFormer idea to speed up execution. Mitigating this can be achieved by optimized hardware setups which enable fast inter-chip communication. This is outside the scope of the current paper though.
• Additional datasets apart from Pile: We thank the reviewer for raising this point. For language modeling, we believe the Pile is a very comprehensive dataset. In fact, some of the well-known contemporary research papers such as “Mamba: Linear Time Sequence Modeling with Selective State Spaces” rely on pre-training experiments on the Pile. Pre-training experiments on a dataset the size of Pile take a significant amount of compute resources already and it is not easy to scale up to larger datasets on a limited compute.
• Comparison with other methods for efficient Transformer inference: Indeed there is plenty of research on making Transformer inference more efficient including the methods the reviewer pointed out such as speculative decoding, quantization, knowledge distillation among others. We don’t view StagFormer as competing with these other methods, rather StagFormer can be used in conjunction with most of these methods (very naturally with techniques such as knowledge distillation and quantization) to give stronger gains. Due to the tremendous amount of research on methods for optimizing Transformer inference, it is hard to perform a comparative analysis with each of them. We will try to include comparative experiments with a few well-known methods in future drafts.
• Effect of depth on decoding speed and model quality: Increasing the depth of each StagFormer stack would make the model of higher quality but it would decode slower.
• Factors influencing the optimal number of stacks: This is a good question. Given a quality threshold that we want to achieve, we believe a small number of stacks (many times just 2) is optimal as too many stacks can start causing severe quality degradations.
• Extending the staggering idea to encoder-decoder transformers: The staggering idea can be used in the decoder of an encoder-decoder architecture naturally.

• Comparative analysis vs simply increasing width of the model: This is an interesting question. However, there is empirical evidence building every day that depth is crucial for reasoning oriented downstream tasks. Although a wider shallower model might achieve a similar log pplx during training it tends to be less performant on downstream tasks compared to deeper models.
• Technique only useful when a single sequence is being decoded: We would like to emphasize that this is false. Our architecture naturally extends to handle processing for a batch of sequences at once.
• Slower to train: Indeed, this is the case. However, in many applications today, it is perfectly acceptable to have a model that is slower to train but more performant during inference since the training cost is a one-time cost. For instance, take the popular Llama 3 series of models wherein the smallest models were trained for much longer than is suggested by data-scaling laws with the intention that they will be used repeatedly for inference so a one-time higher training cost is easily justified.
• Writing typos/ polishing of the paper: We apologize for the typos and missing references. We will fix all of them and also polish the presentation in the paper to make it more detailed and include better explanations and insights from the results presented in our work.

We thank you for your review and feedback on our paper. We hope to address your concerns below.

• Memory bottlenecks during decoding: Thank you for these questions. We will add a detailed discussion of this topic to clarify some possible misconceptions with our current write-up. We agree that memory access can dominate decoding times. Irrespective of whether we are in a small batch size/context length setting or a large batch size/context length setting, instead of operating a size B Transformer model in parallel on C accelerator chips, we propose serving a size B (approximately) StagFormer model on 2C chips which will reduce the latency of decoding each token without adding additional memory load on any given chip. The use of extra hardware chips allows us to be memory usage neutral with respect to the baseline model. Our simulated latency measurement takes this memory load into account. In particular, your claim that “While StagFormer can conceptually parallelize execution of layers, the associated memory access load cannot be parallelized.” doesn’t hold here since the additional hardware chips help with the extra memory access required for the cross attention KV caches.
• Stagformer variant used in table 2: We apologize for the lack of clarity here. The variant we used here is in fact the Separate Weights - Stagformer 2x18L model We don’t run into the memory issues you mentioned because of the explanation given above. Note that we do take a hit compared to the theoretical latency savings of 50% (and only achieve ~33%) because of the additional cross-attention and a few other engineering overheads.
• Misleading task performance of the Recurrent variant: We apologize for not making this more clear. We will explicitly mention that on scoring tasks recurrent decoding doesn’t make a difference and the performance remains the same since prefill still goes through 2 stacks.
• Incomplete references: We apologize for this oversight. We will fix all typos, missing references and also make a pass over the entire paper to polish the writing and clarity of the results.
• Experimental results on model variants are hard to follow: We apologize for the confusion and will re-organize the sections to make the reading more clear. We will also break down the results into more subsections so as to highlight the main insights within each table more directly.
• Architecture details: A vocab size of 256k although uncommon for models of this size, is not completely unheard of. For instance, Gemma which has models of a similar size used 256k vocab as well. Our 18 layer baseline uses 1536 model dimension, 12288 as hidden dimension and 12 attention heads. We will add these details to the paper.
• Confusion on time complexity of computing cross-attention: We apologize for this confusion. We were referring to the time complexity of attention computation during the prefill step wherein we need to compute contextual embeddings for every token in the sequence. This is quadratic normally and can be made linear when choosing a small window size. Consequently, during decoding, the linear cost of attention during decoding steps drops to constant with a constant window size. We will clarify this in the writeup.

We would like to thank the reviewer for their careful reading of our paper and for the feedback provided. We try to address the concerns raised below.

• Typos/Paper organization and polishing: We apologize for the typos and missing references. We will fix them. We will also take a pass to make the presentation of our main results cleaner and more detailed and polished. Can you elaborate on what you mean by “few results for proof of concept”?
• Table 3: Shared Weights Stagformer (1.6B) outperforming 2.8B Baseline on downstream evaluations: While it appears that the shared weights stagformer outperforms the 2.8B baseline model on some evaluations, it is not a universal trend and not hence inconclusive whether it is better than a 2.8B model. Moreover, in general we don’t expect it to outperform a 2x depth baseline. That is why we didn’t focus too much on the specific numbers on a few evals where the stagfomer model was outperforming the 2x depth baseline.
• Table 1: Strong performance of Stagformer 2.9B model in comparison to Baseline 2.8B model: As you correctly point out, the strong performance of the StagFormer 2.9B model in comparison to the Baseline 2.8B is due to the presence of the cross attention layers. In general, we expect it to roughly match the performance of the baseline, i.e. the extra cross-attention layers help offset the quality loss from staggering. Indeed, this can be seen in Table 1 where the improvements we see in some columns are very minor over the 2.8B Baseline.
• Measurement of the decoding time: We apologize for the lack of details in this section. We will add more discussion in this regard. We measured the decoding time using a setup that simulates a Stagformer with 2 stacks running 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 compared to the other contributor to latency under an optimized hardware setup.
• KV-cache question: Yes for half of the layers, we would need an extra KV-cache for the cross-attention. Note however that in the setup described above where we use double the number of chips as a baseline model this would not increase the per chip memory load.

Thanks for clarifying some misleading points.

1. I clarified my words: "The experiments provided seem limited in scope, hinting at strong performance only in a few simplified settings. To further validate the robustness, a more extensive evaluation and delicate ablation studies are recommend like exploring across various model sizes, architectures.".

2. Thanks.

3. I believe that StagFormer has a great potential by looking at the overall results.

4. Thanks.

5. Thanks.

While I truly believe in the potential of this work and its contribution, I maintain my score as I feel it's not yet ready for publication.