Mengueri 3 Miliar Vektor

• Karena masalah kueri 3 miliar vektor dari Jeff Dean, ini adalah catatan eksperimen teknis yang membahas proses mengimplementasikan sendiri solusi map-reduce optimal
• Dimulai dari implementasi naive yang menghitung dot product antara 3 miliar vektor float32 berdimensi 768 dan 1.000 vektor kueri, lalu mencapai peningkatan performa bertahap melalui vektorisasi numpy dan konversi ke float32
• Dengan patokan 3.000 vektor, metode naive memerlukan sekitar 2 detik, setelah vektorisasi menjadi 0,01 detik, dan dengan float32 turun hingga 0,0045 detik
• Saat diskalakan ke 3 miliar, matriks hasil memerlukan sekitar 8,6 TB memori sehingga muncul masalah OOM, dan dibutuhkan optimasi tambahan seperti pemrosesan batch, memory mapping, penulisan ulang dengan Rust/C, serta SimSIMD
• Ditekankan bahwa tantangan inti yang lebih sulit daripada solusi teknis adalah mendefinisikan kebutuhan (bentuk pengembalian, spesifikasi mesin, apakah kompresi diperbolehkan, dan lain-lain)

---

Titik awal masalah

• Berawal dari upaya untuk mengimplementasikan sendiri solusi map-reduce optimal setelah melihat diskusi tentang Jeff Dean dan masalah kueri 3 miliar vektor
• Vektor adalah array angka floating-point berdimensi n, dan pencarian vektor digunakan untuk menemukan kata atau item yang mirip secara semantik
• Ini adalah pola umum pemanfaatan embedding dalam aplikasi pencarian generatif seperti pencarian, rekomendasi, dan Cursor

Implementasi naive

• Asumsi awal: ada 3 miliar vektor dokumen target pencarian dan sekitar 1.000 vektor kueri, semuanya disimpan di disk dalam format .npy
• Dimensi vektor adalah 768, ukuran umum yang digunakan oleh banyak model kueri embedding berbasis kemiripan
• Menggunakan pendekatan loop ganda yang membandingkan setiap vektor kueri dengan seluruh vektor dokumen untuk menghitung dot product dan mengembalikan hasilnya
• Saat diuji lebih dulu dengan 3.000 vektor, hasilnya memerlukan sekitar 2 detik di M2 MacBook (skala 1 juta kali lebih kecil dari 3 miliar)

Optimasi vektorisasi (Vectorized)

• Menerapkan pendekatan vektorisasi dengan mengganti loop ganda menggunakan operasi perkalian matriks numpy (vectors_file @ query_vectors.T)
• Untuk 3.000 vektor, waktu turun drastis menjadi 0,0107 detik
• Saat diperluas ke 3 juta vektor, waktu yang dibutuhkan adalah 12,85 detik

Konversi float32

• Melakukan optimasi tambahan dengan mengonversi data menjadi np.float32
• Untuk 3.000 vektor, waktu berkurang lagi hingga 0,0045 detik
• Karena untuk 3 juta vektor sekitar 13 detik, maka untuk 3 miliar diperkirakan 1.000 kali lipat, yaitu sekitar 3.216 menit

Ringkasan perbandingan performa

• Metode naive (3.000 vektor): 1,9877 detik
• Metode vektorisasi (3.000 vektor): 0,0107 detik
• Metode vektorisasi (3 juta vektor): 12,8491 detik
• Metode float32 (3.000 vektor): 0,0045 detik

Masalah OOM dan arah optimasi lanjutan

• Memori yang dibutuhkan saat memproses 3 miliar objek sebagai float32 (4 byte): sekitar 8,6 TB, sehingga saat dijalankan terjadi OOM
• Diperlukan optimasi tambahan ke arah yang disarankan Jeff Dean:
  • Mengubah kode dengan generator dan operasi perbandingan batch
  • Menulis ke disk setiap beberapa operasi tertentu atau menggunakan memory mapping
  • Menulis ulang kode optimasi level sistem dengan Rust atau C
  • Memanfaatkan library khusus perbandingan kemiripan vektor skala besar seperti SimSIMD

Pentingnya mendefinisikan kebutuhan

• Sebelum optimasi lebih lanjut, ditemukan beberapa bagian masalah itu sendiri yang belum jelas:
  • Apakah ini mengembalikan seluruh hasil setelah satu kueri atas semua 3 miliar vektor dengan satu vektor kueri, atau merupakan pencarian ANN top-k
  • Apakah vektor asli juga harus dikembalikan, dan apakah perlu reranking berdasarkan kemiripan
  • Apakah seluruh data akan dikueri sekaligus dengan 1.000 vektor kueri
  • Apakah vektor sudah ada di memori, dibaca satu per satu dari disk, atau menggunakan pendekatan streaming
  • Apakah dijalankan secara lokal, bagaimana spesifikasi mesinnya, dan apakah GPU bisa digunakan
  • Dampak ukuran embedding terhadap representasi hasil serta ukuran vektor input/output
  • Apakah kompresi dari float64 ke float32 dapat diterima dari sisi akurasi
  • Apakah ini proyek sekali jalan, dan berapa banyak waktu yang tersedia untuk pembuatannya
• Ditekankan bahwa mendefinisikan kebutuhan secara akurat adalah tantangan tersulit, bahkan dibanding solusi teknis itu sendiri