Gated Slot Attention for Efficient Linear-Time Sequence Modeling

Summary

The paper introduces Gated Slot Attention (GSA), an enhancement of Gated Linear Attention (GLA), aimed at improving the efficiency of sequence modeling. GSA incorporates a selective gating mechanism to manage memory updates, leveraging a two-pass GLA structure. This approach allows GSA to be more context-aware and to retain a bounded memory footprint, making it suitable for long sequence tasks. The key improvement is the incorporation of a softmax operation to retain sharp attention distributions, which enhances sequence modeling by reducing dilution.

Strengths

- The GSA's gated update mechanism and two-pass structure offer a straightforward yet effective enhancement over GLA. This design ensures that memory usage remains bounded and manageable, which is critical for handling long sequences efficiently. - The experimental results provided by the authors demonstrate the advantages of the GSA mechanism. The results highlight GSA's ability to improve performance in sequence modeling tasks, validating the practical effectiveness of the proposed approach.

Weaknesses

- 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. - Despite efforts to reduce computational overhead, the softmax operation on $QK^T$ 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. - 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.

Questions

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

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Summary

A major challenge is storing information in a bounded number of memory slots. This work builds on ideas from ABC and gated linear attention. ABC recursively updates the bounded-memory keys and values states over time, and computes a softmax with the queries at timestep t to produce the output at t. Gated linear attention and Mamba use data-dependent decays over time to better decide what information to keep versus throw away given the limited memory. GLA, unlike Mamba, enables the use of GPU tensor cores through its parameterization. GSA retains the selective gating from GLA and the update rules from GSA. The architecture remains chunk-wise parallelizable (like GLA) and provides inductive biases that address limitations of GLA and ABC respectively. GSA is validated up to 2.7Bn parameters from scratch and 7Bn parameters in continual training, providing promising results on standard LM-eval harness benchmark tasks.

Strengths

The writing and contextualization of the contributions with respect to prior work are very clear. The architectural inductive baises from GSA versus prior linear recurrent architectures are compelling and very interesting. For e.g., “In GLA the query vector is only dependent on the current input token, while in GSA, the query vector is the output of the first GLA pass and thereby is aware of the entire historical information and in this sense more context aware”. Two passes over the input could help the model make better decisions about what information to keep versus throw away given the limited recurrent memory. GSA is trained from scratch to relatively large scales (2.7Bn parameters, 100B tokens) and demonstrates high quality on LM-eval harness tasks compared to the baselines, despite using less memory than GLA. GSA can be implemented efficiently by adopting off-the-shelf algorithms from flash linear attention.

Weaknesses

Building on my comment on the compelling properties of GSA, the paper does not extensively show how these properties make a difference in language modeling. 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. OR for the authors to include discussion on what kind of sequences this property might help with. The architectural modifications are not tied to empirical insights beyond lm-eval harness scores in the current submission. GSA is only evaluated on LM-eval harness. It is known that models perform somewhat similarly on these short-context tasks with very little memory (Arora et al, 2024) and mentioned in the paper’s limitations section. The paper does not dive into any tasks that require longer context reasoning or retrieval, making it unclear how GSA behaves. - "GSA only needs half the recurrent state size of GLA and a quarter the recurrent size of RetNet, while having better performance, thus enjoying a lower memory footprint during inference". The claim is not fully validated unless the GSA models are evaluated on tasks that stress-test memory utilization (like retrieval, long-context tasks). 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. There is no analysis as to why this is. 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.

Questions

1. Why do the authors believe the current set of benchmarks is sufficient? Is it possible to include results on retrieval and longer context benchmarks? 2. Can the authors perform error analysis on the continually trained models to help understand why SUPRA and GSA perform poorly on MMLU? 3. Can the authors provide more concrete hypotheses as to where context-aware query tokens (resulting from the first GLA pass) could help in real language modeling settings? In the introduction, “slots” are not defined. They are first defined in section 2.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors do mention that they ignore retrieval style tasks in their analysis (in Section 6), however if the paper makes claims about memory-efficiency then it is important to evaluate quality on this style of tasks. This is because it is known that there are fundamental memory and quality tradeoffs. If the reviewers address these concerns, I will consider increasing my score!

Summary

The paper explores a new variant of attention with bounded memory to reduce the growing memory size and thus mitigates the memory-intensive challenges of Transformers. The key idea is to set a memory bound, with a predetermined number of usable memory slots, and a gating mechanism to select or mix KV vectors from the previous step.

Strengths

The paper proposes a variant of attention mechanism with bounded memory, which seems to be an original contribution.

Weaknesses

* 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. * The writing on the presentation of the proposed method and the evaluation and insightful discussion could be improved.

Questions

* Fig. 1 can be better illustrated and explained how the proposed Gate Slot Attention works. * What is the reason for choosing memory slots to fit in SRAMs? Particularly when the claim is for efficient training. * What is the contribution when the memory footprint as shown in Table 2 is larger than other methods? * What are the overheads, say additional parameters and operations, as introduced by this work? * Some key results could be mentioned in the abstract and introduction. * Fix the extra line-breaker as in line 255? * The quotation marks used in line 128 to 135 are bizarre.

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

The limitation mentioned in the paper is moderate.

Thank you for your response

The new experiments are helpful. Can the authors further explain why GSA should be faster or equal in speed to GLA, despite using two-passes? The writing around this point could be clarified.

Response to All Reviewers

We sincerely thank each reviewer for their thoughtful feedback and recognition of the improvements presented in our rebuttal. We've added new results on recall-intensive tasks and challenging language understanding and reasoning tasks, as well as length extrapolation to address the concerns raised by Reviewers AbD2 and LU5r. Further analysis on time and memory complexity and comparisons of inference performance were also provided in response to common concerns. We are pleased that our additional empirical results, detailed analysis, and explanations have addressed your concerns and contributed to raising your scores. All reviewers have acknowledged our contributions and given positive feedback to our work. We are committed to further enhancing our work. In the next version, we will incorporate all suggested revisions to improve clarity, polish the writing, and rectify any remaining errors. Additionally, we will: * Release our code: Make our codebase publicly available to facilitate further research and validation of our methods, including Triton kernels for training and inference. * Provide pretrained weights: Release all pretrained weights, including 1.3B, 2.7B, and 7B parameters. We sincerely value your input and look forward to sharing our work more widely. Thank you once again for your detailed reviews and assistance in enhancing our submission.