Enhancing Transformers Through Conditioned Embedded Tokens

TL;DR

The paper tackles ill-conditioning in transformer self-attention by linking the conditioning of the first-layer attention to the condition number of embedded tokens and introduces conditioned embedded tokens via a correction term derived from the SVD of the embedding. It provides a theoretical framework with bounds showing that a suitable correction C can reduce $\kappa(X)$ and thereby improve attention conditioning, and demonstrates that $\kappa(X+C) \leq 2$. Empirically, conditioned embedded tokens yield consistent performance improvements across vision (ViT, DeiT, Swin, XCiT, DaViT), object detection (Mask R-CNN with XCiT backbones), and language models (Crammed BERT, GPT-2), including long-range sequence models (Nyströmformer). While offering a practical drop-in enhancement, the work acknowledges limitations in proving optimization guarantees via NTK and notes that the method targets embedded tokens rather than the queries/keys/values directly. The approach holds broad potential for stabilizing and accelerating training in diverse transformer-based systems, with applicability to both standard and efficient attention variants.

Abstract

Transformers have transformed modern machine learning, driving breakthroughs in computer vision, natural language processing, and robotics. At the core of their success lies the attention mechanism, which enables the modeling of global dependencies among input tokens. However, we reveal that the attention block in transformers suffers from inherent ill-conditioning, which hampers gradient-based optimization and leads to inefficient training. To address this, we develop a theoretical framework that establishes a direct relationship between the conditioning of the attention block and that of the embedded tokenized data. Building on this insight, we introduce conditioned embedded tokens, a method that systematically modifies the embedded tokens to improve the conditioning of the attention mechanism. Our analysis demonstrates that this approach significantly mitigates ill-conditioning, leading to more stable and efficient training. We validate our methodology across various transformer architectures, achieving consistent improvements in image classification, object detection, instance segmentation, and natural language processing, highlighting its broad applicability and effectiveness.

Figures (4)

• Figure 1: Schematic representation of conditioned embeddings for a vision transformer: An image is divided into $N$ patches, with each patch then embedded as a high-dimensional vector $X_i \in \mathbb{R}^d$ for $1 \leq i \leq N$. These vectors are then concatenated to form the embedded token matrix $X = [X_1 \cdots X_N]^T$. To improve the conditioning of the self-attention mechanism, a correction term $C = [C_1 \cdots C_N]^T$ is added to $X$, reducing its condition number. The modified matrix is then fed into the first layer of the transformer (positional encoding not shown). Its effect on the self-attention equation is illustrated in the green equation (for simplicity, layer normalization has been omitted).
• Figure 2: Condition number comparison for ViT-B on ImageNet-1k. The left, middle, and right bars show the condition number of embedded tokens, first-layer self-attention (averaged across heads), and self-attention across all layers, respectively, averaged over training epochs. Our model (conditioned) consistently achieves a significantly lower condition number than the original.
• Figure 3: Condition number comparison for GPT-2 training on TinyStories. The left, middle, and right bars show the condition number of embedded tokens, first-layer self-attention (averaged across heads), and self-attention across all layers, respectively, averaged over all epochs. Our model (conditioned) consistently achieves a significantly lower condition number than the original.
• Figure 4: Condition number comparison for Nyströmformer on the LRA benchmark. The left, middle, and right bars show the condition number of embedded tokens, first-layer self-attention (averaged across heads), and self-attention across all layers, averaged over training epochs and tasks. Our model (conditioned) consistently achieves a significantly lower condition number than the original.

Theorems & Definitions (8)

• Definition 4.1
• Proposition 4.2
• Remark 4.3
• Theorem 4.4
• Theorem 4.5
• Remark 4.6
• Remark 4.7
• Remark 4.8