Spectral Conditioning of Attention Improves Transformer Performance

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.

1. 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.

We 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.

We 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.

Together, 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.

2. Recommend extending the analysis to feedforward components to comprehensively validate the applicability of this analytical framework to Transformer architectures.

The 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.

To 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$. We then trained ViT-B on the ImageNet-1k dataset under four settings:

1. No spectral conditioning applied
2. Spectral conditioning applied only to the attention layers
3. Spectral conditioning applied only to the feedforward layers
4. Spectral conditioning applied to both attention and feedforward layers

The 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.

---

These 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.

We 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.

We 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.

These 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.

We repeated the above analysis for the Nyströmformer on the text classification task in the LRA benchmark set.

3. Lack of hyperparameter guidance: The $\lambda$ parameter relies on grid search, lacking theoretical guiding principles / Suggest providing theoretical selection criteria for the $\lambda$ parameter to enhance the method's practicality in large-scale models.

Theorem 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.

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.

1. 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.

This 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. In 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.

---

2. 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.

Theorems 3.3 and 3.4 establish that the condition number of the self‑attention Jacobian is controlled by the projection matrices $(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.

3. 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

The 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

A = \\begin{bmatrix} 1 & 0 \\\\ 0 & 1 \\end{bmatrix} \\quad \\text{and} \\quad B = \\begin{bmatrix} 10.1 & 0 \\\\ 0 & 10 \\end{bmatrix}.

The matrix $A$ has condition number $1$ and the matrix $B$ has condition number $1.01$, thus we see that both 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$.

C = \\begin{bmatrix} 1 & 0 \\\\ 0 & 1 \\end{bmatrix} \\quad \\text{and} \\quad D = \\begin{bmatrix} \\sqrt{2} & 0 \\\\ 0 & 0.1 \\end{bmatrix}.

$C$ has a condition number of $1$ and $D$ has condition number of approximately $14.14$, thus $D$ has a much larger condition number than $C$. However, their Frobenius norms are similar.

This demonstrates that computing the norm of the Jacobian during training does not necessarily provide any information about its condition number.

In 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.

Therefore, 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.

In 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.

We 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.

4. 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.

We 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.

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.

1. The method only aims to minimize the upper bound of the condition number, and the practical implementation further loosens the upper bound being minimized.

We 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. While 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.

In 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.

2. Theorem 3.5 is a trivial result that is unnecessary to derive the final practical method (Theorem 3.8 is all you need).

Thank 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.

3. Self-attention is a special case of attention, and it is not clear why the analysis does not extend to cases like cross-attention.

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$ and not the input tokens $X$, and thus also holds for cross attention. We will include a paragraph explaining this in the appendix.

We 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.

4. 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.

The 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.

We 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.

5. 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.

While 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.

We sincerely thank the reviewer for their thorough examination of both the main text and the appendix. Your detailed feedback is greatly appreciated.

1. If there are any mistakes in this reasoning, please let me know...are there any missing assumptions, that are satisfied in practice, that can help fix the theory?

Thank you so much for point this out. As you have shown inequality (53) in the paper we have written is indeed incorrect. The correct form of the Weyl inequality that we need to use is

$\sigma_{\min}(A + \lambda I) \geq \sigma_{min}( \lambda I) - \sigma_{\max}(A)$.

In order for our result to hold we do need to add a condition namely that

$\lambda - \sigma_{\max}(A) > 1$

We further assume $\sigma_{\min}(A) < 0.4$ and $\sigma_{\max}(A) > 2$.

With these assumptions we can give the proof as follows:

$\kappa(A + \lambda I) = \frac{\sigma_{\max}(A + \lambda I)}{\sigma_{\min}(A + \lambda I)}$

$\leq \frac{\sigma_{\max}(A + \lambda I)}{\lambda - \sigma_{\max}(A)}$ (by above Weyl inequality. Note denominator is positive by our condition)

$\leq \frac{\sigma_{\max}(A) + \lambda}{\lambda - \sigma_{\max}(A)}$ (by Weyl inequality (52) in appendix)

$< \sigma_{\max}(A) + \frac{\lambda}{\lambda - \sigma_{\max}(A)}$ (by our condition above)

$\leq \sigma_{\max}(A) + 1 + \frac{\sigma_{\max}(A)}{\lambda - \sigma_{\max}(A)}$

$< 2\sigma_{\max}(A) + 1$ (by our condition above)

$= 2(\sigma_{\max}(A) + \frac{1}{2})$

$< \frac{\sigma_{\max}(A)}{\sigma_{\min}(A)}$ (by our assumed conditions on $\sigma_{\max}(A)$ and $\sigma_{\min}(A)$)

$= \kappa(A)$.

Given this we feel it is best to write theorem 3.8 as:

Theorem: Let $A \in \mathbb{R}^{m\times n}$ and assume that $\sigma_{\min}(A) < 0.4$ and $\sigma_{\max}(A) > 2$. Let $I_k \in \mathbb{R}^{m\times n}$ denotethe matrix that has 1 on its main $k \times k$ diagonal, where $k = \min(m, n)$ and zero elsewhere. Assume at least one of the following two conditions hold:

(1) $\lambda - \sigma_{\max}(A) > 1$ or

(2) $\sigma_{\min}(A + \lambda I_k) \geq \sigma_{\min}(\lambda I_k) - \sigma_{\min}(A)$ and $\lambda \geq 2$

Then $\kappa(A + \lambda I_k) < \kappa(A)$.

We believe the above stated theorem employs Weyl's inequality the right way and provides conditions under which the theorem holds.

We reanalyzed the saved attention weights and found that condition (2) held for most transformers, even though it need not for arbitrary matrices, while in ViT-B (and similarly DaViT-B) condition (1) held for $W_Q$ and $W_V$ and condition (2) for $W_K$. In every case at least one condition held. We’ll expand the appendix with a detailed discussion and a plot of both conditions over training.

A cleaner narrative can be presented by including results that utilize Theorem 3.5, showing it to perform better than standard attention...

Thank you. In the original draft we only had Theorem 3.5, but its stronger guarantee came at a cost, for example vision transformers took 12–20 hours longer to train. To address this, we derived Theorem 3.8, which substantially cuts training time at the expense of performance. We’ve preserved all experiments with Theorem 3.5 and can include them in the appendix to illustrate this trade-off. Below, we show the training times for the different methods for ViTs.

If the theory holds more broadly, it is better to present the theory more generally and note the special case of self-attention..

We focused on self-attention because it’s the most common form. We can definitely include a general formulation in the appendix and retitle the paper to “Spectral Conditioning of Attention Improves Transformer Performance.” As you noted we already have empirical results for broader attention mechanisms in the paper.

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.

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....

We 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.

(i) Results on CUB-200

---

(ii) Results on CARS-196

---

(iii) Results on iNat2017

---

2. I would suggest incorporating more Detection/Segmentation Models.

We 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.

We 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.

---

3. A potential omission is the lack of evaluation on vision-language transformer models, such as CLIP or SigLIP....

We 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.

4. 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...

We 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.

We 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.