Published on

Google's Titan Architecture: Breaking Transformer Memory Bottlenecks

Authors
  • avatar
    Name
    Ajax
    Twitter

Introducing Titan: A New Architecture From Google

The tech world is buzzing about Titan, a novel architecture emerging from Google. It's designed to challenge the limitations of Transformer models, particularly in how they handle memory. This new architecture is gaining significant attention as a potential successor to the Transformer, especially given its development by a team within Google.

The Memory Challenge In Existing Models

Traditional models like LSTM and Transformer, while innovative, face challenges in simulating human-like memory. These challenges include:

  • Limited Capacity: Data is often compressed into a fixed-size hidden state, restricting the amount of information that can be retained.
  • Computational Overhead: While capable of capturing long-range dependencies, the computational cost increases quadratically with sequence length, making it inefficient for very long sequences.
  • Over-Reliance on Training Data: Simply memorizing training data doesn’t always help with real-world application, where test data can fall outside the training distribution.

Titan's Approach: A Neuro-Inspired Memory Module

The Titan team has taken a different approach, seeking to encode information into the parameters of a neural network. They've developed an online meta-model designed to learn how to remember and forget specific data during testing. This model is inspired by neuro-psychological principles, incorporating the following key elements:

  • Surprise as a Trigger: Unexpected events are more easily remembered. The "surprise" is measured by the gradient of the input to the memory module. The larger the gradient, the more unexpected the input.
  • Momentum and Forgetting Mechanisms: A momentum mechanism accumulates short-term surprises into long-term memory, while a forgetting mechanism erases old memories, preventing memory overflow.
  • Multi-Layer Perceptron (MLP) Based Memory: The memory module is composed of multiple MLP layers, allowing it to store deep abstractions of data, making it more powerful than traditional matrix-based memories.

This online meta-learning approach helps the model to focus on learning how to adapt to new data, rather than merely memorizing training data. The module is also designed for parallel computation, enhancing its efficiency.

Integrating The Memory Module Into Deep Learning Architectures

The Titans research team proposed three variations for incorporating their memory module into deep learning architectures:

  1. MAC (Memory as Context): This method combines long-term and persistent memory (which encodes task knowledge) as context that is input to the attention mechanism.
  2. MAG (Memory as Gate): This approach uses gated fusion of the memory module with a sliding window attention mechanism across two branches.
  3. MAL (Memory as Layer): Here, the memory module is implemented as an independent layer that compresses historical information before feeding it to the attention mechanism.

The team found that each variation has its strengths and weaknesses.

Performance And Advantages Of Titans

Titans has demonstrated superior performance across a variety of tasks, including language modeling, common-sense reasoning, and time-series prediction. It has surpassed state-of-the-art models like Transformer and Mamba. Notably, the long-term memory module (LMM) alone has outperformed baseline models in several tasks, showcasing its independent learning capabilities without short-term memory (attention).

In a "needle in a haystack" test designed to find fine-grained clues in long texts, Titans maintained around 90% accuracy even as sequence lengths increased from 2k to 16k. The team indicates that the standard tests do not fully exhibit Titans’ advantages in handling long texts. Titans also outperformed models like GPT4, Mamba, and even Llama3.1 with RAG in a task requiring inference from facts spread across extremely long documents.

Titans has shown impressive performance in specific areas such as time-series prediction and DNA sequence modeling as well.

The Team Behind Titans

The research was conducted by a team from Google Research NYC algorithms and optimization group, not currently part of Google DeepMind.

  • Ali Behrouz, a Cornell University intern, is the first author of the paper.
  • Zhong Peilin, a Tsinghua University alumnus and a Ph.D. graduate from Columbia University, is a research scientist at Google since 2021. He is notable for having published a first-author paper at STOC 2016 as an undergraduate student.
  • Vahab Mirrokni, a Google Fellow and VP, leads the team.

The team developed Titans using Pytorch and Jax and plan to release the code for training and evaluation soon.