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

1. Il Problema: Sovraccarico di Archiviazione nei Sistemi VDR

Il Visual Document Retrieval (VDR) è fondamentale per sistemi di ricerca aziendale e per il Retrieval-Augmented Generation (RAG) in domini specifici. A differenza dei sistemi tradizionali basati su OCR, il VDR tratta le pagine dei documenti come entità visive per preservare informazioni strutturali e di layout critiche.

Sebbene i recenti Multimodal Large Language Models (MLLM) abbiano migliorato le prestazioni generando embedding multi-vettore (dove ogni "patch" visiva della pagina diventa un vettore), questo approccio presenta un grave limite pratico:

• Sovraccarico di Archiviazione: Rappresentare una pagina con migliaia di token visivi (spesso 1000+) richiede enormi risorse di memoria e rende difficile la scalabilità in ambienti di produzione.
• Limiti delle Soluzioni Attuali: I metodi esistenti per comprimere questi vettori (come il pruning tramite clustering, merging o tecniche come MetaEmbed) spesso comportano una perdita significativa di accuratezza nella ricerca o richiedono budget di token fissi durante l'addestramento, limitando la flessibilità.

La domanda centrale è: È possibile sfruttare la potenza generativa degli MLLM per produrre embedding multi-vettore riducendo drasticamente la ridondanza senza sacrificare le prestazioni?

2. Metodologia: CausalEmbed

Gli autori propongono CausalEmbed, un nuovo paradigma che trasforma la generazione di embedding da un processo parallelo (basato su patch) a un processo auto-regressivo sequenziale.

Concetti Chiave:

• Generazione Auto-Regressiva: Invece di generare tutti i vettori di un documento in parallelo, il modello (un MLLM pre-addestrato fine-tunato) genera una sequenza di vettori latenti uno alla volta. Ogni vettore $z_k$ è condizionato dall'input originale e dalla storia dei vettori precedentemente generati ($z_1, ..., z_{k-1}$).
• Compressione Estrema: Questo approccio permette di rappresentare un'intera pagina con solo decine di token (es. 32) invece di migliaia, ottenendo un rapporto di compressione di 30-155 volte.
• Funzione di Obiettivo di Addestramento: Per evitare il collasso della rappresentazione e garantire qualità, CausalEmbed utilizza una combinazione di tre termini di perdita:
  1. Contrastive Loss ($L_m$): Allineamento standard query-documento basato sul margine.
  2. Progressive Refinement Loss ($L_d$): Incoraggia il modello a migliorare progressivamente la pertinenza man mano che la sequenza di token cresce. Ogni passo aggiuntivo deve contribuire a una migliore distinzione tra documenti positivi e negativi.
  3. Diversity Regularization ($L_q$): Penalizza la similarità tra i token generati nella stessa sequenza query per prevenire l'omogeneizzazione (evitare che il modello ripeta lo stesso vettore).

Analisi Teorica del Flusso di Gradiente

Il paper dimostra teoricamente che l'approccio auto-regressivo è più efficiente nell'addestramento rispetto a quello forward-based. Grazie alla dipendenza causale, i gradienti si propagano attraverso l'intera catena di generazione, permettendo a ogni token di documentare di influenzare l'aggiornamento dei parametri in modo più completo rispetto ai metodi paralleli dove solo i token selezionati ricevono segnali diretti.

3. Contributi Chiave

1. Nuovo Paradigma: Introduzione di un approccio auto-regressivo sequenziale per la generazione di embedding multi-vettore, abbandonando l'encoding parallelo basato su patch.
2. Efficienza e Accuratezza: Dimostrazione che CausalEmbed supera i metodi basati su pruning e clustering, mantenendo prestazioni competitive (o superiori) con un numero di token ridotto di 30 volte.
3. Scalabilità al Test (Test-time Scaling): Identificazione di una proprietà unica: la precisione della ricerca può essere regolata dinamicamente modificando il numero di token generati durante l'inferenza (da pochi token per bassa latenza a molti per alta accuratezza), simile al concetto di Matryoshka Representation.
4. Generalizzazione: Il metodo funziona robustamente su diversi backbones, inclusi modelli più piccoli o meno performanti come PaliGemma, migliorandone le prestazioni anche di oltre il 14% rispetto alla baseline completa.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sui benchmark ViDoRe V1, V2 e V3, utilizzando backbones come Qwen2.5-VL e PaliGemma.

• Prestazioni Superiori: Con soli 32 token, CausalEmbed supera significativamente le migliori baseline di pruning (come SemCluster o K-Means) e si avvicina o supera le prestazioni dei modelli full-token (ColQwen/ColPali) su diverse metriche nDCG@5.
  • Esempio: Su ViDoRe V2, CausalQwen raggiunge un nDCG@5 di 0.516, superando la migliore baseline di pruning (0.455) e avvicinandosi al modello base full-token (0.576) con una frazione dei token.
• Efficienza Computazionale: Sebbene la generazione auto-regressiva introduca un leggero overhead nell'encoding della query, riduce drasticamente il tempo di Late Interaction (MaxSim) (il collo di bottiglia principale) di circa 11.8 volte rispetto ai metodi basati su clustering, grazie alla drastica riduzione del numero di vettori da confrontare.
• Robustezza: Il metodo mantiene alte prestazioni su domini diversificati (medico, legale, finanziario) e in contesti multilingue.

5. Significato e Impatto

CausalEmbed rappresenta un passo fondamentale verso l'adozione pratica degli embedding multi-vettore per la ricerca di documenti visivi.

• Riduzione dei Costi: La riduzione di 30-155 volte del numero di token abbassa drasticamente i costi di archiviazione e di inferenza, rendendo fattibile l'uso di questi modelli in produzione su larga scala.
• Flessibilità Operativa: La capacità di adattare la lunghezza della sequenza al momento dell'inferenza permette di bilanciare dinamicamente latenza e accuratezza in base alle esigenze dell'applicazione.
• Paradigma Generativo: Il lavoro sposta il focus dalla compressione di rappresentazioni esistenti alla generazione nativa di rappresentazioni compatte, aprendo nuove direzioni di ricerca nell'intersezione tra generazione e retrieval multimodale.

In sintesi, CausalEmbed risolve il dilemma "efficienza vs. accuratezza" nel VDR, dimostrando che la generazione sequenziale di embedding latenti è una soluzione superiore rispetto alle tecniche di pruning tradizionali.