Matroschka-Repräsentationslernen erklärt: Die Methode hinter OpenAIs effizienten Text-Embeddings
Die Entwicklung realer Machine-Learning-Modelle beinhaltet immer einen Kompromiss zwischen Kosten und Leistung. Zum Beispiel gilt: Je größer das Modell und der Datensatz sind, die wir während des Trainings verwenden, desto leistungsfähiger wird unser trainiertes Modell letztendlich sein. Diese erhöhte Leistungsfähigkeit geht jedoch mit längerer Zeit und höheren Rechenkosten einher, da das Training deutlich längere Zeiträume erfordert. Dasselbe Prinzip gilt während der Inferenz, bei der größere Modelle tendenziell größere Feature-Repräsentationen erzeugen, die mehr Speicherplatz für die Speicherung benötigen.
Da die Anforderungen während des Machine-Learning-Trainings und der Inferenz von einem Anwendungsfall zum anderen erheblich variieren können, ist es entscheidend, eine Methode zu haben, die es uns ermöglicht, einen kleinen Teil der Leistung eines Modells gegen reduzierte Kosten einzutauschen. Hier kommt eine Methode wie Matryoshka Representation Learning (MRL) ins Spiel. Zum Beispiel nutzt OpenAIs text-embedding-3-small Modell MRL, um Entwicklern zu ermöglichen, Embeddings zu verkürzen und gleichzeitig ihre grundlegenden konzeptrepräsentierenden Eigenschaften beizubehalten. Indem MRL die Anpassung von Embedding-Dimensionen ermöglicht, hilft es, die perfekte Balance zwischen Kosteneffizienz und Modellleistung zu finden.
Dieser Artikel untersucht, wie MRL funktioniert, wie es implementiert wird und wie es skalierbare und effiziente Machine-Learning-Modelle ermöglicht. Beginnen wir mit der Motivation hinter MRL.
Die Motivation hinter Matryoshka Representation Learning (MRL)
Der Kosten-Leistungs-Kompromiss ist immer eine Überlegung, bevor Machine-Learning-Modelle entwickelt oder verwendet werden.
Während des Trainings beeinflussen die Parameter eines Modells und die Menge der Trainingsdaten direkt seine endgültige Leistung. Größere Modelle und Trainingsdatensätze führen tendenziell zu leistungsfähigeren Modellen. Allerdings sind auch die Rechenkosten für das Training dieser Modelle höher.
Während der Inferenz hingegen erzeugen größere Modelle größere Feature-Repräsentationen. Diese Zunahme an Feature-Repräsentationen bedeutet, dass wir mehr Speicher benötigen, um sie zu speichern.
Verwenden wir BERT als Beispiel im Kontext des Information Retrieval. BERT ist eines der frühesten auf Transformer basierenden Deep-Learning-Modelle, das auf mehreren Benchmark-Datensätzen in Textklassifikation, Named Entity Recognition (NER), Fragebeantwortung (QnA) und mehr State-of-the-Art-Leistung erzielte.
Abbildung: Workflow der Embedding-Erzeugung mit dem BERT-Basismodell.
Kurz gesagt nimmt BERT einen Satz oder ein Wort als Eingabe und transformiert ihn bzw. es in ein Embedding fester Größe. Dieses Embedding erfasst die semantische Bedeutung der ursprünglichen Eingabe, und seine Größe hängt von der spezifischen verwendeten Modellvariante ab.
BERT hat zwei Hauptvarianten: das Basismodell und das große Modell. Das Basismodell erzeugt ein 768-dimensionales Embedding, während das große Modell ein 1024-dimensionales Embedding generiert.
Nehmen wir an, wir möchten 10 Millionen Embeddings in einer Vektordatenbank für einen Information-Retrieval-Anwendungsfall speichern. Wenn wir das BERT-Basismodell im FP32-Format verwenden, benötigen wir ungefähr 768×10M×4=30,72 GB Speicher, um diese Embeddings zu speichern. Im Gegensatz dazu würde die Verwendung des großen BERT-Modells im selben Szenario etwa 40,9 GB erfordern. Zusätzlich kann eine größere Embedding-Größe den Retrieval-Prozess verlangsamen, indem sie die Rechenkomplexität von Ähnlichkeitssuchen erhöht. Der Hauptvorteil der Verwendung größerer Embeddings besteht jedoch in ihrer verbesserten Fähigkeit, im Vergleich zu kleineren Embeddings hochrelevante Informationen abzurufen.
Eine ideale Lösung für diese Probleme im Information Retrieval besteht darin, den gesamten Prozess in zwei Teile aufzuteilen: Shortlisting und Reranking. Shortlisting bezieht sich darauf, eine erste Menge von Kandidatendokumenten aus einer riesigen Sammlung in unserer Vektordatenbank wie Milvus abzurufen. Reranking nimmt dann diese in die engere Auswahl genommenen Kandidaten und ordnet sie neu an, um die Relevanz der endgültigen Ergebnisse zu maximieren.
Abbildung- Shortlisting- und Reranking-Workflow
Abbildung: Shortlisting- und Reranking-Workflow.
Um das Shortlisting durchzuführen, berechnet der Algorithmus die Ähnlichkeit zwischen Query-Embeddings und einer großen Sammlung von Embeddings, die in der Vektordatenbank gespeichert sind. Daher kann die Verwendung kleinerer Embedding-Größen diese Berechnung effizienter und schneller machen. Umgekehrt legt Reranking den Schwerpunkt auf Präzision, was bedeutet, dass wir von größeren Embedding-Größen profitieren würden.
Das Problem bei dieser Methode ist, dass, sobald wir ein bestimmtes Modell auswählen, die Größe der Embeddings oder Feature-Repräsentationen, die es erzeugt, festgelegt ist. Dieser Mangel an Flexibilität schränkt unsere Fähigkeit ein, Embeddings unterschiedlicher Größen aus einem einzigen Modell zu verwenden.
Abbildung- Workflow eines Modells, das ein Embedding fester Größe erzeugt..png
Abbildung: Workflow eines Modells, das ein Embedding fester Größe erzeugt.
Der Ansatz des Matryoshka Representation Learning (MRL) bietet eine interessante Lösung für dieses Problem: Modelle, die mit dieser Methode trainiert wurden, können Embeddings verschiedener Größen erzeugen. Wir werden den Mechanismus dieses Ansatzes im nächsten Abschnitt ausführlich untersuchen.
Was ist Matryoshka Representation Learning (MRL)?
Matryoshka Representation Learning (MRL) ist eine Methode zum Trainieren neuronaler Netzwerke, um mehrskalige Repräsentationen innerhalb eines einzigen Modells zu erzeugen. Inspiriert von Matroschka-Puppen, bei denen kleinere Puppen in größere passen, ermöglicht MRL einem Modell, Repräsentationen unterschiedlicher Größen (von grob bis fein) mit einem einzigen Vorwärtsdurchlauf auszugeben. Dieser Ansatz ermöglicht es Deep-Learning-Modellen, komplexe Beziehungen und Nuancen in Daten besser zu verstehen. MRL ist besonders effektiv für Aufgaben wie semantische Suche, Information Retrieval, mehrsprachige Verarbeitung und jede Anwendung, die nuancierte Repräsentationen von Daten über verschiedene Abstraktionsebenen hinweg erfordert.
Visualisierung von Matryoshka-Embeddings mit mehreren Detailebenen
Abbildung: Visualisierung von Matryoshka-Embeddings mit mehreren Detailebenen
Beliebte Embedding-Modelle, die den Ansatz des Matryoshka Representation Learning (MRL) übernommen haben, umfassen OpenAIs text-embedding-3-large, Nomics nomic-embed-text-v1 und Alibabas gte-multilingual-base.
Wie der MRL-Ansatz funktioniert
Der MRL-Ansatz ermöglicht es uns, Feature-Repräsentationen aus jedem Machine-Learning-Modell in mehreren Größen zu extrahieren. Anstatt beispielsweise die ursprünglichen 1024 Dimensionen eines Embeddings zu verwenden, können wir die ersten 16, 32, 64, 128 oder 256 Dimensionen (oder beliebige von uns gewählte Dimensionen) nutzen. Der Schlüssel zu dieser Fähigkeit liegt in der Art und Weise, wie das Modell beim Implementieren des MRL-Ansatzes trainiert wird.
Abbildung- Training der MRL-Verlustfunktion und ihr Anwendungsfall während der Inferenz
Abbildung: Training der MRL-Verlustfunktion und ihr Anwendungsfall während der Inferenz. Quelle.
Während des Modelltrainings mit MRL optimieren wir nicht nur eine Verlustfunktion, wie es beim standardmäßigen Modelltraining üblich ist, sondern mehrere. Jede Verlustfunktion zielt darauf ab, die Merkmalsrepräsentation für eine bestimmte Dimension zu optimieren, die wir im Voraus definieren. Wenn wir fünf verschiedene Merkmalsdimensionen festlegen, müssen wir während des Trainings fünf unterschiedliche Verlustfunktionen optimieren.
Nehmen wir an, wir möchten ein BERT-large-Modell mit dem MRL-Ansatz trainieren. Wie im vorherigen Abschnitt erwähnt, erzeugt dieses Modell ein Embedding der Größe 1024. Wenn wir die Merkmalsdimensionen auf 32, 64, 128, 256 und 1024 festlegen, optimieren wir während des Trainings fünf verschiedene Verlustfunktionen, sodass das Modell lernt, optimierte Embeddings in diesen Dimensionen zu erzeugen.
Die Optimierung über diese verschiedenen Dimensionen hinweg ist unkompliziert: MRL zerlegt die gesamte Verlustfunktion in die Summe der Verluste für jede einzelne Dimension. Anhand unseres obigen Beispiels lässt sich dies wie folgt ausdrücken:
Der Optimierungsprozess für ein mit MRL trainiertes Modell lässt sich formaler wie folgt darstellen:
Abbildung: Optimierungsgleichung des MRL-Ansatzes. Quelle.
Wie Sie sehen, ist die Definition der obigen Verlustfunktion recht allgemein. Daher kann der MRL-Ansatz auf nahezu jedes Modell angewendet werden und ist völlig unabhängig von der Architektur des Modells. Wir können auch vortrainierte Modelle wie BERT oder beliebige andere transformerbasierte Modelle feinabstimmen, um Embeddings verschiedener Größen auszugeben. Wenn wir unser obiges Szenario fortsetzen, können wir nun 1024-dimensionale Embeddings aus unserem trainierten BERT-large-Modell sowie 32-, 64-, 128- und 256-dimensionale Embeddings erzeugen.
Insgesamt tragen nach dem Training eines Modells mit MRL die anfänglichen Dimensionen der vom trainierten Modell erzeugten Merkmale bedeutendere Informationen als spätere Dimensionen. Die ersten wenigen Dimensionen enthalten Details auf hoher Ebene, während die späteren Dimensionen sich auf granularere Informationen konzentrieren und damit nachahmen, wie eine Matrjoschka-Puppe aufgebaut ist.
Dies bedeutet jedoch nicht, dass kürzere Embeddings immer einfach die gekürzten Versionen längerer Embeddings sind. Die Werte der einzelnen Elemente in kürzeren Embeddings können sich von denen in längeren Embeddings unterscheiden, abhängig vom Skalierungsfaktor, der während des Trainings auf jede Merkmalsdimension angewendet wird. Wenn wir jedoch den Skalierungsfaktor in jeder Dimension gleich festlegen, können die Werte der einzelnen Elemente zwischen kürzeren und längeren Embeddings identisch aussehen.
Experimentelle Ergebnisse von MRL
Der MRL-Ansatz wurde an Machine-Learning-Modellen über verschiedene Modalitäten hinweg evaluiert, darunter Text, Bild und Bild-Text. ResNet50- und ViT-Modelle repräsentieren bildbasierte Modelle, BERT repräsentiert textbasierte Modelle, und ALIGN repräsentiert die Kombination aus Bild und Text. Diese Modelle wurden hauptsächlich für zwei häufige Anwendungsfälle bewertet: Klassifikation und Retrieval. Lassen Sie uns zuerst die Klassifikation besprechen.
Beim Vergleich der Leistung des auf dem ImageNet-1K-Datensatz trainierten ResNet50-Modells mit einem unabhängig trainierten Standard-ResNet50 erreicht das MRL-Modell eine vergleichbare Top-1-Genauigkeit bei verschiedenen Größen der Merkmalsrepräsentation.
Um den Nutzen der Merkmalsrepräsentationen für nachgelagerte Aufgaben weiter zu bewerten, wurde auch die Genauigkeit des 1-Nearest Neighbor (1-NN) für jede Größe der Merkmalsrepräsentation gemessen. Der Aufbau zur Berechnung von 1-NN ist wie folgt: Bei 1,3K Bildbeispielen in der Datenbank besteht die Aufgabe darin, für jede der 50K Abfragen den nächsten Nachbarn zu finden. Mit diesem Aufbau ist das mit MRL trainierte ResNet50 bei jeder Größe der Merkmalsrepräsentation bis zu 2 % genauer als sein Pendant mit festen Merkmalen.
Abbildung- Top-1-Genauigkeit und 1-NN-Genauigkeit von ResNet50 auf ImageNet-1K
Abbildung: Top-1-Genauigkeit und 1-NN-Genauigkeit von ResNet50 auf ImageNet-1K. Quelle.
In der Zwischenzeit ist auch die Leistung des mit MRL auf dem JFT-300M-Datensatz trainierten ViT-Modells über alle Repräsentationsgrößen hinweg sehr wettbewerbsfähig. Seine 1-NN-Genauigkeit ist mit der von ViT vergleichbar, das mit Merkmalsrepräsentationen fester Größe trainiert wurde. Wie in der Abbildung unten gezeigt, ist die Leistung des MRL-Modells bei niedrigeren Merkmalsrepräsentationen ebenfalls besser als die seines Pendants mit fester Größe, teilweise weil zufällige Merkmale aus dem Modell fester Größe ausgewählt werden, um niedrigere Dimensionen darzustellen. Ein ähnlicher Trend ist beim ALIGN-Modell zu beobachten, das mit dem MRL-Ansatz trainiert wurde; seine Leistung entspricht der des ALIGN-Modells, das mit Repräsentationen fester Größe trainiert wurde.
Abbildung- 1-NN-Genauigkeit von ViT und ALIGN auf ImageNet-1K
Abbildung: 1-NN-Genauigkeit von ViT und ALIGN auf ImageNet-1K. Quelle.
Der Hauptvorteil der Verwendung des MRL-Ansatzes ist seine Flexibilität, zwischen verschiedenen Größen von Merkmalsrepräsentationen desselben Modells zu wechseln. Um diese Stärke voll auszuschöpfen, wurden auch Tests zur adaptiven Klassifikation mit ResNet50 auf ImageNet-1K durchgeführt.
Dieser Aufbau beinhaltet das Lernen der maximalen Softmax-Wahrscheinlichkeit, um Übergänge von kleineren Merkmalsrepräsentationen zu größeren zu bestimmen. Die Testergebnisse zeigen, dass das MRL-Modell mit einer 37-dimensionalen Merkmalsrepräsentation genauso genau ist wie das ResNet50-Modell, das mit einer festen 512-dimensionalen Merkmalsrepräsentation trainiert wurde, während es nur 0,8 % weniger genau ist als das 2048-dimensionale Modell.
Abbildung- Top-1-Genauigkeit des ResNet50-Modells in der adaptiven Klassifikation (links) und mAP von ResNet50 im Anwendungsfall der Bildsuche (rechts)
Abbildung: Top-1-Genauigkeit des ResNet50-Modells in der adaptiven Klassifikation (links) und mAP von ResNet50 im Anwendungsfall der Bildsuche (rechts). Quelle.
Als Nächstes wollen wir die Suche besprechen. Wie Sie vielleicht wissen, besteht das Ziel der Suche darin, für eine gegebene Abfrage eine Menge ähnlicher Kontexte zu finden. Daher wurde die Suchqualität zwischen ResNet50, das mit MRL trainiert wurde, und Merkmalsrepräsentationen fester Größe unter Verwendung der mean Average Precision (mAP) getestet. Das MRL-Modell zeigt bei jeder Größe der Merkmalsrepräsentation eine Verbesserung der mAP um bis zu 3 % im Vergleich zu seinem Pendant mit fester Größe, wie in der obigen Visualisierung dargestellt.
Da wir mit MRL mehrere Größen von Merkmalsrepräsentationen nutzen können, ist es außerdem besonders interessant zu untersuchen, wie wir den Suchprozess beschleunigen und gleichzeitig mit MRL die Genauigkeit beibehalten können. Wir können dies in einem Szenario testen, das als adaptive Suche bezeichnet wird.
Wie bereits erwähnt, gibt es in einem Suchprozess zwei Hauptphasen: Vorauswahl und Neuranking. Der Aufbau der adaptiven Suche umfasste die Verwendung einer 16-dimensionalen Merkmalsrepräsentation für die Vorauswahl einer Menge von 200 Kandidaten, während eine 2048-dimensionale Merkmalsrepräsentation für das Neuranking verwendet wird. Dieser Ansatz wird mit einer anderen Methode verglichen, die eine 2048-dimensionale Merkmalsrepräsentation sowohl für die Vorauswahl als auch für das Neuranking verwendet.
Bei ImageNet-1K kann die Verwendung des MRL-Modells den Retrieval-Prozess theoretisch um bis zu das 128-Fache beschleunigen. Wir bezeichnen dies als theoretisch, da wir in realen Anwendungen selten einen naiven Nächste-Nachbarn-Algorithmus für Retrieval-Aufgaben verwenden. Stattdessen werden üblicherweise Approximate-Nearest-Neighbor-(ANN)-Algorithmen wie FAISS, ANNOY oder HNSW implementiert, um die Retrieval-Geschwindigkeit zu erhöhen. In diesem Fall erreicht das Setup mit dem MRL-Modell eine 14-fache Beschleunigung im Vergleich zu Retrieval-Prozessen mit dem HNSW-Algorithmus auf identischer Hardware.
Ein ähnliches Ergebnis wurde mit dem ImageNet-4K-Datensatz beobachtet, bei dem die Autoren des Papers ein MRL-Modell mit einer 64-dimensionalen Merkmalsrepräsentation für das Shortlisting und einer 2048-dimensionalen Repräsentation für das Reranking verwendeten. Die Ergebnisse zeigen eine theoretische Beschleunigung um das 32-Fache und eine reale Beschleunigung um das Sechsfache.
Der Trade-off zwischen mAP@10 und MFLOPs/Query für Adaptive Retrieval (AR) auf ImageNet-1K (links) und ImageNet-4K (rechts).
Der Trade-off zwischen mAP@10 und MFLOPs/Query für Adaptive Retrieval (AR) auf ImageNet-1K (links) und ImageNet-4K (rechts). Quelle.
Implementierung des MRL-Ansatzes
In diesem Abschnitt werfen wir einen Blick auf die einfache Implementierung des MRL-Ansatzes. Insbesondere verwenden wir ein MPNet-Basismodell, das mit dem MRL-Ansatz auf dem NLI-Datensatz trainiert wurde, um Embeddings in verschiedenen Größen zu erzeugen. Wie Sie vielleicht bereits wissen, kann das ursprüngliche MPNet-Basismodell nur ein 768-dimensionales Embedding erzeugen. Dieses trainierte MPNet-Modell kann jedoch Embeddings in 768, 512, 256, 128 und 64 Dimensionen erzeugen, wobei das Gewicht jeder Dimension während des Trainings auf 1 gesetzt ist.
Wir können dieses Modell mithilfe der SentenceTransformers-Bibliothek mit dem folgenden Code verwenden.
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]
"""
Im obigen Code verwenden wir das trainierte MPNet-Basismodell, um Embeddings in zwei verschiedenen Größen zu erzeugen: 64 und 768. Da das Gewicht oder der Skalierungsfaktor jeder Dimension auf 1 gesetzt ist, haben die beiden Embeddings identische Elemente, wie Sie in den ersten 10 Elementen der beiden obigen Embeddings sehen können.
Da der Wert jedes Elements identisch ist, wäre die Ähnlichkeit zwischen den beiden Embeddings 1.
from sentence_transformers.util import cos_sim
similarities = cos_sim(short_embedding[0], long_embedding[0][:matryoshka_dim_short])
print(similarities)
# tensor([[1.]])
Wenn Sie Ihr eigenes Modell mit dem MRL-Ansatz trainieren möchten, können Sie sich das offizielle GitHub-Repo dieser Methode ansehen.
Fazit
Der MRL-Ansatz führt eine Lösung ein, um den Kosten-Leistungs-Kompromiss im maschinellen Lernen auszugleichen. Indem jedes maschinelle Lernmodell in die Lage versetzt wird, Merkmalsrepräsentationen unterschiedlicher Größen zu erzeugen, bietet MRL die Flexibilität, je nach Anwendungsfall und Ressourcen entweder auf Geschwindigkeit oder Genauigkeit zu optimieren. Wir haben gesehen, dass diese Anpassungsfähigkeit in Anwendungen wie Klassifikation und Retrieval wertvoll ist, bei denen der Wechsel zwischen kleineren und größeren Repräsentationen die Effizienz erheblich steigern kann, ohne die Leistung massiv zu beeinträchtigen.
Die experimentellen Ergebnisse zeigen, dass MRL die Genauigkeit traditioneller Modelle mit fester Größe über mehrere Bereiche hinweg erreicht und oft übertrifft, einschließlich Text-, Bild- und multimodaler Aufgaben. Außerdem unterstreicht die Kompatibilität des Ansatzes mit bestehenden Architekturen wie ResNet50, ViT und BERT seine Anwendbarkeit über verschiedene Modellmodalitäten hinweg. Mit dem Potenzial für erhebliche Beschleunigungen in realen Szenarien und verbesserter Retrieval-Qualität stellt MRL einen vielversprechenden Fortschritt für effizientere und vielseitigere Lösungen im maschinellen Lernen dar.
Verwandte Ressourcen
Weiterlesen
We spent 8 years making vector databases faster. Then we stopped.
Rarely queried embeddings still need to stay searchable. See how Vector Lakebase enables on-demand vector search without always-on compute costs.
Zilliz Cloud Update: Smarter Autoscaling for Cost Savings, Stronger Compliance with Audit Logs, and More
What's new in Zilliz Cloud? Smarter autoscaling with scale-down, audit logs GA, enhanced SSO, and Milvus 2.6 in Private Preview.
1 Table = 1000 Words? Foundation Models for Tabular Data
TableGPT2 automates tabular data insights, overcoming schema variability, while Milvus accelerates vector search for efficient, scalable decision-making.