Consultar 3 mil millones de vectores

• A partir del problema de consultar 3 mil millones de vectores de Jeff Dean, este es un registro de experimentación técnica sobre cómo implementar directamente una solución óptima de map-reduce
• Comienza con una implementación naive que calcula el dot product entre 3 mil millones de vectores float32 de 768 dimensiones y 1,000 vectores de consulta, y logra mejoras graduales de rendimiento con vectorización en numpy y conversión a float32
• Con 3,000 vectores, el enfoque naive tarda unos 2 segundos; tras vectorizarlo baja a 0.01 segundos, y al usar float32 se reduce hasta 0.0045 segundos
• Al escalar a 3 mil millones, la matriz de resultados requiere cerca de 8.6 TB de memoria, lo que provoca un problema de OOM, por lo que se necesitan optimizaciones adicionales como procesamiento por lotes, memory mapping, reescritura en Rust/C y herramientas como SimSIMD
• Se enfatiza que, más que la solución técnica, el desafío central es definir los requisitos (forma de retorno, especificaciones de la máquina, si se permite compresión, etc.)

---

Punto de partida del problema

• El intento comenzó al ver el intercambio sobre Jeff Dean y el problema de consultar 3 mil millones de vectores, con la idea de implementar personalmente la solución óptima de map-reduce
• Un vector es un arreglo de números de punto flotante de n dimensiones, y la búsqueda vectorial se usa para encontrar palabras o elementos semánticamente similares
• Es un patrón común usar embeddings en aplicaciones de búsqueda generativa como búsqueda, recomendación y Cursor

Implementación naive

• Suposición inicial: 3 mil millones de vectores de documentos a buscar y alrededor de 1,000 vectores de consulta, todos almacenados en disco en formato .npy
• La dimensión del vector es 768, un tamaño común en muchos modelos de consulta de embeddings basados en similitud
• Se usa un enfoque de doble bucle que compara cada vector de consulta contra todos los vectores de documentos para calcular el dot product y devolver el resultado
• Primero se probó con 3,000 vectores y el resultado fue de unos 2 segundos en una MacBook M2 (una escala un millón de veces menor que 3 mil millones)

Optimización vectorizada (Vectorized)

• Se aplicó un enfoque vectorizado que reemplaza el doble bucle con la operación de multiplicación de matrices de numpy (vectors_file @ query_vectors.T)
• Con 3,000 vectores, mejoró drásticamente hasta 0.0107 segundos
• Al escalar a 3 millones de vectores, tomó 12.85 segundos

Conversión a float32

• Se realizó una optimización adicional convirtiendo los datos a np.float32
• Con 3,000 vectores, el tiempo se redujo hasta 0.0045 segundos
• Si con 3 millones de vectores tarda unos 13 segundos, para 3 mil millones se estima unas 1,000 veces más: aproximadamente 3,216 minutos

Resumen comparativo de rendimiento

• Método naive (3,000 vectores): 1.9877 segundos
• Método vectorizado (3,000 vectores): 0.0107 segundos
• Método vectorizado (3 millones de vectores): 12.8491 segundos
• Método con float32 (3,000 vectores): 0.0045 segundos

Problema de OOM y dirección de optimización adicional

• Memoria necesaria para procesar 3 mil millones de objetos como float32 (4 bytes): cerca de 8.6 TB, lo que produce OOM durante la ejecución
• Se requieren optimizaciones adicionales en la línea sugerida por Jeff Dean:
  • Cambiar el código para usar generadores y operaciones de comparación por lotes
  • Escribir en disco cada cierta cantidad de operaciones o usar memory mapping
  • Reescribir el código con optimizaciones a nivel de sistema en Rust o C
  • Aprovechar bibliotecas especializadas para comparación masiva de similitud vectorial como SimSIMD

La importancia de definir los requisitos

• Antes de optimizar más, se detectaron varias ambigüedades en el problema mismo:
  • Si se consulta una sola vez todo el conjunto de 3 mil millones con un único vector de consulta y se devuelven todos los resultados, o si se trata de una búsqueda top-k ANN
  • Si también deben devolverse los vectores originales y si hace falta reranking según similitud
  • Si se consulta todo el conjunto una sola vez con 1,000 vectores de consulta
  • Si los vectores ya están en memoria, si se leen del disco uno por uno o si el enfoque es de streaming
  • Si la ejecución es local, cuáles son las especificaciones de la máquina y si se puede usar GPU
  • El impacto del tamaño del embedding en la representación de los resultados y en el tamaño de los vectores de entrada/salida
  • Si la compresión de float64 a float32 es aceptable en términos de precisión
  • Si es un proyecto de una sola vez y cuánto tiempo hay disponible para generarlo
• Se concluye que, más que la solución técnica en sí, la tarea más difícil es definir con precisión los requisitos