I modelli di embedding multi-vettore, come ColBERT, offrono un’elevata precisione nella ricerca, ma il numero di vettori memorizzati per documento fa salire rapidamente il costo dell’indice. Una tecnica di raggruppamento, il token pooling, riduce questa dimensione, e un addestramento adattato consente, secondo la fonte, di preservare meglio la qualità dei risultati a parità di budget di archiviazione. Con pool_factor=4, un modello addestrato normalmente conserva il 94,6% del suo nDCG@10 non compresso su NanoBEIR, contro il 97,8% per il modello addestrato tenendo conto del pooling.
Perché i modelli multi-vettore memorizzano così tanti vettori
Un modello di embedding classico riassume un intero documento in un unico vettore. I modelli multi-vettore conservano invece un insieme di vettori contestuali, uno per token (un’unità di testo, spesso una parola o parte di una parola), permettendo a diverse parti di una query di corrispondere a diverse parti di un documento. La fonte fa l’esempio di una query su un laptop a lunga autonomia, che può trovare un passaggio sul dispositivo in una parte del documento e un passaggio sull’autonomia della batteria in un’altra.
Questo funzionamento ha un costo. Secondo il calcolo presentato, con una quantizzazione a 4 bit per coordinata residua, 200 vettori di 128 dimensioni rappresentano 12,8 KB per documento, cioè 12,8 GB per un milione di documenti. Ridurre ogni documento a 50 vettori porta questo totale a 3,2 GB, a parità di dimensioni e precisione.
Il metodo si basa su un’osservazione: molti vettori di uno stesso documento possono essere vicini, poiché diversi token descrivono spesso concetti simili. HierarchicalTokenPooling, la libreria Sentence Transformers, raggruppa questi vettori tramite clustering di Ward, che a ogni passo fonde i due gruppi che causano il minor aumento della dispersione interna. Ogni gruppo viene poi sostituito dalla sua media, quindi normalizzato. Con pool_factor=4, un documento conserva circa un quarto dei suoi vettori iniziali, in modo che un documento più lungo mantenga proporzionalmente più vettori.
Addestrare il modello alla compressione
Più forte è la compressione, più è facile perdere un’informazione utile. L’addestramento che tiene conto del pooling mira proprio a conservare l’informazione importante quando i vettori vengono fusi. Ogni batch di addestramento combina due perdite: la ricerca con gli embedding di documenti originali e la ricerca dopo la compressione a circa un quarto dei vettori. Queste due perdite vengono mediate e i due obiettivi aggiornano lo stesso encoder.
Poiché il clustering prende decisioni discrete, la fonte indica che l’addestramento utilizza un stimatore straight-through, che trasmette i gradienti dei vettori raggruppati verso i token che li compongono, senza differenziare l’assegnazione ai cluster. La classe MultiVectorTokenPoolingLoss incapsula una perdita di ricerca esistente, qui CachedMultiVectorMultipleNegativesRankingLoss, con pesi di 0,5 per ogni variante. Il modello di partenza è answerdotai/ModernBERT-base, addestrato su 120.000 tripletti MS MARCO per un’epoca.
Risultati su NanoBEIR
Con pool_factor=4, l’nDCG@10 passa da 49,03 a 51,23 per il modello addestrato con pooling, riducendo al contempo del 74,8% il numero di vettori di documenti rispetto a una ricerca non compressa. Anche il modello non compresso migliora, da 51,84 a 52,39. Per una media di 29,58 vettori per documento, il pooling adattativo conserva il 97,3% del punteggio non compresso del modello addestrato, contro il 93,3% del modello addestrato normalmente.
Per ridurre la dimensione di un indice, la fonte suggerisce prima la quantizzazione degli embedding, poi il token pooling, eventualmente con una messa a punto che tiene conto del pooling, e infine gli embedding Matryoshka per ridurre la dimensione di ogni vettore. Questi approcci possono essere combinati. I lavori su cui si basa l’articolo sono descritti nello studio Learn to Pool di Stefan Josef.
