Self-RAG: адаптивный поиск и самокритика для LLM

Стандартный RAG выполняет поиск каждый раз, независимо от того, помогает ли это. Self-RAG от Asai и др. (ICLR 2024 Oral) задает другой вопрос: что если модель сама будет решать, когда искать информацию, а затем оценивать результат? Оказывается, это имеет большое значение, а сам механизм достаточно понятен, чтобы изучить его детально.

Статья

Основная претензия к «ванильному» RAG (Retrieval-Augmented Generation) заключается в его неразборчивости: поиск фиксированного количества фрагментов для каждого запроса, их добавление в контекст и последующая генерация. Это работает неплохо, когда поиск помогает, но активно мешает, когда фрагменты не релевантны или когда у модели уже есть ответ в ее весах. В статье представлен Self-Reflective Retrieval-Augmented Generation (Self-RAG), авторами которого являются Акари Асаи, Цзэцю У, Ичжун Ван, Авируп Сил и Ханнане Хаджиширзи (Вашингтонский университет и IBM Research).

Ключевой механизм — это набор из четырех специальных токенов рефлексии, встроенных в словарь модели на этапе обучения. Retrieve решает, нужно ли вообще обращаться к поиску. IsRel (relevance) оценивает, содержит ли извлеченный фрагмент полезную информацию для запроса. IsSup (support) проверяет, подтверждено ли сгенерированное утверждение фрагментом полностью, частично или не подтверждено вовсе. IsUse (utility) оценивает общее качество ответа по шкале от 1 до 5. Модель учится выдавать эти токены вместе с обычным текстом, тем самым критикуя собственный поиск и генерацию за один прямой проход.

Обучение проходит в два этапа: сначала модель-критик (дообученная LLaMA 2, 7B) обучается на примерно 4 000–20 000 размеченных примеров для каждого типа токенов, достигая более 90% совпадения с предсказаниями GPT-4. Затем этот критик в автономном режиме аннотирует корпус из 150 000 примеров «инструкция-ответ», и генератор обучается на этих аннотированных данных, где токены рефлексии рассматриваются как обычные слова из словаря. Обучение с подкреплением не требуется.

Ключевые идеи

• Четыре токена рефлексии (Retrieve, IsRel, IsSup, IsUse) обеспечивают модели структурированный внутренний диалог о том, стоит ли доверять доказательствам, а не просто бинарное решение «искать/не искать».
• Self-RAG 13B достигает 55,8% на PopQA, 69,3% на TriviaQA, 74,5% на PubHealth, 73,1% на ARC-Challenge и 80,2 FactScore на биографиях — опережая ChatGPT и Llama2-chat с дополненной выборкой в каждом из этих тестов.
• Исследования абляции на PopQA показывают, что отключение поиска на этапе тестирования снижает результат на 20,8 процентных пункта, тогда как отключение только критика — всего на 2,9 п.п. Это означает, что основной вклад вносит поиск; критика добавляет калибровку.
• При инференсе веса токенов критики можно настраивать, чтобы сбалансировать точность цитирования и плавность речи без переобучения. Это делает поведение модели настраиваемым для различных прикладных задач.
• Программный комитет ICLR 2024 присвоил Self-RAG статус Oral (топ-1%), что отражает признание технического вклада научным сообществом.

Что подтверждается, а что — нет

Результаты абляции выглядят убедительно. Разрыв между режимами «всегда искать» и «никогда не искать» велик (20,8 п.п.); модель явно научилась отличать полезный поиск от шума. Токены IsRel и IsSup добавляют измеримую ценность сверх одного лишь адаптивного поиска. Это значимый результат, а не просто переформулировка.

В чем я менее уверен, так это в заявлениях об обобщающей способности. Все пять задач оценки (PopQA, TriviaQA, PubHealth, ARC-Challenge, ASQA) представляют собой краткие ответы на вопросы или выбор из вариантов — именно те условия, где один извлеченный фрагмент может дать решающий сигнал. Генерация длинных текстов на основе контекста из нескольких документов, что характерно для задач в сфере финансов, изучена меньше. FactScore для биографий (80,2) — наиболее близкий показатель, но биографии относительно хорошо структурированы по сравнению с разветвленным многолетним гроссбухом расходов.

Есть также нюанс с воспроизводимостью: обучающие метки для модели-критика получены от GPT-4. Это ставит качество разметки в зависимость от проприетарной системы и влечет за собой затраты на API, которые не указаны. Позже в CRAG (arXiv:2401.15884) было показано, что оценщик поиска размером 0,77B — гораздо более легкий, чем критик Self-RAG на 7B — может корректировать качество поиска и давать прирост в 19,0 п.п. по сравнению со стандартным RAG на PopQA. Это серьезный вызов архитектуре, даже если сама идея селективного поиска верна.

Наконец, важен базовый уровень для сравнения. Победа над ChatGPT (вероятно, GPT-3.5-turbo конца 2023 года) и Llama2-chat — достойная планка для открытой модели 13B, но передовые модели с тех пор значительно продвинулись. В статье не рассматривается вопрос о том, превзойдет ли адаптивный поиск Self-RAG хорошо спроектированный промпт для GPT-4o с простой схемой «всегда искать» на тех же бенчмарках.

Почему это важно для финансового ИИ

Финансовые агенты, работающие с гроссбухами Beancount, сталкиваются именно с той проблемой избирательности поиска, которую решает Self-RAG. Когда пользователь спрашивает: «каков мой чистый доход за этот месяц?», агент может произвести расчет на основе загруженного контекста — поиск может лишь добавить шума. Когда тот же пользователь спрашивает: «записал ли я счет подрядчика за третий квартал?», агенту нужно просмотреть записи потенциально за несколько лет. Режим «всегда искать» забивает контекст и рискует внести нерелевантные старые транзакции; режим «никогда не искать» пропустит нужные данные.

Токены IsRel и IsSup четко ложатся на логику валидации гроссбуха. IsRel: действительно ли извлеченная запись о транзакции относится к запросу? IsSup: действительно ли извлеченный контекст подтверждает сгенерированную сумму баланса, или число галлюцинировано? Оценка полезности (1–5) может служить индикатором уверенности при записи данных: фиксировать предложенную журнальную запись только тогда, когда модель оценивает собственные рассуждения на 4 или 5, а остальные помечать для проверки человеком.

Проблема воспроизводимости здесь также важна. Для промышленного бухгалтерского агента зависимость от GPT-4 для генерации обучающих меток является операционным ограничением. Если более легкий оценщик (как в CRAG) может обеспечить сопоставимый селективный поиск, то этот путь более предпочтителен для внедрения. Принципы проектирования Self-RAG — «решить перед поиском, оценить после поиска» — остаются ценными, даже если конкретный рецепт обучения токенов будет заменен.

Что читать дальше

• CRAG: Corrective Retrieval Augmented Generation (arXiv:2401.15884) — развивает идею адаптивного поиска Self-RAG с более легким оценщиком и резервным веб-поиском на случай неудачи локального поиска; стоит сравнить напрямую с Self-RAG на пересекающихся бенчмарках.
• RQ-RAG: Learning to Refine Queries for Retrieval Augmented Generation (arXiv:2404.00610) — фокусируется именно на декомпозиции запросов для сложных многоходовых ответов на вопросы, с чем Self-RAG справляется хуже всего.
• FRAMES: Retrieval and Augmentation for Multi-Hop Evaluation (arXiv:2409.12941) — бенчмарк от Google DeepMind для многодокументного RAG, требующий связывания нескольких извлеченных фактов; естественное и более сложное испытание для моделей типа Self-RAG.