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

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

Стаття

Основна претензія до звичайного Retrieval-Augmented Generation полягає в його невибірковості: пошук фіксованої кількості уривків для кожного вводу, додавання їх на початок і генерація. Це працює досить добре, коли пошук допомагає, але активно шкодить, коли уривки не є релевантними або коли модель вже має відповідь у своїх вагах. У статті представлено Self-Reflective Retrieval-Augmented Generation (Self-RAG), авторами якого є Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil та Hannaneh Hajishirzi (Вашингтонський університет та IBM Research).

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

Навчання проходить у два етапи: спочатку модель-критик (LLaMA 2, 7B fine-tuned) навчається на приблизно 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 та FactScore біографій 80,2 — перевершуючи ChatGPT та Llama2-chat з доповненням пошуком у кожному з цих тестів.
• Абляції на PopQA показують, що видалення пошуку під час тестування коштує 20,8 відсоткових пунктів, тоді як видалення лише критика коштує лише 2,9 в.п. — пошук є опорним елементом; критика додає калібрування зверху.
• Під час інференсу ваги токенів критики можна регулювати, щоб знаходити баланс між точністю цитування та плавністю мовлення без будь-якого перенавчання. Це робить поведінку моделі конфігурованою для різних прикладних завдань.
• Програмний комітет ICLR 2024 надав Self-RAG статус Oral (топ 1%), що відображає справжнє визнання технічного внеску науковою спільнотою.

Що підтверджується, а що — ні

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

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

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

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

Чому це важливо для фінансового ШІ

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

Токени 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) — зосереджується саме на декомпозиції запитів для складних багатокрокових QA, що є сценарієм, з яким Self-RAG справляється найменш впевнено.
• FRAMES: Retrieval and Augmentation for Multi-Hop Evaluation (arXiv:2409.12941) — бенчмарк від Google DeepMind для RAG з багатьма документами, що вимагає поєднання кількох знайдених фактів; природний складніший тест для моделей у стилі Self-RAG.