Self-RAG: Adaptívne vyhľadávanie a sebakritika pre LLM

Štandardný RAG vyhľadáva zakaždým, bez ohľadu na to, či to pomôže. Self-RAG od Asai et al. (ICLR 2024 Oral) si kladie inú otázku: čo keby sa samotný model rozhodol, kedy niečo vyhľadať, a potom ohodnotil výsledok? Ukazuje sa, že na tom dosť záleží, a tento mechanizmus je dostatočne čistý na to, aby stál za podrobné preštudovanie.

Článok

Hlavnou výhradou voči základnému generovaniu rozšírenému o vyhľadávanie (RAG) je, že nerozlišuje: vyhľadá fixný počet pasáží pre každý vstup, pripojí ich a generuje. To funguje dostatočne dobre, keď vyhľadávanie pomáha, ale aktívne to škodí, keď sú pasáže nerelevantné alebo keď model už má odpoveď vo svojich váhach. Článok predstavuje Seba-reflexívne generovanie rozšírené o vyhľadávanie (Self-RAG), ktorého autormi sú Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil a Hannaneh Hajishirzi (University of Washington a IBM Research).

Kľúčovým mechanizmom je sada štyroch špeciálnych reflexných tokenov zapracovaných do slovníka modelu počas trénovania. Retrieve rozhoduje, či vôbec volať vyhľadávač. IsRel (relevancia) posudzuje, či vyhľadaná pasáž skutočne obsahuje užitočné informácie pre dopyt. IsSup (podpora) kontroluje, či je vygenerované tvrdenie plne, čiastočne alebo vôbec nepodložené pasážou. IsUse (užitočnosť) hodnotí celkovú kvalitu odpovede na stupnici od 1 do 5. Model sa naučí emitovať tieto tokeny priamo vo svojom bežnom výstupe — takže kritizuje vlastné vyhľadávanie a generovanie v rámci jedného dopredného prechodu.

Trénovanie je dvojfázové: najprv sa natrénuje model kritika (LLaMA 2, 7B fine-tuned) na približne 4 000 – 20 000 označených príkladoch na každý typ tokenu, čím dosiahne viac ako 90 % zhodu s predikciami GPT-4. Tento kritik potom offline anotuje korpus inštrukcií a výstupov so 150 000 príkladmi a generátor sa následne trénuje na týchto anotovaných dátach, pričom s reflexnými tokenmi sa zaobchádza ako s bežnou slovnou zásobou. Nie je potrebné žiadne posilňované učenie (RL).

Kľúčové myšlienky

• Štyri reflexné tokeny (Retrieve, IsRel, IsSup, IsUse) dávajú modelu štruktúrovaný vnútorný dialóg o tom, či sa dôkazom oplatí dôverovať — nejde len o binárne rozhodnutie vyhľadať/nevyhľadať.
• Self-RAG 13B dosahuje 55,8 % v PopQA, 69,3 % v TriviaQA, 74,5 % v PubHealth, 73,1 % v ARC-Challenge a Biography FactScore 80,2 — čím v každom z nich prekonáva ChatGPT a Llama2-chat s rozšíreným vyhľadávaním.
• Ablačné štúdie na PopQA ukazujú, že odstránenie vyhľadávania v čase testovania stojí 20,8 percentuálneho bodu, zatiaľ čo odstránenie samotného kritika stojí len 2,9 p. b. — vyhľadávač je kľúčovým nosným prvkom, kritika pridáva nadstavbu v podobe kalibrácie.
• Pri inferencii je možné upraviť váhy tokenov kritiky, aby sa dosiahol kompromis medzi presnosťou citácií a plynulosťou bez nutnosti pretrénovania. Vďaka tomu je správanie modelu konfigurovateľné pre rôzne následné aplikácie.
• Programový výbor ICLR 2024 udelil Self-RAG štatút "oral" (medzi top 1 %), čo odráža skutočné odborné uznanie technického prínosu.

Čo obstojí — a čo nie

Výsledky ablácie sú presvedčivé. Rozdiel medzi "vždy vyhľadávať" a "nikdy nevyhľadávať" je veľký (20,8 p. b.); model sa jasne naučil rozlišovať užitočné vyhľadávanie od šumu. Tokeny IsRel a IsSup pridávajú merateľnú hodnotu nad rámec samotného adaptívneho vyhľadávania. To je zmysluplný výsledok, nie len premenovanie starých konceptov.

Čím som si menej istý, je tvrdenie o generalizácii. Všetkých päť evaluačných úloh (PopQA, TriviaQA, PubHealth, ARC-Challenge, ASQA) sú QA úlohy s krátkymi odpoveďami alebo výberom z možností — presne ten scenár, kde jediná vyhľadaná pasáž môže poskytnúť rozhodujúci signál. Generovanie dlhých textov v kontextoch s viacerými dokumentmi, čo je doména finančných úloh, je podrobené menšiemu skúmaniu. Biography FactScore (80,2) je najbližším priblížením, ale biografie sú relatívne dobre štruktúrované v porovnaní s rozsiahlymi viacročnými knihami výdavkov.

Existuje tu aj háčik v reprodukovateľnosti: trénovacie štítky modelu kritika pochádzajú z GPT-4. To robí kvalitu štítkov závislou od proprietárneho systému a prináša náklady na API, ktoré nie sú uvádzané. CRAG (arXiv:2401.15884) neskôr ukázal, že 0,77B evaluátor vyhľadávania — oveľa ľahší než 7B kritik Self-RAG — dokázal korigovať kvalitu vyhľadávania a získať 19,0 p. b. oproti štandardnému RAG v PopQA, čo naznačuje, že ťažký fine-tuned kritik nemusí byť nevyhnutný. To je dôležitá výzva pre tento dizajn, aj keď hlavná myšlienka o selektívnom vyhľadávaní zostáva platná.

Nakoniec záleží na porovnávacom základe (baseline). Prekonanie ChatGPT (pravdepodobne GPT-3.5-turbo, koniec roka 2023) a Llama2-chat je rozumná latka pre otvorený 13B model, ale špičkové modely sa odvtedy výrazne posunuli. Otázka, či by adaptívne vyhľadávanie Self-RAG prekonalo dobre napromptovaný GPT-4o s jednoduchým nastavením "vždy vyhľadať" na tých istých benchmarkoch, zostáva nezodpovedaná.

Prečo na tom záleží pre finančnú AI

Finanční agenti nad účtovnými knihami Beancount čelia presne problému diskriminácie vyhľadávania, ktorý Self-RAG rieši. Keď sa používateľ opýta "aký je môj čistý príjem za tento mesiac?", agent ho môže vypočítať z načítaného kontextu — vyhľadávanie by mohlo len pridať šum. Keď sa ten istý používateľ opýta "zaúčtoval som faktúru od dodávateľa za 3. štvrťrok?", agent musí prehľadať záznamy možno aj za niekoľko rokov. Prístup "vždy vyhľadávať" plytvá kontextom a riskuje vloženie irelevantných starých transakcií; prístup "nikdy nevyhľadávať" vynechá potrebné dohľadanie.

Tokeny IsRel a IsSup sa dajú jasne namapovať na logiku validácie účtovnej knihy. IsRel: súvisí vyhľadaný transakčný zápis skutočne s dopytom? IsSup: skutočne vyhľadaný kontext potvrdzuje vygenerovanú sumu zostatku, alebo je to číslo halucinácia? Skóre užitočnosti (1 – 5) by mohlo informovať o dôvere pri spätnom zápise: potvrdiť navrhovaný účtovný zápis len vtedy, keď model ohodnotí vlastnú úvahu známkou 4 alebo 5, a zvyšok označiť na kontrolu človekom.

Obava o reprodukovateľnosť je dôležitá aj tu. Pre produkčného účtovného agenta je závislosť od GPT-4 pri generovaní trénovacích štítkov prevádzkovým obmedzením. Ak ľahší evaluátor (v štýle CRAG) dokáže dosiahnuť porovnateľné selektívne vyhľadávanie, je to cesta schodnejšia pre nasadenie. Princípy dizajnu Self-RAG — rozhodnúť sa pred vyhľadávaním, kritizovať po vyhľadávaní — zostávajú cenné, aj keď sa konkrétny recept na trénovanie tokenov zmení.

Čo čítať ďalej

• CRAG: Corrective Retrieval Augmented Generation (arXiv:2401.15884) — stavia na myšlienke adaptívneho vyhľadávania Self-RAG s ľahším evaluátorom a záložným webovým vyhľadávaním, keď lokálne vyhľadávanie zlyhá; stojí za priame porovnanie so Self-RAG na prekrývajúcich sa benchmarkoch.
• RQ-RAG: Learning to Refine Queries for Retrieval Augmented Generation (arXiv:2404.00610) — zameriava sa špecificky na dekompozíciu dopytov pre zložité "multi-hop" QA, čo je scenár, ktorý Self-RAG zvláda menej elegantne.
• FRAMES: Retrieval and Augmentation for Multi-Hop Evaluation (arXiv:2409.12941) — benchmark od Google DeepMind pre RAG s viacerými dokumentmi, ktorý si vyžaduje reťazenie viacerých vyhľadaných faktov; prirodzený náročnejší test pre modely štýlu Self-RAG.