Quantization-Aware Training di jina-embeddings-v4

La quantizzazione è ampiamente utilizzata per affrontare i problemi di scalabilità nell'IA. Il nome la fa sembrare complicata, ma si tratta semplicemente di arrotondare i numeri per far sì che occupino meno spazio. Ciò significa vettori di embedding più piccoli che occupano meno memoria e spazio di archiviazione e un recupero delle informazioni più veloce perché ci vuole meno tempo per confrontare i vettori. La quantizzazione è una tecnica puramente numerica che non si preoccupa del tipo di dati elaborati dal modello o dei casi d'uso, quindi può apportare miglioramenti senza richiedere molta costosa conoscenza del dominio.

Ci si potrebbe aspettare che la quantizzazione comporti dei buoni vecchi compromessi e che nulla sia gratis, dove dobbiamo sacrificare un po' di precisione. In questo articolo, ti mostreremo un modo per renderla senza perdite tramite il training consapevole della quantizzazione (quantization-aware training, QAT). Questa tecnica è utilizzata in jina-embeddings-v4 per fornire embedding più piccoli richiesti in applicazioni in cui lo spazio è fondamentale.

tagPanoramica delle tecniche di quantizzazione

La quantizzazione del modello di solito significa una di queste quattro cose:

• Quantizzazione post-training (Post-training quantization, PTQ)
• Training per output di embedding quantizzati (Output QAT, Output QAT)
• Training per modelli completamente quantizzati (Full QAT, Full QAT)
• Distillare un nuovo modello quantizzato da uno esistente non quantizzato

La quantizzazione post-training (PTQ) accetta il modello di embedding addestrato così com'è e non lo modifica in alcun modo. Si tratta solo di scartare le cifre meno significative dei valori in virgola mobile prodotti dal modello. Arrotondiamo semplicemente i numeri e a volte li scaliamo a un intervallo.

Output QAT significa ottimizzare il modello di embedding per produrre vettori a precisione ridotta ottimali. Ciò significa modificare il modello, ma non cambia la precisione dei pesi del modello e quindi non ne riduce le dimensioni. Viene ridotta solo la dimensione del vettore di output.

Full QAT inizia con un modello a precisione completa completamente addestrato e abbassa la precisione dei pesi del modello, quindi ottimizza le prestazioni di questo modello modificato. Ciò produce un modello significativamente più piccolo e embedding più piccoli, al prezzo di una messa a punto.

La distillazione è il processo di addestramento di un nuovo modello per corrispondere alle prestazioni di uno esistente. Ciò significa creare un nuovo modello progettato da zero come quantizzato e quindi utilizzare il modello esistente per generare tutti i dati di addestramento necessari per addestrarlo fino a quando non si comporta il più fedelmente possibile al modello esistente.

I vantaggi di questi quattro approcci sono riassunti nella tabella seguente:

Tutti e quattro producono embedding più compatti, ma a parte PTQ, tutti richiedono un ulteriore training, mentre solo Full QAT e Distillazione producono modelli nuovi e più veloci. Full QAT e Distillazione sono molto più costosi da implementare perché richiedono molto più training rispetto a Output QAT.

In questo articolo, esamineremo solo PTQ e Output QAT, che non modificano le dimensioni o la velocità del modello di embedding.

tagConfigurazione sperimentale

Per questi esperimenti, il nostro modello di base è jina-embeddings-v4 con l'adattatore di recupero, che produce vettori in virgola mobile a 32 bit di precisione (FP32) in 2048 dimensioni. Ogni embedding ha quindi una dimensione di 8196 byte o 8 kB.

Abbiamo studiato diverse condizioni sperimentali utilizzando attività di benchmark di recupero di query-documenti dalla suite NanoBEIR benchmark. Il processo di recupero utilizza la similarità del coseno tra i vettori per trovare e classificare i documenti che meglio corrispondono alle query.

• Baseline: le prestazioni dei vettori di embedding jina-embeddings-v4 senza alcuna quantizzazione. Tutti questi esperimenti hanno utilizzato una versione beta del modello e le prestazioni di rilascio sono leggermente migliori.
• PTQ: abbiamo quantizzato i vettori di output in vettori binari senza modificare il modello.
• Output QAT: abbiamo quantizzato i vettori di output e applicato l'ottimizzazione all'adattatore di recupero per migliorarne le prestazioni in condizioni quantizzate.

tagLivelli di quantizzazione

Abbiamo sperimentato con quattro diversi livelli di quantizzazione.

• Interi a 8 bit: i valori FP32 vengono ridotti a interi nell'intervallo da -128 a 127, riducendo gli embedding di 4 volte a 2048 byte.
• Interi a 4 bit: come per gli interi a 4 bit, ma mappiamo all'intervallo da -8 a 7, riducendo le dimensioni del vettore di un fattore 8, a 1024 byte.
• Quantizzazione ternaria: tutti i valori vengono mappati a uno dei tre valori: -1, 0, 1. Memorizzato in modo ottimale, questo riduce ogni dimensione a 1,6 bit, riducendo le dimensioni dei vettori di embedding di circa 40 volte a circa 230 byte.
• Quantizzazione binaria: convertiamo i valori scalari FP32 in un bit, utilizzando il tipo di dati torch.sign, che prevede solo due valori, impiegando un bit per l'archiviazione. Ciò riduce i vettori di embedding a 2048 dimensioni da 8192 byte a 128 byte, una riduzione di 64 volte.

tagScalabilità

Per la quantizzazione binaria, la quantizzazione è molto semplice: se un valore vettoriale è superiore a 0 o positivo, viene mappato a 1. Altrimenti, viene mappato a -1.

Per gli altri scenari di quantizzazione, abbiamo normalizzato i valori a un intervallo e quindi arrotondato al valore più vicino consentito dal livello di quantizzazione. I vettori di embedding sono costituiti da numeri di scala compresi tra -∞ e +∞ (o, in pratica, numeri positivi e negativi molto grandi). Usiamo due numeri, $max$ e $min$, per scalare i valori per la quantizzazione.

Per la quantizzazione ternaria, prendiamo ogni componente vettoriale $v$ e lo traduciamo come segue:

• se $v$ ≥ $max$, $v$ diventa 1.
• se $v$ ≤ $min$, $v$ diventa -1.
• se $min$ < $v$ < $max$, $v$ diventa 0.

Per interi a 4 bit:

• se $v$ ≥ $max$, $v$ diventa 7.
• se $v$ ≤ $min$, $v$ diventa -8.
• se $min$ < $v$ < $max$, $v$ diventa $16*(v - min)/(max - min) - 8$, quindi arrotondato all'intero più vicino. Questo scala il valore all'intervallo $[-8,7]$.

Per interi a 8 bit:

• se $v$ ≥ $max$, $v$ diventa 127.
• se $v$ ≤ $min$, $v$ diventa -128.
• se $min$ < $v$ < $max$, $v$ diventa $256*(v - min)/(max - min) - 128$, arrotondato all'intero più vicino. Questo scala il valore all'intervallo $[-128,127]$.

Per calcolare $max$ e $min$, abbiamo utilizzato due approcci:

• Min/Max: abbiamo elaborato i nostri dati in batch e, per ogni batch, abbiamo identificato il componente vettoriale più alto e più basso, impostando $max$ sul più alto e $min$ sul più basso.
• Media mobile sui batch: per ogni batch, abbiamo calcolato la media e la deviazione standard dei componenti vettoriali. Abbiamo mantenuto una media mobile sia della media che della deviazione standard durante l'elaborazione di tutti i batch. Se $avg$ è la media mobile corrente dei valori medi del batch e $std$ è la media mobile corrente delle deviazioni standard, allora per ogni batch:

$max = avg + std$
$min = avg - std$

tagOttimizzazione QAT

Per gli esperimenti PTQ, abbiamo utilizzato il modello così com'è e abbiamo quantizzato gli embedding prodotti utilizzando i metodi descritti sopra.

Per l'Output QAT, abbiamo ottimizzato il modello utilizzando la stima straight-through. Ciò significa che invertiamo il processo di quantizzazione, ripristinando la precisione completa ai valori, prima di calcolare la perdita (cioè l'errore), e quindi utilizziamo tale metrica di perdita per ottimizzare il modello.

In ogni caso, abbiamo eseguito il fine-tuning per 10.000 passaggi, salvando un checkpoint ogni 500 passaggi. Abbiamo quindi conservato il checkpoint con il punteggio più alto sul benchmark NanoBEIR.

tagQuantizzazione asimmetrica

PTQ e Output QAT riducono le dimensioni dei vettori di embedding, ma non riducono le dimensioni del modello o la velocità di inferenza; tutti i risparmi sono nelle dimensioni degli 向量模型 (embedding) dei documenti memorizzati e nella velocità di recupero.

Di conseguenza, abbiamo testato sia la quantizzazione dei vettori di query sia il fatto di lasciarli non quantizzati al momento del recupero, perché in entrambi i casi non cambia la dimensione dei vettori di embedding memorizzati.

tagRisultati

Abbiamo testato nove condizioni in totale, riassunte nelle tabelle seguenti:

Tabella 2: Condizioni sperimentali

Tabella 3: Punteggio medio (in % corretto) per ogni condizione sui dodici benchmark NanoBEIR.

Dalla tabella sopra si può notare che il fine-tuning per la quantizzazione migliora i punteggi. L'unica differenza tra le condizioni PTQ Binary e QAT Binary è il fine-tuning, e la differenza di punteggio è significativa. Allo stesso modo, vediamo un miglioramento di quasi il 2% nei punteggi tra le condizioni PTQ Binary Docs Only e QAT Binary Docs Only, che si distinguono solo per lo stesso fine-tuning.

Non sorprende che vediamo anche che i punteggi generalmente migliorano quanto meno quantizziamo, con la quantizzazione a 4 bit che ottiene un punteggio migliore della trinaria e la trinaria migliore della binaria. Tuttavia, passare a 8 bit non sembra aver migliorato nulla.

Abbiamo testato solo il fatto di lasciare le query non quantizzate nei casi binari, ma questo sembra migliorare le prestazioni.

Infine, i nostri test suggeriscono che il metodo di scaling della media mobile supera il semplice approccio min/max.

tagConclusione

La quantizzazione ha alcuni importanti vantaggi operativi per i modelli di embedding, riducendo significativamente le dimensioni dei vettori di embedding e accelerando il recupero delle informazioni. Mentre la semplice quantizzazione post-training (PTQ) fornisce vantaggi immediati in termini di memoria e archiviazione, i nostri esperimenti dimostrano che l'addestramento consapevole della quantizzazione (QAT) mitiga in modo significativo le inevitabili perdite di precisione. Il fine-tuning ha prodotto costantemente punteggi migliori.

Il grado di quantizzazione ha un impatto diretto sulle prestazioni, che è ciò che ci si aspetterebbe da un metodo basato sulla riduzione della precisione dei valori. La quantizzazione meno aggressiva (ad esempio, a 4 bit) generalmente supera i metodi più aggressivi (ad esempio, binaria), ma sorprendentemente, non c'è stata una differenza significativa nelle prestazioni tra la quantizzazione a 8 bit e quella a 4 bit. Sembrerebbe che fino a quando non si raggiunge una certa soglia di imprecisione, c'è poca differenza tra una quantizzazione maggiore e minore.

Anche le strategie di scaling sono significative, con il metodo della media mobile che mostra risultati superiori rispetto a un approccio min/max fisso. L'uso di valori di scaling relativi ai dati sembra funzionare significativamente meglio e merita un'ulteriore esplorazione.

La quantizzazione può farti ottenere di più dai tuoi modelli di embedding a un costo inferiore. Anche se questo articolo non esplora tutte le opzioni per la quantizzazione, ne esplora due che sono facilmente accessibili e hanno reali vantaggi da offrire. Stiamo lavorando per perfezionare e migliorare le strategie di quantizzazione in modo da poter ridurre ulteriormente i costi degli utenti e prevediamo di rilasciare il supporto binario per jina-embeddings-v4 nel prossimo futuro.