2025-11-20T08:25:14.880374

Titans: Learning to Memorize at Test Time

Behrouz, Zhong, Mirrokni
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.
academic

Titans: Обучение запоминанию во время тестирования

Основная информация

  • ID статьи: 2501.00663
  • Название: Titans: Learning to Memorize at Test Time
  • Авторы: Ali Behrouz, Peilin Zhong, Vahab Mirrokni (Google Research)
  • Классификация: cs.LG cs.AI cs.CL
  • Дата публикации: 31 декабря 2024 г.
  • Ссылка на статью: https://arxiv.org/abs/2501.00663

Аннотация

В данной работе предлагается новый модуль нейронной долгосрочной памяти, способный обучаться запоминанию исторического контекста и помогающий механизму внимания сосредоточиться на текущем контексте при одновременном использовании долгосрочной информации. Авторы утверждают с позиции памяти, что механизм внимания функционирует как кратковременная память благодаря ограниченному контексту, но точному моделированию зависимостей, тогда как нейронная память действует как долгосрочная, более устойчивая память благодаря своей способности сохранять данные. На основе этих двух модулей авторы представляют новое семейство архитектур Titans и предлагают три варианта для эффективной интеграции памяти в архитектуру. Экспериментальные результаты показывают, что Titans превосходит Transformers и современные линейные рекуррентные модели в задачах языкового моделирования, логического вывода на основе здравого смысла, геномики и анализа временных рядов, эффективно масштабируясь до размеров контекстного окна свыше 2M токенов.

Исследовательский контекст и мотивация

Основные проблемы

Существующие архитектуры последовательного моделирования сталкиваются с компромиссом между эффективностью и производительностью:

  1. Transformers: хотя и способны точно моделировать зависимости, имеют вычислительную сложность O(n²), что ограничивает длину контекста
  2. Линейные Transformers/RNNs: хотя и эффективны, сжимают информацию в состояние фиксированного размера, что приводит к снижению производительности на длинных последовательностях
  3. Отсутствие системы памяти: существующие архитектуры не обладают многоуровневой системой памяти, подобной человеческому мозгу (кратковременная память, долгосрочная память, метапамять и т.д.)

Исследовательская мотивация

Вдохновляясь системой человеческой памяти, авторы полагают, что эффективная парадигма обучения требует:

  • Различных, но взаимосвязанных модулей, каждый из которых отвечает за ключевой компонент процесса обучения
  • Способности активно обучаться из данных и запоминать абстракции прошлой истории
  • Механизма непрерывного обучения и адаптации во время тестирования

Основные вклады

  1. Модуль нейронной долгосрочной памяти: предлагается глубокая нейронная сеть в качестве метамодели, которая обучается во время тестирования запоминанию/сохранению данных в своих параметрах
  2. Механизм управления памятью: разработан механизм обновления памяти на основе "удивления" и адаптивный механизм забывания
  3. Семейство архитектур Titans: предложены три способа интеграции памяти в архитектуры глубокого обучения: Memory as Context (MAC), Memory as Gate (MAG), Memory as Layer (MAL)
  4. Параллельный алгоритм обучения: предоставлен быстрый параллелизуемый алгоритм обучения, обеспечивающий эффективное обучение модулей глубокой памяти
  5. Обширная экспериментальная верификация: верификация эффективности Titans на множестве задач, включая языковое моделирование, логический вывод на основе здравого смысла, геномику и прогнозирование временных рядов

Подробное описание методологии

Определение задачи

В данной работе исследуется задача последовательного моделирования, где входные данные представляют собой последовательность x∈RN×dinx \in \mathbb{R}^{N \times d_{in}}, а цель состоит в обучении модели, способной эффективно обрабатывать длинные последовательности, которая должна:

  • Непрерывно обучаться и запоминать во время тестирования
  • Сбалансировать использование кратковременной и долгосрочной памяти
  • Обладать линейной сложностью при сохранении высокой выразительной способности

Модуль нейронной долгосрочной памяти

Основная идея проектирования

Вдохновляясь долгосрочной памятью человека, события, нарушающие ожидания (удивительные события), запоминаются легче. Авторы используют градиент нейронной сети относительно входных данных для измерения "удивления".

Механизм обновления памяти

Базовое правило обновления:

M_t = M_{t-1} - θ_t ∇ℓ(M_{t-1}; x_t)

Улучшенное правило обновления (с введением момента):

M_t = M_{t-1} + S_t
S_t = η_t S_{t-1} - θ_t ∇ℓ(M_{t-1}; x_t)

Где:

  • S_t: момент удивления, содержащий прошлое удивление и мгновенное удивление
  • η_t: параметр затухания удивления, зависящий от данных
  • θ_t: параметр, контролирующий степень интеграции мгновенного удивления

Механизм забывания

Для обработки длинных последовательностей вводится адаптивный механизм забывания:

M_t = (1 - α_t)M_{t-1} + S_t

Где α_t ∈ [0,1] — механизм вентиля, контролирующий степень забывания.

Целевая функция

Используется функция потерь ассоциативной памяти:

ℓ(M_{t-1}; x_t) = ||M_{t-1}(k_t) - v_t||²₂

Где k_t = x_t W_K, v_t = x_t W_V

Варианты архитектуры Titans

1. Memory as Context (MAC)

  • Память используется как контекст для текущей информации
  • Последовательность обрабатывается блоками, каждый блок запрашивает долгосрочную память для получения релевантной исторической информации
  • Механизм внимания определяет, требуется ли информация из долгосрочной памяти

2. Memory as Gate (MAG)

  • Одна ветвь использует скользящее окно внимания как кратковременную память
  • Другая ветвь использует модуль нейронной памяти как долгосрочную память
  • Две ветви объединяются через механизм вентиля

3. Memory as Layer (MAL)

  • Нейронная память используется как слой глубокой сети
  • Последовательная обработка: сначала через слой памяти, затем через слой внимания
  • Аналогично иерархическому проектированию существующих гибридных моделей

Параллельное обучение

Высокоэффективное параллельное обучение достигается путем переформулирования процесса обучения с использованием операций матричного умножения и суммирования:

  • Последовательность разбивается на блоки размером b
  • Для вычисления членов момента используется параллельное ассоциативное сканирование
  • Быстрое обучение реализуется путем тензоризации стохастического градиентного спуска в малых партиях

Экспериментальная установка

Наборы данных

  • Языковое моделирование: набор данных FineWeb-Edu, 15B/30B токенов
  • Логический вывод на основе здравого смысла: PIQA, HellaSwag, WinoGrande, ARC-easy/challenge, SIQA, BoolQ
  • Задачи с длинным контекстом: RULER benchmark (S-NIAH), BABILong benchmark
  • Временные ряды: ETT, ECL, Traffic, Weather наборы данных
  • Геномика: GenomicsBenchmarks набор данных

Размеры моделей

  • Модели с 170M, 340M, 400M, 760M параметров
  • Длина обучения: 4K токенов
  • Контекстное окно: масштабируется до 2M+ токенов

Методы сравнения

  • Transformers: Transformer++
  • Линейные рекуррентные модели: RetNet, GLA, Mamba, Mamba2, DeltaNet, TTT, Gated DeltaNet
  • Гибридные модели: Samba, Gated DeltaNet-H2
  • Большие модели: GPT-4, Llama3, RecurrentGemma, Mistral

Результаты экспериментов

Производительность языкового моделирования

На модели с 340M параметров:

  • Titans (LMM): перплексия 26.18 (Wiki), 29.97 (LMB)
  • Лучший базовый метод TTT: перплексия 27.44 (Wiki), 34.19 (LMB)
  • Titans (MAG) среди гибридных моделей показывает лучший результат: перплексия 25.07 (Wiki), 28.72 (LMB)

Задачи с длинным контекстом

На задаче S-NIAH (длина последовательности 16K):

  • Titans (MAC): S-NIAH-PK 98.4%, S-NIAH-N 97.4%, S-NIAH-W 95.2%
  • Mamba2: S-NIAH-PK 5.4%, S-NIAH-N 0.0%, S-NIAH-W 0.0%
  • TTT: S-NIAH-PK 88.4%, S-NIAH-N 4.4%, S-NIAH-W 0.0%

Тестирование на BABILong

  • Titans превосходит все базовые методы в условиях few-shot, включая GPT-4
  • В условиях fine-tuning небольшие модели Titans превосходят GPT-4 с параметрами на 70 раз больше

Прогнозирование временных рядов

Модуль нейронной памяти превосходит все базовые методы на всех наборах данных, включая методы на основе Mamba, Transformer и линейных моделей.

Абляционные исследования

Вклад компонентов (в порядке важности):

  1. Затухание весов (механизм забывания)
  2. Механизм момента
  3. Сверточный слой
  4. Постоянная память
  5. Глубокая память vs линейная память

Связанные работы

Линейные рекуррентные модели

  • Первое поколение: RetNet, LRU, RWKV, S4/S5 — использование независимых от данных матриц переходов
  • Второе поколение: Griffin, серия Mamba — введение механизмов вентилей
  • Третье поколение: DeltaNet, TTT, Longhorn — правила обновления на основе метаобучения/онлайн-обучения

Варианты Transformer

  • Оптимизация эффективности: разреженное внимание, линейное внимание, реализация с учетом ввода-вывода
  • Сегментированный Transformer: RMT и другие используют простую векторную память для передачи информации между блоками

Обучение во время тестирования

  • Вдохновлено ранними алгоритмами локального обучения
  • Наиболее связано с MNM и TTT-layer, но Titans обладает механизмом забывания и обновлением момента

Заключение и обсуждение

Основные выводы

  1. Важность системы памяти: многоуровневая система памяти (кратковременная + долгосрочная + постоянная память) критична для последовательного моделирования
  2. Эффективность обучения во время тестирования: непрерывное обучение и запоминание во время тестирования могут значительно улучшить способность обработки длинных последовательностей
  3. Влияние проектирования архитектуры: архитектуры MAC и MAG превосходят традиционный иерархический дизайн MAL
  4. Верификация масштабируемости: Titans эффективно масштабируется до контекстного окна 2M+ токенов

Ограничения

  1. Вычислительные затраты: модули глубокой памяти требуют больше вычислительных ресурсов, чем простые матричные состояния
  2. Компромисс глубины памяти: более глубокие модули памяти показывают лучшие результаты, но обучаются медленнее
  3. Чувствительность параметров: требуется тщательная настройка параметров, связанных с удивлением
  4. Недостаточный теоретический анализ: отсутствуют теоретические гарантии емкости памяти и стратегий забывания

Направления будущих исследований

  1. Оптимизация архитектуры памяти: исследование более эффективных проектов архитектур нейронной памяти
  2. Теоретический анализ: предоставление теоретического анализа емкости памяти и стратегий забывания
  3. Верификация в большом масштабе: верификация эффективности метода на моделях большего размера
  4. Расширение приложений: исследование потенциала применения в большем количестве областей

Глубокая оценка

Преимущества

  1. Сильная концептуальная новизна: переосмысление последовательного моделирования с позиции системы человеческой памяти, предложение новой многоуровневой архитектуры памяти
  2. Всесторонний технический вклад: не только предложение модуля нейронной памяти, но и разработка трех способов интеграции в архитектуру и эффективного параллельного алгоритма
  3. Достаточная экспериментальная верификация: комплексные эксперименты в нескольких областях (NLP, временные ряды, геномика) с убедительными результатами
  4. Прочная теоретическая база: связь обновления памяти с градиентным спуском, моментом и затуханием весов, предоставление теоретического объяснения

Недостатки

  1. Недостаточный анализ вычислительной сложности: хотя утверждается линейная сложность, анализ фактических вычислительных затрат модулей глубокой памяти недостаточно подробен
  2. Чувствительность гиперпараметров: установка нескольких зависящих от данных параметров (α_t, θ_t, η_t) может быть относительно сложной
  3. Ограниченная аналогия с человеческой памятью: хотя вдохновлено человеческой памятью, определение "удивления" относительно простое
  4. Недостаточная верификация на больших моделях: наибольшая модель содержит только 760M параметров, отсутствует верификация на моделях с миллиардами параметров

Влияние

  1. Академическая ценность: предоставление новой перспективы памяти для последовательного моделирования, потенциально вдохновляющее дальнейшие исследования
  2. Практическая ценность: превосходная производительность на задачах обработки длинных последовательностей с потенциалом практического применения
  3. Воспроизводимость: авторы обещают открыть исходный код, способствуя распространению и верификации метода

Применимые сценарии

  1. Обработка длинных документов: применимо к задачам NLP, требующим обработки длинных документов
  2. Анализ временных рядов: особенно подходит для задач прогнозирования, требующих долгосрочной исторической информации
  3. Сценарии онлайн-обучения: применимо к приложениям, требующим непрерывной адаптации во время тестирования
  4. Задачи, требующие интенсивного использования памяти: такие как системы вопросов и ответов, диалоговые системы и другие приложения, требующие запоминания большого объема информации

Библиография

Статья цитирует 138 связанных работ, охватывающих множество областей, включая Transformer, рекуррентные нейронные сети, механизмы внимания, сети памяти, обучение во время тестирования и другие важные работы в смежных областях, обеспечивая прочную теоретическую базу для данного исследования.