Gated Slot Attention for Efficient Linear-Time Sequence Modeling

Q: While GSA presents improvements, it largely builds on existing GLA techniques. The enhancements, though valuable, might be seen as incremental rather than revolutionary, potentially limiting the perceived impact of the work.

A: We appreciate your feedback and respectfully argue that GSA represents substantial advancements over existing techniques:

• GSA is motivated by viewing $\mathbf{K}$ and $\mathbf{V}$ as memories and linearizing standard attention by managing $\mathbf{K}$ and $\mathbf{V}$ into constant number of memory slots. This perspective drives our 2-pass recurrence formulation: $\mathrm{softmax}(\mathbf{Q}\tilde{\mathbf{K}}^\top)\tilde{\mathbf{V}}$. This significantly differ from exising linear attention that conduct linearization by rearranging the computation order of $\mathbf{Q}\mathbf{K}^\top\mathbf{V}$.
• The 2-pass recurrence with context-aware queries demonstrates significant improvements in retrieval tasks (FDA, SQUAD, SWDE), showcasing the practical impact of our approach.
• GSA maintains the general form of standard attention while offering $\mathbf{K}\mathbf{V}$ linearization. The $\mathrm{softmax}$ spikiness [1] notably enhances continual-pretraining results, paving the way for efficient scaling to larger models.

Q: Despite efforts to reduce computational overhead, the softmax operation on $\mathbf{Q}\mathbf{K}^\top$ still retains a quadratic complexity in training. This raises questions about the authors' claim of linear-time sequence modeling in the title, which could be misleading.

A: Thank you for your concern! We would like to clarify that in GSA, the $\mathrm{softmax}$ is applied over $\mathbf{Q}\tilde{\mathbf{K}}^\top$ rather than $\mathbf{Q}\mathbf{K}^\top$ in standard attention (SA), as presented in Section 2.2. Unlike $\mathbf{K}\in \mathbb{R}^{T\times d}$, $\tilde{\mathbf{K}}\in \mathbb{R}^{m\times d}$ is derived by passing $\mathbf{K}$ along with $\mathbf{\phi}\in \mathbb{R}^{T\times m}$ through GLA recurrence, compressing it to a constant size of $m\times d$. This crucial difference enables GSA to reduce the training complexity from quadratic to linear time, supporting our claim of linear-time sequence modeling.

We will ensure this distinction is more explicitly stated in the revised version. Thank you.

Q: The authors do not specify explicitly whether GSA inherits the recurrent architecture of GLA. If this is the case, it would be necessary to compare GSA to other recurrent models like RWKV to provide a comprehensive evaluation of its performance and advantages in sequence modeling tasks.

A: We would like to clarify that: as detailed in Section 3, GSA's output is indeed derived from $\tilde{\mathbf{K}}_t^\top \mathbf{q}_t \in \mathbb{R}^d$, where $\tilde{\mathbf{K}}$ and $\tilde{\mathbf{V}}$ maintain constant sizes regardless of sequence length. Both $\tilde{\mathbf{K}}$ and $\tilde{\mathbf{V}}$ involves one GLA recurrence. We provide hardware-efficient implementations for the 2-pass operations to enhance training efficiency. Full algorithm details are available in Appendix A.

Regarding comparisons with other recurrent models, we include comparisons with RWKV6 in Table 5 in the paper. Our results show that GSA significantly outperforms RWKV6 under controlled settings, while RWKV6 is trained on billions of tokens and GSA uses continual pretraining. It's worth noting that while RWKV6 shares similarities with GLA in terms of data-dependent gating and matrix-formed hidden states, it employs additional short convolutions at the expense of efficiency.

Due to limited computational resources, we haven't provided results for RWKV6 models trained from scratch, nor have we employed the specific techniques used by RWKV6, to ensure fair comparisons. We commit to improving the clarity of our presentation in the next version of our paper.

Q: How does the inference performance of GSA compare to other SOTA models in terms of speed and memory usage?

A: During inference, GSA demonstrates similar theoretical inference performance compared to other SOTA models :

• GSA needs $O(1)$ space for hidden states, similar to GLA and other linear models, while SA requires $O(N)$ KV cache.
• GSA requires $O(1)$ time per step to access memory and generate outputs as in GLA.

Our empirical speed evaluations for 7B models with 2k prefix support these advantages:

• GSA and GLA consume similar 15GiB memory, while SA requires an additional 3GiB due to KV cache requirements for further generation.
• Regarding speed, the latency for GSA is 13s, comparable to GLA (13.4s) and SA (12s), aligning with our findings in Tables 2 & 3 of our paper.

It's important to note that we have not yet fully optimize the recurrent kernel for speeding up the generation stage, nor reducing the I/O overhead through customization. We are committed to exploring these optimizations in future work to fully realize GSA's theoretical inference advantages and enhance its real-world performance. Thank you.

[1] The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry: https://arxiv.org/abs/2402.04347

To provide more clarity, we conducted additional comparisons on an H800 GPU, optimizing our inference kernel for GSA to perform one-by-one generation. This optimization significantly reduces I/O overhead for 2-pass recurrences, which is crucial for auto-regressive generation. Our updated results show:

We compare the inference latency (seconds) as well as the memory consumption of Transformer++, GLA, and GSA for a single sequence. By varying the generation length from 2k to 16k:

• GSA maintains consistent memory usage across different sequence lengths, unlike Transformer++ which consumes up to 4GiB more for 16k sequences.
• GSA shows comparable or slightly better inference speed than GLA, especially for longer sequences, thanks to our optimized fused kernel.

These findings underscore GSA's competitive performance in real-world scenarios, combining the theoretical advantages of linear attention models with practical optimizations. We are committed to further exploring and implementing optimizations to fully realize GSA's potential and enhance its real-world performance.

We hope this clarifies our contribution and strengthens our paper's position.

Thank you for your thoughtful feedback and insightful questions. We promise to revise the paper to address your concerns and make the presentation more clear and comprehensive in the next version.

Q: It would be interesting to know how the "context aware query vectors" help GSA on real-world data by comparing to models without this property. Why do the authors believe the current set of benchmarks is sufficient? Is it possible to include results on retrieval and longer context benchmarks?

A: Thank you for your very constructive suggestions. In the above table, we have additionally included results to validate the effectiveness of "context aware query vectors":

• We conduct experiments on three tasks proposed by Arora et al, 2024 [1]: FDA, SQUAD, and SWDE, showing significant improvements in GSA's performance compared to Mamba, RetNet, and GLA on these recall-intensive tasks. Despite smaller memory capacity, GSA outperforms Mamba, RetNet, and GLA by an average of 5.9, 7.7, and 4.8 points, respectively.
• We report results on long context benchmarks reported by Xiong et al. 2023 [2]. It is clear that GSA can be well extrapolated to 16k give 2k training context length. GSA exhibits better extrapolation than GLA and RetNet on Qasper, NarrativeQA, QuALITY and QMSum.

These results validate GSA's strong performance on information extraction tasks, which we attribute to more effective memory utilization enabled by "context aware query vectors." We are planning to inspect the performance of GSA scaled to larger model size and number of slots.

Q: The MMLU scores of the continually fine-tuned 7B checkpoints remains very low, just like the baseline – the other scores are roughly comparable across models. The paper claims that GSA is drastically better than SUPRA, but the delta is small (<1 point on average), so it is not clear what is meant by this claim. Can the authors perform error analysis on the continually trained models to help understand why SUPRA and GSA perform poorly on MMLU?

A: After carefully diving into the details of SUPRA, we updated the continual pretraining results of GSA models trained with low learning rate of $3×10^{-5}$ combined with slow decay scheduler. From the above table, we observe that GSA significantly outperforms RetNet and GLA across various tasks, including NQ, TriviaQA, and BBH. Notably, for challenging reasoning / language understanding tasks of BBH and MMLU, GSA outperforms GLA by 2.7 and 4.0 points, indicating that GSA's formulation, which is similar to standard attention (SA), allows it to retain more capabilities when finetuned from another SA LLM.

In the bottom lines we further provide the results of continual pretraining 100B tokens: GSA greatly surpasses SUPRA by 2.2, 5.4, 8.7 and 6.0 points, respectively.

[1] Simple linear attention language models balance the recall-throughput tradeoff: https://arxiv.org/abs/2402.18668

[2] Effective Long-Context Scaling of Foundation Models: https://arxiv.org/abs/2309.16039

Thank you for your question and we are happy to provide further clarifications on the speed comparisons during both training and inference:

Training

For sequence length $N$, chunk $C$, and head dimension $d$:

• GLA with chunkwise parallelism requires $O(NC d + N d^2)$ FLOPs and $O(Nd)$ additional memory for forget gates [1].
• GSA with $m$ memory slots (modeled as two-pass GLA recurrences) requires $O(NC m + NC d + 2N md)$ FLOPs and $O(Nm+Nd)$ memory.

In our experiments, we set $m=64$ to leverage tensor core accelerations, as matmuls operating on $64\times 64$ tiles are shown to be highly hardware-efficient [2]. For models > 1B, $d$ is larger than 512. In these cases, despite the two-pass nature, GSA's complexity ($O(NC m + NC d + 2N md)$) can be lower than that of GLA ($O(NC d + N d^2)$) since $2md=128d < d^2$.

Inference

Both GLA and GSA maintain constant time and memory complexity during token-by-token generation. We also implemented fused kernels to reduce the IO overhead for operations like $\mathtt{softmax}$, potentially making GSA faster than GLA implementations in FLA [3].

For comparison, Transformers require quadratic time complexity for training and $O(N)$ time and memory for inference.

To illustrate the efficiency, we provide inference latency (seconds) for a single sentence beyond the training speed / memory analysis in Table 2 of our paper:

Note: We used the basic Huggingface model.generate API for demonstration, leaving room for further optimization.

• Transformer++ underperforms even at moderate lengths (2k) and scales poorly.
• GSA performs comparably to GLA, slightly outperforming it for sequences >8k due to our optimized fused inference implementations.

We will include a detailed speed analysis for both training and inference in our revised manuscript. Thank you for bringing this to our attention.

[1] Gated Linear Attention Transformers with Hardware-Efficient Training: https://arxiv.org/abc/2312.06635

[2] GPUs Go Brrr: https://hazyresearch.stanford.edu/blog/2024-05-12-tk

[3] FLA: A Triton-Based Library for Hardware-Efficient Implementations of Linear Attention Mechanism: https://github.com/sustcsonglin/flash-linear-attention

We sincerely appreciate your thorough review and insightful feedback. We will revise the paper accordingly in the next version to enhance clarity in the abstract, introduction, and discussion sections, polish the writing, fix potential errors, and improve the figures.

Q: The paper could clarify the significance of its proposed method. The results show not so much improvements compared with prior work in terms of performance and memory costs. The effectiveness and efficiency of the proposed method remain unclear. What is the contribution when the memory footprint as shown in Table 2 is larger than other methods?

A: We highlight that Gated Slot Attention (GSA) offers significant advantages over existing methods, as shown in the above tables:

• GSA shows improved performance on recall-focused tasks given limited memory capacity compared with Mamba, RetNet and GLA. Especially on SQUAD and SWDE, GSA outperforms other models by 5.4 and 10.8, respectively.
• GSA's formulation, which mimics the spikiness in standard attention, facilitates effective continual pretraining from well-trained LLMs like Mistral 7B, demonstrating superior performance on challenging tasks like NQ, TriviaQA, BBH, and MMLU.

Despite a slight increase in FLOPs, GSA achieves comparable training throughput to GLA (42K vs. 42.5K tokens/s under 8K context length) and outperforms Flash Attention (36.9K tokens/s) greatly. Regarding memory footprint, the increase is minimal compared to GLA for sufficiently large head dimensions (d). The benefits in performance and versatility outweigh this small trade-off.

We will provide more detailed discussions on complexity and additional results in the next version of our paper to better highlight these contributions.

Q: What is the reason for choosing memory slots to fit in SRAMs? Particularly when the claim is for efficient training.

A: GSA involves two passes of gated recurrent processes, necessitating hardware-efficient implementations to match GLA's efficiency. To address this, we follow the approach in [1], accelerating GSA by fully utilizing tensor cores, which is 16× faster than non-matmul counterparts.

[2] reveals that latencies for tensor-core matmuls of 16×16 remain similar to 64×64 (Figure 1), despite operating on 16×16 tiles. Our preliminary experiments confirm these findings, motivating us to choose a number of slots that can be processed efficiently in one matmul pass.

As shown in the table, we observe a significant throughput degradation when increasing slots from 64 to 128, while there's a substantial perplexity gap between 32 and 64 slots. Consequently, we adopt 64 slots as the default setting, balancing efficiency and performance.

This hardware-efficient implementation allows GSA to match GLA's efficiency while providing improved performance. Detailed implementations are provided in Appendix A. We will include more discussions and comparisons in the next version of our paper. Thank you.

Q: What are the overheads, say additional parameters and operations, as introduced by this work?

A: We appreciate the reviewer's question regarding the overheads introduced by our work. We have included some discussions in Section 3.1 of our paper, but we're happy to provide further clarification here.

• Parameter allocation: For 1.3B parameter models, GSA introduces approximately $dhm \approx 0.125 d^2$ additional parameters for forget gate mappings compared to Llama. This results in a parameter allocation similar to that of GLA.
• Computation Complexity: The computation complexity of GSA and GLA operation is $O(NCd+Nd^2)$ [3] and $O(NCm+NCd+2Ndm)$, respectively, where N is sequence length, C is chunk size, d is head dimension, and m is number of memory slots. While we set $m=64$ and $m\ll d$, the increased complexities are negligible.

Our hardware-efficient implementations demonstrate that GSA performs comparably to GLA in terms of efficiency, as illustrated in Figure 3 in the paper. Importantly, this modest increase in parameters and comparable computational efficiency come with significant benefits, which we've discussed in detail in our response to Q1. These advantages justify the minimal overhead introduced by our approach.

[1] FLA: A Triton-Based Library for Hardware-Efficient Implementations of Linear Attention Mechanism: https://github.com/sustcsonglin/flash-linear-attention

[2] Simple linear attention language models balance the recall-throughput tradeoff: https://arxiv.org/abs/2402.18668

[3] Gated Linear Attention Transformers with Hardware-Efficient Training: https://arxiv.org/abs/2312.06635

Thank you! We're glad our clarifications were helpful. We'll include these new results in the next version.