Quantisierungsbewusstes Training von jina-embeddings-v4

Die Quantisierung ist eine weit verbreitete Methode zur Lösung von Skalierungsproblemen in der KI. Der Name klingt kompliziert, aber es handelt sich lediglich um das Abrunden von Zahlen, damit sie weniger Speicherplatz beanspruchen. Dies bedeutet kleinere Vektorraummodelle (Embeddings), die weniger Speicherplatz benötigen, und eine schnellere Informationsabfrage, da der Vergleich von Vektoren weniger Zeit in Anspruch nimmt. Die Quantisierung ist eine rein numerische Technik, bei der es keine Rolle spielt, welche Art von Daten Ihr Modell verarbeitet oder welche Anwendungsfälle Sie haben. Sie kann also Verbesserungen bringen, ohne dass viel teures Fachwissen erforderlich ist.

Man könnte erwarten, dass die Quantisierung mit den üblichen Kompromissen verbunden ist und nichts umsonst ist - wobei wir etwas Präzision opfern müssen. In diesem Artikel zeigen wir Ihnen einen Weg, dies durch quantisierungsbewusstes Training (QAT) verlustfrei zu gestalten. Diese Technik wird in jina-embeddings-v4 verwendet, um kleinere Vektorraummodelle (Embeddings) bereitzustellen, die in speicherkritischen Anwendungen benötigt werden.

tagÜberblick über die Quantisierungstechniken

Modellquantisierung bedeutet in der Regel eines von vier Dingen:

• Post-Training-Quantisierung (PTQ)
• Training für quantisierte Vektorraummodell-Ausgaben (Output QAT)
• Training für vollständig quantisierte Modelle (Full QAT)
• Destillieren eines neuen quantisierten Modells aus einem bestehenden, nicht quantisierten Modell

Die Post-Training-Quantisierung (PTQ) akzeptiert das trainierte Vektorraummodell (Embedding)-Modell so, wie es ist, und verändert es in keiner Weise. Es geht lediglich darum, die am wenigsten signifikanten Stellen der vom Modell erzeugten Gleitkommawerte zu verwerfen. Wir runden die Zahlen einfach ab und skalieren sie manchmal auf einen bestimmten Bereich.

Output QAT bedeutet, das Vektorraummodell (Embedding)-Modell so zu optimieren, dass es optimale Vektoren mit reduzierter Präzision erzeugt. Dies bedeutet, dass das Modell verändert wird, aber es ändert nichts an der Präzision der Gewichte des Modells und reduziert daher nicht seine Größe. Nur die Größe des Ausgabvektors wird reduziert.

Full QAT beginnt mit einem vollständig trainierten Modell mit voller Präzision und reduziert die Präzision der Modellgewichte. Anschließend wird die Leistung dieses modifizierten Modells feinabgestimmt. Dies führt zu einem deutlich kleineren Modell sowie zu kleineren Vektorraummodellen (Embeddings), allerdings auf Kosten einer gewissen Feinabstimmung.

Destillation ist der Prozess des Trainierens eines neuen Modells, um die Leistung eines bestehenden Modells zu erreichen. Dies bedeutet, dass ein neues Modell erstellt wird, das von Grund auf als quantisiert konzipiert ist, und dann das bestehende Modell verwendet wird, um so viele Trainingsdaten wie nötig zu generieren, um es so lange zu trainieren, bis es so gut wie möglich mit dem bestehenden Modell übereinstimmt.

Die Vorteile dieser vier Ansätze sind in der folgenden Tabelle zusammengefasst:

Alle vier erzeugen kompaktere Vektorraummodelle (Embeddings), aber außer PTQ erfordern alle ein zusätzliches Training, während nur Full QAT und Destillation neue, schnellere Modelle erzeugen. Full QAT und Destillation sind in der Umsetzung viel aufwendiger, da sie ein wesentlich umfangreicheres Training erfordern als Output QAT.

In diesem Artikel werden wir uns nur mit PTQ und Output QAT befassen, die die Größe oder Geschwindigkeit des Vektorraummodell (Embedding)-Modells nicht verändern.

tagExperimenteller Aufbau

Für diese Experimente ist unser Basismodell jina-embeddings-v4 mit dem Retrieval-Adapter, der 32-Bit-Gleitkommavektoren (FP32) in 2048 Dimensionen erzeugt. Jedes Vektorraummodell (Embedding) ist daher 8196 Byte oder 8 KB groß.

Wir haben verschiedene experimentelle Bedingungen anhand von Abfrage-Dokumenten-Retrieval-Benchmark-Aufgaben aus der NanoBEIR-Benchmark-Suite untersucht. Der Retrieval-Prozess verwendet die Kosinusähnlichkeit zwischen Vektoren, um die Dokumente zu finden und zu ordnen, die am besten zu den Abfragen passen.

• Baseline – Die Leistung von jina-embeddings-v4 Vektorraummodellen (Embedding)-Vektoren ohne Quantisierung. Alle diese Experimente verwendeten eine Beta-Version des Modells, und die Release-Leistung ist etwas besser.
• PTQ – Wir haben die Ausgabvektoren zu binären Vektoren quantisiert, ohne das Modell zu verändern.
• Output QAT – Wir haben die Ausgabvektoren quantisiert und den Retrieval-Adapter feinabgestimmt, um seine Leistung unter quantisierten Bedingungen zu verbessern.

tagQuantisierungsstufen

Wir haben mit vier verschiedenen Quantisierungsstufen experimentiert.

• 8-Bit-Ganzzahlen – FP32-Werte werden auf ganze Zahlen im Bereich von -128 bis 127 reduziert, wodurch sich die Vektorraummodelle (Embeddings) um das Vierfache auf 2048 Byte verkleinern.
• 4-Bit-Ganzzahlen – Gleiches gilt wie für 4-Bit-Ganzzahlen, aber wir ordnen dem Bereich von -8 bis 7 zu, wodurch die Vektorgrößen um den Faktor 8 auf 1024 Byte reduziert werden.
• Trinäre Quantisierung – Alle Werte werden einem von drei Werten zugeordnet: -1, 0, 1. Optimal gespeichert, reduziert dies jede Dimension auf 1,6 Bit, wodurch die Größe der Vektorraummodelle (Embedding)-Vektoren um etwa das 40-fache auf etwa 230 Byte reduziert wird.
• Binäre Quantisierung – Wir konvertieren FP32-Skalarwerte in ein Bit, indem wir den Datentyp torch.sign verwenden, der nur zwei Werte bereitstellt, für deren Speicherung ein Bit benötigt wird. Dies reduziert 2048-dimensionale Vektorraummodelle (Embedding)-Vektoren von 8192 Byte auf 128 Byte, eine 64-fache Reduzierung.

tagSkalierung

Bei der binären Quantisierung ist die Quantisierung sehr einfach: Wenn ein Vektorwert über 0 oder positiv ist, wird er 1 zugeordnet. Andernfalls wird er -1 zugeordnet.

Für die anderen Quantisierungsszenarien haben wir die Werte auf einen Bereich normalisiert und dann auf den nächsten Wert gerundet, der durch die Quantisierungsstufe zugelassen wurde. Vektorraummodelle (Embedding)-Vektoren bestehen aus Skalenzahlen zwischen -∞ und +∞ (oder in der Praxis sehr großen positiven und negativen Zahlen). Wir verwenden zwei Zahlen, $max$ und $min$, um die Werte für die Quantisierung zu skalieren.

Für die trinäre Quantisierung nehmen wir jede Vektorkomponente $v$ und übersetzen sie wie folgt:

• wenn $v$ ≥ $max$, wird $v$ zu 1.
• wenn $v$ ≤ $min$, wird $v$ zu -1.
• wenn $min$ < $v$ < $max$, wird $v$ zu 0.

Für 4-Bit-Ganzzahlen:

• wenn $v$ ≥ $max$, wird $v$ zu 7.
• wenn $v$ ≤ $min$, wird $v$ zu -8.
• wenn $min$ < $v$ < $max$, wird $v$ zu $16*(v - min)/(max - min) - 8$, dann auf die nächste ganze Zahl gerundet. Dies skaliert den Wert in den Bereich $[-8,7]$.

Für 8-Bit-Ganzzahlen:

• wenn $v$ ≥ $max$, wird $v$ zu 127.
• wenn $v$ ≤ $min$, wird $v$ zu -128.
• wenn $min$ < $v$ < $max$, wird $v$ zu $256*(v - min)/(max - min) - 128$, gerundet auf die nächste ganze Zahl. Dies skaliert den Wert in den Bereich $[-128,127]$.

Zur Berechnung von $max$ und $min$ haben wir zwei Ansätze verwendet:

• Min/Max – Wir haben unsere Daten in Batches verarbeitet und für jeden Batch die höchste und niedrigste Vektorkomponente ermittelt, wobei wir $max$ auf die höchste und $min$ auf die niedrigste gesetzt haben.
• Rollierende Mittelwertbildung über Batches – Für jeden Batch haben wir den Durchschnitt und die Standardabweichung der Vektorkomponenten berechnet. Wir haben einen gleitenden Durchschnitt sowohl des Durchschnitts als auch der Standardabweichung beibehalten, während wir alle Batches verarbeitet haben. Wenn $avg$ der aktuelle gleitende Durchschnitt der durchschnittlichen Batch-Werte ist und $std$ der aktuelle gleitende Durchschnitt der Standardabweichungen ist, dann gilt für jeden Batch:

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

tagQAT-Feinabstimmung

Für die PTQ-Experimente haben wir das Modell so verwendet, wie es ist, und die von ihm erzeugten Vektorraummodelle (Embeddings) mit den oben beschriebenen Methoden quantisiert.

Für die Output QAT haben wir das Modell mit Straight-Through Estimation feinabgestimmt. Das bedeutet, dass wir den Quantisierungsprozess umkehren und die volle Präzision der Werte wiederherstellen, bevor wir den Verlust (d. h. den Fehler) berechnen, und dann diese Verlustmetrik verwenden, um das Modell feinabzustimmen.

Wir haben in jedem Fall für 10.000 Schritte feinjustiert und alle 500 Schritte einen Checkpoint gespeichert. Wir haben dann den Checkpoint mit der höchsten Punktzahl auf dem NanoBEIR-Benchmark beibehalten.

tagAsymmetrische Quantisierung

PTQ und Output QAT reduzieren die Größe der Vektorraummodelle (Embeddings), reduzieren aber nicht die Modellgröße oder die Inferenzgeschwindigkeit; alle Einsparungen liegen in der Größe der gespeicherten Dokumenten-Vektorraummodelle (Embeddings) und der Abrufgeschwindigkeit.

Infolgedessen haben wir sowohl die Quantisierung der Abfragevektoren getestet als auch die unquantisierten Vektoren zum Zeitpunkt des Abrufs beibehalten, da dies die Größe der gespeicherten Vektorraummodelle (Embeddings) in beiden Fällen nicht verändert.

tagErgebnisse

Wir haben insgesamt neun Bedingungen getestet, die in den folgenden Tabellen zusammengefasst sind:

Tabelle 2: Versuchsbedingungen

Tabelle 3: Durchschnittliche Punktzahl (in % richtig) für jede Bedingung über die zwölf NanoBEIR-Benchmarks.

Wie Sie der obigen Tabelle entnehmen können, verbessert die Feinabstimmung für die Quantisierung die Punktzahlen. Der einzige Unterschied zwischen den Bedingungen PTQ Binär und QAT Binär ist die Feinabstimmung, und der Unterschied in der Punktzahl ist signifikant. In ähnlicher Weise sehen wir eine Verbesserung der Punktzahlen um fast 2 % zwischen den Bedingungen PTQ Binär Nur Dokumente und QAT Binär Nur Dokumente, die sich nur durch die gleiche Feinabstimmung unterscheiden.

Es überrascht nicht, dass wir auch feststellen, dass sich die Punktzahlen im Allgemeinen verbessern, je weniger wir quantisieren, wobei die 4-Bit-Quantisierung besser abschneidet als die trinäre und die trinäre besser als die binäre. Es scheint jedoch, dass eine weitere Umstellung auf 8-Bit nichts verbessert hat.

Wir haben das Weglassen der Quantisierung von Abfragen nur in binären Fällen getestet, aber dies scheint die Leistung zu verbessern.

Schließlich legen unsere Tests nahe, dass die rollierende Durchschnittsskalierungsmethode die einfache Min/Max-Methode übertrifft.

tagFazit

Die Quantisierung hat einige wichtige operative Vorteile für Vektorraummodelle (Embedding), da sie die Größe der Vektorraummodelle (Embedding) erheblich reduziert und die Informationsbeschaffung beschleunigt. Während die einfache Post-Training-Quantisierung (PTQ) unmittelbare Vorteile in Bezug auf Speicher und Lagerung bietet, zeigen unsere Experimente, dass das quantisierungsbewusste Training (QAT) die unvermeidlichen Präzisionsverluste erheblich reduziert. Die Feinabstimmung führte durchweg zu besseren Ergebnissen.

Der Grad der Quantisierung wirkt sich direkt auf die Leistung aus, was man von einer Methode erwarten würde, die auf der Reduzierung der Präzision von Werten basiert. Weniger aggressive Quantisierung (z. B. 4-Bit) übertrifft im Allgemeinen aggressivere Methoden (z. B. binär), aber überraschenderweise gab es keinen signifikanten Unterschied in der Leistung zwischen 8-Bit- und 4-Bit-Quantisierung. Es scheint, dass es bis zum Erreichen einer gewissen Schwelle der Ungenauigkeit kaum einen Unterschied zwischen größerer und geringerer Quantisierung gibt.

Auch Skalierungsstrategien sind von Bedeutung, wobei die rollierende Durchschnittsmethode im Vergleich zu einem festen Min/Max-Ansatz überlegene Ergebnisse zeigt. Die Verwendung von Skalierungswerten, die relativ zu den Daten sind, scheint deutlich besser zu funktionieren und verdient weitere Untersuchungen.

Die Quantisierung kann Ihnen mehr aus Ihren Vektorraummodellen (Embedding) für weniger herausholen. Obwohl dieser Artikel nicht alle Optionen für die Quantisierung untersucht, werden zwei Optionen untersucht, die leicht zugänglich sind und echte Vorteile bieten. Wir arbeiten daran, die Quantisierungsstrategien zu verfeinern und zu verbessern, damit wir die Kosten der Benutzer weiter senken können, und erwarten, in naher Zukunft binäre Unterstützung für jina-embeddings-v4 zu veröffentlichen.