Multi-Token Attention (MTA) Enables Efficient Retrieval of Contextual Information

This paper introduces Multi-Token Attention (MTA), an advanced attention mechanism that conditions attention weights simultaneously on multiple query and key vectors.

Large Language Models (LLMs) have significantly benefited from attention mechanisms, which enable the effective retrieval of contextual information. However, traditional attention methods primarily depend on single token attention, where each attention weight is calculated from a single pair of query and key vectors.

This design inherently constrains the model's ability to discern contexts that require the integration of multiple token signals, ultimately limiting its effectiveness on complex linguistic dependencies. For instance, identifying sentences that simultaneously contain both "Alice" and "rabbit" poses a challenge because conventional attention mechanisms struggle to combine multiple separate attention signals efficiently without substantially increasing model complexity.

To address this limitation, researchers from Meta AI have introduced Multi-Token Attention (MTA), an advanced attention mechanism that simultaneously conditions attention weights on multiple query and key vectors. MTA integrates convolution operations over queries, keys, and attention heads, thus enhancing the precision and efficiency of contextual information retrieval.

MTA framework consists of two convolutional components:

1) key-query convolution, which aggregates multiple token signals within individual attention heads, and

2) head mixing convolution, which facilitates information sharing among different attention heads. MTA is implemented using group normalization with depth-dependent scaling to stabilize gradient flow, further improving model training stability and efficacy.

At a technical level, MTA modifies standard attention calculations by incorporating a two-dimensional convolution operation on the attention logits before softmax normalization. This convolution allows adjacent queries and keys to influence attention scores mutually, enabling the attention mechanism to identify contextual relationships more precisely. Consequently, the model efficiently aggregates local token interactions without significantly increasing the number of parameters or the dimensionality of attention vectors.

MTA promotes effective knowledge transfer among attention heads, selectively amplifying relevant context signals while attenuating less pertinent information. These enhancements collectively yield a more robust attention mechanism capable of capturing complex multi-token interactions.

Empirical evaluations validate the efficacy of MTA across several natural language processing (NLP) benchmarks. In a structured motivating task explicitly designed to illustrate the shortcomings of single-token attention mechanisms, MTA demonstrated near-perfect performance, achieving an error rate of only 0.1% in tasks with 4 x 1024 token sequences. In contrast, standard Transformer models exhibited error rates greater than 50%.

Further large-scale experiments involved an 880M-parameter model trained on 105 billion tokens using MTA and baseline architectures. MTA achieved superior validation perplexity scores across diverse datasets such as arXiv, GitHub, and Wikipedia.

MTA outperformed standard Transformer models in tasks requiring extended context comprehension, such as the Needle-in-the-Haystack and BabiLong benchmarks. In the Needle-in-the-Haystack task with 4K token contexts containing multiple needles, MTA achieved accuracies ranging from 67% to 97.6%, surpassing standard models by substantial margins. These results highlight the potential of MTA for enabling LLMs to efficiently process very long-range dependencies.

In summary, Multi-Token Attention (MTA) presents a refined advancement in attention mechanisms by addressing fundamental limitations of traditional single-token attention. Leveraging convolutional operations to concurrently integrate multiple query-key interactions, MTA enhances the ability of language models to handle intricate contextual dependencies.

These methodological improvements facilitate more precise and efficient performance, particularly in scenarios involving complex token interactions and long-range contextual understanding. Through targeted modifications to standard attention mechanisms, MTA contributes meaningfully to the evolution of more sophisticated, accurate, and computationally efficient language models.