Introduzione ai modelli multi-vettore con sentence transformers

Sentence Transformers è una libreria Python progettata per l'utilizzo e l'addestramento di modelli di embedding e reranker applicabili a una vasta gamma di casi d'uso, come la generazione aumentata da recupero (RAG), la ricerca semantica, la somiglianza testuale semantica e molto altro. L'aggiornamento alla versione 6.0 rappresenta un momento cruciale nello sviluppo della libreria, introducendo un quarto tipo di modello: MultiVectorEncoder, specificamente ottimizzato per il recupero con late interaction nello stile di ColBERT. Insieme a questa novità, viene presentato un approccio di addestramento completo dedicato a questo tipo di architettura.

Questo articolo approfondisce come utilizzare MultiVectorEncoder per affinare un modello multi-vettore che supera i retriever generici su dati specifici del vostro dominio. La metodologia presentata permette non solo l'affinamento di modelli esistenti, ma anche l'addestramento da zero di nuovi modelli multi-vettore estremamente performanti. Tutto ciò che viene descritto può essere eseguito con il semplice comando pip install -U "sentence-transformers[train]".

Concetti fondamentali: embedding densi vs. multi-vettore

Prima di approfondire gli aspetti tecnici dell'addestramento, è essenziale comprendere le differenze fondamentali tra i due approcci principali ai modelli di embedding.

Un modello di embedding denso comprime l'intero testo in un singolo vettore, e la somiglianza viene calcolata come un singolo prodotto scalare tra due di questi riassunti vettoriali. Questo approccio è efficiente dal punto di vista dello spazio di memoria e della velocità computazionale, ma comporta una compressione significativa dell'informazione.

Un modello multi-vettore (noto anche come late interaction o modello nello stile ColBERT) segue una filosofia completamente diversa. Anziché comprimere il testo, mantiene un vettore piccolo per ogni token. La similitudine tra una query e un documento viene calcolata mediante l'operatore MaxSim, dove ogni token della query trova il token del documento che corrisponde meglio e i punteggi vengono sommati. Questo approccio token-by-token preserva esattamente i segnali granulari che un singolo vettore deve mediare, il che generalmente significa un recupero più robusto, al costo di un indice più grande.

L'articolo complementare sui modelli di embedding multi-vettore copre in dettaglio l'architettura, la codifica, il punteggio e l'indicizzazione, quindi qui ci concentreremo sui componenti di addestramento specifici.

Perché affinare modelli multi-vettore

L'affinamento di modelli multi-vettore migliora significativamente le loro prestazioni di recupero nel vostro dominio specifico. Il vocabolario, lo stile delle query e la nozione di rilevanza differiscono notevolmente tra la ricerca web, la scoperta legale, la ricerca di codice e la revisione della letteratura scientifica. Poiché le query e i documenti sono abbinati token per token, i modelli multi-vettore acquisiscono segnali di dominio granulari che i modelli a singolo vettore tendono a mediare, e rispondono molto bene anche a modest quantità di dati di affinamento in-dominio.

Oltre a questo aspetto teorico, la maggior parte dei modelli di recupero rilasciati pubblicamente è stata configurata per passaggi brevi. I checkpoint classici di ColBERT troncano i documenti a 180 o 300 token, e molti modelli densi popolari a 256 o 512, perché i loro dati di addestramento in stile MS MARCO raramente superano tali limiti. Se i vostri documenti sono lunghi, questi modelli scartano silenziosamente la maggior parte di ogni documento prima di calcolarne il punteggio. Nella valutazione medica dell'autore, con passaggi che hanno una media di 941 token, è stato misurato che questo troncamento costa fino a 0,24 NDCG@10, considerevolmente più di qualsiasi differenza tra architetture di modelli. Quando addestrate il vostro modello, configurate la lunghezza del documento che i vostri dati richiedono effettivamente.

LightOn ha incontrato la stessa dinamica nel recupero di codice, dove il modello LateOn generale non era sufficiente e hanno addestrato LateOn-Code. Il vostro dominio, che sia medico, legale, finanziario o documenti interni dell'azienda, probabilmente non avrà un modello ufficiale pre-addestrato. Questo articolo mostra come costruirlo autonomamente, in poche ore, su una singola GPU di consumo.

Componenti dell'addestramento multi-vettore

L'addestramento di modelli MultiVectorEncoder coinvolge diversi componenti interconnessi:

Il modello stesso (architettura e checkpoint iniziale)
Dataset per l'addestramento e la valutazione
Funzioni di perdita appropriate al compito
Argomenti di addestramento (learning rate, batch size, etc.)
Valutatori per misurare le prestazioni
La classe Trainer che orchestra tutto il processo

Esamineremo ciascuno di questi componenti in dettaglio, accompagnati da esempi pratici di come possono essere utilizzati per affinare modelli multi-vettore robuste.

Scelta del checkpoint iniziale

L'addestramento multi-vettore offre una vera scelta di punto di partenza, e questa scelta è più importante di quanto si potrebbe aspettare.

Se desiderate affinare ulteriormente un modello multi-vettore esistente, non dovrete preoccuparvi dell'architettura: il checkpoint porta con sé la sua ricetta. Tuttavia, la prima cosa da verificare è la configurazione della lunghezza, poiché molti checkpoint rilasciati limitano i documenti a 180-512 token, mentre i passaggi medici possono arrivare a 1.400 token. La famiglia mLateOn serve già il contesto completo di 8192 token della backbone, ma se il vostro checkpoint iniziale ha limitazioni, potete alzarle configurando il model_max_length al momento del caricamento.

Un'altra modifica utile è aggiungere una lista di esclusione (skiplist) di punteggiatura che esclude i token di punteggiatura dal punteggio e dall'archiviazione nel documento. In un'ablazione a 4 vie (nessuno, punteggiatura, stopword, entrambi), questo ha mostrato un modesto miglioramento della qualità e riduce l'indice dei documenti del 9,6% su questi dati senza costi aggiuntivi.

È possibile anche puntare MultiVectorEncoder su qualsiasi trasformatore di base, e una proiezione a livello di token fresco e inizializzato casualmente viene aggiunta automaticamente. Questo è il pipeline classico di ColBERT: un Transformer che produce embedding di token contestualizzati, un Dense a livello di token che proietta ciascuno di essi fino a 128 dimensioni, un MultiVectorMask che decide quali token contano durante il punteggio, e un Normalize a livello di token.

La proiezione inizia casualmente, quindi l'addestramento è necessario prima che questo modello sia utile. Interessantemente, questo funziona bene anche con backbone di embedding densi forti. Una proiezione fresca su Alibaba-NLP/gte-modernbert-base ha raggiunto entro 0,03 dei punti di partenza dei checkpoint esistenti negli esperimenti dell'autore, derivando solo dalla proiezione e da 25.000 coppie di addestramento.

Configurazione del modello e personalizzazione

I trucchi di tokenizzazione classici di ColBERT sono tutti disattivati per impostazione predefinita e configurabili:

[MASK] query expansion
[Q]/[D] prefix tokens
Document length cap
Punctuation skiplist

Per una descrizione completa di tutte le opzioni, consultate la sezione "Creating Custom Models". Vale la pena notare che l'espansione di query [MASK] è stata testata in quattro configurazioni per il dominio medico e nessuna di esse ha fatto una differenza misurabile, quindi non sentitevi obbligati a raggiungere per la ricetta classica.

Esperimenti con punti di partenza diversi

Per determinare il punto di partenza ottimale, è stato condotto un esperimento diretto prendendo sei punti di partenza e addestrando ciascuno con la ricetta identica su 25.000 coppie di domande-passaggi medici da MIRIAD, quindi valutando su 1.000 domande trattenute rispetto a un corpus di 50.000 passaggi.

Il risultato è sorprendente e si è replicato su due famiglie di modelli: i checkpoint non supervisionati si adattano a un nuovo dominio molto meglio dei loro fratelli finiti, superandoli nonostante inizino più in basso. Questi checkpoint si situano dopo il preaddestramento contrastivo su larga scala ma prima dell'affinamento supervisionato sul recupero generale, quindi contengono tutta la struttura di late interaction senza la sintonizzazione per uso generale che l'addestramento di dominio deve poi annullare. I checkpoint finiti, al contrario, hanno mostrato un movimento minimo o addirittura regressione, a ogni learning rate provato.

Quindi, se la famiglia di modelli che vi piace pubblica un checkpoint pre-supervisionato, iniziate da lì. Se no, una proiezione fresca su una backbone forte pre-addestrata per il recupero è un candidato molto vicino. Continuare da un checkpoint completamente finito è l'opzione più debole per l'adattamento di dominio, nonostante sia quella che sembra più naturale.

Gestione dei dataset

MultiVectorEncoderTrainer utilizza istanze di datasets.Dataset o datasets.DatasetDict per l'addestramento e la valutazione. È possibile caricare dati da Hugging Face Datasets Hub o utilizzare dati locali in qualsiasi formato preferito (CSV, JSON, Parquet, Arrow, SQL).

Molti dataset pubblici che funzionano perfettamente con Sentence Transformers sono stati etichettati con "sentence-transformers" su Hugging Face Hub, permettendovi di trovarli facilmente su https://huggingface.co/datasets?other=sentence-transformers. Vale la pena navigare attraverso questi per trovare dataset pronti all'uso che potrebbero essere utili per i vostri compiti, domini o lingue.

Il dataset utilizzato in questo articolo comprende 4,4 milioni di domande mediche da MIRIAD, ciascuna abbinata al passaggio sorgente che contiene la sua risposta (con una media di 941 token). Semplici coppie (query, passaggio rilevante) come queste sono i dati di addestramento per il recupero più facili da raccogliere per il vostro dominio, e come vedrete, sono tutto ciò che serve.

Caricamento da Hugging Face Datasets Hub

Potete usare la funzione load_dataset per caricare dati da dataset su Hub:

from datasets import load_dataset
dataset = load_dataset("your-dataset-name")

Caricamento da dati locali

Potete anche usare load_dataset per caricare dati locali in formati file comuni:

dataset = load_dataset("csv", data_files="your_data.csv")
dataset = load_dataset("json", data_files="your_data.json")

Creazione di dataset da dizionari

Se i vostri dati locali richiedono pre-elaborazione, potete usare datasets.Dataset.from_dict per inizializzare il vostro dataset con un dizionario di liste:

from datasets import Dataset
data_dict = {
    "query": ["domanda 1", "domanda 2"],
    "passage": ["passaggio 1", "passaggio 2"]
}
dataset = Dataset.from_dict(data_dict)

Compatibilità dataset-loss function

È importante che il formato del vostro dataset corrisponda alla vostra funzione di perdita (o che scegliate una funzione di perdita che corrisponda al vostro formato dataset). Verificare se un formato dataset funziona con una funzione di perdita comporta due passaggi:

Leggere la documentazione della funzione di perdita per comprendere i campi richiesti
Assicurarsi che il vostro dataset contenga esattamente questi campi con i nomi corretti

Ci sono due convenzioni specifiche per i modelli multi-vettore:

I nomi dei campi devono corrispondere alle aspettative della funzione di perdita
I dati devono essere formattati coerentemente (es. query, passaggio positivo, passaggi negativi opzionali)

Funzioni di perdita per modelli multi-vettore

Le funzioni di perdita quantificano quanto bene un modello si comporta per un dato batch di dati, permettendo a un ottimizzatore di aggiornare i pesi del modello per produrre valori di perdita più favorevoli (cioè inferiori). La funzione di perdita giusta per il vostro compito dipende dai dati che avete e da ciò che state cercando di ottenere. È possibile trovare un elenco completo delle opzioni nella panoramica delle funzioni di perdita.