Orthogonal Self-Attention

Leo Zhang, James Martens

Track: tiny paper (up to 4 pages)

Keywords: Transformers, Representation Learning, Lie Groups

TL;DR: We design a novel attention mechanism: Orthogonal Self-Attention for constructing Transformers which are able to be trained without skip connection and normalisation layers.

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.

Anonymization: This submission has been anonymized for double-blind review via the removal of identifying information such as names, affiliations, and identifying URLs.

Presenter: ~Leo_Zhang1

Format: Yes, the presenting author will definitely attend in person because they attending ICLR for other complementary reasons.

Funding: Yes, the presenting author of this submission falls under ICLR’s funding aims, and funding would significantly impact their ability to attend the workshop in person.

Serve As Reviewer: ~Leo_Zhang1

Submission Number: 66