UniDoc-RL: Coarse-to-Fine Visual RAG with Hierarchical Actions and Dense Rewards

Probleemstelling

Retrieval-Augmented Generation (RAG) voor grote visuele-taalmodellen (LVLMs) heeft tot doel externe visuele kennis te integreren om complexe redeneringen mogelijk te maken. Echter, bestaande systemen voor visuele RAG kampen met drie fundamentele beperkingen:

1. Onnauwkeurige Retrieval: Bestaande methoden vertrouwen vaak op generieke zoeksignalen (zoals embedding-similariteit) die de fijne, taakspecifieke visuele semantiek missen die nodig is voor complexe redenering. Dit leidt tot irrelevante context en hallucinaties.
2. Passief Visueel Gebruik: Modellen verwerken vaak volledige afbeeldingen passief, wat resulteert in het meenemen van veel ruis en redundante achtergrondinformatie. Ze missen het vermogen om actief te focussen op informatie-dichte regio's (zoals specifieke delen van een grafiek of tabel).
3. Schaars Beloningssysteem: Reinforcement Learning (RL) methoden gebruiken meestal alleen een 'outcome reward' (beloning op basis van het eindantwoord). Dit creëert een 'credit assignment'-probleem: het model krijgt geen feedback over tussentijdse beslissingen (zoals welke afbeelding te selecteren of welk gebied te bijsnijden), waardoor de optimalisatie inefficiënt is.

Methodologie: UniDoc-RL

UniDoc-RL is een unificerend RL-framework waarin een LVLM-agent gezamenlijk zoekt, rangschikt, actief visuele perceptie uitvoert en redeneert. De kern van de methode bestaat uit drie componenten:

1. Hiërarchische Actieruimte (Coarse-to-Fine)

In plaats van één grote stap te nemen, doorloopt het agent een sequentieel besluitvormingsproces met drie lagen van verfijning:

• Zoekactie (Search): Het agent genereert een zoekopdracht om een breed scala aan kandidaat-afbeeldingen uit een externe corpus op te halen (coarse-grained).
• Selectieactie (Precise Selection): Het LVLM evalueert de gevonden kandidaten en selecteert de meest relevante afbeeldingen op basis van semantische uitlijning met de query. Dit filtert ruis op documentniveau voordat de redenering begint.
• Visuele Perceptieactie (Active Visual Perception): In plaats van de hele afbeelding te verwerken, leert het agent om actief gebieden te lokaliseren, te bijsnijden (crop) en in te zoomen op specifieke regio's (bijv. een klein diagram in een rapport). Dit imiteert menselijke visuele aandacht en behoudt hoge resolutie-details terwijl ruis wordt verwijderd.

2. Dicht Multi-Beloningssysteem (Dense Multi-Rewards)

Om het credit assignment-probleem op te lossen, introduceert UniDoc-RL een gedetailleerd beloningssysteem dat elke stap in het proces superviseert:

• Zoekbeloning (NDCG-based): Meet de kwaliteit van de initiële zoekresultaten ten opzichte van de 'golden' relevante afbeeldingen.
• Selectiebeloning (Select-based): Belooft de juiste selectie van relevante afbeeldingen uit de zoekresultaten.
• Perceptiebeloning (IoU-based): Gebruikt Intersection over Union (IoU) om de nauwkeurigheid van de geselecteerde bijsnijdingen (bounding boxes) te meten ten opzichte van de grondwaarheid.
• Patroonbeloning: Straft formatfouten (bijv. incorrect gebruik van XML-tags zoals <search>, <select>, <bbox>).
• Outcome Beloning: Een model-based beloning die de uiteindelijke juistheid van het antwoord evalueert.
  De totale beloning is een gewogen som van deze componenten, wat zorgt voor synergetische optimalisatie van het hele proces.

3. Trainingsparadigma

• Data Synthese: Een hoogwaardig dataset van redeneertrajecten met fijne actielabels is gegenereerd met behulp van een 'teacher' agent (Qwen3-VL-235B) en tools voor layout-analyse (Mineru).
• SFT + GRPO: Het model ondergaat eerst Supervised Fine-Tuning (SFT) om de actievormaten te leren. Vervolgens wordt Group Relative Policy Optimization (GRPO) toegepast. GRPO aligneert het gedrag van de agent met de meerdere doelen zonder een apart waarde-netwerk (value network) te vereisen, wat de training efficiënter maakt.

Kernbijdragen

1. UniDoc-RL Framework: Een unificerend RL-framework dat zoekopdrachten, rangschikking, actieve visuele perceptie en redenering integreert in één besluitvormingsproces.
2. Nieuwe Dataset: Het publiceren van een uitgebreide dataset van 12.621 samples voor SFT en 5.537 samples voor RL, rijk aan gedetailleerde actielabels voor visuele RAG-taken.
3. Hiërarchische Acties & Dichte Beloningen: De introductie van een "Zoek-Selecteer-Perceptie"-strategie gecombineerd met een gedicht beloningssysteem dat de optimalisatie van tussentijdse stappen mogelijk maakt.

Resultaten

De prestaties van UniDoc-RL zijn getest op drie uitdagende benchmarks: SlideVQA, ViDoSeek en MMLongBench.

• Superieure Prestaties: UniDoc-RL overtreft consistent state-of-the-art baselines, waaronder andere RL-methoden zoals VRAG-RL.
• Verbetering: Er wordt een verbetering van tot 17,7% behaald ten opzichte van eerdere RL-gebaseerde methoden.
• Ablatie Studies:
  • De Selectieactie verbetert vooral de prestaties op datasets die zwaar redeneren vereisen (SlideVQA, ViDoSeek) door de contextkwaliteit te verhogen.
  • De Visuele Perceptieactie (bijsnijden) levert de grootste winst op bij taken die fijne visuele details vereisen (MMLongBench).
  • Het gebruik van dichte beloningen leidt tot een significant betere nauwkeurigheid in tussentijdse stappen dan methoden die alleen op het eindresultaat belonen.
• Actief Gedrag: Na RL-training toont het model een veel actiever gedrag (meer frequent bijsnijden) en leert het om alleen te bijsnijden wanneer dit nodig is voor nauwkeurigheid, in plaats van passief de hele afbeelding te lezen.

Betekenis en Impact

UniDoc-RL markeert een belangrijke stap in de evolutie van visuele documentbegrip. Het bewijst dat het integreren van actieve visuele perceptie (het leren van waar te kijken) en hiërarchische besluitvorming essentieel is voor complexe visuele RAG-taken. Door het probleem van schaarse beloningen op te lossen via een gedicht multi-reward systeem, maakt het end-to-end training mogelijk van agenten die niet alleen antwoorden genereren, maar ook leren hoe ze effectief informatie moeten zoeken, filteren en visualiseren. Dit heeft grote implicaties voor toepassingen zoals het analyseren van wetenschappelijke rapporten, financiële documenten en complexe infographics.