Tucker Attention: A generalization of approximate attention mechanisms

The pursuit of reducing the memory footprint of the self-attention mechanism in multi-headed self attention (MHA) spawned a rich portfolio of methods, e.g., group-query attention (GQA) and multi-head latent attention (MLA). The methods leverage specialized low-rank factorizations across embedding dimensions or attention heads. From the point of view of classical low-rank approximation, these methods are unconventional and raise questions of which objects they really approximate and how to interpret the low-rank behavior of the resulting representations. To answer these questions, this work proposes a generalized view on the weight objects in the self-attention layer and a factorization strategy, which allows us to construct a parameter efficient scheme, called Tucker Attention. Tucker Attention requires an order of magnitude fewer parameters for comparable validation metrics, compared to GQA and MLA, as evaluated in LLM and ViT test cases. Additionally, Tucker Attention~encompasses GQA, MLA, MHA as special cases and is fully compatible with flash-attention and rotary position embeddings (RoPE). This generalization strategy yields insights of the actual ranks achieved by MHA, GQA, and MLA, and further enables simplifications for MLA.

Modern artificial intelligence systems such as large language models and vision transformers rely on an operation called attention, which helps the model decide which parts of the input are most relevant. However, attention can require a large amount of memory and computation, especially as models grow. This limits how easily such models can be trained, deployed, or run on long inputs. This paper introduces Tucker Attention, a new way to make attention layers smaller while keeping their performance competitive. Instead of compressing individual pieces of the attention mechanism separately, Tucker Attention views the full attention layer as a structured object and compresses it in a more flexible way. This approach also explains several existing efficient attention methods as special cases of a single broader framework. Experiments on language models and vision transformers show that Tucker Attention can use far fewer attention parameters than common alternatives while achieving similar validation performance. It also remains compatible with practical tools used in modern models, including key-value caching, rotary position information, and FlashAttention. These properties make Tucker Attention a promising method for reducing the memory cost of transformer models during both training and inference, without requiring a major change to the surrounding model architecture.