From Tokens to Concepts: Leveraging SAE for SPLADE

Probleemstelling

Learned Sparse Retrieval (LSR) modellen, zoals SPLADE, bieden een uitstekende afweging tussen efficiëntie en effectiviteit voor informatieretrieval (IR). Ze lossen het probleem van lexicale mismatch op door term-expansie te leren via contextuele embeddings, terwijl ze de voordelen van Bag-of-Words-modellen behouden (zoals exacte matching en invertie-indexen).

Echter, deze modellen hebben fundamentele beperkingen die voortvloeien uit hun afhankelijkheid van de vaste vocabulaire van de onderliggende Pre-trained Language Model (PLM), zoals BERT of DistilBERT:

1. Polysemie en Synoniemen: De modelprestaties worden gehinderd doordat tokens meerdere betekenissen kunnen hebben (polysemie) of dat één betekenis door meerdere tokens wordt uitgedrukt (synoniemen).
2. Multilinguale uitdagingen: Het aanpassen aan nieuwe talen of modaliteiten is complex omdat het model gebonden is aan de specifieke token-vocabulaire van de PLM.
3. Suboptimale lagen: De sparse representatie wordt afgeleid van de laatste laag van de Transformer, die niet noodzakelijk de meest geschikte laag is voor IR-taken.

Methodologie: SAE-SPLADE

De auteurs stellen SAE-SPLADE voor, een nieuw architectuurconcept dat de traditionele "MLM-head" (Masked Language Model head) van SPLADE vervangt door een Sparse Auto-Encoder (SAE). In plaats van een vocabulaire van vaste tokens te gebruiken, leert het model een latente ruimte van semantische concepten.

Het trainingsproces verloopt in twee fasen:

1. SAE Pre-training: Een SAE wordt getraind om de contextuele token-embeddings van een PLM (in dit geval DistilBERT) te reconstrueren. De encoder van de SAE projecteert de dichte embeddings naar een hoge-dimensionale, sparse vector van latente concepten ($z \in \mathbb{R}^M$). Er wordt gebruik gemaakt van de TopK SAE variant, waarbij alleen de $k$ sterkst actieve latenten worden behouden om sparsiteit te garanderen zonder expliciete $L_1$-regularisatie (wat vaak leidt tot "shrinkage bias").
2. SAE-SPLADE Fine-tuning: De decoder van de SAE wordt verwijderd. De getrainde SAE-encoder fungeert nu als de projectie-laag voor de IR-taak, vervangend voor de originele MLM-head. Het model wordt vervolgens gefine-tuned op de IR-taak (gebaseerd op MS MARCO) met een doelwit dat bestaat uit:
   • KL-divergentie en MarginMSE distillatie (om de ranking te optimaliseren).
   • FLOPs-regularisatie (om de sparsiteit van de documentrepresentaties te controleren).
   • Een token-level TopK-constraint tijdens zowel training als inferentie.

De auteurs onderzoeken ook welke laag van de PLM het beste als input voor de SAE dient (ze kiezen voor de laatste laag vóór de transformatie-head) en de impact van de grootte van het vocabulaire ($M$) en de waarde van $k$.

Belangrijkste Bijdragen

1. Architecturale Innovatie: De eerste toepassing van een SAE om de vocabulaire-projectie in een LSR-model (SPLADE) volledig te vervangen, waardoor een "semantisch vocabulaire" van concepten ontstaat in plaats van een lexicaal vocabulaire van tokens.
2. Nieuwe Evaluatiemetric (E2): De introductie van een nieuwe scorefunctie E2 (Efficiency-Effectiveness), die de trade-off tussen MRR (effectiviteit) en QD-FLOPs (efficiëntie, oftewel het aantal te verwerken posting-lijsten) kwantificeert.
3. Uitgebreide Experimenten: Een grondige analyse van hyperparameters (zoals $k$, vocabulairegrootte, en SAE-varianten) en een vergelijking met state-of-the-art modellen op diverse datasets (MS MARCO, TREC-DL, BEIR, LoTTE).
4. Multilinguale Validatie: Experimenten die aantonen dat SAE-SPLADE beter presteert in multilinguale settings dan traditionele SPLADE-modellen, dankzij de hogere overlap van geactiveerde latenten over verschillende talen.

Resultaten

De experimenten tonen aan dat SAE-SPLADE concurrerende resultaten behaalt ten opzichte van SPLADE en andere baselines:

• Effectiviteit: SAE-SPLADE bereikt prestaties die vergelijkbaar zijn met SPLADE op in-domain (MS MARCO) en out-of-domain (BEIR, TREC-DL) taken. Op sommige datasets (zoals BEIR) presteert het zelfs significant beter.
• Efficiëntie: Het model is aanzienlijk efficiënter. De QD-FLOPs (Query-Document FLOPs) zijn lager, wat betekent dat er minder posting-lijsten nodig zijn tijdens de zoekopdracht. Bijvoorbeeld, de variant met $k=8$ heeft een QD-FLOP van 0.67 tegenover 1.47 voor de standaard SPLADE.
• Multilinguaal: In multilinguale tests (mMARCO en MIRACL) overtreft SAE-SPLADE de token-gebaseerde SPLADE, zowel in effectiviteit als efficiëntie. De analyse toont aan dat SAE-latenten een hogere overlap hebben tussen verschillende talen, wat wijst op betere cross-linguale semantische representatie.
• Ablatie Studies:
  • Het gebruik van een TopK-constraint tijdens training en inferentie is cruciaal voor het behalen van een goede balans tussen efficiëntie en effectiviteit.
  • Het trainen van de SAE op de laatste laag van de encoder (zonder de MLM-transformatie) levert de beste resultaten op.
  • Complexere SAE-varianten (zoals Hierarchical of Matryoshka TopK) bieden geen significant voordeel ten opzichte van de eenvoudige TopK SAE voor deze specifieke toepassing.

Betekenis en Conclusie

Dit paper bewijst dat het vervangen van de lexicaal gebonden vocabulaire van PLM's door een geleerde semantische vocabulaire van concepten (via SAE) een haalbare en krachtige strategie is voor Learned Sparse Retrieval.

De belangrijkste implicaties zijn:

• Overcoming Vocabulary Limits: Het lost problemen op zoals polysemie en synoniemen op een meer natuurlijke manier dan token-expansie.
• Multilinguale Schaalbaarheid: Het biedt een robuustere basis voor multilinguale retrieval zonder de noodzaak van complexe cross-linguale alignment technieken die vaak nodig zijn bij token-gebaseerde modellen.
• Efficiëntie: Het behoudt de snelheid en lage opslagkosten van sparse retrieval, maar verbetert deze aanzienlijk ten opzichte van bestaande LSR-modellen.

De auteurs concluderen dat SAE-pretraining een geldig alternatief is voor MLM-pretraining in de context van IR, en dat de combinatie van SAE en SPLADE een nieuwe richting opent voor efficiënte en schaalbare semantische zoekopdrachten. Toekomstig werk richt zich op het integreren van geavanceerde cross-linguale strategieën en het combineren van concept-gebaseerde en lexicaal-gebaseerde signalen.