Spectral Conditioning of Attention Improves Transformer Performance
Authors: Hemanth Saratchandran
Simon Lucey
Venue: NeurIPS 2025 poster ()
Keywords: spectral conditioning
attention
Primary area: deep_learning
TL;DR: We propose spectral conditioning of attention layers to improve Jacobian conditioning, leading to more stable and efficient optimization with negligible computational overhead and consistent gains across diverse transformer architectures.

Abstract: We present a theoretical analysis of the Jacobian of a attention block within a transformer, showing that it is governed by the query, key, and value projections that define the attention mechanism. Leveraging this insight, we introduce a method that systematically alters the spectral properties of each attention layer to reduce the Jacobian’s condition number, thereby improving the overall conditioning of the attention layers within a transformer network. We empirically show that this improved Jacobian conditioning translates to enhanced performance in practice. Our approach is simple, broadly applicable, and can be easily integrated as a drop-in replacement for a wide range of existing attention mechanisms. We validate its effectiveness across diverse transformer architectures and tasks, demonstrating consistent improvements in performance.

Spectral Conditioning of Attention Improves Transformer Performance
Hemanth Saratchandran
Simon Lucey
Australian Institute for Machine Learning
Australian Institute for Machine Learning Adelaide University simon. leeenn.uceyGadlliidd.. edu.an
Adelaide University
hemanth. sawatcchandran@allllliaao. edu.
Abstract
We present theoretical analysis of the Jacobian of a attention block within a transformer, showing that is governed by the query, key, and value projections that define the attention mechanism. Leveraging this insight, we introduce a method that systematically alters the spectral properties of each attention layer to reduce the Jacobian's ss condition number, thereby improving the overall conditioning of the attention layers within a transformer network. We empirically show that this improved Jacobian conditioning translates to enhanced performance in practice. Our approach is simple, broadly applicable, and can be easily integrated as a drop- in replacement for wide range of existing attention mechanisms. We validate effectiveness across diverse transformer architectures and tasks, demonstrating consistent improvements in performance.
1 Introduction
Since its introduction in [3], the modern transformer has emerged as one of the most influential architectures in modern machine learning, yielding advances across wide range of domains,
computer vision 7 20 33 S), and robotics 27 [23] At the heart of its success lies the attention mechanism, which enables the model to capture complex relationships by computing pairwise interactions between all input tokens. By dynamically assigning importance different tokens, attention allows transformers model global dependencies, making them a foundational architecture in deep learning.
including natural language processing (NIP) 34 42 41
In this work, investigate the conditioning of the Jacobian associated with the attention layer in transformer models. The conditioning of a matrix is measured by its condition number, the ratio of its largest to smallest singular value, with high values indicating ill-conditioning. Poorly conditioned Jacobians can hinder performance for gradient- based optimizers (25]. Recent advances in feedforward neural networks have shown that improving Jacobian conditioning can lead to better optimization and generalization performance. For instance, (30) proposed a weight normalization method to address ill-conditioning, while [19] demonstrated that reducing the condition number of the neural tangent kernel (NTK) leads to better convergence. Despite these insights, the conditioning of Jacobians in transformer models, particularly within attention layers remains largely unexamined. Addressing this gap is the central focus of our work.
We present theoretical framework showing that the conditioning of the Jacobian in attention layer is influenced by the conditioning of the queries, keys, and values matrices that constitute the attention matrix within that layer. Building the spectrum of these matrices by adding a carefully designed correction term, yielding a new attention block that we call spectral conditioned atfention, resulting in significantly improved Jacobian conditioning. Our analysis shows that the optimal correction can derived using the singular value decomposition (SVD) of the query, key, and value matrices. However, computing the SVD directly can
this insight, introduce method that modifies
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

as well within the context of fine-tuning be prohibitively expensive in large-scale models. To address this, we also propose computationally efficient approximation that achieves comparable improvements in conditioning with substantially lower overhead.
While full theoretical proof linking our methodology improved convergence in gradient-based optimization would require analyzing the NTK of a transformer, an extremely difficult task, we provide strong empirical evidence that our approach improves transformer performance across a wide range of applications, including image classification, object detection, instance segmentation, and natural language processing. A key advantage of our approach is its compatibility with broad class which we confirm empirically. Our main 6 38
of different attention mechanisms 20 33 contributions include:
39
1. A theoretical framework that analyses the condition number of the Jacobian of attention showing its dependence on the condition number of the query, key and value matrices comprising an attention layer.
2. The introduction of spectral conditioned attention that adds correction term each of the query, key and value matrices resulting in a better conditioned attention Jacobian.
We validate the above contributions across a range of transformer architectures and applications, in- cluding image classification, object detection, instance segmentation, and natural language processing, consistently demonstrating improved performance in each setting.
2 Related Work
Conditioning. Several recent works have highlighted the critical role of conditioning in the opti- mization and generalization of neural networks. In [30], the authors focus the condition numbers of weight matrices in feedforward architectures, showing that well-conditioned weights correlate with improved performance across tasks. They propose a preconditioning strategy that multiplies each weight matrix by a carefully constructed transformation, effectively reducing its condition number and facilitating more stable training. In complementary line of work, (19) investigates neural network optimization from the perspective of the neural tangent kernel (NTK). The authors demonstrate that improving the conditioning of the NTK leads better convergence, especially in the infinite-width regime where the NTK governs training dynamics 14]. In 12D) it is shown that normalization improves the conditioning of various neural network architectures. Similarly, 40 demonstrate that normalizing attention weights enhances training convergence, which can be interpreted as an alternative means of improving the conditioning of attention layers. In different approach. [ shows that simply increasing the depth of feedforward networks can enhance their conditioning, thereby improving the effectiveness of gradient- based optimization methods. Recent work on transformers has examined conditioning from multiple perspectives, including tokenization, activation functions, and optimization 28 17
15
29
[6 3 22 (2).
Attention. A range of methods have been proposed to improve the efficiency and scalability of Transformer architectures, often by rethinking the design of attention mechanisms reducing computational complexity. For example, the Data-Efficient Image Transformer (DeiT) 333 introduces distillation tokens achieve competitive performance in vision tasks without relying on large- scale datasets. The Cross-Covariance Image Transformer (XCiT) [4] reformulates attention by leveraging cross-covariances of spatial features, enabling efficient spatial interactions with reduced overhead. Swin Transformer (20) introduces a hierarchical architecture and novel shifted window- based attention mechanism, significantly improving efficiency and effectiveness in vision tasks. DiVIT 6 extends traditional vision transformers by incorporating both spatial and channel attention mechanisms, enhancing feature extraction across multiple dimensions. The Nyströmformer [38 takes a different approach by approximating full self-attention using the Nyström method, significantly lowering the quadratic complexity to near-lines while maintaining the core properties of attention, making it especially effective for long men-sequence modeling.

---

3 Methodology
3.1 Preliminaries
In this section, we define the transformer architecture by describing the structure of transformer layer, along with establishing notation for various mathematical quantities.
A transformer layer can be represented as mapping
T NND + NXD
(I)
which is formally defined by
T(X) = ((A(X) +X,,
(2)
where denotes a feed forward neural network with a residual connection, and A represents an attention mechanism. In general, layer normalization added however for simplicity omit layer normalization for this discussion.
For the theoretical framework we will primarily focus on self-attention, which is one of the most common forms of attention. Self-attention composed of three learnable matrices, query We Dxd key WK E RDdd, and value Wy E Dxd defined for an input sequence X E NR by
A(X) softmax X W (Wo
where softmax the softmax activation that acts row-wise on a matrix 266 Note that then
A(X) Nxx When training transformer the self-attention parameters are given by the weight matrices Wo. Wx and Wv. For further details on transformers readers may consult [26].
The self-attention map of layer in transformer A (XX) has parameters given by those parameters in from the previous layer and those given by Wo. Wx and Wv that define A(X). Our work will consider the Jacobian of A(X) with respect to the parameters within the layer of AX)), namely Wo. WK and Wv. Therefore, when we speak of the Jacobian of ( it will with respect to Wo. WK. K. Wy. We will denote this Jacobian by (X)))) and note that it is defined by
J(A(X)) (XA(X) OA(X) 8A(X)
4
aWQ aWk dwy
Given matrix W m we denote the vectorization of by vec(W) E R.mxx for such a matrix there is transformation Tmn Rmmxxmn mn such that Tan Tecvec(W vec(W) = where W denotes the transpose of W. The matrix Tra is known permutation matrix [22). The maximum singular value of a matrix W will be denoted by 7 (W) and the minimum singular value by min (W). We will use the standard terminology SVD to denote the singular value decomposition of matrix. Given vector R the notation ||=||2 will denote the vector n-rorm of z.
Note that (cc(W)) commutation matrix and is
22
3.2
Main Theorems
In this section, we analyze the conditioning of the Jacobian of the self-attention matrix in transformer. We will show that the condition number of the Jacobian depends on the condition number of the queries, keys and values matrices. This then motivates of the queries, keys and values matrices can lead to a lower condition number for the Jacobian of the self-attention matrix in each layer of transformer, which we empirically verify in Section 4 We have chosen to focus our theory on self-attention so as to yield concrete formulas. However, our theory goes through for more general attention mechanisms, Section A. for discussion. Proofs of all lemmas and theorems in this section can be found in Section A.1
key insight: reducing the condition number
Definition 3.1. Let A be an x matrix of full rank. The condition number of A, denoted by K, is defined as
T max (m)x(()
<(A)
Omin (A)
where Tmax (A) denotes the maximum singular value of and mmin( the minimum singular value of A, which know is non-zero A is full rank.

---

Our objective is to analyze the condition number of the self-attention block within transformer. We demonstrate that the condition number of its Jacobian is influenced by the condition numbers of the query, key, and value weight matrices. By appropriately conditioning these matrices, we show that the Jacobian's condition number can be significantly reduced, leading to a more stable and well-conditioned attention mechanism suitable for applications.
To establish this, we first examine the derivatives of the self-attention block with respect to the parameters Wo. WK. and Wv, yielding Theorem 3.3 which is then used to establish a condition
number bound on the Jacobian given in Theorem 3.4 As starting point, Theorem 3.2 presents the derivative of the softmax function, which plays a central role in the analysis.
Lemma 3.2. Let : + R denote the function A(z) Diag(z) 2 xT. We then have that
@softmax
(2) = A(softmaxx2))
Òr
Theorem 3.3. Let ( denote a self-attention matrix with input X as defined by Equation (3
Then
0A(X) (Wix* @ IN) A(sofmax(xx(W777)) (XWx @ x
DWQ
0A(X) (WTX* @ IN) (|ssotmax(xWoWWZZ) (X XWQ) Tpa
(8)
OWk
(A(X)
la @ softmax(XWo
aWy
where A is defined in Theorem 3.2 and Da is the commutation matrix satisfying TDa)
vec(Wk) (see Section 3.1
Theorem 3.4. Let A(X) denote a self-attention matrix, as defined in Equation ( with input
X and let (A))) denote its Jacobian with respect the parameter matrices Wo. Q. WK and
Wv as defined in Equation 4h. Assume that J (A(()) has full rank so that (A is finite. Then
x(J(A(X))) (Msot))((saa ((WW)((W))) =(Wk)) (10)
+ xx
where A is defined in Theorem 3.2
Theorem 3.4 shows that x(J(A(X))) is bounded above by a sum of two terms:
-
(W x x(sottmax((XWWWWXX))
k(Wk))
(11)
(12)
Observation: Directly conditioning the Jacobian of the self-attention map with respect to the parameters Wo. WK. k. and Wv during training would require computing the Jacobian at every iteration, which is prohibitively expensive. However, Theorem (heorrem3.44offer
more efficient alternative: reducing the upper bound in Equation (1O). Since we have access to Wq. WK. , and Wy, lowering their condition numbers reduces the quantity in Equation (TT). thereby tightening the bound in Equation (0).. The following theorem shows how to do this.
Theorem 3.5. Let Wq. WK and Wv denote the parameters of self-attention matrix A(X).
Then there exist matrices Co. Cx and Cv such that
x(Wq + Cq). n(Wx Ck). x(Wy + Cv) 52
(13)

---

Self-attention (Viswani et al.)
Feed Forward Network
XW W X
softmax
XV
vd
Layer Norm
Spectral conditioned self-attention
X(Wo + Co) Cx)
Spectral Conditioned Self-attention
softmax
x(Wr+
vd
Layer Norm
An illustration of spectrally conditioned self-attention within transformer layer. At each layer, the self-attention weights Wo. WK. and Wy are modified by adding correction terms Cq. CK+ and Cv, respectively. The correction terms Co. Ck, and Cy are initialized before training using Theorem 3.8 and remain fixed throughout training.
Figure I:
Overview of proof of Theorem 3.5
The same proof works for Wq. WK and Wy. Thus will give the main steps for We. As in Section 3.1 let We RDxd
1. Using the SVD write Wq USVT where U E DxD is the left singular vectors, S E Dxd is the matrix of singular values on the main diagonal and V E Rdxd are the right singular vectors.
2. Define new matrix Cq = USV where 3 has mmax (Wq) on its main diagonal and zeros everywhere else.
3. The matrix Wº + Cq then satisfies x(Wq + CQ) 2.
3.3 Spectral Conditioned A Attention
In Theorem 3.5
showed that adding correction terms Co. Ck, and Cy the weight matrices Wo. Wx. and Wv. respectively, significantly reduces their condition numbers k(WQ), k(Wx) and n(Wv). This motivates the following definition.
Definition 3.6. We define spectral conditioned attention by
SpecA(X) softmax(X((Q + CQ) + +Cx) +C)
(14)
where Co. CK and Cy the correction terms given in Theorem
3.5
We call the process of adding correction matrices Co. Ck and Cv to Wo, WK and Wv respectively lower their condition number spectral conditioning.
together with Theorem 3.4 shows that spectrally conditioned attention reduces the upper bound in Equation (TO). thereby yielding tighter bound on the condition number of the Jacobian of the self-attention layer.
Remark 3.7. Theorem 3.5
The proof overview for Theorem showed the correction term Co (and Ck. Cv) derived by taking SVD. Computing the of each of Wq. Wx. Wy at each iteration of training would lead to significant memory overhead. We therefore provide an implementation friendly form of spectral conditioned attention, motivated by the following theorem (proof given in Section Tb.
5

---

Theorem 3.8. Let A E _ and let I E RRuxn be the matrix that has 1 on its main k x k diagonal, where min n), and zeros elsewhere. Let À 2 2 be a fixed constant and assume that Amax (A)+2 v Tmax (A) Then k( (A + A1&) < <(A).
X-0mm (4)
The key difference between Theorem 3.5 and Theorem 3.8 when applied Wq. WK. and Wv is that
Theorem 3.5 guarantees a condition number strictly less than 2, whereas Theorem 3.8 only ensures
a reduction relative to the original i.e. <(WQ + X < K: (We), and similarly for WK and Wy.
However, the advantage of Theorem 3.8 over Theorem 3.5 is that the correction term X does not
require computing the SVD of Wq. Wx or Wv and hence is much more memory efficient.
Implementation: We initialized correction matrices Cg. Ck, and Cv as A/& with À = 10 (see Section 1 for an ablation study on A), and used spectral conditioned attention where the modified weights were Wq + Co. WK + Ck, and Wv + Cv. These correction matrices fixed during training and are not updated, incurring no additional memory overhead during backpropagation. An overview the resulting attention architecture shown in Figure
!
For comparison of spectral conditioned attention using Theorem 3.5 and Theorem 3.8 on vision transformers see Section A.2.
4 Experiments
In this section, we evaluate our insights from Section on a variety of transformer applications. For each application will consider original transformer architecture used within the literature and compare it with one employing spectral conditioning on its attention layer.
Implementation. In all cases, we used the implementation described in Section 3.3 correction terms Co. Ck, and Cy are added to the query, key and value matrices We. WK and Wy respectively within each attention layer, with À 10 (see Section A.2.1 for an ablation on A). These terms are not updated during training and therefore do not introduce any additional trainable parameters. For more details see Section A.2
where fixed
4.1 Image Classification
Vision transformers. We applied spectral conditioned attention to variety of modern vision transformers: the original vision transformer (ViT-B) 77. Swin transformer (Swin-B) (20) Cross-
Covariance image transformer (XCIT-M) [4) Data efficient image transformer (DeiT-B) 333 Dual
attention vision transformer (DaViT-B) (6) for image classification on ImageNet-IK, Ik. Each of these
vision transformers uses different attention block to the standard self-attention used in the ViT-B architecture. However, spectral conditioned attention (see Theorem 3.6) works by adding a correction term each query, key and value matrix and hence can easily be plugged into these modern variants of self-attention. We include these latter models to illustrate that, while our theoretical analysis in Section 3 centers self-attention, spectral conditioned attention can be readily incorporated into modern transformer architectures, including those with more complex designs.
Validating the theory. We validate the theoretical results from Section 3 on VIT-B model trained on the ImageNet-II 1k dataset. For each experiments we ran five trials and took the mean and standard deviations. For each attention head in every transformer layer, we compute the minimum singular value of the query (We), key (Wx) and value (Wv) weight matrices, and then average these values across all heads and layers. To improve the conditioning, we add correction terms Cq. Ck. and Cy as described after Theorem 3.8 resulting in spectral conditioned attention matrices Wo + Cq- WK + CK, and Wy + Cv. The left figure in Figure 2 shows the average minimum singular value during training and in each case we see that the minimum singular value of the corrected weights is higher. Similar plots for the maximum singular value is shown in Section A.2. The middle plot in Figure @egure@comper compares the average condition numbers of the original and spectral conditioned query, key and value matrices, demonstrating that the latter are substantially better conditioned, thus empirically verifying the claim of Theorem 3.8 Finally, the right plot reports the average condition number of the

---

Jacobian of the self-attention matrices, both for the standard ViT-B and for the spectral conditioned version. also includes the theoretical upper bound from Theorem 3.4 While the spectral correction was motivated by this theoretical bound, the results clearly show that it leads to better-conditioned Jacobians in practice, supporting use in designing more stable attention mechanisms. Note that each plot shows the mean over five trials. Standard deviations cannot be seen due to the log plot of the y-axis but can be found in separate plots in Section A.2.1] We repeated the experiment using the XCTT-M transformer 4 As shown in Figure 3 the results exhibit the same trend observed with ViT-B, further supporting the theoretical findings in Section 3
100
3371
100
2300
2041
30
16
Spectall ceodtorddd sturtie
0g
150
227
200 822 650
Figure 2: Analysis for ViT-B. Left: Average minimum singular value of the query, key, and value projection matrices (Wq. WK. Wv) and their spectrally conditioned counterparts (WQ + Co. WK + CK. + Cv) throughout training. Middle: Condition numbers of Wq. WK. and Wy, and their spectrally conditioned forms during training. Right: Average condition number of the self-attention Jacobian over the course of training, before and after spectral conditioning, along with the theoretical bound from Equation 10
305
M4 Doditon
30-4
10-4
30
300
350
200
353
300
253
400
100
250
200
250
300
200
0
53
230
150
238
250
300
300
438
Figure Analysis for XCTT-M. Left: Average minimum singular value of the query, key, and value projection matrices (Wq. WK. Wv) and their spectrally conditioned counterparts (WQ Cq- WK + Ck. Wy + Cv) throughout training. Middle: Condition numbers WQ. WK. and Wy. and their spectrally conditioned forms during training. Right: Average condition number of the self-attention Jacobian over the course of training, before and after spectral conditioning, along with the theoretical bound from Equation (10)
Rank assumption. In the statement of Theorem ensuring the condition number in Equation TO) remains finite. Empirically, we observed this assumption holds across all vision transformer architectures considered. For both ViT-B and XCYT-M, the final plots in Figure 2 and Figure confirm that the average condition number remains finite and stable throughout training, indicating no rank deficiency in the Jacobian of each layer.
assumed that the Jacobian full rank,
Results on ImageNet-1k: We trained two versions of each transformer: a baseline model and one incorporating spectral conditioned attention. For each base architecture, we observed that the minimum singular values of the query, key. and value matrices remained below 1 throughout training (see Figure for ViT-B and Figure XCTT-M), allowing us to apply Theorem (3. implement spectral conditioning as described in Section 3.3, All models were trained using the AdamW optimizer, see Section A.2.I We trained each model five times with five different random seeds. The final results, summarized in Table [ show the mean test accuracy with the standard deviation in brackets. We observe that in every case, the spectral conditioned variant achieves higher test accuracy than its unmodified counterpart. Hardware, memory cost and FLOPS analysis is shown
7

---

in Section A.2.1 An ablation on À is given in Section A.3 Furthermore, spectral conditioning was applied together with layer normalization as shown in Figure A discussion comparing the two is given in Section A.3
Table I: Comparison of Vision Transformers with their original attention architecture versus one with spectrally conditioned attention (Spec. cond.), pre-trained on the ImageNer- Ik dataset. We report Top-1 % classification accuracy. In each case, spectral conditioning improves performance over the original attention mechanism.
ViT-B
DeïT-B
Swin-B
XCIT-M
DaViT-B
Original
80.7 (#:0.41)
81.6 (=0.30)
83.4 (+:0228
82.6 (+0.39)
84.3 (#0.26)
Spec. cond.
81.7 (+0.38)
82.6 (#0.32)
84.1 (4:.25)
83.5 (+0.35)
84.9 (+0.21)
4.2
Object Detection and Instance Segmentation
In this section, we evaluate our methodology on two downstream tasks: object detection and instance segmentation. The goal is to assess the effectiveness of spectral conditioned attention in a fine-tuning setting. We begin by pretraining an XCiT architecture [4] on the ImageNet-Ik Ik dataset, followed by fine-tuning the COCO 2017 dataset [18]. The XCIT models used as backbones within the Mask R-CNN framework [33, enhanced with a Feature Pyramid Network (FPN) to improve multi-scale feature representation. To integrate XCiT with FPN, modify its column structure to extract intermediate features from multiple layers. Specifically, we use the 12-layer XCiT-Small (XCiT-S) adapting their its strides from the original fixed stride of 16 to [4, 8, 16, 32) match the FPN hierarchy. Downsampling is performed via max pooling, and upsampling is achieved using a single transposed convolution layer. We conduct this procedure for both the original XCIT-S and one incorporating spectral conditioned attention blocks.
Downstream results. The results for both object detection and instance segmentation are shown in T We ran five trials each with a different seed and plotted the mean metric with standard deviation in brackets. The reported metrics include Ab (Average Precision for bounding box predictions), AP 50/73 (Average Precision at IoU thresholds of 0.50 and 0.75 for bounding boxes), APM (Average Precision for mask predictions), and APM 50/75 (Average Precision at IoU thresholds of 0.50 and 0.75 for mask predictions). In each case see the XCIT employing spectral conditioned (spec. cond.) out performs the original architecture. Hardware and memory overheads are discussed in Section A.2.2
Table Performance evaluation of object detection and instance segmentation on the COCO dataset. For each metric, our spectrally conditioned architecture (Spec. cond.) outperforms the original.
Model
APb
AP 50
AP
Apm
APM
AP
original 44.9 (1:0.33) 66.1 1 (1-0.31) 48.9 0=33) 40.1 (1:0.29) 63.1 (100.1) 42.8 (1:0.34)
Spec. cond. 45.6 (+0.31) 66.7 (+:036) 49.6 (=.3)) 40.5 (4:0.3) 63.4 (+.0288 43.3 (+.032)
4.3 Nyströmformer on LRA Benchmark
To evaluate the effectiveness of our method, we applied it to the Nyströmformer architecture [38]. which specifically designed handle long-range dependencies efficiently. Experiments were conducted on the Long-Range Arena (LRA) benchmark [32], suite of tasks targeting the ability of models to process extended input sequences. We trained two variants of the Nyströmformer: one baseline and one with spectral conditioned attention.
Validating the theory. We validate the theoretical results from Section 3 on a Nyströmformer trained on the LRA text classification task. Each experiment was repeated five times, reporting the mean and standard deviation. We computed the minimum singular values and condition numbers of W. WK. Wy, and their spectral-conditiooee forms Wq + Co. WK + Ck, and Wy + Cy. We
8

---

also evaluated the average condition number of the attention Jacobian, its conditioned variant, and the theoretical bound from Equation (TO). The observed trends closely follow those for ViT-B and XCTT-M in Section 4.1 supporting our theoretical findings. Results are shown in Figure 4 standard deviations and maximum singular value plots are provided in Section A.2.T
10"
10-4
10"
Specir cOSTORor Theoniccal bound
C
2006 5003 7506 10083 22030 20003 37500 20003
25:00 5600 7540
20006
32903 12406 17500 20036
0 2500 5068
7500 10000
23500 15000 37500 2006
Figure 4: Analysis for Nyströmformer on text classification task. Left: Average minimum singular value of the query, key, and value projection matrices (Wq. WK. Wy) and their spectrally conditioned counterparts (Wq + Co. WK + Ck. Wy+ + Cv) throughout training. Middle: Condition numbers of We. WK. K. and Wv. and their spectrally conditioned forms during training, Right: Average condition number the attention Jacobian over the course training, before and after spectral conditioning, along with the theoretical bound from Equation (10)
Results. Table 3 shows the results of the experiment, As can be seen from the table the Nyström- former that employed spectral conditioned attention outperformed the original Nyströmformer from 38) in every task in the LRA benchmark. We used the cxact training setup and hyperparameters from [38]. Hardware and memory overheads are discussed in Section A.3 and an ablation on is given in Section A.3 A discussion comparing spectral conditioning and layer normalization is given
in Section A.3
Table 3: Comparison of a Nyströmformer using its original architecture (original) with a Nyström- former using spectral conditioned attention (Spec. cond.) on the LRA benchmark. We report evaluation accuracy (%). As shown, our methodology consistently improves performance across all
tasks.
Model
ListOps
Text
Retrieval
Image
Pathfinder
Original
37.1 (#:0211)
63.8 (+0.24)
79.8 (+:0226
39.9 (+0.20)
72.9 (+0.255)
Spec. cond.
37.8 (#:0233)
64.8 (+0.20)
80.6 (+0.27)
40.2 (=0.22)
73.7 (=0.23)
4.4
Language Modeling with Crammed BERT
We applied our insights from Section to a Crammed BERT language model, trained entirely from scratch using masked language modeling following the approach IT]. The model consists of 110 million parameters, with 12 transformer layers and 12 attention heads per layer. For this experiment, we train two versions from scratch: the original Crammed BERT baseline model from [I] and a variant incorporating spectral conditioned attention in each layer, following the implementation from Section 3.3 Both models are trained on The Pile dataset 9 a large-scale corpus designed for language model training. following the pretraining regime carried out in [II]. After pretraining, we evaluate the performance on the GLUE benchmark (35] following the evaluation methodology outlined in [11 Each model was trained five times on the Pile dataset [9] with five different random seeds. Table 4 shows the mean results with standard deviations shown in brackets. As can be seen from the table, spectral conditioned attention outperforms the baseline model in every downstream task in the GIUE benchmark. For the Hardware and memory overheads we refer the reader to Section A.4

---

Table 4: Evaluation of pre-trained Crammed BERT on the GLUE benchmark using the original architecture (Original) and our spectrally conditioned variant (Spec. cond.). Spectral conditioning improves performance across all tasks.
MNLI
QNLI
QQP
Avg.
Original
83.8 (1:021) 92.3 (10.18) 86.3 (2221) 55.1 (1:0200) 90.1 (1:0233 87.3 (15..15) 85.0 (1:0255 48.9 (::0.26) 78.6
Spec. coad. 84.3 20) 92.7 (20.9) 86.6 (=0.22 (#0024) (10.22) (2:0.3)) (#0.23) 51.7 (#:024) 79.4
5 Limitations
Our spectral conditioning approach is derived by optimizing an upper bound on the condition number of the self-attention Jacobian, rather than directly minimizing the Jacobian's condition number itself. While this bound provides useful theoretical guidance and aligns with empirical improvements, remains an indirect proxy. Developing methods to efficiently estimate and control the exact Jacobian conditioning during training would be valuable direction for future work. Additionally, due to computational resource constraints, our experiments were limited models with up 100 million parameters. Whether spectral conditioning offers similar benefits for large-scale transformer models with billions of parameters remains an open question.
Conclusion
We presented a theoretical framework linking the conditioning of self-attention Jacobians to the spectral properties of the query, key, and value matrices transformer architectures. Building on this insight, we introduced spectral conditioned self-attention, simple and broadly applicable method that modifies these matrices with carefully designed correction terms to improve Jacobian conditioning. Our empirical results demonstrate that this approach consistently yields good performance across diverse transformer models and tasks, including image classification, object detection, language modeling, and long-range sequence learning.
10

---

References
[1] Naman Agarwal, Pranjal Awasthi, and Satyen Kale. A deep conditioning treatment of neural networks. In Algorithmic Learning Theory, pages 249-305. PMLR, 2021.
[2] Paul Albert, Frederic Z Zhang, Hemanth Saratchandran, Anton van den Hengel, and Ehsan Abbasnejad. Towards higher effective rank in parameter-effiiiet fine-tuning using khatri-rao product. arXiv preprint
2025.
[3] Paul Albert, Frederic Zhang, Hemanth Saratchandran, Cristian Rodriguez- Opazo, Anton van den Hengel, and Ehsan Abbasnejad. Randlora: Full-rank parameter-efffcient fine-tuning of large models. arXiv preprini urXiv:2502.00987. 2025.
[4] Alaaeldin Ali, Hugo Touvron, Mathilde Caron, Piotr Bojanowski, Matthijs Douze, Armand Joulin, Ivan Laptev, Natalia Neverova, Gabriel Synnaeve, Jakob Verbeek, et al. Xcit: Cross- covariance image transformers. Advances in neural information processing systems, 34:20014- 20027, 2021.
[5] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference on computer vision, pages 213-229. Springer, 2020.
[6] Mingyu Ding, Bin Xiao, Noel Codella, Ping Luo, Jingdong Wang, and Lu Yuan. Davit: Dual attention vision transformers. In European conference on computer vision, pages 74-92. Springer, 2022.
[7] 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. arXiv preprint 20Ii::0000.11929, 2020.
[8] Joel Franklin. Matrix theory. Courier Corporation, 2000.
[9] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Aadi Thite, Eric Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint 20000000077 2021.
[10] Jonas Geiping, Cramming,
piappiiiinnraannngggg........... 2023.
[11] Jonas Geiping and Tom Goldstein. Cramming: Training a language model on a single gpu in one day. In International Conference on Machine Learning, pages 11117-11143, PMLR, 2023.
[12] Cameron Gordon, Yiping Ji, Hemanth Saratchandran, Paul Albert, and Simon Lucey. Com- pressing sine-activated low-rank adapters through post-training quantization. arXiv preprint arXiv:2505.218955 2025.
[13] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In Proceedings of the IEEE international conference on computer vision, pages 2961-2969, 2017.
[14] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[15] Yiping Ji, James Martens, Jianqiao Zheng, Ziqin Zhou, Peyman Moghadam, Xinyu Zhang, Hemanth Saratchandran, and Simon Lucey. Cutting the skip: Training residual-free transformers. arXiv preprint arXh::2510.00345 2025.
[16] Yiping Ji, Hemanth Saratchandran, Cameron Gordon, Zeyu Zhang, and Simon Lucey. Sine activated low-rank matrices for parameter efficient learning. CoRR, 2024.
[17] Yiping Ji, Hemanth Saratchandran, Peyman Moghadam, and Simon Lucey. Always skip attention. arXiv preprint arXiv:2505.01990 2025.

---

[18] Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollär, and C Lawrencce Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part 13, pages 740-755. Springer, 2014.
[19] Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over- parameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 59:85-116 16, 2022.
[20] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012-10022, 2021.
[21] Lachlan MacDonald, Jack Valmadre, Hemanth Saratchandran, and Simon Lucey. On skip connections and normalisation layers in deep optimisation. Advances in Neural Information Processing Systems, 36: 147444 2023.
[22] Jan R Magnus and Heinz Neudecker. Matrix differential calculus with applications in statistics and econometrics. John Wiley & Sons, 2019.
[23] Abhisek Maiti, Sander Oude Elberink, and George Vosselman. Transfusion: Multi-modal fusion network for semantic segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6536-6546, 2023.
[24] Matterport. Mask r-cnn implementation, 2017.
[25] Jorge Nocedal and Stephen J Wright. Numerical optimization. Springer, 1999.
[26] Simon JD Prince. Understanding deep learning. MIT press, 2023.
[27] Tim Salzmann, Boris Ivanovic, Punarjay Chakravarty, and Marco Pavone. Trajectron++: Multi- agent generative trajectory forecasting with heterogeneous data for control. arXiv preprint arXiv:2001.03093, 2, 2020.
[28] Hemanth Saratchandran and Simon Lucey. Enhancing transformers through conditioned embedded tokens. arXiv preprint MxXi::2505./27888 2025.
[29] Hemanth Saratchandran, Damien Teney, and Simon Lucey. Leaner transformers: More heads, less depth. arXiv preprint wxxiv22505.2080 2025.
[30] Hemanth Saratchandran, Thomas X Wang, and Simon Lucey. Weight conditioning for smooth optimization of neural networks. In European Conference on Computer Vision, pages 310-325. Springer, 2025.
[31] Hemanth Saratchandran, Jianqiao Zheng, Yiping Ji, Wenbo Zhang, and Simon Lucey. Rethink- ing softmax: Self-attention with polynomial activations. arXiv preprint arXiv:24100 18..8883,, 2024.
[32] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers. arXiv preprint arXiv:201 1.04066, 2020.
[33] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pages 10347-10357. PMLR, 2021.
[34] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[35] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 2018.
12

---

[36] Ross Wightman. Pytorch pytorch- image- models 2019.
image models.
httpss//ggthuuu..o////gggggggi
[37] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Fei Tan, Glenn Fung, Vikas Singh, Xiaodong Yuan, Sungsoo Ahn Wang, Dimitris Papailiopoulos, and Katerina Fragkiadaki. Github repository, 2021.
[38] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 14138-14148, 2021.
[39] Li Yuan, Qibin Hou, Zihang Jiang, Jiashi Feng, and Shuicheng Yan. Volo: Vision outlooker for visual recognition. IEEE transactions on pattern analysis and machine intelligence, 45(5):6575- 6586, 2022.
[40] 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. International Conference on Machine Learning, pages 40770- 40803. PMLR, 2023.
[41] Q Zhen, W Sun, H Deng, D Li, Y Wei, B Lv, J Yan, L. Kong, and Y Zhong. cosformer: rethinking softmax in attention. In International Conference on Learning Representations. 2022.
[42] Liu Zhuang. Lin Wayne, Shi Ya, and Zhao Jun. A robustly optimized bert pre-training approach with post-training, In Proceedings of the 20th chinese national conference on computational linguistics, pages 1218-1227, 2021.
13

---

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

---

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

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We use publicly available code from works within the literature that clearly cite.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurlPS code and data submission guidelines (https://nips.cm
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ http:/.cccpubllc/uuuueeee Codeeubmii ssiinn crrionPPoiccy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We use the exact setup from works in the literature that we explicitly cite and whose GitHubs publicly available. We also clearly state the hardware used in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We ran five trials for each experiment and showed the mean result along with standard deviations. Error plots are then shown in the Section A.2
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute and hardware is explicitly stated for each experiment in Section A.2
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: //tner//nsusppee htcp://pblccttccccciieeeeel..eea
Answer: [Yes]
Justification: We read the NeurIPS Code of Ethics and can confirm that our paper coforms those ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: For this work there is societal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
17

---

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

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines: The paper does not involve crowdsourcing nor research with human subjects. The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
20

---

A Appendix
A.1 Theoretical Framework
In this section we will give the proofs of the lemma and theorems in Section 3.2 and Section 3.3
To begin with will need the following standard facts derivatives of matrices. Lemma A.1. Les A E R=xm, B E R*xl and C E mR Then
OACB ac
B" & A.
(15)
Proof. We start by using a well known vectorization identity 22
vec(ACB) = (B" & A)vec(C)
(16)
where vec denotes the vectorization operator which takes matrix and maps a vector by stacking its columns on top of each other, see [22]. We then differentiate Equation (16) to obtain
Ovec(ACB) B" 2 A.
(17)
@vec(C)
The result the lemma follows.
Lemma A.2. Let A E gaxm so that AT E [Rmx", Then
vec(A") Tm vec(A)
(18)
Ovec(A") T mn
(19)
Ōvec(A)
where Tmn is commutation matrix.
Proof. The first equation follows from the definition of the transpose of matrix 22 The second
equation then follows from the first.
We now give the proof of Theorem 3.2
Proof of Theorem 3.2 Let [21, ...,..., 2n denote vector in R". Then by definition
soft 2n])
(20)
E
From the above equation we can compute the partial derivative and find
i soflmax(z); softmax(=)) (6.g softmax(2))).
(21)
Ø2j
The term @softmax(z)) is precisely the component of the matrix msoftmax(s) Putting each of these dzj
az
ij terms into an n matrix we find
@softmax(z) Diag(z) TT
(22)
Oz
which proves the lemma.
We can use the above lemmas to give the proof of Theorem 3.3
Proof of Theorem 3.3 We start by establishing the derivative formula the term OA(X)
TW We will
use the notation used Section
3.1
Note that by definition 8A(X)
E RddNxdD Using Equation
aWa
A A(X) Lxsoftmax XXW
(23
21

---

where IN is the N x N identity matrix. This is done so that we can apply Theorem A.1 compute
We then
(A(X)
B xX))xW
(24)
@We
aWe
Osoftmax
XWWW
=WTX" @ IN)
using Theorem A.1
(25)
aWe
=W?Xx* @ (1)
using Theorem
3.2
oW
(26)
XWTXT @ (s (XW x
(27)
which proves the firs equality in Theorem 3.3
To compute (A(X) OW*
E RRdxxDD we proceed in a similar way.
(X()
X WT (
(28)
OWk
OWk
@softmax (XW W XT)
(WFX* @ IN)
using Theorem A.1
(29)
OWk
-WTXT @ (1))(5fttmmx((WWWWX))) OXW W
using T Theorem
3.2
OWK
(30)
(WTX* @ Ix)MssftmaxxxXWWWW @ XWQ)Tna
(31)
where the last equality follows from Theorems A.1 and
A.2
This establishes the second equality in
Theorem 3
To prove the identity for aA(X) 3Wv
E pddNNdd xdD we write
A(X) = softmax XX
(32)
where la is the d x d identity matrix. Then we simply apply Theorem A.
to obtain
AA(X) @(softmax (X XWvla) dWr
(33)
OWv
Ta @ softmax WT
(34)
which proves the final equality in Theorem 3.3
Proof of Theorem 3.4 The proof of Theorem 3.4 follows from using Theorem 3.3 and the definition of the Jacobian of A X) with respect to Wo. Wx and Wy given by
J(A(X)) = (((A())) A(A())) ((A())) aWQ
(355)
OWK
Ow
We recall that the condition number is defined x(J(A(X)))
Omaz ((((X))
Amis (AA))) where
m (J(A (X))) the maximum singular value of J(A(X)) and Tmin (J(A(X)) the minimum singular value which we know is non-zero because of the assumption that J(A(X)) has full rank.
Note that, using the notation in Section 3.1 we have that J(A(X)) E R RRdN xd as 8W
@A(X)
OA(X) 2A(X) E xdD For each of notation we will write Aq = oW
OA(X)
0A(X)
Aς = OW
DW* BWy
8A(X)
Ay = DWv
22

---

We will start by computing bound for the maximum singular value. We have for any vector & E R
we have
T
[(A(X)) 2(A(X))
OWK (2), ((A(X))
(x)=21
oWo (z),
OWv (z)
(36)
=
(37)
[|AQ(=) + [IAxxx))11 + [1
(38)
S (o (Aq)² Imax (Ax)² + max (Av)²) ||=12.
(39)
This implics that
[(A
max maxx (J(A)))
max
(40)
=70
||
S v mmax (A²)2 + max (2x) + max (Av)
(41)
S maxx (Aq) T (Ax) + o Imaax (Av).
(42)
The next step is to compute lower bound for the minimum singular value o min (X)) The approach is similar to the above, using the fact that amin (J((())(())) = min I= (X)))((2).
We
can then use the inequality
((((X)) 0(A(X))
(z), ((A(X))
(=(A
(z),
(43)
OW
OWK
Dwv
((((X))
((A(X))
((A()))
max
(44)
aWQ
aWk
dWv
Then minimizing the above over the constraint Z E IRDD such that
||=||
we obtain
Omin (J(A(X))) 2 max { {mmn (8(A(X))
(8(A(X))
(O(A(XI))
}
aWQ
T'iin
OWK
..ii
Owv
(45)
Combing the bounds max (X()(X)) and mmin (J(A(X))) we obtain
(O(A(X)) + (((A(X)) + (((A(X))
<(((A(X))) S (
(46)
OW*
The final step is to get bound on the condition numbers for each term on the right hand side in the above inequality. This is done using two facts: Firstly, given two matrices C' and such that the product CD is full rank then x(CD) < xD)) and the second that x(C @ D) <(C')(D), Using these two facts we can then use Theorem 3.3 to obtain the bound of Theorem 3.4 and the proof is finished.
We will now give the proof of Theorem 3.5
Proof of Theorem 3.5 analogous.
We will give the proof for Wq as the proof for Wx and Wy are exactly
Take the SVD of Wq to obtain We = USVT where U E DDD is the left singular vectors, E Dxd is the matrix of singular values having the singular values of We on its main diagonal, V e ddd is the right singular vectors. We then define new matrix Co by the formula
Co:= USV
(47)
where 3 E Dxd has Fmaa (Wo) on its main diagonal and zeros elsewhere. We then have that
We + Cq = USS US
(48)
23

---

We can then compute max Q + Cq) = 20 (WQ) by definition of Co and m (WQ +C)
miin (W2) + mmxx (Wq). This implies
max mmax(WQ CQ)
x(WQ + CQ)
(49)
min ( + Ci))
2a (W²)
(50)
min (WQ) + 0max (WQ)
(51)
which finishes the proof of the theorem.
key issue for implementing the approach in Theorem 3.5 into the attention layer is that it requires the computation of the SVD of each of the matrices W q. WK K and Wv. Computing the SVD each time for large matrix extremely memory intensive. This is exactly the premise of Theorem 3.8 whose proof we now give.
Proof of Theorem 3.8 to obtain
The proof starts by making use of the Weyl inequalities for singular values
max (A + x Ix) S mmax m + 0 max (A Is)-
(52)
min (A + X- 1x) 2 m (A 1x) mmax (A).
(53)
We then compute
K(A + x = mmax (A + Ax)x)
(54)
mmin (A 1 X Ix)
mmax (A) + Fmax (A 1x)
(55)
mmin (X Tx)) -mmin (A)
mmxx (A) +
(56)
A mmin (A)
maxx (A)
by our assumption
(57)
min (A)
x(A)
(58)
which completes the proof.
Other forms of Attention. The theory developed Section was presented for the self-attention mechanism defined in Section to allow for concrete analytical expressions. However, as evident from the derivations, the proposed corrections to the weight matrices Wq. WK. and Wy apply to any attention layer. Since the results depend only the derivatives of the attention function with respect to these matrices, analogous theorems can established for more general mechanisms such cross-attention. Moreover, in Section 4 empirically demonstrate that our theoretical results extend to a broad class of attention architectures.
A.2 Experimental Analysis
Details of implementation of spectral conditioned attention. For all of the experiments in Section we implemented spectral conditioned attention following the strategy outlined Section via Theorem 3.8 We describe this implementation in detail. For each attention layer in transformer fix the query, key and value weights as Wo. WK. Wv. We then form three matrices denoted by Co. Cx and Cy with each being defined by the equation
Co = Ale
(59)
Ck AIK
(60)
Cv = Alv
(61)
with À and where Co. Cx and Cv have the same shape as Wo. Wx and Wy and Tq. Ix and Iv has the same shape as Cq, CK and Cv with I's on their main diagonal and zeros elsewhere.

---

The matrices Co Cx and Cv, which we call the correction matrices, are fixed throughout training the transformer and never change. They are not updated and hence do not add any memory the backpropagation algorithm and the only memory comes from storing them. During the training of the transformer, we add each of these correction matrices to the weights Wq. Wx k and V during each forward pass forming. Thus after initializing the weights W W and W we add the correction terms Co, Cx and Cv as follows
1. - Wº + Co
(62)
2. K W Ck
(63)
3.W W + Cv
(64)
perform forward pass using W + Co. W + CK and WQ + Cv. then perform a backward pass and during the backward pass only the weights W W and Wpp are updated forming W W K
and W.. This process then repeats, add the fixed correction terms to W W and W3 forming
1 + W2 + Co 2. + W + CK
(65)
(66)
3.W. - W + Cv
(67)
then backward pass updates W W and W and this process continues. Note that Theorem 3.8 allows any À 2 2 when forming Through ablations we found the best value is À = 10.
A.2.1 Vision transformers on Imagenet-Ik
Validating the theory. In Section 4.1 we validated the theory given in Section 3 on a ViT-B and XCTT-M architecture. For each experiment ran five trials with five different random seeds and plotted the mean. The plots in Section 4.1 were shown in a log scale where the standard deviations were not visible. We have plotted each plot new scale that clearly shows the standard deviations. The plots for ViT-B can be seen in Figures/Sjooff Band the plots for XCÝT-M can be seen from Figures P to T2 Furthermore, Figure 6 validates the assumption Theorem 3.8 on the maximum singular
value for ViT-B and Figure 10 for the maximum singular for XCIT justifying the use of Theorem 3.8
Ablation on À from Theorem 3.8 In Section we used the implementation for the correction
terms given in Section 3 This required choosing a factor À > 2. For the experiments in Section 4 we chose through a grid search treating it as a hyperparameter. An ablation for on the ViT-B architecture trained on ImageNet- 1k is shown in Table What we found was that À 2 10 gives the best result.
Table 5: Ablation on for VYT-B
A
2
4
6
8
10
12
14
16
ViT-B Acc.
80.8
81.1
81.4
81.6
81.7
81.7
81.6
81.5
Relation to Normalization. Our conditioning methodology can be seen as heing related to normal- ization. By adding correction terms to the weights Wq. WK and Wy to improve their spectrum we are in effect producing a normalization for their spectrum. We therefore decided to test our spectral conditioning when we remove layer normalization in the vision transformers to understand whether spectral conditioning can replace layer normalization. Table 6 shows the results on the ViTs on ImageNet- 1k. We ran each experiments times and plot the mean and standard deviation in brackets. We found that if we remove layer normalization and keep spectral conditioning (with X = 10) the corrected
performance drops. We noticed that this was due to the fact that the weight values terms Wq + Co. Wx + Cx and Wy + Cv were much larger than without layer normalization which can lead to issues with backpropagation and suggesting that layer normalization was still crucial to maintain weights from getting too large.
25

---

Table 6: Performance comparison of spectral conditioning with and without layer normalization.
ViT-B
DeiT-B
Swin-B
XCIT-M
DaViT-B
Original (with only layer norm.)
80.7 (+:0.41)
81.6 (+0.30)
83.4 (+0.28)
82.6 (++0.39)
84.3 (++0226
Spec. cond. + layer norm.
81.7 (++0.38)
82.6 (+0.32)
84.1 (+0.25)
83.5 (+0.35)
84.9 (+0.21))
Spee. cond. layer norm.
79.8 (+0.31)
80.4 (+0.33)
80.5 (+0.23)
80.0 (+0.35)
80.3 (+0.21)
Comparing Theorem 3.5 and Theorem 3.8 In Section 3 we that Theorem 3.5 gave
methodology that could reduce the condition number of the query weights, key weights and value weights to less than 2. However, the reason we did not use this theorem the basis of our experiments is that it requires an SVD computation which would significantly increase training times when compared to the implementation friendly theorem Theorem 3.8 We carried out comparison on the vision transformer for both methods, As can be seen from Table 7 the vision transformer employing the methodology from Theorem 3.3 take much longer train.
Table 7: Comparison of accuracy and training time across transformer variants with spectral condi-
tioning of attention following Theorem 3.5 and Theorem 3.8
Model
Ace.
Training time (hrs:mins)
ViT-B (original)
80.7
29:29
VYT-B spec. cond. (Thm. 3.5)
82.0
41:38
VIT-B spec, cond. (Thm. 3.8)
81.7
29:33
DeiT-B (original)
81.6
26:16
DeiT-B B spec. cond. (Thm. 3.5)
82.9
37:54
DeiT-B spec. cond. (Thm. 3.8)
82.6
26:24
Swin-B (original)
83.4
53:12
Swin-B spec. cond. (Thm. 3.5)
84.3
68:32
Swin-B spec. cond. (Thm. 3.8)
84.1
53:26
XCTT-M (original)
82.6
91:03
XCTT-M spec. cond. (Thm. 3.5)
83.9
109:12
XCTT-M spec. cond. (Thm. 3.8)
83.5
91:18
DaViT-M (original)
84.3
75:09
DaViT-M spec. cond. (Thm. 3.5)
85.2
93:12
DaViT-M spec. cond. (Thm. 3.8)
84.9
75:28
Hardware and implementation. The image classification experiments in Section 4.1 of the paper were done on Nvidia A100 GPUs. The implementation of the ViTs were all done using the Timm code base [6. The architectures were all trained from scratch on the ImageNet- Ik dataset using the AdamW optimizer following the hyperparameters used in the original papers
20
We analyze the memory overhead introduced by using spectrally conditioned weight matrices Wq + Cq. WK + Ck. and Wv + Cv in the forward pass for the VÝT-B architecture. similar analysis gives similar profile for all the other ViT architectures we consider in Section 4.1 The correction matrices Co. Ck. and Cv are fixed diagonal matrices with constant entries (value 10) on the main diagonal, are non-trainable, and have gradients disabled (requires_grg __radd False).
Memory Cost.
Computing X X(Wq + Cq) expands as: CC) XW2 + XCo.
(68)
where XWQ W is a standard matrix multiplication and X XCq corresponds to scaling each column of X by 10. This operation is efficiently implemented as a column-wise scaling, requiring no explicit storage of Co beyond scalar or small constant vector.
Parameters: No trainable parameters are introduced; Cq. Ck, and Cy are fixed and non-trainable.
Gradient Storage: Since the correction matrices not participate in backpropagation, no additional gradient memory is used.
26

---

Activation Memory: The computation of XCo involves element-wist scaling and summa- tion, which reuses existing output tensors without requiring new memory allocations.
In practice, the additional memory overhead from spectral conditioning is negligible, as it involves no new trainable parameters, no gradient storage, and no extra activation tensors.
Table 8: Comparison of memory overhead for spectral conditioning relative to standard attention projections.
Memory Aspect
Original Attention
With Spectral Conditioning
Trainable Parameters
3 x Dd
3 Dd
Gradient Storage
3 x Dd
3 x Dd
Correction Matrices (Cq, Ck. Cv)
negligible (constant scalar) 3 x Nd (no extra tensors)
Activation Memory
3 x Nd
FLOPS Cost. We analyze the additional FLOPS overhead introduced by using spectrally condi- tioned weight matrices Wq + Co. WK + CC, and Wv + Cv in the forward pass of the attention mechanism in the forward pass for the ViT-B architecture. similar analysis gives a similar profile for all the other VIT architectures we consider in Section 4. Consider the self-attention formulation:
softmaxx ( + CQ) + Cx) X(Wy + Cv).
(69)
where X E RNXD is the input, and Wq, WK. Wv E Pxd are the learnable projection matrices. The correction matrices Co. Ck, and Cy are fixed diagonal matrices with constant entries 10 along their main diagonal.
Computing the term XWQ CQ) expands as:
X(W + CQ)= XWQ + XCO.
(70)
The first term, X W.Q corresponds to a standard matrix multiplication with computational complexity Dd)),, requiring 2.N Dd FLOPS. The second term, XCq, corresponds to scaling each column of X by 10, which requires only Nd FLOPS (one multiplication per entry).
Thus, the total FLOPS for computing X X Cq) is:
2N Dd + Nd.
(71)
identical cost arises for X(Wx + CR) and X(Wv + Cv). Therefore, the additional FLOOPS introduced by the correction terms across all three projections is:
3dd.
(72)
Compared to the original cost of 6N Dd FLOPS for the query, key, and value projections without spectral conditioning, the overhead introduced by spectral conditioning is negligible:
3Nd 1
(73)
6.N Dd
2D
which is insignificant for typical values of D used in practice (e.g., 512, 768, 1024 are common dimension for ViTs).
Table 9: FLOPS comparison for the forward pass of the attention projection layers with and without spectral conditioning. N denotes the sequence length, D the embedding dimension, and d the projection dimension per head.
Operation
FLOPS (Original)
FLOPS (With Spectral Conditioning)
Query projection (X
2N Dd
2N Dd + Nd Dd Nd
Key projection (X (XWR)
2N Dd
Value projection (XWv)
2N Dd
2N + Nd
Total Projection FLOPS
6N Dd
Dd + 3Nd
27

---

10.0
5.0
10'4
1.5
1C
150
252
400
50 350 200
302
400
Epacha
Esoche
Figure 5: Left: Average minimum singular value of the query, key, and value projection matrices (We. WK. Wv) for a ViT-B during training. We plot the mean over five trials and the standard deviation. Right: Average minimum singular value of the corrected query, key, and value projection matrices (Wq + Cq, WK + Ck. Wv + Cv) for a ViT-B during training. We plot the mean over five trials and the standard deviation.
30
mama angule
200
Wr
10
197
Figure 6: Left: Average maximum singular value of the query, key, and value projection matrices o. WK. Wy) for ViT-B during training. We plot the mean five trials and the standard deviation. Right: Average maximum singular value of the corrected query, key, and value projection matrices (Wq- + Cq. Wx Ck. Wv Cv) for VIT-B during training. We plot the mean five trials and the standard deviation.
12P
221
169
50
100
150
200
230
200
250
400
50 100
130
200
230
300
350
400
Figure 7: Left: Average condition number of the query, key, and value projection matrices (Wq. W WK- Wv) for a ViT-B during training. We plot the mean over five trials and the standard deviation. Right: Average condition number of the corrected query, key, and value projection matrices (Wa Co. Wx+ Wv + Cv) for a ViT-B during training. We plot the mean over five trials and the standard deviation.
A.2.2 Object detection and instance segmentation
Hardware and Implementation:
The experiments for Section 4.2 of the paper on object detection and instance segmentation were carried out on Nvidia A100 GPUs. The implementation followed B3. We used the code base given by the GitHub 24] following their exact training regime.
28

---

10
Condition number
10
1010
10
Artention
Spectral conditioned actention
107
Theoretical bound
50
100
150
200
250
300
350
400
Figure 8: Average condition number of the self-attention Jacobian of ViT-B over the course of training, before and after spectral conditioning, along with the theoretical bound from Equation (TO). ran five trials with five different random seeds with plots showing the mean and standard deviations.
7.
1019
6.5
C.
10-4
130
150
200
100 200 300
430
Figure 9: Left: Average minimum singular value of the query, key, and value projection matrices (Wo. WK. Wv) for a XCTT-M during training. We plot the mean over five trials and the standard deviation. Right: Average minimum singular value of the corrected query, key, and value projection matrices (WQ + Cq. WK + CK. Wy Cv) for a XCIT-M during training. We plot the mean over five trials and the standard deviation.
max wingatarreee
35
20
W
20
R,
152
200
1.0 150
200
250
100
400
Figure 10: Left: Average maximum singular value of the query, key, and value projection matrices (WQ. WK. Wv) for XCTT-M during training. We plot the mean over five trials and the standard deviation. Right: Average maximum singular value of the corrected query, key, and value projection matrices (WQ + Cq. WK + Ck. Wy Cv) for XCiT-M during training. We plot the mean over five trials and the standard deviation.
For this experiment we used pre-trained XCIT-S and since, as explained in Section A.2.1 the correction terms are not trained and fixed throughout training there was no real memory overhead.
Memory and Overheads:
29

---

120
2775
2.00
1.73
1.50
1.00
100 150
205
250
300
350
400
50 120
150
203
250
300
350
400
Epocis
Esechs
Figure 11: Left: Average condition number of the query, key, and value projection matrices (WQ. WK. Wy) for a XCiT-M during training. We plot the mean over five trials and the standard deviation. Right: Average condition number of the corrected query, key, and value projection matrices (Wo + Co. WK + Ck. Wv + Cv) for a XCTT-M during training. We plot the mean over five trials and the standard deviation.
1011
1010
Condition number
10a
10a
10
10a
Attention Spectral conditioned attention
10a
Theoretical bound
100
150
200 Epochs
250
300
350
400
Figure 12: Average condition number of the self-attention Jacobian of a XCTT-M over the course of training, before and after spectral conditioning. along with the theoretical bound from Equation (TO). We ran five trials with five different random seeds with plots showing the mean and standard deviations.
Further results: We also ran the object detection and instance segmentation results using a Swin- base vision transformer as the backbone. The result are shown Table I0
Table 10: Comparison of Swin-hase and spectrally conditioned (spec. cond.) Swin-base on COCO benchmarks.
Model
APb
AP
AP
AP
AP
AP
Swin-base
45.9
67.1
50.1
41.1
63.7
43.8
Spec. cond. Swin-base
46.8
68.1
50.7
41.8
A.3 Nyströmformer on LRA benchmark
Validating the theory. Section 4.3 we validated the theory given in Section Nyströmformer on the text classification task in the LRA benchmark. For each experiment we ran five trials with five different random seeds and plotted the mean. The plots Section 4.3 were shown in log scale where the standard deviations were not visible. We have plotted each plot in new scale that clearly shows the standard deviations. The plots are shown from Figures [3 to T6, We also ran similar empirical analysis for the listops task in the LRA benchmark. The results of which can be seen in
30

---

Figures 17 to 21 Furthermore, Figure 14 and Figure 19 validates the assumption in Theorem 3.8 on the maximum singular value justifying the use of Theorem 3.8
Ablation on from Theorem 3.8 search treating
For the experiments in Section 4.3 we chose À through grid
hyperparameter. An ablation for À on the Nyströmformer trained on the text classification task on the LRA benchmark is shown in Table 11 Once again we found that À 2 10 gives the best result.
Table 11: Ablation À for Nyströmffrmmer on text classification task.
10
16
Acc. 63.7 64.1 64.3
64.6
64.8
64.8
64.7
64.7
Relation to Normalization. We compared spectral conditioning a Nyströmformer with and without layer normalization, As can be seen Table 12 we see that layer normalization is important for spectral conditioning. Once again when we removed layer normalization and just applied spectral conditioning we found that the weights were large resulting in poorer performance.
Table 12:
Comparing spectral conditioning with and without layer normalization for the Nyström- former on text classification.
Acc. (%)
Original (with only layer norm.)
63.8 (#0.24)
Spec. cond. + layer norm.
64.8 (=0.200)
Spec. cond. - llaee norm
62.3 (d: 0,23)
Hardware. All the experiments for the Nyströmformer on LRA benchmark results in Section 4.3 were carried out Nvidia A100 GPUs following the implementation and hyperparameter settings given 37
Memory and Overheads: The analysis carried out FLOPS and memory in Section A.2. apply
in this case as once again the correction matrices Co. Cg and Cy are fixed throughout training and only added for forward passes.
5.8
C
15000
20000
Esodhs
Esechs
Figure 13: Left: Average minimum singular value of the query, key, and value projection matrices (Wo, WK. Wv) for Nyströmformer on the text classification task during training. We plot the mean over five trials and the standard deviation. Right: Average minimum singular value of the corrected query, key, and value projection matrices (Wq + Cq. WK Cx. Wy Cv) for a Nyströmformer on the text classification during training. We plot the mean over five trials and the standard deviation.
A.4 Language Modeling with Crammed BERT
Hardware. The language modeling experiment in Section 4.4 were all carried out Nvidia A6000 GPU. The Crammed-Ber was implemented following the original paper [I]] and the original GitHub [00. The training regime follows [10].
31

---

max stgulea unlln
40
Au
102
200
150 200
300 400
Epochs
Epocts
Figure 14: Left: Average maximum singular value of the query, key, and value projection matrices (We. WK. Wv) for Nyströmformer on the text classification task during training. We plot the mean over five trials and the standard deviation. Right: Average maximum singular value of the corrected query, key, and value projection matrices (Wq + Cq. W + Ck. Wv + Cv) for a Nyströmformer on the text classification during training. We plot the mean over five trials and the standard deviation.
24
22
2.0
1.8
1.4
Al
14
180
12
1.0
2500
$900
7500
10000
12500
15003
17503
20000
2500
5000
7500
10000
12590
15009
17500
20000
Esochs
Figure 15: Left: Average condition number of the query, key, and value projection matrices (Wq- K. Wv) for a Nyströmformer on the text classification task during training. We plot mean over five trials and the standard deviation. Right: Average condition number of the corrected query, key, and value projection matrices (Wa + Co. WK + CK. Wy + Cv) for Nyströmformer on the text classification task during training. We plot the mean over five trials and the standard deviation.
10a
Condition number
107
10a
10a
10c
conditioned attention
Theoretical bound
10a
2500
5000
7500
10000
12500
15000
17500
20000
Epochs
Figure 16: Average condition number of the self-attention Jacobian Nyströmformer on the text classification task over the course of training, before and after spectral conditioning, along with the theoretical bound from Equation (10). We ran five trials with five different random seeds with plots showing the mean and standard deviations.
32

---

397
133
10-4
23'
Epode
Eoocit
Figure 17: Analysis for Nyströmformer on listops task. Left: Average minimum singular value of the query, key, and value projection matrices (Wq, WK. Wy) and their spectrally conditioned counterparts (WQ + Co. WK Cx. Wy+Cy) throughout training. Middle: Condition numbers of Wo. WK. and Wy: v. and their spectrally conditioned forms during training. Right: Average condition number of the attention Jacobian over the course of training, before and after spectral conditioning, along with the theoretical bound from Equation 10
109
6.5
22-4
6.0
5.5
5.0
C
10
1000
5002
Figure 18: Left: Average minimum singular value of the query, key, and value projection matrices (Wq. WK. Wv) for Nyströmformer on the listops task during training. We plot the mean over five trials and the standard deviation. Right: Average minimum singular value of the corrected query. key, and value projection matrices (Wq + Co. WK + CR. Wv Cv) for a Nyströmformer on the listops task during training. We plot the mean over five trials and the standard deviation.
50
max
20
50 303
155
200
250
330
350
400
50
100
150
200
250
300
350
400
Epocts
Epects
Figure 19: Left: Average maximum singular value of the query, key, and value projection matrices (Wq. WK. Wv) for a Nyströmformer on the listops task during training. We plot the mean over five trials and the standard deviation. Right: Average maximum singular value of the corrected query, key, and value projection matrices (Wo + Co. WK + CK. Wy Cy) for a Nyströmformer on the listops task during training. We plot the mean five trials and the standard deviation.
Memory and Overheads: The analysis carried out on FLOPS and memory in Section A.2. apply in this case as once again the correction matrices Co. CK and Cy are fixed throughout training and only added for forward passes.
33

---

2.2
2.0
21.8
2.6
3.4
32
C
W.
2303
4000
3050
5000
Epechs
Esechs
Average condition number of the query, key, and value projection matrices (WQ. Q- Wk. Wv) for a Nyströmformer on the listops task during training. We plot the mean over five trials and the standard deviation. Right: Average condition number of the corrected query, key, and value projection matrices (Wq + Co. Wx + Cx. Wv + Cv) for Nyströmformer on the listops task during training. We plot the mean over five trials and the standard deviation.
Figure 20: Left:
107
10c
Condition
10
10a
Attention
Spectral conditioned attention
10a
Theoretical bound
0
1000 2000
3000
4000
5000
Epochs
Figure 21: Average condition number of the self-attention Jacobian of Nyströmformer on the listops task over the course of training, before and after spectral conditioning, along with the theoretical bound from Equation (10). We five trials with five different random seeds with plots showing the mean and standard deviations.
34


Reviews:
[{"review_id": "Gs7ZKd9s9l", "reviewer": "Reviewer_kbKM", "summary": "This paper introduces a **theoretically grounded** and **computationally efficient** method to enhance Transformer models by improving the *conditioning* of their self-attention layers. The authors show that the condition number of the Jacobian of the self-attention block—an important factor affecting optimization stability—depends on the *spectral properties of the query, key, and value projection matrices* (WQ, WK, WV). To mitigate ill-conditioning, they propose spectral conditioning, which adds **fixed correction terms** to these matrices, significantly reducing their condition numbers and, in turn, improving the conditioning of the self-attention Jacobian.\n\nThe method is simple to implement and compatible with a wide range of attention variants and Transformer architectures. It does not introduce additional trainable parameters and incurs minimal computational overhead. The authors validate their approach across diverse tasks including **image classification** (ViT, Swin, XCiT, etc.), **object detection** and **segmentation** (COCO dataset), **long-range sequence modeling** (LRA benchmark with Nyströmformer), and **language modeling** (Crammed BERT on GLUE). In all cases, spectral conditioning consistently improves performance and leads to more stable training dynamics.\n\nWhile the proposed method optimizes an *upper bound* on the Jacobian condition number rather than the *true value* itself, the empirical gains support the theoretical motivation. The technique’s **ease of integration** and **general applicability** position it as a valuable tool for improving Transformer robustness without modifying their architecture or training procedure.", "questions": "**Questions and Suggestions for the Authors**\n\n1. Additional Classification Datasets.\n2. More Detection/Segmentation Models.\n3. Vision-Language Models.\n4. Discussion on Applicability to Cross-Attention.\n\nMy current rating is **borderline accept**, and *I lean towards accepting the paper* due to its solid theoretical contributions and promising empirical results. Addressing the questions above—particularly expanding the experimental validation—could strengthen my recommendation.", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n- This paper is **well-written** and **accessible**, offering a clear presentation of both theoretical and empirical contributions. The mathematical analysis is rigorous and well-motivated, with theorems and proofs presented in a concise and structured manner. The accompanying figures, especially those analyzing singular values and condition numbers, are clean and effectively support the theoretical claims. Overall, the exposition is strong and reinforces the soundness of the proposed method.\n\n- Despite the simplicity of the idea, the core contribution is **conceptually elegant** and **broadly applicable**. The method—spectral conditioning of self-attention—targets a foundational property of Transformers (the conditioning of the Jacobian), and can be applied as a drop-in modification to virtually any self-attention mechanism. This generality strengthens its practical impact. In terms of novelty, the work stands out by focusing specifically on self-attention layers, where previous conditioning methods (e.g., reference [21]) have typically targeted fully connected layers. Thus, I would characterize the level of novelty as medium to high.\n\n- The experimental section covers both vision and language domains, evaluating the method across diverse tasks. In vision, the authors explore not only image classification but also object detection and instance segmentation. The inclusion of results on long-range sequence modeling (via the LRA benchmark) and language modeling (Crammed BERT on GLUE) further supports the generality of the approach.\n\n**Weaknesses and Suggestions**\n\n- While the experiments are diverse in terms of task types, the breadth of datasets and model variants could be improved to strengthen the empirical evidence. For image classification, all experiments are limited to ImageNet-1k. Given current standards, it would be beneficial to include evaluations on additional datasets, such as Places365, iNaturalist, or fine-grained benchmarks like CUB-200 or CARS-196. These additions would test the robustness and generalization of spectral conditioning beyond ImageNet and increase the confidence in its effectiveness across visual domains.\n\n- Similarly, for detection and segmentation tasks, results are only provided for XCiT within the Mask R-CNN framework. To better demonstrate the generality of the method across vision backbones, I would suggest incorporating one or two additional models in this setup.\n\n- A potential omission is the lack of evaluation on vision-language transformer models, such as CLIP or SigLIP. These models also rely on self-attention, and a simple adaptation of spectral conditioning could be tested on tasks like zero-shot image recognition. Given the popularity and widespread use of such models, this would be a valuable addition and highlight the method’s potential in multimodal settings.\n\n- Lastly, while the focus is on self-attention, it would be helpful for the authors to briefly discuss whether the proposed conditioning method could extend to cross-attention mechanisms as well. Even a short theoretical or empirical remark on this could broaden the relevance of the work and preempt reader questions.", "was_revised": true, "final_justification": "Authors addressed my comments. They provided extra experiments with consistent results. In general they had a good rebuttal. Thus, I raise my score +1.", "rebuttal": "We sincerely thank the reviewer for taking the time to review our paper and for providing thoughtful and constructive feedback. Below, we address each of the points and questions raised in the review.\n\n**1. While the experiments are diverse in terms of task types, the breadth of datasets and model variants could be improved to strengthen the empirical evidence. For image classification, all experiments are limited to ImageNet-1k. Given current standards, it would be beneficial to include evaluations on additional datasets....**\n\nWe have run spectral conditioning on a ViT-B architecture and a fine-grained transformer architecture TransFG, as in He et al., on the CUB-200, CARS-196 and the iNat2017 datasets. The results shown below clearly show that spectral conditioning helps in fine grained image recognition.\n\n(i) Results on CUB-200\n\n| **Model**               | **Test Acc. on CUB-200** |\n|-------------------------|---------------------------|\n| ViT-B                   | 90.4                      |\n| **Spec. Cond. ViT-B**   | **91.5**                  |\n| TransFG                 | 91.7                      |\n| **Spec. Cond. TransFG** | **92.9**                  |\n---------------------------------------------------------------\n\n(ii) Results on CARS-196\n\n| **Model**               | **Test Acc. on CARS-196** |\n|-------------------------|----------------------------|\n| ViT-B                   | 93.6                       |\n| **Spec. Cond. ViT-B**   | **94.8**                   |\n| TransFG                 | 94.8                       |\n| **Spec. Cond. TransFG** | **95.8**                   |\n-----------------------------------------------------------------\n\n(iii) Results on iNat2017\n\n| **Model**               | **Test Acc. on iNat2017** |\n|-------------------------|----------------------------|\n| ViT-B                   | 68.9                       |\n| **Spec. Cond. ViT-B**   | **71.1**                   |\n| TransFG                 | 71.8                       |\n| **Spec. Cond. TransFG** | **72.7**                   |\n--------------------------------------------------------------------------\n\n2. **I would suggest incorporating more Detection/Segmentation Models.**\n\nWe ran a Swin-base model for both object detection and instance segmentation. The table below shows that spectral conditioning yields better performance. We are happy to include this in the appendix as an extra experiment.\n\n| Model                   | APᵇ  | APᵇ₅₀ | APᵇ₇₅ | APᵐ  | APᵐ₅₀ | APᵐ₇₅ |\n|-------------------------|------|--------|--------|------|--------|--------|\n| Swin-base               | 45.9 | 67.1   | 50.1   | 41.1 | 63.7   | 43.8   |\n| Spec. Cond. Swin-base   | **46.8** | **68.1** | **50.7** | **41.8** | **64.2** | **44.6** |\n\nWe also ran a Swin-small architecture on both object detection and instance segment. The table below shows the results one again showing spectral conditioning yields better performance.\n\n| Model                   | APᵇ  | APᵇ₅₀ | APᵇ₇₅ | APᵐ  | APᵐ₅₀ | APᵐ₇₅ |\n|-------------------------|------|--------|--------|------|--------|--------|\n| Swin-small              | 44.8 | 66.3   | 49.0   | 40.3 | 63.0   | 42.9   |\n| Spec. Cond. Swin-small  | **45.7** | **66.9** | **49.8** | **40.7** | **63.4** | **43.5** |\n------------------------------------------------------------------------------------------------------------------\n\n3. **A potential omission is the lack of evaluation on vision-language transformer models, such as CLIP or SigLIP....**\n\nWe thank the reviewer for this suggestion. While vision-language transformer models are very interesting we believe our experimental analysis is already very thorough. We performed experiments on image classification, object detection and instance segmentation, long range sequence modelling (LRA benchmark) and language modelling. Plus we have above carried out 3 different fine grained image classification tasks as well as doing another model on both object detection and instance segmentation. In each case we have shown spectral conditioned attention yields better results clearly showing our methodology is useful to the community.\n\n4. **Lastly, while the focus is on self-attention, it would be helpful for the authors to briefly discuss whether the proposed conditioning method could extend to cross-attention mechanisms as well...**\n\nWe thank the reviewer for their question. Yes our conditioning method does extend to cross attention. This is because cross attention uses the same kernel as regular self-attention. The difference between the two is that cross attention allows an input stream of tokens to interact with a different stream of tokens while self-attention only allows tokens to interact with each other within each stream of input. However, both forms of attention use the same kernel defined by eq. (3) in the paper. Theorem 3.3 and 3.4 of the paper is only dependent on the kernel used, as we are taking derivatives with respect to $W_Q$, $W_K$ and $W_V$ parameters and not the input tokens $X$, and thus also holds for cross attention. We will include a paragraph explaining this in the appendix. \n\nWe would also like to point out to the reviewer that we empirically demonstrated spectral conditioning works on a variety of different attention forms such as cross covariance attention used in XCiT, Windowed attention in Swinformer, and Dual attention in DaViT. These results were all shown in Table 1 in the paper. In each case spectral conditioning the associated attention form led to improved performance."}, {"review_id": "PhSbKRS6tF", "reviewer": "Reviewer_E8ez", "summary": "- The paper presents a theoretical analysis of the conditioning of the Jacobian of a self-attention layer of transformers.\n- Citing previous work, the authors argue that poorly conditioned Jacobians can hinder performance and improving Jacobian conditioning can lead to better optimization and generalization performance.\n- A theoretical analysis is performed, showing how the conditioning of the Jacobian of a self-attention layer is influenced by the conditioning of the query (Q), key (K), and value (V) matrices.\n- Based on this analysis, the authors propose spectral conditioned self-attention, which adds correction terms to the Q, K, V matrices during self-attention.\n- Since computing the correction terms using SVD at each iteration is prohibitive, the paper proposes an implementation-friendly form of the correction terms as $\\lambda \\mathbf{I}$, for some scalar hyperparameter $\\lambda \\geq 2$ (in the experiments, $\\lambda=10$ is chosen).\n- The authors empirically show improved conditioning and generalization on ImageNet classification, COCO object and instance segmentation, LRA text classification, and the GLUE benchmark.", "questions": "Is it possible to generalize the paper's analysis and method beyond self-attention (e.g., cross-attention)?", "limitations": "Yes, except the limitation to the self-attention case of attention.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strengths\n\n- Quality\n    - The proposed method is simple and technically sound.\n    - Experiments are performed across multiple domains and tasks.\n    - Assumption of full rank in the theorems is empirically validated.\n- Clarity\n    - The paper is clear and easy to follow.\n- Significance\n    - Despite being simple, the method leads to generalization improvements across diverse settings.\n- Originality\n    - The proposed spectral conditioned self-attention layer is novel.\n\n### Weaknesses\n\n- Quality\n    - The method only aims to minimize the upper bound of the condition number, and the practical implementation further loosens the upper bound being minimized.\n    - Theorem 3.5 is a trivial result that is unnecessary to derive the final practical method (Theorem 3.8 is all you need).\n    - Self-attention is a special case of attention, and it is not clear why the analysis does not extend to cases like cross-attention.\n- Significance\n    - The paper evaluates relatively small models and the generalization improvements are relatively small. There is no clear signal that the generalization improvements will persist in larger-scale settings.\n- Originality\n    - The idea behind the final practical method (Theorem 3.8), adding values to the diagonal to improve condition number, is not novel and is a fairly standard approach in linear algebra.", "was_revised": true, "final_justification": "In the response to the authors' rebuttal, I pointed out counterexamples to Theorem 3.8, which, if left unresolved, would have led to a clear recommendation of rejection.\n- The authors were able to fix the theorem and its proof without invalidating the paper. Based on the updated proof provided by the authors, I recommended simplifying and strengthening the theorem statement to not require constants like 0.4 and 2, which the authors agreed to incorporate into the paper.\n\nOne of my concerns was that the paper aims to minimize the upper bound of the condition number, but the bound used in the final method is so loose that one of their own intermediate upper bound theorems (3.5) is made unnecessary by the practical upper bound theorem (3.8), which they argue leads to a faster algorithm. In addition to a larger perceived gap between theory and practice, it also leads to a confusing presentation where a notable part of the paper feels disconnected from the rest.\n- The authors provided an empirical link to reduce this perceived gap: baseline is fast and performs the worst, the method based on theorem 3.5 is the slowest but performs the best, and the method based on theorem 3.8 is comparably fast as the baseline with performance between that of the other two settings.\n\nAnother concern, echoed by other reviewers, was that self-attention is a special case of attention, and it is not clear why the analysis does not extend to cases like cross-attention.\n- The authors already have empirical results where the method is applied to more general forms of attention. During the discussion, they stated that the theoretical results are also general and that they can generalize the paper's narrative accordingly.\n\nFinally, my original review noted that the models and generalization improvements are relatively small.\n- The authors were limited by resources and pointed out that this is stated in their limitations. This is entirely reasonable.\n- Regarding the generalization improvements being relatively small, the performance improvements using a method based on theorem 3.5 (discussed above) are larger, providing better support for the more general message for spectral conditioning of attention.\n\nOverall, I am satisfied with the discussions, and I increase my score to 5 (Accept).", "rebuttal": "We sincerely thank the reviewer for taking the time to review our paper and for providing thoughtful and constructive feedback. Below, we address each of the points and questions raised in the review.\n\n1. **The method only aims to minimize the upper bound of the condition number, and the practical implementation further loosens the upper bound being minimized.**\n\nWe thank the reviewer for their comment. However, we would like to note that this point is already acknowledged in the limitations section (Section 5) of the paper. We kindly ask the reviewer to follow the NeurIPS reviewer guidelines, see https://neurips.cc/Conferences/2025/ReviewerGuidelines, where it clearly states that authors should not be punished for being up front about the limitations of their work. \nWhile our method minimizes an upper bound on the condition number, our experimental results demonstrate that doing so is nonetheless effective, leading to improved performance across a range of tasks.\n\nIn particular, Figures  2, 3, and 4 (rightmost subplots) show that spectral conditioning substantially reduces the condition number of the self-attention Jacobian, by 2 to 3 orders of magnitude, and that this reduction is maintained consistently throughout training. These findings suggest that, even though our approach targets an upper bound, the resulting improvements in conditioning have meaningful and beneficial consequences for performance.\n\n2. **Theorem 3.5 is a trivial result that is unnecessary to derive the final practical method (Theorem 3.8 is all you need).**\n\nThank you for the suggestion. We introduced Theorem 3.5 to show that a tighter bound on the condition number can be obtained when one explicitly applies the singular‑value decomposition (SVD). Because computing an SVD at every iteration is prohibitively expensive, Theorem 3.8 provides a practical, SVD‑free alternative. We believe highlighting this trade‑off is helpful for readers who may not be familiar with such techniques, especially given NeurIPS’ broad, general‑machine‑learning audience. That said, we are happy to move Theorem 3.5 to the appendix or restate it as a proposition if the reviewer feels that would be better.\n\n3. **Self-attention is a special case of attention, and it is not clear why the analysis does not extend to cases like cross-attention.**\n\nOur conditioning method does extend to cross attention. This is because cross attention uses the same kernel as regular self-attention. The difference between the two is that cross attention allows an input stream of tokens to interact with a different stream of tokens while self-attention only allows tokens to interact with each other within each stream of input. However, both forms of attention use the same kernel defined by eq. (3) in the paper. Theorem 3.3 and 3.4 of the paper is only dependent on the kernel used, as we are taking derivatives with respect to $W_Q$, $W_K$ and $W_V$ and not the input tokens $X$, and thus also holds for cross attention. We will include a paragraph explaining this in the appendix. \n\nWe would also like to point out to the reviewer that we empirically demonstrated spectral conditioning works on a variety of different attention forms such as cross covariance attention used in XCiT, Windowed attention in Swinformer, and Dual attention in DaViT. These results were all shown in Table 1 in the paper. In each case spectral conditioning the associated attention form led to improved performance.\n\n4. **The paper evaluates relatively small models and the generalization improvements are relatively small. There is no clear signal that the generalization improvements will persist in larger-scale settings.**\n\nThe models we trained were at most roughly 100 million parameters primarily because our resources only allowed us to experiment with such scale models. We would also like to point out that we clearly stated this as a limitation in our limitations section in section 5 of the paper. We would kindly ask the reviewer to follow the NeurIPS 2025 reviewer guidelines: https://neurips.cc/Conferences/2025/ReviewerGuidelines where it says that authors should not be punished for being up front about the limitations of their paper, see point 4 under the Main Task heading and point 8 under the Reviewer Form heading.\n\nWe believe the insights derived at the roughly 100M‑parameter scale are still valuable to the community: they reveal consistent trends and provide a solid stepping‑stone for future work on larger models.\n\n5. **The idea behind the final practical method (Theorem 3.8), adding values to the diagonal to improve condition number, is not novel and is a fairly standard approach in linear algebra.**\n\nWhile the technique itself is rooted in linear‑algebraic reasoning, our literature search uncovered no prior work, either in machine learning or in the linear‑algebra community, that applies it to attention mechanisms within Transformer architectures. If we have overlooked a relevant citation, we would greatly appreciate the reviewer’s guidance and will gladly include any suggested references to ensure the paper is both comprehensive and fair."}, {"review_id": "fukfnC1tSm", "reviewer": "Reviewer_njoQ", "summary": "This paper focuses on analyzing the conditioning of the self-attention mechanism in Transformers to enable more stable training. An upper bound on the condition number of the Jacobian of the self-attention function is derived, providing a principled way to control Jacobian conditioning during training. Specifically, the paper introduces a technique called spectral conditioning, which can be applied not only to the classical attention mechanism but also to more advanced variants such as shifted-window attention (e.g., Swin Transformer) and low-rank approximations (e.g., Nyströmformer). Empirical results demonstrate that spectral conditioning improves the conditioning of attention modules during training. Experiments across diverse learning tasks confirm the enhanced performance brought by this technique.", "questions": "While I do appreciate the work, there are some related concerns:\n1. Low-rank approximation is closely related to the conditioning of the attention mechanism. Both Theorem 3.3 and Theorem 3.4 involve the term $\\text{softmax}(XW_QW_K^\\top X^\\top)$, which is a core target of many low-rank Transformer variants, such as Nyströmformer, Linformer [1], Performer [2]. First, to what extent do the theoretical bounds provided in Theorems 3.3 and 3.4 generalize to these low-rank attention variants? From my perspective, these theorem may be applicable to these methods since some spectral components can be controlled under the low-rank condition. Second, is it sufficient to control only the weight matrices $(W_Q, W_K, W_V)$ in these low-rank attentions in order to reduce their Jacobian's condition number? A broader discussion and numerical comparison across different low-rank self-attention mechanisms is necessary, as it could reveal how spectral conditioning complements or overlaps with structural approximations in improving training stability and performance.\n2. To better evaluate the impact of Jacobian conditioning, I feel it is necessary to add experiments that directly reflect numerical stability during training, such as a tracking on gradient norm, losses, etc.\n3. There are still some unclear notations and typos in the manuscript. Eq (4), $J$ should be removed from the partial derivation; Table 1, it is not immediately clear whether “Spec. cons.” refers to “Spec. cons. + the corresponding attention mechanism” and also the subsequent experiments. Caption of Figure 14, ViT-B should be corrected to Nyströmformer. Please double check the typos and expressions.\n\nI would like to consider raising the score if these concerns are resolved.\n\nReferences\n\n[1] Wang et al., Linformer: Self-Attention with Linear Complexity, arXiv 2020\n\n[2] Choromanski et al., Rethinking Attention with Performers. ICLR 2021", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength: Conditioning analysis in Transformer models remains relatively under-explored compared to other network architectures. Hence, this paper could inspire conditioning-aware design that improves the numerical stability of attention-based models.\n\nWeakness: A discussion and comparison with low-rank attention mechanisms may be necessary.", "was_revised": true, "final_justification": "The authors have addressed my concerns with extra experimental supports. Hence, I would like to increase the rating.", "rebuttal": "We sincerely thank the reviewer for taking the time to review our paper and for providing thoughtful and constructive feedback. Below, we address each of the points and questions raised in the review.\n\n1. **Low-rank approximation is closely related to the conditioning of the attention mechanism. Both Theorem 3.3 and Theorem 3.4 involve the term $softmax(XW_QW_K^TX^T)$, which is a core target of many low-rank Transformer variants, such as Nyströmformer, Linformer [1], Performer [2]. First, to what extent do the theoretical bounds provided in Theorems 3.3 and 3.4 generalize to these low-rank attention variants? From my perspective, these theorem may be applicable to these methods since some spectral components can be controlled under the low-rank condition.**\n\nThis is a very good question and we thank the reviewer for asking it. The reviewer is exactly right. Since both theorem 3.3 and 3.4 involve the term $softmax(XW_QW_K^TX^T)$ and this is the critical term used in any low-rank variant of self-attention the theoretical bounds in theorems 3.3 and 3.4 generalize to a low rank approximation of such attention terms. The key point is that in deriving theorem 3.3 and 3.4 we make use of differentiating the term $softmax(XW_QW_K^TX^T)$ with respect to $W_Q$ and $W_K$ which holds for low-rank attention variants.\nIn fact this is exactly why spectral conditioning worked for the Nyströmformer on LRA benchmark shown in section 4.3 of the paper. To further show the reviewer that spectral conditioning helps in the low-rank setting we applied it to the Performer architecture, as referenced by the reviewer, on the LRA benchmark. As can be seen from the table below spectral conditioning yields better performance.\n\n|                  | Listops | Text   | Retrieval | Image  | Pathfinder |\n|------------------|---------|--------|-----------|--------|------------|\n| Performer        | 18.04   | 65.40  | 53.83     | 42.77  | 77.05      |\n| **Spec. Cond. Performer** | **19.08** | **66.61** | **54.65** | **43.79** | **78.21**    |\n------------------------------------------------------------------------------------------------------------------\n\n2. **Second, is it sufficient to control only the weight matrices $(W_Q, W_K, W_V)$ in these low-rank attentions in order to reduce their Jacobian's condition number? A broader discussion and numerical comparison across different low-rank self-attention mechanisms is necessary, as it could reveal how spectral conditioning complements or overlaps with structural approximations in improving training stability and performance.**\n\nTheorems 3.3 and 3.4 establish that the condition number of the self‑attention Jacobian is controlled by the projection matrices \n$(W_Q, W_K, W_V)$. Theorem 3.3 gives the explicit relationship. Low‑rank variants such as Nyströmformer, Performer, and Linformer still compute attention in the form $softmax(XW_QW_K^TX^T)$ or compute an approximation of it that still uses the terms $(W_Q, W_K, W_V)$. Consequently, their Jacobians inherit the same dependence on $(W_Q, W_K, W_V)$. By conditioning these matrices, we can therefore reduce the Jacobian’s condition number in both the standard and low‑rank settings. We will add a short derivation outlining this point in the appendix.\n\n3. **To better evaluate the impact of Jacobian conditioning, I feel it is necessary to add experiments that directly reflect numerical stability during training, such as a tracking on gradient norm, losses, etc**\n\nThe tracking of gradient norm only captures the magnitude of the gradient, and since the Jacobian is the transpose of the gradient, this is equivalent to tracking the norm of the Jacobian. However, we would like to emphasize that conditioning is fundamentally different from merely normalizing the size of the Jacobian norm. In general, two matrices can have the same condition number but very different norms. As a concrete example, consider the matrices\n\n$$\nA = \\\\begin{bmatrix}\n1 & 0 \\\\\\\\\n0 & 1\n\\\\end{bmatrix}\n\\\\quad \\\\text{and} \\\\quad\nB = \\\\begin{bmatrix}\n10.1 & 0 \\\\\\\\\n0 & 10\n\\\\end{bmatrix}.\n$$\n\nThe matrix $A$ has condition number $1$ and the matrix $B$ has condition number $1.01$, thus we see that\nboth matrices have condition number close to $1$. However, the Frobenius norm of $B$ is significantly larger than that of $A$. There are also examples where the norms of the matrices can be similar but their condition numbers are very different. For example consider the matrix $C$ and $D$.\n\n$$\nC = \\\\begin{bmatrix}\n1 & 0 \\\\\\\\\n0 & 1\n\\\\end{bmatrix}\n\\\\quad \\\\text{and} \\\\quad\nD = \\\\begin{bmatrix}\n\\\\sqrt{2} & 0 \\\\\\\\\n0 & 0.1\n\\\\end{bmatrix}.\n$$\n\n$C$ has a condition number of $1$ and $D$ has condition number of approximately $14.14$, thus\n$D$ has a much larger \ncondition number than $C$. However, their Frobenius norms are similar.\n\nThis demonstrates that computing the norm of the Jacobian during training does not necessarily provide any information about its condition number.\n\nIn our paper, we apply spectral conditioning by adding predefined matrices $C_Q$, $C_K$, and $C_V$ to the weight matrices $W_Q$, $W_K$, and $W_V$, respectively. According to Theorem 3.3, this increases the Frobenius norm of the self-attention Jacobian. However, all transformer architectures used in the paper employ Layer Normalization, which regulates the scale of the weight norms and ensures that Jacobian norms do not explode during backpropagation. \n\nTherefore, while spectral conditioning may increase the Jacobian norm, Layer Normalization helps mitigate this effect and prevents the norm from becoming excessively large. We explicitly discuss this interaction in the paper: on page 211, we refer the reader to Appendix A.3, where we present an ablation study on spectral conditioning and Layer Normalization in the context of the Nyströmformer. See lines 539–542 and Table 8 in Appendix A.3. The table shows that removing Layer Normalization while retaining spectral conditioning leads to worse performance. This is because spectral conditioning alone can increase the Jacobian norm, and without Layer Normalization to regulate it, training becomes unstable.\n\nIn summary, spectral conditioning improves the conditioning of the singular value spectrum but does not necessarily reduce the Jacobian norm. In contrast, Layer Normalization helps control the norm of the Jacobian. Thus, spectral conditioning should be viewed as complementary to Layer Normalization, not a replacement. We will add this clarification to Appendix A.3 to make the distinction clear. \n\nWe would also like to add that we did compute the condition numbers of the Jacobian of the attention during training for the vision transformers and the Nyströmformer both with and without spectral conditioning, see Fig. 2, 3, 4 (right most figure). In each case you can see that spectral conditioning reduces the condition number of the Jacobian by almost 2-3 orders of magnitude clearly showing the stability of our methodology. \n\n4. **There are still some unclear notations and typos in the manuscript. Eq (4), $J$ \n should be removed from the partial derivation; Table 1, it is not immediately clear whether “Spec. cons.” refers to “Spec. cons. + the corresponding attention mechanism” and also the subsequent experiments. Caption of Figure 14, ViT-B should be corrected to Nyströmformer. Please double check the typos and expressions.**\n\nWe sincerely thank the reviewer for pointing this out. The $J$ in Equation (4) was a typo and will be removed from the partial derivative. In Tables 1-4, Spec. Cond consistently refers to applying spectral conditioning to the matrices $W_Q$, $W_K$, and $W_V$ used in the attention mechanism of the corresponding architecture. We will clarify this explicitly for the reader. The caption of Figure 14 will also be corrected to Nyströmformer."}, {"review_id": "8pfEEeaP4b", "reviewer": "Reviewer_GB8p", "summary": "This paper presents a theoretical analysis of the Jacobian matrix of self-attention layers in Transformers, revealing its dependence on query, key, and value projection matrices. Based on this insight, the authors propose a method to systematically adjust the spectral properties of self-attention layers to improve the overall conditioning of self-attention layers within Transformer networks. The method demonstrates consistent performance improvements across various architectures.", "questions": "1. Suggest experimental analysis to verify whether attention layers in Transformers exhibit similar condition number problems as feedforward networks.\n\n2. Recommend extending the analysis to feedforward components to comprehensively validate the applicability of this analytical framework to Transformer architectures.\n\n3. Suggest providing theoretical selection criteria for the $\\lambda$ parameter to enhance the method's practicality in large-scale models.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n**Solid theoretical foundation**: Establishes the mathematical relationship between the condition number of self-attention layer Jacobian matrices and the condition numbers of WQ, WK, WV matrices, transforming Transformer optimization into a matrix condition number optimization problem.\n\n**Simple and practical method**: The proposed spectral conditioning method serves as a drop-in replacement, requires no additional parameters, and is applicable to various attention mechanisms with good generalizability.\n\n**Comprehensive experimental validation**: Validates the method's effectiveness across multiple tasks including image classification, object detection, instance segmentation, and nlp.\n\n## Weaknesses:\n**Insufficient motivation**: While the authors mention the effectiveness of Jacobian conditioning optimization in feedforward networks [1], they do not adequately demonstrate whether similar problems exist in Transformers or the necessity of condition number optimization.\n\n**Lack of hyperparameter guidance**: The $\\lambda$ parameter relies on grid search, lacking theoretical guiding principles.\n\n[1] Weight conditioning for smooth optimization of neural networks. ECCV 2025.", "was_revised": true, "final_justification": "The authors have addressed my concerns regarding the research motivation and have further supported the necessity of Jacobian conditioning optimization in the Transformer architecture through additional experiments. They have also effectively clarified the role of the FFN component in Transformer blocks.\n\nHowever, the hyperparameter selection strategy remains insufficiently addressed, which may limit the applicability of the proposed method to larger-scale models. Therefore, I have raised my score from 3 to 4.", "rebuttal": "We sincerely thank the reviewer for taking the time to review our paper and for providing thoughtful and constructive feedback. Below, we address each of the points and questions raised in the review.\n\n1. **Insufficient motivation: While the authors mention the effectiveness of Jacobian conditioning optimization in feedforward networks [1], they do not adequately demonstrate whether similar problems exist in Transformers or the necessity of condition number optimization / Suggest experimental analysis to verify whether attention layers in Transformers exhibit similar condition number problems as feedforward networks.**\n\nWe thank the reviewer for this comment, though we find it somewhat surprising. In Figures 2 and 3, we compute the condition number of the Jacobian of the self-attention block for the ViT-B and XCiT models, respectively (see the rightmost plots in each figure). For ViT-B, the condition number reaches approximately $\\mathcal{O}(10^{10})$, while for XCiT it is around $\\mathcal{O}(10^8)$, both of which are extremely high. These same figures also show the effect of applying spectral conditioning: the condition number for ViT-B is reduced to about $\\mathcal{O}(10^7)$, and for XCiT it drops to approximately $\\mathcal{O}(10^5)$. In both cases, this represents a reduction of nearly three orders of magnitude, clearly demonstrating that the Jacobian of self-attention can become severely ill-conditioned during training, and that spectral conditioning provides a simple and effective remedy.\n\nWe further conducted this analysis on the Nyströmformer using the text classification task from the LRA benchmark. As shown in Figure 4, the condition number of the attention block’s Jacobian is around $\\mathcal{O}(10^6)$, and spectral conditioning brings it down to $\\mathcal{O}(10^4)$. Additional results are provided in Appendix A.3, where we carry out a similar analysis for the Nyströmformer on the ListOps task.\n\nTogether, these results provide strong empirical evidence that the Jacobians of attention blocks in various transformer architectures are highly ill-conditioned, and that spectral conditioning offers a simple and broadly effective method for improving their conditioning.\n\n2. **Recommend extending the analysis to feedforward components to comprehensively validate the applicability of this analytical framework to Transformer architectures.**\n\nThe primary focus of our paper is to investigate the condition number of the Jacobian of the self-attention block in Transformer architectures and to show that it is often highly ill-conditioned. Motivated by this observation, we propose a simple and theoretically grounded method, spectral conditioning, that effectively reduces the Jacobian’s condition number. We then demonstrate that applying this technique improves performance across a variety of Transformer-based applications. We believe that this is more than enough for a NeurIPS especially given the theoretical derivations and the comprehensive experimental validation.\n\nTo address the reviewer's comment, we conducted an additional set of experiments applying spectral conditioning to the feedforward (FFN) layers within each Transformer block. Specifically, for each weight matrix $W$ in the feedforward layers, we applied a spectral correction matrix $C_W$ as defined by Theorem 3.8, forming the updated matrix $W + C_W$. Note that Theorem 3.8 ensures that the conditioned matrix $W + C_W$ has a lower condition number than the original $W$.\nWe then trained ViT-B on the ImageNet-1k dataset under four settings:\n\n1. No spectral conditioning applied  \n2. Spectral conditioning applied only to the attention layers  \n3. Spectral conditioning applied only to the feedforward layers  \n4. Spectral conditioning applied to both attention and feedforward layers\n\nThe results are summarized in the table below. We observe that applying spectral conditioning only to the feedforward layers yields a minor performance gain. In contrast, conditioning only the attention layers results in an approximate 1% improvement. Applying spectral conditioning to both components further boosts accuracy by about 1.1% overall. \n\n| Model                                           | Accuracy |\n|------------------------------------------------|----------|\n| ViT-B (original)                                | 80.7     |\n| ViT-B with Spec. Cond. on Attention             | 81.7     |\n| ViT-B with Spec. Cond. on Feedforward           | 80.9     |\n| ViT-B with Spec. Cond. on Attention and Feedforward | 81.8     |\n--------------------------------------------------------------------------------------\n\nThese findings support our main claim: for ViT-B models trained on ImageNet-1k, applying spectral conditioning to the self-attention layers significantly improves performance, and conditioning the feedforward layers provides a small additional benefit. This reinforces the importance of controlling the Jacobian’s condition number in attention mechanisms and further highlights the effectiveness of our proposed method.\n\nWe also computed the condition number of the Jacobian for both the self-attention and feedforward layers of a ViT-B model during training on ImageNet-1k. For the self-attention layers, we extracted the condition number of the Jacobian for each head in each layer, then averaged across all heads, all layers, and across the entire training pass to obtain a single average condition number representing the self-attention block over training. Similarly, for the feedforward layers, we computed the condition number of the Jacobian for each layer and averaged these values across layers and training steps to obtain an overall average for the feedforward blocks.\n\nWe repeated this procedure after applying spectral conditioning. The results are shown in the table below. Notably, we observe that the average condition number of the self-attention block is significantly higher than that of the feedforward block. Applying spectral conditioning to the attention layers reduces their condition number by nearly three orders of magnitude, whereas applying it to the feedforward layers results in a reduction of roughly one order of magnitude.\n\nThese results clearly indicate that the attention mechanism is the most ill-conditioned component of the Transformer architecture, and that spectral conditioning is particularly effective in addressing this issue.\n\n|                            | Average condition number |\n|----------------------------|--------------------------|\n| Attention                  | 9.8e9                   |\n| Feedforward                | 2.6e4                   |\n| Spec. Cond. Attention      | 2.8e7                   |\n| Spec. Cond. Feedforward    | 3.2e3                   |\n\nWe repeated the above analysis for the Nyströmformer on the text classification task in the LRA benchmark set.\n\n\n| Model                                                       | Accuracy |\n|-------------------------------------------------------------|----------|\n| Nyströmformer (original)                                    | 63.8     |\n| Nyströmformer with Spec. Cond. on Attention                 | 64.8     |\n| Nyströmformer with Spec. Cond. on Feedforward               | 63.9     |\n| Nyströmformer with Spec. Cond. on Attention and Feedforward | 64.9     |\n\n\n\n| Component                | Average Condition Number |\n|--------------------------|--------------------------|\n| Attention                | 9.1e6                    |\n| Feedforward              | 4.6e3                    |\n| Spec. Cond. Attention    | 2.1e4                    |\n| Spec. Cond. Feedforward  | 8.9e2                    |\n\n3. **Lack of hyperparameter guidance: The $\\lambda$\n parameter relies on grid search, lacking theoretical guiding principles / Suggest providing theoretical selection criteria for the $\\lambda$ \n parameter to enhance the method's practicality in large-scale models.**\n\nTheorem 3.8 provides a theoretical lower bound for the $\\lambda$ value. While we agree that a deeper theoretical explanation for selecting $\\lambda$ would be valuable, we believe that such an analysis falls outside the scope of the current work and intend to explore it in future research. That said, we do provide a concrete ablation study in Appendix A.2.1, treating $\\lambda$ as a tunable hyperparameter. This study illustrates how performance varies as $\\lambda$ is adjusted, offering clear empirical insight into its effect."}]

Author rebuttal:

Meta-review: This paper studies the conditioning of an attention layer but studying its Jacobian. The authors identify the Q K V matrices dominate the condition number, and this proposes a new parameterization of the attention layer that makes it better conditioned. All reviewers have found this work interesting and its findings relevant to the ML community. The AC agrees this consensus and recommends accept. 

One note, this work's finding has high similarity to the SigmaReparam work [1] where they apply spectral norm to the Q K V matrices instead. The authors are encouraged cite and discuss their relation in the camera ready version.

References

[1] Stabilizing Transformer Training by Preventing Attention Entropy Collapse, Zhai et al, ICML 2023
Decision: Accept (poster)