QUEST: A robust attention formulation using query-modulated spherical attention

Hariprasath Govindarajan; Per Sidén; Jacob Roll; Fredrik Lindsten

cs.LGMar 31, 2026

QUEST: A robust attention formulation using query-modulated spherical attention

Authors: Hariprasath Govindarajan, Per Sidén, Jacob Roll, Fredrik Lindsten

Organizations: Linköping University, Sweden · Qualcomm Auto Ltd Sweden Filial

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.

Figures & tables

Figure 1: Class-activation maps for an image from the Macaw class in ImageNet, generated using AG-CAM ( Leem and Seo, 2024 ) . Standard attention concentrates on few bird instances (see first row) and mis-classifies 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 in standard attention. Hence, when the top part of the image is noised, the model focuses on the birds in the bottom part of the image since they are the most salient object in the image then, but results in a misclassification. QUEST attention attends evenly to 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 in 4.1.2 .

Figure 2: Illustration of toy example.

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 of the data. The QKNorm methods obtained $\sim$ 0% overall success rate and their results are available in Figure A1 .

Panel of Figure 4 — 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.

Attention	Scaling	IN-val	IN-v2	IN-ReaL	IN-C	IN-A
		Top-1	Top-1	Top-1	MCE $\downarrow$	Top-1	Top-5
Standard	$1/\sqrt{D_{h}}$	$72.6$	$60.6$	$80.4$	$55.7$	$8.2$	$32.9$
QUEST	-	$\bm{73.4}$	$\bm{61.1}$	$\bm{81.2}$	$\bm{55.0}$	$\bm{8.5}$	$\bm{34.6}$
QNorm	-	$72.7$	$60.8$	$80.6$	$55.3$	$8.2$	$34.5$
QKNorm-HS	† ${\bm{C}}\in\mathbb{R}^{L\times H}$	$72.5$	$60.8$	$80.5$	$56.4$	$7.9$	$33.3$
QKNorm-DS	‡ ${\bm{C}}_{q},{\bm{C}}_{k}\in\mathbb{R}^{L\times D_{h}}$	$71.6$	$59.7$	$79.6$	$57.4$	$7.2$	$31.5$

Table 1: Ablation of different QK normalization methods for ImageNet classification (ViT-Tiny model trained using DeiT for 300 epochs) [ † Liu et al. (2022) , ‡ Dehghani et al. (2023) ]. The Top-1 validation accuracies for ImageNet are the means over 3 independent training runs. The detailed statistics of these results are reported in Table A1 .

Model	Attention	Epochs	IN-val	IN-v2	IN-ReaL	IN-C	IN-A
			Top-1	Top-1	Top-1	MCE $\downarrow$	Top-1	Top-5
DeiT-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$	$\bm{80.2}$	$\bm{68.9}$	$\bm{86.2}$	$\bm{43.2}$	$\bm{20.4}$	$\bm{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$

Table 2: ImageNet Classification using DeiT and DeiT-3 training recipes [ ‡ Dehghani et al. (2023) ]. For longer trainings with larger ViT models (denoted as † ), 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.

Attention	Clean Data			FGSM			PGD
	Top-1	Top-5	NLL	Top-1	Top-5	NLL	Top-1	Top-5	NLL
Standard	$72.50$	$91.45$	$1.190$	$54.23$	$85.28$	$1.827$	$43.65$	$78.18$	$2.503$
QUEST	$\bm{73.33}$	$\bm{91.91}$	$\bm{1.160}$	$\bm{56.90}$	$\bm{86.63}$	$\bm{1.745}$	$\bm{45.26}$	$\bm{79.33}$	$\bm{2.448}$
Elliptical	71.53	90.70	1.254	55.96	85.53	1.746	46.30	80.05	2.231
Elliptical-QUEST	$\bm{72.48}$	$\bm{91.20}$	$\bm{1.214}$	$\bm{56.39}$	$\bm{85.94}$	$\bm{1.741}$	$\bm{47.25}$	$\bm{80.61}$	$\bm{2.211}$

Table 3: Robustness to adversarial attacks using ViT-Ti/16 model trained using DeiT

Attention	SPSA			Auto
	Top-1	Top-5	NLL	Top-1	Top-5	NLL
Standard	$47.17$	$83.95$	$2.027$	$26.57$	$67.60$	$3.230$
QUEST	$\bm{50.70}$	$\bm{85.49}$	$\bm{1.904}$	$\bm{27.29}$	$\bm{67.98}$	$\bm{3.200}$
Elliptical	57.96	86.92	1.654	27.35	67.28	3.014
Elliptical-QUEST	$\bm{59.15}$	$\bm{87.44}$	$\bm{1.613}$	$\bm{28.54}$	$\bm{68.10}$	$\bm{2.965}$

Table 4: Robustness to adversarial attacks using ViT-Ti/16 model trained using DeiT

Model	Attention	Clean data	Corrupted data
		mIoU	mIoU
ViT-Ti/16	Standard	$37.34$	$32.19$
ViT-Ti/16	QUEST	$\bm{38.87}$	$\bm{33.55}$
ViT-S/16	Standard	$43.43$	$38.45$
ViT-S/16	QUEST	$\bm{44.13}$	$\bm{39.19}$

Table 5: ADE20K Image Segmentation

Backbone	Attention	Method	Epochs	Pixel accuracy (%) $\uparrow$	mIoU $\uparrow$	DICE Score $\uparrow$
ViT-B/16	Standard	DeiT-3	$400+20$	$65.22$	$35.53$	$0.4870$
ViT-B/16	QUEST	DeiT-3	$400+20$	$\bm{70.76}$	$\bm{53.54}$	$\bm{0.6597}$

Table 6: Model explainability using AG-CAM

Attention	ZINC	CIFAR10	PATTERN	CLUSTER
	MAE $\downarrow$	Accuracy $\uparrow$	Accuracy $\uparrow$	Accuracy $\uparrow$
Standard	$0.070\pm 0.004$	$72.298\pm 0.356$	$86.685\pm 0.059$	$\bm{78.016\pm 0.180}$
QUEST	$\bm{0.069\pm 0.002}$	$\bm{72.843\pm 0.526}$	$\bm{86.760\pm 0.046}$	$77.894\pm 0.205$

Table 8: GNN Benchmarks using GraphGPS (mean

\pm

s.d over 10 runs)

Attention	PascalVOC-SP	COCO-SP	Peptides-func	Peptides-struct	PCQM-Contact
	F1 score $\uparrow$	F1 score $\uparrow$	AP $\uparrow$	MAE $\downarrow$	MRR $\uparrow$
Standard	$\bm{0.375\pm 0.011}$	$0.341\pm 0.004$	$0.654\pm 0.004$	$\bm{0.250\pm 0.001}$	$0.334\pm 0.001$
QUEST	$0.373\pm 0.003$	$\bm{0.349\pm 0.004}$	$\bm{0.662\pm 0.004}$	$0.251\pm 0.002$	$\bm{0.346\pm 0.001}$

Table 9: Long Range Graph Benchmarks using GraphGPS (mean

\pm

s.d over 4 runs)

Method	Model size (Parameters)	Attention	Clean Data PPL $\downarrow$		Contaminated Data PPL (Corruption %) $\downarrow$
	Model size (Parameters)		Val	Test	Test (1.5%)	Test (2.5%)	Test (5.0%)
Transformer	Small (44M)	Standard	$33.073$	$34.076$	$41.566$	$46.380$	$\bm{60.944}$
Transformer	Small (44M)	QUEST	$\bm{32.966}$	$\bm{33.966}$	$\bm{41.522}$	$\bm{46.299}$	$61.062$
Transformer	Medium (90M)	Standard	$27.441$	$28.851$	$36.234$	$40.866$	$55.012$
Transformer	Medium (90M)	QNorm	$27.233$	$28.688$	$36.262$	$40.853$	$55.451$
Transformer	Medium (90M)	QKNorm-DS	$27.376$	$28.624$	$36.018$	$40.715$	$54.899$

Table 10: Language Modeling on WikiText-103

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Figure A1: Success rates of learning the correct solution to toy example. Models are trained with different hyperparameter combinations with 5 different weight initializations and 5 different realizations of the data.

Figure A2: Distribution of training accuracies at the end of the training for the toy example.

Figure A3: Distribution of test accuracies at the end of the training for the toy example.

Figure A4: Norm of unbiased answer token keys and non-answer token keys in standard attention. In terms of the key norms, we do not observe any distinction between unbiased answer tokens and non-answer tokens.

Figure A5: Norm of unbiased answer token keys and non-answer token keys in QNorm attention. In terms of the key norms, we do not observe any distinction between unbiased answer tokens and non-answer tokens.

Panel of Figure 20 — Figure A6: Maximum logits and their associated query and key norms in different layers for a ViT-Base model trained using DeiT. The model with standard attention crashes around 20000 iterations, which is not observed in the other models. The rapid increase in maximum logits (especially in intermediate and initial layers) is accompanied by an 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 to normalize both queries and keys to ensure training stability.

Attention	Scaling	ImageNet-val
		Top-1 (%)	Top-5 (%)	NLL $\downarrow$
Standard	$1/\sqrt{D_{h}}$	$72.65\pm 0.10$	$91.49\pm 0.03$	$1.187\pm 0.003$
QUEST	-	$\textbf{73.36}\pm\textbf{0.16}$	$\textbf{91.86}\pm\textbf{0.06}$	$\textbf{1.160}\pm\textbf{0.004}$
QNorm	-	$72.73\pm 0.11$	$91.44\pm 0.02$	$1.188\pm 0.004$
QKNorm-HS	† ${\bm{C}}\in\mathbb{R}^{L\times H}$	$72.49\pm 0.10$	$91.43\pm 0.09$	$1.198\pm 0.009$
QKNorm-DS	‡ ${\bm{C}}_{q},{\bm{C}}_{k}\in\mathbb{R}^{L\times D_{h}}$	$71.87\pm 0.18$	$90.96\pm 0.11$	$1.231\pm 0.008$

Appendix

Table A1: Ablation of different QK normalization methods for ImageNet classification (ViT-Tiny model trained using DeiT for 300 epochs) [ † Liu et al. (2022) , ‡ Dehghani et al. (2023) ]. Means and standard deviations of results are reported over 3 independent runs with different seeds.

Data (%)	Attention	ImageNet-val
		Top-1	Top-5	NLL
5%	Standard	$40.82\pm 0.81$	$63.92\pm 0.98$	$3.326\pm 0.074$
5%	QUEST	$\textbf{41.96}\pm\textbf{0.15}$	$\textbf{65.24}\pm\textbf{0.07}$	$\textbf{3.254}\pm\textbf{0.014}$
10%	Standard	$56.02\pm 0.58$	$78.52\pm 0.66$	$2.201\pm 0.042$
10%	QUEST	$\textbf{57.51}\pm\textbf{0.33}$	$\textbf{79.75}\pm\textbf{0.17}$	$\textbf{2.136}\pm\textbf{0.016}$
25%	Standard	$68.87\pm 0.45$	$88.86\pm 0.25$	$1.386\pm 0.020$

Appendix

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.

Learning rate	Epochs	Standard attention		QUEST attention
		Top-1	Top-5	Top-1	Top-5
$1\mathrm{e}-4$	$50$	$62.8$	$84.9$	$63.9$	$85.7$
$2\mathrm{e}-4$	$50$	$66.7$	$87.9$	$67.8$	$88.5$
$5\mathrm{e}-4$	$50$	66.9	87.9	68.5	88.7
$1\mathrm{e}-3$	$50$	–training crashed–		$67.9$	$88.5$
$2\mathrm{e}-3$	$50$	–training crashed–		$67.0$	$87.8$

Appendix

Table A3: Learning rate sensitivity of DeiT training evaluated using a ViT-Small model on ImageNet validation performance

Model	Attention	Epochs	ImageNet-val
			Top-1	Top-5	NLL
CrossViT-9-Dagger	Standard	$300$	$76.4$	$93.4$	$1.009$
CrossViT-9-Dagger	QUEST	$300$	$77.0$	$93.6$	$0.991$
CrossViT-Small	Standard	$300$	$80.5$	$95.5$	$0.834$
CrossViT-Small	QUEST	$300$	$80.9$	$95.6$	$0.826$
CrossViT-18-Dagger	Standard	$300$	–training crashed–

Appendix

Table A4: CrossViT ImageNet Classification

Model	Attention	Epochs	FGSM		PGD		Auto
			Top-1	NLL	Top-1	NLL	Top-1	NLL
ViT-S/16	Standard	$200$	$65.90$	$1.362$	$54.47$	$2.050$	$38.22$	$2.699$
ViT-S/16	QUEST	$200$	$\bm{67.02}$	$\bm{1.328}$	$\bm{57.51}$	$\bm{1.886}$	$\bm{40.98}$	$\bm{2.540}$
ViT-B/16 †	Standard	$400+20$	$69.41$	$1.229$	$52.55$	$2.080$	$-$	$-$
ViT-B/16 †	QUEST	$400+20$	$\bm{70.67}$	$\bm{1.194}$	$\bm{54.64}$	$\bm{2.033}$	$-$	$-$
ViT-L/16	QKNorm-LN	$100$	$58.95$	$1.849$	$51.79$	$2.355$	$-$	$-$

Appendix

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	Parameters	Attention	Epochs	IN-val
				Top-1	Top-5	NLL $\downarrow$
ViT-H/14	600M	Standard	$100$	$68.3$	$88.8$	$1.545$
ViT-H/14	600M	QUEST	$100$	$\bm{69.2}$	$\bm{89.4}$	$\bm{1.495}$
ViT-H/14	600M	Standard	$100+20$	$76.2$	$93.3$	$1.042$
ViT-H/14	600M	QUEST	$100+20$	$\bm{76.8}$	$\bm{93.6}$	$\bm{1.002}$
ViT-2B/16	2.4B	Standard	$10$	$3.1$	$9.4$	$6.153$

Appendix

Table A6: DeiT-3 ImageNet Classification with large-scale ViT models

Panel of Figure 27 — 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.1 . 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.

Panel of Figure 28 — Figure A8: Examples showing class activation maps (CAM) for images containing two distinct objects. The CAM is 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 4.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 to the entire object or all relevant regions.

Dataset	Standard	Standard + [CLS]
EthanolConcentration	28.14	29.28
FaceDetection	67.99	65.24
Handwriting	38.24	42.00
Heartbeat	77.07	77.56
JapaneseVowels	97.84	98.38
PEMS-SF	84.97	83.82

Appendix

Table A7: Impact of using a [CLS] token for UEA Multivariate Time Series Classification using a standard Transformer

Dataset	Standard	QNorm	QKNorm	QUEST
EthanolConcentration	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
JapaneseVowels	98.38	98.92	98.38	98.38
PEMS-SF	83.82	80.92	79.19	80.92

Appendix

Table A8: Ablation experiment for UEA Multivariate Time Series Classification using different forms of attention

Model	Parameters	Attention	nuScenes val. mIoU (%)
PointTransformer-V3	46.2 M	Standard	80.40
PointTransformer-V3	46.2 M	QKNorm-DS	80.37
PointTransformer-V3	46.2 M	QUEST	80.83

Appendix

Table A9: Pointcloud segmentation using PointTransformer-V3 model on the nuScenes dataset. We report the mIoU on the validation set.

Hyperparameter	ZINC	CIFAR10	PATTERN	CLUSTER
# GPS Layers	10	3	6	16
Hidden dim	64	52	64	48
GPS-MPNN	GINE	GatedGCN	GatedGCN	GatedGCN
GPS-GlobAttn	Transformer	Transformer	Transformer	Transformer
# Heads	4	4	4	8
Dropout	0	0	0	0.1

Appendix

Table A10: GPS hyperparameter setup for standard GNN benchmarks ( Dwivedi et al., 2023 )

Hyperparameter	PascalVOC-SP	COCO-SP	PCQM-Contact	Peptides-func	Peptides-struct
# GPS Layers	4	4	4	4	4
Hidden dim	96	96	96	96	96
GPS-MPNN	GatedGCN	GatedGCN	GatedGCN	GatedGCN	GatedGCN
GPS-SelfAttn	Transformer	Transformer	Transformer	Transformer	Transformer
# Heads	8	8	4	4	4
Dropout	0	0	0	0	0

Appendix

Table A11: GPS hyperparameter setup for long-range graph benchmarks ( Dwivedi et al., 2022 )

Attention	kNN	Linear	Finetuning	Few-shot
				1 img/cls	2 imgs/cls	5 imgs/cls	1% imgs
Standard	$78.7$	$80.3$	$84.1$	$51.6\pm 0.1$	$61.1\pm 0.7$	$68.3\pm 0.3$	$72.3$
QUEST	$\bm{79.0}$	$\bm{80.5}$	$\bm{84.4}$	$\bm{52.5\pm 0.2}$	$\bm{61.7\pm 0.6}$	$\bm{69.0\pm 0.1}$	$\bm{72.7}$

Appendix

Table A12: Self-supervised pre-training on ImageNet with iBOT and evaluating on ImageNet tasks

Attention	Transfer linear probe	Image Retrieval				VOS
	Average	RParis		ROxford		DAVIS-2017
		M	H	M	H	$\mathcal{J\&F}_{m}$	$\mathcal{J}_{m}$	$\mathcal{F}_{m}$
Standard	$81.5$	$65.4$	$38.1$	$38.1$	$13.8$	$63.1$	$\bm{61.9}$	$64.2$
QUEST	$\bm{81.9}$	$\bm{65.8}$	$\bm{39.3}$	$\bm{38.6}$	$\bm{15.3}$	$\bm{63.2}$	$61.8$	$\bm{64.5}$

Appendix

Table A13: Self-supervised pre-training on ImageNet with iBOT and evaluating on transfer tasks

Attention	Acft.	Cal101	DTD	Flwrs.	Food	Pets	SUN	Avg.
Standard	$58.1$	$\bm{95.5}$	$\bm{74.7}$	$94.8$	$83.6$	$93.9$	$\bm{70.2}$	$81.5$
QUEST	$\bm{59.9}$	$95.1$	$74.4$	$\bm{95.8}$	$\bm{83.9}$	$\bm{94.3}$	$69.8$	$\bm{81.9}$

Appendix

Table A14: Self-supervised pre-training on ImageNet-1K with iBOT-vMF and evaluating with transfer linear probes

Explore similar work

QSV: Quat-Sphere-Vision for Coupled Quaternion Attention on Spherical Lattices
Sep 24, 2026Nicholas Foley, Devin Marinelli, Donny Moore +2Attention Mechanisms Vision Transformer
Department of Computer Science, University of Texas at San Antonio
Query Expansion and Key Specialization in Transformer Attention Geometry
Sep 28, 2026Vidit Gupta, Siddhesh Nadkarni, Mihik Chaudhari +2Neural Representation Geometry Transformer Attention
Dwarkadas J. Sanghvi College of Engineering Mumbai, India
Why Attend to Everything? Focus is the Key
Mar 12, 2026Hengshuai Yao, Xing Chen, Ahmed Murtadha +8Attention Mechanisms Efficient Attention
Sapient Intelligence · University of Alberta · Tsinghua University

Sep 24, 2026J/K move · Enter open · S save

QSV: Quat-Sphere-Vision for Coupled Quaternion Attention on Spherical Lattices

Nicholas Foley, Devin Marinelli, Donny Moore, Diego Enriquez+1

Department of Computer Science, University of Texas at San Antonio

Attention Mechanisms Vision Transformer Sparse Attention Image Classification

Open paper AskPDF arXiv

In standard attention, three separately learned projections decide how strongly a token attends to each neighbor (

W_Q

W_K

) and how the attended features are transformed before aggregation (

W_V

). We study Quat-Sphere-Vision (QSV), a sparse spherical vision model that replaces this projection triple with a single learned unit quaternion per token: the relative quaternion

r_{ij} = q_i^{*} \otimes q_j

supplies both the attention logit

\operatorname{Re}(r_{ij})

and a sandwich-product feature transport

x \mapsto r_{ij} \otimes x \otimes r_{ij}^{*}

, with messages passed over sparse kNN graphs on concentric Fibonacci spheres. Ablations that change only the targeted component show the two roles to be asymmetric. Removing the transport reduces test accuracy by about four percentage points on CIFAR-10 and CIFAR-100 (single runs per CIFAR-100 variant), while replacing the learned attention weights with uniform averaging leaves it essentially unchanged. Parameter-matched controls then remove the geometry itself: standard attention on the same graph exceeds QSV (mean

87.3\%

vs.

85.9\%

), and the same model on a flat 2D lattice reaches

91.1\%

, within

2.1

points of a ResNet-20 trained under the same pipeline (single run). In the coupled kernel, nearly all of the learned pairwise computation resides in the transport channel.

Abstract

Explore similar work

QSV: Quat-Sphere-Vision for Coupled Quaternion Attention on Spherical Lattices

Query Expansion and Key Specialization in Transformer Attention Geometry

Why Attend to Everything? Focus is the Key