Preconditioned Attention: Enhancing Efficiency in Transformers

Abstract

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

Figures (4)

• Figure 1: Schematic representation of preconditioned self-attention. Left: A layer of a general transformer employing a preconditioned self-attention block. Right: Self-attention (top) and preconditioned self-attention (bottom) are compared. The key difference is that preconditioned self-attention applies a multiplication by a diagonal preconditioner matrix C, which depends on the query $Q$, key $K$, and value $V$.
• Figure 2: Total number of epochs required for each preconditioned model to reach the accuracy of the baseline model. For each ViT we see that the preconditioned model requires roughly 20-30% less epochs.
• Figure 3: Average condition number of a ViT-B and a preconditioned ViT-B during training on the ImageNet-1k dataset.
• Figure 4: Conditioning analysis of the Nyströmformer on text classification and ListOps. The preconditioned variant achieves consistently lower condition numbers (left and middle) and converges faster, requiring fewer iterations to reach the baseline model’s final accuracy (right).

Theorems & Definitions (10)

• Definition 4.1
• Theorem 4.2: guggenheimer1995simple
• Theorem 4.3
• Theorem 4.4
• Lemma A.1
• proof
• proof : Proof of \ref{['thm:condition_self_attn']}
• Lemma A.2
• proof
• proof : Proof of \ref{['thm:precond_attn']}