Efficient Multivector Retrieval with Token-Aware Clustering and Hierarchical Indexing

1. Enunciato del Problema

I modelli di recupero multivettore (ad esempio, ColBERT) hanno raggiunto un'efficacia all'avanguardia nel Recupero di Informazioni (IR) codificando i documenti come insiemi di vettori a livello di token e utilizzando l'"interazione tardiva" (MaxSim) per calcolare la rilevanza. Tuttavia, il loro dispiegamento è ostacolato da:

• Elevati Costi Computazionali e di Memoria: L'archiviazione e la ricerca di miliardi di vettori di token richiedono risorse massicce.
• Clustering Inefficiente: Le soluzioni attuali si affidano al clustering $k$-means standard per generare vettori centroidi per la compressione e la raccolta dei candidati.
  • Colli di Bottiglia nella Scalabilità: $k$-means scala male con la dimensione del dataset e il numero di cluster ($k$), richiedendo spesso decine di ore o GPU costose per addestrare anche insiemi moderatamente grandi (ad esempio, 262k centroidi).
  • Squilibrio dei Token: $k$-means standard tratta tutti i vettori in modo uguale. Di conseguenza, i token frequenti (parole vuote) dominano la funzione di perdita e consumano la maggior parte dei centroidi, mentre i token rari e altamente discriminativi sono sottorappresentati, degradando la qualità del recupero.

2. Metodologia

Gli autori propongono Tachiom, un sistema di recupero che comprende due innovazioni fondamentali: Clustering Consapevole dei Token (Tac) e un Indice Gerarchico per un recupero ottimizzato.

A. Clustering Consapevole dei Token (Tac)

Tac affronta le limitazioni di $k$-means standard sfruttando la proprietà strutturale dell'identità del token negli embedding multivettore. Invece di clusterizzare globalmente tutti i vettori, Tac scompone il problema in sottoproblemi indipendenti per ogni token.

• Strategia di Allocazione dei Centroidi: Tac alloca il budget dei centroidi in base alla frequenza del token e alla varianza semantica, piuttosto che alla distribuzione uniforme dei vettori.
  • Gestione della Coda: I token estremamente rari ($n < \mu$) hanno garantito almeno 1 centroide; i token piccoli ($\mu \le n < \tau$) ne ottengono 2.
  • Punteggio Smorzato: Per i token attivi, viene calcolato un peso $w_j$ come $\sqrt{n_j} \cdot s_j$, dove $n_j$ è la frequenza e $s_j$ è la varianza semantica (varianza componente per componente). La radice quadrata smorza l'influenza dei token ad alta frequenza, impedendo loro di monopolizzare il budget.
  • Delimitazione e Riconciliazione: Soglie minime e massime rigide assicurano una rappresentazione minima per i token rari e prevengono l'allocazione eccessiva. Il budget finale viene riconciliato per corrispondere al target $k$.
• Parallelizzazione: Risolvendo problemi indipendenti di $k_j$-means per ogni tipo di token, Tac abilita una massiccia parallelizzazione, riducendo drasticamente la complessità computazionale rispetto a $k$-means globale.
• Compressione dei Residui: Dopo il clustering, i residui (vettore - centroide) vengono normalizzati e compressi utilizzando la Quantizzazione del Prodotto (PQ).

B. Indice Gerarchico per il Recupero (Tachiom)

Tachiom sfrutta i centroidi ad alta risoluzione generati da Tac per ottimizzare la pipeline di recupero in due fasi:

1. Fase di Raccolta (Valutazione Solo Centroidi):
   • Viene costruito un grafo Hnsw sui centroidi.
   • Ogni centroide mantiene una lista inversa di ID documento.
   • Per una query, il sistema attraversa il grafo Hnsw per trovare i centroidi più vicini, recupera i documenti associati e calcola un punteggio MaxSim approssimato utilizzando solo le interazioni dei centroidi.
   • Questo bypassa completamente i costosi calcoli a livello di token durante la selezione iniziale dei candidati.
2. Fase di Raffinamento (MaxSim Ottimizzato PQ):
   • I migliori candidati vengono raffinati utilizzando sia i centroidi che i residui compressi PQ.
   • Layout Ottimizzato per la Cache: Gli autori introducono un layout di memoria specializzato per le tabelle delle distanze PQ. Invece di accessi alla memoria dispersi per ogni token della query, le distanze sono organizzate in Macro-blocchi (sottospazi), Blocchi di centroidi e Micro-blocchi di token di query. Questo permette di recuperare le distanze per tutti i token della query rispetto a un centroide specifico come un singolo blocco di memoria contiguo, migliorando significativamente la località della cache e l'utilizzo SIMD.

3. Contributi Chiave

1. Clustering Consapevole dei Token (Tac): Un algoritmo innovativo che scompone il clustering globale in sottoproblemi per token. Garantisce teoricamente un'accelerazione con limite inferiore dato dal rapporto tra il peso totale del vocabolario e il peso massimo di un singolo token.
2. Architettura Tachiom: Un sistema di recupero che combina la raccolta di centroidi basata su Hnsw con un layout PQ ottimizzato per la cache, adattato per l'interazione tardiva, abilitando una valutazione completa efficiente.
3. Validazione Teorica ed Empirica: Il documento fornisce un limite inferiore teorico per l'accelerazione di Tac e dimostra che il metodo scala a milioni di centroidi mantenendo o migliorando l'efficacia del recupero.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su Ms Marco-v1 (8,8M passaggi) e LoTTE (2,4M passaggi) utilizzando embedding ColBERTv2.

• Efficienza del Clustering:
  • Tac ha raggiunto un addestramento fino a 247$\times$ più veloce rispetto a $k$-means standard (Faiss) e 84$\times$ più veloce rispetto a Faiss con accelerazione MKL.
  • Tac ha clusterizzato con successo 600M vettori in 262k cluster in 8 minuti, mentre i concorrenti hanno raggiunto un timeout di 24 ore prima di arrivare a 131k–524k cluster.
  • Tac scala senza soluzione di continuità a oltre 4 milioni di centroidi in ~100 minuti.
• Efficacia del Recupero:
  • Tac ha ottenuto un MRR@10 uguale o superiore rispetto a $k$-means standard con budget di centroidi fissi, dimostrando che dare priorità ai token rari migliora la qualità dell'approssimazione.
• Velocità di Recupero:
  • Tachiom ha raggiunto fino a 9,8$\times$ un recupero end-to-end più veloce rispetto ai sistemi all'avanguardia (Emvb, Warp, Igp).
  • Ha eguagliato l'efficacia massima di Emvb (che utilizza varianti PQ supervisionate costose) utilizzando PQ standard e ottenendo un accelerazione di 5,5$\times$ rispetto a Emvb.

5. Significato

Questo lavoro sposta fondamentalmente il paradigma per l'efficienza del recupero multivettore:

• Scalabilità: Rimuove il collo di bottiglia computazionale del clustering, permettendo ai sistemi di utilizzare milioni di centroidi (granularità più fine) che in precedenza non erano fattibili da addestrare.
• Compromesso Qualità-Efficienza: Allineando l'allocazione dei centroidi agli obiettivi di recupero (favorendo i token rari discriminativi) piuttosto che al puro errore di ricostruzione, raggiunge una migliore efficacia con meno risorse.
• Dispiegamento Pratico: La combinazione di addestramento rapido e inferenza ultra-veloce (tramite raccolta Hnsw e PQ ottimizzato) rende il recupero multivettore ad alte prestazioni fattibile per ambienti di produzione su larga scala.

Gli autori hanno rilasciato la loro implementazione in Rust di Tac e Tachiom su GitHub per garantire la riproducibilità.