CausalEmbed: Auto-Regressive Multi-Vector Generation in Latent Space for Visual Document Embedding

Probleemstelling

Visual Document Retrieval (VDR) is een cruciale taak voor het terugvinden van relevante pagina's uit grote documentcollecties, essentieel voor systemen zoals Retrieval-Augmented Generation (RAG). Traditionele methoden gebruiken vaak Optical Character Recognition (OCR), wat echter belangrijke lay-out- en structurele informatie verliest. Moderne benaderingen maken gebruik van Multimodal Large Language Models (MLLMs) om visuele documenten als entiteiten te behandelen.

De huidige staat van de kunst, zoals vertegenwoordigd door modellen als ColPali, gebruikt een multi-vector paradigma. Hierbij wordt een documentpagina opgesplitst in honderden of duizenden visuele "patches" (token-level representaties), waarbij elke patch een apart vector-embeddings krijgt. Hoewel dit uitstekende prestaties levert, introduceert het een aanzienlijke opslag- en rekentijd-overhead. Het opslaan en doorzoeken van duizenden vectoren per pagina is onpraktisch voor schaalbare, real-world toepassingen. Bestaande compressiemethoden (zoals clustering of pruning) proberen dit op te lossen door vectoren te verwijderen of samen te voegen, maar dit gaat vaak ten koste van de zoeknauwkeurigheid en is beperkt door de kwaliteit van de oorspronkelijke dichte representaties.

Methodologie: CausalEmbed

Het paper introduceert CausalEmbed, een nieuw paradigma dat de generatieve kracht van MLLMs gebruikt om multi-vector embeddings auto-regressief te genereren in plaats van ze parallel te encoderen.

Kernprincipes:

1. Auto-Regressieve Generatie: In plaats van een document in één keer om te zetten naar een lange reeks vectoren (zoals bij ColPali), genereert CausalEmbed een compacte reeks latent vectoren sequentieel (token voor token). Het model leert een dynamisch generatiebeleid waarbij elke nieuwe vector afhankelijk is van de eerdere context en de oorspronkelijke input.
2. Training met Margin Loss: Tijdens het contrastieve training wordt een iteratieve margin loss geïntroduceerd. Dit moedigt het model aan om compacte en goed gestructureerde representaties te leren.
3. Training Objectief: Het totale verlies bestaat uit drie componenten:
   • Contrastive Loss ($L_m$): Standaard loss om relevante query-documentparen dichter bij elkaar te brengen dan niet-relevante paren.
   • Progressive Refinement Loss ($L_d$): Deze term zorgt ervoor dat elke extra auto-regressieve stap bijdraagt aan een betere retrieval. Het model wordt gestimuleerd om de volledige sequentielengte te benutten om positieve matches te verbeteren en negatieve voorbeelden eerder in de reeks te onderscheiden.
   • Diversity Regularization ($L_q$): Een anti-homogenisatie constraint die voorkomt dat de gegenereerde vectoren identiek of repetitief worden. Het straalt de cosine-similariteit tussen verschillende tokens in de gegenereerde query-sequentie af.

Theoretisch Voordeel:
De auteurs analyseren de gradiëntstroom en tonen aan dat auto-regressieve generatie leidt tot een efficiëntere gradiëntpropagatie. Bij parallelle patch-encoding krijgen alleen de geselecteerde document-tokens directe gradiëntsignalen. Bij auto-regressieve generatie hangt elke token af van de volledige voorgaande context, waardoor gradiënten door de hele generatieketen kunnen stromen, wat resulteert in een hogere trainingsefficiëntie.

Belangrijkste Bijdragen

1. Nieuw Paradigma: Introductie van een auto-regressief paradigma voor visuele document-embedding dat afwijkt van de conventionele parallelle patch-based encoding.
2. Superieure Compressie: Het behalen van een optimale afweging tussen opslag-efficiëntie en retrieval-nauwkeurigheid. CausalEmbed bereikt een 30x tot 155x reductie in het aantal tokens (van duizenden naar slechts enkele tientallen) terwijl het de prestaties van geclusterde baselines overtreft.
3. Test-time Scaling: Een unieke eigenschap waarbij de precisie dynamisch kan worden aangepast door het aantal gegenereerde tokens tijdens de inferentie te variëren. Het model vertoont een "Matryoshka-achtig" gedrag: meer tokens leiden tot betere resultaten, maar zelfs met zeer korte sequenties presteert het goed.
4. Robuuste Generalisatie: De methode werkt effectief op diverse backbones, inclusief minder krachtige modellen zoals PaliGemma, waar het zelfs de volledige ColPali-baselines overtreft.

Resultaten

De methodologie is uitgebreid getest op de ViDoRe benchmarks (V1, V2 en V3), die diverse domeinen (academisch, medisch, industrieel) en talen omvatten.

• Prestaties onder Extreme Compressie: CausalEmbed (bijv. CausalQwen met 32 tokens) overtreft bijna alle andere methoden onder dezelfde compressie-eisen. Op ViDoRe V2 en V3 behaalt het aanzienlijk hogere nDCG@5 scores dan de beste pruning- en clustering-baselines (zoals SemCluster of BiQwen2).
• Vergelijking met Volledige Modellen: Opmerkelijk is dat CausalPali (gebaseerd op PaliGemma) met slechts 32 tokens de prestaties van de volledige ColPali-baseline (met ~1000 tokens) op alle drie de ViDoRe-versies met bijna 5% overtreft.
• Efficiëntie: Hoewel auto-regressieve generatie een lichte overhead introduceert bij het coderen van de query, compenseert dit ruimschoots door de enorme versnelling in de "Late Interaction" (MaxSim) fase. De totale end-to-end latency is lager dan bij methoden die clustering vereisen, vooral bij schaalbare zoekopdrachten.
• Test-time Scaling: Experimenten tonen aan dat de prestaties lineair verbeteren naarmate het aantal tokens tijdens de inferentie toeneemt (van 1 tot 32+), wat flexibele implementatiestrategieën mogelijk maakt (bijv. snelheid vs. nauwkeurigheid).

Significantie

CausalEmbed vertegenwoordigt een fundamentele verschuiving in het veld van visuele documentretrieval. Het lost het "storage bottleneck"-probleem van multi-vector systemen op door te laten zien dat generatieve modellen in staat zijn om compacte, hoogwaardige representaties te synthetiseren zonder de nuance van visuele informatie te verliezen.

De methode maakt multi-vector VDR praktisch toepasbaar in real-world scenario's waar opslag en latency kritiek zijn (zoals enterprise search en RAG-systemen). Door de introductie van test-time scaling biedt het ook een nieuwe flexibiliteit voor adaptieve inferentie. Het paper schetst de weg voor een generatief paradigma binnen multimodale documentretrieval, waarbij efficiëntie en prestaties niet langer als tegenstrijdige doelen worden gezien.