Spectral Conditioning of Self-Attention Improves Transformer Performance
Hemanth Saratchandran, Simon Lucey
We present a theoretical analysis of the Jacobian of a self-attention block within a transformer, showing that it is governed by the query, key, and value projections that define the attention mechanism. Leveraging this insight, we introduce a method that systematically alters the spectral properties of each self-attention layer to reduce the Jacobian’s condition number, thereby improving the overall conditioning of the self-attention layers within a transformer network. We empirically show that this improved Jacobian conditioning translates to enhanced performance in practice. Our approach is simple, broadly applicable, and can be easily integrated as a drop-in replacement for a wide range of existing attention mechanisms. We validate its effectiveness across diverse transformer architectures and tasks, demonstrating consistent improvements in performance.