Apprendimento di rappresentazioni Matryoshka spiegato: il metodo alla base degli embedding di testo efficienti di OpenAI
Lo sviluppo di modelli di machine learning reali comporta sempre un compromesso tra costo e prestazioni. Ad esempio, più grandi sono il modello e il dataset che utilizziamo durante l’addestramento, più capace sarà in definitiva il nostro modello addestrato. Tuttavia, questa maggiore capacità comporta tempi più lunghi e costi computazionali più elevati, poiché l’addestramento richiede durate significativamente maggiori. Lo stesso principio si applica durante l’inferenza, dove i modelli più grandi tendono a generare rappresentazioni delle feature più grandi che richiedono maggiore memoria per l’archiviazione.
Poiché le esigenze durante l’addestramento e l’inferenza del machine learning possono variare significativamente da un caso d’uso all’altro, è essenziale disporre di un metodo che ci consenta di scambiare una piccola parte delle prestazioni di un modello in cambio di costi ridotti. È qui che entra in gioco un metodo come Matryoshka Representation Learning (MRL) . Ad esempio, il modello text-embedding-3-small di OpenAI utilizza MRL per consentire agli sviluppatori di accorciare gli embedding mantenendo le loro proprietà fondamentali di rappresentazione dei concetti. Consentendo la regolazione delle dimensioni degli embedding, MRL aiuta a trovare il perfetto equilibrio tra efficienza dei costi e prestazioni del modello.
Questo articolo esplorerà come funziona MRL, la sua implementazione e come consente modelli di machine learning scalabili ed efficienti. Iniziamo dalla motivazione alla base di MRL.
La motivazione alla base di Matryoshka Representation Learning (MRL)
Il compromesso tra costo e prestazioni è sempre una considerazione prima di sviluppare o utilizzare qualsiasi modello di machine learning.
Durante l’addestramento, i parametri di un modello e la quantità di dati di addestramento influenzano direttamente le sue prestazioni finali. Modelli e dataset di addestramento più grandi tendono a produrre modelli più capaci. Tuttavia, anche il costo computazionale per addestrare questi modelli è più elevato.
Nel frattempo, durante l’inferenza, i modelli più grandi producono rappresentazioni delle feature più grandi. Questo aumento delle rappresentazioni delle feature significa che abbiamo bisogno di più memoria per archiviarle.
Usiamo BERT come esempio nel contesto del recupero di informazioni. BERT è uno dei primi modelli di deep learning basati su Transformer ad aver raggiunto prestazioni all’avanguardia su diversi dataset di benchmark nella classificazione del testo, nel Named Entity Recognition (NER), nel question answering (QnA) e altro ancora.
Figura: Workflow della generazione di embedding utilizzando il modello BERT base.
In poche parole, BERT prende una frase o una parola come input e la trasforma in un embedding di dimensione fissa. Questo embedding cattura il significato semantico dell’input originale e la sua dimensione dipende dalla specifica variante del modello utilizzata.
BERT ha due varianti principali: il modello base e il modello large. Il modello base produce un embedding a 768 dimensioni, mentre il modello large genera un embedding a 1024 dimensioni.
Supponiamo di voler archiviare 10 milioni di embedding in un vector database per un caso d’uso di recupero di informazioni. Se utilizziamo il modello BERT base in formato FP32, avremo bisogno di circa 768×10M×4=30,72 GB di memoria per archiviare questi embedding. Al contrario, utilizzare il modello BERT large nello stesso scenario richiederebbe circa 40,9 GB. Inoltre, una dimensione maggiore degli embedding può rallentare il processo di recupero aumentando la complessità computazionale delle similarity searches. Tuttavia, il vantaggio principale dell’utilizzo di embedding più grandi è la loro maggiore capacità di recuperare informazioni altamente rilevanti rispetto a embedding più piccoli.
Una soluzione ideale a questi problemi nel recupero delle informazioni consiste nel suddividere l'intero processo in due parti: selezione preliminare e riordinamento. La selezione preliminare si riferisce al recupero di un insieme iniziale di documenti candidati da una vasta raccolta nel nostro database vettoriale come Milvus. Il riordinamento prende quindi questi candidati selezionati preliminarmente e li riordina per massimizzare la pertinenza dei risultati finali.
Figura- Flusso di lavoro di selezione preliminare e riordinamento
Figura: Flusso di lavoro di selezione preliminare e riordinamento.
Per eseguire la selezione preliminare, l'algoritmo calcola la similarità tra gli embedding della query e una vasta raccolta di embedding archiviati nel database vettoriale. Pertanto, l'utilizzo di embedding di dimensioni più piccole può rendere questo calcolo più efficiente e rapido. Al contrario, il riordinamento enfatizza la precisione, il che significa che trarremmo vantaggio da embedding di dimensioni maggiori.
Il problema di questo metodo è che, una volta selezionato un modello particolare, la dimensione degli embedding o delle rappresentazioni delle caratteristiche che produce è fissa. Questa mancanza di flessibilità limita la nostra capacità di utilizzare embedding di dimensioni diverse da un singolo modello.
Figura- Flusso di lavoro di un modello che genera un embedding di dimensione fissa..png
Figura: Flusso di lavoro di un modello che genera un embedding di dimensione fissa.
L'approccio Matryoshka Representation Learning (MRL) offre una soluzione interessante a questo problema: i modelli addestrati con questo metodo possono produrre embedding di varie dimensioni. Esploreremo in dettaglio il meccanismo di questo approccio nella prossima sezione.
Che cos'è Matryoshka Representation Learning (MRL)?
Matryoshka Representation Learning (MRL) è un metodo per addestrare reti neurali a produrre rappresentazioni multi-scala all'interno di un singolo modello. Ispirato alle bambole Matryoshka, in cui le bambole più piccole si inseriscono dentro quelle più grandi, MRL consente a un modello di produrre rappresentazioni di dimensioni variabili (da grossolane a fini) usando un singolo passaggio in avanti. Questo approccio consente ai modelli di deep learning di comprendere meglio relazioni complesse e sfumature nei dati. MRL è particolarmente efficace per attività come la ricerca semantica, il recupero delle informazioni, l'elaborazione multilingue e qualsiasi applicazione che richieda rappresentazioni sfumate dei dati a diversi livelli di astrazione.
Visualizzazione degli embedding Matryoshka con più livelli di dettaglio
Figura: Visualizzazione degli embedding Matryoshka con più livelli di dettaglio
I modelli di embedding più diffusi che hanno adottato l'approccio Matryoshka Representation Learning (MRL) includono text-embedding-3-large di OpenAI, nomic-embed-text-v1 di Nomic e gte-multilingual-base di Alibaba.
Come funziona l'approccio MRL
L'approccio MRL ci consente di estrarre rappresentazioni delle caratteristiche da qualsiasi modello di machine learning in diverse dimensioni. Ad esempio, invece di utilizzare le 1024 dimensioni originali di un embedding, possiamo utilizzare le prime 16, 32, 64, 128 o 256 dimensioni (o qualsiasi dimensione scegliamo). La chiave di questa capacità risiede nel modo in cui il modello viene addestrato quando si implementa l'approccio MRL.
Figura- Addestramento della funzione di perdita MRL e suo caso d'uso durante l'inferenza
Figura: Addestramento della funzione di perdita MRL e suo caso d'uso durante l'inferenza. Fonte.
Durante l’addestramento del modello con MRL, ottimizziamo non una sola funzione di perdita, come avviene tipicamente nell’addestramento standard dei modelli, ma diverse. Ogni funzione di perdita mira a ottimizzare la rappresentazione delle feature per una dimensione specifica che definiamo in anticipo. Se impostiamo cinque diverse dimensioni delle feature, dovremo ottimizzare cinque funzioni di perdita distinte durante l’addestramento.
Supponiamo di voler addestrare un modello BERT large utilizzando l’approccio MRL. Come menzionato nella sezione precedente, questo modello genera un embedding di dimensione 1024. Se impostiamo le dimensioni delle feature a 32, 64, 128, 256 e 1024, ottimizzeremo cinque diverse funzioni di perdita durante l’addestramento, in modo che il modello impari a generare embedding ottimizzati in quelle dimensioni.
L’ottimizzazione su queste varie dimensioni è semplice: MRL scompone la funzione di perdita complessiva nella somma delle perdite per ciascuna singola dimensione. Usando il nostro esempio sopra, ciò può essere espresso come:
Il processo di ottimizzazione per un modello addestrato con MRL può essere rappresentato più formalmente come segue:
Figura: Equazione di ottimizzazione dell’approccio MRL. Fonte.
Come puoi vedere, la definizione della funzione di perdita sopra è piuttosto generale. Pertanto, l’approccio MRL può essere applicato a quasi qualsiasi modello ed è del tutto indipendente dall’architettura del modello. Possiamo anche effettuare il fine-tuning di modelli pre-addestrati come BERT o qualsiasi altro modello basato su transformer affinché producano embedding di varie dimensioni. Continuando con il nostro scenario sopra, ora possiamo generare embedding a 1024 dimensioni dal nostro modello BERT large addestrato e embedding a 32, 64, 128 e 256 dimensioni.
Nel complesso, dopo aver addestrato un modello con MRL, le dimensioni iniziali delle feature generate dal modello addestrato contengono informazioni più significative rispetto alle dimensioni successive. Le prime dimensioni contengono dettagli di alto livello, mentre le dimensioni successive si concentrano su informazioni più granulari, imitando la struttura di una matrioska.
Tuttavia, ciò non significa che gli embedding più corti siano sempre semplicemente versioni troncate di quelli più lunghi. I valori di ciascun elemento negli embedding più corti potrebbero differire da quelli negli embedding più lunghi a causa del fattore di scala applicato a ciascuna dimensione delle feature durante l’addestramento. Tuttavia, se impostiamo il fattore di scala in ciascuna dimensione in modo che sia uguale, i valori di ciascun elemento tra embedding più corti e più lunghi potrebbero apparire identici.
Risultati sperimentali di MRL
L’approccio MRL è stato valutato su modelli di machine learning in diverse modalità, tra cui testo, visione e visione-testo. I modelli ResNet50 e ViT rappresentano modelli basati sulla visione, BERT rappresenta modelli basati sul testo e ALIGN rappresenta la combinazione di visione e testo. Questi modelli sono stati valutati principalmente per due casi d’uso comuni: classificazione e recupero. Parliamo prima della classificazione.
Confrontando le prestazioni del modello ResNet50 addestrato sul dataset ImageNet-1K con un ResNet50 standard addestrato indipendentemente, il modello MRL raggiunge un’accuratezza top-1 comparabile con varie dimensioni di rappresentazione delle feature.
Per valutare ulteriormente l’utilità delle rappresentazioni delle feature per i task downstream, è stata misurata anche l’accuratezza del 1-nearest neighbor (1-NN) per ciascuna dimensione della rappresentazione delle feature. La configurazione per calcolare il 1-NN è la seguente: dati 1,3K campioni di immagini nel database, il compito è trovare il vicino più prossimo per ciascuna delle 50K query. Con questa configurazione, ResNet50 addestrato con MRL è fino al 2% più accurato per ciascuna dimensione della rappresentazione delle feature rispetto alla sua controparte a feature fisse.
Figura- Accuratezza Top-1 e accuratezza 1-NN di ResNet50 su ImageNet-1K
Figura: Accuratezza Top-1 e accuratezza 1-NN di ResNet50 su ImageNet-1K. Fonte.
Nel frattempo, anche le prestazioni del modello ViT addestrato con MRL sul dataset JFT-300M sono molto competitive in tutte le dimensioni della rappresentazione. La sua accuratezza 1-NN è comparabile a quella di ViT addestrato con rappresentazioni delle feature di dimensione fissa. Come mostrato nella figura seguente, le prestazioni del modello MRL sono anche migliori della sua controparte a dimensione fissa con rappresentazioni delle feature più piccole, in parte perché vengono selezionate feature casuali dal modello a dimensione fissa per rappresentare dimensioni inferiori. Una tendenza simile si osserva con il modello ALIGN addestrato utilizzando l’approccio MRL; le sue prestazioni corrispondono a quelle del modello ALIGN addestrato con rappresentazioni a dimensione fissa.
Figura- Accuratezza 1-NN di ViT e ALIGN su ImageNet-1K
Figura: Accuratezza 1-NN di ViT e ALIGN su ImageNet-1K. Fonte.
Il principale vantaggio dell’utilizzo dell’approccio MRL è la sua flessibilità nel passare tra diverse dimensioni delle rappresentazioni delle feature dallo stesso modello. Per sfruttare appieno questo punto di forza, sono stati condotti anche test sulla classificazione adattiva utilizzando ResNet50 su ImageNet-1K.
Questa configurazione prevede l’apprendimento della probabilità softmax massima per determinare le transizioni da rappresentazioni delle feature più piccole a quelle più grandi. I risultati dei test mostrano che il modello MRL con una rappresentazione delle feature a 37 dimensioni è accurato quanto il modello ResNet50 addestrato con una rappresentazione delle feature fissa a 512 dimensioni, pur essendo solo dello 0,8% meno accurato rispetto al modello a 2048 dimensioni.
Figura- Accuratezza Top 1 del modello ResNet50 nella classificazione adattiva (sinistra) e mAP di ResNet50 nel caso d’uso del recupero di immagini (destra)
Figura: Accuratezza Top 1 del modello ResNet50 nella classificazione adattiva (sinistra) e mAP di ResNet50 nel caso d’uso del recupero di immagini (destra). Fonte.
Ora parliamo del recupero. Come forse sai, l’obiettivo del recupero è trovare un insieme di contesti simili per una determinata query. Pertanto, è stata testata la qualità del recupero tra ResNet50 addestrato con MRL e rappresentazioni delle feature di dimensione fissa utilizzando la mean Average Precision (mAP). Il modello MRL mostra un miglioramento della mAP fino al 3% rispetto alla sua controparte a dimensione fissa per ogni dimensione della rappresentazione delle feature, come illustrato nella visualizzazione sopra.
Poiché con MRL possiamo utilizzare diverse dimensioni della rappresentazione delle feature, è anche particolarmente interessante esplorare come possiamo accelerare il processo di recupero mantenendo l’accuratezza con MRL. Possiamo testarlo in uno scenario chiamato recupero adattivo.
Come accennato in precedenza, ci sono due fasi principali in un processo di recupero: shortlist e reranking. La configurazione del recupero adattivo prevedeva l’utilizzo di una rappresentazione delle feature a 16 dimensioni per creare una shortlist di 200 candidati, mentre per il reranking viene utilizzata una rappresentazione delle feature a 2048 dimensioni. Questo approccio viene confrontato con un altro metodo che utilizza una rappresentazione delle feature a 2048 dimensioni sia per la shortlist sia per il reranking.
Su ImageNet-1K, l’utilizzo del modello MRL può teoricamente accelerare il processo di retrieval fino a 128 volte. Lo definiamo teorico perché nelle applicazioni reali raramente utilizziamo un algoritmo ingenuo di nearest neighbor per le attività di retrieval. Invece, algoritmi di approximate nearest neighbor (ANN) come FAISS, ANNOY o HNSW vengono comunemente implementati per migliorare la velocità di retrieval. In questo caso, la configurazione che utilizza il modello MRL raggiunge un’accelerazione di 14 volte rispetto ai processi di retrieval che utilizzano l’algoritmo HNSW su hardware identico.
Un risultato simile è stato osservato con il dataset ImageNet-4K, dove gli autori del paper hanno utilizzato un modello MRL con una rappresentazione delle feature a 64 dimensioni per la shortlisting e una rappresentazione a 2048 dimensioni per il reranking. I risultati indicano un’accelerazione teorica di 32 volte e un’accelerazione reale di sei volte.
Il compromesso tra mAP@10 e MFLOPs/Query per Adaptive Retrieval (AR) su ImageNet-1K (sinistra) e ImageNet-4K (destra).
Il compromesso tra mAP@10 e MFLOPs/Query per Adaptive Retrieval (AR) su ImageNet-1K (sinistra) e ImageNet-4K (destra). Fonte.
Implementazione dell’approccio MRL
In questa sezione, daremo un’occhiata alla semplice implementazione dell’approccio MRL. In particolare, useremo un modello base MPNet addestrato con l’approccio MRL sul dataset NLI per generare embedding di diverse dimensioni. Come forse già sai, il modello base MPNet originale può generare solo un embedding a 768 dimensioni. Tuttavia, questo modello MPNet addestrato può generare embedding a 768, 512, 256, 128 e 64 dimensioni, dove il peso di ciascuna dimensione durante l’addestramento è impostato a 1.
Possiamo usare questo modello con l’aiuto della libreria SentenceTransformers con il seguente codice.
from sentence_transformers import SentenceTransformer
matryoshka_dim_short = 64
matryoshka_dim_long = 768
text = ["The weather is so nice!"]
short_embedding = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim_short).encode(text)
long_embedding = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim_long).encode(text)
print(f"Shape: {short_embedding.shape, long_embedding.shape}")
print(short_embedding[0][0:10])
print(long_embedding[0][0:10])
"""
Output:
Shape: ((1, 64), (1, 768))
[-0.33891088 0.01647538 -0.29915053 0.24952686 -0.04321517 -0.31616145
-0.12996909 -0.05221268 0.02296597 0.07074839]
[-0.33891088 0.01647538 -0.29915053 0.24952686 -0.04321517 -0.31616145
-0.12996909 -0.05221268 0.02296597 0.07074839]
"""
Nel codice sopra, utilizziamo il modello base MPNet addestrato per generare embedding in due dimensioni diverse: 64 e 768. Poiché il peso o fattore di scaling di ciascuna dimensione è impostato a 1, i due embedding hanno elementi identici, come puoi vedere nei primi 10 elementi dei due embedding sopra.
Poiché il valore di ciascun elemento è identico, la similarità tra i due embedding sarebbe 1.
from sentence_transformers.util import cos_sim
similarities = cos_sim(short_embedding[0], long_embedding[0][:matryoshka_dim_short])
print(similarities)
# tensor([[1.]])
Se desideri addestrare il tuo modello con l’approccio MRL, puoi dare un’occhiata al repo GitHub ufficiale di questo metodo.
Conclusione
L’approccio MRL introduce una soluzione per bilanciare il compromesso tra costo e prestazioni nel machine learning. Consentendo a qualsiasi modello di machine learning di produrre rappresentazioni delle feature di dimensioni variabili, MRL offre la flessibilità necessaria per ottimizzare la velocità o l’accuratezza a seconda del nostro caso d’uso e delle risorse disponibili. Abbiamo visto che questa adattabilità è preziosa in applicazioni come la classificazione e il retrieval, dove passare da rappresentazioni più piccole a più grandi può migliorare significativamente l’efficienza senza sacrificare enormemente le prestazioni.
I risultati sperimentali dimostrano che MRL eguaglia e spesso supera l’accuratezza dei tradizionali modelli a dimensione fissa in diversi domini, inclusi testo, visione e attività multimodali. Inoltre, la compatibilità dell’approccio con architetture esistenti come ResNet50, ViT e BERT evidenzia la sua applicabilità a diverse modalità di modello. Con il potenziale di accelerazioni sostanziali in scenari reali e una migliore qualità del retrieval, MRL rappresenta un promettente progresso verso soluzioni di machine learning più efficienti e versatili.
Risorse correlate
Continua a leggere
Introducing Loon: A New Storage Engine for Vector Data That Never Stops Changing
Loon is a new storage engine for Milvus 3.0 and Zilliz Vector Lakebase, built to manage evolving vector datasets with ColumnGroups, row ID alignment, and Manifests.
Why and How to Migrate from Self-Hosted Milvus to Zilliz Cloud
A simple, step-by-step guide to migrating from Milvus to Zilliz Cloud. Learn both endpoint and backup methods for a smooth, scalable vector database migration.
How to Use Anthropic MCP Server with Milvus
MCP + Milvus: Streamline AI agent development with standardized data access, eliminating integration hassles while enhancing context and flexibility.