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

Титаны: Учимся запоминать во время тестов Оригинал: https://arxiv.org/pdf/2501.00663v1

Неофициальная реализация Titans Architecture: https://github.com/lucidrains/titans-pytorch

1, Предыстория и мотивация: ограничения трансформера и вдохновение человеческой памяти

1. Ограничения трансформера: узкое место при обработке длинных последовательностей

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

• Высокая вычислительная сложность ограничивает масштабируемость модели :.
  • Механизм самовнушения требует вычисления каждого жетон сходство со всеми остальными лексемами в последовательности, с временной и пространственной сложностью O(N²)который N длина последовательности.
  • Это означает, что вычисления и потребление памяти растут в квадрате при увеличении длины последовательности, что сильно ограничивает возможности модели по работе с длинными последовательностями. Например, Transformer часто оказывается не в состоянии справиться с такими задачами, как длинный текст, понимание видео или предсказание долгосрочных временных рядов.
  

  Рисунок 1: Вычислительный процесс механизма самовнушения.

• Ограниченное контекстное окно для захвата зависимостей на больших расстояниях :.
  • Чтобы облегчить вычислительную нагрузку, трансформеры обычно используют контекстное окно фиксированной длины (например, 512 или 1024), что означает, что модель может сосредоточиться только на информации в пределах текущего окна.
  • Однако для решения многих реальных задач требуются модели, способные улавливать зависимости в более длительных временных масштабах, например, для понимания контекстной информации в длинных текстах или диалогах, интеграции информации из разных моментов времени в видео, а также для прогнозирования на основе долгосрочных тенденций и закономерностей в исторических данных.

2. Компромисс линейного трансформатора: эффективность против производительности

Чтобы решить проблему узкого места в вычислительной системе Transformer, исследователи предложилиЛинейный трансформаторОсновными улучшениями являются:

• Замена softmax на ядерную функцию. Замена вычисления softmax в механизме самовнушения на функцию ядра снижает вычислительную сложность до O(N).
• Параллелизуемые рассуждения. Вычислительный процесс линейного трансформатора может быть представлен в циклической форме, что позволяет проводить более эффективные рассуждения.

Однако линейный трансформатор также имеет некоторые ограничения:

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

  Рисунок 2: Процесс обновления памяти для линейного трансформатора.

3. Вдохновение человеческой системы памяти: создание более прочных механизмов долговременной памяти

Для того чтобы преодолеть вышеуказанные трудности, авторы началисистема памяти человекаЧерпая вдохновение в

• Взаимосвязь между памятью и обучением: в
  • Диссертация опирается на определения памяти и обучения из нейропсихологической литературы, рассматривая память как нейронное обновление, вызванное входными данными, и определяя обучение как процесс приобретения эффективных и полезных воспоминаний с учетом поставленной цели.
  • Это означает, чтоЭффективное обучение невозможно без сильного механизма памяти.
• Многоуровневая природа человеческой памяти:
  • Система памяти человека не является единой структурой, а состоит из множества подсистем, таких как кратковременная память, рабочая память и долговременная память, каждая из которых имеет различные функции и организационную структуру и способна работать независимо.
  • Такая многоуровневая природа позволяет людям эффективно хранить, извлекать и управлять информацией.
• Недостатки существующих моделей.
  • Существующие архитектуры нейронных сетей (от сетей Хопфилда до LSTM и трансформеров) испытывают трудности с обобщением, экстраполяцией длины и выводом - возможностями, которые критически важны для многих сложных задач реального мира.
  • Все эти архитектуры, хотя и черпали вдохновение в человеческом мозге, не обладали способностьюэфемерная памятьЭффективное моделированиеМногоуровневая природа системы памятиСимуляция.

2, Основная инновация: модуль нейронной долговременной памяти и архитектура титанов

Основываясь на этих размышлениях, авторы предлагают следующие нововведения:

1. Нейронный модуль долговременной памяти

(1) Концепция дизайна.

• Механизмы метаконтекстного обучения.
  • Модуль разработан какметамодель, научившись запоминать/сохранять данные в своих параметрах во время тестирования.
  • такойОнлайн-обучениеПодход позволяет модели динамически адаптировать свою память к текущим входным данным, а не полагаться на память, которая была зафиксирована во время предварительного обучения.
• Обновление памяти на основе сюрпризов : The
  • Авторы опираются на механизм человеческой памяти, согласно которому "неожиданные события запоминаются с большей вероятностью", и предлагают метод, основанный настепень удивленияМеханизм обновления памяти.
  • степень удивленияРассчитывая нейронную сеть относительно входного сигналаградиентчем больше градиент, тем больше входные данные отличаются от исторических и тем больше они заслуживают внимания.
    • Этот метод эффективен для захвата ключевой информации в данных и сохранения ее в долговременной памяти.
    • В отличие от этого, Linear Transformer может выполнять только линейное преобразование на основе текущих входных данных, что затрудняет эффективное отражение дальних зависимостей.
  

  Рисунок 3: Механизм обновления памяти на основе сюрпризов.

(2) Ключевые технологии.

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

(3) Структуры памяти.

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

2. Архитектура титанов: интеграция долговременной и кратковременной памяти

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

(1) Три суперголовки, работающие в тандеме.

• Ядро.
  • Состоящая из кратковременной памяти, она отвечает за основной процесс обработки данных.
  • пользоватьсяМеханизмы внимания для окон конечного размераНапример, Sliding Window Attention (SWA) или Fully Connected Attention (FCA).
  • Кратковременная память может рассматриваться каккратковременная память, который используется для определения зависимостей в текущем контексте.
• Долговременная память.
  • Отвечает за хранение/запоминание информации из далекого прошлого.
  • Используется нейронный модуль долговременной памяти, описанный выше.
  • Долговременная память может рассматриваться какэфемерная памятьОн используется для хранения и извлечения информации в течение длительного времени.
• Постоянная память.
  • это набор обучаемых, но не зависящих от данных параметров, которые кодируют априорные знания о задаче.
  • Аналогичны параметрам полностью связанного слоя в Transformer, но имеют другую функциональность.
  • Постоянную память можно рассматривать какметапамять, который используется для хранения знаний, связанных с задачей, таких как правила грамматики, знания здравого смысла и т.д.
  

  Рисунок 4. Принципиальная схема архитектуры Titans (вариант MAC).

(2) Три различных типа интеграции.

• Память как контекст (MAC).
  • Подключение длительных и устойчивых воспоминаний к входным последовательностям в качестве дополнительной информации к текущему контексту.
  • Механизмы внимания определяют, какая информация должна быть сохранена в долговременной памяти.
  • Во время тестирования параметры постоянной памяти оставались фиксированными, а модуль внимания был настроен на контекстное обучение, в то время как модуль долговременной памяти продолжал изучать/запоминать информацию.
  • Такая конструкция позволяет модели гибко использовать информацию из долговременной памяти в зависимости от текущих входных данных.
• Память как контроль (MAG).
  • Использование внимания со скользящим окном в качестве кратковременной памяти и нейронных модулей памяти в качестве долговременной памяти.
  • Комбинирование с помощью механизмов разделения, например, нормализация двух показателей с помощью обучаемых векторно-значных весов и последующее применение нелинейной функции активации.
  • Этот дизайн можно рассматривать как многоголовочную архитектуру, в которой разные головки имеют разную структуру.

  
  Рисунок 5: Различные варианты архитектуры Titans (MAC и MAG).

• Память как слой (MAL).
  • Использование модуля нейронной памяти в качестве слоя глубокой нейронной сети сжимает прошлую и текущую контекстную информацию перед модулем внимания.
  • Этот дизайн чаще встречается в литературе, например, модель H3.

(3) Сильные стороны.

• Более гибкое управление памятью.
  • Используя модули памяти в качестве контекстных или управляемых ветвей, архитектура Titans способна динамически использовать информацию долговременной памяти в зависимости от текущих входных данных.
  • Это обеспечивает большую гибкость по сравнению с традиционным подходом, когда модули памяти используются в качестве слоев.
• Более сильное выражение лица.
  • Синергия трех суперголовок позволяет архитектуре Titans более эффективно обрабатывать длинные последовательные данные и объединять преимущества кратковременной, долговременной и постоянной памяти.
• Масштабируемость.
  • По сравнению с Transformer, архитектура Titans обладает лучшей масштабируемостью при работе с длинными последовательностями и способна поддерживать высокую производительность при большем контекстном окне.

3, Результаты экспериментов и анализ: проверка архитектуры Titans

Авторы провели обширные эксперименты на нескольких задачах, чтобы оценить производительность архитектуры Titans и ее вариантов:

1. Лингвистическое моделирование и рассуждения на основе здравого смысла.

• Экспериментальная установка.
  • Были использованы три модели Titans разного размера (340M, 400M, 760M параметров), а также несколько базовых моделей, включая Transformer++, RetNet, GLA, Mamba, Mamba2, DeltaNet, TTT и Gated DeltaNet.
  • В качестве обучающих данных использовался набор данных FineWeb-Edu.
• Ключевые выводы.
  • В несмешанной модели модуль нейронной долговременной памяти показал наилучшие результаты как по показателям сложности, так и по точности.
  • Все три варианта Titans (MAC, MAG, MAL) превзошли Samba (Mamba + внимание) и Gated DeltaNet-H2 (Gated DeltaNet + внимание).
  • MAC лучше справляется с зависимостями на больших расстояниях, а MAG и MAC превосходят вариант MAL.
  

  Рисунок 6: Сравнение производительности Титанов с базовой моделью в задачах языкового моделирования и рассуждений.

2. Задание "Иголка в стоге сена".

• Экспериментальная установка.
  • С помощью задачи Single NIAH (S-NIAH) в бенчмарке RULER оценивалась способность модели к извлечению информации из последовательностей длиной 2K, 4K, 8K и 16K.
• Ключевые выводы.
  • Модуль "Нейронная долговременная память" показал наилучшие результаты во всех трех заданиях.
  • Варианты Titans также показали хорошие результаты, причем вариант MAC оказался лучшим.

3. Бенчмаркинг BABILong.

• Экспериментальная установка.
  • Задача требует, чтобы модель рассуждала о распределенной фактической информации в очень длинных документах.
  • Существуют настройки уменьшения выборки и настройки тонкой настройки.
• Ключевые выводы.
  • В условиях меньшей выборки Titans превосходит все базовые модели, включая модели с большим числом параметров, такие как GPT-4 и GPT4o-mini.
  • Titans также превосходит все модели в настройках тонкой настройки, даже такие крупные модели, как GPT-4.
  • По сравнению с моделями памяти на основе трансформаторов (RMT), Titans демонстрирует более высокую производительность, в основном благодаря своим мощным возможностям памяти.
  

  Рисунок 7: Сравнение производительности Titans и базовой модели на бенчмарке BABILong.

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

• Экспериментальная установка.
  • Используя фреймворк Simba, модуль Mamba был заменен на нейронный модуль долговременной памяти.
  • Оценивается на эталонных наборах данных ETT, ECL, Traffic и Weather.
• Ключевые выводы.
  • Модуль нейронной долговременной памяти превосходит все базовые архитектуры, включая Mamba, линейную модель и архитектуру на основе трансформатора.

5. Моделирование ДНК.

• Экспериментальная установка.
  • Оцените производительность предварительно обученных моделей в последующих задачах на GenomicsBenchmarks.
• Ключевые выводы.
  • Титаны (LMM) конкурентоспособны в различных задачах геномики и не уступают современным методам.

6. Анализ эффективности.

• Ключевые выводы.
  • Модуль нейронной долговременной памяти обучается немного медленнее, чем другие рекуррентные модели, в основном за счет более глубокой памяти и более сложного процесса преобразования, а также за счет высоко оптимизированного ядра, реализованного в Mamba2.
  • Titans (MAL) быстрее базового уровня, а также модуля памяти, в основном за счет использования высокооптимизированного ядра FlashAttention.

7. Исследования абляции.

• Ключевые выводы.
  • Все компоненты конструкции нейронной памяти вносят положительный вклад в производительность, причем наибольший вклад вносят такие компоненты, как затухание веса, импульс, свертка и постоянная память.
  • Архитектурный дизайн также оказывает значительное влияние на производительность: MAC и MAG демонстрируют близкие результаты в задачах языкового моделирования и рассуждений, в то время как MAC лучше справляется с задачами с длинным контекстом.

4, Инновационные моменты и сильные стороны диссертации

1. Предлагается новый нейронный модуль долговременной памяти:
   • Она использует ключевые элементы механизмов человеческой памяти, такие как удивление, импульс и забывание, чтобы добиться более эффективного обновления и хранения памяти.
   • Глубокая нейронная сеть используется в качестве модуля памяти, чтобы придать модели большую выразительность.
2. Архитектура Titans была разработана таким образом, чтобы объединить долговременную и кратковременную память.
   • Предлагаются три различных подхода к интеграции, обеспечивающие гибкие возможности для различных сценариев применения.
   • Синергия трех супернаправлений - ядра, долговременной памяти и постоянной памяти - позволяет модели более эффективно обрабатывать длинные последовательные данные.
3. Отлично справляется с несколькими задачами.
   • При моделировании языка, рассуждений, предсказании временных рядов и моделировании ДНК архитектура Titans демонстрирует высокую производительность, превосходящую существующие модели трансформаторов и линейных циклов.
4. Масштабируемость.
   • Способность поддерживать высокую производительность в большом контекстном окне открывает возможность обработки очень длинных последовательностей.

5, прогноз на будущее

Несмотря на то, что архитектура Titans позволила добиться впечатляющих результатов по ряду направлений, дальнейших исследований заслуживают следующие направления.

1. Изучение более сложных архитектур модулей памяти: исследование
   • Например, введение иерархических структур памяти или объединение модулей памяти с другими моделями, такими как графовые нейронные сети.
2. Разработка более эффективных механизмов обновления и хранения памяти.
   • Например, для уменьшения потребления памяти и вычислительных затрат используются методы разрежения или квантования.
3. Применение архитектуры Titans в более широком спектре областей.
   • В качестве примера можно привести понимание видео, управление роботами, рекомендательные системы и т.д.
4. Изучите более эффективные стратегии обучения.
   • Например, внедрение более совершенных алгоритмов оптимизации или использование метаобучения для ускорения обучения моделей.
5. Исследование интерпретируемости архитектуры титанов.
   • Более глубокое понимание того, как Титаны хранят и используют информацию в долговременной памяти, может дать новые идеи для создания более мощных систем искусственного интеллекта.

6, резюме

Основной вклад данной работы заключается в следующем:

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