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.
Figures & tables
Figure 1: Average condition number of the attention Jacobian during training under three common initialization schemes, shown alongside the theoretical bound from eq. 8 .
ViT-B
DeiT-B
Swin-B
XCiT-M
DaViT-B
Original
80.3
81.6
83.4
82.6
84.3
Mimetic
80.5
81.6
83.5
82.6
84.4
Conditioned (ours)
81.5
82.7
84.6
83.5
85.3
Table 1: Comparison of Vision Transformers with different initializations pretrained on ImageNet-1k. We report Top-1% classification accuracy. In each case, conditioned initialization improves performance over the default and mimetic initializations.
Figure 2: Total number of epochs required for each initialization to reach the final accuracy of the default initialization reported in table 1 , across ViT-B, DeiT-B, Swin-B, and XCiT-M. In all cases, conditioned initialization converges more quickly, requiring fewer epochs than both default and mimetic initialization.
Pets
Flowers
CIFAR-10
CIFAR-100
Default
26.7
64.5
92.4
71.7
Mimetic
47.7
71.6
93.6
75.0
Conditioned (ours)
47.7
72.1
94.1
75.3
Table 2: Comparison of ViT-T pretrained on four small scale datasets with three different initializations. We report Top-1% classification accuracy. In each case, conditioned initialization improves performance over the default and mimetic initializations.
Figure 3: Average condition number of the attention Jacobian during training under three common initialization schemes, shown alongside the theoretical bound from eq. 8 .
Model
APb
AP50b
AP75b
APm
AP50m
AP75m
Default
44.9
66.1
48.9
40.1
63.1
42.8
Mimetic
44.8
66.0
49.1
40.2
63.1
42.9
Conditioned (ours)
45.5
66.8
49.5
40.6
63.5
43.3
Table 3: 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
ListOps
Text
Retrieval
Image
Pathfinder
Default
37.1
63.8
79.8
39.9
72.9
Mimetic
37.1
64.0
79.9
40.2
73.2
Conditioned (ours)
37.9
64.9
80.8
40.4
73.9
Table 4: Nyströmformer with three different initializations on the LRA benchmark. We report evaluation accuracy (%). As shown, our initialization improves performance across all tasks.
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
Mimetic
84.1
92.5
86.5
55.1
90.3
87.5
85.1
50.1
78.9
Conditioned
84.8
92.9
86.9
55.5
91.1
87.7
86.0
51.7
79.6
Table 5: We evaluate a pretrained Crammed BERT with three different initialization schemes on the GLUE benchmark, and find that conditioned initialization outperforms the other two.
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: Total number of epochs required for each initialization to reach the final accuracy of the default initialization reported in table 1 for DaViT-B architecture. Conditioned initialization converges more quickly, requiring fewer epochs than both default and mimetic initialization.
Food-101
CIFAR-10
CIFAR-100
ImageNet-1k + Default
85.5
96.6
79.7
ImageNet-1k + Mimetic
86.4
96.3
79.9
ImageNet-1k + Conditioned (ours)
87.3
97.1
81.0
ImageNet-21k + Default
87.1
97.1
81.1
Appendix
Table 6: Comparison of ViT-T pretrained on four small scale datasets with three different initializations. We report Top-1% classification accuracy. In each case, conditioned initialization improves performance over the default and mimetic initializations.
Figure 5: Training loss curves for different initializations for ViT-B (left) and DeiT-B (right).
Figure 6: Training loss curves for different initializations for Swin-B (left) and XCiT-M (right).
Figure 7: Training loss curves for different initializations for DaViT-B.
Perplexity
Default
2.47
Mimetic
2.40
Conditioned (ours)
2.29
Appendix
Table 7: GPT-2 models trained on the TinyStories dataset. We initialize a model in three different ways. Conditioned initialization achieves a perplexity than the other two initializations.
Model
Perplexity
Default
2.39
Mimetic
2.32
Conditioned (ours)
2.20
Appendix
Table 8: Perplexity of 472M GPT-2 model pretrained on TinyStories.
Model
Perplexity
Default
44.2
Mimetic
43.8
Conditioned (ours)
42.7
Appendix
Table 9: Perplexity of 472M GPT-2 model pretrained on WikiText-103.
Model
Perplexity
Default
4.15
Mimetic
4.25
Conditioned (ours)
4.03
Appendix
Table 10: Perplexity of 1.72B GPT-2 model pretrained on TinyStories.
Model
Perplexity
Default
48.1
Mimetic
48.5
Conditioned (ours)
46.9
Appendix
Table 11: Perplexity of 1.72B GPT-2 model pretrained on WikiText-103.
ViT-L
DeiT-L
Swin-L
XCiT-L
DaViT-L
Original
79.6
80.7
82.6
81.5
83.3
Mimetic 2
79.7
80.6
82.4
81.5
83.2
Conditioned (ours)
80.7
81.5
83.7
82.4
84.4
Appendix
Table 12: Large scale ViTs with different initializations pretrained on ImageNet-1k. We show the Top1% accuracy.
ViT-G
DaViT-G
Original
78.8
81.9
Mimetic 2
78.6
82.0
Conditioned (ours)
79.9
82.9
Appendix
Table 13: 1B scale ViTs with different initializations pretrained on ImageNet-1k. We show the Top1% accuracy.
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.
Hemanth Saratchandran
Australian Institute for Machine Learning (AIML), Adelaide University
We present a theoretical analysis of the Jacobian of an 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.
Hemanth Saratchandran, Simon Lucey
Australian Institute for Machine Learning, Adelaide University
Transformer attention is typically implemented using softmax normalization, which enforces attention weights with unit sum normalization. While effective in many settings, this constraint can limit flexibility in controlling attention magnitudes and may contribute to overly concentrated or unstable attention patterns during training. Prior work has explored modifications such as attention sinks or gating mechanisms, but these approaches provide only limited or indirect control over attention reweighting. We propose Affine-Scaled Attention, a simple extension to standard attention that introduces input-dependent scaling and a corresponding bias term applied to softmax-normalized attention weights. This design relaxes the strict normalization constraint while maintaining aggregation of value representations, allowing the model to adjust both the relative distribution and the scale of attention in a controlled manner. We empirically evaluate Affine-Scaled Attention in large-scale language model pretraining across multiple model sizes. Experimental results show consistent improvements in training stability, optimization behavior, and downstream task performance compared to standard softmax attention and attention sink baselines. These findings suggest that modest reweighting of attention outputs provides a practical and effective way to improve attention behavior in Transformer models.