Selective Attention: Enhancing Transformer through Principled Context Control
Authors: Xuechen Zhang
Xiangyu Chang
Mingchen Li
Amit Roy-Chowdhury
Jiasi Chen
Samet Oymak
Venue: NeurIPS 2024 poster (2024)
Keywords: attention mechanism
transformer
sparsity
language model
architecture
Primary area: deep_learning_architectures

Abstract: The attention mechanism within the transformer architecture enables the model to weigh and combine tokens based on their relevance to the query. While self-attention has enjoyed major success, it notably treats all queries $q$ in the same way by applying the mapping $V^\top\text{softmax}(Kq)$, where $V,K$ are the value and key embeddings respectively. In this work, we argue that this uniform treatment hinders the ability to control contextual sparsity and relevance. As a solution, we introduce the Selective Self-Attention (SSA) layer that augments the softmax nonlinearity with a principled temperature scaling strategy. By controlling temperature, SSA adapts the contextual sparsity of the attention map to the query embedding and its position in the context window. Through theory and experiments, we demonstrate that this alleviates attention dilution, aids the optimization process, and enhances the model's ability to control softmax spikiness of individual queries. We also incorporate temperature scaling for value embeddings and show that it boosts the model's ability to suppress irrelevant/noisy tokens. Notably, SSA is a lightweight method which introduces less than 0.5\% new parameters through a weight-sharing strategy and can be fine-tuned on existing LLMs. Extensive empirical evaluations demonstrate that SSA-equipped models achieve a noticeable and consistent accuracy improvement on language modeling benchmarks.

Selective Attention: Enhancing Transformer through
Principled Context Control
Xuechen Zhang University of Michigan zxuechen@@ umiich.du
Xiangyu Chang
Mingchen Li University of Michigan milii @umich.edu
University of California, Riverside cxian008 @uc.edu
Amit Koy-Chowdhury
Jiasi Chen
Samet Oymak
University of California, Riverside amitre @ ucc.ucr.edu
University of Michigan @iasi@ mass@uumichedu
University of Michigan oymak@ caraak@umicchcii
Abstract
The attention mechanism within the transformer architecture enables the model to weigh and combine tokens based on their relevance to the query. While self- attention has enjoyed major success. it notably treats all queries q the same way by applying the mapping V softmax(Kg), where V, K are the value and key embeddings respectively. In this work, we argue that this uniform treatment hinders the ability to control contextual sparsity and relevance. As solution, we introduce the "Selective Self-Attention" (SSA) layer that augments the softmax nonlinearity with principled temperature scaling strategy. By controlling temperature, SSA adapts the contextual sparsity of the attention map to the query embedding and its demonstrate position in the context window. Through theory and experiments, that this alleviates attention dilution, aids the optimization process, and enhances the model's ability to control softmax spikiness of individual queries. We also incorporate temperature scaling for value embeddings and show that it boosts the model's ability suppress irrelevant/noisy tokens. Notably, SSA is lightweight method which introduces less than 0.5% new parameters through a weight- sharrng strategy and can be fine-tuned existing LLMs. Extensive empirical evalua- tions demonstrate that SSA-equipped models achieve a noticeable and consistent accuracy improvement language modeling benchmarks.
Introduction
Attention is pivotal mechanism in modern machine learning that allows the model to focus and retrieve different parts of the data, enhancing its ability to capture contextual relationships across time and space. While it was originally developed for NLP tasks through the transformer architecture, has enjoyed widespread success in other domains such as computer vision, sequence modeling, and reinforcement learning [45, 35, 9, 39..
The canonical self-attention mechanism a sequenneeto--sequence map that outputs - S(QKT) V where $(.) denotes the row-wise softmax nonlinearity and Q. K, V are the query, key, and value embeddings obtained through linear projections the input sequence X. Through this process, for each query, the model creates query-dependent composition of the input context. Importantly, the model has to accomplish two objectives: namely, capturing semantic similarity between tokens and also adjusting the contextual sparsity. Here, semantic similarity can be quantified through the angle between key-query embeddings and the contextual sparsity through the spikiness of the attention map. While the importance of the former is clear, the latter is equally important given the fact that attention maps tend to be sparse in practice [8, 43, 37, 6].
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

In this paper, we argue that these two objectives can be at odds and, as a result, the self-attention layer may struggle achieve both objectives simultaneously due its relatively inflexible parameterization. To address this issue, we propose the Selective Self-Aftention (SSA) layer that aims to decouple semantic similarity from contextual sparsity. SSA relies on a principled application of temperature- scaling (TS) to query and value embeddings. For instance, given query embedding q. rather than computing q) SSA computes S(T(g)) Kq) where r(q) is the learnable inverse-temperature Intuitively, this allows for better control of the context window because r(g) can control contextual sparsity while the projection matrices Ws, W can fully focus on controlling semantic similarity. Figure shows example of the learned token temperatures when training the Pythia model with SSA. In summary, we make the following theoretical and empirical contributions:
Query selectivity. We prove that introducing TS to the query embeddings enhances the model's capability to express target attention map with smaller parameter norms (Proposi- tion 1). TS particularly helps when attention maps exhibit large variations in spikiness across different queries. Real and synthetic experiments corroborate that TS enables spikier/sharper attention maps and mitigates attention dilution. See Figure 3 as an illustration.
Value selectivity. We formalize the benefit of TS on value embeddings through denoising perspective. Namely, we describe denoising task where the linear value projection fails to filter the noisy tokens, and demonstrate how nonlinear scaling can boost denoising capability.
Positional temperature. We incorporate a term that adjusts the to the position in the context window. We show that this term can mitigate the dilution of attention scores caused by the increasing context length.
according
Modularity and parameter-efficiencyy of SSA. Selective Attention is accomplished by introducing a parameter-efficient temperature module that can be easily integrated into existing attention models. In practice, this introduces 5% additional parameters to the model. We also introduce a weight sharing strategy that reduces the number of parameter overhead to less than 0.5% while maintaining the benefits of SSA. We reuse the attention weights within the temperature module, which results in negligible inference/latency overhead since no additional matrix multiplication is required. These methods only involve vector cot-products (at the output layer of the temperature module) and elementwise scaling of matrices.
Empirical benefits. Our evaluations on the NLP benchmarks of Wikitext [27], Lam- bada [32], Piqa [4], Hella [51], Winogrande [38], Arc-E, and Arc-C [10] demonstrate that Selective Attention noticeably improves language modeling performance. These benefits are consistent across various models including GPT-2 [34], Pythia [3], Llama [44] and
Llama3 [16], as well as during both fine-tuning and pre-training, as shown in Table 3. Additionally, evaluations on the passkey retrieval task [33, 29] reveal that SSA substantially
enhances the retrieval capabilities of the transformer, shown in Table 4.
People think focus means saying yes to the thing you ve got to focus on. But that sot what it means at all. It means saying no to the hundred other good ideas that there are. You have to pick carefully
0.8
Token Temperature
0.6
0.4
0.2
Figure 1: A quotation by Steve Jobs. We highlight tokens according to their temperatures learned by
the SSA layer. Darker colors correspond to lower temperatures and receive a sparser attention map.
2 Related Work
Temperature Scaling (TS): TS is a fundamental method for controlling model behavior, influencing aspects such as stochasticity of generative LLMs, calibration and uncertainty, and imbalanced data, as highlighted in several studies [26, 22, 52]. Related to us, previous research [33, 49, 7] has also proposed utilizing a temperature term in the softmax function to enhance the length extrapolation capabilities of transformers. For instance, Yarn [33] scales the attention logits as function of the sequence length and shows that this improves the perplexity when extending the context window. Our work provides a formal justification for the temperature scaling rule proposed in Yarn (see Proposition 2) and also highlights the value of adapting temperature to the individual positions. Importantly, our
approach is differentiable and obviates the need for grid search required by prior works. Since
2

---

don't focus on length generalization, we have found that position-awand temperature has much smaller benefit compared to token- aware temperature, which our primary contribution.
Gating mechanisms and selectivity: Various strategies have been developed to mitigate the impact of uninformative inputs in model training and processing. Gating mechanisms, originally introduced through LSTMs [19], have been proposed to selectively filter or scale down the input sequence [48, 13, 14, 25, 40J. Very recent sequence models such Mamba (a.k.a. selective state-space model) and Griffin also incorporate gating to boost language modeling [18, 46, 54, 14, 21J. These models leverage imput-dependent gating ensure parallellizable training and enjoyed noticeable success. These methodologies inspired our approach, which incorporates TS to augment the selection capabilities of the attention layer. Specifically. TS can be viewed an instance of gating that selectively passes or suppresses tokens to provide better control of contextual sparsity and relevance. In this light, our work also provides a mechanistic understanding of how gating mechanism can aid self-attention to improve its expressive capabilities. Finally, we highlight the concurrent work [50] which utilizes a differential softmax parameterization to promote spiky attention maps.
Mechanistic understanding of transformers: The importance of transformer-be based models led to many research efforts on developing a stronger understanding of various aspects of transformer and attention [30, 47, 15J. While it is impossible to cover all of these works, it is evident that capability to select relevant features and promote contextual sparsity crucial for the ability of language models to perform complex tasks such as reasoning [23, 1, 43, 533. These have provided inspiration for us to pursue an enhanced modeling of attention's 's spikiness (e.g. as Figure 3). The experiments in Figure 3 are inspired by the recent work [20] which characterizes the learnability of a ground-truth attention model via the next-token prediction objective in terms of the associated Markov transition matrix.
3 Methodology: Selective Attention Layer
Let us recap the self- mechanism in Transformer [45]. Canonical softmax attention admits an input sequence X = [xx xy." e RRddd of length L with embedding dimension d. We then project X to obtain key, query, and value embeddings (K XW2, Q XW.. q. V = XW,) and compute the
VV. Here S()) : R - R4 denotes the softmax nonlinearity that applies row-wise and Wg. Wt, W, Raxdd are learnable weight matrices. In this paper, we mainly focus on casual language modeling where each token can only attend to previous tokens in the input.
output of the dot-product attention as An( K, V) S( QK
va
The uniform treatment of all tokens through the same softmax map could hinder the ability to control contextual sparsity and relevance. For instance, it has been observed that current Transformer language models suffer from an attention dilution issue: the longer the input sequence, the flatter the attention distribution [49, 7J. A natural solution the dispersed attention issue sharpen the self-attention distribution. Selective Attention aims to provide a general strategy to control spikiness of the softmax adaptive to the query and value embedding, as well as the position of the token.
Definition 1 (Selective Self-Attention (SSA)). Let X [x] ... xx] E R be an input sequence. functions for keys, queries, and values, respectively.
Let Teigiv()) : Rd - Ie' be the Then the embeddings for keys (K), queries (Q), and values (V) are computed as follows:
= Tx)) XW Q = xx(X0 XW g. V T
where @ denotes the elementwise product that assigns temperature to individual tokens. Selective Self-Attentoon (SSA) is then computed as 3( QK"
V.
vd
In essence, SSA incorporates a temperature modulation mechanism into the attention framework to enhance selectivity and context control. The temperatue function T(-) data-dependent, allowing for dynamic adjustment of attention across different parts of the input sequence. In practice, we choose Tkigív be a scalar valued function as pector-valued temperature does not provide a significant advantage. It is also worth mentioning that we don't 't restrict Tkigiv to be non-negative. As a result, our temperature scaling strategy can be seen as an application of scalar gating on K/Q/V embeddings, and hence, the SSA layer could also be referred to as Scalar-Gated Atlention (SGA) layer. The GitHub repo containing SSA implementation is provided in https://g// thub. com/ umi ch- sot-ssottsslleee swwsctieeeeeeeeeeeeeeeeeeeeeeeee attention. Below, we discuss the design choices underlying SSA.
3

---

Temperature scaling for query and value tokens. In an attention mechanism, the concepts of keys (K), queries (Q), and values (V) play distinct roles in deternining how information is weighted and combined across a sequence. Temperature functions can be applied to all of those components, designated as Key-temperature Tx(+), Query-temperature T,()) and Value-temperature T,O) We explore the advantages of each temperature function in in Appendix B.1. In practice, we employ Query-temperature Ty(+) and Value-temperature T,(-) but don't It touch the original key embeddings. The query-lemperature T adjusts the spikiness of the attention map associated with the query according to its embedding and position in the context window. The value-temperature Tv enhances the model's ability to suppress irrelevant noisy tokens, ensuring refined aggregation of context window. In Section 4, we provide insights into theoretical and empirical benefits of keep the keys unmodified, guided by the intuition from word
incorporating these terms. terms.While embeddings of [28] suggests that the similarity between a (key, query) pair should align with their cosine similarity. That is, cos(key), query)) > cos(key2, query) should ideally imply that the query attends more to key; compared to keyz- Assigning temperature/gating to scale the query vector does not change this order. However, if we assign distinct scalings key: and key2, we will end up with scenarios where attention scores are flipped i.c. Ti * key, query < T2 s key, 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 (Section 4.1 and Proposition l).
Token-aware and position-aware temperature scaling. The data-dependent inverse-temperature function is composed of two distinct components r(x) = -lo (x) + (x) x is a token within the sequence X: Token-aware Temperature Scaling (-) and Position-aware Temperature Scaling (.))()) Token-aware Temperature Scaling (+) is devised to modulate the influence of individual tokens within the sequence. The formula for this component given by (0)(x)) tanh(f(x)), where f(+) represents a trainable function that adjusts the impact of the token x. The activation function tanh()) is used to enable the scaling function to output both positive and negative temperatures; for instance, if we want to have the option to fully-suppress a token r'x can attain ~ O. To address the issue of dispersed attention, where increasing length of the input sequence leads to a flatter attention distribution, we introduce Position-award Temperature Scaling. This is defined by (x) 1 + (og))))))) where n denotes the position of the token x within the sequence =x ... xx]T= E Riad E [L]. We remark that reflects the token length when computing the temperature of token x,, aligning with our focus on causal attention where each token restricted to attending only to previous tokens in the sequence. a is parameter designed to modify the scale of the factor. The non-linearity o--) is the sigmoid function, employed to control the range of poor and ensure the stability of the training process.
Weight sharing. We introduce weight sharing strategy to reduce the number of parameter overhead below 0.5% (10x fewer) while maintaining the benefits of SSA. Specifically, the Position- aware Temperature Scaling term, (x) only includes a single parameter a, whereas the Token- aware Temperature Scaling term (x) = tanh(f(x)), relies on a trainable function f(.) defined as Ge x) involves separate trainable parameters W mp and Wiig: which increases parameter load. To improve efficiency, we (re)use the attention weights Weig/t for the temperature module by setting f(x) = W mmp W (eLg((Wigg)a)))) Here, SSA only adds the output layer of the MLP., a vector with few parameters. The approach only stores 3 vectors (not matrices) per attention head. This also have negligible inference/latency overhead because we don't '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. Other strategies can also be deployed to reduce the computational overhead. We describe feature-based approach which use simple token-level statistics, such as their frequencies in training corpus. Only constant parameters per head need to be stored that reduce the number of parameter overhead below 0.1%. The deatils are shown in B.2.
Finally, we discuss conceptual connections to sparse attention methods in Appendix D.
4 Theoretical Insights into Selective Attention
Selective attention computes the query temperature based on the embedding and the position of the query. also computes the value temperature based on the value embedding. In what follows, we discuss how these three components provably enhance expressivity of the attention mechanism.

---

4.1 The benefits of incorporating query embedding
Decoupling semantics from specificity. Consider two words: "Hinton" and "Scientiss", The former is specific instance of the latter. As a result, while we expect token embeddings of these two words to have high cosine similarity, they might benefit from different attention maps. Specifically, "Hinton" refers to a specific person and we expect it to have a more targeted attention to the context associated with it. We argue that query-temperature can aid optimization by retaining semantic similarity while allowing for distinct specificity. More formally, by specificity we are referring to the contextual sparsity level of query. Denoting the combined key-query weights to be W W, W. as problem-agnostie measure of specificity, we will consider the magnitude of the query embedding. That is, given query token q, define specw (q) == ||W 4l12- It is well-established [43] that in order for attention map to be more sparse (hence higher specificity), the norm of the query embedding, or more generally the operator norm of W, has to grow larger. justifying this definition. The following Lemma shows that, without TS, the attention weights within softmax have to be lower bounded by the ratio of specificity difference to semantic distance.
Lemma 1. Lei W = Wy W Raddd be the combined query-key matrix. Let a, b Rd be unit norm token embeddings associated with the specific and general token respectively. Suppose wish to achieve specificities specw(a)> 2 Lo and sb)cw(b) Lb- Then, the associated W obeys [|
L--4 Ta-bli
Above Lu L is the s difference whereas I|a bll2 is the semantic distance. The proof follows from the triangle inequality ||W]| 2 (a-(a-bile [W'ale aWll-W/-WW'' Alla
L.-Is Ila-bil:"
|aa-bll:
la-bil
Comparison to Selective Attention. In SSA, the effective attention weight matrix for a query q is W = T(q)- W, WT. To achieve the same specificity in Lemma with SSA, we can set the temperatures as r(a) Lu- r(b) Le, and KQ-weights as ||W]| 1 (e.g. via W I2). This achieves the desired specificities while maintaining that effective weights are upper bounded as ||Wall. [|Wall s max(/ L4). In other words, the required norm growth is entirely decoupled from the semantic distance between the queries.
In essence, this highlights that without query-selectivity, the model weights have to grow excessively to assign different specificity to similar words. In practice, this is expected to create performance bottlenecks: (1) As the weights grow, optimization may slow down along certain directions due to vanishing softmax derivative and, (2) even if the optimization is successful, the final model could overfit be overly sensitive to small perturbations in the context, hindering test accuracy.
This is also verified by our experiments. To study the norm growth of attention weights, we train Pythia from scratch, trainig with the SlimPajama dataset [41][our pre- training setting) and evaluate on Wikitext dataset. We examine the average norm of combined query-key matrix weight [[W]] from the average of all layers within the model. Additionally, we quantify the spik- iness of the attention map computed as the ratio of the I-normm to the squared 12-norm and normalized by the length, de- [ll,, fined as [ s where s is the softmax prob- ability vector. It takes values from 0 to 1. smaller value indicates a sparser vector. We compute the average of the first 1000 tokens of the Wikitext dataset. The results
0.4
SSA
Prebabiity
0.3
0.0
024.
0.40.6
0.6
0.8.1
1-1.2
x1.2
Ranges
Figure 2: The operator norm of W with and without Query-temperature scaling, scaled by x103 The fig- ure depicts the distribution across 1000 tokens. The dashed line is the average norm. Notably, the norm of the vanilla attention layer is approximately three times larger than that of Addahhed red line compare to green line). Furthermore, the vanilla attention layer exhibits lower spikiness score (0.39) compared to SSA (0.26), where lower value indicates higher spikiness.
shown in Figure 2 align with the theory. The attention weights for selective atten-
tion are smaller than the original ones, while the attention is sparser. Expressivity benefits of query-selectivity. A closely related consideration is whether query-selectivity can enhance expres- sivity. We expect that through tnerrttemperrture, the same attention head will have an easier time expressing sparse and dense attention maps associated with distinet queries. To formalize this we investigate the ability of a single (selective) attention head to express a target attention map between
5

---

(b) Ground-truth token transitions P.
(c) P learned SSA
(d) P learned by Self-Attention
(a) The graph.
Figure 3: We compare 1-layer SSA and I-layer attention when solving next-token prediction on a small vocabulary of size 8. (a) is the graph associated to the token transition dynamics. (b) is the the pairwise token transition matrix this vocabulary. Each row of P. represents an attention map where a particular token is the query and all tokens in the vocabulary serve as keys (see Sec 4.1 for details). The transition matrix P estimated by SSA in (c) is sharper and more closely resembles the optimal P.. SSA achieves a smaller cross-entropy loss compared vanilla attention, 0.009 vs 0.0126. The i approximation error of the attention map of SSA is also smaller than that of vanilla attention, 0.358 vs 0.543.
all tokens in discrete vocabulary. LV e,, be a vocabulary of K tokens. To capture all K² pairwise interactions of these tokens, we first form the sequence E [e; ex]" e Rx where each token appears uniquely and then study K attention maps associated with individual queries, ie., att((,e,,, e,) for S i K. Stacking these together as rows, we study the K x K attention matrix att(E). For standard attention with weights W, this is given by att(E, W) S(EWE"), whereas for query-selective attention, att(E, W) = S(T(E) O EWET).
Thanks to the softmax nonlinearity, att(E) is stochastic matrix where rows add up to 1. This matrix can be viewed as Markov chain transition between different tokens, which motivates a fundamental question: Can query-selective atlention help express a larger class of stochastic matrices? Intuitively, we expect that if stochastic matrix P., which we wish to express via att(E), exhibits lot of spikiness variation across its rows (i.e., different queries), selectivity can better capture these.
This can be verified with a token generation experiments. Recall that we expect "bacteria" attend to more words compared to "salmonella". We might expect more general words to have a larger number of neighbors in graph. Accordingly, we abstract the vocabulary, which comprises words with various levels of specificity, into simple undirected graph. This is depicted in Section 4.1. Additionally, the stochastic matrix P. can be derived from this graph, with the results displayed in Figure 3()). To build the estimation of the stochastic matrix P. training, we conduct next token prediction experiments.
Token generation selling: Let X E V be sequence of length I. drawn from V. Suppose X ends with q:= Xy. The token Y = XL+1 that follows X will be drawn uniformly from q or one of the neighbors of This neighborhood is parameterized via the latent attention map P. which will govern the generation process. Let E = [ei ex]" be the token embeddings associated with the vocabulary V. Assume elements of E have unit & norm. In data generation, we simple sample input sequences containing each token in the vocabulary precisely once, and sample the next token according to the attention map P.. that is, the row of P. that corresponds to the final query token. We then fit lne-r self-attention or SSA model f(X) to approximate this latent dynamics. Concretely, we predict the next token Ý of f(X) according to the distribution g(X)= 3(C)(X) R Here C'e e Nxd the linear prediction head. As loss measure on how well we to the latent P, dynamics, use the cross entropy distance between g(X) and the true label Y. Through this, we wish to formalize and visualize the intuitions on why "salmonella" deserves lower temperature than "bacteria". Further experimental details are described in Appendix A.
In our experiments, besides smaller cross-entropy loss, we find that Selective Attention achieves a better approximation of P. as shown in Figure 3. To evaluate the similarity between the attention map P. and we also define the I distance between the attention maps, namely,
err_ma 1
We find that the err arr_mapssA is also much lower than err_ arr_mapvanill (0.358 vs 0.543). Additionally, SSA naturally assigns lower temperatures to tokens with fewer neighbors. This line with our expectations as fewer neighbors imply sparser attention map. The results are shown in Table 1.

---

Table Temperature for each depth. Nodes with the same of neighbors share the same temperature. # of neighbors(includin itself)
2
3
4
nodes index
7
4,5,6
1,2,3
0
temperature
0.002
0.019
0.152
0.751
To further formalize this, we revisit Lemma in terms of softmax map. Let K = 2 and P.
1-y y
0 be the target pairwise attention map. Here second token is highly specific (only selects
itself) whereas the first token is less specific when 0 < 1. The following proposition establishes variation of Lemma when approximating P..
Proposition 1. Suppose the embeddings e1,e2 have unit l2 norm with correlation p e.e2.
0 & min(y. y) and I log ! ){. For any obeying IIP, -GEWE" )l we have
that [| 2
|le |e!l" (log (da) -)). Conversely, Selective Attention can achieve this 6-approximation 22-2
with weights bounded as r(e1.) || s lle: -e max (t).
4.2 The benefits of incorporating query position
The need for position-dependent scaling arises from the fact that, for fixed weight matrix W Wg W. the attention scores 8 S(XW x2) become diluted as sequence length grows. Specifically, for retrieval-type tasks, the model may want to concentrate softmax scores s on a single token. However, assuming unit norm tokens, the top probability in s is upper bounded via ||s I+ This implies that, enforce [|s"Ile, be constant, require the spectral norm lower growth rate of [|W]| 2 0.5 log + 0(1). This motivates our logarithmic scaling strategy which was also proposed by [33, 7].
Here provide more formal justification on the optimal temperature scaling rule by describing simple yet insightful task which is not solvable by a single attention head unless temperature scaling is employed. Specifically, we consider setting where the sequence exhibits feature imbalances where frequent tokens start dominating the context and potentially overwhelm the less frequent but relevant tokens.
Imbalanced token setup: Suppose the input sequence X = [x] x is composed of a minority token E Rd and majority token b E Rd, that is, x; E (a, b) for all i E [Z]. For each position, we will simply ask the model to output a target mixture of and b, namely, y = aa + (1 arb for some E (0. 1). Thus, using 1 I-layer causal attention, study the following objective by calculating the loss between target and each attention output:
L(W) Σ -b XWT
XWTx,)i x,)...
(1)
Above, Tu is the inverse-temperatuuee for the Th position. Here, no is burn-in period simplify our exposition: ng is the smallest number such that both and appear at least once within the first no tokens', Additionally, let na be the number of tokens x; that are equal within [n]. We have the following theorem.
Proposition 2. Assume b are unit Euclidean norm and linearly independent. Define the imbalance ratio K, (n na)/na for n E [L]. There W. such that, setting Tn log K, log 1-a" L(W.) minimizes the risk (1) to achieve L(W.) 0.
Conversely, consider the problem instance with target mixture of a 1/2, second-quadrant imbalance of 2 Kn 2 for L/4 S n S L/2 and fourth-quadrant imbalance of K, 2 4 for 2 3L/4, If we employ flat temperature T 1 for all E [L]. for any choice of attention weights W Raxdd we have the lower bound L(W)= > 1/500.
'Tii is without loss of generality, otherwise, causal attention would output fixed vector (either a or b) regardless of the attention weights.
7

---

Table 2: We apply normalization attention output and compute the MSE risk.
Vanilla Value-selective 1.390
Naive averaging Bayes optimal estimator
0.071
2.058
0.003
Proposition 2 inspired our design of position-aware temperature scaling, Intuitively, increases, the sequence may include less related tokens, leading to an increase in KG- When Kn follows power- law K, noox, recover the logarithmic temperature scaling rule of Tn const pow log Consequently, our Position-awarg Temperature Scaling function Tn is designed as r**(x) = + (og(n))))) is the position length, is the trainable parameter, or is the non-linearity function sigmoid. The function is motivated by, other paper's 's rules [33, 26, 22, 52].
4.3
The benefits of incorporating value embedding
Within attention, value embeddings (V) are transformed using only linear projection. Consequently, each token's s contribution to the output a weighted sum based on the attention scores, with these weights adjusted linearly. In sequences with many tokens, irrelevant or noisy tokens can negatively influence the attention mechanism. Because value embeddings are linearly projected, they may not be able to fully distinguish between relevant and irrelevant tokens. The value-temperaturg scaling acts as a nonlinear scalar weighting function. By adjusting the temperature, we aim to control the impact of each token, suppressing the influence of irrelevant noisy tokens. This helps emphasize more relevant tokens, thereby improving the quality of the context representation. We motivate the potential benefits of TS on value embeddings through the following synthetic denoising task.
Denoising task Let [K] be the token alphabet with embeddings (e)) Assume d K and e are standard basis. Consider the following data distribution (X.y) ~ DD where X [x] x2]" E Rixd is the input sequence and E Rd is the target label.
Draw q Unif([K]). Set = eg.
Let (z)) be IID noise vectors with o20
X,. = e, Zı. For E [4.- 1], x; determined by Bernoulli distribution with a parameter of selecting between eq + and Consequently, a of the tokens are signal tokens ey + Z,-
The denoising objective is minimizing the MSE risk
L() Bollyy norm(f)ll
where norm(f)= 9/1l 9 is the output of model ʃ(-). (=)
To solve this task, the attention model f(X) should intelligently combine the tokens within X to approximate the denoised target eg. Importantly, the model will strictly benefit from eliminating the pure noise tokens, i.c., instances with x; z, Note that the value projection of the attention matrix will not suffice to denoise the input sequence. The reason is that is uniform, and signal tokens span the whole space. Thus, we will benefit from nonlinear denoising procedure.
To test this intuition, we use a 1-layer single-head attention model, denoted as different f() to minimize the denoising objective. We compare the model with value-selectivity to the following baselines:
1. Vanilla Attention: The standard 1-layer single-head attention model, Jar Att(X)
2. Value-selective relf-attention: 1-layer Selective Self- Self-Attention (SSA). ISSA = SSA(X). Since this a synthetic task, as proxy for the token-aware temperature scaling, we use the selection function max je(d] X, 2 1/2. Intuitively, when noise or $ 1/ Vlogd. d, thresholding with the largest entry will detect the signal tokens.
3. Naive averaging: Directly average the tokens, Inuive = EL xy.
4. Bayes optimal estimator: Pops J Eies X, where S c [L] is the ground-truth of signal tokens distributed as e, Zi-
The resulting MSE risks are displayed Table We set d = k 8 and a = J. With the addition of the value-selection function, the model achieved loss comparable to the optimal estimator, indicating
8

---

Table 3: Experiment results for model pretraining and finetuning. For perplexity (ppl), lower is better, and for accuracy (acc), higher is better.
Model
Wikited ppll
Lambadi_sst
Lamboda openaii
Toll
Lambada_sst a00?
Lamboda acef
Hella
Winogrande acsT
E
Are
Averupe
Finerune
GPPT2
36.503
51.631
29.13 10.134
0.540
01.51
0.584
0.313
0.476
0.457
0.221
0.406
+SSA base
34.118
50.412
27.23 235
0.361
01,669
0.630
0.358
0.512
0.479
0.249
01431
.SSA weight shering
35,147
50.832
17.95
0.357
0.465
0.603
0.354
0.500
0.472
0.243
0.425
Pythia- (60m
26.681
47.996
10.
0.383
0.494
0.674
0.362
0.542
0.500
0.277
0.462
-SSA base
26.314
23.956
0,388
0.515
0.688
0.375
0.557
0.291
0.477
-SSA weight sharing
26.780
0.586
0.553
0.285
P
42.68
0,418
0.758
base
19976
0.430
0.381
0.558
0.572
0.202
+SSA weight sharing
20.190
0.428
0.380
0.551
0.566
Elamm
19.764
16.513
0.436
0.574
0.704
0.377
0.540
0.995
0.802
0.504
+SSA base
19.305
27/627
15.660
0.428
01.91
0.710
0.388
0.562
0.618
0.356
0.518
+SSA weight sharing
19.512
27
16.038
0.426
0.579
0.708
0.385
0.557
0.608
0.331
0.513
3-80
12.416
24.002
13.954
0.481
0.684
0.772
0.544
0.698
0.790
0.463
0.632
+SSA base
10582
23:671
12.002
0.489
0,690
0.779
0.550
0.703
0.787
0.472
0.639
+SSA weight sharing
11.498
23.905
10. 164
0.487
0.687
0.776
0.548
0.701
0.764
0.471
0.636
Pretrain
GPTT2
35,113
104 225
42.15 12.187
0,216
0.304
0.608
0.309
0.462
0.359
0.186
0.349
+SSA hase
33.528
103.933
40 960
0.231
0.318
0.631
0.317
0.80
0.365
0.201
0.362
+SSA weight shariig
34,601
104 004
41.323 11.326
0,219
0.312
0.623
0.312
0.469
0.365
0.197
0.356
Pythia 160m
27.943
15.487
34,406
0,279
0.351
0
0.348
0.498
0.801
0.219
0.389
-SSA base
26.912
22.391
33.126
0.294
0.360
0
0.359
0.508
0.426
0.250
0.405
.SSA weight staring
27.046
33.814
0,291
0.352
0.503
0.221
Pythia
22,516
base
21,402
11.269
0.237
4SSA weight sharing
21,580
69.041
31.458
0.331
0.384
0.658
0.362
0.534
0.445
0.237
0.422
successful suppression of noisy tokens. In contrast, while vanilla softmax self-attention performs similarly to naive averaging, it fails to sufficiently denoise, resulting much larger loss compared to our value-selective attention.
5
Empirical Evaluations
5.1 Standard Benchmarks
Drawing theoretical insights, we assess the performance of SSA NLP tasks by integrating
SSA into established models such as GPT-2 [34], Pythia [3], Llama [44] and Llama3 [16]. Our methodology includes both pre-training and fine-tuning evaluate SSA's performance and efficiency. For the pre-training evaluation, we train the model from scratch on the SlimPajama dataset [41]. Subsequently, we evaluate the model on various downstream zero-shot tasks, including Wikitext [27), Lambada [32], Piqa [4], Hella [5 Winogrande [38], Arc-E, and Arc-C [10]. This approach is widely used for measuring the performance and generalization capabilities of pretrained large language models across diverse tasks [2. 3, 18]. For the fine-tuning evaluation, we start by loading the official pre-trained model and then fine-tune it on the downstream tasks. Unlike pre-training, where the downstream tasks are unseen during training, fine-tuning involves direct training on the tasks. This allows the model to better approximate the token distribution and understand the text domain. Details of the models are provided in Appendix A.
Our primary results are shown in Table 3. Based on the theoretical insights and ablation study results, we conduct both Token-aware and Position-aware Temperature Scaling on query Q. and value V. We observe that across various models and datasets, incorporating SSA consistently enhances performance. Notably, experiments with larger and more recent models, such as Llama3-8B and Pythia 410M, confirm that SSA improves accuracy across across model scales and architectures. We further introduce a weight sharing strategy that reduces the number of parameter overhead to less than 0.5% while preserving the benefits of SSA and still outperforming the standard transformer. This underscores the value of selectivity irrespective of its precise implementation. Thus, our improvements are not arising from an increase in the parameter count, but rather from the strategic integration of SSA. Additionally, we have also explored a Seature-based method to further enhance SSA's parameter efficiency. nutshell, rather than training an MLP, we select the temperature as a function of token-level features, such as the frequency of token in the training corpus, by fitting single scalar parameter. This process requires only O(1) additional weights ( (<0.01 % of total). Further details and results are provided in Appendix B.2.

---

For the ablation study, we fine-tuned the models on the Wikitext dataset to compare the in- fluence of each component. using the same dataset and training configurations as those in the real experiments. The results are shown in Appendix B.1. Among the results, we observe that deploying both Token-aware and Position-aware Temperature Scaling on Q and inde- pendently could achieve significant improvement, aligning with our theoretical insights. Addi- tionally, combining Key and Query temperatures can achieve additional improvement. More- over, between token-aware and position-aware temperature scaling, the latter demonstrates more consistent improvement across different scenarios, while combining them can achieve the best overall result. We also compare with more baselines including [26, 22, 52] and the results are shown in Appendix B.3. Our method consistently outperforms the baselines.
Additionally, SSA can accelerate the training process by achieving comparable performance with fewer tokens. This efficiency not only reduces the demand on compu- tational resources but also shortens the time required to effectively train models. We illustrate this efficiency by plotting the training results when fine-tuning the Llama model on the Wikitext dataset, both with vanilla attention layer SSA, in Figure 4. The results indicate that SSA can accelerate training, achieving similar performance with 1.45x reduction in pretraining steps.
37.5
35.0
32.5
27.5
25.0
Spcod up
22.5
20.0
0.0
05
1.0
15
2.0
25
3.5
toleens(M)
Figure 4: Comparison of training curves. SSA provides reasonable benefits in terms of training speedup.
5.2 Passkey Retrieval
We also examines the perfromance on the passkey retrieval task as defined in [33, 2his is synthetic task to measure model's ability to retrieve a simple passkey (i.e., a five-digit number) within large amount of otherwise meaningless text. We performed 10 iterations of the passkey retrieval task with the passkey placed at a random location uniformly distributed across the evaluation context window. Intuitively, SSA could better solve this task by assigning different token-level temperatures to digits vs words. For our evaluation of the fine-tuned Pythia, SSA leads to substantial improvement (from 56.9% to 74.4%), as seen in Table 4.
Table 4: Passkey retrieval performance of various models.
Model
Original
+SSA
+SSA(weight sharing)
Pythia- 160m
56.89
74.41 89.53
66.90
Llama
77.62
89.45
Conclusions, Limitations, and Future Directions
We have introduced the Selective Self-Attention layer, which augments the softmax nonlinearity with a principled cemperature-scalinn strategy. SSA shows consistent benefits and augments the performance of existing transformer-based models such Pythia and Llama 2. We also provide theoretical insights into the benefits of query, value, and positional selectivity.
Future research. Based on SSA, there are several interesting research avenues to pursue. Firstly, our method can extend to linear attention strategies. While we can use the same method for value embeddings, for queries, can train additive bias term on attention similarities rather than using temperature scaling. Secondly, based on the visual benefits of SSA on Figure 3, would be interesting to explore how SSA can help the interpretability and quality of the attention maps. Overall, SSA has the potential to assist in more principled use of transformers in language, vision, and other modalities.
Limitations. Our work focuses on the canonical softmax-attentin mechanism, which suffers from the quadratic computation bottleneck. As mentioned above, extending our method to linear attention can mitigate computational costs. Another direction to enhance efficiency building stronger connections to sparsity and understanding how SSA can benefit and be integrated with sparse attention algorithms.
10

---

Acknowledgements
This work was supported in part by the National Science Foundation grants CCF-2046816, CCF- 2403075, CCF-2008020, the Office of Naval Research award N00014241220 an Adobe Data Science Research award, and gifts by Open Philanthropy and Google Research.
References
[1] Emmanuel Abbe, Samy Bengio, Aryo Lotfi, and Kevin Rizk. Generalization the unseen, logic reasoning and degree curriculum. International Conference on Machine Learning, pages 31-60. PMLR, 2023.
[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 tecoul-throughput tradeoff. arXiv preprint arXiv:2402. 240..66002/.888 2024.
[3] Stella Biderman, Hailey Schoelkopf. Quentin Gregory Anthony, Herbie Bradley, Kyle O' 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.
[4] Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about phys- ical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7439-7439, 2020.
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901. 2020.
[6] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. Advances Neural Information Processing Systems, 34:17413-17426, 2021.
[7] Ta- Ta-Chung Chi, Ting-Hap Han Fan, and Alexander Rudnicky. Attention aligmment and flexible posi- tional embeddings improve transformer length extrapolation. arXiv preprint arXiv:2311.0088 2023.
[8] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[9] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1- 113, 2023.
[10] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 1803:0880000055, 2018.
[11] Tri Dao. Flashattention-22 22 Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08661, 2023.
[12] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
[13] Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference machine learning, pages 933-941. PMLR, 2017.
[14] Soham De, Samuel Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint 2rXiv:2402.19427, 2024.
11

---

[15] Yihe Dong, Jean- Baptiste Cordonnier, and Andreas Loukas, Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 2793-2800 PMLR, 2021.
[16] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Ilama 3 herd of models. arXiv preprint arXiv:240..21783, 2024.
[17] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.0000 2020.
[18] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint urXiv:23/2.00752, 2023.
[19] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9 1997.
[20] M Emrullah Iidiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self-attention markov models: Unveiling the dynamics of generative transformers. Interna- lional Conference on Machine Learning, 2024.
[21] Tobias Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. arXiv preprint arXiv:2311.01977 2023.
[22] Mingchen Li, Xuechen Zhang, Christos Thrampoulidis, Jiasi Chen, and Samet Oymak. Au- tobalance: Optimized loss functions for imbalanced data. Advances Neural Information Processing Systems, 34:3163-3111 2021.
[23] Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Exposing attention glitches with flip-flop language modeling. Advances in Neural Information Processing Systems, 36, 2024.
[24] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXhv:1711.05101 2017.
[25] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur, Long range language modeling via gated state spaces. In International Conference on Learning Representations, 2023.
[26] Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, and Sanjiv Kumar. Long-tail learning via logit adjustment. ICLR, 2021.
[27] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv: /60.00/60900843 2016.
[28] Tomas Mikolov. Efficient estimation of word representations in vector space. arXiv preprint 1701, 2013.
[29] Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. arXiv preprint arXiv:2305. 16300, 2023.
[30] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.1
2022.
[31] Matteo Pagliardini. Daniele Paliotta, Martin Jaggi, and François Fleuret. Fast attention over long sequences with dynamic sparse flash attention. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[32] Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint
2016.
12

---

[33] Bowen Peng. Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.0007,, 2023.
[34] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog. 188):, 2019.
[35] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang. Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified lext-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
[36] Liliang Ren, Yang Liu, Shuohang Wang, Yichong Xu, Chenguang Zhu, and ChengXiang Zhai. Sparse modular activation for efficient sequence modeling. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[37] Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Unraveling attention via convex duality: Analysis and interpretations of vision transformers. In International Conference on Machine Learning, pages 19050-19088. PMLR, 2022.
[38] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge scale. Communications of the ACM, 64(9):99- 106, 2021.
[39] Victor Sanh, Albert Webson, Colin Raffel, Stephen H Bach, Lintang Sutawika, Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Teven Le Scao, Arun Raja, al. Multitask prompted training enables zero-shot task generalization. arXiv preprint arXiv:2110.082000 2021.
[40] Noam Shazeer. Glu variants improve transformer. arXiv preprint urXiv:200..05202, 2020.
[41] Zhiqiang Shen, Tianhua Tao, Liqun Ma, Willie Neiswanger, Joel Hestness, Natalia Vassilieva, Daria Soboleva, and Eric Xing. Slimpajama-de Understanding data combinations for llm training. arXiv preprint arXiv:2309, 10818, 2023.
[42] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
[43] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023.
[44] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13977, 2023.
[45] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[46] Junxiong Wang, Tushaar Gangavarapu, Jing Nathan Yan, and Alexander M Rush, Mambabyte: Token-free selective state space model. arXiv preprint arXiv:2401.13661 2024.
[47] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021.
[48] Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training, arXiv preprint arXiv:2312.06633 2023.
[49] Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli, editors, Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing ( (Volume 1: Long Papers), pages 3770-3785, Online, August 2021. Association for Computational Linguistics.
[50] Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. Differential transformer. arXiv preprint arXiv:2410.05258, 2024.
13

---

[51] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint arXiv:1905.07830, 2019.
[52] Xuechen Zhang, Mingchen Li, Jiasi Chen, Christos Thrampoulidis, and Samet Oymak. Class- attribute priors: Adapting optimization to heterogeneity and fairness objective. appear at AAAI, 2024.
[53] Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms transformers learn? a study in length generalization. arXiv preprini arXiv:2310. 16028, 2023.
[54] Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint 40X1v:240.009117,200 2024.
14

---

Table 5: Fine-tuning experiment results for language models on the Wikitext dataset, showeasing
baseline and variations with different components (Q, K, V).
Configuration
Pythia
GPT2
Baseline
28.781
36,503
Q
27.416
34.832
K
28.715
36,443
27.980
35.857
Q. V
26.514
34.618
K.
Q
26.603
34.609
A Implementation details
For the next token prediction experiment that substantiates the expressivity benefit of query-selectivity detailed in Figure 3 and Table 1, employ the Adam optimizer to train model. This model consists of a single-layer, single-head attention mechanism, accompanied by tokenizer and a fully connected layer. The tokenizer embeds the discrete sequence to continuous embedding E. The fully connected layer used as the classifier to predict the node index. We set the learning rate le The training loss is the cross-entropy loss. In experiments, /. = N 8. For SSA, implement Token-aware Temperature Scaling for the query matrix Q. We assign a scaling parameter each group of nodes that share the same number of neighbors. To have better visualization,
do normalization to plot the attention map P. and P. For the experiments shown in Table 2, also use the Adam optimizer and learning rate le But the objective is the MSE risk.
For our empirical evaluation, we utilize several models. We employ GPT-2, which has 124 million parameters, and use the official OpenAl GPT-2 checkpoints that were pre-trained on the WebText dataset [34] for our finetuning experiments. For Pythia, our experiments are conducted with a model size of 160 million and 410 million parameters, using the official checkpoint pre-trained on the Pile dataset [17]. Lastly, for Llama, utilize the smallest variant available, with billion parameters, and similarly fine-tune using the official pre-trained model. As the training configuration, we train with 3.5 million tokens for fine-tuning and 15B tokens for pre-training. We always use the AdamW optimizer [24], B1 0.9 and B2 = 0.95. We set learning rate le with no weight decay and no warmup. The pre- pre-training takes about hours using 4 A40 and fine-tuning takes about 2 days. We FlashAttention [11] to accelerate the training. For weight sharing, each head shares the same funtion. All the experiments conducted with 4 or 8 A40 or L40S. We can directly reuse FlashA Attentottntion|! 122, which significantly improves model efficiency.
B
Additional experiments
B.1
Ablation study
For the ablation study, we conducted fine-tuning on the Wikitext dataset to compare the influence of each component, using the same dataset and training configurations as those in the real experiments. The models are evaluated with perplexity (ppl).
B.1.1 Key-temperature, Query-temperature, and Value-temperature
To evaluate the benefits of applying temperature scaling K. Q. and V. we conducted ablation study, examining each component individually and in combination. For a fair comparison, both position-awe aware and token-awan aware temperature scaling were applied to all components. The results, detailed in table 5, indicate that modifying Q. and V independently yields clear benefits, whereas alterations to K result in performance that is similar to, or even worse than, the baseline vanilla attention layer. The results align well with the theoretical analysis presented in section 4. However, when Q and V are combined, we observe consistent improvements, These findings led us to develop final algorithm, which applies temperature scaling to both @ and V.
15

---

Table Investigate the benefits of Token- aware Temperature Scaling, Position Scaling.
-ware aware Temperature
model
vanilla
T + plo
T
Too
T
7'0
T
Pythia
28.781
27.416
27.995
27.503
27.980
28.342
27.975
GPT2
36.503
34.832
34.970
35.064
35.857
36.320
35.617
Table 7: Comparing different SSA parameterizations
Wikitext pplj
Lambada std Lambada openai ppll
Lambada ace?
Lambada openai acc?
Model
ppl
Finetune
Pythia
26.681
47.99
24.102
0.383
0.494
Pythia +SSA base
26.514
47.945
23.956
0.388
0,513
Pythia +SSA weight sharing
26.780
47.961
24.027
0.386
0,509
Pythia + SSA cature--basdd
27.048
47.966
24.114
0.387
0.499
Pretrain
Pythia Pythia +SSA base
27.943
75.487
34,406
0.279
0.351
26.912
72.891
33.126
0.294
0.360
Pythia +SSA weight sharing
27.046
73.071
33.814
0.291
0.360
Pythia +SSA feature-based
27.281
73.614
33.794
0.287
0.357
B.1.2 Token-aware Temperature Scaling, Position-aware Temperature Scaling
We also conduct experiments investigate the benefits of Token-aware and position-aware Tem- perature Scaling applied to Query q and value v. The results are shown in table 6. Token-aware Temperature Scaling positively impacts both Query q and value whereas Position-aware Tempera- ture Scaling shows smaller improvement on value v, aligning with our theoretical insights. Further- more, when compared GPT-2, Pythia-which features demonstrates fewer improvements. This suggests that while new strategies may mitigate the dis- persed attention issue, our Selective Self-Attention (SSA) method still offers additional improvements.
advanced positional encoding[42]
B.2 Parameter-efficient SSA: Weight sharing and featurization
Here, also introduce feature-based approach improve parameter efficiency. In nutshell, rather than training an MLP, we select the temperature as a function of token-level features, such the frequency of token in the training corpus, by fitting single scalar parameter. This process requires only O(1) additional weights per attention head (<0.01% % total). This is inspired by the logit adjustment strategy of [26] which sets the cross-entropy temperature function of class frequencies.
Our evaluations on feature-based SSA are provided in Table 7. We find that, while the feature-based method is beneficial and highly parameter-effiiient, can be sensitive to feature selection and exhibits more variability across datasets.
B.3 Ablation of Different Parameterizations
In addition to the functions we propose for temperature design, we also explore alternative approaches. Instead of employing our Position-awae aware Temperature Scaling function, we use constant parameter, as suggested in other studies [26, 22, 52]. We also compare with the temperature scaling method proposed by [33]. Furthermore, in place of our Token-awart Temperature Scaling, we adopt simpler approach by directly utilizing token frequency and training only a scale parameter. These experiments were conducted using the Pythia model, fine-tuned on the Wikitext dataset. The outcomes of these comparative analyses are presented in Table 8. Among those baselines, we consistently outperform their results.
16

---

Table Conducting different functions.
Configuration
Pythia
Vanilla
28.781
Yarn [33]
27.602
Token
Constant
28.058
Position
Frequency
27.360
Position+ Token
SSA
26.514
C
Proofs
C.1
Proof of Proposition
Proof.
Given embeddings e1.e2 with unit Eg norm, their correlation is eeje2.
First, consider the approximation error bound | S(EWE Ie S &. To achieve this, the weight matrix W must satisfy the inequality.
To derive lower bound on WiI, observe that:
|P+ S(EWE aa S &
IP. S(EWE' Ma
EWET P.
8
Using the fact that e and e unit vectors and p e e2, have:
1
JEWE Il 2
48
Now, the norm [[W] is given by:
||
|le ezll V2 2p²
log
r
48
Conversely, to achieve s-approximation using Selective Attention, the weights need to be bounded such that:
T(e1.2) ||]] S |le: - - max log
Therefore, the selective attention avoids the 1/ -I-p² p²
dependence on the log(1 term, decoupling the high-specificity requirement (small from the semantic similarity of the tokens.
This completes the proof of Proposition
C.2 Proof of Proposition 2
Proof. We first show the success direction. Set W such that b = and a" Wa = a' Wb = 1. Such W exists thanks to the linear independence of a, b. Now plugging this W into (1), for each position, regardless of whether X, = or X, = b. we obtain
nn
(n na)e b.
X STn XWxx)= nteto na)
+
nge ng
We wish ensure
n
+ K, a.
no (n no)
1 + (n/n, 1.
This in turn implies
1
Kg To
To log Kn + log
a
1 a

---

Next, we discuss the failure case of flat temperature. To do so, we will lower bound the loss over the queries x, b. Set M = c(b-)) wb Following same argument as above, for fixed temperature, W will output non-adaptive composition of the form
gn (XWb) = + MK₃ a + MKę b.
1 MKn
Thus, the loss function will be lower bounded by (accounting for the prediction error in terms and their orthogonality)
L(W)
min
Σ (0.5 =
МKę MK, ) + (0.5 1 Mr,
) min Σ 2. 1 MK,
Now since Kn over both second, at least 1/2 of the queries X,, b. Similarly, least 4/5 of the queries are X b over the last quadrant. We will lower bound the loss over the two scenarios depending M 1/3 not.
First, suppose M 2 1/3, in that case, using Ka 2 4 over the last quadrant, the loss is lower bounded by
L(W) min 1 N
Σ
2..5.
2 (0.5 + MKn )
1 4/3 ) 0.002.
5
H23N/4,x,-
=b
where we used the fact that there are 2
queries with x, b over the last quadrant.
Similarly, suppose M < 1/3, in that case, using Kn - 2 over second quadrant, the loss is lower bounded by
L(W) Z min MSO N N/22n2N/4xx,5
Σ
(0.5 - MKn
1 (0.5 2² 4
1 - 0.0025.
1 )²=
where we used the fact that there are at least N/8 queries with x, = b over the second quadrant. Combining these two cases, we found that, for any choice of W, the loss is lower bounded by 0.002.
D Further Discussion on the Sparsity and Temperature Connection
Connections between sparsity and temperature.
To formally study the sparsity and temperature connection, let us consider fixed attention row n and introduce:
(T) = STa(T XWx,), the scaled attention scores with inverse-temperatuuee 0. 5(k) = Sga(XXx,,, x) denote the sparse attention scores where the top-xn entries are retained and the are to 0 where 0 s 1.
The connection between sparsity and temperature scaling is clear. For instance, the top entry of S(T) will be decreasing whereas the entropy of S(T) will be increasing. Here, would like to establish how temperature scaling rule can be mapped to sparsity rule. We will do this under power-law relevance assumption on the attention scores. Here, we assume that the attention scores admit two values and the fraction of larger/relevant attention scores follow a power-law as context window grows.
Assumption 1 (Power-law relevance). Consider the vector of raw attention scores XWx,, Each entry of is either e or c, = + for some > Additionally, n pow fraction of the entries are equal c. for some pow
Above c. is the score attained by the salient tokens, is the score advantage of salient tokens over rest of the tokens, and pow dictates the fraction of the salient tokens. To proceed, we have the following lemma which identifies condition under which TS and sparse-attention exhibit the same softmax temperature behavior.
Lemma 2. For any choice of = a log and corresponding sparsity == Poy we have that [is(r) [|(()))
POW
18

---

This reveals the clear connection between temperature scaling and sparsification rules. Simplifying the above, this lemma advocates that the sparsification rule should follow the power law decay of K a n => Consistent with this lemma, our experiments demonstrate that sparsification with power-law results in respectable performance.
Proof. We first compute the top entry of S(T) follows
ex
1 nI-pow n11 DOM yo
(n n² ppoN
We similarly compute the top sparse attention score as
Si
ex
E n powe + (K 1- pow)n nl-pow + (K e 7n
Combining these, the top softmax probabilities are matched by setting
e-((--
y = (1 DoN dou e-- K - ppo pow (k - 11 pON y
+ 12 DOW
We need to clarify that our method is not about sparse approximation of the attention map and instead aims to control the "spikiness of attention". The "spikiness of attention" can be viewed as an "effective sparsity" which can be quantified through Lo norm, L 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 from (hard) sparsity targeted in [31,6]].
E
Theoretical Considerations
E.1 Hierarchical vocabulary
Hierarchical vocabulary. Consider k-ary tree of depth D: Each node has exactly & children, except depth d. Such tree has + + +.... kDD N (KeD+1 1)/(k 1) nodes. The tree will correspond to the words/tokens in the vocabulary V of size N.
Token generation rule: Let X E v be sequence of length drawn from V. Suppose X ends with == Xy. The token = XL-1 that follows will be drawn from or the children of available in the context window. If q depth including itself. Let Dxy denote the data distribution (Y, X) where is drawn uniformly from one of the child lokens of xı. available in the context window X.
can attend
total of (K-+1-1__ 1)/(k unique tokens,
The claims below aim to formalize the benefits SSA for modeling the hierarchical token generation Jd be the token embeddings associated the vocabulary V. Assume
process. Let E = [e] ... eN]" elements of are unit & norm. During training, we embed the discrete sequence X into X (ex, }.
Claim (Benefits on attention map). Consider the attention map map(X) X) = S(T(x))) XWxx). Note that map a function of W, E. Define the ideal attention map X to be (X) which uniformly attends to the children of xı, and assigns zero probability to other tokens. Define the population error
err err_map(E, W. r) = Bx-Dx [Immp(XX map" -map"(X0l1]
(XIlI)-
Under suitable assumptions (see remark below), QSSA is provably better than vanilla self celf-attention Lee. having T(x) improves atlention capability by reducing err_map(F,, W.t).
Claim 2 (Benefits on prediction). Let f(X) be attention layer (SSA vanilla). Suppose we sample the next token Y from f(X) according to the distribution g(X) = SXCf(X) RN Here Ce Nxd the linear prediction head. As loss measure, use the expected total-variation (TV) distance between g(X) and the true label Y. Under suitable assumptions, g(X) with QSSA f(X) fits better to the hierarchical distribution compared to vanilla f(X).
19

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
We introduce the studied question and list our contribution in Section 1.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitation of our work in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
20

---

We provide our theory and related theoretical definitions in Section 4. We provide the related lemmas and proofs in Appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide the design of our algorithm in Section 3 and implementation details in Section 5. We also provide main experimental results and detailed analysis in Section 5 and additional results Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
21

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We release our training and evaluation code zip file.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: http:puulllggddds/eeeeee
/
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We explain the training and testing details Section 5
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
experiments representative datasets and compare models performance with several baselines. The experiments results are runs average and we provide error bar in the table.
Justification: We
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
22

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: In Section we state the hardware to run the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't I make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cmm http///wwwpp..ccuuuiiic/hhccttat
Answer: [Yes]
Justification: There is no potential harms caused by the research process and potential harmful societal impact.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We state the potential social benefit of our work in Section 6.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
23

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [Yes]
Justification:
We state the safeguards of our work in Section 6.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the dataset and also all the models used in our paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in package should be provided. For popular datasets, paperswi twwoee com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
24

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We include our experiment code and data generation code in our supplementary
files.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our method and experiments only relate to next token generation, which does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our method and experiments only relate to next token generation, which does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
25

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
26


Reviews:
[{"review_id": "7zhAF5bkZy", "reviewer": "Reviewer_NV9t", "summary": "This work introduces Selective Self-Attention (SSA), an additional trainable module for standard transformer self-attention that applies temperature scaling to its three components: queries, keys, and values. By performing both token-based and position-based scaling to modulate the influence of different tokens and positions dynamically, SSA aims to improve control over contextual sparsity and relevance in self-attention mechanisms. The authors empirically demonstrate that SSA can be integrated into existing transformer architectures such as Pythia and LLaMA, adding approximately 5% more trainable parameters. In terms of training efficiency, SSA shows potential for accelerating training achieving similar performance with fewer tokens.", "questions": "Please see weaknesses section.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1)  The paper focuses on an interesting area of understanding attention patterns and developing attention mechanisms to enhance contextual sparsity given a particular query.\n\n2) The paper is well-written and majority of sections are easy to follow.", "weaknesses": "1. Computational overhead: the proposed selective self-attention modules lead to an approximate 5% increase in model parameters. The authors should address the additional computational cost at inference time, which could be substantial for large-scale deployments. In the case of larger models with billions of parameters, even a 5% increase could significantly impact both training and inference costs.\n\t\n2. Baseline model comparisons: The vanilla models evaluated in this work are relatively smaller in terms of parameter size, potentially making the comparison unfair.\n\n3. Need for pre-training: The SSA-based LLMs studied in the work are pre-trained from scratch. Can the authors provide a justification for this approach?\n\n4. Evaluation tasks: The evaluations are mainly on language modeling tasks (Wikitext and Lambada). Given the research motivation, I believe it is important to demonstrate SSA's effectiveness on other NLP tasks where the learned contextual sparsity can address practical issues such as reducing inference latency over longer sequences, KV cache reduction for improving prefill latency, etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and helpful feedback.\n\n**W1: Computational overhead**\n\nIn 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.\n\n**W2: Baseline model comparisons**\nBeyond 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.\nBeyond 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.\n\n| Model                      | Wikitext ppl ↓ | Lambada_std ppl ↓ | Lambada_openai ppl ↓ | Lambada_std acc ↑ | Lambada_openai acc ↑ |\n|----------------------------|----------------|-------------------|----------------------|-------------------|----------------------|\n| **Finetune**               |                |                   |                      |                   |                      |\n| Pythia                     | 28.781         | 51.620            | 25.173               | 0.373             | 0.497                |\n| +SSA base(+5.26%)        | 26.514         | 47.945            | 23.956               | 0.388             | 0.513                |\n| +LoRA(+5.63%)       | 27.944         | 49.852            | 24.653               | 0.375             | 0.501                |\n|+SSA shared $W_{k/q/v}$ (+0.47%)|26.681|47.996|24.102|0.383|0.504|\n|+SSA feature-based (<0.01%)|27.048|48.906|24.814|0.379|0.499|\n\n**W3 :Need for pre-training**\n\nOur 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].\n\n[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.\n\n[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.\n\n[15] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.\n\n**W4: Evaluation tasks**\n\nTable 3 in the common response also examines the accuracy on the **passkey retrieval** task as defined in [[Mohtashami and Jaggi NeurIPS’23]](https://arxiv.org/pdf/2305.16300). 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Nb9HpHA34Z", "reviewer": "Reviewer_Lp6e", "summary": "The paper proposes token-aware and position-aware gating for query, key and value vectors in self-attention layers.", "questions": "1. For token-aware \"temperature scaling\", why not using a simpler formulation of tanh() but 2*sigmoid()-1? Do you have some empirical results supporting such formulation?\n2. Do you have some theoretical explanations for why Key vector gating is not used in practice?", "limitations": "Yes.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 4, "strengths": "1. The position-aware gating is interesting, and, in theory, can alleviate dispersed attention and thus may lead to better length extrapolation.\n2. The inductive bias introduced in the paper can lead to faster convergence when finetuning a small language model with less than 200M parameters.", "weaknesses": "1. The terminology and the positioning of the paper is misleading. The whole paper claims that they are doing qkv selection but there is actually no sparse attention involved. The author should look at some papers [1,2] focusing on real selection of the input for self-attention. Also, the paper claims that they are doing temperature scaling but actually their temperatures can be negative, so it is basically doing gating of the query, key and value vector because of its element-wise product nature. The authors also use the terminology, \"soft sparsity\", which is self-contradictory and unclear. I suggest a better terminology can be used as \"spikiness of attention\" instead of \"sparsity of softmax\".\n2. No actual wall-time efficiency benefits of the proposed method as a result of it essentially being a gating mechanism. \n3. The scale of the experiments are rather small with LMs below than 200M parameters.\n4. Lack of the experimental supports for the claim that position-aware gating can alleviate dispersed attention.\n\n---\n\n[1] Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention. https://openreview.net/forum?id=UINHuKeWUa\n\n[2] Sparse Modular Activation for Efficient Sequence Modeling. https://openreview.net/forum?id=TfbzX6I14i", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1: The terminology and the positioning of the paper is misleading.**\n- **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].\n- **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.\n\n|Model|Wikitext ppl↓|Lambada_std ppl↓|Lambada_openai ppl↓|Lambada_std acc↑|Lambada_openai acc↑|\n|-|-|-|-|-|-|\n|**Finetune**||||||\n|Pythia|28.781|51.620|25.173|0.373|0.497|\n|+SSA(original)|26.514|47.945|23.956|0.388|0.513|\n|+SSA(gating)|26.509|48.067|24.460|0.383|0.515|\n\n**W2: No actual wall-time efficiency benefits.**\n\nOur 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.\n\n\n**W3: The scale of the experiments are rather small.**\n\nWe 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.\n\n**W4: Lack of experimental support for the claim that position-aware gating can alleviate dispersed attention.**\n\nIn 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. \n\n||128|4096|\n|-|-|-|\n|**$P_{\\text{max}}$(spikiness)**|||\n|Pythia|0.38|0.24|\n|Pythia+SSA|0.41|0.36|\n|**$H$ (entropy)**|||\n|Pythia|2.83|4.86|\n|Pythia+SSA|2.76|3.02|\n\n**Q1: why not using tanh() but 2*sigmoid()-1**\nNote 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.\n\n**Q2: why Key vector gating is not used**\nThe 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).\n\nWe 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.\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "28XwXW0MD5", "reviewer": "Reviewer_qMbv", "summary": "This paper presents Selective Self-Attention (SSA), which introduces trainable temperature functions to adapt the contextual sparsity of attention weights. SSA uses a query temperature to adjust the sparsity for each query and its context position, and a value temperature to suppress irrelevant or noisy tokens. Experimental results show that SSA reduces perplexity in language modeling and improves optimization, with only a slight increase in parameters.", "questions": "1. How does SSA's performance vary with different parameter numbers? With more parameters, attention weight dilution might be alleviated, potentially weakening SSA's effect.\n2. There is no detailed description of the accuracy mentioned in Table 3. Could you clarify what this accuracy refers to?", "limitations": "Please refer to the weaknesses.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written and easy to follow.\n2. The motivation behind the method is clear and reasonable.", "weaknesses": "1. The evaluation is limited to perplexity and accuracy in language modeling; conducting experiments on downstream tasks would provide a more comprehensive evaluation.\n2. The perplexity improvement over the original models is modest.\n3. The paper lacks comparisons with other selective or sparse attention methods from the existing literature.\n4. The claim that SSA benefits long text scenarios lacks supporting evidence. There are no studies on performance relative to text length.\n5. There is no detailed description of the datasets used, such as average text length or the number of training instances, which could help explain performance differences.\n\nMinor comment:\nIn the abstract, it should be Q and K rather than V and K.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**W1,2: Conducting experiments on downstream tasks, The perplexity improvement over the original models is modest.** \n\nAs 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. \n\nWe also ran an experiment on the **passkey retrieval** task introduced in [[Mohtashami and Jaggi NeurIPS’23]](https://arxiv.org/pdf/2305.16300). 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.\n\n**W4: The claim that SSA benefits long text scenarios lacks supporting evidence.**\nIn 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. \n|                           | 128  | 4096 |\n|---------------------------|------|------|\n| **$P_{\\text{max}}$**      |      |      |\n| Pythia                    | 0.38 | 0.24 |\n| Pythia +SSA               | 0.41 | 0.36 |\n| **$H$**                   |      |      |\n| Pythia               | 2.83 | 4.86 |\n| Pythia +SSA          | 2.76 | 3.02 |\n\n\n**W5: There is no detailed description of the datasets used.**\n\nWe have provided all implementation details in Appendix A. We will add a remark in the revised version.\n\n\n**Q1: How does SSA's performance vary with different parameter numbers** \n\nWe 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.\n\n\n**Q2: There is no detailed description of the accuracy mentioned in Table 3. Could you clarify what this accuracy refers to?**\n\nLAMBADA 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]\n\n**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.”\n\n**Target sentence**: “Do you honestly think that I would want you to have a ____ ?”\n\n**Target word**: miscarriage.\n\nAccuracy is evaluated based on the model's ability to correctly predict the target word.\n\n[1] Paperno, Denis, et al. \"The LAMBADA dataset: Word prediction requiring a broad discourse context.\" arXiv preprint arXiv:1606.06031 (2016).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "atJxj2bl8f", "reviewer": "Reviewer_TP8j", "summary": "The authors introduce a modification to the widely-used self-attention mechanism, which modulates the queries and values in the attention computation with a per-token temperature parameter. They demonstrate that this allows the attention layer to model certain observations better for toy problems, and they also adapt several large language models to use this Selective Self-Attention (SSA), demonstrating that it moderately boosts overall perplexity & accuracy scores.", "questions": "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? Or can we instead argue this by looking at model scaling work that shows that a lot more parameter scaling is required to achieve similar drops in perplexity?", "limitations": "Attention is used extremely broadly, and the only real-world task used to evaluate SSA was the language modeling task. I would mention this as an important limitation on the claims of the paper.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The technical contribution represents an interesting and intuitive addition to standard attention. The addition of learnable query and value temperature clearly benefits the attention mechanism, both on illustrative toy examples and on real-world data. The mathematical exploration is comprehensive, leaving the reader with solid reasoning for why SSA solves several problems. The approach not only improves the performance of attention on standard modeling problems, but promises to improve on long-context tasks and positional encoding as well.", "weaknesses": "My main concern is the claim that the performance improvement on the LM task is due to \"strategic integration of SSA\" rather than an increase in parameter count. The improvements in perplexity/accuracy look statistically significant, but I don't think it's obvious that the improvement couldn't come from the increase in parameter count. There at least needs to be more argumentation to support this claim.\n\nAnother weakness is that the mathematical sections are pretty dense at times. Terms and mathematical propositions were introduced quickly without a lot of background. I found it difficult to understand the motivation for the lemma and propositions. I also think the denoising task, \"feature imbalance\", \"effective weights\", and \"specificity\" needed more introduction.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time, effort, and positive assessment.\n\n**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?**\n\nIn 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.\n\n| Model                      | Wikitext ppl ↓ | Lambada_std ppl ↓ | Lambada_openai ppl ↓ | Lambada_std acc ↑ | Lambada_openai acc ↑ |\n|----------------------------|--------------------|-----------------------|--------------------------|-----------------------|--------------------------|\n| **Finetune**               |                |                   |                      |                   |                      |\n| Pythia                     | 28.781         | 51.620            | 25.173               | 0.373             | 0.497                |\n| +SSA base (+5.26%)        | 26.514         | 47.945            | 23.956               | 0.388             | 0.513                |\n| +LoRA(+5.63%)       | 27.944         | 49.852            | 24.653               | 0.375             | 0.501                |\n|+SSA shared $W_{k/q/v}$ (+0.47%)|26.681|47.996|24.102|0.383|0.504|\n|+SSA feature-based (<0.01%)|27.048|48.906|24.814|0.379|0.499|\n\n**W2: The mathematical sections are pretty dense at times.** \n\nThank you for this suggestion. We will make sure to add additional explanation and motivation in the text.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
## Common Response

We thank all reviewers for their detailed and constructive feedback. To recap, our method Selective Self-Attention (SSA) enhances the approximation capability of attention through a learnable temperature scaling (TS) parameter. We theoretically establish the benefit of choosing temperature in a position- and token-aware fashion and evaluations also corroborate the value of SSA. Most reviewers found the method to be intuitive and interesting. They also raised excellent questions on the practical viability. Here, we address their shared concerns and demonstrate how SSA can indeed benefit state-of-the-art language models with negligible compute and parameter overhead.

Our new experiments focus on three aspects:

- (1) Improving parameter efficiency of SSA

- (2) Demonstrating the benefit of SSA on additional benchmarks

- (3) Benefit of SSA on larger LLMs such as Llama3-8B.

The full tables are in the pdf attachment. We also add partial tables in the text.

**(1) Making SSA parameter-efficient:** Our base approach adds ~5% parameters to the model. Reviewers NV9t, TP8j, and qMbv motivated us to examine more efficient parameterization to push this overhead below 0.5%. In Table 1, we tried several approaches by utilizing weight sharing and feature-based temperature selection. 

|Model|wikitext ppl↓|lambada_std ppl↓|lambada_openai ppl↓|lambada_std acc↑|lambada_openai acc↑|
|-|-|-|-|-|-|
|**Fine-tune**||||||
|Pythia|28.781|51.620|25.173|0.373|0.497|
|+SSA base (+5.26%)|26.514|47.945|23.956|0.388|0.513|
|+SSA shared $W_{k/q/v}$ (+0.47%)|26.681|47.996|24.102|0.383|0.504|
|+SSA feature-based (<0.01%)|27.048|48.906|24.814|0.379|0.499|
|**Pretrain**||||||
|Pythia|27.943|75.487|34.406|0.279|0.351|
|+SSA base (+5.26%)|26.912|72.891|33.126|0.294|0.360|
|+SSA shared $W_{k/q/v}$ (+0.47%)|27.046|73.071|33.814|0.291|0.360|
|+SSA feature-based (<0.01%)|27.281|73.614|33.794|0.287|0.359|

Above the parameter-efficient approaches are:

- **Shared $W_{k/q/v}$:** We use the attention weights $W_{k/q/v}$ as the input MLP layer of the temperature module. This  reduces the parameter count **below 0.5% (10x fewer) while maintaining on-par performance as full SSA**. Here, SSA only adds the output layer of the MLP, a vector with few parameters. Note that, this also makes the inference overhead negligible.

- **Feature-based:** Rather than learning an MLP, we tune a single scalar based on **token level statistics**. Specifically, we set the temperature as a function of the occurrence frequency of the token within the corpus. This is inspired by the logit adjustment strategy of [[Menon et al. ICLR’21]](https://arxiv.org/pdf/2007.07314) which sets the cross-entropy temperature as a function of class frequencies. This approach adds only *constant parameters (<0.01% overhead)*.

Finally, all other parameterizations we tried (in the pdf) also outperform the standard transformer, underscoring the value of selectivity irrespective of its precise implementation.

**(2) Evaluations on additional benchmarks:** Based on Reviewers NV9t and qMbv, we ran new evaluations in Tables 2 and 3. Table 2 provides results on 5 new benchmarks. We find that adding SSA uniformly improves the accuracy and that parameter-efficient SSA methods perform on par with baseline SSA. The improvement is particularly substantial on Piqa, Hellaswag, Arc-E, and Arc-C benchmarks. Notably, feature-based SSA, which has essentially no overhead, also works well across the board.

|Model|Piqa acc↑|Hella acc_norm↑|Winogrande acc↑|Arc-E acc↑|Arc-C acc_norm↑|
|-|-|-|-|-|-|
|Pythia|0.630|0.318|0.498|0.401|0.219|
|+SSA base (+5.26%)|0.661|0.359|0.508|0.426|0.230|
|+SSA feature-based (<0.01%)|0.656|0.350|0.500|0.417|0.224|

Table 3 examines the accuracy on the **passkey retrieval** task as defined in [[Mohtashami and Jaggi NeurIPS’23]](https://arxiv.org/pdf/2305.16300). This is a synthetic task to measure 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 could better solve this task by assigning different token-level temperatures to digits vs words.

|Pythia|+SSA base (+5.26%)|+SSA feature-based (<0.01%)|
|-|-|-|
|56.89|74.41|72.78|

**(3) Evaluations on larger LLMs:** Following Reviewers NV9t and Lp6e, we ran additional experiments with Llama3-8B and Pythia 410M as summarized in Table 4. These demonstrate that SSA improves the accuracy of larger and more recent models as well (e.g. compared to our main evals on GPT2, Pythia-160M, Llama2).

|Model|Wikitext ppl↓|Lambada_std ppl↓|Lambada_openai ppl↓|Lambada_std acc↑|Lambada_openai acc↑|
|-|-|-|-|-|-|
|**Fine tune**||||||
|Llama3-8b|12.416|24.002|13.954|0.481|0.684|
|Llama3-8b+SSA|10.982|22.671|11.052|0.489|0.690|
|**Pretrain**||||||
|Pythia-410m|22.516|69.814|32.781|0.321|0.371|
|Pythia-410m+SSA|21.980|69.041|31.458|0.331|0.384|

**Summary:** We hope that these new evaluations address the concerns of the reviewers and further corroborate the value of SSA. We look forward to engaging with the reviewers during the discussion period and thank them again for raising excellent points.

Meta-review: This paper applies temperature scaling to K/Q/V, and by that improves control over sparsity and relevance, resulting in efficiency gains. Reviewers found the work of interest and clear, and the results promising and potentially useful. Concerns regarding the experimental setup, overhead, and some claims not being backed up by experiments have largely been addressed by the authors.
Decision: Accept (poster)