MemAgent: Революция в обработке длинных контекстов в LLM с помощью обучения с подкреплением

Проблемы обработки длинных документов в LLM

Большие языковые модели (LLM) сталкиваются с серьёзными трудностями при работе с очень длинными документами. Традиционные методы, такие как экстраполяция длины и разреженное внимание, часто приводят к снижению производительности или высоким вычислительным затратам.

Существующие подходы и их ограничения

Современные методы моделирования длинного контекста делятся на три категории:

• Методы экстраполяции длины (NTK, PI, YaRN, DCA): расширяют окно контекста за счёт манипуляций с позиционными эмбеддингами, но страдают от проблем масштабирования и снижения качества.
• Разреженное и линейное внимание: снижают сложность внимания до O(n), но требуют дообучения и зависят от фиксированных или вручную заданных шаблонов.
• Сжатие контекста: используют токен-уровневые или внешние модули памяти для сокращения входных данных, но часто ухудшают качество генерации и плохо работают с очень длинными текстами. Ни один из этих методов не обеспечивает одновременно поддержку неограниченной длины, стабильную точность и эффективную линейную сложность.

MemAgent: стратегия памяти, похожая на человеческую

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

Основные особенности:

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

Обучение с использованием многоконтекстного обучения с подкреплением (GRPO)

MemAgent рассматривает взаимодействия с каждой частью документа как независимые диалоги. Используется Group Relative Policy Optimization (GRPO) в рамках многодиалогового RL-пайплайна DAPO для управления обновлением памяти на основе наград.

Ключевые компоненты:

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

Результаты работы

MemAgent был обучен на контексте длиной 8 тысяч токенов и успешно экстраполировал до 3,5 миллионов токенов, используя бенчмарки RULER и синтетические датасеты HotpotQA и SQuAD.

Он поддерживал точность выше 95% на RULER (от 8K до 512K токенов) и стабильно превосходил другие методы с длинным контекстом и дистилляционные модели.

Пример: многозадачный вопрос-ответ

Для вопроса «Режиссёр романтической комедии «Big Stone Gap» базируется в каком районе Нью-Йорка?» MemAgent обработал три части текста, успешно выделил релевантную информацию и проигнорировал нерелевантные данные, корректно обновляя память. Итоговый ответ: Гринвич-Виллидж, Нью-Йорк.

Теоретические основы и сложность

MemAgent формулирует авторегрессионную модель с латентными переменными памяти (m₁…mₖ):

p(x_{1:N}) = \sum_{m_{1:k}} \prod_k p(c_k | m_{k-1}) \cdot p(m_k | c_k, m_{k-1})

Это обеспечивает вычислительную сложность O(N) и позволяет получать человекочитаемые промежуточные состояния памяти. RL необходим, поскольку дискретные обновления памяти нельзя оптимизировать через обратное распространение ошибки.

Применение и значение

MemAgent подходит для любых Transformer-моделей без изменений архитектуры. Его можно использовать для работы с длинными документами, систем памяти агентов, обзора юридических документов, анализа научных публикаций и принятия решений в реальном времени на основе больших объёмов данных.

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