Selective Attention: Enhancing Transformer through Principled Context Control

We thank the reviewer for their time, effort, and positive assessment.

W1: demonstrate the performance improvement on the LM task is due to "strategic integration of SSA" rather than an increase in parameter count. Q1: Could we bolster the claim that SSA packs more benefit than just extra parameters by testing models that are scaled up with normal attention by equivalent amounts?

In the general response (1) Making SSA parameter-efficient: We examine more efficient parameterization to push this overhead below 0.5%. In Table 1 in the common response, we tried several approaches by utilizing weight sharing and feature-based temperature selection. Those parameterizations also outperform the standard transformer. We also fine-tuned the model using LoRA as illustrated in the Table below. By controlling the dimension of LoRA, we made sure that LoRA and SSA contributed a comparable number of additional parameters. Although LoRA demonstrated performance similar to directly fine-tuned Pythia, SSA notably outperformed LoRA, creating a distinct gap.

W2: The mathematical sections are pretty dense at times.

Thank you for this suggestion. We will make sure to add additional explanation and motivation in the text.

We thank the reviewer for their time and effort. We hope that our response below addresses their concerns and we would be grateful to respond to further questions during the discussion period.

W1,2: Conducting experiments on downstream tasks, The perplexity improvement over the original models is modest.

As we discussed in Common Response: Evaluations on additional benchmarks Based on Reviewers NV9t and qMbv, we ran new evaluations. Table 2 provides results on 5 new benchmarks: Piqa, Hellaswag, Winogrande, Arc-E, and Arc-C. We also ran additional experiments with Llama3-8B and Pythia 410M as summarized in Table 4. For all settings, SSA exhibits clear and consistent improvements.

We also ran an experiment on the passkey retrieval task introduced in [Mohtashami and Jaggi NeurIPS’23]. The results are shown in Table 3. This is a synthetic task that measures a model’s ability to retrieve a simple passkey (i.e., a five-digit number) within a large amount of text. SSA leads to substantial improvement (from 56.9% to 74.4%). Intuitively, SSA is able to better solve this task by assigning different token-level temperatures to digits vs words.

W4: The claim that SSA benefits long text scenarios lacks supporting evidence. In the table below, our experiment was conducted on sequence lengths of 128 and 4096, where $P_{max}$ represents the average maximum probability and H is the average entropy. From the results, one can observe for the vanilla Pythia models, an increase in sequence length from 128 to 4096 leads to a decrease in maximum probability of attention and an increase in entropy, implying a more dispersed attention. However, when incorporating SSA, the maximum probability for the 4096-length sequence is almost consistent with that of the 128-sequence length, and the entropy is significantly smaller in comparison to the vanilla models.

W5: There is no detailed description of the datasets used.

We have provided all implementation details in Appendix A. We will add a remark in the revised version.

Q1: How does SSA's performance vary with different parameter numbers

We have conducted further experiments using larger models, specifically Llama3-8B and Pythia-410m. The details are shown in Common Response: Evaluations on larger LLMs: in the general response. In short, the approach continues to be beneficial for these models. We agree with the reviewer that more parameters could alleviate the attention map’s dilution. However, for this to happen, the embedding dimension of the K/Q/V should be large (ideally square weight matrices). However, while token embeddings grow with the model size, in practice, the number of attention heads also increases resulting in slower growth in K/Q/V’s embedding dimension and decaying aspect ratio for K/Q/V weight matrices. For instance, Llama3-8B has 32 attention heads with 4096 dim token embeddings resulting in only 128 dim K/Q/V embeddings whereas a smaller GPT2 has 768 embedding dim, 12 heads, and 64 K/Q/V dim.

Q2: There is no detailed description of the accuracy mentioned in Table 3. Could you clarify what this accuracy refers to?

LAMBADA introduces a word prediction task where the target item is challenging to guess (for English speakers) when only the sentence containing it is provided, but becomes easier with broader context. Here is a example from [1]

Context: “Yes, I thought I was going to lose the baby.” “I was scared too,” he stated, sincerity flooding his eyes. “You were?” “Yes, of course. Why do you even ask?” “This baby wasn’t exactly planned for.”

Target sentence: “Do you honestly think that I would want you to have a ____ ?”

Target word: miscarriage.

Accuracy is evaluated based on the model's ability to correctly predict the target word.

[1] Paperno, Denis, et al. "The LAMBADA dataset: Word prediction requiring a broad discourse context." arXiv preprint arXiv:1606.06031 (2016).

Thanks again for your valuable suggestions. Before the discussion period is over, we wanted to kindly remind you that we have conducted comprehensive evaluations, which we believe, address your concerns. Specifically, our evals demonstrate that SSA facilitates noticeable improvements on multiple benchmarks (addressing W1,2,4) and we have provided two simple methods to reduce the computational/parameter cost of SSA (only <0.5% or <0.01% extra parameters) (addressing Q1). We appreciate your time.

W1: The terminology and the positioning of the paper is misleading.

• Sparsity ve Spikiness: We appreciate the feedback. We will clarify that our method is not about sparse approximation of the attention map and instead aims to control the “spikiness of attention” as the reviewer mentions. We will also clarify that “spikiness of attention” can be viewed as an “effective sparsity” which can be quantified through $L_\infty$ norm, $L_1/L_2$ ratio, or inverse-entropy of the softmax map. This discussion will also better clarify what is meant by “contextual sparsity” throughout the paper and distinguish it from (hard) sparsity targeted in [1,2].
• Gating vs Temperature: As we discuss in our related work section, temperature scaling (TS) is indeed a special case of gating. However, it is also a simple and powerful method by itself (use cases in uncertainty quantification, long tailed data, etc). Through TS, we are able to improve the performance of the model by scaling all entries of q/k/v using a single scalar (which is why we refer to it as temperature) and provide a clear explanation on how this provably improves the model’s expressivity. In practice, gating mechanisms are often more complex as we can gate each weight with a distinct learnable scalar. Based on the reviewer’s feedback, we have also ran additional experiments to see if such more complex gating can improve the performance. Specifically, we make $\tau$ a learnable $d$ dimensional vector and elementwise multiply with k/q/v (as described in our Def 1). As shown in the table below, using a scalar temperature parameter is as good as learning individual scalars for each coordinate (called “SSA gating”). This also suggests that improvements are indeed arising from controlling “contextual sparsity/spikiness” rather than a more sophisticated gating mechanism.

W2: No actual wall-time efficiency benefits.

Our work aims to improve expressivity rather than wall-time. That said, the benefit of token- and position-aware temperature strongly indicate that token- and position-aware sparsification should similarly help, which can unlock wall-time benefits through sparsity. For instance, Appendix D provides a theoretical connection between position-aware temperature scaling and sparsification and establishes a 1-1 map between the temperature choice and the associated sparsity level to maintain an identical spikiness level of the attention map.

W3: The scale of the experiments are rather small.

We appreciate your feedback. Our submission had GPT2, Pythia-160M, and Llama2-7B. In addition to these, we have conducted further experiments using larger and state-of-the-art models, namely, Llama3-8B and Pythia-410M. The details are provided in the Common Response: Evaluations on larger LLMs. SSA noticeably improves the perplexity and accuracy for both of these larger models, suggesting that its benefit persists across different model scales.

W4: Lack of experimental support for the claim that position-aware gating can alleviate dispersed attention.

In the table below, our experiment was conducted on sequence lengths of 128 and 4096, where $P_{max}$ represents the average maximum probability and H is the average entropy. From the results, one can observe for the vanilla Pythia models, an increase in sequence length from 128 to 4096 leads to a decrease in maximum probability of attention and an increase in entropy, implying a more dispersed attention. However, when incorporating SSA, the maximum probability for the 4096-length sequence is almost consistent with that of the 128-sequence length, and the entropy is significantly smaller in comparison to the vanilla models.

Q1: why not using tanh() but 2*sigmoid()-1 Note that $\tanh(x) = 2 \cdot \text{sigmoid}(2x) - 1$, which is essentially the same function we use, differing only in scale. Since $x$ is trainable, the results should essentially be same.

Q2: why Key vector gating is not used The intuition from word embeddings of [Mikolov et al.’13] suggests that the similarity between a (key, query) pair should align with their cosine similarity. That is, $cos(key_1, query)>cos(key_2, query)$ should ideally imply that the $query$ attends more to $key_1$ compared to $key_2$. Assigning temperature/gating to scale the query vector does not change this order. However, if we assign distinct scalings to $key_1$ and $key_2$, we will end up with scenarios where attention scores are flipped i.e. $\tau_1 \cdot key_1^\top query<\tau_2\cdot key_2^\top query$. In other words, our intuition is that assigning gating on keys will end up influencing their relative semantic similarities to queries (which could perhaps be better achieved via attention weights). This is in contrast to query-scaling which helps decouple the semantic similarity and contextual sparsity and the associated theoretical benefits (Sec 4.1.1 and Proposition 1).

We have also provided ablation experiments in Table 4 of the appendix of the submission. These show that SSA with query and value independently yields clear benefits (in line with our theory), whereas SSA on keys result in performance that is similar to the baseline vanilla attention layer (in line with the above intuition). We will revise and clarify the discussion of key-gating in Section 3.

We greatly appreciate the reviewer’s detailed and constructive feedback. We hope that this response has addressed your concerns. We would be happy to engage further during the discussion week.

We thank the reviewer for their time and helpful feedback.

W1: Computational overhead

In our Common Response, we describe two strategies that reduce the number of parameter overhead below 0.5% while maintaining the benefits of SSA. These approaches are based on weight-sharing (0.47% overhead) and feature-based (<0.01% overhead). The former (re)uses the attention weights within the temperature module whereas the latter uses simple token-level statistics, such as their frequencies in training corpus. The former approach only stores 3 vectors (not matrices) per attention head whereas the latter is constant parameters per head. Both approaches also have negligible inference/latency overhead because we don’t require additional matrix multiplication. These methods only require vector dot-products (at the output layer of the temperature module) and elementwise scaling of matrices.

W2: Baseline model comparisons Beyond the discussion in our general response about employing multiple strategies to reduce the number of parameters while still achieving significant improvements, we also fine-tuned the model using LoRA as illustrated in the Table below. By controlling the dimension of LoRA, we made sure that LoRA and SSA contributed a comparable number of additional parameters. Although LoRA demonstrated performance similar to directly fine-tuned Pythia, SSA notably outperformed LoRA, creating a distinct gap. Beyond the discussion in our general response about employing multiple strategies to reduce the number of parameters while still achieving significant improvements, we also fine-tuned the model using LoRA as illustrated in the Table below. By controlling the dimension of LoRA, we made sure that LoRA and SSA contributed a comparable number of additional parameters. Although LoRA demonstrated performance similar to directly fine-tuned Pythia, SSA notably outperformed LoRA, creating a distinct gap.

W3 :Need for pre-training

Our methodology includes both pre-training and fine-tuning to evaluate. SSA’s outperforms the baseline on both settings. For fine-tuning evaluation, we start by loading the official pre-trained model and then fine-tune it on the downstream tasks. For the pre-training evaluation, we train the model from scratch on the SlimPajama dataset. Subsequently, we evaluate the model on various downstream zero-shot tasks. This approach is widely used for measuring the performance and generalization capabilities of pretrained large language models across diverse tasks [2,4,15].

[2] Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff. arXiv preprint arXiv:2402.18668, 2024.

[4] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O’Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pages 2397–2430. PMLR, 2023.

[15] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.

W4: Evaluation tasks

Table 3 in the common response also examines the accuracy on the passkey retrieval task as defined in [Mohtashami and Jaggi NeurIPS’23]. This synthetic task measures a model’s ability to retrieve a simple passkey (a five-digit number) from a large amount of text. The details are shown in the general response. Our evaluation tasks primarily focus on language modeling tasks like Wikitext and Lambada. However, potential future applications of SSA, like enhancing linear attention by adjusting its 'spikiness', and improving the interpretability of the attention maps, were indeed discussed in the paper. It's important to note that our main focus was to propose and understand the theoretical effectiveness of SSA. While envisaging its wider scope in the realm of NLP wasn't central to this work, we certainly acknowledge it as an intriguing direction to explore in future studies.

Thanks again for your valuable suggestions. Before the discussion period is over, we wanted to kindly remind you that we have conducted comprehensive evaluations, which we believe, address your concerns. For instance, we provided two simple methods to reduce the computational cost of SSA (only 0.5% or 0.01% extra parameters). We appreciate your time.