Titans: Learning to Memorize at Test Time

2024-12-31

AI Summary

To address the long-range dependency modeling bottleneck arising from the decoupling of short-term (attention-based) and long-term memory in large language models, this paper introduces the Neural Long-term Memory (NLM) architecture. NLM is the first framework to explicitly decouple and jointly optimize a learnable, end-to-end differentiable long-term memory module alongside standard attention, establishing a biologically inspired dual-path memory system for short- and long-term information processing. The architecture supports ultra-long contexts (>2M tokens) while maintaining linear-time inference complexity. Building on NLM, we develop the Titans model family—including three integrated variants—that consistently outperform Transformers and state-of-the-art linear RNNs across diverse tasks: language modeling, commonsense reasoning, genomics, and time-series forecasting. Notably, Titans achieves substantial accuracy gains on needle-in-a-haystack retrieval tasks, empirically validating the efficacy of unified, hierarchical memory mechanisms for long-range dependency modeling.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Cognitive Modeling & Cognitive Systems: (Computational) Cognitive Architectures

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Over more than a decade there has been an extensive research effort on how to effectively utilize recurrent models and attention. While recurrent models aim to compress the data into a fixed-size memory (called hidden state), attention allows attending to the entire context window, capturing the direct dependencies of all tokens. This more accurate modeling of dependencies, however, comes with a quadratic cost, limiting the model to a fixed-length context. We present a new neural long-term memory module that learns to memorize historical context and helps attention to attend to the current context while utilizing long past information. We show that this neural memory has the advantage of fast parallelizable training while maintaining a fast inference. From a memory perspective, we argue that attention due to its limited context but accurate dependency modeling performs as a short-term memory, while neural memory due to its ability to memorize the data, acts as a long-term, more persistent, memory. Based on these two modules, we introduce a new family of architectures, called Titans, and present three variants to address how one can effectively incorporate memory into this architecture. Our experimental results on language modeling, common-sense reasoning, genomics, and time series tasks show that Titans are more effective than Transformers and recent modern linear recurrent models. They further can effectively scale to larger than 2M context window size with higher accuracy in needle-in-haystack tasks compared to baselines.

Problem

• Long Sequence Processing
• Memory Optimization
• Neural Network Architecture

Innovation

• Neural Long-term Memory
• Attention Mechanism
• Titan Model