Проблема: длинный контекст упирается в память
Чем больше текста подаётся модели на вход, тем тяжелее ей обрабатывать каждую новую подсказку. Классический подход in-context learning требует, чтобы модель «перечитывала» весь контекст при каждом запросе. Это означает, что и время ответа, и потребление GPU-памяти растут вместе с длиной входного текста. На практике это заметно уже при десятках тысяч токенов, а на 128K токенов ресурсы расходуются крайне неэффективно.
Хуже того, у обычных трансформеров есть нативное окно контекста, за пределами которого качество резко падает. Даже если модель формально принимает длинный вход, она начинает «забывать» детали из середины текста. Особенно наглядно это видно на задачах вроде needle-in-a-haystack, когда нужно найти один важный факт в большом объёме информации.
Идея MoNe: вынести контекст в отдельную память
Метод MoNe (Modular Neural Memory) предлагает избавиться от прямой зависимости между длиной контекста и стоимостью запроса. Это лёгкая модульная нейронная память, которая подключается к уже готовой, замороженной трансформер-модели. Дообучать базовую модель не нужно — достаточно добавить внешний модуль, который берёт на себя работу с контекстом.
Ключевая хитрость в том, как именно память устроена. MoNe обрабатывает входной текст не целиком, а сегментами фиксированного размера. Внутри модуля используются сети быстрых весов — они обучаются прямо на этапе тестирования, обновляя свои параметры локальными градиентами на уровне слоёв. Это звучит сложно, но по сути означает: память адаптируется к конкретному тексту без вмешательства в веса основной модели.
Двухфазная архитектура: отдельно предобработка, отдельно запрос
Инференс разбит на два независимых этапа. На первом этапе выполняется предобработка: модель читает контекст сегментами, формирует из них ключи и значения и сохраняет их во внешней памяти. Этот этап линейный — стоимость растёт пропорционально длине текста, но это разовая операция.
На втором этапе — когда приходит запрос — память уже не обращается к исходному контексту. Вместо этого она генерирует ключи и значения только из токенов запроса. Контекстные токены повторно не считываются, а значит, время ответа вообще не зависит от длины исходного текста. Формально это выражается как O(N) на предобработку и O(1) на запрос, где N — длина контекста.
Благодаря такому разделению пиковое использование GPU-памяти перестаёт расти с ростом N. Это важное преимущество: даже при очень длинном контексте модель не пытается одновременно удерживать в памяти все входные токены.
Что это даёт на практике
Авторы проверили подход на бенчмарке RULER, который включает сценарии с «иголкой в стоге сена» и извлечением конкретных слов. Именно там стандартный in-context learning заметно деградирует при увеличении длины контекста. MoNe, напротив, показывает стабильные результаты и обобщается на длины, которые значительно превышают нативное окно базовой модели.
Отдельно стоит подчеркнуть эффективность. При работе со 128K токенов MoNe сокращает вычислительные затраты и пиковую GPU-память примерно на 80% по сравнению с обычным ICL. При этом дополнительные параметры модуля составляют всего 6.4% — то есть накладные расходы на память для хранения самой модели минимальны.
Итоги
MoNe — это свежий взгляд на проблему длинного контекста. Вместо того чтобы наращивать вычислительные мощности или придумывать сложные механизмы внимания, исследователи предлагают вынести работу с контекстом в отдельный модуль. Он не требует дообучения трансформера, добавляет немного параметров, зато делает ответ на запрос константным по времени и не даёт пиковой памяти расти вместе с текстом.
Пока это препринт на arXiv от международной группы исследователей. Но направление выглядит перспективным: если метод подтвердится на более широком спектре задач, он может стать практичным способом превращать обычные замороженные модели в системы, комфортно работающие с очень длинными документами без серьёзных изменений в их архитектуре.
