Multi-Token Attention

There is a fundamental limitation in traditional attention mechanisms used in large language models, where attention weights are determined by comparing only single query and key token vectors. This “single token attention” restricts the model’s ability to identify relevant context that requires multiple token associations.

To solve this problem, the authors propose Multi-Token Attention (MTA), which applies convolution operations across keys, queries, and attention heads, allowing neighboring tokens to influence each other’s attention weights. This approach enables the model to condition its attention on multiple vector pairs simultaneously, facilitating more precise context location in complex scenarios.

### Understanding Multi-Token Attention (MTA) Mechanism > Multi-Token Attention solves a fundamental limitation in traditional attention mechanisms by allowing LLMs to consider multiple tokens simultaneously when deciding where to focus. In standard attention, each attention value depends solely on comparing a single query vector with a single key vector. This creates a bottleneck when the model needs to find content that contains multiple elements together (like a sentence mentioning both "Alice" and "rabbit"). > > MTA introduces three new steps in this process: > > 1. **Key-Query Convolution**: Instead of looking at token pairs in isolation, MTA applies a sliding window (convolution) over the attention matrix before or after the softmax operation. This allows nearby queries and keys to influence each other's attention weights. For example, when searching for "Alice" and "rabbit," this convolution helps the model focus on areas where both appear in proximity. > > 2. **Head Mixing Convolution**: MTA groups attention heads together and allows them to share information through another convolution operation. If one head finds "Alice" and another finds "rabbit," this mixing helps combine these findings to locate where both terms appear together. > > 3. **Group Normalization with Depth Scaling**: This helps maintain balanced gradients throughout the network, preventing the attention signals from being overwhelmed by the residual stream as they flow through deeper layers. > > The result is an attention mechanism that can effectively use richer contextual information to locate relevant content. Rather than being limited to what can be encoded in a single vector, MTA enables the model to consider patterns across multiple tokens, making it particularly effective for tasks requiring complex information retrieval from long contexts.