Conditioned Initialization for Attention
Authors: Hemanth Saratchandran
Simon Lucey
Venue: ICLR 2026 Poster (2026)
Keywords: spectral conditioning transformers
spectral properties of attention
Primary area: other topics in machine learning (i.e., none of the above)

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

Published as a conference paper at ICLR 2026
CONDITIONED INITIALIZATION FOR ATTENTION
Hemanth Saratchandran Australian Institute for Machine Learning Adelaide University
Simon Lucey
Australian Institute for Machine Learning Adelaide University
ABSTRACT
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 mimetic initialization, which imitates weight patterns from converged models, and weight selection, which transfers weights from a teacher model. In this paper, argue that initialization can in- troduce 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 num- ber of the attention Jacobian, leading to more stable optimization. Empirically, it accelerates convergence and improves generalization across diverse applications. highlighting conditioning as critical yet underexplored area for advancing Trans- former performance. Importantly, conditioned initialization is simple apply and integrates seamlessly into wide range of Transformer architectures.
INTRODUCTION
Transformers Vaswani et al. 2017) have rapidly become cornerstone of modern machine learning,
driving progress in fields as diverse as natural language processing (Vaswani et al., 2017 Zhuang
al. 2021 Zhen et al. 2022), computer vision (Dosovitskiy et al. 2020 Liu et al. 2021 Touvron
et al. 2021 Carion et al. 2020), and robotics (Salzmann et al.) 2020 Maiti et al., 2023) A key
factor behind this versatility is the self-attention mechanism, which models interactions between tokens by comparing them pairwise and dynamically weighting their contributions. This ability to capture both long-range and global dependencies has positioned Transformers as foundational architecture across a wide spectrum of learning tasks.
While substantial effort has been devoted to improving the efficiency, scalability, and expressive- ness of attention mechanisms (Ali et al. 2021 Xiong al. 2021b Ding al. 2022), relatively little focus has been placed on a more basic but equally important aspect: how attention weights are initialized. Initialization plays a critical role in shaping optimization landscapes of deep neu- ral networks. Classical schemes such as Xavier (Glorot & Bengio, 2010) and Kaiming (He et al.) 2015) initialization showed that carefully chosen scaling of weights at the start of training can dra- matically improve gradient optimization and stability in deep networks. These insights were crucial for enabling the training of very deep architectures such as ResNets (He et al. 2016). where poor initialization could otherwise cause vanishing or exploding gradients. Yet, despite their depth and complexity, Transformers have received far less theoretical scrutiny on this front. Given that self- attention relies query, key, and value projections whose interplay directly governs the stability of token interactions, it is natural to ask whether initialization schemes tailored specifically to attention could offer similar benefits.
Currently, Transformers typically adopt simple random initializations, without consideration of the unique structure of attention layers. Recent alternatives such mimetic initialization (Trockman & Kolter, 2023), which transfers statistical patterns from converged models, and weight selection (Xu al. 2023), which reuses pretrained weights from larger teacher models, highlight growing recog-

---

Published as a conference paper at ICLR 2026
nition that initialization matters. However, these methods remain heuristic and lack principled connection to the conditioning of the attention mechanism itself.
In this work, revisit Transformer initialization from theoretical perspective. We show that the stability of optimization in self-attention layers is closely tied to the conditioning of their Jacobians, which in turn depends the spectral properties of the query, key, and value projections. Building on this insight, we introduce conditioned initialization, a principled scheme designed to improve the spectral conditioning of attention blocks the start of training. Rather than directly modifying training objectives, our method intervenes at initialization, providing an inductive bias that promotes more stable optimization dynamics.
Our contributions are threefold:
1. Theoretical framework: We establish a connection between the conditioning of self- attention Jacobians and the spectral structure of the query, key, and value matrices, mo- tivating initialization schemes that explicitly target this property.
2. Conditioned initialization: We propose simple initialization method that reduces the upper bound on the attention Jacobian's condition number, thereby biasing training toward stable optimization.
3. Empirical validation:
Through experiments on diverse benchmarks, spanning image clas- sification, object detection, instance segmentation, language modeling, and long-range se- quenee learning, we show that conditioned initialization consistently accelerates conver- gence and improves generalization, and can be easily integrated into variety of different Transformer architectures.
By highlighting initialization as a critical yet underexplored component of Transformer design, our work opens up a new perspective how principled conditioning can be leveraged to improve opti- mization stability and downstream performance.
2 RELATED WORK
Initialization. Initialization strategies play pivotal role in determining how efficiently deep net- works can be trained. Early advances such as Xavier (Glorot & Bengio 2010) and Kaiming initial- ization (He al. 2015) demonstrated that properly scaling weights the outset helps preserve vari- ance across layers. preventing issues like vanishing or exploding gradients. These ideas were foun- dational in enabling the successful training of very deep architectures, most notably ResNets (He) eet al.) 2016). In the context of Transformers, however, initialization has typically been treated in a more hoe manner, with standard practice relying on simple schemes from normal or trun- cated normal distributions. More recent efforts have begun to acknowledge the unique structure of attention layers. For example, mimetic initialization (Trockman Kolter 2023) introduces induc- tive bias by imitating the statistical patterns of trained networks, while weight selection (Xu et al.) 2023) transfers weights from larger pretrained teacher models to provide a stronger starting point for smaller architectures. In our experiments compare these two methods.
Conditioning. A growing body of work has underscored the importance of conditioning for both the trainability and generalization of neural networks. In particular, [Saratchandran et al. (2025) showed that networks with better-conddtioned weights tend achieve superior performance, and proposed a matrix preconditioning method to explicitly control condition numbers during training. From a different angle. Liu et al. (2022) analyzed optimization through the lens of the neural tangent kernel (NTK), demonstrating that conditondd NTKs lead to faster and more reliable conver- gence, particularly the infinite- -ihth regime where the NTK dominates learning dynamics (Jacot et al. 2018). This has been extended to the transformer setting in Yangl (2020). Other studies have pointed out structural factors that affect conditioning. For example, Agarwal et al. (2021) observed that increasing network depth can itself improve conditioning, thereby aiding gradient-based meth- ods. In the case of Transformers, Ji et al. (2025) argued that skip connections act as an implicit conditioning mechanism, stabilizing the optimization of deep Transformers. Our work departs from these approaches by asking whether initialization itself can be designed to yield conter-conditioned attention layers from the outset.
2

---

Published as a conference paper at ICLR 2026
3 THEORETICAL FRAMEWORK
3.1 PRELIMINARIES
For the theoretical framework we will primarily focus on self-attention, which one of the most common forms of attention in a Transformer. Self-attention is composed of three learnable matrices, query Wq E IRDx key WK E Dxd, and value Wy Dxd defined for an input sequence X IRRN N.D.
(1) where softmax is the softmax activation that acts row-wise on a matrix (Prince) 2023). Note that then A(X) RRd In general. Transformers employ multiple heads i.c. attention matrices A, for 1 < i S h where h is the number of heads. These are then concatenated together to form a multi-head attention layer [AI...., A,). We point out that in some references the notation A(X) is reserved for only the term softmax X W. & However, in this paper as we will be concerned with the whole attention layer we use A (X) to denote the whole layer output as defined in eq. 4
A(X) = softmax (acllmaxxxW(lllil X
For further details on Transformers readers may consult Prnsuet(Prince (2023).
The self-attention map of a layer in Transformer A ( has parameters given by those parameters in X from the previous layer and those given by Wq. WK. Wv that define A(X). Our work will consider the Jacobian of A(X) with respect to the parameters within the layer of A(X), namely Wa. WR K. Wy. Therefore, when speak of the Jacobian of A(X) it will be with respect to WQ. WK. Wy. We will denote this Jacobian by J((()) and note that it is defined by
[OA(X) AA(X) (A(X) T
J(A(X))
LaW OWK dwv
Given a matrix Z E R=X"" we denote the vectorization of Z by vec(Z) E R R Magnus & Neudecker 2019). Note that for such matrix there transformation Tmn Rmmxxm such that vec Z) vec
where z" denotes the transpose of Z. The matrix Tmn is known as a commutation matrix and is a permutation matrix (Magnus & Neudeeker 2019). The maximum sin- gular value of a matrix Z will be denoted by Tmax (2) and the minimum singular value by m (Z). We will use the standard terminology SVD to denote the singular value decomposition of matrix. Given a vector E R" the notation ||æ||| will denote the vector 2-norm of e. Finally, we will let
Xoxn denote the rectangular identity matrix that has all I's on its main diagonal and O m x as the real X semi-orthogonal matrices.
3.2 MAIN THEOREMS
In this section, we present the main theorem of the paper, which motivates the development of a sim- ple yet effective strategy, conditioned initialization, designed to reduce the condition number of the Jacobian of the attention layer at initialization. As shown in section 4 this scheme is straightforward implement while providing significant optimization benefits.
Definition 3.1. Let Z be N x d matrix of full rank. The condition number of Z, denoted by K, is defined as
max ZZ)
<(Z)
(3)
Tmin (Z)
where mmax (Z) denotes the maximum singular value of Z and mmin (Z) the minimum singular value of 7, which know is non-zero Z is full rank.
Our objective is to analyze the condition number of the self-attention layer in a Transformer. We show that the condition number of its Jacobian depends on the condition numbers of the query, key, and value weight matrices. Furthermore, we demonstrate that initializing Wo. WK. and Wr with low condition numbers imparts an inductive bias into the Transformer architecture that leads to more effective optimization.
We begin by examining the derivatives of the self-attention layer with respect to the parameters Wo. W k. and Wy. We will need the following lemma.
3

---

Published as a conference paper at ICLR 2026
Lemma 3.1. Lef A R + Raxi denote the function (2) Diag(z) - 2. =" We then have that
Ōsoftmax
(z) A(softmax(2)).
(4)
dx
Proposition 3.1. Let A(X) denote self-attention matrix with inpur X as defined by eq.
Then
0A(X) (W x & (N) M(soltmxx(XWWW2X33 (XWx x
(5)
OWQ
0A(X) (WTx* & INN A(sottma((XWWWXXX) (X XWQ) TDa
(6)
OWK
(A(X) Ia soflmax(XWooWia
dWy
where A is defined lemma 3.1 and Tdd is the commutation matrix satisfying Taavoo (oo)
vcc(WX) (see section 3.7.
From proposition matrix.
3.1 we obtain a bound on the condition number of the Jacobian of the attention
Theorem 3.1. Let A(X) denote self-attention matrix, defined in eq. (1. with input X and let (X))) denote its Jacobian with respect to the parameter matrices Wo. WK and Wy as defined in eq. (2). Assume that J (A( has full rank so that k(J(A(X))) is finite. Then
x(J(A(X))) x(X))* x ((softmax(X (
+((K)) (8)
+ x(X [[ssssssmxxxxxXWWW
where A is defined in lemmal
Theorem 3.1 shows that K (A(())))) is bounded above by a sum of two terms:
k(X)s K ( (Assofmmax(XW W ( + n(Wx))
(9)
((sotmax((XWWW>>
(10)
Observation. Theorem 3.1 provides strategy for reducing the condition number of the Jacobian of A through the upper bound in eq. (8). Since we directly control Wq. WK. and Wy, reducing their condition numbers decreases to show that this can be achieved at initialization and to confirm empirically in section 4 that introduces a bias which improves both optimization and performance.
term eq. (9). thereby tightening the bound. Our next step
3.3 CONDITIONED INITIALIZATION
Our goal in this section is to design simple and effective initialization for the query, key, and value matrices Wo. WK. and Wy that lowers the condition number of their singular value spectra, thereby improving the conditioning of the Jacobian J(A).
We begin with two key observations. There are two families of m x n matrices with condition number 1:
1. scalar multiples the identity Xmxxn with + 0, and
2. semi-orthogonal matrices Omxa (matrices with orthonormal rows columns).
From theorem
3.1 the condition number of J(A) admits the surrogate upper bound
B(J(A)) x(X)3 <(A(softmax(XXW & W WxXT)) x(Wv) (m(W)) s(W)))
(11)
+ xXX) x(sottmax(XW
(12)
Unlike the true condition number x{J(A)), the bound ((J(A)) can be directly influenced at initial- ization by controlling the conditioning Wo. WK. and Wy. Standard practice initializes these matrices from Gaussian uniform distributions, which do not enforce good conditioning.
4

---

Published as a conference paper at ICLR 2026
The following proposition shows that initializing Wo. WK. and Wy from either of the families above yields a strictly better surrogate bound at initialization.
Let ( denote attention matrix with Wq. WK. and Wy initialized from a Gaussian uniform distribution, and let A(X) denote one where W Q W K. and W v are initialized from either Dx or Oxxa- Then
Proposition 3.2.
B(J(A)) S B(J(A)).
Remark. The quantity (A)( is only an upper bound on x()())) Thus, proposition 3.2 guaran- tees tighter bound but does not by itself imply improved conditioning of the Jacobian. Nonetheless, as we demonstrate in section [ the initialization in proposition 3.2 consistenty lowers the condition number during training and leads more stable optimization and improved performance.
Initialization Strategy. Although proposition 3.2 shows that initializing the matrices Wo, WK. and Wv using either (Apxxa} or Oxx can potentially lower the condition number of J(A), it does not specify which choice is most suitable for each matrix. We note that We. W and Wy play distinct algebraic roles in attention, and this motivates different treatments. For the value map Wv, which enters linearly into the output
initializing with the rectangular identity preserves the scale of the input representations (XWy X), keeps x(Wv) = 1. and avoids unnecessary distortion of the Jacobian. In contrast, the query and key maps interact bilinearly through
S = XW"
and initializing them as rectangular identities can bias projections toward coordinate subspaces, yielding anisotropic logits and unstable softmax dynamics. A semi-orthogonal initialization for Wo and Wx k instead provides ncar-isometric embeddings, giving each head balanced representations of X and supporting more diverse and stable attention patterns.
Following the above design principle, in the experiments of section initialized the value matrices Wy for each head as rectangular identities. For the queries and keys, we initialized each wx()) and w in the i-th head with independent semi orth-onthogonal projections,
Implementation.
(w ( y Ias
y Ia-
for = 1,..., ......,, h. where h is the number of heads. This produces near-isometric embeddings into distinct subspaces, thereby diversifying the logits 8(0 attention patterns. See section A. for concrete way to carry out this procedure. We refer to this initialization as conditioned initialization.
and the resulting
Different forms of attention. The formulation section describes the classical self-attention layer used in Transformers. In practice, many recent architectures have proposed variations of atten- tion to improve efficiency and effectiveness (Touvron et al. 2021| Ali et al. 2022 Liu et al. |021 Ding et al. 2022 Xiong et al. 2021b). In particular, many of these newer forms of attention apply normalization to the query, key and values. Our conditioned initialization readily applicable to these generalized forms of attention, including those with normalization (see Heenry al. (2020);
Dehghani et al. (2023); [Zhang & Sennrich (201)), and we empirically demonstrate in section that consistently yields strong performance.
EXPERIMENTS
In this section, we evaluate the theoretical insights from section across range of Transformer applications. For each setting, we compare our conditioned initialization against the standard de- fault schemes commonly used in the literature, well more recent alternatives such mimetic
5

---

Published as a conference paper at ICLR 2026
Vit-B during training
XCiT-M during training
Average condition number
2012
Average condition number
2031
2014
10a
10a
10"
107
Theoretical
Theonetical bound
Delauk
10a
10f
Mimesc
Conditiored (ours)
Conditioned (ours)
10a
10a
100
150
200
250
50
100
150
200
250
300
Epochs
Epochs
Figure I: Average condition number of the attention Jacobian during training under three common initialization schemes, shown alongside the theoretical bound from eq. (8)
initialization (Trockman & Kolter 2023) and weight selection (Xu et al. 2023). In all experiments, conditioned initialization is implemented as described section 3.3 The goal is to test our initial- ization wide range of architectures at different parameter levels. For larger scale experiments we refer the reader to section A.4
4.1 VITS FOR IMAGE CLASSIFICATION
Vision Transformers. We applied conditioned initialization to the attention layer of a variety of modern vision Transformers: ViT-Base (ViT-B) (Dosovitskiy al., 2020), Swin-Base (Swin-B)
(Liu et al. 2021). a XCiT-Medium (XCIT-M) (Ali et al. 2021). a DeiT-Base (DeiT-B) (Touvron
al. 202 and DaViT-Base (DaViTBB) (Ding al. 2022 for image classification on ImageNet- lk. Each model is initialized in three ways: (i) the default truncated normal initialization (Hugging Kolter 2023 and (iii) our conditioned ini- tialization from section 3.3 We point out that each of these vision Transformers uses different attention layers to the standard self self-attention used in the VYT-B architecture. However, as mentioned in section 3.3 conditioned initialization applies to more general forms of attention.
Face 2025b), (ii) mimetic initialization (Trockman
Results on ImageNet-1k. We trained three versions of each Transformer: a baseline model with default initialization (Hugging Face 2025b), one with mimetic initialization (Trockman Kolter, 2023) and one incorporating conditioned initialization, see section A.3. for training details. The final results, summarized table [ show the test accuracy for each initialization on each model. We observe that in every case, conditioned initialization outperforms the other two.
Validating the theory. We validate the theoretical results of section 3 using ViT-B and XCIT-M models trained on ImageNet-Ik. Figure reports the average condition number of the Jacobian of the attention matrices for ViT-B (left) and XCTT-M (right), during training, under the above men- tioned three initializations, alongside the theoretical upper bound from theorem 3.1 The results demonstrate that conditioned initialization consistently yields better-conditioned Jacobian, pro- viding empirical support for its role in enabling more stable attention mechanisms.
Table I: 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.
ViT-B
DefT-B
Swin-B
XCTT-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

---

Published as a conference paper at ICLR 2026
140
Swin-B
391
394
VIT-B
316
DeiTB
221
288
Eg
206
200
Dehutt
Winet: andttireddee frand
Sehatt
Minett Grdditoon Juant
Dehuut
Wiret: Centileett loun)
Delutt
Miedt:: Conditoent (oun))
Figure 2: Total number of epochs required for each initialization to reach the final accuracy of the default initialization reported in table 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.
Optimization efficiency. As shown in table conditioned initialization achieves higher accuracy under the same training regime compared both default and mimetic initialization. To further assess training efficiency, we measured the number of epochs required by mimetic and conditioned initialization match the final accuracy attained by the default initialization across ViT-B, DeiT-B, Swin-B, and XCIT-M. Figure 2 reports these results, demonstrating that conditioned initialization consistently converges 20-30% faster the same accuracy level. Similar analysis for DaViT-B is given in section A.3.
4.1.1 SMALL SCALE DATASETS
Following section 4. we assessed conditioned initialization on small-scale image classification
datasets: Flowers Nilsback & Zisserman, 2008). Pets (Vedaldi) 2012), CIFAR-10. and CIFAR-
100 (Krizhevsky et al. 2009). Experiments were conducted with the ViT-Tiny (ViT-T) architecture,
common choice for such benchmarks (Trockman & Kolter 2023 Xu al. 2023). We compared conditioned initialization with the default truncated normal scheme Hugging Face (025b) and with mimetic initialization (Trockman & Kolter, (203). Because ViTs lack strong inductive bias. they typically perform poorly on small datasets; mimetic initialization has been shown to partially remedy this by providing a more suitable inductive bias. As shown in table 2 conditioned initialization reliably improves over the default and performs on par with mimetic initialization.
Weight selection.
We compared our initialization with the weight selection method from Xu et al. 2023 This can be found in section A.3.1
Table 2: Comparison of ViT-T pretrained on four small scale datasets with three different initializa- tions. We report Top- classification accuracy. each case, conditioned initialization improves performance over the default and mimetic initializations.
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
4.2 OBJECT DETECTION AND INSTANCE SEGMENTATION
In this section, test conditioned initialization in fine-tuning setting two downstream tasks: object detection and instance segmentation. We first pretrain an XCTT architecture (Ali et al.) (02T) on ImageNet- K and then fine-tune COCO 2017 (Lin al. 2014). The XCTT models used

---

Published as a conference paper at ICLR 2026
classification
Listops
Textt diseeFication
condition number
LutOps
Epochs
Epochs
Figure 3: Average condition number of the attention Jacobian during training under three common initialization schemes, shown alongside the theoretical bound from eq. (8)
as backbones within Mask R-CNN framework (He et al. 2017) equipped with Feature Pyramid Network (FPN) for multi-scale features. To connect XCiT to FPN, we adapt its column structure to output intermediate representations, using the 12-layer XCiT-Small (XCIT-S) with strides adjusted from 16 to [4, 8, 16, 32) for FPN compatibility. Downsampling is done with max pooling and up- sampling with single transposed convolution. We evaluate this setup across three initializations: default truncated normal, mimetic, and our conditioned initialization.
We Results. The results for object detection and instance segmentation are presented in table 3 report AP (Average Precision for bounding boxes), AP 30/75 (Average Precision at IoU thresholds of 0.50 and 0.75 for bounding boxes), AP" (Average Precision for masks), and 505775 (Average Precision IoU thresholds of 0.50 and 0.75 for masks). Across all metries, XCIT models with conditioned initialization consistently outperform both alternatives.
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
App
AP
50
APł
Apm
APm
APm
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
4.3 LONG RANGE SEQUENCE MODELING
Long-range sequences are essential for Transformers, enabling the integration of information across distant tokens. We assess our initialization scheme on the Long-Range Arena (LRA) bench- mark (Tay et al. 2020). designed to evaluate models extended inputs. For this, we use the Nystrômformer (Xiong al. (02Tb), which achieves efficient long-range modeling via near-linear attention. We train three variants:
with default truncated normal initialization (Xiong et al., 2021b Hugging Face 2025a), one with mimetic initialization (Trockman Kolter, 2023), and one with our conditioned initialization from section 3.3 following the setup of Xiong et al. b)
Results. From table we see that across all tasks in the LRA benchmark, conditioned initializa- tion consistently outperforms both default and mimetic initialization. We validate the theoretical results of section the text classification and ListOps task in the LRA benchmark suite. Figure reports the average condition number of the Jacobian of the attention matrices for a Nyströmformer on the text classification task (left) and Nyströmformer ing, under the above mentioned three initializations, alongside the theoretical upper bound from theorem 3.1 The results demonstrate that conditioned initialization consistently yields a better- conditioned Jacobian, providing empirical support for its role enabling more stable attention mechanisms. Furthermore, measured the number of epochs required by mimetic and condi-
the ListOps task (middle), during train-
8

---

Published as a conference paper at ICLR 2026
tioned initialization match the final accuracy attained by the default initialization. Figure (right) reports these results demonstrating that conditioned initialization consistently converges approxi- mately 25% faster to the same accuracy level.
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.
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
4.4 LANGUAGE MODELING
We apply the insights from section to the Crammed BERT language model (Geiping & Goldstein 2023), trained with masked language modeling. We consider three variants: the original model with default normal initialization (μr = 0, 0.02), one with mimetic initialization, and one with our conditioned initialization. All models are pretrained on The Pile (Gao et al., (02I]) following the setup of Geiping & Goldstein (2023). and evaluated the GLUE benchmark (Wang et al. 2018). As shown in table 5 conditioned initialization yields the strongest performance, with mimetic also outperforming the default baseline. Additional results for GPT-2 are provided in section A.3.4
Table We evaluate a pretrained Crammed BERT with three different initialization schemes on the GLUE benchmark, and find that conditioned initialization outperforms the other two.
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
LIMITATIONS
Our conditioned initialization derived by optimizing an upper bound on the condition number of the self-attention Jacobian, rather than minimizing the Jacobian's condition number directly. The motivation was to examine whether such a bound-based initialization could induce more favorable optimization bias for training Transformers. While this approach offers useful theoretical guidance and consistent with empirical gains we observe, remains an indirect proxy. Developing methods that can efficiently estimate and control the exact Jacobian conditioning during training would therefore be a valuable direction for future work.
6 CONCLUSION
In this paper, introduced theoretical framework that relates the conditioning of self-attention Ja- cobians to the spectral properties of the query, key, and value matrices in Transformer architectures. Building this insight, we proposed simple initialization strategy, conditioned initialization, that aims to reduce the condition number of the attention Jacobian at initialization, thereby providing favorable inductive bias for optimization. Extensive experiments show that this approach consistently improves performance across a wide range of Transformer models and tasks, including image classification, object detection, language modeling, and long-range sequence learning.
Digital tools were used for grammar and formatting only. No large language models contributed to the research, and all findings are original work by the authors.

---

Published as a conference paper at ICLR 2026
REFERENCES
Naman Agarwal, Pranjal Awasthi, and Satyen Kale. A deep conditioning treatment of neural net- works. In Algorithmic Learning Theory, pp. 249-305. PMLR, 2021.
Alaaeldin Ali, Hugo Touvron, Mathilde Caron, Piotr Bojanowski, Matthijs Douze, Armand Joulin, Ivan Laptev, Natalia Neverova, Gabriel Synnaeve, Jakob Verbeek, et al. Xcit: Cross-covariane image transformers. Advances in neural information processing systems, 34:20014-20022) 2021.
Lukas Bossard, Matthieu Guillaumin, and Luc Van Gool. Food- food-101-an 100-1001-mminng discriminative compo- nents with random forests. In European conference on computer vision, pp. 446-461. Springer, 2014.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zogoruyko. End-to-end object detection with transformers. In European conference on computer vision, pp. 213-229 Springer, 2020.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski. Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International conference on machine learning, pp. 7480-7512. PMLR, 2023.
Mingyu Ding, Bin Xiao, Noel Codella, Ping Luo, Jingdong Wang, and Lu Yuan. Davit: Dual attention vision transformers. In European conference computer vision, pp. 74-92. Springer, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprint (xXii:23000007559 2023.
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 arXiv:2101.00027, 2021.
Jonas Geiping. Cramming.
htwp:/thub..coo/o/onoaaaaaaa wiga:///ooaaaeeiiingggggoooo/e piram/rannniiiiiiimmimmmmmmm.mia cerammiaannn ng
2023.
Jonas Geiping and Tom Goldstein. Cramming: Training a language model on a single gpu one day. In International Conference on Machine Learning, pp. 11117-11143, PMLR, 2023.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256. JMLR Workshop and Conference Proceedings, 2010.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In Proceedings of the IEEE international conference on computer vision, pp. 2961-2969, 2017.
Alex Henry, Prudhvi Raj Dachapally, Shubham Shantaram Pawar, and Yuxuan Chen. Query- normalization for tranfformers. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 4246 2020.
Hugging Face. Hugging face transformers nystromformer documentation. https: // huggingface co/ wtt.aa/nooroeeeeeee/oeeeeeeedeed Accessed: September 2025.
2025a.
10

---

Published as a conference paper at ICLR 2026
Hugging Face. Pytorch image models (timm). bctps://qithue wwt. com/hugginngaace/ myytorch-imgg-moonn models 2025b. Version 1.0.20 (or whichever you used); Apache-2.0 Li-
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel:
Convergence and gen- eralization in neural networks. Advances in neural information processing systems, 31, 2018.
Yiping Ji, Hemanth Saratchandran, Peyman Moghadam, and Simon Lucey. Always skip attention. arXiv preprint (xXi::5505009996 2025.
Alex Krizhevsky, Geoffrey Hinton, Learning multiple layers of features from tiny im- a2es. 2009.
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 740-755. Springer, 2014.
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, 19:85-116 16, 2022.
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 computer vision, pp. 10012-10022, 2021.
Jan R Magnus and Heinz Neudecker. Matrix differential calculus with applications in statistics and economelrics. John Wiley & Sons, 2019.
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, pp.
2023.
Matterport. Mask r-cnn implementation, 2017. httphttpss//ttthb.cm///aatteeppp
Mask RCNN
Maria Eleia- Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. In 2008 Sixth Indian conference computer vision, graphics & image processing, pp. 722-729. IEEE, 2008
Simon JD Prince.
Understanding deep learning. MIT press, 2023.
Xianbiao Qi, Yelin He, Jiaquan Ye, Chun-Guang Li, Bojia Zi, Xili Dai, Qin Zou, and Rong Xiao. Taming transformer without using learning rate warmup. arXiv preprint arXiv:2505.21910, 2025.
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.
Hemanth Saratchandran, Thomas X Wang. and Simon Lucey. Weight conditioning for smooth optimization of neural networks. In European Conference Computer Vision, pp. 310-325.
Springer, 2025.
A Steiner, A Kolesnikov, X Zhai, R Wightman, J Uszkoreit, and L Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arxiv 2021. arXiv preprint arXiv:2106.10270, 5, 2016.
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:2011.044
2020.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data data-efficient image transformers & distillation through attention. In International conference machine learning, pp. 10347-10357. PMLR, 2021.

---

Published as a conference paper at ICLR 2026
Asher Trockman and J Zico Kolter. Mimetic initialization of self-aatention layers. In International Conference on Machine Learning, pp. 94456-34468 PMLR, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Andrea Vedaldi. Cats and dogs. In Proceedings of the 2012 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 3498-3505, 2012.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint ar orXiv:/80..07461. 2018.
Ross Wightman. Pytorch image models.
https://githab.. com wwt.w//wwww/aaaa..a.c
pytorch- www.mmaaaaeeeeeee 2019.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Fei Tan, Glenn Fung, Vikas Singh, Xi- aodong Yuan, Sungsoo Ahn Wang, Dimitris Papailiopoulos, and Katerina Fragkiadaki. Github repository, 2021a.
com. pen
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, pp. 14138-14148, 2021b.
Zhiqiu Xu, Yanjie Chen, Kirill Vishniakov, Yida Yin, Zhiqiang Shen, Trevor Darrell, Lingjie Liu, and Zhuang Liu. Initializing models with larger ones. arXiv preprint arXiv:2311.1882. 2023.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006. 14548, 2020.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in neural infor- mation processing systems, 32, 2019.
Zhen, W Sun, H Deng, Li. Wei. B Lv, J Yan, L Kong, and Zhong. cosformer: rethinking softmax in attention. In International Conference on Learning Representations, 2022.
Liu Zhuang, Lin Wayne, Shi Ya, and Zhao Jun. A robustly optimized bert pre-training approach with post-training. Proceedings of the 20th chinese national conference on computational linguistics, pp. 1218-120 1227, 2021.
12

---

Published as a conference paper at ICLR 2026
A APPENDIX
ETHICS STATEMENT
This work relies solely on publicly available datasets and does not involve human subjects, person- ally identifiable information, or sensitive data. The proposed methods are developed exclusively to advance fundamental research in machine learning.
REPRODUCIBILITY STATEMENT
All experiments in this work were designed with reproducibility in mind. References are provided for any external codebases employed, and full details of training protocols and hardware are de- scribed in the appendix. Complete proofs of all theoretical results are also included allow inde- pendent verification.
USE OF LLMS
This manuscript was prepared with the assistance of LLMs for grammar checking only. No language models were used conducting the research or drafting the scientific content.
A.1 THEORETICAL ANALYSIS
In this section, give the proofs of the lemma and theorems from section 3.2 and the proposition from section 3.3
Notation. For the convenience of the reader we restate the notation we used section
Given a matrix Z E mR we denote the vectorization of Z by vec(Z) E [maxx1 Magguss & Neudecker 2019). Note that for such matrix there transformation Tmn R.maxmen such that Tnn.Vec Z) vec( (TT) where z" denotes the transpose of Z. The matrix Tmn is known commutation matrix and is a permutation matrix (Magnus & Neudecker 2019). The maximum sin- gular value of matrix Z will be denoted by mmax (Z) and the minimum singular value by amin (Z). We will use the standard terminology SVD to denote the singular value decomposition of matrix. Given vector
R" the notation ||0||| will denote the vector 2-norm of e. Finally, will let Imxn denote the rectangular identity matrix that has all I's on its main diagonal. When we are dealing with square matrices will often simply write J, with the understanding that I is the n x n square identity matrix. Context will make it clear whether we are in the square or non-square regime. We also let Omxn as the real m n semi-orthogonal matrices.
We start with the following standard facts on derivatives of matrices. We point out to the reader that Qi et al. (025) also computes various Jacobians of the self-attention layer using Kronecker factorizations however their computations are slightly different to ours and we therefore give explicit details of how obtain the proofs of lemma 3.1 and proposition 3.1
Lemma A.I. Let A Rexm, B E R and C B Then
θACB
ac BT & A.
(13)
Proof. We start by using a well known vectorization identity (Magnus & Neudecker 2019)
vec(ACB) = (B" @ A))ec(C))
(14)
where vec denotes the vectorization operator which takes matrix and maps to a vector by stacking its columns on top of each other, see Magnus & Neudecker (2019). We then differentiate eq. (14) to obtain
@vec (Pec(ACB) BT A Ovec(C))
(15)
The result of the lemma follows.
n
13

---

Published as a conference paper at ICLR 2026
Lemma A.2. Let A E R that A E Rmxx. Then
voc(A") T mn veo( A)
(16)
Öveo( (AT) @vec(A)
mm
(17)
where Tmn is commutation matrix.
Proof. The first equation follows from the definition of the transpose of a matrix Magnus & Neudecker 2019). The second equation then follows from the first.
The proof lemma 3.1 is given as follows.
Proof of lemma 3.1 Let (. [n] denote row vector in R" Then by definition
e
(e1..ax((2.......
(18)
Di=
2i=1
From the above equation we can compute the partial derivative and find
softmax(2): softmax(=)) (6ig softmax(=)))
(19)
Zzy
The term 2 @sofmmax()); ij terms into an n x n matrix we find
precisely the ij component of the matrix @softmax(s)
Putting each of these
Bay
82
@softmax(2) Diag(z) x
(20)
Oz
which proves the lemma.
We can use the above lemmas to give the proof of proposition 3.1
Proof of proposition 3.1 We start by establishing the derivative formula for the term
8A(X).
We
BWQ
will use the notation used section 3.1 Note that by definition aA(X) RdNNxD Using eq.
aWQ
A(X) Insoftmax(X WWW
(21)
where IN is the N x N identity matrix. This is done so that we can apply lemma A. compute
We then
OA(X) (rnssottaxxxWooolo xx)
(22)
OWq
dWq
#softmax(XWqll W x
-WTXx @ IN)
using temma A.1
(23)
DWq
W(X (
= =WWXx" @ (x)sssotmmax(WWWWW
using lemma
3.1
awQ
(24)
= =WWXX" 8 IN) [x))(ssftmxx(XXW W X
(X 8 X)
(25)
which proves the firs equality in proposition 3.1
To compute 8A(X) E RdNN proceed in similar way.
oW*
OA(X) @L
(26)
aWk
aW* Ōsoftmax (XWc WT
= -WTXT @ IN)
using lemaa A.1
(27)
0W
=wW xT & Tx) ( a(XWoWF
using lemma 3.1
OWk
(28)
=(WFXT @ (s))sssfma (X & XWQ)
(29)
14

---

Published as a conference paper at ICLR 2026
where the last equality follows from lemmas A.1 and A.2 proposition 3.1
This establishes the second equality in
To prove the identity for 8A(X)
E R RdNxxD xdD
we write
BWy
(
X
(30)
where La is the d x d identity matrix. Then we simply apply lemma A.1
to obtain
aA(X) D[softmax(XWaW/ XWVvTa))
(31)
dwv
awv
La softmax(X (W W X
(32)
which proves the final equality in proposition 3.1
We also give the proof of theorem 3.1
Proof of theorem 3.7 The proof of theorem
follows from using proposition and the defini-
tion of the Jacobian of A(X) with respect to WQ. Wi and Wv given by
J(A(X)) [((A()) ((A(X)) B(A(K))
(33)
aW OWK
Owv
We recall that the condition number defined as x(J(A(X)))
Omin (J(A(X))) where max (J)A(X)) is the maximum singular value of (((()) and a (A(X))) the minimum
Trowe (J(A()))
singular value which we know is non-zero because of the assumption that (A(X)) has full rank.
Note that, using the notation section 3.1 we have that J(A(X) E Tg3dNxDD as 0A(X) AA(X)
aWq JWx
0A(X) E gddNNel xdD For each of notation we will write Aq 8A(X), aWv
BWq AK 8A(X)
Av 8Wv
8A(X).
OWN
We will start by computing bound for the maximum singular value. We have for any vector
E RRD we have
2
((A(X)) ((A())) 8(A(X))
[A((X))
awQ (z), OWK (z),
Owr (2)
(34)
Av())]] |AQ(2)]12 + [|Ax(=))l [|Av(=)112
(35)
(36)
(omax (AQ)² Tmax (Ax)² )2 + mmxx (Av)²) ||=1|2.
(37)
This implies that
max (J(AX))) = max |](A())=
(38)
270 |
/ mmax (AQ)2 + m (Ax) + (n
(39)
(max((A) o (Ax) Tmax (Av).
(40)
The next step is compute lower bound for the minimum singular value o miin (J(A(X))) approach is similar to the above, using the fact that min (J(((X))) = min |= ((A(X can then use the inequality
The
We
O(A(X)) O(A(x))
(z), O(A(X))
(z),
(z)
(41)
aWq
OWK
OWv
((A(X)
((A()))
|| ((A(X))
(2)
(42)
aw
OWk
dwr
Then minimizing the above over the constraint z E RdD such that ||||| we obtain
((A(X))
(((A(X))
((((X)))
Owr )
min (J((X))) 2 max
mmin
aWe
min
(Ww
Fmin
(43)
15

---

Published as a conference paper at ICLR 2026
Combing the bounds max
m (J(A(X))) and minn (J(A(X))) we obtain ((((x))) +
(0(A(X))
(OLACX))
x(J(A(X))) S
+
(44)
aWa
aWK
OWv
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 D such that the product CD is full rank then x(CD) I x(C))(() and the second that x(C² @ D) = <(C")(()) Using these two facts we can then use proposition 3.1 obtain the bound of theorem and the proof is finished.
Using theorem
3.1 the proof of proposition straightforward.
We first observe that by definition of the condition number of a general m x n matrix M must always satisfy ss(M) 2 1. For matrices with I's on the diagonal and zero elsewhere the condition number 1. If M is semi-orthogonal matrix then either
Proof of proposition 3.2
1. M Imxm m if m 2 n
(45)
2. T M Inxo x if n 2 m.
(46)
Since the singular values of M are precisely the eigenvalues of MM i 2 or M M if 2 m. It follows that the singular values of M must all be and hence it has condition number 1.
Therefore, if the Wo. Wx and Wy matrices are initialized as Dxd or as a matrix in Oxa we have that their condition number 1. Therefore, we must have
B(A) s ((A)
(47)
and the proposition is proved.
Remark A. In the implementation strategy in section 3.2 we saw that we initialized the values matrix Wy with rectangular identity matrix I Dxd- However, from the theory of that section we could have also initialized with X d for X + O. In general, we found empirically that this could be done but that if À got too large we noticed some instability in training due to Wy having weights that were too large. Therefore. opting for the identity D was what we found worked well.
Why Conditioning the Jacobian Aids Optimization. The rationale for improving the condition- ing of the self-attention Jacobian is connected to established results on the Neural Tangent Kernel (NTK). Prior work, see Liu et al. (2022), has shown that conditiond NTKs lead to faster and more reliable convergence to a global minima during gradient-basse optimization. Since the singular values of network's ' Jacobian correspond to the positive square roots of the eigenvalues of its NTK, improving the conditioning of the Jacobian directly enhances the conditioning of the NTK. This connection provides theoretical basis for why controlling the spectral structure of the self-attention Jacobian can benefit optimization. Furthermore, recent extensions of NTK theory to transformer architectures (e.g., Yang 2020) support the relevance of these insights in the attention setting. Our initialization scheme leverages this relationship by explicitly targeting improved con- ditioning at initialization, which consistent with the empirical performance gains observed across tasks.
A.1.1 IMPLEMENTATION DETAILS
In this section, we give the details of how we implement the seti-orthogonal initialization of the We and Wy matrix of each head.
We recall from section 3.3 that our initialization for the queries and keys, proceeded by initializing each Wi and wi( in the i-th head with independent semi- ortioorthooona projections.
( = La,
W Las
for 1, 1..... h where h is the number of heads. To do this suppose each T( and w(O)) are D x d and let min(D) D, d) for each head we form two random matrices R. and R. of shape D) x d and then take the truncated SVD
(1( (r) and UK (r)sk (r)(v:(F(())
(48)
16

---

Published as a conference paper at ICLR 2026
where each of the 0,9()) E R pxr and (V,))r) Rxxd and similarly for the K ones. We then observe that 0 = (f) (r) and O == U (r) (r)
(49)
are semi- orthogooolo....... Doing this for each different head we get different semi-orthogonas matrices for each We and Wx for each head. The implementation of Wr is the same for each head and simply done by fixing Wv = /pxd- Remark A.2. We note that in the above we used the SVD to obtain the initializations of each Wq and Wy. One can also use the QR decomposition (Magnus & Neudecker 2019) and PyTorch has a built in way do this via nn. init. orthogonal.
A.2 DISCUSSION ON NORMALIZATION
This section provides additional clarification on how normalization and stabilization mechanisms interact with the Jacobian conditioning analysis and with the proposed conditioned initialization scheme. Several modern transformer architectures apply normalization directly to the queries, keys, or values (e.g.. RMSNorm, QKNorm), while Layer Normalization (LN) typically applied to the inputs of each block. The discussion below outlines how these components relate to the theoretical results in section and to the practical behaviour observed in section
Theoretical setting. In the original self-attention formulation Vaswani et al. (2017), no normaliza- tion is applied directly to the query, key, or value weight matrices. The only modification the fixed scaling factor 1/vd, determined by the head dimension. Since this factor does not depend on the model parameters, differentiating the attention map simply scales the Jacobian by constant. A con- stant scalar multiplication does not change the condition number of a matrix; therefore, this scaling has effect the conditioning analysis developed in Section 3. The derivations for vanilla self- attention thus remain mathematically correct and serve as a foundation for the initialization scheme proposed later.
Compatibility with modern normalization layers. Many contemporary transformer variants in- corporate normalization directly into the attention pathway, for example through RMSNorm or QKNorm applied to queries, keys, or values. The conditioned initialization introduced in this paper is designed operate top of these mechanisms rather than replacement for them. In all experiments, the architectures retain their original normalization layers exactly as implemented in the publicly released codebases. For example, the DeiT-B model applies QKNorm both queries and keys. When applying conditioned initialization, these QKNorm layers remain in place. This ensures that comparisons in section 4 are consistent with the established baseline implementations in the literature.
Normalization versus conditioning. Normalization and conditioning address different aspects of stability. Normalization controls the scale of activations and gradients, whereas conditioning con- cerns the spectral structure of the Jacobian. These effects are distinct. A matrix may have small norm but poor conditioning, or it may have large norm yet be well-conditioned. For example,
07
10.1
A =
B
0
1
10
have very different Frobenius norms but nearly identical condition numbers, while
1 0 C = 0 1
2
0
D=
0.1
have similar Frobenius norms but condition numbers 1 and approximately 14.14, respectively. Nor- malization schemes such as LN, RMSNorm, and QKNorm regulate magnitudes, while conditioned initialization directly shapes the spectral behaviour of the Jacobian. These mechanisms therefore complement one another.
Ablation Partial Initialization. We conducted several ablation studies in which the condi- tioned initialization was applied to only a subset of the attention projection matrices, such as ini- tializing only Wo. only Wx or only Wv. An initial hypothesis that initializing Wy alone
17

---

Published as a conference paper at ICLR 2026
might be sufficient, since the softmax operation effectively performs row-wise normalization the attention scores. However, normalization and conditioning target fundamentally different prop- erties: matrix may have small norm yet still exhibit large condition number. This observation led us to a more complete analysis, formalized in theorem 3.1 which shows that the conditioning of the self-attention Jacobian depends the joint spectral structure of Wq. W K, and Wy. Consistent with this theoretical insight, we found that applying conditioned initialization to all three projections yields the most stable behaviour and the strongest empirical performance.
Discussion the Output Projection. We also examined whether the conditioned initialization should applied to the output projection matrix Wo. given its close relationship to Wv. Empir- ically, found that once Wy is initialized using the proposed scheme, the initialization of Wo has negligible effect. Applying the conditioned initialization to Wo in addition to Wy did not al- ter training behaviour final performance. Consequently, our method focuses on conditioning the query, key, and value projections, while leaving Wo with its standard initialization.
A.3 EXPERIMENTS
A.3.1 VISION TRANSFORMERS
The image classification experiments in section 4.1 of the paper were done on Nvidia A100 GPUs. The implementation of the ViTs was all done using the Timm code base (Wightman 2019). The architectures were all trained from scratch on the ImageNet- Ik dataset using the AdamW W optimizer following the hyperparameters used the original papers (Dosovisskiy al.] [020 Steiner et al.] 22066; Liu et al.) 2021; Ali et al. 2021 Touvron et al. 2021 Ding et al.] 2022). For the case of VIT-T on the Pets, Flowers, CIFAR-10 and CIFAR-
Hardware and implementation.
datasets used Wightman (2019) for the architecture and the training hyperparameters from Xu ct al. 2023).
Optimization analysis for DaVIT-B. shown table [ conditioned initialization achieves higher accuracy under the same training regime compared to both default and mimetic initializa- tion on DaViT-B architecture. epochs required by mimetic and conditioned initialization to match the final accuracy attained by the default initialization across on the DaViT-B architecture. Figure reports these results, demonstrat- ing that conditioned initialization consistently converges approximately 25% faster the same accuracy level.
further assess training efficiency, we measured the number of
DaViT-B
293
250
221
Epochs
200
150
100
50
Default
Mimetic
Conditioned (ours)
Figure 4: Total number of epochs required for each initialization to reach the final accuracy of the default initialization reported in table [ for DaVIT-B architecture. Conditioned initialization converges more quickly, requiring fewer epochs than both default and mimetic initialization.
18

---

Published as a conference paper at ICLR 2026
Weight Selection Xu al. (Xu et al. 2023 introduced the weight selection method, showing that weights transferred from an ImageNet-2IK-poe preetrrine ViT-Small (ViT-S) model provide a strong in- ductive bias for initializing a ViT-Tiny (ViT-T) small-scale datasets. Here, examine whether weight selection remains effective when pretraining is performed on ImageNet-1K: To this end, pretrained ViT-S models ImageNet-1K under three initialization schemes: default truncated normal (Hugging Face 2025b), mimetic (Trockman & Kolter, 2023), and our conditioned initial- ization from section 3.3 Using the weight selection procedure (Xu et al. 2 we then derived corresponding ViT-T initializations and trained them on Food-101 (Bossard et al. 2014), CIFAR- 10, and CIFAR-100 (Krizhevsky et al. 2009). Table summarizes the results. Each entry labeled "ImageNet-IKK I indicates that the ViT-S model was pretrained on ImageNet-1K with the specified initialization before its weights were transferred to ViT-T via weight selection. For comparison, the final row reports the performance of VIT-T initialized from an ImageNet-21 K-pretrineee ViT-S with default truncated normal initialization. The results highlight key finding: replacing ImageNet- 21K with ImageNet-In IK can yield comparable performance, provided that the initialization is chosen carefully. In particular, conditioned initialization on ImageNet- 1 K achieves accuracy on par with de- fault initialization pretrained on ImageNet-21K, underscoring its effectiveness in data-limited pre- training regimes.
Table Comparison of ViT-T pretrained on four small scale datasets with three different initializa- tions. We report Top-1% 1% classification accuracy. In each case, conditioned initialization improves performance over the default and mimetic initializations.
Food-101
CIFAR-10
CIFAR- 100
ImageNet- 1k Default
85.5
96.6
79.7
ImageNet-lk: + Mimetic
86.4
96.3
79.9
ImageNet- Ik Conditioned (ours)
87.3
97.1
81.0
ImageNet-2lk + Default
87.1
97.1
81.1
Training loss curves. We plot the training curves for each of the vision transformer experiments with each different initialization. fig. 5 fig. 6 fig. 7 we plot the training loss for the different initializations for the ViT-B, DeīT-B, Swin-B, XCIT-M and DaViT-B architectures respectively. In each case, we see conditioned initialization has trains stably and converges faster when compared to the default and Mimetic initializations.
Training Loss for VIT-B
Training Loss for DeVIT
Conditiore (ours)
Condl (ours)
Delaut
Defauk
LosS
3
Epochs
Epochs
Figure 5: Training loss curves for different initializations for ViT-B (left) and DeiT-B (right).
A.3.2 OBJECT DETECTION AND INSTANCE SEGMENTATION
The experiments for section 4.2 of the paper on object detection and instance segmentation were carried out Nvidia A100 GPUs The implementation followed He et (017)) We used the code base given by the GitHub Matterport (2017) following their exact training regime.
Hardware and Implementation:
19

---

Published as a conference paper at ICLR 2026
Training Lass for Swin-8
Training Lass for XCTTM
Conditfoned
Conditioned
LOss
1.00
200
250
300
350
50
150
200 Epochs
250
300
350
400
Epochs
Figure 6: Training loss curves for different initializations for Swin-B (left) and XCÝT-M (right).
Training Less for DEVIT
Conditioned
100
300
Epochs
Figure 7: Training loss curves for different initializations for DaViT-B.
A.3.3 LONG RANGE SEQUENCE MODELING WITH NYSTRÔMFORMER
Hardware and Implementation. All the experiments for the Nyströmformer LRA benchmark results in section 4.3 were carried out Nvidia A100 GPUs following the implementation and hyperparameter settings given in Xiong et (2021a).
A.3.4 LANGUAGE MODELING
Crammed BERT. The BERT language modeling experiment section 4.4 were all carried out on Nvidia A6000 GPU. The Crammed-Bert was implemented following the original paper Geiping Goldstein (2023) and the original GitHub Geiping (2023). The training regime follows Geiping (2023).
GPT-2 TinyStories. We train autoregressive Transformer, namely GPT-2 architecture trained on the TtoySeo dataset (Eldan & Li 2023). Once again we compared three initializations: one with the default normal initialization (μ 0, 0.02), one with mimetic and third with conditioned initialization. As shown in table conditioned initialization achieves a lower perplexity than both the default and mimetic initializations showing that boosts performance the setting of autoregressive Transformers. We used the training regime from (Eldan & Li 2023).
A.4 LARGER SCALE EXPERIMENTS
In this section we tested our initialization on much larger models in both the vision and language setting.
20

---

Published as a conference paper at ICLR 2026
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.
Perplexity
Default
2.47
Mimetic
2.40
Conditioned (ours)
2.29
We larger GPT-2 model, 472 million parameters, on the WikiText-103 dataset and the TinyS- tories dataset. The results are shown in table S and table 9 below clearly shows our initialization obtains better performance by obtaining the lowest perplexity.
Table 8: Perplexity of 472M GPT-2 model pretrained on TinyStories.
Model
Perplexity
Default
2.39
Mimetic
2.32
Conditioned (ours)
2.20
Table 9: Perplexity of 472M GPT-2 model pretrained on WikiText-103.
Model
Perplexity
Default
44.2
Mimetic
43.8
Conditioned (ours)
42.7
We repeated this experiment with 1.72 billion parameter GPT-2 model. In this witnessed overfitting in all initialized cases. However, this to be expected as a billion parameter model too large for the TinyStories and WikiText-IO3 datasets. Yet even in this case our initialization performed much better as can be seen from table I0]and and table [I
Table 10: Perplexity 1.72B GPT-2 model pretrained on TinyStories.
Model
Perplexity
Default
4.15
Mimetic
4.25
Conditioned (ours)
4.03
Table 11: Perplexity of 1.72B GPT-2 model pretrained on WikiTextt-003
Model
Perplexity
Default
48.1
Mimetic
48.5
Conditioned (ours)
46.9
We also ran experiments on large scale ViTs on the ImageNet- dataset, where these models range from 200-300M parameters. Once again for these larger models witnessed overfitting (this has also been observed in Dosovitskiy et al. (2020); Ding et al. (2022)) yet even this case our initial- ization outperformed the other two standard ones, table 12
We then ran two billion parameter models namely a ViT-Giant (ViT-G) and a DaViT-Giant (Da VTT- G) these where some of the standard billion parameter vision transformers could find in
21

---

Published as a conference paper at ICLR 2026
Table 12: Large scale ViTs with different initializations pretrained on ImageNet-IN Ik. Top1 % accuracy.
We show the
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
Hugging Face (202bb). We pretrained them each with the different initializations on ImageNet-Ie lk. As can be seen from table our our initialization out performs the other two.
Table 13: 1B scale ViTs with different initializations pretrained ImageNet-Ik: Top1 % accuracy.
We show the
VIT-G
DaViT-G
Original
78.8
81.9
Mimetic 2
78.6
82.0
Conditioned (ours)
79.9
82.9
22


Reviews:
[{"review_id": "VkKlNBW6nC", "reviewer": "Reviewer_ip1U", "summary": "The paper presents a spectrum inspired approach to initialize the attention matrices to reduce the conditioning number of the attention Jacobian matrices with partial theoretical analysis. Extensive experiments demonstrate that using the presented initialization approach leads to consistently superior performance over the baselines.", "questions": "- Please refer to the weaknesses.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "+ The paper proposes a simple but effective intialization approach, which might have a better condition number of the Jacobian of the attention matrix with respect to the parameter matrices $W_Q$, $W_K$ and $Q_V$.   \n+ Extensive experiments demonstrate consistent improvements over the two baseline methods.", "weaknesses": "- The results in Lemma 3.1 and Theorem 3.1 cannot be listed as the (theoretical) contributions of the paper. These results are merely the formal expression of the gradients of a matrix function with respect to the parameter matrices via the Kronecker product. Similar results can be found in prior work, e.g., the appendix in [a]. \n\n- The results for Jacobian matrices developed in the paper might not complete due to ignorance of the stablization structure, e.g., Layer Norm (LN), RMSNorm, QKNorm. Thus, it is also questionable whether the upper bound of the condition number makes any sense in practical. If either LN, or RMSNorm, or QKNorm is introduced, can the proposed approach still yield improved performance comparing to the counterpart baseline methods? \n\n- While the condition number of the proposed initialization strategy is reduced, it is merely a heuristic way to form the initialization for the attention matrices. Does it enable the training process stable? How about the effects of using the conditioned initialization on the learning curves?  In practice, when one of LN, or RMSNorm, or QKNorm or some combination of them is introduced, what aobut the learning curves? \n\n- In previous work, there are many attempts to design stablized optimization algorithm for training Transformer. It would be more interesting if some evaluations to connect the proposed initialization strategy with the stablized algorithms. The reviewer is curious about that whether or not the proposed initialization still works when stablized optimization algorithms (or stablized structure, e.g., LN, RMSN, QKNorm, etc.) used.\n\n\n[a] Taming Transformer Without Using Learning Rate Warmup, ICLR'25. \n\n[b] Learning deep transformer models for machine translation. arXiv preprint arXiv:1906.01787, 2019.\n\n[c] Query-key normalization for transformers. EMNLP 2020\n\n[d] Scaling vision transformers to 22 billion parameters. ICML 2023.\n\n[e] Root mean square layer normalization. NeurIPS 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3BaP2APKmQ", "reviewer": "Reviewer_sReQ", "summary": "This paper proposes conditioned initialization, an initialization scheme for attention weights that explicitly targets the conditioning of the attention Jacobian. The authors present a theoretical analysis showing that the condition number of Jacobian can be upper-bounded in terms of the condition numbers of the query, key, and value matrices. Making these matrices well-conditioned at initialization is expected to stabilize optimization. Extensive experiments on various downstream tasks demonstrate that this initialization improves final accuracy and accelerates convergence compared to baseline methods on small models. The approach is simple to apply and does not require changes to existing training pipelines.", "questions": "1. Does conditioned initialization continue to provide performance gains and faster convergence when hyperparameters such as learning rate, warmup steps, and weight decay are re-tuned for each baseline? Could baseline methods catch up with light hyperparameter tuning?\n2. Is the proposed method still effective for large-scale models?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper provides a clear theoretical motivation and analysis, deriving an explicit upper bound related to attention optimization stability. This is both novel and well justified.\n2. The paper is well written and easy to follow. The proposed method is straightforward to implement and architecture-agnostic.\n3. Experimental results across various downstream tasks are promising. The method consistently improves performance and accelerates convergence.", "weaknesses": "1. The theoretical analysis optimizes an upper bound on the condition number of the Jacobian rather than the condition number itself. Although empirical results support the approach, the gap between the bound and the actual condition number is not fully characterized. It remains unclear whether a tighter bound would yield further improvements.\n2. All experiments are conducted on relatively small models. It is unclear whether the benefits of conditioned initialization extend to large-scale models.\n3. There is a typo on line 869: “The implementation of the ViTs were” should be “The implementation of the ViTs was.”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sASb52klwU", "reviewer": "Reviewer_44Pd", "summary": "This paper proposes to initialize the query and key weights of self-attention layers as semi-orthogonal matrices, and the value weights as rectangular identity matrices. This scheme is based on an analysis of the condition number of the Jacobian of self-attention layers, supported by previous research suggesting that bounding this condition number leads to smoother/faster convergence.", "questions": "Did you try anything at the intersection of conditioned and mimetic initializations? For example, you could maybe pick a close orthogonal pair of matrices. Or you could use conditioned init for W_Q, W_K and mimetic init for W_V, W_O. \n\nDid you try any ablations? (Like only initializing W_Q, W_K or only initializing W_V?)\n\nDid you consider the output projection at all? It seems like it would be closely tied to W_V.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- Comprehensive empirical validation on a variety of vision tasks, showing clear positive benefit for large and small scale image classification, detection and segmentation\n- (Somewhat less) comprehensive empirical validation on language tasks, showing improved performance on LRA and on 100m-param-scale language modeling\n- Demonstrated benefit over mimetic initialization\n- Addresses the ViT small-scale-data issue as well or better than mimetic initialization while being similarly simple\n- Theoretically grounded\n- Works for a variety of attention mechanisms", "weaknesses": "- Relatively small-scale language model evaluation (e.g., experiments requiring one GPU while training ViT-B presumably took multiple GPUs). But I still find the results convincing and promising. I know it's hard to do this on an academic budget. And the vision results alone justify my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Wy698R8rgX", "reviewer": "Reviewer_jfhc", "summary": "This paper focuses on the initialization of transformers, specifically the initialization of the attention mechanism. Specifically, arguing that for better optimization of the network, the initialization should bound the condition number of the Jacobian matrix of the attention computation.\n\nThe paper suggests a practical implementation for this by initializing the $\\mathbf{W}_q$ and $\\mathbf{W}_k$ matrices as orthogonal matrices, and $\\mathbf{W}_v$ matrix as an identity one.\n\nThrough a series of experiments across language and vision tasks, the initialization is shown to achieve better performance, while achieving the same accuracy as baselines with significantly less compute, highlighting its efficiency.", "questions": "See W1", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "S1: The paper is clear and can be followed easily.\n\nS2: The practical implementation of the initialization is easy to use, making it conducive to wide adoption.\n\nS3: The method shows consistent performance gains compared to the baselines across different tasks and modalities (language and vision).\n\nS4: The suggested initialization method is highly efficient, achieving the same performance as baselines with much less compute. This is an advantage that could be highlighted further by the authors, perhaps through showing detailed training loss curves of the different models and adding more figures like Figure 2 for all tasks.", "weaknesses": "W1: The main weakness is that while the paper demonstrates the proposed initialization bounds the condition number of the Jacobian matrix of the attention computation, it does not sufficiently explain the theoretical basis for why this is a desired property for better optimization. Although better performance is shown in practice, a proof or a more robust theoretical explanation for this optimization benefit is expected, especially since the paper outlines how to achieve the bound.\n\nW2: This is a minor point, but using the notation $\\mathbf{A}(\\mathbf{X})$ for the output of the attention computation is potentially confusing, as this symbol is often reserved for the attention matrix itself. A symbol change is suggested for clarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)