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.

翻译：Softmax自注意力（SSA）是Transformer架构的核心组件。然而，当应用于旨在提升表征学习的无跳跃连接架构时，近期研究指出SSA存在内在不稳定性，因其会引发秩崩溃及病态雅可比矩阵。本文设计了一种新颖的注意力机制：正交自注意力（OSA），旨在规避SSA的上述问题，从而使（非因果）Transformer无需跳跃连接与归一化层即可更易于训练。具体而言，OSA通过将查询-键值构成的斜对称矩阵映射至矩阵指数，将注意力矩阵参数化为正交矩阵。我们证明该机制可通过利用查询-键值的低秩结构实现高效计算，使得OSA的计算复杂度与内存开销随序列长度呈线性增长。此外，我们推导出一种初始化方案，并证明该方案能确保OSA的雅可比矩阵处于良态条件。