Preconditioned Attention: Enhancing Efficiency in Transformers

Hemanth Saratchandran

cs.LGMar 28, 2026

Preconditioned Attention: Enhancing Efficiency in Transformers

Authors: Hemanth Saratchandran

Organizations: Australian Institute for Machine Learning (AIML), Adelaide University

Abstract

Central to the success of Transformers is the attention block, which effectively models global dependencies among input tokens associated to a dataset. However, we theoretically demonstrate that standard attention mechanisms in transformers often produce ill-conditioned matrices with large condition numbers. This ill-conditioning is a well-known obstacle for gradient-based optimizers, leading to inefficient training. To address this issue, we introduce preconditioned attention, a novel approach that incorporates a conditioning matrix into each attention head. Our theoretical analysis shows that this method significantly reduces the condition number of attention matrices, resulting in better-conditioned matrices that improve optimization. Conditioned attention serves as a simple drop-in replacement for a wide variety of attention mechanisms in the literature. We validate the effectiveness of preconditioned attention across a diverse set of transformer applications, including image classification, object detection, instance segmentation, long sequence modeling and language modeling.

Figures & tables

Figure 1: Schematic representation of preconditioned self-attention. Left: A layer of a general transformer employing a preconditioned self-attention block. Right: Self-attention (top) and preconditioned self-attention (bottom) are compared. The key difference is that preconditioned self-attention applies a multiplication by a diagonal preconditioner matrix C, which depends on the query $Q$ , key $K$ , and value $V$ .

	Models
	ViT-base	XciT-medium	DeiT-base	Swin-base
Original	80.2	82.6	81.6	83.4
Preconditioned	81.4	83.5	82.7	84.5

Table 1: Comparison of vision transformers with their original attention mechanism versus one with a preconditioned attention mechanism pretrained on the ImageNet-1k dataset. We report the classification top-1% accuracy.

Figure 2: Total number of epochs required for each preconditioned model to reach the accuracy of the baseline model. For each ViT we see that the preconditioned model requires roughly 20-30% less epochs.

Figure 3: Average condition number of a ViT-B and a preconditioned ViT-B during training on the ImageNet-1k dataset.

	Train Time (hrs:mins)	Memory (GB)
ViT-Base	29:28	176.1
Precond. ViT-Base	30:45	178.5
DeiT-Base	26:16	173.6
Precond. DeiT-Base	27:25	175.1
Swin-Base	53:11	150.4
Precond. Swin-Base	54:24	152.6

Table 2: Training time/memory overhead for ViTs

Model	$AP^{b}$	$AP^{b}_{50}$	$AP^{b}_{75}$	$AP^{m}$	$AP^{m}_{50}$	$AP^{m}_{75}$
XCiT-S12	44.9	66.1	48.9	40.1	63.1	42.8
Preconditioned XCiT-S12	45.4	66.4	49.5	40.5	63.3	43.2
XCiT-M24	45.7	66.8	49.6	40.8	63.6	43.3
Preconditioned XCiT-M24	46.0	67.2	49.9	41.2	63.9	44.0

Table 3: Performance on COCO object detection and instance segmentation (mini-val set) using Mask R-CNN with ImageNet-1k pretrained backbones. Reported metrics include bounding box AP (

AP^{b}

), AP at IoU thresholds 0.50 and 0.75 (

AP^{b}_{50}

AP^{b}_{75}

), mask AP (

AP^{m}

), and mask AP at IoU thresholds 0.50 and 0.75 (

AP^{m}_{50}

AP^{m}_{75}

Model	ListOps	Text	Retrieval	Image	Pathfinder
Nyströmformer	37.1	63.8	79.8	39.9	72.9
Preconditioned Nyströmformer	37.9	64.8	80.7	40.8	73.8

Table 4: Comparison of a Nyströmformer with a preconditioned Nyströmformer on LRA benchmark. We report evaluation accuracy (%).

Figure 4: Conditioning analysis of the Nyströmformer on text classification and ListOps. The preconditioned variant achieves consistently lower condition numbers (left and middle) and converges faster, requiring fewer iterations to reach the baseline model’s final accuracy (right).

	MNLI	SST-2	STSB	RTE	QNLI	QQP	MRPC	CoLA	Avg.
Default	83.8	92.3	86.3	55.1	90.1	87.3	85.0	48.9	78.6
Preconditioned	84.9	92.9	87.5	55.9	90.9	88.4	85.9	50.0	79.6

Table 5: Evaluation of pretrained Crammed BERT using the default baseline and a variant with preconditioned attention (as defined in Algorithm 1 ). The preconditioned model consistently outperforms the baseline across GLUE tasks.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

	Listops (GB)	Text Class. (GB)	Retrieval (GB)	Image Class. (GB)	Pathfinder (GB)
Original	5.6	3.6	13.8	9.6	9.6
Preconditioned	5.8	3.6	15.8	9.6	9.6

Appendix

Table 6: Memory used for training Nyströmformer and preconditioned Nyströmformer on LRA benchmark.

	Total train time on LRA (hrs:mins)
Nyströmformer	8:30
preconditioned Nyströmformer	8:37

Appendix

Table 7: Total training time for Nyströmformer and preconditioned Nyströmformer on LRA benchmark.

Explore similar work

Conditioned Initialization for Attention
Sep 7, 2026Hemanth Saratchandran, Simon LuceyNeural Network Initialization Transformer Attention
Australian Institute for Machine Learning Adelaide University
Spectral Conditioning of Attention Improves Transformer Performance
Mar 7, 2026Hemanth Saratchandran, Simon LuceyTransformer Attention
Australian Institute for Machine Learning, Adelaide University
A Mechanistic Study of Transformers Training Dynamics
Oct 31, 2024Ambroise Odonnat, Wassim Bouaziz, Vivien CabannesTransformer Interpretability Attention Mechanisms
Inria, Univ. Rennes 2 · Mistral AI; work done while at Meta. · FAIR at Meta.

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

Conditioned Initialization for Attention

Hemanth Saratchandran, Simon Lucey

Australian Institute for Machine Learning Adelaide University

Neural Network Initialization Transformer Attention

Open paper AskPDF arXiv

Transformers are a dominant architecture in modern machine learning, powering applications across vision, language, and beyond. At the core of their success lies the attention layer, where the query, key, and value matrices determine how token dependencies are captured. While considerable work has focused on scaling and optimizing Transformers, comparatively little attention has been paid to how the weights of the queries, keys and values are initialized. Common practice relies on random initialization or alternatives such as mimetic initialization, which imitates weight patterns from converged models, and weight selection, which transfers weights from a teacher model. In this paper, we argue that initialization can introduce an optimization bias that fundamentally shapes training dynamics. We propose conditioned initialization, a principled scheme that initializes attention weights to improve the spectral properties of the attention layer. Theoretically, we show that conditioned initialization can potentially reduce the condition number of the attention Jacobian, leading to more stable optimization. Empirically, it accelerates convergence and improves generalization across diverse applications, highlighting conditioning as a critical yet underexplored area for advancing Transformer performance. Importantly, conditioned initialization is simple to apply and integrates seamlessly into a wide range of Transformer architectures.

Abstract

Explore similar work

Conditioned Initialization for Attention

Spectral Conditioning of Attention Improves Transformer Performance

A Mechanistic Study of Transformers Training Dynamics