PaperlayerView exact revision on arXiv

Source-linked AI summary

Conditioned Initialization for Attention

Hemanth Saratchandran, Simon Lucey

arXiv:2609.07086v1cs.LGcs.CV

TL;DR

Attention-weight initialization has received less attention than other aspects of Transformer optimization, despite the role of query, key, and value projections in attention. The paper proposes conditioned initialization to target attention's spectral conditioning and reports improved performance across a range of Transformer models and tasks.

  • Problem

    Transformer attention initialization has received comparatively little focus, despite the role of query, key, and value projections in attention-layer stability.

    Source: §1 · p. 1

  • Method

    Conditioned initialization uses the spectral properties of query, key, and value projections to target a lower upper bound on the attention Jacobian's condition number at initialization.

    Source: §1 · p. 2

  • Results

    Conditioned initialization consistently improves performance across a wide range of Transformer models and tasks, including image classification, object detection, language modeling, and long-range sequence learning.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The findings position principled conditioning of initialization as an underexplored component of Transformer design relevant to optimization stability and downstream performance.

    Source: §1 · p. 2

    The initialization optimizes an upper bound as an indirect proxy; efficiently estimating and controlling exact Jacobian conditioning during training remains future work.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Transformer research has focused more on attention efficiency, scalability, and expressiveness than on how attention weights are initialized. This paper introduces conditioned initialization to target attention-layer spectral conditioning and reports faster convergence and better generalization across diverse benchmarks.

1 INTRODUCTION

  • Initialization design is an underexplored question for attention, whose query, key, and value projections govern token interactions and optimization stability.Earlier Xavier and Kaiming schemes showed the importance of initialization for deep-network optimization, but Transformers have received less theoretical scrutiny.Source: §1 · p. 1
  • Common practice uses random initialization; mimetic initialization transfers trained-model statistics and weight selection reuses larger teachers' pretrained weights, but these alternatives remain heuristic.The cited alternatives lack a principled connection to attention conditioning.Source: §1 · pp. 1–2
  • Conditioned initialization targets attention's spectral properties because self-attention Jacobian conditioning depends on the query, key, and value projections.The method intervenes at initialization rather than modifying training objectives, providing an inductive bias for more stable optimization.Source: §1 · p. 2
  • Experiments across image classification, detection, segmentation, language modeling, and long-range sequence learning report faster convergence and improved generalization.The method is also reported as readily integrable across Transformer architectures.Source: §1 · p. 2
  • The paper frames principled conditioning as an underexplored component of Transformer design with potential relevance to optimization stability and downstream performance.This perspective centers initialization as a design choice rather than only a training-time concern.Source: §1 · p. 2

2 RELATED WORK

Prior Transformer initialization work ranges from standard random schemes to mimetic initialization and teacher-weight transfer. Other conditioning research studies weights, kernels, depth, skip connections, and architectural changes; this work instead asks whether initialization can improve attention conditioning from the outset.

2 RELATED WORK

  • Transformer initialization commonly uses normal or truncated-normal schemes, while mimetic initialization imitates trained-network statistics and weight selection transfers larger teachers' weights.The paper compares against the latter two approaches in its experiments.Source: §2 · p. 2
  • Conditioning studies link better-conditioned weights or neural tangent kernels to performance or convergence, and examine effects of depth, skip connections, and architectural modifications.The cited examples include NTK analyses and skip connections as an implicit conditioning mechanism in Transformers.Source: §2 · p. 2
  • Unlike work focused on training-time preconditioning or architectural changes, this paper investigates designing initialization for better-conditioned attention layers from the outset.The paper positions initialization itself as the focus of its contribution.Source: §2 · pp. 2–3

3 THEORETICAL FRAMEWORK

The framework links self-attention Jacobian conditioning to the spectral properties of the query, key, and value matrices. It proposes conditioned initialization to tighten a Jacobian condition-number bound through matrix-specific initialization choices.

3.2 MAIN THEOREMS

  • The condition number of the self-attention Jacobian depends on the condition numbers of the query, key, and value matrices, and low-condition-number initialization imparts an inductive bias toward more effective optimization.The analysis considers the Jacobian with respect to WQ, WK, and WV.Source: §3.2 · p. 3, §3.1 · p. 3
  • Theorem 3.1 bounds the attention Jacobian’s condition number, and controlling the condition numbers of WQ, WK, and WV tightens that upper bound.The bound is a surrogate: tightening it alone does not establish that the true Jacobian is better-conditioned.Source: §3.2 · p. 4, §3.3 · p. 5

3.3 CONDITIONED INITIALIZATION

  • Conditioned initialization uses rectangular identities for value matrices and semi-orthogonal projections for query and key matrices.The value identity preserves input scale and keeps κ(WV) = 1; semi-orthogonal query and key maps provide near-isometric embeddings and support more diverse, stable attention patterns.Source: §3.3 · p. 5
  • Query and key maps interact bilinearly, so rectangular identities can bias projections toward coordinate subspaces and yield anisotropic logits and unstable softmax dynamics.The proposed semi-orthogonal alternative provides near-isometric embeddings instead.Source: §3.3 · p. 5
  • Conditioned initialization is applicable to generalized attention forms, including those that normalize queries, keys, and values.The paper reports empirical demonstrations of strong performance for these forms.Source: §3.3 · p. 5

4 EXPERIMENTS

Experiments compare conditioned initialization with default and mimetic schemes across vision, detection and segmentation, long-range sequence modeling, and language modeling. Conditioned initialization improves reported performance across these settings and reaches matched accuracy faster in the reported image-classification comparisons.

4.1 VITS FOR IMAGE CLASSIFICATION

  • On ImageNet-1k, conditioned initialization improves Top-1 accuracy over default and mimetic initialization for every evaluated vision Transformer.The comparison covers ViT-B, Swin-B, XCiT-M, DeiT-B, and DaViT-B; across four architectures, it also reaches default initialization's final accuracy 20–30% faster.Source: §4.1 · pp. 6–7, Table 1 · p. 7, Figure 2 · p. 7

4.1.1 SMALL SCALE DATASETS

  • On Flowers, Pets, CIFAR-10, and CIFAR-100, experiments with ViT-T show conditioned initialization improves Top-1% classification accuracy over both default truncated-normal and mimetic initialization in each case.Source: §4.1.1 · p. 7, Table 2 · p. 7

4.2 OBJECT DETECTION AND INSTANCE SEGMENTATION

  • Conditioned initialization outperforms default and mimetic initialization across all reported COCO detection and instance-segmentation metrics.Table 3 covers bounding-box and mask average precision.Source: §4.2 · p. 8, Table 3 · p. 8

4.3 LONG RANGE SEQUENCE MODELING

  • On the LRA benchmark, conditioned initialization improves Nyströmformer evaluation accuracy across all tasks and yields a better-conditioned attention Jacobian on text classification and ListOps.It reaches the accuracy attained by default initialization approximately 25% faster.Source: §4.3 · pp. 8–9, Table 4 · p. 9

4.4 LANGUAGE MODELING

  • On GLUE, conditioned initialization yields the strongest performance among the three tested Crammed BERT initialization schemes.Mimetic initialization also outperforms the default baseline.Source: §4.4 · p. 9, Table 5 · p. 9

5 LIMITATIONS

Conditioned initialization optimizes an upper bound on the self-attention Jacobian’s condition number, not the condition number directly. The authors identify efficient estimation and control of exact Jacobian conditioning during training as future work.

5 LIMITATIONS

  • The method optimizes an upper bound on the self-attention Jacobian’s condition number rather than minimizing the condition number directly.The authors describe this bound-based target as an indirect proxy.Source: §5 · p. 9
  • Efficiently estimating and controlling exact Jacobian conditioning during training remains a direction for future work.The proposed approach provides theoretical guidance and is consistent with the observed empirical gains.Source: §5 · p. 9
  • The approach examines whether optimizing the bound can induce a more favorable optimization bias for Transformer training.The authors frame this as the motivation for using a bound-based initialization.Source: §5 · p. 9

6 CONCLUSION

The conclusion relates self-attention Jacobian conditioning to the spectra of the query, key, and value matrices, and presents conditioned initialization as a strategy to improve it. Experiments report consistent performance gains across several Transformer tasks.

6 CONCLUSION

  • Conditioned initialization aims to reduce the attention Jacobian’s condition number at initialization, providing a more favorable optimization bias.The framework relates Jacobian conditioning to the spectral properties of query, key, and value matrices.Source: §6 · p. 9
  • The theoretical framework connects self-attention Jacobian conditioning with the spectral properties of Transformer query, key, and value matrices.Source: §6 · p. 9
  • Experiments report consistent performance improvements across image classification, object detection, language modeling, and long-range sequence learning.The conclusion characterizes the evaluation as extensive and spanning a wide range of Transformer models and tasks.Source: §6 · pp. 9–10

ETHICS STATEMENT

The ethics statement says the work uses publicly available datasets and does not involve human subjects, personally identifiable information, or sensitive data. It characterizes the methods as fundamental machine-learning research.

ETHICS STATEMENT

REPRODUCIBILITY STATEMENT

The reproducibility statement says the authors cite external codebases and provide training-protocol and hardware details in the appendix. It also reports that complete proofs are included for independent verification.

REPRODUCIBILITY STATEMENT

USE OF LLMS

The manuscript states that LLMs assisted with grammar checking only, not research or scientific drafting.

USE OF LLMS

  • LLMs were used only for grammar checking, not to conduct the research or draft its scientific content.Source: USE OF LLMS · p. 14

A.1 THEORETICAL ANALYSIS

The appendix outlines proofs connecting the attention Jacobian’s condition number to the spectral structure of attention weights, under a full-rank assumption. It also notes that identity or semi-orthogonal initializations give projection matrices condition number 1, while scaling the value initialization too highly can destabilize training.

Proof of Theorem 3.1

  • The Jacobian condition number is the ratio of its maximum to minimum singular values, with the minimum assumed nonzero because the Jacobian has full rank.Source: §A.1 · p. 16
  • The proof bounds the Jacobian’s singular values and uses condition-number bounds for full-rank matrix products and Kronecker products to establish Theorem 3.1.Source: §A.1 · pp. 16–17

Proof of Proposition 3.2

  • Initializing WQ, WK, and WV with identity or semi-orthogonal matrices gives each matrix condition number 1.Source: §A.1 · p. 17

Remark A.1

  • Scaling WV by a nonzero λ is theoretically possible, but empirically large λ caused training instability, so identity initialization worked well.Source: §A.1 · p. 17

Why Conditioning the Jacobian Aids Optimization

  • The NTK connection motivates the analysis: better-conditioned NTKs are associated with faster, more reliable convergence, and Jacobian singular values correspond to square roots of NTK eigenvalues.Source: §A.1 · p. 17

A.1.1 IMPLEMENTATION DETAILS

The appendix describes semi-orthogonal query and key initialization, discusses how conditioning differs from normalization, and reports results across vision and language settings. Across the reported experiments, conditioned initialization is associated with stable or faster convergence and improved performance, with overfitting noted for some larger models.

A.1.1 IMPLEMENTATION DETAILS

  • For each attention head, the query and key projections are initialized with independent semi-orthogonal matrices; the value projection is set to a rectangular identity, with SVD or QR available for construction.Source: §A.1.1 · pp. 17–18

A.2 DISCUSSION ON NORMALIZATION

  • Conditioning targets Jacobian spectral structure rather than activation and gradient scale, so it complements normalization layers retained in the tested architectures.The appendix reports that applying conditioned initialization to WO in addition to WV did not alter training behaviour or final performance; the method therefore leaves WO at standard initialization.Source: §A.2 · pp. 18–19
  • Conditioned initialization on all three projections yields the most stable behaviour and strongest empirical performance, consistent with the Jacobian’s dependence on their joint spectral structure.Source: §A.2 · p. 19

A.3.1 VISION TRANSFORMERS

  • On DaViT-B, conditioned initialization reaches the default initialization’s final accuracy approximately 25% faster and achieves higher accuracy than default and mimetic initialization under the same training regime.The training-loss curves also report stable training and faster convergence than default and mimetic initialization across the listed vision-transformer experiments.Source: §A.3.1 · pp. 19–20, Figure 4 · p. 19, Figure 5 · p. 20
  • For ViT-T weight selection, ImageNet-1K pretraining with conditioned initialization achieves accuracy on par with default initialization transferred from ImageNet-21K.The comparison covers Food-101, CIFAR-10, and CIFAR-100.Source: §A.3.1 · p. 20, Table 6 · p. 20

A.3.4 LANGUAGE MODELING

A.4 LARGER SCALE EXPERIMENTS

  • For 472M-parameter GPT-2 models, conditioned initialization obtains the lowest perplexity on TinyStories and WikiText-103; all tested 1.72B models overfit, though the conditioned initialization performs better than the alternatives.Source: §A.4 · p. 22, Table 8 · p. 22, Table 9 · p. 22
  • On ImageNet-1k, the authors’ initialization outperformed two standard initializations for 200–300M-parameter large-scale ViTs and for 1B-scale ViTs; overfitting was reported for the former.Source: §A.4 · pp. 22–23, Table 13 · p. 23
Loading 2609.07086v1…