Millionen Wikipedia-Seiten offline per „Vibe“ durchsuchen
(leebutterman.com)- Implementiert Offline-Echtzeit-Embedding-Suche ohne Server direkt im Browser und ermöglicht so die interaktive lokale Exploration von 6 Millionen Artikeln der englischen Wikipedia
- Die ursprünglichen 384-dimensionalen float32-Embeddings sind mit rund 9 GB zu groß für eine Browser-App; deshalb wird die Größe auf Basis von all-minilm-l6-v2 und gemittelten Embeddings von Seiten-Chunks reduziert
- Durch Anwendung von Product Quantization in 8-dimensionalen Unterräumen wird der Datensatz auf etwa 288 MB verkleinert; 96 MB können rund 2 Millionen Embeddings enthalten
- Die Suche berechnet Distanzen ohne Dekomprimierung in der komprimierten Domäne; durch Korrektur des Konvertierungsergebnisses mit dem ONNX-Opcode GatherElements wird die Berechnung etwa 4-mal schneller
- Auf modernen Geräten sind MiniLM-Inferenz und Distanzberechnung für 100k Embeddings in kurzer Zeit abgeschlossen, sodass die React-UI Ergebnisse im Takt von 100–300 ms aktualisieren kann
Offline-Vektorsuche direkt im Browser
- Ziel ist zu zeigen, dass Millionen Wikipedia-Seiten auch ohne komplexe serverseitige Suchinfrastruktur in Echtzeit im Browser durchsucht werden können
- Das Kerndesign lässt sich auf drei Punkte verdichten
- Vektorsuche kann allein mit Product Quantization und linearem Scan schnell genug sein
- ONNX Runtime wird genutzt, Ineffizienzen aus der PyTorch-Konvertierung können aber auf Opcode-Ebene behoben werden
- Der Browser ist schnell genug, um mit WASM Echtzeit-Inferenz auszuführen; WebGPU bleibt zudem eine Möglichkeit für die Zukunft
- Die Suchbedingungen kombinieren Vektorähnlichkeit der Text-Embeddings mit einer Gleichheitsbedingung für den ersten Buchstaben des Titels
- Das ähnelt einer Datenbankabfrage, bei der einige Spalten über Vektorähnlichkeit und andere über Gleichheitsbedingungen durchsucht werden
- Wegen der Einschränkung als Offline-Browser-App müssen die Embeddings klein sein, das Embedding-Modell muss Open Source sein, und auch Modellgröße sowie Rechenaufwand müssen gering bleiben
- Stand 2023 bedeutet ein leichtgewichtiges Modell üblicherweise unter 100 MB
Daten der englischen Wikipedia und Embedding-Größe
- Der Datensatz ist die englische Wikipedia und wird mit dem Modell all-minilm-l6-v2 eingebettet
- Es gibt etwa 6 Millionen Dokumente; Seiten werden in Chunks aufgeteilt, anschließend werden die Embeddings der einzelnen Chunks gemittelt, um ein Seiten-Embedding zu erzeugen
- Seiten werden zuerst nach Länge angeordnet
- Beim schrittweisen Laden der Datenbank sollen lange Seiten zuerst angezeigt werden, weil in sie vermutlich mehr menschlicher Aufwand geflossen ist
- Die Embedding-Dimension beträgt 384 Dimensionen
- Als float32 gespeichert ergibt das
6M * 384 * 4 = 9GB - In 96 MB passen nur etwa 64k Embeddings
- Als float32 gespeichert ergibt das
- Die ursprüngliche float32-Kodierung ist zu groß, um zu einer browserbasierten Offline-Suche zu passen
Embeddings mit Product Quantization komprimieren
- Product Quantization speichert Gleitkommawerte nicht direkt, sondern erstellt mehrere Paletten und ersetzt Eingabewerte durch Palettenindizes
- Paletten enthalten typischerweise höchstens 256 Werte, sodass jeder Index maximal 1 Byte belegen kann
- Implizite Paletten verwenden eine verschobene und skalierte Zahlengerade wie -127 bis 127 oder 0 bis 255
- Explizite Paletten speichern 256 Zahlen der Reihe nach
- Product Quantization verwendet explizite Palettisierung
- In der Minimalvariante, bei der jede der 384 Dimensionen durch einen 1-Byte-Index ersetzt wird, ergibt sich ein Datensatz von
6M * 384 = 2.25GB- In 96 MB lassen sich etwa 256k Embeddings speichern
- Auch dieser Ansatz ist für das Ziel noch ineffizient
- Quantisiert man in Einheiten von 2-dimensionalen Punkten, kann ein 384-dimensionales Embedding durch 192 Indizes dargestellt werden, was eine 2-fache Platzersparnis ermöglicht
- Die tatsächliche Konfiguration nutzt Einheiten von 8-dimensionalen Punkten
- 384 Dimensionen werden durch 48 Indizes dargestellt
- Die Palettengröße bleibt bei etwa 384 KB
- Die Datensatzgröße beträgt
6M * 48 = 288M - In 96 MB lassen sich etwa 2 Millionen Embeddings speichern
- Für jede Quantisierungsstufe lassen sich Precision und Recall auswerten
Direkt suchen, ohne zu dekomprimieren
- Ein Vorteil von Product Quantization ist, dass sich im komprimierten Zustand suchen lässt, ohne zu dekomprimieren
- Wenn n Dimensionen zu n/2 zweidimensionalen Punkten gruppiert werden, kann die Distanz zwischen jedem Palettenpunkt und dem entsprechenden zweidimensionalen Punkt der Query vorab berechnet werden
- Anschließend schlägt jedes Embedding Distanzwerte über Palettenindizes nach und summiert sie zur Gesamtdistanz
- Das ist weniger Arbeit, als zunächst den gesamten n-dimensionalen Punkt zu rekonstruieren und dann die Distanz jeder Dimension zu berechnen
- In der realen Konfiguration mit 8-dimensionalen Punkten ist die Distanzberechnung mit etwa 1/8 des Aufwands gegenüber unkomprimierten Embeddings möglich
Arrow als Datenaustauschformat verwenden
- Da Kosten für Datenbewegung und Konvertierung hoch sind, ist es vorteilhaft, komprimierte Embeddings in einem direkt nutzbaren Format zu serialisieren
- Arrow ist ein spaltenorientiertes Format, das für diesen Zweck gut geeignet ist
- Embeddings und Seitentitel werden als Arrow-Tabellen gespeichert
- Palettenindizes werden wie zweidimensionale Arrays behandelt, Titel wie eindimensionale String-Arrays
- Der Fokus kann darauf liegen, Bits an die richtige Stelle zu kopieren, statt auf Parsing oder Laden
- Da das Arrow-Array-Format nur eindimensionale Daten speichert, werden zwei Schemata verwendet, um 48-dimensionale Embeddings zu verarbeiten
- Das Metadatenschema arbeitet in Einheiten von 100.000 Zeilen
- Das Embedding-Schema arbeitet in Einheiten von
100.000 * 48Zeilen - Beim Laden werden die Embeddings wieder per reshape umgeformt
- safetensors folgt denselben Designprinzipien
- JSON serialisiert Arrays mit 48 Elementen als ASCII-Zeichen variabler Länge, was die Ladephase komplizierter macht
- Protocol Buffers verwenden für Ganzzahlen unter 32 Bit Base-128-Varints; dieses Integer-Format wird derzeit von Compute-Kernels nicht gut unterstützt
Wikipedia parsen und Embeddings erzeugen
- Zum Parsen von Wikipedia-Markup wird mediawiki parser from hell, also mwparserfromhell, verwendet
- Das Wikipedia dataset von OLM holt den neuesten Daten-Dump einer Wikipedia in einer bestimmten Sprache und parst ihn in Zeilen wie Titel und Text
- Dieser Datensatz ist weniger harmlos serialisierte Daten als vielmehr Code, der ausgeführt werden muss; dafür ist Ausführungsvertrauen nötig
- Das Parsing läuft auf allen Kernen einer Maschine
- Moderne Maschinen haben Dutzende Kerne und eignen sich damit für Systeme mit weniger als 100 Millionen Dokumenten, wie die aktuelle englische Wikipedia mit 6 Millionen Seiten
- Als Embedding-Modell wird all-minilm-l6-v2 verwendet
- Es ist ein leistungsfähiger und leichter Sentence Transformer für Englisch
- Queries und Dokumente werden in denselben Raum eingebettet
- Mit 22 Millionen Parametern ist es ein sehr leichtes Modell
- Für die Ausführung von JavaScript im Browser wird auch eine 8-Bit-quantisierte Version über ONNX und
transformers.jsgenutzt - Das Modell wurde mit Sequenzen von 128 Tokens trainiert, sodass sein Kontextfenster deutlich kürzer ist als die durchschnittliche Seitenlänge
- Seiten werden in Chunks aufgeteilt, und die Embeddings der einzelnen Chunks werden gemittelt, um ein gemitteltes Seiten-Embedding zu erzeugen
pq.js und Facet-Suche per linearem Scan
- Wenn Millionen Dokumente lokal durchsucht werden, ist ein komplexes Indexierungssystem nicht unbedingt nötig
- Ziel ist es, die nächstgelegenen Top-Ergebnisse zu einem bestimmten Punkt zu erhalten, zum Beispiel die Top 12 Distanzen
- Bei einem Distanzarray mit 10 Millionen Einträgen, einer Facet-Spalte mit 10 Millionen Einträgen und einem Facet-Wert wird gefiltert, indem bei Übereinstimmung 0 und andernfalls Infinity addiert wird, bevor top-k bestimmt wird
- Auf modernen Smartphones lassen sich die Top 100 in unter 10 ms finden
- Die Implementierung befindet sich in filtered-topk und
pq.js
- Die Distanzberechnung für Product-Quantization-komprimierte Embeddings kann in PyTorch geschrieben werden
- Es gibt
subspaceCountPaletten, und jede Palette enthältcodewordCountPunkte mitsubspaceDimDimensionen
- Es gibt
- Die Indexierung des aus PyTorch exportierten ONNX-Modells nimmt eine etwas unbeholfene Form an
- Der ONNX-Opcode GatherElements führt die benötigte Operation direkt aus
- Mit Werkzeugen wie ONNX-modifier lassen sich Knoten im Datenflussgraphen eines exportierten ONNX-Modells hinzufügen oder entfernen
- Ersetzt man mehrere Indexierungsschritte durch den korrekten einzelnen Opcode, wird die Distanzberechnung etwa 4-mal schneller
Streaming-Berechnung und interaktive UI
- Die Zahl der Embeddings, für die Distanzen berechnet werden, ist nicht fest vorgegeben
- Distanzen zwischen der Query und einigen Embeddings können im Streaming berechnet werden
- Nach ausreichend Zeit wird top-k auf den bis dahin berechneten Distanzen ausgeführt, um die Suchergebnisse zu aktualisieren
- Da alle Embeddings lokal vorliegen, ist die Latenz zur Suchinfrastruktur praktisch 0 ms
- Ziel für die UI-Reaktionsfähigkeit ist, dass Ergebnisse innerhalb von 100–300 ms nach einer Bedienaktion sichtbar werden
- MiniLM kann in unter 100 ms ausgeführt werden
- Je nach Edge-Gerät sind auch unter 15 ms möglich
- Die Distanzberechnung für 100k Embeddings ist in etwa 10 ms möglich
- Wenn der Datensatz von langen zu kurzen Artikeln sortiert wird, erscheinen die meisten endgültigen Top-Suchergebnisse schnell in den frühen Streaming-Ergebnissen
- Denn in lange Artikel ist beim Schreiben oft mehr Aufmerksamkeit und Aufwand geflossen, und sie passen häufig besser zu Suchanfragen als die zahlreichen Stub-Artikel
- Da Distanzen global berechnet werden und anschließend Facet-Filter sowie top-k angewendet werden, muss beim Ändern von Facet-Werten oder der Anzahl der Suchergebnisse für dieselbe Query nur das Filtering erneut ausgeführt werden
- Dieses Filtering dauert unter 10 ms und fühlt sich daher unmittelbar an
Wiederverwendbare pq.js-Komponenten
- Viele Bibliotheksfunktionen der vollständigen Wikipedia search app können in wiederverwendbare pq.js-Komponenten überführt werden
- Derzeit sind viele ONNX-Shapes vorab festgelegt
- Unterstützung für verschiedene Quantisierungsstufen und unterschiedliche Embedding-Dimensionen würde eine breitere Wiederverwendung ermöglichen
1 Kommentare
Hacker-News-Meinungen
Definitiv interessant, aber als ich versucht habe, ein paar Begriffe aus Philosophie und Psychologie zu beschreiben, tauchten die gesuchten Einträge alle erst um Platz 20 herum auf
Berühmtere, aber weniger genaue Einträge lagen darüber; wenn ich zum Beispiel versuchte, eine bestimmte Form der Psychotherapie zu definieren, war „psychotherapy“ immer auf Platz 1, egal was ich eingab
Umgekehrt habe ich schon einmal mit ChatGPT nach einem engen Teilgebiet gesucht, dessen Name mir nicht einfiel, und es hat ihn jedes Mal getroffen
Die Idee eines AI-Dienstes, der anhand einer Beschreibung den Namen einer Sache findet, ist gut, aber ich bin mir nicht sicher, ob die Beschränkung auf Wikipedia bzw. Wikipedia-Artikeltitel der richtige Ansatz ist; allgemeine große Sprachmodelle scheinen das schon ziemlich gut zu können
Trotzdem ist es ein Proof of Concept, und dass es lokal im Browser laufen kann, ist wirklich cool
Es gibt offenbar viele Parameter, an denen man drehen kann. Etwa ob man nur den ersten Absatz eines Artikels oder den ganzen Artikel verwendet, oder ob man innerhalb eines Bereichs nahe an einem bestimmten Artikel sucht; daran wurde noch nicht gearbeitet
Wikipedia ist ein großartiges Demo-Dataset, und ich würde gern auch andere Datasets hinzufügen. Zum Beispiel wäre es interessant, mit einem multimodalen Modell wie CLIP mehrere Datasets zu durchsuchen, so wie in iPhoto bei der Suche nach „mountain“ Fotos mit Bergen erscheinen
Große Sprachmodelle wirken wie die besten Fuzzy-Suchmaschinen und funktionieren auf eine ziemlich eigene, aber komplementäre Weise zu traditionellen Suchmaschinen
Das Konzept gefällt mir, aber die Ergebnisse waren nicht gut
Ich habe „weird looking monkey“ eingegeben und Ergebnisse wie Nasenaffe oder Goldstumpfnase erwartet, bekam aber nur Artikel wie „Pet monkey“, „List of individual monkeys“, „Ethnoprimatology“ und „Monkey“
Wenn man dieselbe Anfrage in Google eingibt, kommen genau die erwarteten Ergebnisse heraus; das war schade, denn ich wollte seltsam aussehende Affen entdecken, die ich noch nicht kannte
Der Fokus der Demo lag darauf, die Embedding-Datenbank zu zeigen, aber die Embeddings selbst sind zumindest etwas nützlich
Es werden keinerlei Analysedaten darüber gespeichert, was Nutzer auf der Seite finden oder nicht finden, daher ist man nicht darauf vorbereitet, die Suchergebnisse zu verbessern
Dadurch wird es schwierig, aus diesem Wissenskorpus gute Ergebnisse zu bekommen
Die Implementierung ist wirklich gut, und dass man das offline machen kann, ist cool. Allerdings wirkt die Qualität der Embeddings noch unzureichend
Ein Trick, der helfen könnte, wäre, nicht den gesamten Wikipedia-Artikel einzubetten, sondern nur einen definitionsnahen Satz oder normalerweise den ersten Satz bzw. ersten Absatz. Ich weiß nicht genau, welcher Teil derzeit verwendet wird
Meine Site OneLook bietet seit 2003 unter https://onelook.com/thesaurus/ eine ähnliche Funktion, um Wörter und Begriffe anhand von Beschreibungen zu finden
Anfangs war es eine reine umgekehrte Wörterbuchsuche, aber in den letzten 20 Jahren habe ich Wort-Embeddings, Satz-Embeddings und zuletzt auch große Sprachmodelle ausprobiert; heutzutage erzeugt GPT Kandidaten für Eingaben, die das System nicht selbst beantworten kann
Bei dieser Aufgabe sind große Sprachmodelle den früheren Ansätzen so deutlich überlegen, dass meine Motivation, diesen Teil von OneLook weiter zu verbessern, etwas gesunken ist. Ich sehe oft Leute, die sagen, dass Reverse-Definition-Search der Hauptgrund ist, warum sie ChatGPT nutzen
Ich habe es etwas spät gesehen, aber Text-Embeddings sind, zumindest nach den in diesem Beitrag verwendeten, normalerweise nicht besonders gut darin, nach Vibe zu suchen
Meist geht es eher darum, überlappende Wörter zu vergleichen oder Inhalte zu finden, die der Anfrage ähnlich sind
Es gibt allerdings ein aktuelles Paper, das genau dieses Problem angeht: „Retrieving Texts based on Abstract Descriptions“ (Ravfogel et al., 2023) https://arxiv.org/abs/2305.12517
Das Paper enthält viele Beispiele für Suchen mit abstrakten Beschreibungen wie „Architekt, der Gebäude entwirft“, „Unternehmen, das Teil eines anderen Unternehmens ist“ oder „Buch, das die Entwicklung eines Genres beeinflusst hat“
Die dortigen Embeddings scheinen solche Suchen deutlich besser zu unterstützen; es wäre interessant, die im verlinkten Beitrag beschriebene Offline-Wikipedia-Suche mit diesem neuen Typ von Embedding erneut auszuprobieren
Die Seite funktioniert in meiner Umgebung gerade nicht;
model_quantized.onnxwird nicht geladenWährend ich das eintippe, wurden schon 19,2 MB heruntergeladen, bei etwa 50 KB/s. Wenn jeder Besucher das auslöst, könnte das Lee Buttermans Bandbreitenkosten Schreckliches antun
Was es macht, ist an sich sehr beeindruckend, aber die Qualität der Suchergebnisse wirkt nicht gut
Aus Erfahrung weiß ich, dass es wirklich schwierig ist, die Qualität von Suchergebnissen von Hand zu bewerten. Man kann sehr nah an einem ausgezeichneten Ergebnis liegen und trotzdem ein wesentlich schlechteres Match zurückgeben
Mit neueren Satz-Embeddings dürften die Ergebnisse besser werden, und ich muss mehr Daten sammeln
Die Technik ist sehr beeindruckend, die Ergebnisse waren es aber nicht
Eine Suche nach „pointy building in Paris“ lieferte Tourism in Paris, Bourse de commerce (Paris), Grands Projets of François Mitterrand, List of tallest buildings and structures in the Paris region, List of tourist attractions in Paris, Palais des congrès de Paris, Landmarks in Paris, Palais de la Bourse, Lyon, Outline of Paris, Architecture of Paris
Das berühmteste spitze Bauwerk in Paris war überhaupt nicht zu sehen
Für solche Anwendungen sind Satz-Embeddings des gesamten Dokuments vielleicht nicht die beste Lösung
Ich habe gerade den Artikel geprüft: Das Wort „building“ kommt 19-mal vor, aber meistens als Verb, und danach geht es um das „Chrysler Building“
Zumindest, sofern es nicht noch ein anderes berühmtes spitzes Bauwerk gibt, das mir nicht einfällt
Ein Teil der Magie von Suchmaschinen liegt darin, die Embeddings von Seiten, die auf die jeweilige Seite verlinken, oder klassische Information-Retrieval-Keywords gewichtet nach Klickzahlen und Autoritätswerten beizumischen
Ohne dieses Signal wird viel nützliche Information ignoriert, und die Ergebnisse fühlen sich nicht magisch an
Trotzdem eine beeindruckende und interessante Demo
Ich wollte es mögen, aber bei den Suchen, die ich ausprobiert habe, kamen kaum relevante Ergebnisse heraus
Bei „The wizard in The Lord of the Rings“ gab es weder Gandalf noch Saruman, sondern nur LOTR-bezogene Bücher
Bei „Protagonist of Scorsese's Taxi Driver“ fehlte Travis Bickle
Bei „A person that plants trees for a living“ war aus irgendeinem Grund gardener nicht in der Liste
Bei „Curly-haired painter on TV“ tauchte Bob Ross überhaupt nicht auf
Bei „Unusually shaped modern art museum in Spain“ kam Bilbao zwar auf Platz 4, der Rest hatte aber keine ungewöhnliche Form
Bei „Dog shaped like a sausage“ würde ich erwarten, dass dachshund unter den Top-Ergebnissen ist
Wenn es keine Artikel gäbe, wäre das Fehlen weniger merkwürdig, tatsächlich existieren sie aber alle
„Vibes“ trifft es viel besser als „sentence embeddings“. Vielleicht sollte ich diesen Ausdruck auch verwenden :)
Der Autor des Originalbeitrags hat nicht erklärt, warum er dieses Wort gewählt hat, und es passt zu keiner Verwendung von „vibe“, die ich kenne
Vielleicht war „gist“ einfach nicht buzzword-tauglich genug