StagFormer:  A Staggered Transformer for Decoding Layers in Parallel

Authors: Dylan J Cutler, Arun Kandoor, Nishanth Dikkala, Xin Wang, Nikunj Saunshi, Rina Panigrahy
Primary area: foundation or frontier models, including LLMs
Keywords: decoder only language models, transformers, staggered execution, pipelining, parallel decoding, efficiency

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.

Decision: reject
Reviewer scores: 3, 3, 3, 5

Reviews and discussion:
<thread>
creation date:
1730595831791

thread_type:
official_review

<thread_object>
creator: reviewer_apqd

<content>
summary:
This paper proposes a novel transformer architecture that effectively reduces the number of sequential steps (layers) during the decoding process by staggering the computation across different time-steps. This allows for improved parallelism during decoding individual sequences, providing speedups during inference.

strengths:
- staggered computation leads to significant improvements in per-time-step decoding speeds while slightly improving performance
- provides results and analysis of different variants of staggered transformers that further explores the architecture's efficacy

weaknesses:
- Biggest critique is that it lacks comparative analysis of staggering computation vs. simply increasing the width of the model and lowering the number of layers, as this increases per layer parallelism while decreasing the number of layers leading to a similar improvement in decoding speed.
- This technique is possibly only useful for speeding up decoding when only a single sequence is being decoded. A non-staggered model could in theory process twice the batch size as it has half the parallelism (and hence half the per layer memory requirement) of a model staggered with p=2. 
- StagFormer is possibly slower to train (as inferred by its slower pre-filling speed)
- Paper could be further refined (minor critique): 
    - Some references are configured incorrectly (Table ?? in page 5, "TODO" in page 8)
    - Plots have unnecessary information (Figure 4 doesn't need texts like /1/summarize/train)

questions:
Addressing the weaknesses outlined above would improve the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730641449042

thread_type:
official_review

<thread_object>
creator: reviewer_tmdv

<content>
summary:
This paper proposes a novel Transformer architecture called StagFormer designed to improve the efficiency of decoding in Transformer-based language models by enabling the parallel execution of layers along the depth axis.

strengths:
1. StagFormer introduces a unique method to break the sequential dependency of layers in Transformers, enabling parallel execution.
2. Experiments demonstrate significant latency reduction while maintaining or even exceeding the quality of a standard Transformer.
3. The paper investigates different StagFormer variants, offering flexibility and adaptability to various scenarios and resource constraints.
4. The paper effectively explains the StagFormer concept and its variants, supported by clear diagrams and algorithms.

weaknesses:
1. Limited exploration of p > 2. While the paper explores StagFormer with more than two stacks, it acknowledges performance degradation and the need for further research in this area.
2. The paper mentions the communication cost associated with parallel execution but doesn't offer concrete solutions to mitigate it.
3. While the Pile dataset is comprehensive, evaluating on additional datasets would strengthen the generalizability of the findings.
4. Comparing StagFormer with other methods for efficient Transformer inference, such as speculative decoding, would provide a more comprehensive perspective.

questions:
1. How does varying the depth of individual stacks in StagFormer affect the trade-off between decoding speed and model quality?
2. What factors determine the optimal number of stacks for a given application, balancing computational efficiency and performance?
3. Could the staggering concept be extended to encoder-decoder Transformers, like those used in machine translation?
4. How well could StagFormer be combined with other techniques, like quantization or knowledge distillation, to further enhance decoding efficiency?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730651863158

thread_type:
official_review

<thread_object>
creator: reviewer_os5p

<content>
summary:
The authors present the Staggered Transformer (StagFormer) and its variants which relieve sequential dependancies in the decoding pipeline to enable higher levels of parallel execution.

Consider a transformer with two stacks of layers, A (bottom half) and B (upper half). In vanilla transformers, the input token embedding is passed to stack A. Then, the output of stack A is passed to stack B. All layers apply self-attention on outputs of the previous layer.

In the baseline StagFormer (`Separate-Weights`), stack A is the same. However, stack B takes in the input token embedding rather than the output of stack A.
To supplement this, stack B applies cross-attention on the final outputs of stack A, up until the previous token. In other words, stack B cross-attends to the outputs of *all previous input tokens* from stack A, instead of directly inputting that of the *current* input token. This relieves the dependency of stack B on stack A, within a single decoding step, thus both A and B can be computed simultaneously.

The authors investigate many variants of this design:
1. `Shared-Weights`: this is where stack A and stack B share the same model parameters (excluding the cross-attention layers which are unique to stack B).
2. `Recurrent, Shared-Weights`: this is a unique decoding method for the `Shared-Weights` trained model. In `Shared-Weights` stack A and B are identical, except that stack B applies cross-attention to outputs from stack A. Essentially, the shared stack S (= A = B) is first forwarded without cross-attention, and then forwarded a second time *with* cross-attention, attending to outputs from the first forward pass. The `Recurrent` setting refers to that where the first forward pass is skipped, and cross-attention in the second pass attends to outputs of the "second" pass from the previous decoding step.
3. `p > 2`: this is where more than two stacks are considered.

When compared to vanilla transformers pretrained from scratch, StagFormers show various advantages, mainly:
- `Shared-Weights 2x18L`: StagFormer outperforms the vanilla 18L baseline (with roughly same parameters) in both perplexity and average task performance. Using recurrent decoding (roughly matching 18L baseline computation), average task performance lies between the two. StagFormer underperforms the vanilla 36L baseline with roughly same computation in perplexity, but performs comparably on tasks.
- `Separate-Weights 2x18L`: StagFormer outperforms the vanilla 36L baseline (with roughly same parameters and compute) in both perplexity and task performance.

strengths:
1. The idea and architecture design are very novel
1. The authors propose numerous variants which showcase the potential extension of the idea across various axes–parallel execution, weight sharing, recurrent computation.
1. The architecture shows clear advantages over vanilla transformers across its variants
1. The writing is easy to follow and visual depiction of the architecture and its variants are superb.

weaknesses:
1. **Memory  bottlenecks during decoding may hinder benefits of parallel execution, which is not discussed**: LM decoding is typically bottlenecked by memory rather than compute (see references below). When batch size x context length is small, memory access is dominated by model parameter access. Otherwise, memory access is dominated by KV cache access. While StagFormer can conceptually *parallelize* execution of layers, the associated memory access load cannot be parallelized. In fact, the cross-attention layer will add additional KV cache access overhead. These are critical to assessing the actual wallclock time benefits of decoding with StagFormers compared to vanilla transformers, but is not discussed.
    1. Different variants of StagFormers will have different memory bottlenecks. Examples:
        1. All variants: cross-attention is added in half of layers. Therefore, the overall KV cache access overhead will increase by 50% (relative to that of self-attention, used in all layers). This will have a larger effect on decoding time as batch size x sequence length becomes large.
        1. `Separate-Weights`: both stacks can be executed in parallel, but the memory load is identical as the parameters of both stacks must be retrieved from memory. This means that wall-clock time should typically be identical to vanilla transformers, as decoding is bottlenecked by memory access. `Shared-Weights` can solve this issue.
    1. **It is unclear which StagFormer variant is used in Table 2, raising questions on the performance vs latency comparison**: While Table 2 states that a "comparable quality StagFormer" is 33% faster than baseline transformer during decoding, the exact variant is unclear. Given the reasons above, it seems likely that this is the `Shared-Weights 2x18L` variant. While its average task performance is comparable to baseline 36L, its PPL is in the middle of that between vanilla 18L and 36L. It would be misleading to describe this variant as "comparable quality" to vanilla 36L.
    1. **Missing comparison of performance vs latency across model variants**: Expanding on the point above, a comparison of prefill/decode time across model variants will provide a clear picture on the performance vs latency benefits of each model variant. This could take the form of a single table that lists the PPL, task performance, and prefill/decode time for each model. In the case of  `p > 2, Shared-Weight` variants, I believe this may actually reveal some advantages in terms of latency.
    1. **The additional KV cache overhead of cross attention may slow down decoding for longer contexts**: Since KV cache overhead is quadratic to context length, the decode time advantages as shown in Table 2 will likely diminish with longer contexts, especially in batch decoding. Given the relatively short context length of 1024 tokens considered in this study, compared to modern LLMs with 8K+ context, measurement on longer contexts and larger batch sizes can help gauge the potential of the architecture.
1. **Misleading task performance of `Recurrent` variant**: In Table 3 (for example), the performance of various tasks are identical between the `Shared-Weights 18L` model and its `Recurrent` counterpart. This is likely because the tasks are measured in a teacher-forcing setting, where the outputs of the prefill stage are used for evaluation. This does not represent the task performance of the `Recurrent` setting, as recurrence is only applied to decoding, as explained in Section 3.2.
1. **Experimental results on model variants are hard to follow**: The organization of the results section could be improved to make the comparison between different model variants more clear.
    1. Within tables, variations could be better indicated with separate columns, task names could be shortened for space, latency metrics could be included, etc.
    1. Results on different variants are presented in multiple tables without a clear organization.
1. **Incomplete writing**: "(TODO)" in Line 385, the reference error "??" in Line 267, and numerous typos suggest that this is an incomplete manuscript that is not ready for review.

References on memory bottlenecks during inference
- [Efficiently Scaling Transformer Inference](https://arxiv.org/abs/2211.05102)
- [LLM Inference Unveiled: Survey and Roofline Model Insights](https://arxiv.org/abs/2402.16363v4)
- [Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve](https://arxiv.org/abs/2403.02310)
- [Block Transformer: Global-to-Local Language Modeling for Fast Inference](https://arxiv.org/abs/2406.02657)

questions:
1. Can you describe the architecture shape (vocab size, qkv heads, embedding dimensions) and its justification? The vocab size of 256K is quite high for models of this size.
1. In Lines ~499-501, the authors mention that cross-attention is linear to input length instead of quadratic with window size 1. Isn't it linear with any fixed window size? Considering that the cost of attention mainly stems from KV cache IO during decoding, I think the constant factor with a window size as small as 128 makes the cost of cross-attention negligible compared to self-attention (especially when expanding to modern context lengths of 8K or more).
    1. However, the *increase* in performance when going from full cross-attention (1024) to windowed attention with window size 512 and 128 is strange. Can the authors justify this increase in performance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730810878741

thread_type:
official_review

<thread_object>
creator: reviewer_vlv7

<content>
summary:
This paper proposed a new architecture StagFormer, which stagger the time dependency between the lower and upper layers. The overall design seems a little non-intuitive, but has a lot of potential for throughput and performance. For example, parameter sharing or local cross-attention could yield better throughput.

strengths:
- StagFormer architecture is interesting, and has very good potential for both performance and throughput.
- The idea of parameter sharing and recurrent decoding looks good.

weaknesses:
- I like the concept and potential of this paper, but I believe that this paper is not well-organized, and looks like unfinished work yet. For example, there is missing reference in L.267 (I guess this refers to Table 3), there are a few results for proof of concept.
- Table 3 is showing few-shot results for gray, blue, red lines in Figure 4 (correct me if I’m wrong.) I wonder why shared-weights StagFormer (blue) outperforms Baseline 2.8B params (red) in some benchmarks, even though it shows higher loss values.
- What makes StagFormer 2.9B to outperform Baseline 2.8B params in Table 1? Is it due to cross-attention in upper layers? This looks somewhat interesting and also confusing because I thought the changed structure (using previous timestep’s intermediate activations) could degrade performance a lot.
- How did the authors measure the decoding time in Table 2? Running separate parameters in parallel is not trivial, I believe. Is it actual time or hypothetical time by assuming parallel execution of them?

questions:
- For KV-caches, the total KV caches are a little increased by the amount of one layer for cross-attention in upper layers, rights?
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors proposed an interesting idea that can possibly result in throughput gains via staggered computation and parameter sharing. However, the manuscript is far from being polished, with missing references, ambiguous variant comparisons, and inadequate consideration of memory overhead that compromises parallel decoding gains. The unclear performance vs. latency tradeoffs, along with limited exploration of longer contexts and direct comparisons to other efficient approaches, further weaken the conclusions. Overall, I think the idea seems interesting, but a much more thorough study is needed. Also, the downsides of the proposed approach (reduced depth and how it affects reasoning, etc.) needs to be carefully studied.