QUEST: A robust attention formulation using query-modulated spherical attention
Authors: Hariprasath Govindarajan
Per Sidén
Jacob Roll
Fredrik Lindsten
Venue: ICLR 2026 Poster (2026)
Keywords: attention
transformers
model robustness
Primary area: foundation or frontier models, including LLMs
TL;DR: New attention formulation obtained by normalizing only the keys. Produces stable trainings, improved performance and robustness.

Abstract: The Transformer model architecture has become one of the most widely used in deep learning and the attention mechanism is at its core. The standard attention formulation uses a softmax operation applied to a scaled dot product between query and key vectors. We explore the role played by norms of the queries and keys, which can cause training instabilities when they arbitrarily increase. We demonstrate how this can happen even in simple Transformer models, in the presence of easy-to-learn spurious patterns in the data. We propose a new attention formulation, QUEry-modulated Spherical aTtention (QUEST), that constrains the keys to a hyperspherical latent space, while still allowing individual tokens to flexibly control the sharpness of the attention distribution. QUEST can be easily used as a drop-in replacement for standard attention. We focus on vision applications while also exploring other domains to highlight the method's generality. We show that (1) QUEST trains without instabilities and (2) produces models with improved performance (3) that are robust to data corruptions and adversarial attacks.

Published as a conference paper at ICLR 2026
QUEST: A ROBUST ATTENTION FORMULATION USING QUERY-MODULATED SPHERICAL ATTENTION
Hariprasath Govindarajan" 1.2 Per Sidén?
Jacob Roll? Fredrik Lindsten 'Linköping University, Sweden 2 Qualcomm Auto Ltd Sweden Filial
[hargov, psiden, jrolll @qti qualcomm. com, fredrik. 1 lindsteng wiw.stteeeneieeeeee
ABSTRACT
The Transformer model architecture has become one of the most widely used in deep learning and the attention mechanism is at its core. The standard attention formulation uses softmax operation applied scaled dot product between query and key vectors. We explore the role played by norms of the queries and keys, which can cause training instabilities when they arbitrarily increase. We demonstrate how this can happen even in simple Transformer models, in the presence of easy- to-learn spurious patterns in the data. We propose new attention formulation, mudry-modulated Spherical a a'Ttention (QUEST), that constrains the keys to a hyperspherical latent space, while still allowing individual tokens to flexibly control the sharpness of the attention distribution. QUEST can be easily used as drop- in replacement for standard attention. We focus on vision applications while also exploring other domains to highlight the method's 's generality. We show that (1) QUEST trains without instabilities and (2) produces models with improved performance (3) that robust to data corruptions and adversarial attacks.
INTRODUCTION
Transformers (Vaswani et al. 2017) are one of the most widely used model architectures across many domains in recent times. Each domain has adapted Transformers to build domain-specttn variants such as Vision Transformers (ViT) (Dosovitskiy et al. (2020) in computer vision, GPT (Radford et al. 2018) in natural language processing. P and Conformer (Gulati et al.) (020) in speech recognition, to name few. A core building block of the Transformer that is common across such variants the attention mechanism (Bahdanau et al.. 2015 Britz et al. 2017 Luong et al.) 2015). Although several different variants of the Transformer have been developed, they commonly use the vanilla attention mechanism consisting of scaled dot-product followed by softmax operation. Despite their success, training Transformer models can be challenging due to training instabilities (Chowdhery et al. 2023)] Dehghani et al.)
Zhao et al., 2021 for 3D pointeloud data
2023 Li et al.
2022a Wortsman et al.. 2024 Davis et al., [2021;Lin Liu al. 2020b] /hhii al. [0233] Many training
techniques consisting of initialization methods (Kedia et al. 2024 Huang et al. 2020). specific
hyperparameter schedules (Liang al. [2022] Kobayashi et al. 2024), normalizations (Dehghani
ct al. 2023 Wang et al. 2019 Xiong et | and optimization strategies (Qi et al.
2025) have
been developed to mitigate these issues. While this has limited such training instabilities great extent, why this issue occurs is still not fully understood (Hajra 2025).
We study the scaled dot-product attention formulation and identify the roles of different components of this attention. We find that the arbitrary vector norms of the queries and keys may potentially cause the exploding attention logits, that is known cause training instabilities (Zhai et al.. 2023 Dehghani et al. 2023). Through toy example, we demonstrate a scenario where the norms of these vectors increase and can lead to the model being stuck at a suboptimal solution. Even in stably trained models, we show that the model can concentrate attention on a few tokens instead of relying all relevant tokens (see Figure [ for an example using a ViT). We propose a new formulation, called Query- modulated Spherical Attention (QUEST), that displays improved training robustness across different hyperparameters. QUEST is a very simple modification of the standard attention and can casily be used drop-in replacement in any Transformer. Nevertheless, through extensive experiments

---

Published as a conference paper at ICLR 2026
on Transformers used in different domains, show that our proposed attention formulation bring consistent performance gains by learning robust patterns in the data. This is further reflected in improved robustness to adversarial attacks and data corruptions.
2 ATTENTION, PLEASE!
The attention mechanism is the core com- ponent of Transformers and operates set- like or sequential data. In this section, we first present background on scaled dot product attention (SDPA), which is the most widely used formulation of attention. Then, we provide an interpretation of the different components of this form of attention. Using attention formulation and finally, we use a toy ex- ample to demonstrate how it improves over SDPA and other related attention variants.
Image
QUEST
Standard
Pred: Macaw
Pred: Macaw
that motivation,
propose
Scaled dot product attention: In the self- attention paradigm, the attention mecha- nism transforms an input sequence (of length N) X E ND to an output se- quence E RNNDD, Each item in this sequence is referred to as token. Intu- itively, this can be viewed as a process of relating each input token to the other to- kens and aggregating some relevant infor- mation from these related tokens. In multi- head self-attention, this is repeated over sev- eral heads, to obtain H different outputs Z₁ E N Dy, where DH = D/H. The output Zí, obtained as:
Pred: Macaw
Pred: Macaw
Pred: Macaw
Pred: Spoonbil
Figure I: Class-activation maps for an image from the Macaw class in ImageNet, generated using AG-CAM (Leem & Seo 2024). Standard attention concentrates on few bird instances (see first row) and mis-classifee the image if the region containing those instances is noised (see third row). This indicates that the birds in the bottom half of the image do not contribute to the correct prediction when the top part of the image is noised, the model focuses on the birds in the bottom part of the image since they the most salient object in the image then, but results in a misclassification QUEST attention attends evenly different bird instances and classifies the image correctly even if some of the bird instances are noised. A more diverse attention can make the models more robust to input data variations, which can be observed in the improved model robustness 4.1.2
Z An V. softmax KQ Vn.
where the queries Qn. keys K₁ and val-
standard attention. Hence,
ues V₁ for head h are obtained as Q XWA K.h and V. Vi respectively. The softmax is ap- plied row-wise to the matrix CQ K Typ- ically, the scaling factor a constant, C 1/vDH. For the sake of brevity, we will ignore the head index h from the notation henceforth and denote vanilla attention as: A softmax (CQKT).
X TWT
X XW Q. K₁
2.1 AN ALTERNATIVE VIEW ON ATTENTION
Let v |/o where || is the norm of the vector and ø is unit vector. The attention corresponding to token can be written as:
A; softmax (Cq. K = softmax (Cllg: q: K
exp[C
q. K: (q. k.)] q. ky
exp
C C
KN
(( kx)] ky
exp
(g. ky)]
exp
q
[1(
ky)]
Althoug we focused on self-attention, the interpretation in this section and the proposed QUEST attention also applicable to other attention paradigms like cross-atiention. We explore this in A.5.T
2

---

Published as a conference paper at ICLR 2026
It can be noted that the norms of the queries and keys perform distinct roles in attention. The query norm | scales all the attention logits and thus, controls the sharpness of the attention distribution for that token. A higher query norm results in sharper attention and focuses on fewer tokens whereas a smaller query norm results in a softer distribution, aggregating information from a larger number of tokens. The dot product (q: kj) denotes the similarity vector alignment between each pair of query and key tokens. The key norm | controls the contribution to "general" attention from the key J. For query token that is uniformly distributed in the query-key latent space, the queries are more likely to attend to key tokens that have higher norms. The token that gets the most attention from a query token depends on a combination of its vector alignment and its key norm, [ [[ ky).
During training, if the information from a token helps reduce the loss objective, its key norm and the query norm of the token (e.g. CLS token) that aggregates information from that key grow larger. This can result in an attention logit explosion and attention collapse as noted by ZZai et al. (2023). But this information can sometimes be a spurious correlation. In the attention operation, higher key norms increase the attention towards that key token and reduce the attention to other key tokens. The gradients through the attention operation are weighted by the attention probabilities (Katz & Wolf 2025). Hence, tokens with lower key norms and thereby lower attention probabilities, also contribute less to parameter updates. Further, the parameter updates to the queries depend on a linear combination of the keys and vice versa. This indicates a cross-play where large key norms can further cause related query norms to increase, potentially leading to attention entropy collapse which contributes to training instabilities. We provide a more detailed theoretical analysis of this through the gradient updates corresponding to the queries and keys in A. Models initially learn features which are casy-to-learn, which can sometimes be spurious. If attention solely focuses on these features to solve the task, it is harder for the model to unlearn these spurious features and learn other useful features by attending to other tokens. We demonstrate this using toy example in the section below, where we study the training of simple Transformer model by introducing some spurious patterns in the training data.
In standard attention Vaswani et al. 2017), the constant scaling factor C. = 1/ /D was proposed to prevent large attention logit values, which were observed when DH was large. However, recent efforts to scale Vision Transformer models have shown that this formulation can still be unstable for large Transformer models due to arbitrarily growing attention logits (Dehghani et al. 2023). The proposed solution was to &g normalize both the queries and keys and scale each feature dimension in the queries and keys by learnable parameters (unique to each layer but shared across the heads), C, Cg E R Par We refer to this as QKNorm-DS attention. Another related formulation, QKNorm-HS (Liu al. 2022) instead scales each head with a learnable scalar C E R. DS and HS denote dimension and head scaling, respectively. We list the limitations of these attention variants below:
1. Standard attention is known to have training instabilities arising from large attention logits (Zhai et al. (003). Arbitrarily increasing key and query norms is one mechanism through which this happens.
2. QKNorm attentions (both QKNorm-HE and QKNorm -DfoommDDDD-- enable stable training but scale all tokens all heads by the same scaling factor which limits the expressivity of attention since all tokens are constrained to have the same sharpness.
A natural middle ground is to normalize either queries or keys. This would break the cross dependence between query (key) norms and key (query) gradient, with the potential of stabilizing the training. Neither of these options have however, to the best of our knowledge, been proposed in the literature. We will evaluate both options below, but what we propose is a new formulation of attention obtained by normalizing the keys while keeping the queries unnormalized. The intuition behind this choice is to allow each token to individually control the sharpness of its softmax distribution, and to prevent that large key norms "steal attention globally". We call this attention formulation, Query-modulated Spherical Attention (QUEST), and compute attention as: A = softmax (QK"). where K denotes lg fg-normalized keys. Note that we do not use any additional scaling, i.e. C 1. This is an easily interpretable attention variant where the rank order of the attention distribution is purely defined by the vector alignment between queries and keys in the hyperspherical latent space (cosine similarity). The query norms allow each query to independently control the sharpness of its attention.
3

---

Published as a conference paper at ICLR 2026
2.2 SPURIOUS ATTENTION ISSUES
We construct a simple toy example to demonstrate how standard attention can get stuck on spuriously correlated data patterns and find difficult learn the true and more consis- tent patterns in the data. Given a sequence of N vectors X = "el,...,.N'' where æ; E RDD the task is to retrieve information (an answer) from one of the vectors in [f1....,,,, The vectors consist of two parts: real-valued vector æ; and a one-hot encoded vector x.. All the real-valued vectors are sampled from a certain distribution except the one random answer location L. A correctly learned model should learn to identify this "out of distribution" vector zÿ and extract the answer æy, from its location. This is a robust signal that is always true but we introduce biased signal that also solve the task, but only for a subset of the samples (~50%).
TARGET
ANSWER
Specifically, independently for each training sample, let (nt(N((,,))) [t)) denote the location of the answer token and let Bernoullipp = 0.5) denotes the sample biased or not. For all non-answer tokens we sample x L ~N I)
(regardless of u), whereas for the answer token we sample atz (0, Σ) if u 0 (unbiased) or æt ~ (b, 0.11) if
Figure 2: Illustration of toy example.
" = (biased). Here, the bias vector b ~ (O Σ) is shared for all biased samples and * (details in A.3). Note that introduced an additional bias in the answer location (by sampling from a normal distribution with mean (4t and standard deviation o1) to ensure that the inputs are still biased after the addition of positional embedding. The answer part of the vectors, x!' are all sampled as uniform one-hot vectors over C' classes, but it is the one-hot vector at location L that is defined as the correct answer.
We consider a simple Transformer model = f(X; O) with parameters 0 which takes an input sequence X produce classification output y. We use a single Transformer layer (with learnale positional embeddings), which is sufficient to solve this task, and one head, which enables us to easily study the effect of different query and key norms. We use the [CLS] token features and use linear layer to produce the class logits. More implementation details of the Transformer model is provided in the appendix [A.3] We run this experiment with 5 different realizations of the data and 5 different weight initializations for the model weights. We train the model using the Adam W (Loshchilov) Hutter, 2018) optimizer for 50 epochs with a batch size of 32 and use learning rate values {0.0005, 0.001, 0.0025, 0.005, 0.0075, 0.01} } and weight decay values (0.0, 0.01, 0.02, 0.05, 0.1}.
QUEST (ours) [everall=58%]
Standard [overall=25%)
QNorm [overall=49%]
0.0005
0.48
0.48
0.52
Leaming rates
0.001
0.52
0.64
0.64
0.64
0.72
0
0.04
0
0.08
0
0.2
0.24
0.04
0.8
0.0025
0.84
0.92
0.92
0.8
0.92
0.68
0.6
0.68
0.84
0.84
0.8
0.8
0.88
0.84
0.8
0.005
0.84
0.88
0.84
0.84
0.6
0.68
0.72
0.6
0.68
0.28
0.76
0.92
0.96
0.92
0.68
0.4
0.0075
0.76
0.8
0.52
0.44
0.2
0.36
0.24
0.16
0
0.96
0.92
0.76
0.6
0.16
0.01
0.24
0.12
0.08
0.04
0.76
0.56
0.64
0.28
0.0
0.02
0.05
0.1
0.01
0.02
0.05
0.0
0.01
0.02
0.05
0.1
0.0
Weight decays
Weight decays
Weight decays
Figure 3: Success rates of learning the correct solution to the toy example. Models are trained with different hyperparameter combinations with 5 different weight initializations and 5 different realizations the data. The QKNorm methods obtained ~0% overall success rate and their results are available in Figure AT
Based on the performance on the training and test sets, we can categorize the learned models as degenerate (both training and test accuracy are random chance), biased (training accuracy of 50 ~ 80% but test accuracy 20 40%) and correct (training and test accuracy > 90%). The biased solution could learn a combination of the position and the vector bias to achieve a test accuracy 20 - 40%. We observed that degenerate solutions were common among QKNomm methods, which
4

---

Published as a conference paper at ICLR 2026
Epoch=0
Epoch=50
Epoch-O
Epoch 50
Blased Unbiesed
Biased Unblased
Count
1000
Count
400
Norm answer Loken keys
Norm answer Loken keys
Norm of answer taken keys
Norm of answer token keys
(a) Standard attention
(b) QNomm attention
Figure 4: Norms of answer key tokens for biased and unbiased samples: A common failure case for standard and QNorm attention involves the key norms of the biased answer token increasing as the training progresses. The model relies on looking up the bias vector to identify the answer.
is not surprising since they completely discard the information about the atypical distribution of the answer location contained in the norm of this vector. In general, the biased solution easier to learn and only requires the model to look up the answer corresponding to the biased target vector b. Specifically, the weight matrix for the keys will align with the vector b in the sense that the stretch (or amplification) factor is large in the direction of b, allowing the key norms grow, This will concentrate attention globally on the answer location, but only when the bias vector is present. In standard attention, this is accelerated by the cross- cross-play between query and key parameter updates. In QNorm, the queries are Ez-normalized, which helps in reducing this cross-play effect. Indeed, investigating the failed trainings of standard and QNomm attention, we found that the norms of the biased answer keys grew as the training progressed (see Figure 4). These models only learn to attend to the biased vector b but fail to learn the correct solution. QUEST mitigates this effect by ensuring that individual tokens are unable to "steal attention" globally. In Figure S we show the success rates of learning the correct solution using different hyperparameter combinations and different attention formulations. We observe that the proposed QUEST attention displays a better overall success rate of 58% that also works well across a wider range of hyperparameter values.
3 RELATED WORK
In this work, we focus on improving the core attention mechanism and, hence, restrict our discussion to other works that explored this. We consider this work as an orthogonal contribution other improvements in architecture design, training techniques and optimization. With the motivation of improving efficiency, class of linear complexity Transformers without the softmax operation have been proposed (Wang et al. 2020 Choromanski et al.) 2021 Kitaev et al.) [2020] Katharopoulos et al.) 2020 Han et ] 2023a) However, this work explores softmax-based attention only. Probabilistic interpretations of attention have connected attention to Nadaraya-Watee Watson regression with Gaussian isotropic kernels (Nguyen et al., 20226 Han et al.) 2023b), mixture models (Nguyen et al. 2022a), and asymmetric kernels (Chen et a 2023). Elliptical attention (Nielsen al.) 2024) a recent work that extends the Gaussian isotropic kernel interpretation of standard attention to hyper-
using a Mahalanobis metric. QUEST attention uses keys in the hyperspherical latent space but show that they are synergetic and can be combined as Elliptical QUEST. This uses an elliptical metric instead of cosine similarity between queries and keys (see A.5.2 for further discussion).
Prior works have studied the use of LayerNorm (Ba et al. 2016) in attention (Xiong et al.) 2020
and its positioning (Wang al. 2019). main contributor to the training instability of Transformer
models is the attention logit explexion (Zhai et al. 2023). This is directly related to the properties of queries and keys (Bao et al. 2024). Explicit by normalization of the queries and keys, proposed as QKNorm (Dehghani et al. 2023; Liu et al. 2022), is the closest to our work but were mainly aimed at scaling Vision Transformers. MMngaras al. (2025) used a similar formulation to QKNorm, but in the context of softmax-free attention. While this makes large models stable to train, we show that this limits the expressivity of attention, shown by worse performance on smaller models compared to standard attention. Our QUEST attention is applicable to Transformers in general, stable to train up to the scales that we have attempted and shows improved performance compared to QKNorm. Recent works on linear attention have also identified similar attention entropy collapse in the context of linear attention (Meng et al.) 20)5a]b). We leave the extension and study of our proposed QUEST attention to linear attention variants for future work.
5

---

Published as a conference paper at ICLR 2026
Table I: Ablation of different QK normalization methods for ImageNet classification (ViT-Tiny model trained using DeiT for 300 epochs) [ et al. (2022). DDehghani et al. (2023)). The Top-1 validation accuracies for ImageNet are the means over independent training runs. The detailed statistics of these results are reported in AT
Attention
Scaling
IN-val Top-1
IN-v2 Top-1
IN-ReaL. Top-1
IN-C MCE 4
IN-A
Top-1
Top-5
Standard
1/VDi
72.6
60.6
80.4
55.7
8.2
32.9
QUEST Qorm
73.4
61.1
81.2
55.0
8.5
34.6
72.7
60.8
80.6
55.3
8.2
34.5
QKNorm-HB
CERUxH
72.5
60.8
80.5
56.4
7.9
33.3
QKNorm-DS
Cy. C E RLxDr
71.6
59.7
79.6
57.4
7.2
31.5
QKNorm
Cq. Cs E R RL
71.9
59.3
79.0
58.1
7.0
31.0
EXPERIMENTS
We primarily focus on applications using Vision Transformers but also conduct experiments on Transformers used in other domains such as language modeling, graph Transformers, general time series and pointclouds (pointcloud segmentation experiment is included in the appendix A.5.7). Our goal is to demonstrate broad applicability and effectiveness offered by simple modification. We therefore study the impact of replacing standard attention with QUEST in popular attention-based architectures in different settings across multiple domains. Further, provide ablation experiments comparing QUEST with QNomm and QKNorm attentions in the image classification, language modeling and time series experiments.
4.1 VISION APPLICATIONS
4.1.1 CLASSIFICATION
We conduct image classification experiments using the DeiT (Touvron et al. 202 12 training method. Firstly, we perform ablation experiments different QK-normalization patterns in attention by training ViT-Tiny model on ImageNet- 1 K dataset for 300 epochs and report the results in Table I In addition evaluating on the ImageNet validation set, we also evaluate on the validation data from ImageNet-v2 (Recht ct al., 2019), ImageNet-Reaal (Beyer et al. 22020). ImageNet-Adversarii (Hendrycks al. 2021) and
Ablation:
(Hendrycks & Dietterich 2019). We report validation accuracies on all datasets except for IN-C, which is evaluated using the mean corruption error (MCE) over 16 corruptions. We observe that the proposed QUEST attention performs clearly better than standard attention but also compared to alternative methods to normalize queries and keys. While NKNmmmDDS shown to be stable for larger ViTs, we observe that performs worse than standard attention in smaller models it limits the expressivity of attention. Further, in Table AT we show that these observations are statistically significant by repeating this experiment three times. A theoretical discussion about how QUEST able to perform favorably, based on the gradients in the optimization process, is provided in appendix A On this basis, we consider standard attention as the baseline in other experiments. If training is unstable with standard attention, then we use other QKNorm variants as the baseline.
Next, we consider experiments larger sizes of Vision Transformers (Small, Base and Large). DeiTT training with standard attention is unstable and divergent for ViT-Base and Large. In those cases, we use QKNorm-DS attention the baseline. The shorter 100 epoch trainings on ViT-Base and ViT-Large models show that QUEST attention is stable to train in contrast to standard attention and also achieves better performance compared QKNorm attention (see Table A6 for an additional experiment on a 2B parameter ViT which shows consistent results). We empirically analyze the progression of maximum attention logits and the corresponding query/key norms in Figure A6 Although only the key tokens were normalized in QUEST, we observe that this also stabilizes the query norms and consequently the attention logits, thus enabling stable training. DeiT-3 proposed an improved and stable training recipe for larger ViTs by using better augmentation strategies and techniques such as stochastic depth (Huang et al.) 2016) and LayerScale (Touvron et al.) (021b)) We also evaluate QUEST attention using DeiT-3 training for the Base, Large and Huge ViT', that
DeiT and DeiT-3:

---

Published as a conference paper at ICLR 2026
Table 2: ImageNet Classification using DeïT and DeiT-3 training recipes [*Dehghani et al. (2023) ].
For longer trainings with larger ViT models (denoted as 1), the batch size and input image resolution in the first phase of training are changed compared to Touvron et al. (2022) in order to train on a
limited compute infrastructure, see A.5.1 for details.
Model
Attention
Epochs
IN-val Top-1
IN-v2 Top-1
IN-ReaL. Top-1
IN-C MCE
IN-A Top-1
Top-S
DeiT-1 1
Touvron et al.
(2021a)
trainings
ViT-S/16
Standard
200
79.6
68.4
85.8
44.8
18.2
50.7
ViT-S/16
QUEST
200
80.2
68.9
86.2
43.2
20.4
53.5
ViT-B/16
Standard
100
training crashed-
ViT-B/16
QKNorm-DS*
100
79.0
67.5
84.9
44.4
17.7
50.0
ViT-B/16
QUEST
100
79.7
68.7
85.7
42.9
19.2
51.9
VIT-L/16
Standard
100
training crashed-
ViT-L/16
QKNorm- D
100
72.5
58.5
78.2
54.4
8.9
31.2
VIT-L/16
QUEST
100
74.9
61.4
80.6
50.3
11.1
35.6
DeiT-3 Touvron ct al.
2022)
trainings
ViT-B/16
Standard
400 20
82.7
72.4
88.0
36.5
34.8
67.7
ViT-B/16
QUEST
400 20
83.2
73.3
88.2
35.7
37.6
69.9
ViT-L/16
Standard
400 20
83.9
74.0
88.7
32.5
44.1
75.3
ViT-I/16
QUEST
400 20
84.1
74.3
88.9
32.3
44.5
76.0
VIT-H/14
Standard
400 20
83.2
73.5
88.6
34.1
45.9
77.7
ViT-H/14
QUEST
100 + 20
83.4
74.0
88.7
33.7
46.2
78.0
were training for a larger number of epochs until convergence. The results are reported in Table 2 For the considered model sizes, we find that QUEST consistently achieves better performance. We observe larger performance improvements in IN-C and IN-A evaluations, showing that QUEST attention produces more robust models. We further explore this through experiments on adversarial robustness and explainability in 4.1.2 and 4.1.4 We show additional experiments in A.5. to evaluate the sensitivity of these results to changes in the training data and hyperparameters like learning rate. We found the performance improvement be consistent when training with different limited data subsets. The models with QUEST attention are stable to train at different learning rates.
Additional experiments: We experiment with an alternate Transformer architecture, CrossViT (Chen et al. (2221). to demonstrate that QUEST attention is compatible with cross-attention (see A.5.1). Since commonplace to use self-supervised models, we also evaluate QUEST attention for pre-training in A.5.9 and find that it improves over standard attention in downstream performance.
4.1.2 ROBUSTNESS
In addition improved robustness corruptions and adversarially curated images observed above, we evaluate the robustness of QUEST attention to adversarial attacks. For adversarial attacks under image perturbations, we adopt the experimental setup Nielsen al. (0024) We report the validation performance under the following adversarial attacks: fast gradient sign method (FGSM) (Goodfellow) al. 2015), PGD attack based on projected gradient descent (Madry al. 2018), SPSA attack based on simultaneous perturbation stochastic approximation (Uesato et al.) 2018) and Auto attack (Croce Hein 2020). The Auto attack is an ensemble of auto PGD-Cross Entropy, auto PGD-targeted fast adaptive boundary-targeter and Square attacks. The complete experimental details are provided in A.5.2 We report the validation accuracies and NL.,, after adversarial perturbations for the ViT-Tiny (see Table A for similar results for larger VIT models). Firstly, we observe that QUEST more robust than standard attention across adversarial attacks. QUEST attention relevant object regions whereas standard attention concentrates on only few object parts or object instances (see further discussion below in 4.1.4). Elliptical attention (Nielsen et al. 2024) is SOTA model for robustness but this is achieved at the expense of classification accuracy (7 77.53% vs 72.50%). We show that QUEST can be orthogonally added to Elliptical attention to further improve robustness, while also achieving better classification accuracy clean data.
model in Table 3 and and
focuses more evenly
7

---

Published as a conference paper at ICLR 2026
Table 3: Robustness to adversarial attacks using ViT-Ti/16 model trained using DeiT
Attention
Clean Data
FGSM Top-5
PGD Top-5
Top-1
Top-5
NLL
Top-1
NLL
Top-1
NLL
Standard QUEST
72.50
91.45
1.190
54.23
85.28
1.827
43.65
78.18
2.503
73.33
91.91
1.160
56.90
86.63
1.745
45.26
79.33
2.448
Elliptical Elliptical- BUEittcal-QUEST
71.53
90.70
1.254
55.96
85.53
1.746
46.30
80.05
2.231
72.48
91.20
1.214
56.39
85.94
1.741
47.25
80.61
2.211
Table 4: Robustness to adversarial attacks using ViT-Ti/16 model trained using DeiT
Attention
SPSA Top-5
Auto Top-5
Top-1
NLL
Top-1
NLL
Standard QUEST
47.17
83.95
2.027
26.57
67.60
3.230
50.70
85.49
1.904
27.29
67.98
3.200
Elliptical
57.96
86.92
1.654
27.35
67.28
3.014
Elliptical- QUEST
59.15
87.44
1.613
28.54
68.10
2.965
4.1.3 SEGMENTATION
We evaluate image segmentation using the Segmenter approach (Strudel 2021)) using a Mask Transformer decoder, where we initialize the backbone ViT model with the DeiT weights from above and finetune the entire model (encoder and decoder) for semantic segmentation on the ADE20K dataset (Zhou et al., 2019 2017). We also evaluate the segmentation models for bustness under 16 different types of im-
Table 5: ADE20K Image Segmentation
Model
Attention
Clean data mloU
Corrupted data mloU
VIT-Ti/16
Standard QUEST
37.34
32.19
VIT-Ti/16
38.87
33.55
VIT-S/16
Standard
43.43
38.45
VIT-S/16
QUEST
44.13
39.19
age corruptions, following the experimen-
tal setup of Zhou et al. (022). The segmentation results are reported in Table 5 Based on the
commonly used mloU metric (mean Intersection over Union), QUEST attention performs better than standard attention and displays better robustness to corruptions.
4.1.4 EXPLAINABILITY
Attention based models are also beneficial from a model explainability standpoint. AG-CAM (Leem Seo 2024) is recent explainability method that combines attention maps and gradient information to produce class activation maps (CAMs) for image classifiers. Following the same evaluation protocol, we apply probability threshold of 0.5 to the CAMs and evaluate pixel accuracy, mloU and DICE score by comparing with ground truth localization labels in ImageNet. We consider the ViT-B model trained using the DeiT-3 training recipe and report the results for QUEST and standard attention in Table 6 In Figure 1 and 5 show how model trained with QUEST produces a better CAM for different classes. Specifically, standard attention concentrates on few object parts or instances whereas QUEST attends them more evenly (see A.5.3 for additional qualitative examples). By relying on only few aspects of an object, QUEST can perform more robustly as observed in
Table 6: Model explainability using AG-CAM
Backbone
Attention
Method
Epochs
Pixel accuracy (%) 1
mloU
DICE Score
VIT-B/16
Standard
DeïT-3
400 + 20
65.22 70.76
35.53
0.4870
VIT-B/16
QUEST
DeiT-3
400 + 20
53.54
0.6597
8

---

Published as a conference paper at ICLR 2026
4.2 TIME-SERIES CLASSIFICATION
We conduct general sequence classification experiments using the UEA multivariate time series classification suite (Bagnall al. 2018). We train Transformer models using the experimental setup in Wang et al. (2024) (see results in Table more details in A.5.6). With the exception of datasets, QUEST attention performs better or the same as standard attention. Through an ablation experiment in Table A8 we show that QUEST also outperforms QNorm and QKN based on the overall average performance. The overall average accuracy of QUEST attention surpassed the SOTA (73.17%) achieved by Crossformer Zhhng & Yan (023).
4.3 GRAPH TRANSFORMERS
Recently, several Transformer-basee models have been proposed for tasks involving graph-structured data. We adopt the experimental setup of GraphGPS (Rampášek al. 2022) evaluate on standard GNN benchmarks from Dwivedi et al. (2023) and long-range graph benchmarks (LRGB) from Dwivedi et al. (2022). We use the same optimal hyperparameter setups for each dataset as in GraphGPSS (see A.3.8 for detailed experimental setups for each dataset) and evaluate QUEST attention as a drop-in replacement for standard attention used in those graph Transformers. The results for the standard and LRGB benchmarks are reported in Tables 8 and 9
respectively. For the standard benchmarks, standard attention and QUEST perform on par for all datasets (within significance range). On the long-range benchmark, and PCQM-Contact, whereas we perform on par with standard attention for PascalVOC-SY and Peptides-struct.
significant improvements for COCO-SP, Peptides-func
Table 8: GNN Benchmarks using GraphGPS (mean + s.d over 10 runs)
Attention
ZINC
CIFAR10
PATTERN
CLUSTER
MAE І
Accuracy 1
Accuracy
Accuracy t
Standard QUEST
0.070 + 0.004
72.298 + 0.356
86.685 + 0.059
78.016 + 0.180
0.069 + 0.002
72.843 + 0.526
86.760 + 0.046
77.894 + 0.205
Table 9: Long Range Graph Benchmarks using GraphGPS (mean + s.d over 4 runs)
Attention
PascalVOC-SP
COCO-SP
Peptides-func
Peptides-struct
PCQM-Contact
Fl score f
Fl score f
AP
MAE +
MRR
Standard
0.375 + 0.011 0.373 = 0.003
0.341 + 0.004 0.349 + 0.004
0.654 0.004
0.250 + 0.001
0.334 + 0.001
QUEST
0.662 + 0.004
0.251 = 0.002
0.346 + 0.001
Figure 5: Class activation maps for Elephant and Zebra. Model with QUEST attention shows better coverage of the different instances of the animals than standard attention.
Table 7: UEA Multivariate Time Series Classifi- cation using Transformers
Dataset
Standard
QUEST
FthanolConcentraaioon
29.28
30.42
Image
QUEST
Standard
FaceDetection
65.24
65.83
Query. Zebra
Handwriting
42.00
49.18
Heartbeat
77.56
78.54
Japanese Vowels
98.38
98.38
PEMS-ST
83.82
80.92
SelfRegulationSCCI
88.05
88.05
Query: Elephont
SeliRegulationSCCP
58.89
58.33
SpokenArabicDigitss
98.86
99.41
UaaeeeestureLLbbrryy
86.88
86.25
Average
72.90
73.53

---

Published as a conference paper at ICLR 2026
4.4 LANGUAGE MODELING
We conduct language modeling experiments using the WikiText-103 dataset (Merity et al. 2017 and from|Nielsen use the experimental setup from Nielsen et al. (2024) and (Schlag et al. (021) for the Transformer model. We also consider larger model, the Transformer-XL (Dai et al. 2019) and evaluate it using their experimental setup. Additionally, we evaluate the robustness of these language models using the Word Swap Attack from Nielsen et al. (2024) that replaces random words with the "AAA" token with specified rates. The perplexity (PPL. metrics for these models using clean and contaminated data (different attack rates) are reported in Table [0, We observe that QUEST attention generally produces marginally better performance on both clean and contaminated data. When comparing QUEST with QNomm and QKNorm the Transformer-Meddum model, observe that QUEST performs favorably.
Table 10: Language Modeling on WikiText-103 103
Method
Model size (Parameters)
Attention
Clean Data PPL
Contaminated Data PPL (Coruption %)
Val
Text
Test (1.5%)
Test (2.5%)
Test (5.0%)
Transformer
Small (44M)
Standard
33.073
34.076
41.566
46.380
60.944
Transformer
Small (44M)
QUEST
32.966
33.966
41.522
46.299
61.062
Transformer
Medium (90M)
Standard
27.441
28.851
36.25 234
40.866
55.012
Transformer
Modium (90M)
QNemm
27 27.233
28. .8.688
36, 36.262
10.853
55.451
Transformer
Medium (9UM)
QKNom-DS
27.376
28
40.715
54.899
Transformee
Modium (90M)
QUEST
26.980
28.478
35.849
40.499
54.531
Transformer-Ke XI.
Base (15IM)
Standard
22.650
23.592
29.627
33.386
44.168
Transformee-XL
Base (151M)
QUEST
22.436
28.320
29.339
33.008
43.511
CONCLUSION
The instabilities in training Transformers are well known and occur when attention collapses due to arbitrarily increasing query and key norms. We demonstrate how spuriously correlated features in certain tokens can contribute to such behavior. Unlike prior works which argued that such issues only occurred in larger models, we showed that they can also limit small models, resulting in reduced performance even if the training does not diverge. We propose simple drop-in replacement called QUEST attention that considers a hyperspherical latent space for attention while still allowing individual tokens to flexibly and independently control the sharpness of their attention distributions. Through extensive experiments on several domains, we demonstrate broad applicability. QUEST attention trains more robustly and produces models that typically perform better and that are more robust to corruptions and adversarial attacks than standard attention and its variants like QKNorm. We also show that QUEST attention improves the robustness of the trained models and can orthogonally improve SOTA methods like Elliptical attention. While we have shown consistent improvements when using QUEST with Transformers in multiple domains, we have primarily focused on vision applications and leave more in-depth evaluation on other domains and incorporation of QUEST into comain--specffe SOTA architectures for future work. Given the effectiveness of operating in the hyperspherical latent space, exploring more geometrically aligned operations and optimization methods, e.g. Riemannian gradient descent (Kasai al. 2019), is another promising avenue for future work.
ACKNOWLEDGMMENTS
This research is financially supported by the Swedish Research Council (project no: 2024-05011), 1), the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation, and the Excellence Center Linköping Lund in Information Technology (ELLIIT). The computations were enabled by the Berzelius resource provided by the Knut and Alice Wallenberg Foundation at the National Supercomputer Centre.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We include an algorithmic implementation QUEST attention and standard attention in Algorithm/A] to clearly illustrate our proposed modification. In all experiments, use standard evaluation protocol and use publicly available code repositories. We introduced QUEST attention as a simple drop-in replacement for standard attention the different experiments that we conducted. We provide details in the Appendix (see|A.5) A.5) for all the code repositories used and clarify any changes to hyperparameter configurations.
REFERENCES
Mahmoud Assran, Mathilde Caron, Ishan Misra, Piotr Bojanowski, Florian Bordes, Pascal Vincent, Armand Joulin, Mike Rabbat, and Nicolas Ballas. Masked siamese networks for label-efficient learning. In ECCV, 2022.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.064550 2016.
Anthony Bagnall, Hoang Anh Dau, Jason Lines, Michael Flyun, James Large, Aaron Bostrom, Paul Southam, and Eamonn Kcogh. The UEA multivariate time series classification archive, 2018. arXiv preprint arXiv: 001818111000555 2018.
Dzmitry Bahdanau, Kyung Hyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. In ICLR, 2015.
Han Bao, Ryuichiro Hataya, and Ryo Karakida. Self-attention networks localize when gk- eigenspectrum concentrates. ICML, 2024.
Hangbo Bao, Li Dong. Songhao Piao, and Furu Wei. BEit: BERT pre-training of image transformers. In ICLR, 2022.
Lucas Beyer, Olivier Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and Aäron van den Oord. Are done with imagenet? arXiv preprint arXiv:2006.07159, 2020.
Lukas Bossard, Matthieu Guillaumin, and Luc Van Gool. Food-101 mining discriminative compo- nents with random forests. In ECCV, 2014.
Denny Britz, Anna Goldie, Minh-Thaaa Thang Luong, and Quoe Le. Massive exploration of neural machine translation architcctures. In EMNLP. 2017.
Holger Caesar, Varun Bankiti, Alex H Lang, Sourabh Vora, Venice Erin Liong, Qiang Xu, Anush Krishnan, Yu Pan, Giancarlo Baldan, and Oscar Beijbom. nuScenes: multimodal dataset for autonomous driving. In CVPR, 2020.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. ICCV, 2021.
Chun-Fu Richard Chen, Quanfu Fan, and Rameswar Panda. Crossvit: Cross-attention multi-scale vision transformer for image classification. In ICCV, 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. ICML, 2020.
Yingyi Chen, Qinghua Tao, Francesco Tonin, and Johan Suykens. Primal-attention: Self-attention through asymmetric kernel SVD in primal representation. NewrIPS, 2023.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Ben- jamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In ICLR, 2021.
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. JMLR, 2023.

---

Published as a conference paper at ICLR 2026
Mircea Cimpoi. Subhransu Maji, lasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi. De- scribing textures in the wild. In CVPR, 2014.
Pointcept Contributors. Pointcept: A codebase for point cloud perception research. 2023.
https
Francesco Croce and Matthias Hein. Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks. In ICML, 2020.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-Xe XL: Attentive language models beyond fixed-length context. In ACL, 2019.
Jared Q Davis, Albert Gu, Krzysztof Choromanski, Tri Dao, Christopher Re, Chelsea Finn, and Percy Liang. Catformer: Designing stable transformers via sensitivity analysis. In ICML, 2021.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In ICML, 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In CVPR, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020.
Vijay Prakash Dwivedi, Ladislav Rampāšek, Mikhail Galkin, Ali Parviz, Guy Wolf, Anh Tuan Luu, and Dominique Beaini. Long range graph benchmark. In NeurIPS Datasets and Benchmarks Truck, 2022.
Vijay Prakash Dwivedi, Chaitanya K Joshi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. JMLR, 2023.
Linus Ericsson, Henry Gouk, and Timothy M Hospedales. How well do self-supervised models transfer? CVPR, 2021.
Mark Everingham, Lue Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman. The PASCAL visual object classes (VOC) challenge. IJCV, 2010.
lan J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. In ICLR, 2015.
Hariprasath Govindarajan, Per Sidén, Jacob Roll, and Fredrik Lindsten. DINO as a von Mises-Fisher mixture model. In ICLR, 2023.
Anmol Gulati, James Qin, Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: Convolution-augmenn transformer for speech recognition. Interspeech, 2020.
Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo
Suvadeep Hajra. Short-range dependency effects on transformer instability and a decomposed attention solution. arXiv preprint arXiv:2505.15545 2025.
Dongchen Han, Xuran Pan, Yizeng Han, Shiji Song, and Gao Huang. Flatten transformer: Vision transformer using focused linear attention. In ICCV, 2023a.
Xing Han, Tongzheng Ren, Tan Nguyen, Khai Nguyen, Joydeep Ghosh, and Nhat Ho. Designing robust transformers using robust kernel density estimation. NeurIPS. 2023b.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. ICLR. 2019.
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. CVPR, 2021.
12

---

Published as a conference paper at ICLR 2026
Charles Tapley Hoyt, Max Berrendorf, Mikhail Galkin, Volker Tresp, and Benjamin M Gyori. A unified framework for rank-based evaluation metrics for link prediction in knowledge graphs. In Workshop Graph Learning Benchmarks at The WebConf, 2022.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In ECCV, 2016.
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In ICML, 2020.
Hiroyuki Kasai, Pratik Jawanpuria, and Bamdev Mishra. Riemannian adaptive stochastic gradient algorithms on matrix manifolds. ICML, 2019.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In JCML, 2020.
Shahar Katz and Lior Wolf. Reversed attention: On the gradient descent of attention layers in GPT. In NAACL, 2025.
Akhil Kedia, Mohd Abbas Zaidi, Sushil Khyalia, Jungho Jung, Harshith Goka, and Haejun Lee. Transformers get stable: An end-to-end signal propagation theory for language models. In ICML, 2024.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In ICLR, 2020.
Seijin Kobayashi, Yassir Akram, and Johannes Von Oswald. Weight decay induces low-rank attention layers. NeurlPS, 2024.
Saebom Leem and Hyunseok Seo. Attention guided CAM: visual explanations of vision transformer guided by self-attention. In Proceedings of the AAAI conference on artificial intelligence, 2024.
Conglong Li, Minjia Zhang, and Yuxiong He. The stability-effccinncy dilemma: Investigating sequence length warmup for training GPT models. NeurlPS, 2022a.
Fei-Fei Li, Marco Andreeto, Marc Aurelio Ranzato, and Pietro Perona. Caltech 101, 2022b. URL
Chen Liang, Haoming Jiang, Simiao Zuo, Pengcheng He, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Tho Zhao. No parameters left behind: Sensitivity guided adaptive learning rate for training large transformer models. In ICLR, 2022.
Tsung- Yi Lin, Michael Maire, Serge Belongie. James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft COCO: Common objects in context. In ECCV, 2014.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In ICLR, 2020a.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In EMNLP. 2020b.
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In CVPR, 2022.
Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in Adam. 2018.
Minh-Than Luong, Hieu Pham, and Christopher D Manning. Effective approaches to attention-basee neural machine translation. In EMNLP, 2015.
Aleksander Madry. Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. ICLR, 2018.
Subhransu Maji, Esa Rahtu, Juho Kannala, Matthew Blaschko, and Andrea Vedaldi. Fine-grained visual classification of aircraft. Technical report, 2013.
13

---

Published as a conference paper at ICLR 2026
Weikang Meng, Yadan Luo, Liangyu Huo, Yingjian Li, Yaowei Wang. Xin Li, and Zheng Zhang. Norm x Direction: Restoring the missing query norm in vision linear attention. arXiv preprint arXiv:2506.21137. 2025a.
Weikang Meng, Yadan Luo, Xin Li, Dongmei Jiang, and Zheng Zhang. PolaFormer: Polarity-aware linear attention for vision transformers. In ICLR, 2025b.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In ICLR, 2017.
Gabriel Mongaras, Trevor Dohm, and Eric Larson. Cottention: Linear transformers with cosine attention. In Intelligent Computing-Proceedung the Computing Conference, 2025.
Tam Minh Nguyen, Tan Minh Nguyen, Dung DD Le, Duy Khuong Nguyen, Viet-Anh Tran, Richard Baraniuk, Nhat Ho, and Stanley Osher. Improving transformers with probabilistic attention keys. In ICMI., 2022a.
Tan Nguyen, Minh Pham, Tam Nguyen, Khai Nguyen, Stanley Osher, and Nhat Ho. FourierFormer: Transformer meets generalized Fourier integral theorem. NeurlPS, 2022b.
Stefan Nielsen, Laziz Abdullaev, Rachel SY Teo, and Tan Nguyen. Elliptical attention. NeurIPS, 2024.
Maria-Elena Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. 2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing, 2008.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy V. Vo, Marc Szafraniee, Vasil Khalidov, Pierre Fernandez, Daniel HAZIZA, Francisco Massa, Alaaeldin El-Nouby, Mido Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Yao Huang, Shang-Went -We Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Herve Jegou, Julien Mairal, Patrick Labatut, Ar- mand Joulin, and Piotr Bojanowski. DINOv2: Learning robust visual features without supervision. TMLR, 2024.
Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and CV Jawahar. Cats and dogs. In CVPR, 2012.
Jordi Pont-Tuset, Federico Perazzi, Sergi Caelles, Pablo Arbeláez, Alexander Sorkine-Hornung, and Luc Van Gool. The 2017 DAVIS challenge on video object segmentation. 2017.
Xianbiao Qi, Yelin He, Jiaquan Ye, Chun-Guang Li, Bojia Zi, Xili Dai, Qin Zou, and Rong Xiao. Taming transformer without using learning rate warmup. In ICLR, 2025.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Do- minigue Beaini. Recipe for general, powerful, scalable graph transformer. NewrIPS, 2022.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do ImageNet classifiers generalize to ImageNet? In ICML, 2019.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In ICML, 2021.
Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation. In ICCV, 2021.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-cfficient image transformers & distillation through attention. In ICML, 202 la.
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In ICCV, 2021b.
14

---

Published as a conference paper at ICLR 2026
Hugo Touvron, Matthieu Cord, and Hervé Jégou. DeIT III: Revenge of the ViT. In ECCV, 2022.
Jonathan Uesato, Brendan O'donoghue, Pushmeet Kohli, and Aaron Oord. Adversarial risk and the dangers of evaluating against weak attacks. In ICML, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurlPS, 2017.
Qiang Wang, Bei Li, Tong Xiao, Jingho Zhu, Changliang Li, Derek Wong, and Lidia S. Chao. Learning deep transformer models for machine translation. In ACL, 2019.
Sinong Wang. Belinda Z Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprini arXiv:2006.04777 2020.
Yuxuan Wang, Haixu Wu, Jiaxiang Dong, Yong Liu, Mingsheng Long, and Jianmin Wang. Deep time series models: comprehensive survey and benchmark. 2024.
Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie E Everett, Alexander A Alemi, Ben Adlam, John D Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-Dickstein, Kelvin Xu, Jaehoon Lee, Justin Gilmer, and Simon Kornblith. Small-scale proxies for large-scale transformer training instabilities. In ICLR, 2024.
Xiaoyang Wu, Li Jiang, Peng-Shuai Wang, Zhijian Liu, Xihui Liu, Yu Qiao, Wanli Ouyang, Tong He, and Hengshuang Zhao. Point transformer v3: Simpler faster stronger. In CVPR, 2024.
Jianxiong Xiao, James Hays, Krista A Ehinger, Aude Oliva, and Antonio Torralba. SUN database: Large-scale scene recognition from abbey zoo. CVPR, 2010.
Jianxiong Xiao, Krista A Ehinger, James Hays, Antonio Torralba, and Aude Oliva. SUN database: Exploring large collection of scene categories. IJCV, 2016.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In ICML, 2020.
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M Susskind. Stabilizing transformer training by preventing attention entropy collapse. In /CML, 2023.
Yunhao Zhang and Junchi Yan. Crossformer:
Transformer utilizing cross-dimension dependency for multivariate time series forecasting. In ICLR, 2023.
Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip HS Torr, and Vladlen Koltun. Point transformer. In ICCV, 2021.
Bolei Zhou, Hang Zhao. Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ADE20K dataset. In CVPR, 2017.
Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ADE20K dataset. IJCV, 2019.
Daquan Zhou, Zhiding Yu, Enze Xie, Chaowei Xiao, Animashree Anandkumar, Jiashi Feng, and Jose M Alvarez. Understanding the robustness in vision transformers. In ICML, 2022.
Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong, Image BERT pre-training with online tokenizer, In ICLR, 2021.
15

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 THEORETICAL ANALYSIS OF ATTENTION GRADIENT UPDATES
In the context of Transformer gradient updates, Katz Wolf 2025
defined a reverse attention term,
R as follows:
È = =WW VT NRN
== A ET diag AET)
E RNNN
V DH
where A is the attention computed in the forward pass, W. is the output projection weight and A E RNDD is the Vector Jacobian Products (VJP) of W Note that when attention, A concentrates on specific tokens, the reverse attention R also concentrates on those tokens (contributions from other tokens approach 0). This term then used to derive the VJPs for the query and key gradient updates (for token j) as:
8 R E ERR PM
8. = R Q E R
The VJPs of the queries and keys are linear combinations the keys and queries respectively. These VJPs are dominated by the tokens where attention concentrated and the contributions from the other tokens are diminished. In the attention operation, higher key norms increase the attention towards that key token and reduce the attention to other key tokens. Hence, tokens with lower key norms and thereby lower attention probabilities, also contribute less to parameter updates. When key norms of these tokens grow, this can consequently cause the query norms attending to that token to grow as well. This cross-play causes query and key norms to feed off each other and continue growing, resulting in attention logit collapse.
Based on the above observation, fz-normalizing at least one of them can mitigate this effect. This also be empirically observed in Figure A6 in the paper note that the query norms and max logits in QUEST are stable and do not increase as in the case of standard attention. This provides an explanation as to why QUEST, QNomm and QKNorm are all able to produce stable trainings. Note that, among these, QKrm is the only option that has previously been proposed in the literature (Dehghani et al. 2023 Liu al. 2022) as a stabilizing modification of attention, to the best of our knowledge. QNomm and QUEST provide additional flexibility in the attention mechanism and hence perform better compared to QKNorm. QNomm, though stable, can still allow key norms of specific tokens to grow and hence, "steal attention" from other tokens. On the other hand, the query norms in QUEST can only influence the sharpness of attention and not which token should be attended to. We believe that this enables QUEST to learn better attention distributions as shown in the class-activation maps (see Figures
A7 and A8 and perform robust under corruptions and adversarial attacks.
A.2 QUEST ATTENTION IMPLEMENTATION
The implementation of both standard attention and QUEST attention is illustrated in Algorithm AT The key modification in QUEST attention is to fz-normalize the keys in lines 16-17. In any method that currently uses standard attention, QUEST attention can be used as a drop-in replacement. For other variants of attention that still use similar scaled dot-product formulation, QUEST attention variant can be obtained by normalizing the keys.
A.3 IMPLEMENTATION DETAILS OF THE TOY EXAMPLE
In this section, provide additional details about the construction of the toy example and the Transformer model used in our experiments.
Toy example construction: In the toy example, the real-valued vectors at the answer location æf are sampled differently depending whether the sample is biased or not (as per the Binomial random variable u). For the unbiased case, we define Σ = SST where all elements of S are sampled from a standard normal distribution. For the position bias, sampled the answer positions from normal distribution, (144 σT) and converted them into integer indices. We used μ = 10 and o; = 2. The most frequently sampled position occurs approximately 20% of the time in the data. The real-valued
16

---

Published as a conference paper at ICLR 2026
Algorithm A1 Computation of standard and QUEST attention
1: Input:
2: Tensor Qh E ND
Queries for N tokens head
3: Tensor K₁ NR
Keys for N tokens in head
4: Tensor V42 E RN
Values for N tokens in head Head dimension
5: integer DH N
6: function STANDARDATTENTIONN(),, Kh, Vi,)
7:
C
1
D constant scaling factor
VDw
8:
attention attenttionloogits + + C
Qn x KT
9:
attention poftmax(attente ion_iogits) attention + dropout(attent ion.)
softmax along the keys dimension > attention dropout
10:
output + attention V4
12:
output + out wwwpprcctiooouutpppppp
linear output projection output dropout
13:
output + dropout(output)
14:
return output
15:
function QUESTATTENTION(Q Kn, VT))
NormiKfaactor diag |K,, [R..
Calculate normalization factor of N keys
K₁ NormooooKBBB factor K₁.
Normk_factoor
A> Compute normalized keys
18:
iatention_logitte Qn x KT
tttpaaudaaateeiiinooiiieeee attention +
attention dropout(attenttion)
softmax along the keys dimension attention dropout
20:
output + attention V/
22:
output +- (ropppprcectoon)utttt)))
linear output projection A output dropout
23:
output +- dropout(outpput)
24:
return output
17

---

Published as a conference paper at ICLR 2026
vectors and the one-hot encoded vectors are both of 10 dimensions. As result the task is 10-way classification problem.
Transformer model setup: We use one-layer Transformer model for the experiments involving the toy example as it is sufficient to solve the task. We use only one head to enable easier interpretation of the norms of keys and queries belonging to specific positions such as [CLS] and answer location. The Transformer model setup is illustrated in Algorithm A and follows the standard implementations that are commonly used. The embedding dimensions of the Transformer model the same as the input data dimensions (= 20). The same number of dimensions is also used in the MLP hidden layer the Transformer. For the different attentions, the only change is to use different attention functions (see examples for standard and QUEST attention in Algorithm AT AT
Algorithm A2 Transformer model setup in the toy example
1: Input:
2: Tensor X E RN
Input data samples Learnable positional embeddings for the N positions
3: Tensor P E [
4: Tensor X E RIID
Learnable CLS token
5: function TRANSFORMER(X, XCLS, P)
6:
X + PrependCL.SToee
b. Prepend a learnable CLS token to the data
7:
X X +
Add positional embeddings
8:
Y < LayerNorml (X))
9:
Q XW
query projection
10:
K X W
> key projection
11:
V = X W
D value projection
12:
Y < Y + AttentionnQ.... K, V)
13:
Y - - - LayerNorm2(Y))
14:
Y X MLP(Y)
15:
output + classi i fier (YCLs)
Apply linear classifier to the CLS token features
16:
return output
A.4 EXTENDED ANALYSIS OF THE TOY EXAMPLE
We provide some additional results and analysis using the toy example in this section. In the main paper, we only showed the test success rates for the 3 attention formulations that performed reasonably well (QUEST, Standard and QNomm attentions). the attention formulations. We find that the QKNorm attentions largely fail achieving ~0% overall success rate. We attribute this to the fact that they discard useful information by normalizing the queries and keys. The input vector norms provide initial clue to solving the task.
Figure AT we show the success rates for all
The distributions for the training and test accuracies for the toy example are shown in Figures and A3 respectively. We clearly identify the biased and correct solutions based on their training and test accuracies. biased solution achieves a training accuracy of ~50-80% and a test accuracy of
20-40%. A correct solution achieves test accuracy greater than 90%. Degenerate solutions characterized by random chance test accuracy 10%. In Figure A4 and A5 we show the comparison of norms of unbiased answer token keys and non-answer token keys. For both standard and QNorm attention, the models do not seem to distinguish between the unbiased answer tokens and non-answer tokens in terms of norms of their keys. On the other hand, they assign a much higher norm the biased answer token keys, as shown in Figure from the main paper.
A.5 ADDITIONAL. EXPERIMENTAL DETAILS
A.5.1 IMAGE CLASSIFICATION
Investigating divergent DeÎT trainings: We found the DeiT training be unstable for the ViT-Base model when trained on 4 NVIDIA A100 40GB GPUs using an overall batch size of 1024 as proposed in Touvron et al. (2021a). We use PyTorch 2.1 for this experiment. In Figure A6 provide
18

---

Published as a conference paper at ICLR 2026
(ours) [overall=58%]
Standard [overall=25%]
QKNerm [overall=0%]
1,0
0.8
0.52
0.64
0.04
Learning rates
0.0025
0.84
0.8
0.84
0.6
0.005
0.84
0.88
0.84
0.84
0.6
0.68
0.72
0.6
0.68
0.28
0
0.4
0.0075
0.76
0.8
0.52
0.44
0.2
0.36
0.24
0.16
0
0.2
0.01
0:24
0.2
0.12
0.12
0.08
0.12
0.04
0
0.0
0.0
0.02
0.05
0.1
0.01
0.05
QNorm [overall - 49%)
OKNonm-HS -3%)
OKNorm-DS loverall -2%]
0.0005
0
0
0.001
0.08
0.2
0.24
0.04
0
0
0.08
0
D
Learring rates
0.0025
0.8
0.8
0.88
0.84
0.8
0
0.04
0.12
0.12
D
0.005
0.76
0.92
0.96
0.92
0.68
0
0.04
0.04
0
0.04
0.12
0.0075
0.96
0.92
0.76
0.6
0.16
0.08
0.04
0.12
0
0.01
0.76
0.28
0.12
0:16
0.2
0.04
0.04
0.0
0.01
0.02
0.05
0.1
0.01
0.02
0.05
0.1
0.0
0.01
0,02
0.05
0.1
Weight decays
Weight
Weight decays
Figure A1: Success rates of learning the correct solution to toy example. Models are trained with different hyperparameter combinations with different weight initializations and different realizations of the data.
(ours)
Probability
0.2
0.0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Train accuracy
Train accuracy
Train accuracy
QRNorm-Swin
OKNorm-VIT228
Probability
0.2
0.0
0.2
0.4
0.5
0.8
1.0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Train accuracy
Train accuracy
Train accuracy
Figure A2: Distribution of training accuracies the end of the training for the toy example.
analysis of the maximum logits and its associated query and key norms for ViT-Base model trained using DeiT.
Repeating DeiT ablation experiments: For the ablation experiment shown in section 4.1.1 and Table of the paper, we conduct three independent runs and show the aggregate results (mean and standard deviation) in Table On running a 2-group t-test to check the statistical significance, we obtained p- p-values lower than 0.01 compared to the second best for all metrics in Table AI
Training with limited training data: In Table A2 we show the results for training DeiT-Tin -Tiy models using limited subsets the ImageNet training data. The data subsets are uniformly sampled and the reported results aggregated over 3 such random subsets. All models trained for 300 epochs with the same training setup as a standard DeiT model trained on the full ImageNet training dataset. We observe that QUEST consistently performs better than standard attention by 1.0-1.5%.
Training with different learning rates: In Table A3 we evaluate DeiT training using a ViT-Small model using different learning rates. We find standard attention be unstable at larger learning rates.
19

---

Published as a conference paper at ICLR 2026
QUEST (ours)
Standard
QKNorm
Probabillty
0.4
0.0
1.0
Test accuracy
Test accuracy
Test accuracy
QNorm
QKNorm-Swin
QKNorm-ViT228
Probability
0.4
0.2
0.0
0.2
1.0
0.2
1.0
0.4
0.6
1.0
Test accuracy
Test accuracy
Test accuracy
Figure A3: Distribution of test accuracies at the end of the training for the toy example.
Epoch= D
Epoch=25
5p
200
Unblased answer
Unblased
100
answer
2
Norm answer token
Norm token keys
Norm answer token
Figure A4: Norm of unbiased answer token keys and non-answer token keys in standard attention. In terms the key norms, do observe any distinction between unbiased answer tokens and non-answer tokens.
Epoch=0
Epoch=25
Epoch=
150
Unblesed answer
100
Unblased
1.5
2,0
2.5
3.0
3.5
2.5
5.0
7.5
10.0
0.0
2.5
5.0
7.5
10.0
answer token keys
answer token keys
Norm answer keys
Figure A5: Norm of unbiased answer token keys and non-answer token keys in QNorm attention. In terms the key norms, we do not observe any distinction between unbiased answer tokens and non-answer tokens.
QUEST attention produces stable training at different learning rates and consistently displays better performance.
Training with large-scale VIT models: In Table A6 we train large-scale ViT models, namely, ViT-Huge (600 Million parameters) and VIT-2B (2.4 Billion parameters) for smaller number of epochs of 100 (+20 finetuning epochs) and 10 respectively using the stable DeiT-3 training recipe. We reduce the global batch sizes to 1024 for the ViT-Huge model and 64 for the ViT-2B model, to fit on a single NVIDIA A 110 node. For the ViT-Huge model, we use an input image resolution of 154 x 154 as in Touvron al. (2022) but for the ViT-2B model, we use a reduced input image resolution of 96 x 96. Similarly. for the longer experiments (DeiT-3 trainings for 400+20 epochs) shown Table for ViT-Large and ViT-Huge, we made the following modifications compared the standard training recipe:
VIT-L/16: Batch size reduced from 2048 to 1024.
VIT-H/14: Batch size reduced from 2048 to 1024, input image resolution in the first training phase reduced from 154 x 154 to 126 126.
20

---

Published as a conference paper at ICLR 2026
ViTssViT, an architecture using cross-attention: CiossViT (Chen et al. 2021) is extension to ViT that uses two branches of Transformer layers using different patch size. This is followed by additional Transformer layers which use cross-attention instead of self-attention. We train the models using a 8 NVIDIA A100 40GB GPUs. We use the same experimental configuration as in Chen al.] (202T) for the different models and adapt the batch sizes fit our GPU setup. We consider the Cross ViT--VVT----T ViT-9-Daggel Cross ViT-Sall and Cross VTT- ViT-I8-Dagger models. We use a batch CrossViT-Small and Cross 1iT-18-Daggee models respectively. We report the results using standard and QUEST attention in Table A4 We found the Cross ViT- 18 8-Daggee model be unstable with standard attention and the training crashed
size of 2048, 1024 and 1024 for the
Last layer
Intermediate layer
First layer
Max togla laper
Standard
@Kltem-ts
QUEST
OKNorm D5
10000
5000
10000
20000
5000
10000 15000
(a) Progression of maximum logits in the first, intermediate and last layers of ViT-Base model trained using DeiT and using different attention formulations. We observe that the max logits are highly stable different layers using QUEST attention. The max logits rapidly increase for standard attention before the training crashes around 20000 iterations.
Last
Intermediate layer
First layer
Lop No
5006
10000
15000
20000
5000
10000
15000
20005
5000
10000
15000
20000
Heeraton
tteaation
Trenst
(b) Progression of log of the key norms (corresponding to the maximum logit token above) in the first, interme- diate and last layers of ViT-Base model trained using DeiT and using different attention formulations. For QUEST and QKomm variants, the norms prior to the Ez-normalization and do not have any impact on their attention logits.
Last layer
Intermediate layer
First layer
Lag of totee
Standard
quNorr-tt
QUEST
0104orm-5
19030
15600
20000
5000
10d00
15000
5000
15600
20000
Itoration
Itsration
(c) Progression of log of the query norms (corresponding the maximum logit token above) in the first, intermediate and layers VTT-Base model trained using DeiT and using different attention formulations. For QKNmmm variants. the norms are prior to the fg-normalization and do not have any impact on their attention logits. Note that the query norms are not Eg-normalized in QUEST attention.
Figure A6: Maximum logits and their associated query and key norms in different layers for ViT- Base model trained using DeiT. The model with standard attention crashes around 20000 iterations, which not observed in the other models. The rapid increase maximum logits (especially in intermediate and initial layers) is accompanied by underlying increase in the corresponding query and key norms. Both maximum logits and query norms are stable for QUEST attention, showing that it is not necessary normalize both queries and keys to ensure training stability.
21

---

Published as a conference paper at ICLR 2026
whereas were able to train with QUEST attention without any training issues. We find QUEST attention to perform better than standard attention, showing that QUEST attention can also be useful in coss-attention setups. This also demonstrates the potential for QUEST attention to be used as drop-in replacement for standard attention, orthogonally with other Transformer developments.
Table A1:
Ablation of different QK normalization methods for ImageNet classification ( (ViT-Tiny model trained using DeiT for 300 epochs) ["L al. (2022). (Dehghani et al. [202)]] Means and standard deviations of results are reported over 3 independent runs with different seeds.
Attention
Scaling
ImageNet-val
Top-1 1 (%)
Top-5 (%)
NLL +
Standard
1/vDb
72.65 + 0.10
91.49
0.03
1.187 0.003
QUEST
73.36 + 0.16
91.86 0.06
1.160 + 0.004
QNomm
72.73 + 0.11
91.44 + 0.02
1.188 0.004
QKNorm-HS
C E ERLH
72.49 + 0.10
91.43 + 0.09
1.198 + 0.009
QKNorm-DS
Cg. Cs E RLxDs
71.87 0.18
0.18
90.96 + 0.11
1.231 + 0.008
Table A2: ImageNet classification validation accuracies with DeiT-Tiny trained using different
amounts of training data. Results are averaged over 3 different training data subsets.
Data (%)
Attention
ImageNet-val
Top-1
Top-5
NLL
Standard
40.82 + 0.81
63.92 + 0.98
3.326 + 0.074
5%
QUEST
41.96 0.15
65.24 + 0.07
3.254 + 0.014
Standard QUEST
56.02 + 0.58
78.52 + 0.66
2.201 + 0.042
10%
57.51 + 0.33
79.75 + 0.17
2.136 + 0.016
Standard QUEST
68.87 0.45
88.86 + 0.25
1.386
+ 0.020
25%
69.96 + 0.26
89.50 + 0.11
1.337
+ 0.006
Standard
72.50
91.45 91.91
1.190
100%
QUEST
73.33
1.160
Table A3: Learning rate sensitivity of DeiT training evaluated using a ViT-Small model on ImageNet validation performance
Learning rate
Epochs
Standard attention
QUEST attention
Top-1
Top-5
Top-1
Top-5
le 4
50
62.8
84.9
63.9
85.7
2e 4
50
66.7
87.9
67.8
88.5
5e 4
66.9
87.9
68.5
88.7
le 3
50
-training crashed-
67.9
88.5
2e 3
50
-training crashed-
67.0
87.8
A.5.2 ROBUSTNESS IN IMAGE CLASSIFICATION
We evaluate the adversarial robustness of our models trained using DeiT and DeiT-3 using different adversarial attacks using the experimental protocol of Nielsen et al. (024)) and their public codebase? We consider FGSM, PGD and Auto attacks with a perturbation budget 1/255 and the SPSA attack with perturbation budget of 0.1 (under -nomm..
Prior works on adversarial robustness mainly benchmarked on the ViT-Tiny model. We also evaluated larger ViT models compare standard and QUEST attention. These results reported in Table A5
22

---

Published as a conference paper at ICLR 2026
Table A4: Cross ImageNet Classification
Model
Attention
Epochs
ImageNet-val Top-5
Top-1
NLL
Cross Vig-9-Dagger
Standard
300
76.4
93.4
1.009
Cross ViT-e-
QUEST
300
77.0
93.6
0.991
Cross ViT-Small
Standard
300
80.5
95.5
0.834
Cross ViT-Small
QUEST
300
80.9
95.6
0.826
Cross ViT-I8------- D-Dgger
Standard
300
-training crashed-
Cross ViT-I1-Daager
QUEST
300
82.8
96.1
0.780
For this evaluation, we use the DeiT and DeiT-3 models described in Section 4.1.1 (the clean data results for these models are shown in T) We do not consider Elliittial-QQUEST for these larger ViT's as Nielsen et al. al. (2024) only experimented with ViT-Tiny models and optimal training recipes for larger models are unavailable.
Elliptical-QUESTY Elliptical attention proposes to compute the attention logits as QMK /VDH where M is Mahalanobis factor. The matrix M is diagonal matrix, where the diagonal elements scale the different feature dimensions in the QK product. In Illiptical-QUESS attention, we obtain the logits as SQMK* KT, where S denotes diagonal matrix containing the query token norms and Q. K denote the Ez-oomalleeed queries and keys respectively. In QUEST attention, the product QK" represents cosine similarity. In Elliptical QUEST, the metric can be interpreted as an elliptical analogue, where each dimension is weighted differently (according to M) in a cosine similarity.
Table A5: Robustness to adversarial attacks using larger VIT models trained using DeiT. Note that prior works like Elliptical attention only considered ViT-Tiny models for adversarial robustness evaluation ( DeiT-3
Model
Attention
Epochs
FGSM
PGD
Auto
Top-1
NLL
Top-1
NLL
Top-1
NLL
VIT-S/16
Standard QUEST
200 200
65.90
1.362
54.47
2.050
38,22
2,699
VIT-S/16
67.02
1.328
57.51
1.886
40.98
2.540
VYT-B/16'
Standard
400 + 20
69.41
1.229
52.55
2.080
VIT-B/16
QUEST
400 + 20
70.67
1.194
54.64
2.033
VIT-L/16
QKNorm-LN
100
58.95
1.849
51.79
2.355
VIT-L/16
QUEST
100
62.85
1.669
53.54
2.316
A.5.3 IMAGE CLASSIFICATION EXPLAINABILITY
We evaluate explainability of model using recent method, AG-CAM (Leem & Seo 2024) that combines attention maps with gradient information. We adapt the code from their public repository B to DeiT-3 models. Based on the ViT-B model trained using DeiT-3, we show additional qualitative examples in Figures A7 and A8
A.5.4 IMAGE SEGMENTATION
The image segmentation models are trained on the ADE20K dataset (Zhou et al. 2019: 2017)
following the setup of Segmenter (Strudel al. 202T). We use the same training configurations as in Nielsen et al. (2024) for both ViT-Ti and VIT-S backbones. The VIT backbones are initialized with the weights obtained from the DeiT trainings in Section 4.1. and then, we finetune the entire model, both the encoder and the decoder. This experiment is conducted GPU, We train these segmentation models for 160K iterations with a global batch size of 8. An SGD optimizer with a starting learning rate of 0.001 and polynomial learning rate scheduling is used.
single NVIDIA A100 40GB
Fxplaan-- Vision- Tranofo-rT.
Guided--y-SSlf-Attteee
23

---

Published as a conference paper at ICLR 2026
Table A6: DeiT-3 ImageNet Classification with large-scale ViT models
Model
Parameters
Attention
Epochs
IN-val Top-5
Top-1
NLL +
ViT-H/114
600M
Standard
100
68.3
88.8
1.545
VIT-H/114
600M
QUEST
100
69.2
89.4
1.495
ViT-H/14
600M
Standard QUEST
100 20
76.2
93.3
1.042
ViT-H/14
600M
100 + 20
76.8
93.6
1.002
ViT-2B/16
2,4B
Standard OKNorm-DS
10
6.153
ViT-2B/16
2.4B
10
5.2
14.9
5.740
VIT-2B/16
2,4B
QUEST
10
11.9
27.8
5.011
Image
QUEST
Standard
Pred: Rookcase
Pred: Rookcase
(a) Example from the "Bookcase" class. QUEST attention focuses similarly on most of the books and shelves. Standard attention focuses only few of the instances.
Image
QUEST
Standard
Pred: Cockatoo
Pred: Cockatoo
(b) Example from the "Cockatoo" class, Standard attention only focuses on specific part of the birds. QUEST attention evenly attends to the entire birds.
Image
QUEST
Standard
Pred: Macaw
Pred: Macaw
(c) Example from the "Macaw" class, Standard attention only focuses specific part of the birds. QUEST attention evenly atlends the entire birds.
Figure A7: Examples showing class activation maps (CAM) for images from the ImageNet dataset.
The CAM is obtained with the AG-CAM method using the DeiT-3 models shown in 4.1.I Standard attention concentrates attention on specific parts of the object. On the other hand, QUEST attention attends more evenly to the entire object or all relevant regions.
24

---

Published as a conference paper at ICLR 2026
A.5.5 LANGUAGE MODELING
We follow the experimental protocol of Nguuen al. (202b);; Nielsen et al. (2024): Schlag et al. (202T) and train Transformer models of Small (44M parameters) and Medium (90M parameters) sizes on the clean WikiText- 103 dataset (Merity et al. 2017). The models are trained using Adam and using the code and hyperparameters from the Elliptical Attention repository The Transformer-Smi Small model is trained for 120 epochs with a batch size of 96, starting learning rate of 0.00025 and cosine model is trained with batch size of 56, starting learning
scheduling. The Transformer-Mediun rate of 0.00025 and cosine scheduling. Following the standard evaluation setting in |Schlag et al. (2021), process the text sequence using a sliding window (256 for Transformer-Small and 384 for Transformer-Mediuee
The perplexity is computed based on the last position except for the first segment, where it is evaluated for all positions. We also experiment with the larger Transformer-X Standard (151M parameters) model (note that the model is called Base in the repository) from Dai al. (2119) using their public codebase 2
We use the hyperparameter setup from the repository and this uses the same evaluation protocol described above.
A.5.6 TIME SERIES CLASSTFICATION
We use the training and evaluation protocol from the Time Series Library (TSLib) repository The default Transformer model in the repository concatenates the output token features and uses linear classification layer. The final classification layer can become arbitrarily large since the weights WCLS E R where N is the sequence length, D is the embedding dimensionality and C' is the number classes. Instead, we found that using [CLS] token consistently improved the standard Transformer results (see Table A7). Hence, use Transformers with [CLS] token and the classification layer only depends on the CLS token output from the Transformer. Both standard and QUEST attention this same setup in our experiments. We the same training hyperparameters for both attention types, based on the default configuration for a standard Transformer provided in the repository. This consists of training 3-layer Transformer with a model dimension of 128, batch size of 16, learning rate of 0.001 using the RAdam (Liu et al., 2020a) optimizer for 100 epochs (with an early stopping patience of 10 epochs). In Table A8 comparing with QNomm and QKNorm attention.
we show additional ablation experiment
Table A7: Impact of using a [CLS] token for UEA Multivariate Time Series Classification using standard Transformer
Dataset
Standard
Standard + [CLS]
EthanolConcentrattion
28.14
29.28
FaceDetection
67.99
65.24
Handwriting
38.24
42.00
Heartbeat
77.07
77.56
Japanese Vowels
97.84
98.38
PEMS-SF
84.97
83.82
SellRegulationSCPI
90.44
88.05
SellRegulationSCP2
55.00
58,89
SpokenArabbeeDigits
98.41
98.86
UWaveGestureLLibrrry
86.25
86.88
Average
72.44
72.90
A.5.7 POINTCLOUD SEGMENTATION
We consider the pointcloud segmentation using PointTransformer-VV model architecture (Wu al. 2024) and use the Pointcept (Contributors) 2023) training framework. For this task, we use the nuScenes (Caesar al. 2020) dataset with the same training configurations in the Pointcept
Wiikitxx
thaps://gthub..om///i http:/shwww.ccmom///m/mmeeeeeeee Stwt://ieeeeeSeeeiiiirrrrre
25

---

Published as a conference paper at ICLR 2026
Table A8: Ablation experiment for UEA Multivariate Time Series Classification using different forms
of attention
Dataset
Standard
QNorm
QKNorm
QUEST
EthanolCooncentration
29.28
29.28
29.28
30.42
FaceDetection
65.24
64.81
64.73
65.83
Handwriting
42.00
41.06
49.18
49.18
Heartbeat
77.56
75.61
77.56
78.54
Japanese Vowels
98.38
98.92
98.38
98.38
PEMS-SF
83.82
80.92
79.19
80.92
SellRegullatonSCCPI
88.05
87.71
88.76
88.05
SellRegulationSCCP2
58,89
56.67
54.44
58.33
SpoieeAraaiieDigits
98.86
99.41
99.45
99.41
UWaveGestureL.ibrayy
86.88
85.00
86.81
86.25
Average
72.90
71.94
72.78
73.53
repository. We train PointTransformer-V V3 models using standard, QKNorm and QUEST attentions for 50 epochs and the validation mloU scores are reported in Table A9
Table A9;
Pointeloud segmentation using PointTransformer- V3 model the nuScenes dataset. We report the mloU on the validation set.
Model
Parameters
Attention
nuScenes val. mloU (%)
PointTransformer--V
46.2 M
Standard
80.40
PoinTTransformerr-V. V3
46.2
QKNorm-DS
80.37
PointTransformer-Ve V3
46.2 M
QUEST
80.83
A.5.8 GRAPH TRANSFORMER BENCHMARKS
We consider the ZINC, CIFAR1O, PATTERN and CLUSTER tasks from the standard Graph Neural
Network (GNN) benchmarks (Dwivedi et al. 2023), detailed as follows:
ZINC regression task for molecular property and it is evaluated using the Mean Absolute Error (MAE).
CIFAR1O is graph classification task based on superpixel graphs and it is evaluated using classification accuracy.
CLUSTER and PATTERN are node classification tasks and they evaluated using wlagh-weighted classification accuracy.
From the long-range graph benchmarks (Dwivedi et al. 2022). we consider the PascalVOCCSP. COCO-SP, Peptides-func, Peptides-strue struct and PCQM-Contact tasks, detailed as follows:
VOS-SS and COCO-SP are node classification tasks based on the Pascal-VOC dataset (Everingham et al., 2010) and the MS-COCO (Lin al. (014) datasets respectively. For both tasks, the macro weighted F1 score is used as the performance metric.
PCQM-Contact is a link prediction task and it is evaluated using the Mean Reciprocal Rank (MRR) (Hoyt et al. 2022).
Peptides-fune is multi-label classification task with 10 classes and it is evaluated using the unweighted mean Average Precision (AP).
Peptides-struct is multi-label regression task for graph-level properties and it evaluated using the Mean Absolute Error (MAE).
The hyperparameter setups to train GPS models for standard GNN benchmarks is shown in Table A10
Similarly, the hyperparameter setups to train GPS models for long-range graph benchmarks is shown
26

---

Published as a conference paper at ICLR 2026
in Table A1l The Graph Transformers trained and evaluated using the GraphGPS repository
and the shared configuration files for each dataset.
Table A10: GPS hyperparameter setup for standard GNN henchmarks (Dwivedi et al.
2023
Hyperparameter
ZINC
CIFARIO
PATTERN
CLUSTER
# GPS Layers
Hidden dim
64 GINE
$2
64
48
GatedGCN Transformer
GatedGCN Transformer
GatedGCN
Gos--GooAttnn
Transformer
Transformer
# Heads
4
Dropout
Attention dropout
0.5
0.5
0.5
0.5
Graph pooling
sum
mean
Positional Encoding
RWSE-20
LapPE-8
LapPE-16 16
LapPE-10
PE dim
28
8
16
PE encoder
linear
DeepSet
DeepSet
DeepSet
Batch size
32
16
32
16
Learning Rate
0.001
0.001
0.0005
0.000S
# Epochs
2000
100
100
100
# Warmup epochs Weight decay
50
5
5
5
lc.5
lc:5
lc-5
le-5
# Parameters
423,717
112,726
337,201
502,054
PE precompute
23s
2.55min
28s
67s
Time (epoch/total)
21s/
64s 1.78h
32s 0.89h
86s 2.40h
Table A11: GPS hyperparameter setup for long-range graph benchmarks
Dwivedi et al. 2022
Hyperparameter
PascalVOC-SP
COCO-SP
PCQM-Contact Peptides-fune
Peptides-struct
GPS Layers
Hidden dim
96
96
96
96
GPS-MPNN
GatedGCN Transformer
GatedGCN
GatedGCN Transformer
GatedGCN
GatedGCN
GPS-Sel[Attn
Transformer
Transforner
Transformer
Heads
8
4
4
4
Dropout
Attention dropout
0.5
0.5
0.5
0.5
0.5
Graph pooling
mean
mean
Positional Encoding
LapPE-10 16
LapPE-10 16 DeepSet
LapPE-10 16
LapPE-10
LapPE-10
PE dim
16
16
PE encoder
DeepSet
DeepSet
DeepSet
DeepSet
Batch size
32
32
256
128
128
Learning Rate
0.0005
0.0005
0.0003
0.0003
0.0003
Epachs
300
200
200
200
Warmup epochs
10
10
10
5
5
Weight decay
0
0
0
0
Parameters
510,453
516,273
512,704
504,362
504,459
PE precompute
8.7min
1h 34min
5.23mi
73s
73s
Time (epoch/total)
17.5s / 1.46h
213s / 17.8h
154s 8.54h
6.36s 0.35h
6.15s 0.34h
A.5.9 SELF-SUPERVISED LEARNING
We also conduct self-supervised learning (SSL) experiment using the iBOT Zhou et 2021 method (which is the foundation for SOTA SSL models like DINOv2 (Oquab et al.) 2024)). This experiment was ran on single node of NVIDIA A100 40GB GPUs. When we attempted to reproduce DINOv2 pre-training on ImageNet- 1K using the ViT-Large model (standard attention)
27

---

Published as a conference paper at ICLR 2026
with smaller batch size of 512 (instead of 2048), observed significant drop in performance to 68.2% KNN accuracy (vs 81.6 % reported in their repository). Since training DINOv2 with smaller batch sizes was non-trivial, we instead opted for the iBOT method. We detail the pre-training and downstream evaluations below.
Pre-training: We pre-train a ViT-Base/16 model using the iBOT (Zhou et al. 2021 method and the vMF normalization (Govindarajan et al. 2023) for 400 epochs on the ImageNet-Ik IK dataset (Deng) al. 2009) using the code from the iBOT repository a and with the exact same hyperparameters as iBOT. The pre-training is carried out single node consisting of 8 NVIDIA A100 40GB GPUs.
ImageNet downstream tasks: We evaluate the KNN performance using the same protocol as DINO (Caron al.) [202]). We use weighted k-NN (temperature 0.07) and report the best result among those obtained with & = (5. 10, 20, 100, 200). We generally found & = 10 to produce the best result for both standard and QUEST attention. For linear classification, we follow the evaluation protocol of iBOT (Zhou et al. (021) 1) and report the best results among those obtained using different learning rates. The linear classifier is trained on the features obtained by concatenating the [CLS] features and average pooling of the patch features. The training is run for 100 epochs using the same hyperparameter setup as in iBOT. For evaluations with limited training data, we follow the evaluation protocol of Assran al. (2022) and use the same data subsets. We report the average validation accuracy over different data subsets for the 1, 2 and 5 images per class settings. For finetuning on ImageNet, we train for 100 epochs with layer-wise learning rate decay of 0.65, following the effective recipe from BelT [aao et al. 2022). We report the best performance after considering 1, 3, 2e 3]. These results on ImageNet downstream tasks are reported in Table A12 We observe consistent improvements on all the considered settings. Finetuning models from SSL pre-trained weights common in recent times and we highlight that QUEST attention can also bring improvements in this setting (84.4 % vs 84.1 %).
different learning rates from 8e 4, 9e
Transfer linear probing: We conduct transfer linear probing evaluation by freezing the pre-trained model and training linear classifier on the [CLS] features output by the model. We follow the evaluation protocol of Ericsson et al. (202)) and Chen et al. (2020) and train éz-regularized linear classifiers. We select the regularization strength among set of 45 values spaced linearly in the range [-6, 5] in log-space and compute the standard evaluation metric for each dataset. The dataset suite includes the following datasets: Aircraft (Maji al., 2013). Caltech (Li et al.) 2022b), Describable Textures Dataset (DTD) (Cimpoi et al. 2014), Flowers (Nilsback & Zisserman 2008). Food (Bossard et al. 2014), Pets (Parkhi et al. 2012) and SUN397 (Xiao et al.
2010 2016) datasets. The detailed transfer linear probing results are shown in A144 QUEST performs significantly better than standard attention on Aircraft and Flowers datasets. On other datasets, the results are somewhat mixed. Nevertheless, QUEST performs better than standard attention in terms of the overall average.
Table A12: Self-supervised pre-training on ImageNet with iBOT and evaluating on ImageNet tasks
Attention
KNN
Linear
Finetuning
Few-shot
1 img/cls
2 imgs/cls
5 imgs//ls
1% imgs
Standard 78.7
80.3 80.5
84.1
51.6 + 0.1
61.1 1 + 0.7
72.3 69.0 + 0.1
68.3 - 0.3
QUEST
79.0
84.4
52.5 + 0.2
61.7 + 0.6
72.7
Table A13: Self-supervised pre-training on ImageNet with iBOT and evaluating on transfer tasks
Attention
Transfer linear probe
Image Retrieval
Vos
Average
RParis
ROxford
DAVIS- 2017
M
H
M
H
J&FF
Jm
Fm
Standard QUEST
81.5 81.9
65.4
38.1 1
38.1
13.8
63.1
61.9
64.2
65.8
39.3
38.6
15.3
63.2
61.8
64.5
28

---

Published as a conference paper at ICLR 2026
Table A14: supe-suee pre-training ImageNet-IK with iBOT-vMF and evaluating with transfer
linear probes
Attention
Actt.
Cal101
DTD
Flwrs.
Food
Pets
SUN
Avg.
Standard
58.1
95.5
74.7
94.8
83.6
93.9
70.2
81.5
QUEST
59.9
95.1
74.4
95.8
83.9
94.3
69.8
81.9
Other transfer learning tasks: For image retrieval experiments, follow the evaluation protocol of DINO and evaluate on the face-blurred versions (v1.0) of the Oxford and Paris datasets. We perform image retrieval based nearest neighbors and report the mean Average Precision (mAP) on the medium (M) and hard (H) data splits for each dataset. For video object segmentation (VOS), we use the evaluation protocol of DINO and evaluate VOS performance using standard metrics on the DAVIS-200 benchmark dataset (Pont-TTuse et al. (017). These transfer learning results are reported in Table We find that QUEST attention performs better on image retrieval and similar to standard attention on video object segmentation.
29

---

Published as a conference paper at ICLR 2026
Image
QUEST
Standard
Query: Bottie cap
Query: Bottle
(a) Image queried with the target class of Bottle and Bottlecap.
Image
QUEST
Standard
Query: Zebra
Query: Elephant
(b) Image queried with the target class of Elephant and Zebra.
Image
QUEST
Standard
Query: Dog
Query: Cat
(c) Image queried with target class of Dog and Cat.
Figure A8: Examples showing class activation maps (CAM) for images containing two distinct objects. The CAM obtained with the AG-CAM method by querying for the specified classes (i.e. using the specified class as the target) using the DeiT-3 models shown in 1.1. Standard attention concentrates attention on specific parts of the object or fewer instances of the object. On the other hand, QUEST attention attends more evenly the entire object all relevant regions.
30


Reviews:
[{"review_id": "MECStkA4r6", "reviewer": "Reviewer_5zrP", "summary": "The authors propose a novel attention mechanism, QUEST (QUEry-modulated Spherical aTtention) to mitigate instability (attention collapse) during the training. The authors argue that the instabilities may come from spurious correlation, or stealing attention due to norm of a certain key. QUEST attempts to solve this by normalizing the key vectors. This coulb be interpreted between standard attention and QK-Norm attention, which focuses on the cosine similarity rather than the dot product. This paper demonstrates that QUEST stabilizes training and improves the performance through various experiment.", "questions": "1. Why does query normalization perform worse?\n2.  Please clarify the elliptical-quest about scaling.\n3. Please provide deeper comparison with QK-norm method.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper clearly reformulates the problem. To calculate attention logit, which shows the relevance between query and key, the two obvious choices are  dot-product vs cosine similarity. The authors suggest a solution in between.\n\n2. The authors present a broad experiments results.\n\n3. Quest shows performance gain in IN-C experiment, which implies that QUEST can aggregate information from broad region.", "weaknesses": "1. toy example is too limite. Though the authors devote large portion of the paper into the toy example, the model does not have query, key projection head. The role of projection head is to extract information from embedding vectors. Analyzing attention without the learnable projection head is not compelling.\n\n2. The theoretical analysis is not clear. The auithors argue that QUEST can  mitigate the trainig instability but does not show this metric with other method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "af3ZIKSACw", "reviewer": "Reviewer_Bs2Y", "summary": "For transformer architecture, in order to solve training instabilities when the norm of queries and keys arbitrarily increase, this paper propose  a new attention  (QUEST). This mechanism map the keys to a hyperspherical latent space and single token are still able to sharpen the attention distribution. Empirical experiments show the effectiveness of the proposed methods.", "questions": "See above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well-written and easy to follow. It has clear motivation and it identifies an attention instability that arises from growing norms of queries and keys. Based on that, this paper proposes an elegant solution driven by mathematical intuition and compare with relevant methods.\n\nThe proposed method is easy to be implemented.\n\nThis paper validate proposed method on broad experiments(vision, language, graph transformers and time-series.) and achieve superior performance. Moreover, in vision area, the explainability and robustness analysis are provided.", "weaknesses": "Training instability is not very proper since this paper does not analyse the optimization process but only after optimization. \n\nTechnical depth is limted since it is conceptual but not theoretical. The analysis is mostly based on observation and training instability is affected by the network Jacobian which is lack in the paper.\n\nThe scope is moderate", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "MP3XOb1tQK", "reviewer": "Reviewer_MR6x", "summary": "This paper revisits the instability issues in Transformer training, which arise from uncontrolled growth of query/key norms in standard softmax attention. The authors propose a simple yet effective variant, QUEST, that normalizes only the keys while keeping queries unnormalized. This breaks the mutual amplification between query and key norms, stabilizing training and allowing each query to control its own attention sharpness. The method is a drop-in replacement for standard attention and is evaluated across diverse domains. Experiments show that QUEST improves training stability, robustness to corruptions and adversarial attacks, and sometimes accuracy.", "questions": "1. Can the authors provide a more formal theoretical justification or analytical evidence beyond heuristic explanations?\n\n2. In NLP tasks, does QUEST consistently outperform QKNorm or LayerNorm-based stabilizations?\n\n3. More modality/field need to be test. Can QUEST be extended to multi-modal or diffusion Transformer architectures?\n\n4. Previous [1] attributes training instability mainly to attention logit explosion and proposes QKNorm as a normalization-based remedy, how does QUEST theoretically or empirically differ from QKNorm in addressing this issue? Specifically, can the authors clarify whether QUEST’s key-only normalization provides comparable stability at large model scales, and if so, why this asymmetric normalization would outperform the symmetric approach proposed in [1]?\n\n5. In NLP tasks, does QUEST consistently outperform QKNorm or LayerNorm-based stabilizations?\n\n\nThis is an interesting work, and I will raise my score if authors address all my concerns.\n\n[1] Scaling Vision Transformers to 22 Billion Parameters", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Well-motivated**: Identifies a concrete cause of Transformer instability (growing norms) and proposes an intuitive fix requiring minimal code change.\n\n**Strong Empirical Results**: Demonstrated consistent gains on multiple benchmarks (ImageNet, ADE20K, WikiText-103, GraphGPS, UEA datasets).\n\n**Robustness Improvements**: Shows better performance under adversarial attacks (FGSM, PGD, AutoAttack) and data corruptions.\n\n**General Applicability**: Works as a drop-in replacement across architectures (ViT, CrossViT, Transformer-XL, Graph Transformers).\n\n**Theoretical & Empirical Insight**: Offers a helpful interpretation of attention norm dynamics and a controlled toy experiment demonstrating the effect of spurious correlations.", "weaknesses": "**Experiments:** Most experiments emphasize vision tasks; evaluations in NLP and other domains are relatively light.\n\n**Novelty:** Normalizing keys is a simple modification; some may view it as an extension of prior QKNorm works rather than a fundamentally new paradigm.\n\n**Ablation:** The comparison between normalizing only queries vs. only keys could be expanded.\n\n**Analysis**:Theoretical analysis is mostly heuristic and lacks rigorous mathematical derivation or verified data support.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "dIiWi60Ms6", "reviewer": "Reviewer_bMvw", "summary": "This paper presents QUEST, a technique designed to enhance training stability and feature learning in Attention mechanisms. The authors investigate how attention logit magnitudes and the norms of Query (Q) and Key (K) matrices influence training dynamics. They find that applying global normalization to the K matrix can cause the model to prematurely focus on less significant features during early training phases, leading to suboptimal convergence. Additionally, since K's norm influences gradient magnitudes, this approach creates training instability. These findings are supported by prior research and demonstrated through a simplified example.\nThe proposed solution involves applying L2-normalization exclusively to the K matrices while leaving Q matrices unnormalized. This design prevents K from dominating the training process while preserving Q's ability to modulate attention sharpness on a per-token basis through softmax temperature control. To validate their approach, the authors perform experiments training vision transformer models both with and without QUEST, demonstrating its effectiveness.", "questions": "See Weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) The paper addresses a significant issue affecting transformers across diverse applications and offers a straightforward solution to enhance their stability and performance. \n\n2) The paper provides a thorough analysis of the attention mechanism and the function of its components. Claims regarding current limitations of Scaled Dot-Product Attention and proposed alternatives are validated through prior research and a toy example, which reinforce the paper's arguments and empirically demonstrate QUEST's effectiveness in addressing these issues.\n\n3) QUEST's impact is illustrated through visualizations, which support the assertion that it prevents converging to less meaningfull attention weights.\n\n4) The authors evaluate QUEST across multiple transformer architectures and tasks, demonstrating superior performance in nearly all experimental settings.", "weaknesses": "1) The main weakness of this work is that the reported improvements are relatively small, often around 1% or less. This raises some concern, especially since it appears that each experiment was run only once. Given the number of experiments conducted, it would be helpful to include a statistical summary of the results. Could the authors provide some form of statistical analysis of the improvements across tasks? Even a simple aggregate analysis could offer a clearer picture of the overall gains.\n\n2) The paper assumes that the primary issue arises from the normalization of the K vectors, suggesting that Q should remain unnormalized during optimization. However, experiments involving only Q normalization also seem to yield positive effects. In the toy experiment, Q appears less stable than K but still more stable than the baseline, and in some cases even achieves better accuracy, particularly for certain learning rate values, where Q-Norm seems easier to optimize in the QUEST setup. This trend is also reflected in the ablation study, where Q-Norm improves over the baseline and performs comparably to K-Norm. Did the authors further investigate Q-Norm beyond the results presented in the paper?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "ya9W6kEUbe", "reviewer": "Reviewer_5T4g", "summary": "This paper analyzes the role of the norm of queries and keys in the attention module, proposes an alternative variant to the standard attention mechanism, and demonstrates the effectiveness of the proposed attention variant across multiple domains.", "questions": "1.**Question regarding Fig. 1**. In the third row of Fig. 1, the standard ViT attends to the bird, yet it still makes an incorrect prediction. I hope the authors can provide a clear explanation to help readers understand this phenomenon.\n\n2.**The baselines in Table 2 appear relatively low**. First, when using DeiT and MAE recipes for supervised training standard ViT-Base/Large on ImageNet, I did not encounter training instability issues. Second, using the DeiT-3 training method (fourth row), the authors report a baseline accuracy of 82.7%, which is lower than the official DeiT-3 result of 83.8% (which I have successfully reproduced). Third, I request the authors to provide the training curves.\n\n3.**Lack of large-scale experiments**. Although the effectiveness of the proposed method has been validated across multiple domains (cf. Tables 2, 5, 10), the model parameters used are relatively small. As a simple replacement for standard attention, I believe more large-scale experiments should be conducted.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The designed attention variant serves as a drop-in replacement for the standard approach, requiring only simple modifications.\n\n2. Experiments were conducted across multiple domains, including image classification, language modeling, Graph Transformers, and time series tasks.", "weaknesses": "see questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)