Google представила Titans и MIRAS - новое поколение ИИ для сверхдлинных контекстов

Google представила новую архитектуру для работы с длинными последовательностями Titans и теоретическую основу MIRAS, которая объединяет скорость рекуррентных сетей с точностью трансформеров. Компания утверждает, что такой подход помогает моделям лучше запоминать важные детали на очень длинных контекстах - от больших документов до геномных данных.

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

Исследовательское сообщество уже пробовало обойти это ограничение через более быстрые линейные архитектуры - например, эффективные RNN и модели класса state space models вроде Mamba-2. Они сжимают прошлый контекст в вектор фиксированного размера, благодаря чему вычислительные затраты растут линейно. Но именно из-за жесткого ограничения по размеру памяти такие модели плохо справляются с действительно богатыми и длинными последовательностями, где важно не потерять тонкие связи между фактами.

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

Главная идея Titans в том, что у модели есть несколько уровней памяти, похожих на краткосрочную и долговременную у человека. Краткосрочную часть по-прежнему обеспечивает механизм внимания, который хорошо работает на локальных отрезках текста. А вот долговременная память реализована не как простой вектор или матрица фиксированного размера, а как полноценная глубинная нейросеть - многослойный перцептрон. Благодаря этому модель может гораздо богаче кодировать прошлый контекст и "понимать" историю, а не просто хранить набор заметок.

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

Чтобы система не реагировала только на разовые всплески, в Titans добавили два элемента. Во-первых, "моментум" - модель учитывает не только текущий сюрприз, но и недавнюю историю, чтобы захватывать цепочки связанных событий. Во-вторых, механизм забывания через адаптивный weight decay. Это некий "клапан", который не дает памяти переполняться при работе с экстремально длинными последовательностями и постепенно вытесняет устаревшую информацию.

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

Отдельный акцент Google делает на том, что MIRAS выходит за рамки привычных евклидовых метрик вроде среднеквадратичной ошибки и скалярного произведения. В большинстве современных моделей именно MSE и dot-product лежат в основе и "пристрастий" модели, и механизмов регуляризации памяти. Это удобно, но делает систему чувствительной к выбросам и ограничивает класс допустимых решений. MIRAS предлагает более богатое пространство целей и регуляризаторов, опираясь на идеи из оптимизации и статистики, и таким образом открывает путь к архитектурам с нестандартными, неевклидовыми функциями потерь.

В рамках этой рамки исследователи Google разработали три конкретные модели без явного внимания - YAAD, MONETA и MEMORA. YAAD использует более мягкое наказание за ошибки через Huber loss и лучше переносит шумные данные, где возможны опечатки или единичные артефакты. MONETA экспериментирует с более сложными нормами и проверяет, могут ли строгие математические "правила" для запоминания и забывания улучшить стабильность долгосрочной памяти. MEMORA пытается максимально стабилизировать память, заставляя ее вести себя как вероятностное распределение, чтобы каждое обновление состояния было контролируемым и аккуратным.

В экспериментах Titans и модели MIRAS сравнивали с современными архитектурами вроде Transformer++, Mamba-2 и Gated DeltaNet. Тестирование проводили на стандартных датасетах для языкового моделирования, таких как C4 и WikiText, а также на задачах нулевого шага по здравому смыслу, включая HellaSwag и PIQA. По заявлениям Google, новые модели показывают более низкую перплексию и лучшую точность по сравнению с базовыми архитектурами сопоставимого размера, при этом сохраняя линейную масштабируемость и возможность параллельного обучения.

Особенно заметно преимущество Titans в задачах с экстремально длинным контекстом. На бенчмарке BABILong, где нужно делать выводы на основе фактов, разбросанных по очень длинным документам, архитектура Titans обгоняет все baseline-модели, включая крупные системы вроде GPT-4, хотя сама содержит гораздо меньше параметров. Авторы также отмечают, что Titans удается масштабировать контекстное окно до величин свыше 2 миллионов токенов и при этом удерживать качество.

Отдельные исследования показали, что глубина модуля памяти играет критическую роль. При одинаковом "размере" памяти, но разной глубине, более глубокие варианты дают устойчиво меньшую перплексию и лучше масштабируются при росте длины последовательности. Это дополнительный аргумент в пользу того, что богатая, глубинная долговременная память действительно помогает моделям ориентироваться в очень длинных последовательностях.

В итоге Google позиционирует Titans и MIRAS как шаг к новому поколению моделей, которые умеют учиться "на лету", хранить важные детали на огромных промежутках контекста и при этом остаются достаточно эффективными для практического использования. Теоретическая рамка MIRAS связывает воедино онлайн-оптимизацию, ассоциативную память и архитектурный дизайн, а сама архитектура Titans демонстрирует, как можно совмещать скорость RNN с выразительной мощностью трансформеров в эпоху длинноконтекстного ИИ.