Transformers have become the backbone of deep learning models for tasks that require sequential data processing, such as natural language understanding, computer vision, and reinforcement learning. These models rely heavily on self-attention mechanisms, allowing them to capture complex relationships within input sequences. However, as tasks and models scale, the demand for longer contextual windows increases significantly. Managing this extended context window efficiently is crucial because it affects performance and computational cost. Despite their strength, transformers face challenges in maintaining efficiency while handling long context inputs, making this an active area of research.
One of the important challenges is to balance performance with resource efficiency. Transformers store previously computed representations in a cache known as the key-value (KV) cache, allowing them to reference past inputs efficiently. However, this KV cache grows exponentially for long-context tasks, consuming a substantial amount of memory and computational resources. Existing approaches attempt to reduce the size of the KV cache by removing less important tokens, but these methods rely on manually designed heuristics. The limitations of these approaches are evident: they often lead to performance degradation, as detokening strategies are not optimized to retain essential information for subsequent tasks.
Current tools such as the H2O and L2 methods attempt to alleviate this problem by introducing metrics such as L2 norms and entropy to quantify the importance of tokens. These approaches aim to selectively remove tokens from the KV cache, reducing memory usage while preserving model performance. Despite some success, these methods introduce an inherent trade-off: reduced memory usage results in a performance loss. Models using these techniques have difficulty generalizing across tasks, and their heuristic-based design prevents significant improvements in both performance and efficiency simultaneously.
A research team from Sakana ai, Japan has introduced Neural Attention Memory Models (NAMM). NAMMs are a new class of memory management models that dynamically optimize the KV cache in transformers. Instead of relying on manually designed rules, NAMMs learn symbolic importance through evolutionary optimization. By conditioning on the attention matrices of the transformers, NAMMs allow each layer to retain only the most relevant tokens, improving both efficiency and performance without altering the base architecture of the transformer. This universality makes NAMMs applicable to any transformer-based model, since their design depends solely on features extracted from attention matrices.
The methodology behind NAMMs involves extracting meaningful features from the attention matrix using a spectrogram-based technique. Researchers apply the Short-Time Fourier Transform (STFT) to compress attention values into a spectrogram representation. This compact representation captures how symbolic importance evolves over the attention span. The spectrogram features are then reduced using an exponential moving average (EMA) operation to minimize complexity. NAMMs use a lightweight neural network to evaluate these compressed features and assign a selection score to each token. Tokens with low selection scores are removed from the KV cache, freeing memory and ensuring performance is not compromised.
A key innovation in NAMMs is the introduction of backward attention mechanisms. This design allows the network to compare tokens efficiently, preserving only the most relevant occurrences and discarding redundant ones. By leveraging inter-token communication, NAMMs optimize memory usage dynamically across all layers, ensuring that transformers retain crucial long-range information for each task.
The performance of NAMMs was rigorously evaluated across multiple benchmarks, showing its superiority over existing methods. On the LongBench benchmark, NAMMs improved normalized performance by 11% and reduced the KV cache size to 25% of the original model. Similarly, on the challenging InfiniteBench benchmark, where average input lengths exceed 200,000 tokens, NAMMs outperformed benchmark models by increasing performance from 1.05% to 11%. This result highlights the ability of NAMMs to scale effectively for long-context tasks without sacrificing accuracy. Furthermore, the memory footprint of NAMMs in InfiniteBench was reduced to approximately 40% of the original size, demonstrating its efficiency in handling long sequences.
The researchers further validated the versatility of NAMMs through zero-shot transfer experiments. NAMMs trained exclusively on natural language tasks were applied to new transformers and input modalities, including computer vision and reinforcement learning models. For example, when tested with a Llava Next Video 7B model on long video comprehension tasks, NAMMs outperformed the base model while maintaining a reduced memory footprint. In reinforcement learning experiments using decision transformers in continuous control tasks, NAMMs achieved an average performance gain of 9% across multiple tasks, demonstrating their ability to discard useless information and improve decision-making capabilities.
In conclusion, NAMMs provide a powerful solution to the challenge of long context processing in transformers. By learning efficient memory management strategies through evolutionary optimization, NAMMs overcome the limitations of manually designed heuristics. The results demonstrate that NAMM-equipped transformers achieve superior performance while significantly reducing computational costs. Its universal applicability and success in various tasks highlight its potential to advance transformer-based models in multiple domains, marking a significant step towards efficient long-context modeling.
Verify he Paper and <a target="_blank" href="https://sakana.ai/namm/” target=”_blank” rel=”noreferrer noopener”>Details. All credit for this research goes to the researchers of this project. Also, don't forget to follow us on <a target="_blank" href="https://twitter.com/Marktechpost”>twitter and join our Telegram channel and LinkedIn Grabove. Don't forget to join our SubReddit over 60,000 ml.
Trending: LG ai Research launches EXAONE 3.5 – three frontier-level bilingual open-source ai models that deliver unmatched instruction following and broad context understanding for global leadership in generative ai excellence….
Nikhil is an internal consultant at Marktechpost. He is pursuing an integrated double degree in Materials at the Indian Institute of technology Kharagpur. Nikhil is an ai/ML enthusiast who is always researching applications in fields like biomaterials and biomedical science. With a strong background in materials science, he is exploring new advances and creating opportunities to contribute.
<script async src="//platform.twitter.com/widgets.js” charset=”utf-8″>
