Expanding Expressivity in Transformer Models with M\"obiusAttention

Attention mechanisms and Transformer architectures have revolutionized Natural Language Processing (NLP) by enabling exceptional modeling of long-range dependencies and capturing intricate linguistic patterns. However, their inherent reliance on linear operations in the form of matrix multiplications limits their ability to fully capture inter-token relationships on their own. We propose MöbiusAttention, a novel approach that integrates Möbius transformations within the attention mechanism of Transformer-based models. Möbius transformations are non-linear operations in spaces over complex numbers with the ability to map between various geometries. By incorporating these properties, MöbiusAttention empowers models to learn more intricate geometric relationships between tokens and capture a wider range of information through complex-valued weight vectors. We build and pre-train a BERT and a RoFormer version enhanced with MöbiusAttention, which we then finetune on the GLUE benchmark. We evaluate empirically our approach against the baseline BERT and RoFormer models on a range of downstream tasks. Our approach compares favorably against the baseline models, even with smaller number of parameters suggesting the enhanced expressivity of MöbiusAttention. This research paves the way for exploring the potential of Möbius transformations in the complex projective space to enhance the expressivity and performance of foundation models. At the heart of their success lies the attention mechanism (Vaswani et al., 2017), a powerful tool that enables them to identify relationships between different parts of the data, be it words in a sentence or image patches in a scene. Despite their remarkable impact, current transformers face limitations. A key constraint is the inherent linearity of the attention mechanism, which primarily relies on weights learned through linear transformations, matrix multiplications, and the softmax function. While softmax is a non-linear operation, it is only used to produce a probability distribution over the elements signaling their relative importance in comparison to the others, and not to introduce non-linear interdependencies. Predominantly linear operations restrict the ability of models to capture complex linguistic dependencies, leading to potential information loss within each attention layer as shown by recent research (Zhang, 2023). Figure 1: Various Möbius transformations: Each sub-figure shows flows from a single point after successive transformations. Elliptic Möbius has two fixed points at the centers of two circular flows.