Google Titan: новая архитектура для преодоления ограничений памяти Transformer

Представляем Titan: новую архитектуру от Google

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

Проблема памяти в существующих моделях

Традиционные модели, такие как LSTM и Transformer, хотя и являются инновационными, сталкиваются с проблемами в моделировании человеческой памяти. Эти проблемы включают:

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

Подход Titan: нейро-вдохновленный модуль памяти

Команда Titan применила другой подход, стремясь закодировать информацию в параметры нейронной сети. Они разработали онлайн-метамодель, предназначенную для обучения тому, как запоминать и забывать конкретные данные во время тестирования. Эта модель вдохновлена нейропсихологическими принципами, включая следующие ключевые элементы:

Неожиданность как триггер: Неожиданные события легче запоминаются. "Неожиданность" измеряется градиентом входных данных в модуль памяти. Чем больше градиент, тем неожиданнее вход.
Механизмы импульса и забывания: Механизм импульса накапливает кратковременные неожиданности в долговременную память, а механизм забывания стирает старые воспоминания, предотвращая переполнение памяти.
Память на основе многослойного перцептрона (MLP): Модуль памяти состоит из нескольких слоев MLP, что позволяет ему хранить глубокие абстракции данных, делая его более мощным, чем традиционные матрицы памяти.

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

Интеграция модуля памяти в архитектуры глубокого обучения

Исследовательская группа Titan предложила три варианта интеграции своего модуля памяти в архитектуры глубокого обучения:

MAC (Memory as Context): Этот метод объединяет долгосрочную и постоянную память (которая кодирует знания о задаче) в качестве контекста, который вводится в механизм внимания.
MAG (Memory as Gate): Этот подход использует управляемое слияние модуля памяти с механизмом внимания скользящего окна по двум ветвям.
MAL (Memory as Layer): Здесь модуль памяти реализован как независимый слой, который сжимает историческую информацию перед подачей ее в механизм внимания.

Команда обнаружила, что каждый вариант имеет свои сильные и слабые стороны.

Производительность и преимущества Titan

Titan продемонстрировал превосходную производительность в различных задачах, включая языковое моделирование, рассуждения на основе здравого смысла и прогнозирование временных рядов. Он превзошел современные модели, такие как Transformer и Mamba. Примечательно, что один только модуль долговременной памяти (LMM) превзошел базовые модели в нескольких задачах, демонстрируя свои независимые возможности обучения без кратковременной памяти (внимания).

В тесте "иголка в стоге сена", предназначенном для поиска мелких деталей в длинных текстах, Titan сохранял точность около 90% даже при увеличении длины последовательности с 2k до 16k. Команда отмечает, что стандартные тесты не полностью демонстрируют преимущества Titan в обработке длинных текстов. Titan также превзошел такие модели, как GPT4, Mamba и даже Llama3.1 с RAG, в задаче, требующей вывода из фактов, разбросанных по очень длинным документам.

Titan также показал впечатляющие результаты в конкретных областях, таких как прогнозирование временных рядов и моделирование последовательности ДНК.

Команда разработчиков Titan

Исследование проводилось командой из группы алгоритмов и оптимизации Google Research NYC, которая в настоящее время не входит в Google DeepMind.

Али Бехруз, стажер из Корнельского университета, является первым автором статьи.
Чжун Пейлин, выпускник Университета Цинхуа и доктор философии Колумбийского университета, является научным сотрудником Google с 2021 года. Он известен тем, что опубликовал статью в качестве первого автора на STOC 2016, будучи студентом бакалавриата.
Вахаб Миррокни, научный сотрудник и вице-президент Google, возглавляет команду.

Команда разработала Titan, используя Pytorch и Jax, и планирует в ближайшее время выпустить код для обучения и оценки.