3 Milliarden Vektoren abfragen

• Ausgelöst durch Jeff Deans Problem mit 3 Milliarden Vektoren dokumentiert dieser technische Erfahrungsbericht den Versuch, selbst eine optimale Map-Reduce-Lösung zu implementieren
• Ausgangspunkt ist eine naive Implementierung zur Berechnung des dot product von 3 Milliarden 768-dimensionalen float32-Vektoren und 1.000 Query-Vektoren; anschließend werden schrittweise Performance-Verbesserungen durch NumPy-Vektorisierung und die Umstellung auf float32 erzielt
• Bei 3.000 Vektoren sinkt die Laufzeit von etwa 2 Sekunden mit der naiven Methode auf 0,01 Sekunden nach der Vektorisierung und auf 0,0045 Sekunden mit float32
• Hochskaliert auf 3 Milliarden Vektoren benötigt die Ergebnismatrix etwa 8,6 TB Speicher, was zu OOM-Problemen führt; dafür sind weitere Optimierungen wie Batch-Verarbeitung, Memory Mapping, eine Neufassung in Rust/C oder SimSIMD nötig
• Es wird betont, dass nicht die technische Lösung, sondern die Definition der Anforderungen (Rückgabeformat, Maschinenspezifikation, zulässige Kompression usw.) die schwierigere Kernaufgabe ist

---

Der Ausgangspunkt des Problems

• Nach einem Austausch über Jeff Dean und das Problem mit 3 Milliarden Vektoren entstand der Versuch, selbst eine optimale Map-Reduce-Lösung zu implementieren
• Vektoren sind Arrays aus Fließkommazahlen mit n Dimensionen; die Vektorsuche wird verwendet, um semantisch ähnliche Wörter oder Elemente zu finden
• Ein verbreitetes Muster bei Suche, Empfehlungen und generativen Suchanwendungen wie Cursor ist der Einsatz von Embeddings

Naive Implementierung

• Anfangsannahme: 3 Milliarden Dokumentvektoren als Suchbasis und etwa 1.000 Query-Vektoren, alle auf der Festplatte im .npy-Format gespeichert
• Die Vektordimension beträgt 768, eine gängige Größe in vielen embeddingbasierten Query-Modellen
• Verwendet wird ein Doppelschleifen-Ansatz, bei dem jeder Query-Vektor mit allen Dokumentvektoren verglichen wird, um das dot product zu berechnen und das Ergebnis zurückzugeben
• Ein erster Test mit 3.000 Vektoren benötigt auf einem M2 MacBook etwa 2 Sekunden – bei einer Größenordnung, die eine Million Mal kleiner ist als 3 Milliarden

Vektorisierungs-Optimierung

• Die Doppelschleife wird durch einen vektorisierten Ansatz mit NumPys Matrixmultiplikation (vectors_file @ query_vectors.T) ersetzt
• Bei 3.000 Vektoren verbessert sich die Laufzeit drastisch auf 0,0107 Sekunden
• Skaliert auf 3 Millionen Vektoren dauert die Berechnung 12,85 Sekunden

Umstellung auf float32

• Durch die Konvertierung der Daten zu np.float32 wird weiter optimiert
• Bei 3.000 Vektoren sinkt die Laufzeit auf 0,0045 Sekunden
• Wenn 3 Millionen Vektoren etwa 13 Sekunden benötigen, ergibt sich für 3 Milliarden eine um den Faktor 1.000 größere geschätzte Laufzeit von etwa 3.216 Minuten

Zusammenfassung des Performance-Vergleichs

• Naive Methode (3.000 Vektoren): 1,9877 Sekunden
• Vektorisierte Methode (3.000 Vektoren): 0,0107 Sekunden
• Vektorisierte Methode (3 Millionen Vektoren): 12,8491 Sekunden
• float32-Methode (3.000 Vektoren): 0,0045 Sekunden

OOM-Problem und weitere Optimierungsrichtungen

• Für die Verarbeitung von 3 Milliarden Objekten als float32 (4 Byte) wird ein Speicherbedarf von etwa 8,6 TB berechnet; beim Ausführen tritt ein OOM-Fehler auf
• Zusätzliche Optimierungen in der von Jeff Dean vorgeschlagenen Richtung sind nötig:
  • Code-Änderung auf Generatoren und Batch-Vergleichsoperationen
  • Ergebnisse in regelmäßigen Abständen auf die Festplatte schreiben oder Memory Mapping verwenden
  • Neuschreiben des Codes in Rust oder C für systemnahe Optimierungen
  • Nutzung spezialisierter Bibliotheken wie SimSIMD für großskalige Vektor-Ähnlichkeitsvergleiche

Die Bedeutung klar definierter Anforderungen

• Noch vor weiteren Optimierungen zeigen sich Unklarheiten im eigentlichen Problem:
  • Soll für einen einzelnen Query-Vektor einmal über alle 3 Milliarden Vektoren gesucht und das vollständige Ergebnis zurückgegeben werden, oder handelt es sich um eine Top-k-ANN-Suche?
  • Müssen auch die ursprünglichen Vektoren zurückgegeben werden, und ist ein Reranking nach Ähnlichkeit erforderlich?
  • Sollen 1.000 Query-Vektoren jeweils einmal gegen den gesamten Bestand ausgeführt werden?
  • Liegen die Vektoren bereits im Speicher, werden sie einzeln von der Festplatte gelesen, oder erfolgt die Verarbeitung als Streaming?
  • Erfolgt die Ausführung lokal, welche Maschinenspezifikation liegt vor, und ist die Nutzung einer GPU möglich?
  • Welchen Einfluss hat die Embedding-Größe auf die Darstellung der Ergebnisse sowie auf Ein- und Ausgabegrößen der Vektoren?
  • Ist die Kompression von float64 auf float32 hinsichtlich der Genauigkeit akzeptabel?
  • Handelt es sich um ein einmaliges Projekt, und wie viel Zeit steht für die Erstellung zur Verfügung?
• Schwieriger als die technische Lösung selbst ist es, die Anforderungen präzise zu definieren