Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Orthogonal Self-Attention

Leo Zhang, James Martens

TL;DR

The paper addresses instability in Softmax Self-Attention (SSA) when training skipless Transformers by proposing Orthogonal Self-Attention (OSA), which enforces an orthogonal attention matrix through a matrix exponential of a skew-symmetric matrix S = (α / sqrt(d_v)) (QK^T − KQ^T). OSA enables linear-time computation by exploiting a low-rank structure and yields A(X) ∈ SO(N), improving stability for non-causal architectures. The authors provide theoretical results showing that OSA preserves rank and offers initialization schemes that make the input–output Jacobian well-conditioned, along with practical methods (reduced QR or Newton–Schultz) to compute the exponential. Empirical results on MNIST demonstrate that replacing SSA with OSA in a ViT while removing skip connections and LayerNorm yields competitive generalisation and training efficiency, indicating OSA’s potential to enable stable, efficient skipless Transformers. Overall, OSA presents a principled, scalable approach to stabilising skipless Transformer architectures with implications for representation learning and model design.

Abstract

Softmax Self-Attention (SSA) is a key component of Transformer architectures. However, when utilised within skipless architectures, which aim to improve representation learning, recent work has highlighted the inherent instability of SSA due to inducing rank collapse and poorly-conditioned Jacobians. In this work, we design a novel attention mechanism: Orthogonal Self-Attention (OSA), which aims to bypass these issues with SSA, in order to allow for (non-causal) Transformers without skip connections and normalisation layers to be more easily trained. In particular, OSA parametrises the attention matrix to be orthogonal via mapping a skew-symmetric matrix, formed from query-key values, through the matrix exponential. We show that this can be practically implemented, by exploiting the low-rank structure of our query-key values, resulting in the computational complexity and memory cost of OSA scaling linearly with sequence length. Furthermore, we derive an initialisation scheme for which we prove ensures that the Jacobian of OSA is well-conditioned.

Orthogonal Self-Attention

TL;DR

The paper addresses instability in Softmax Self-Attention (SSA) when training skipless Transformers by proposing Orthogonal Self-Attention (OSA), which enforces an orthogonal attention matrix through a matrix exponential of a skew-symmetric matrix S = (α / sqrt(d_v)) (QK^T − KQ^T). OSA enables linear-time computation by exploiting a low-rank structure and yields A(X) ∈ SO(N), improving stability for non-causal architectures. The authors provide theoretical results showing that OSA preserves rank and offers initialization schemes that make the input–output Jacobian well-conditioned, along with practical methods (reduced QR or Newton–Schultz) to compute the exponential. Empirical results on MNIST demonstrate that replacing SSA with OSA in a ViT while removing skip connections and LayerNorm yields competitive generalisation and training efficiency, indicating OSA’s potential to enable stable, efficient skipless Transformers. Overall, OSA presents a principled, scalable approach to stabilising skipless Transformer architectures with implications for representation learning and model design.

Abstract

Softmax Self-Attention (SSA) is a key component of Transformer architectures. However, when utilised within skipless architectures, which aim to improve representation learning, recent work has highlighted the inherent instability of SSA due to inducing rank collapse and poorly-conditioned Jacobians. In this work, we design a novel attention mechanism: Orthogonal Self-Attention (OSA), which aims to bypass these issues with SSA, in order to allow for (non-causal) Transformers without skip connections and normalisation layers to be more easily trained. In particular, OSA parametrises the attention matrix to be orthogonal via mapping a skew-symmetric matrix, formed from query-key values, through the matrix exponential. We show that this can be practically implemented, by exploiting the low-rank structure of our query-key values, resulting in the computational complexity and memory cost of OSA scaling linearly with sequence length. Furthermore, we derive an initialisation scheme for which we prove ensures that the Jacobian of OSA is well-conditioned.
Paper Structure (24 sections, 15 theorems, 56 equations, 1 figure, 2 tables, 1 algorithm)

This paper contains 24 sections, 15 theorems, 56 equations, 1 figure, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Orthogonal Self-Attention
  3. Implementation Details
  4. Kernel Analysis
  5. Basis Computation
  6. Complexity and Memory Analysis
  7. Initialisation for OSA
  8. Value-Output Initialisation
  9. Query-Key Initialisation
  10. $\alpha$ Initialisation and Jacobian Analysis
  11. MLP Initialisation
  12. Experiments
  13. Conclusion
  14. Further Details on Complexity and Memory Analysis
  15. Sampling
  16. ...and 9 more sections

Key Result

Theorem 2.1

Let $\mathbf{B}(\mathbf{X})\in\mathbb{R}^{N\times r}$ be a matrix where the columns provide an orthonormal basis for the subspace $U$ spanned by the columns of $\mathbf{Q}, \mathbf{K}$. Then we have: where

Figures (1)

  • Figure 1: Train and test loss curves for OSA-Transformer and ViT models trained on MNIST for classification.

Theorems & Definitions (31)

  • Theorem 2.1
  • Theorem 2.2
  • Theorem 3.1
  • Theorem 4.1
  • Theorem 4.2
  • Theorem 4.3
  • Theorem 4.4
  • Remark 4.5
  • proof
  • proof
  • ...and 21 more