Retrieve-for-Train: RL-tanulással gyorsított, set-alapú keresés diffúziós retrieverrel

A kutatás célja a keresési és ajánló rendszerek egyik növekvő elvárásának kielégítése: nem egyetlen legjobb találatot, hanem koherens, kiegészítő eredményhalmazt kell visszaadni. Például a „kempingfelszerelés” keresésnél a felhasználó nem tíz, kissé eltérő négyfős sátorra kíváncsi, hanem egy olyan, egymást kiegészítő csomagra, amely sátorból, hálózsákból, főzőeszközből és fejlámpából áll.

Miért nem elég a sima, zero-shot LLM a set-kereséshez?

A szokványos, off-the-shelf nagy nyelvi modellek (LLM-ek) két alapvető problémába ütköznek, ha adatbázis-érzékeny query-dekompozíciót várunk tőlük:

• Parafrázis kollapszus: adatbázis-tudatosság hiányában a zero-shot LLM-ek gyakran redundáns, közel szinonim al-kereséseket generálnak (például „bohemian festival fashion” és „bohemian festival clothes”), így homogén, nem informatív slate-et hoznak létre.
• Autoregresszív késleltetés: a komplex, kiegészítő al-keresések előállításához a modellek hosszú láncnyi gondolkodási (chain-of-thought) tokenre támaszkodnak, ami nagy tesztidőbeli számítási költséget és lineárisan növekvő késleltetést eredményez — ez ellentmond a production-öszerű, alatti másodperces válaszidőknek.

A Retrieve-for-Train megközelítés (ICML 2026)

Az „Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion” című, ICML 2026-ban bemutatott munka egy offline, jutalom→adat (reward-to-data) fordítási keretrendszert mutat be. Ahelyett, hogy a modellt minden lekérdezésnél hosszasan gondolkodtatnák, a Retrieve-for-Train egyszeri offline megerősítéses tanulással (RL) felfedezi a jutalomhoz igazodó fan-outokat, majd ezeket felhasználói felügyeletté (supervision) alakítja. A megszerzett viselkedést egy könnyű, diffúziós retrieverbe desztillálják, amely egyetlen, nem-autoregresszív passzban adja meg a teljes al-kérdés készletet.

A pipeline három fő lépésre bontható:

1. Fan-out nyelvi modell RL-tanítása: a fan-out modell megerősítéses tanulással olyan, set-szintű tulajdonságellenőrzésen alapuló al-kereséseket (sub-queries) bocsát ki, amelyek a teljes eredményhalmazt jutalmazó szabályok szerint értékelhetők.
2. Felügyelet szintézise: a befagyasztott fan-out modell offline generál (query → target-set) párokat felügyeletes tanuláshoz, emberi címkézés nélkül.
3. Diffúziós retriever tanítása: egy kompakt, 53.9M paraméteres diffúziós modell megtanulja a lekérdezés embeddingjét közvetlenül a teljes cél-embedding halmazra leképezni, egyetlen nem-autoregresszív lépésben.

Set-szintű célok: kompozit jutalomrendszer

A rendszer teljesítménye attól függ, hogyan definiáljuk a „jó” keresést. A hagyományos pontonkénti relevanciaértékelés (learning to rank) nem fedi le a halmazokra jellemző, nem-dekomponálható tulajdonságokat, mint a sokféleség vagy kiegészítettség. Ennek megfelelően a kutatás 4B-paraméteres nyílt forráskódú modelleket (Gemma3-4B és Qwen3-4B) finomhangolt RL-lel egy szigorú, matematikailag meghatározott kompozit jutalom mentén.

A kompozit jutalom három összetevőből áll:

• Groundedness: bünteti a generált al-keresések és az adatbázis sokaság közötti távolságot, biztosítva, hogy minden al-keresés valódi, visszakereshető tételre mutasson.
• Diversity: a teljes al-keresés-halmaz Vendi Score szerinti mérése, amely a szemantikai sokféleségre kényszeríti a modellt.
• Alignment: rögzíti az al-kereséseket az eredeti, széles prompt szándékához, hogy elkerülje a sematikus jobbra-eltolódást.

Ezek a komponensek egymást kölcsönösen korlátozó „ellen-ankerként” működnek, megakadályozva a jutalom-kihasználó (reward-hacking) stratégiákat: csak groundedness optimalizálása nonszensz kimenetekhez vezetne, míg csak alignment esetén parafrázis-kollapszust kapnánk. A Vendi Score ezt a rövidutat lezárva kényszeríti a modellt érdemi, sokféle, de földhözragadt al-keresések felfedezésére.

Tréningmódszer: Soft-GRPO és egyéb részletek

A fan-out modell optimalizálásához a szerzők group relative policy optimization (GRPO) módszert alkalmaztak soft proximal policy optimization (soft PPO) regularizációval — ezt a megközelítést Soft-GRPO néven említik. A fan-out modelleknél pontosan 10 al-keresés generálása volt a cél minden fő lekérdezésre.

Kísérletek: adathátterek és értékelési beállítások

A rendszert két set-értékű visszakeresési feltételben tesztelték:

• Open-ended abstract retrieval: nincs egyedi „gold” válasz, a minőséget csak set-szintű mutatók (diverzitás, alignment, groundedness) határozzák meg.
• Weakly supervised compositional retrieval: a lekérdezésekhez gyenge, referencia-szerű target halmazok tartoznak, amelyek csak egy lehetséges megvalósítást jelentenek.

A multimodális embedding hátterek két doménben szerepeltek: egy nagyméretű, felhasználó-összeállította outfit divatadatbázis (CLIP alapú retrieverrel értékelve) és egy ipari, szakértők által összeállított zenei playlist-adatbázis (MuLan alapú értékelés). A fan-outot Gemma3-4B és Qwen3-4B 4B modellek végezték, minden lekérdezésre pontosan 10 al-keresést generálva.

Eredmények

• Lekérdezési minőség: a Retrieve-for-Train minden vizsgált feladatban felülmúlta a hagyományos single-query keresést, a zero-shot expanziót és a Best-of-N baseline-t. A korábbi zero-shot LLM-eknél megfigyelt redundáns parafrázisok helyett a módszer sokféle, adatbázishoz kötött és jelentésileg eltérő al-kereséseket produkált.

• Késleltetés és skálázhatóság: közvetlenül az RL-tuningolt nyelvi modell jó minőséget adott, de megőrizte az autoregresszióból fakadó késleltetést. A viselkedés 53.9M-paraméteres diffúziós modellbe történő leszűrése azonban megszüntette ezt a szűk keresztmetszetet: mivel a diffúziós modell párhuzamosan, egyetlen nem-autoregresszív passzban generálja az összes célirányt az embedding-térben, 12–20× sebességnövekedést értek el az autoregresszív megoldásokhoz képest. Míg az autoregresszív fan-out késleltetése nagy kontextusmennyiségek esetén lineárisan növekedve közel 50 másodpercig is felmehet, a Retrieve-for-Train-Diffusion rendszer válaszideje alatti másodpercektől néhány másodpercig terjed, ami éles használatra alkalmassá teszi.

• Ablációk: a jutalomoptimalizációs vizsgálatok feltárták, hogy a diversity komponens nélkül a modell gyorsan „degenerálódó”, vektor-kihasználó stringeket tanult (például ismétlődő, értelmetlen token-sorozatokat). A Vendi Score bevezetése e rövidutakat lezárta, stabil régióba kényszerítve a politikát, ahol a magas jutalom eléréséhez érdemi, szakértői viselkedést kellett felfedezni.

Következtetés

A munka megmutatja, hogy az RL hatékonyan alkalmazható egyszeri „célátalakítóként” (objective transducer), amellyel el lehet választani a drága jutalomvezérelt viselkedés felfedezését a valós idejű inferencia során futó modell struktúrájától. Az RL-felfedezéssel létrehozott felügyeleti adatok diffúziós priorba történő leszűrése lehetővé teszi a rendszerek számára, hogy magasabb szintű, set-szintű tulajdonságokat (diverzitás, alignment, groundedness) hatékonyan optimalizáljanak alacsony késleltetéssel és kis számítási költséggel. Retrieve-for-Train így skálázható, adat-hatékony csővezetéknek kínál megoldást különleges vagy multimodális területeken, ahol az emberi címkézés költséges vagy ritka.

A részletek és a teljes módszertan az ICML 2026-on bemutatott „Efficient, Property-Aligned Fan-Out Retrieval via RL-Compiled Diffusion” cikkben találhatók.