RAG мёртв? Хайп вокруг 1M-моделей

На Реддите и в твиттере время от времени появляются такие посты

И такие

У людей, которые уже успели повыкатывать приложения с LLM под капотом в продакшн, немедленно начинает бомбить в комментариях. Почему? Давайте разбираться.

Эра большого контекста

Проблему увеличения контекста жаждут решить с самого появления архитектуры трансформера. В отличие от RNN, которые в теории могут проносить контекст сквозь всю последовательность независимо от её длины, трансформеры ограничены тем контекстом, который может поместиться в вычислительную память. Большой документ можно разбить на куски, обрабатывать их отдельно и потом как-то агрегировать результаты, но такой способ подойдёт не для всех задач. В бенчмарке RULER, созданном специально для оценки качества работы LLM на длинном контексте, можно видеть, насколько сильно падает качество работы моделей, которые обучались только на коротком контексте без специальных архитектурных трюков.

Одной из заметных попыток решения проблемы контекста в этеншн-моделях в своё время стала статья Transformer XL. Её основные идеи:

• Давайте дробить длинный текст на сегменты, но будем кэшировать hidden state всех предыдущих кусков. Градиент до них доходить не будет, но по крайней мере у сети будет вся информация из предыдущих сегментов текста.
• Нужно использовать относительный, а не абсолютный positional encoding, иначе модель не будет нормально работать на документах длиннее, чем в трейн-сете.

При таком подходе память рано или поздно всё равно закончится - ведь мы кэшируем hidden state всех предыдущих сегментов. Метод Infini-Attention, который вроде как использовался в первой версии Gemini с миллионным контекстом, теоретически может работать с бесконечным контекстом за счёт использования компрессивной памяти фиксированного размера. Идея концептуально напоминает RNN-модели:

• Последовательность делится на сегменты фиксированного размера - скажем, 2048 токенов
• Вместо хранения всех предыдущих ключей (K) и значений (V), мы используем "память" фиксированного размера, которая модифицируется на основе текущих K и V. Если информация из текущих KV не сильно отличается от уже сохранённой, то память можно особо не обновлять, нам интересна только новая информация
• При калькуляции этеншна для текущего токена используется и локальный этеншн, и этеншн c памятью. Они агрегируются простым суммированием с весами. По итогам обучения по этим весам можно увидеть, что образуется два типа этеншн-голов - "специализированные", которые берут информацию только из локального контекста или только из памяти, и "миксованные", которые смешивают информацию

Ребята из HF решили проверить, смогут ли они воспроизвести впечатляющие результаты из этой статьи. Первые эксперименты показали, что модель в принципе использует память, но не справляется с задачкой поиска спрятанного пароля в раннем сегменте. Кроме того, все этеншн-головы получились "миксованные", а не "специализированные". Разные трюки и подкрутки гиперпараметров улучшили ситуацию, но в целом авторы делают вывод, что при обучении длинноконтекстных LLM лучше пока придерживаться стандартных, проверенных техник.

Именно так и поступила команда Qwen при обучении Qwen2.5-1M. Простой, рабоче-крестьянский и надёжный способ улучшения работы моделей на задачах, требующих длинного контекста - дообучение ровно на таких задачах. Как же добыть такие данные? Квеновцы придумали прикольный способ:

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

Только дообучения, однако, недостаточно для достижения по-настоящему больших контекстов. Обучение на 256k токенов - уже очень недешёвое удовольствие, как же "раздвинуть" контекст до миллиона токенов?

Одна из важных причин деградации качества моделей на длинных контекстах - необученные позиционные эмбеддинги. Модель знает, что значит, что токены находятся на расстоянии 250k друг от друга, но никогда не видела токены, которые находятся в 700k токенах. Традиционный метод исправления ситуации - это Position Interpolation. Идея проста - сжимаем все индексы таким образом, чтобы они ложились в исходный диапазон. Рекомендуется короткий файн-тюнинг для адаптации модели к интерполированным эмбеддингам.

В Qwen2.5-1M используется Dual Chunk Attention, которая вообще не требует никакого дообучения. Длинная последовательность делится на чанки размером меньше, чем максимальный контекст модели при обучении. Внутри чанка используется стандартный этеншн с обычным RoPE. Для остальных чанков (кроме соседнего) используется константная большая позиция (например, максимальный контекст модели), чтобы модель "понимала", что это какие-то дальние куски текста.

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

Даже если модель успешно справляется с большим контекстом, но инференс требует нереальной кучи GPU, то в реальной жизни такую модель не поиспользуешь. Для оптимизации инференса в Qwen2.5-1M интегрированы идеи из MInference. Например, если посмотреть на веса этеншна LLM, можно увидеть несколько интересных паттернов. Токены часто обращают внимание:

• на начало последовательности и на соседние токены. Инструкция или описание задачи обычно располагаются в начале, а близкие токены часто взаимосвязаны. Это статический паттерн, он постоянен для любой последовательности.
• на некоторые ключевые токены и на токены, расположенные на равных фиксированных интервалах от текущего (здесь и здесь можно почитать, почему при обучении образуется такой паттерн). Это динамический паттерн, поэтому для вычисления важных токенов используются последние n векторов запросов (Q) и ключей (K).
• на определённые блоки из всей матрицы внимания. Это самый динамический паттерн.

Эти паттерны позволяют значительно ускорить prefill-стадию работы LLMи. Кстати, есть вот такой балдёжный репозиторий от NVIDIA, в котором собраны и реализованые разные способы компресии KV-кэша, что очень важно как раз для работы с большими контекстами.

Об остальных фишках можно почитать в их отличном техническом репорте.

Минусы больших размеров

Ну так что же, RAG отправляется на свалку истории или нет? Кажется, нет:

• В куче реальных бизнес-кейсов объём документов, с которыми надо работать, сильно превышает даже 1 миллион токенов. Даже в небольших компаниях все чаты, базы данных, код, документация легко переваливают за эту цифру. Это я ещё не учитываю документы из интернета.
• Огромная стоимость инференса с большим контекстом. На одной RTX 4090 24GB максимальный контекст 14-миллиардной модели в квантизации Q6_K составляет порядка 48к токенов. У всех есть кластер из H100 под рукой?
• Скорость инференса - вытащить даже с помощью какого-нибудь хитрого RAG на стероидах 10 релевантных документов и загнать их в LLM сильно быстрее, чем промалывать миллион токенов.
• Реальный контекст моделей на реальных задачах всегда сильно-сильно меньше заявленного. Часто контекстные бенчмарки используют искусственные задачи типа поиска "иголки в стоге сена", а в недавней статье решили немного усложнить задачу. К примеру для иголки "Жека работает у БКЗ на Лиговском" вопрос будет звучать не "Кто из упомянутых людей работает у БКЗ на Лиговском?", а "Кто из упомянутых людей живёт в Санкт-Петербурге?". Модели нужно не просто найти фрагмент в длинном тексте, а ещё и сделать логический переход ("БКЗ на Лиговском, скорее всего, находится в Петербурге"). Результаты, скажем так, впечатляющие:

• На практике тоже всегда заметно, что на контекстах даже в 50% от заявленного начинается явная деградация качества и забывание.

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

Вывод простой и вечный - всегда думайте, а что за задачу я решаю, какой инструмент тут подойдёт больше? К примеру, для задачи саммеризации жирной взаимосвязанной кодовой базы большой контекст действительно важен. В нашей ИИ-вики используется и RAG (для поиска нужных для ответа на вопрос источников), и большой контекст (для саммеризации больших документов при индексации).