Technische Samenvatting: DiffRetriever

Probleemstelling

Recent werk, specifiek PromptReps (Zhuang et al., 2024), toonde aan dat autoregressieve taalmodellen (AR LM's) kunnen fungeren als zero-shot retrievers door hen te prompten om representatieve tokens (dichte vectoren en spaarse logits) te genereren voor queries en passages. Het uitbreiden van deze aanpak naar multi-token retrieval (het genereren van $K$ representanten in plaats van één) bleek echter inefficiënt en vaak ineffectief met AR-modellen. Omdat AR-modellen tokens sequentieel genereren, vereist het produceren van $K$ representanten $K$ forward passes, waardoor de encodeerlatentie lineair schaalt met $K$. Eerdere pogingen tot multi-token PromptReps lieten geen consistente prestatiewinst zien ten opzichte van single-token decoding, ondanks deze aanzienlijke latentiekost. Hierdoor bleef het onduidelijk of de beperking het multi-token concept zelf was of de bottleneck van sequentiële generatie.

Methodologie

De auteurs stellen DiffRetriever voor, een retrievalkader dat specifiek is ontworpen voor Diffusion Language Models (DiffLM's) om de bottleneck van sequentiële generatie te overwinnen.

Kernmechanisme

In tegenstelling tot AR-modellen die tokens één voor één genereren, vullen DiffLM's gemaskeerde posities gezamenlijk in onder bidirectionele aandacht. DiffRetriever maakt hier gebruik van door:

1. Promptconstructie: Het toevoegen van $K$ gemaskeerde posities ([MASK]) aan een prompt voor representatieve tokens (bijv. "De woorden zijn [MASK]...").
2. Parallelle Decoding: Het invoeren van de volledige sequentie (prompt + $K$ masks) in de DiffLM in een enkele bidirectionele forward pass.
3. Extractie van Representaties: Het gelijktijdig extraheren van $K$ dichte hidden states en $K$ spaarse logit-vectoren uit de $K$ gemaskeerde posities.

Scoren en Training

• Scoren: Het kader adopteert het hybride score-recept van PromptReps:
  • Dichte Score: Bereikt via late interactie in ColBERT-stijl (MaxSim) tussen hidden states van query en passage.
  • Spaarse Score: Bereikt via max-gepoolde, ReLU-geactiveerde logits.
  • Hybride Score: Een lineaire interpolatie van genormaliseerde dichte en spaarse scores.
• Fine-tuning: De auteurs passen contrastieve fine-tuning toe (InfoNCE-verlies) op zowel dichte als spaarse scores, toepasbaar op zowel AR- als DiffLM-backbones.

Experimentele Opzet

De studie vergelijkt vier backbones (7B–8B parameters):

• Autoregressief (AR): LLaMA3-8B, Qwen2.5-7B.
• Diffusie (Diff): Dream-7B (geïnitieerd vanuit Qwen2.5), LLaDA-8B (vanaf nul getraind).
• Baselines: BM25, PromptReps (AR), DiffEmbed (DiffLM als mean-gepoolde encoder), en RepLLaMA (contrastief gefinetuned AR).
• Datasets: MS MARCO, TREC DL 2019/2020, en BEIR-7 (zeven diverse datasets buiten het domein).

Belangrijkste Resultaten

1. Efficiëntie en Effectiviteit van Multi-Token

• Diffusie versus Autoregressie: Multi-token DiffRetriever verbetert aanzienlijk ten opzichte van single-token DiffRetriever op elke geteste diffusie-backbone. Daarentegen tonen multi-token AR-modellen vlakke of negatieve winsten ten opzichte van single-token AR-modellen.
• Latentie: De encodeerkost voor multi-token DiffRetriever blijft ongeveer constant ongeacht $K$ (enkele forward pass), terwijl de AR multi-token-kost lineair schaalt. In de zero-shot setting ($K \le 20$) levert AR multi-token een latentiestraf van $\approx 15\times$ op ten opzichte van diffusie, zonder consistente prestatiewinst.
• Prestatieomkering: In zero-shot settings presteren AR-modellen aanvankelijk beter dan DiffLM's bij $K=1$. Echter, bij $K>1$ overtreffen DiffLM's (specifiek Dream en LLaDA) hun AR-tegenhangers. Bijvoorbeeld, Dream multi-token verdubbelt bijna de MS MARCO-prestaties ten opzichte van zijn single-token versie, en passeert Qwen2.5.

2. State-of-the-Art Prestaties

• Zero-Shot: DiffRetriever op LLaDA behaalt de hoogste gemiddelde BEIR-7-score onder zero-systemen in de vergelijking, en presteert beter dan PromptReps op AR-backbones en de DiffEmbed-baseline.
• Gefinetuned: Na supervised fine-tuning wordt DiffRetriever op Dream de sterkste BEIR-7-retriever overall, en passeert gefinetuned PromptReps, DiffEmbed en de contrastief gefinetuned RepLLaMA.
• Synergie van Backbone: De Dream-backbone (geïnitieerd vanuit Qwen2.5) toont een grotere prestatieschommeling van zero-shot naar gefinetuned in vergelijking met LLaDA, wat suggereert dat AR-pretraining-priors gecombineerd met bidirectionele diffusiedecoding zeer effectief zijn voor retrieval.

3. Budgetanalyse en Oracle-headroom

• Vast versus Adaptief: De studie identificeert dat het gebruik van een vast budget $(K_q, K_p)$ voor alle queries aanzienlijke prestaties laat liggen.
• Oracle-analyse: Een "per-query oracle" (die de optimale $K$ voor elke query selecteert met behulp van ground truth) op het bevroren basismodel presteert beter dan contrastieve fine-tuning bij een vast budget voor alle combinaties.
• Voorspelbaarheid: Het optimale budget correleert positief met goedkope querykenmerken (lengte en Shannon-entropie), wat suggereert dat een geleerde router voor adaptieve budgetselectie een haalbare toekomstrichting is.

Betekenis en Claims

Het artikel claimt de ambiguïteit rondom multi-token retrieval op te lossen:

1. De Bottleneck is Sequentiële Generatie: Het falen van eerdere multi-token AR-retrievers was niet te wijten aan het multi-token concept zelf, maar aan de kosten van sequentiële decoding. DiffRetriever bewijst dat multi-token retrieval zowel effectief als goedkoop is wanneer het ontkoppeld wordt van sequentiële generatie via diffusiemodellen.
2. Training-uitgelijnde Doelen: Het bevragen van het diffusiemodel in zijn voorgetrainde vorm (het voorspellen van gemaskeerde posities parallel) transferreert beter naar retrievaltaken dan het hergebruiken van het model als een standaard BERT-stijl encoder (zoals gedaan in DiffEmbed).
3. Adaptief Potentieel: Het bestaan van een grote kloof tussen prestaties bij een vast budget en per-query oracle-prestaties benadrukt adaptieve budgetselectie als een kritieke, onbenutte weg voor het verbeteren van retrieval-efficiëntie en -effectiviteit.

De auteurs concluderen dat DiffRetriever een nieuwe state-of-the-art vestigt voor op diffusie gebaseerde retrieval, en aantoont dat parallelle representatieve tokens zowel single-token diffusie als sequentiële autoregressieve aanpakken kunnen overtreffen, terwijl ze lage latentie behouden.