From Tokens to Concepts: Leveraging SAE for SPLADE

Titolo

Dai Token ai Concetti: Sfruttare gli SAE per SPLADE

1. Il Problema

I modelli di Recupero dell'Informazione (IR) appresi e sparsi (Learned Sparse Retrieval - LSR), come SPLADE, offrono un eccellente compromesso tra efficienza ed efficacia. Tuttavia, questi modelli presentano limitazioni intrinseche legate alla loro dipendenza dal vocabolario sottostante del modello linguistico pre-addestrato (PLM), solitamente basato su token lessicali:

• Polisemia e Sinonimia: L'uso diretto dei token può portare a ambiguità (un token con più significati) o a una mancata cattura delle relazioni semantiche tra parole diverse (sinonimi).
• Sfide Multilingua e Multimodali: Adattare questi modelli a nuove lingue o modalità è complesso a causa della rigidità del vocabolario lessicale.
• Rappresentazione Semantica: I modelli attuali derivano rappresentazioni sparse dagli stati nascosti dell'ultimo strato del Transformer, che potrebbe non essere lo strato più adatto per compiti specifici di IR.
• Espansione dei Termini: L'espansione dei termini per gestire il mismatch lessicale (es. sinonimi) e la frammentazione delle menzioni di entità non sono sempre soddisfacenti.

2. Metodologia

Gli autori propongono SAE-SPLADE, un'architettura che sostituisce il "testa" (MLM head) di proiezione lessicale di SPLADE con uno spazio latente di concetti semantici appresi tramite Sparse Auto-Encoders (SAE).

Architettura e Pipeline di Addestramento

Il processo avviene in due fasi distinte (illustrate nella Figura 1 del paper):

1. Pre-addestramento dell'SAE:

   • Viene utilizzato un encoder PLM (es. DistilBERT) per generare embedding contestualizzati dei token.
   • Un Sparse Auto-Encoder (SAE) viene addestrato per ricostruire questi embedding densi. L'SAE mappa gli stati latenti in un vettore sparso di dimensioni $M$ (dove $M \gg d$, la dimensione originale).
   • Viene utilizzato un TopK SAE, che seleziona le $k$ attivazioni più forti per ogni embedding, garantendo la sparsità senza bisogno di regolarizzazione L1 esplicita (evitando problemi di "shrinkage bias").
   • L'obiettivo è imparare un vocabolario semantico dove ogni latente rappresenta un concetto specifico, non un token fisso.

2. Fine-tuning per IR (SAE-SPLADE):

   • Il decoder dell'SAE viene rimosso.
   • L'encoder SAE addestrato sostituisce la testa MLM originale di SPLADE.
   • Il modello viene fine-tunato sull'obiettivo di IR (distillazione da un cross-encoder) utilizzando la stessa strategia di SPLADEv3 (combinazione di KL-divergence e MarginMSE), ma proiettando le rappresentazioni nello spazio dei concetti SAE invece che in quello lessicale.
   • Viene applicata una regolarizzazione FLOPs per controllare ulteriormente la sparsità della rappresentazione del documento/query.

Scelte Progettuali Chiave

• Livello di Input: Gli esperimenti mostrano che applicare l'SAE all'ultimo strato dell'encoder (prima della trasformazione finale della testa MLM) offre il miglior compromesso.
• Dimensione del Vocabolario: Viene utilizzato un vocabolario di latenti di dimensione $M=2^{16}$ (circa il doppio del vocabolario BERT).
• Vincolo Top-K: Sia durante l'addestramento dell'SAE che durante il fine-tuning per IR, viene applicato un vincolo Top-K ($k$) per mantenere la rappresentazione sparsa.

3. Contributi Chiave

1. Nuova Architettura (SAE-SPLADE): Sostituzione del vocabolario lessicale con un vocabolario semantico appreso tramite SAE, permettendo di superare le limitazioni dei token fissi.
2. Analisi di Efficienza ed Efficacia: Introduzione di una nuova metrica E2 (Efficiency-Effectiveness) per valutare il compromesso tra prestazioni di recupero (MRR/nDCG) e costo computazionale (QD-FLOPs).
3. Risultati Multilingua: Dimostrazione che l'uso di latenti SAE migliora l'adattabilità multilingua rispetto ai modelli basati su MLM, riducendo il mismatch lessicale tra lingue diverse.
4. Analisi Qualitativa: Studio approfondito della relazione tra token e latenti, mostrando come i latenti SAE possano disambiguare la polisemia e raggruppare i sinonimi in modo più efficace rispetto ai token tradizionali.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su dataset in-domain (MS MARCO, TREC-DL) e out-of-domain (LoTTE, BEIR, MIRACL).

• Prestazioni In-Domain e Out-of-Domain:

  • SAE-SPLADE raggiunge prestazioni di recupero comparabili a SPLADE (e in alcuni casi superiori su BEIR) mantenendo un'efficienza significativamente migliore.
  • La variante con $k=8$ ottiene un miglioramento dell'efficienza (QD-FLOPs ridotti da 1.47 a 0.67) con una perdita minima di efficacia rispetto a SPLADE.
  • La variante con $k=16$ raggiunge prestazioni quasi identiche a SPLADEv3 (lo stato dell'arte) ma con un indice più piccolo e un'efficienza superiore.

• Efficienza:

  • SAE-SPLADE è più efficiente di SPLADE tradizionale, richiedendo meno accessi alle liste di posting (QD-FLOPs inferiori) e occupando meno spazio di archiviazione (Avg. Doc Len inferiore).
  • Supera i modelli densi (come ColBERTv2) in termini di latenza di recupero e uso del disco, pur avendo prestazioni leggermente inferiori in termini di MRR assoluto su alcuni task ID.

• Risultati Multilingua:

  • Su dataset multilingue (mMARCO, MIRACL), SAE-SPLADE supera SPLADE in termini di efficacia ed efficienza.
  • L'analisi mostra una maggiore sovrapposizione di latenti attivati tra lingue diverse rispetto alla sovrapposizione di token, indicando una migliore capacità di generalizzazione semantica cross-lingua. Tuttavia, una parte significativa dei latenti rimane specifica per lingua.

• Analisi Ablativa:

  • L'uso di un vincolo Top-K durante l'addestramento e l'inferenza è cruciale per bilanciare sparsità ed efficacia.
  • Il pre-addestramento SAE è superiore all'inizializzazione casuale per i modelli LSR.
  • Varianti SAE più complesse (es. Hierarchical, Matryoshka) non hanno mostrato vantaggi significativi rispetto al semplice TopK SAE in questo contesto.

5. Significato e Implicazioni

Il lavoro dimostra che è possibile sostituire la proiezione lessicale standard dei modelli di recupero con uno spazio semantico appreso, mantenendo i vantaggi dei modelli sparsi (velocità, indice invertibile) mentre si mitigano i problemi di polisemia e sinonimia.

• Versatilità: L'approccio offre una via promettente per il recupero multilingua e multimodale, riducendo la dipendenza da vocabolari lessicali fissi.
• Efficienza: Conferma che i modelli sparsi appresi possono essere ulteriormente ottimizzati per l'efficienza senza sacrificare drasticamente l'efficacia, rendendoli candidati ideali per sistemi di recupero su larga scala.
• Futuro: Il paper suggerisce che l'integrazione di strategie di allineamento cross-lingua più avanzate e l'uso di backbones più grandi potrebbero ulteriormente migliorare le prestazioni, aprendo nuove direzioni per la ricerca sui modelli di recupero sparsi basati su concetti.

In sintesi, SAE-SPLADE rappresenta un passo avanti significativo verso modelli di IR che comprendono il significato (concetti) piuttosto che limitarsi a manipolare token, offrendo un'alternativa robusta ed efficiente ai metodi basati su vocabolario lessicale.