Spectral Conditioning of Attention Improves Transformer Performance
This work addresses the problem of unstable training and suboptimal performance in transformer networks by improving the conditioning of attention layers, which could benefit researchers and practitioners working with these models.
This paper analyzes the Jacobian of an attention block and finds it is governed by query, key, and value projections. They introduce a method to alter the spectral properties of each attention layer, reducing the Jacobian's condition number, which leads to improved transformer performance.
We present a theoretical analysis of the Jacobian of an 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 attention layer to reduce the Jacobian's condition number, thereby improving the overall conditioning of the 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.