Привет! Мы — команда Memory‑Augmented models в составе лаборатории Cognitive AI Systems AIRI. В ходе своих исследований мы стараемся разобраться, насколько хорошо LLM могут обрабатывать большой объем данных и решать задачи на основе них.
Разработчики современных языковых моделей соревнуются в длине контекста, и счёт уже идёт на миллионы токенов. Но насколько эффективно LLM пользуются информацией из этого контекста?
Чтобы выяснить это, мы вместе с коллегами из МФТИ и Лондонского института Математических Наук создали новый бенчмарк под названием BABILong, который привезли на NeurIPS в этом году. Он оценивает то, насколько успешно современные модели умеют искать информацию в собственных гигантских контекстах. Оказалось, что зачастую главное — это не размер, а умение пользоваться.
В этой статье расскажем подробнее о наших экспериментах, а также о том, как эффективно использовать длинный контекст.
Зачем языковым моделям большой контекст?
Контекст в языковых моделях — это размер текста, который модель может обработать одновременно. Чтобы решить задачу или ответить на вопрос, вся релевантная информация должна поместиться в контекст, чтобы модель смогла правильно ответить. Зачем современным языковым моделям увеличивать длину контекста? Давайте разберемся.
Первое и очевидное преимущество больших контекстов — это добавление в модель новых знаний. Представьте себе большую языковую модель, обученную однажды на гигантском массиве данных. Все знания, которые она извлекла из этих текстов, записаны в её параметры. Но что делать, если нам нужно, чтобы модель обрабатывала совершенно новую информацию, например, свежие финансовые отчёты, недавно вышедшие книги или новости?
Здесь‑то и приходит на помощь большой контекст. Мы можем загрузить эти новые данные прямо в контекст модели, и она сможет их использовать для ответов на вопросы и решения задач.
Другой важный аргумент в пользу длинного контекста — это демонстрация решений задач прямо внутри модели. Большие языковые модели отлично умеют учиться на лету через In‑Context Learning или обучение внутри контекста. Например, мы можем показать модели несколько примеров в формате «вход‑выход», а затем попросить её самостоятельно решить аналогичную задачу. Исследования показывают, что, чем больше примеров мы предоставляем в контексте, тем выше качество работы модели. Более того, иногда оно может соперничать с качеством полноценной дообученной модели!
Есть и ещё один интересный аспект: длинные контексты позволяют модели совершать пошаговое рассуждение. Например, прежде чем дать финальный ответ, модель может «проговорить» промежуточные шаги. Недавно выпущенная OpenAI модель o1, прежде чем ответить, строит длинные цепочки рассуждений. Чем больше длина контекста, тем сложнее размышления и, как следствие, более точный результат.
Ну и, конечно, не стоит забывать про работу с другими модальностями: кодом, видео, аудио. Все это требует расширенных контекстных возможностей. Ведь будущее языковых моделей — это универсальность, где текст — это лишь одна из форм взаимодействия.
Размер контекста у современных моделей: в чем проблема?
Контекст языковых моделей значительно вырос за последние годы. Этот путь отлично иллюстрирует график ниже:
Но как же удалось прийти к текущим, довольно‑таки гигантским размерам контекста, и какие трудности были на этом пути?
История начинается с трансформеров — архитектуры, которая была представлена в статье Attention is All You Need в 2017 году и стала основой для языковых моделей. Первые версии трансформеров умели обрабатывать меньше тысячи токенов — это всего пара страниц текста. Однако с тех пор длина контекста постоянно увеличивалась, а разработчики внедряли всё новые подходы, чтобы модели могли воспринимать всё больше данных.
На сегодняшний день крупные модели от ведущих разработчиков работают со стандартной длиной контекста в 128 тысяч токенов и более. Для масштаба: это примерно весь текст книги «Хоббит, или Туда и обратно». Более продвинутые модели вроде GPT-4 или Gemini 1.5 Pro уже справляются с миллионами токенов, что сопоставимо с текстом всех книг о Гарри Поттере. А самые амбициозные проекты начинают шагать за пределы 10 миллионов токенов — эквивалент полного собрания сочинений Стивена Кинга.
Цена длинного контекста
Разумеется, такое масштабирование не бесплатно. Главная проблема заложена в самой архитектуре трансформеров. Механизм внимания (attention), на котором основана эта архитектура, обладает квадратичной сложностью. Это означает, что увеличение длины контекста в 10 раз требует примерно в 100 раз больше вычислительных ресурсов для работы со слоями внимания.
Ещё одна техническая загвоздка — это так называемый key‑value кэш. Для обработки текста трансформер сохраняет состояния всех токенов, которые он уже пропустил через себя. С увеличением длины контекста этот кэш разрастается, съедая всё больше памяти.
Может ли GPT-4 найти иголку в стоге сена?
Увеличение длины контекста у языковых моделей — значительный шаг вперёд, но использование этого контекста оказалось сложнее, чем ожидалось. В июле прошлого года статья Lost in the Middle выявила удивительную проблему: то, где в контексте находится информация, критически влияет на способность модели её использовать.
Авторы статьи провели простой, но наглядный эксперимент: документ с ответом на вопрос помещался в начало, середину или конец контекста. Результаты оказались неожиданными: если нужный текст находился в начале или конце, модель справлялась с задачей неплохо. Однако если документ располагался в середине контекста, модель часто просто «теряла» его.
Бенчмарк Needle in a Haystack 🔎📚🪡📚❓
На основе этой идеи был предложен бенчмарк Needle in a Haystack. Его суть отражает название: «иголка» — это ключевая информация (в оригинальной версии — рекомендация для туристов в Сан‑Франциско), а «стог сена» — длинные документы (в данном случае эссе Пола Грэма). Идея бенчмарка — проверить, насколько хорошо модель может найти эту «иголку» в контексте, растянутом на десятки или сотни тысяч токенов.
Результаты тестирования моделей, например GPT-4, на этом бенчмарке показали характерный U‑образный график. Если ключевая информация находится ближе к началу или концу контекста, модель работает сравнительно хорошо. Но если информация расположена в середине, особенно при длинных контекстах, качество резко падает.
Важность Needle in a Haystack была быстро признана разработчиками крупных моделей. Например, для Claude 2 и Gemini 1.5 Pro, заявляющих поддержку сотен тысяч и миллионов токенов соответственно, такие бенчмарки стали своего рода «обязательной справкой» при релизе.
Однако это послужило разработчикам мотивацией включать в обучающую выборку схожие задачи, что быстро привело к переобучению на них и вырождению Needle in a Haystack. Научные публикации теперь часто сопровождаются тепловыми картами (heatmaps), демонстрирующими, что модель якобы не теряет информацию в длинных контекстах.
Однако реальность не всегда столь радужна: проблемы с пониманием «иголок» в длинных текстах остаются, особенно в сложных задачах. Чтобы это отследить, нужные другие несинтетические бенчмарки. Они есть (например, LongBench), но большинство из них не работает на контексте более 16 тысяч токенов.
Промежуточные выводы
-
Needle in a Haystack стал стандартным sanity check для языковых моделей. Если модель проваливается на этом бенчмарке, вероятно, она не справится и с другими задачами на длинных контекстах.
-
Современные большие языковые модели научились работать с контекстами длиной в сотни тысяч токенов благодаря инженерным улучшениям, но до полной надежности ещё далеко.
-
Возрастающие длины контекстов создают вызов для бенчмарков: они не успевают адаптироваться и не всегда позволяют объективно оценить качество работы моделей.
Мы тоже задумались над этими вопросами и предложили собственный подход, чтобы точнее оценивать возможности моделей в работе с большими контекстами.
BABILong — логический бенчмарк для длинных задач
Результатом нашей работы стал BABILong — усложнённая версия бенчмарка Needle in a Haystack, которая проверяет способность моделей выполнять задачи логического вывода в длинных контекстах.
Идея бенчмарка
В основе бенчмарка — задачи из датасета bAbI (2015), предложенного как тест для интеллектуальных систем, отсюда и название BABILong. Условие каждой задачи состоит из нескольких фактов, таких как перемещения персонажей («Мэри пошла в ванную», «Джон пошёл в холл»). Модели же требуется ответить на вопросы вроде: «Где находится Мэри?». Всего бенчмарк включает 20 задач, проверяющих выполнение базовых вычислений, объединение информации из нескольких фактов и логическое .
BABILong симулирует ситуацию, когда релевантная информация распределена по контексту. Предложения bAbI случайно распределяются в случайных местах нерелевантного длинного текста, выполняющего ролю шума. Мы собрали фоновый текст из художественной литературы — датасета PG19.
BABILong доступен на Hugging Face, а также можно сгенерировать версию с собственными задачами и шумом с помощью кода на GitHub.
Результаты
Мы протестировали на BABILong более 35 современных языковых моделей с длиной контекста более 4000 токенов, включая как проприетарные решения GPT-4o (128k) и Gemini (10M), так и свежие модели с открытым кодом — семейства Llama 3.1, Qwen 2.5 и другие. В таблице модели отсортированы по заявленной длине контекста, а цветом обозначена оценка, полученная в бенчмарке (доля верно решенных задач данной длины в %). Модели с «~» в начале — экспериментальные архитектуры, обученные на BABILong, о них поговорим дальше.
Получившаяся таблица очень наглядна, кратко её подытожим:
-
С увеличением длины контекста качество у всех моделей падает, хотя интенсивность деградации варьируется.
-
Большинство моделей эффективно работают лишь с 10–20% от заявленной длины контекста.
-
Даже лучшие модели, такие как Gemini 1.5 и Qwen 2.5 показывают падение качества с 80-90% до 60% и меньше на более длинных задачах.
BABILong демонстрирует, что работа даже с простыми задачами с длинными контекстами — всё ещё сложна даже для лидирующих моделей. Заявленный контекст в сотни тысяч токенов вовсе не означает, что модель способна его воспринять.
А как же RAG?
Одним из наиболее популярных методов расширения контекста LLM является Retrieval‑Augmented Generation (RAG). В этом подходе длинный контекст разбивается на сегменты или предложения, каждое из которых кодируется, после чего по запросу выбираются и подаются в контекст только релевантные для задачи тексты.
Мы протестировали на BABILong два различных RAG‑пайплайна. Первый основан на LangChain, где для ответа на вопросы применялась модель GPT-4, а эмбеддинги создавались с помощью OpenAI. Второй пайплайн использовал Llama 3 с эмбеддингами от NVIDIA, следуя аналогичному процессу.
Результаты показали, что качество поиска релевантных предложений (retrieval quality) на задаче QA1 стабильно высокое. Однако простой механизм семантического поиска, используемый в RAG, не способен сохранять информацию о порядке фактов. Это приводит к низкой точности даже на коротких контекстах. Интересно, что качество RAG стабильно даже при увеличении размера контекста и превосходит GPT-4 на длинах более 32к токенов!
Ключевая сложность RAG заключается в том, что за одну итерацию трудно извлечь всю необходимую информацию из контекста, что влияет на его производительность в более сложных сценариях.
А можно ли вообще решить BABILong?
Одно из перспективных направлений, которое мы исследуем, — это рекуррентность на уровне сегментов. Представьте, что длинный текст разбивается на сегменты (например, по 512 токенов), которые подаются на вход трансформеру. Вместе с сегментами передаются токены памяти, обновляющиеся на каждом шаге. Затем обновленные токены памяти передаются следующему сегменту, превращая трансформер в рекуррентную сеть. В токенах памяти сохраняется информация о прошлом, добавляя в трансформер механизм долгосрочной памяти. На этом принципе основана наша модель Recurrent Memory Transformer (RMT).
Мы протестировали RMT, основанный на GPT-2, обучив его на задачах из BABILong с длиной контекста в 16 000 токенов (32 сегмента по 512). На инференсе модель показала впечатляющие результаты: до длины контекста 100 000 токенов качество практически не снижается, а при 1 млн токенов достигает 80%. Это лучше, чем у большинства подходов RAG. Кроме того, RMT можно запускать на обычных видеокартах, например, GTX 1080 Ti.
Мы также предложили улучшенную архитектуру — ассоциативный RMT. Здесь память отдельная для каждого слоя и дополнена ассоциативным блоком. Этот блок использует матрицу памяти, куда данные записываются через внешнее произведение ключей и значений, а обновление происходит с учетом гейта, удаляющего устаревшую информацию.
Результаты впечатляют: улучшенная модель поддерживает качество 80% вплоть до длины контекста 50 миллионов токенов. Это в пять раз больше, чем у базового RMT, и ставит ассоциативный RMT в лидеры по работе с экстремально длинными контекстами.
Выводы
Подводя итоги, хочется отметить следующее. Стандартные бенчмарки, такие, как Needle in a Haystack, не могут полноценно оценить растущий контекст больших языковых моделей. BABILong же позволяет более точно выявлять сильные и слабые стороны моделей при работе с длинными задачами. И такая проверка показывает, что простые методы «растяжения» контекста моделей не дают значительных улучшений. Чтобы справляться с большими контекстами, необходимы более сложные и специализированные архитектуры.
Базовый и очень популярный RAG не справляется с задачами BABILong. Важно улучшать методы Retrieval‑Augmented подходов, внедряя итеративный retrieval и комбинируя их с рекуррентными архитектурами, например, — RMT. Такой подход позволяет более эффективно работать с информацией из прошлого, извлекая её более точно, и сохраняя информацию на протяжении десятков миллионов токенов.
Детально наши эксперименты описаны в статье, которую мы в прошлую среду представляли на NeurIPS 2024.
Будем рады ответить на ваши вопросы!
Автор: booydar
