La ampliación de los modelos de lenguaje grande (LLM) está cada vez más limitada por la sobrecarga de comunicación de la memoria entre la memoria de alto ancho de banda (HBM) y SRAM. Específicamente, el tamaño de la caché de valor clave (KV) se escala tanto con las dimensiones del modelo como con la longitud del contexto, lo que crea un cuello de botella significativo para la inferencia de contexto largo. El equipo de investigación de Google ha propuesto TurboQuant, un marco de cuantificación ajeno a los datos diseñado para lograr tasas de distorsión casi óptimas para vectores euclidianos de alta dimensión y al mismo tiempo abordar tanto el error cuadrático medio (MSE) como la distorsión interna del producto.
Abordar el muro de la memoria con VQ ajeno a los datos
La cuantización vectorial (VQ) en el espacio euclidiano es un problema fundamental arraigado en la teoría de codificación fuente de Shannon. Los algoritmos VQ tradicionales, como la cuantificación de productos (PQ), a menudo requieren un preprocesamiento fuera de línea extenso y una capacitación en libros de códigos dependiente de los datos, lo que los hace inadecuados para los requisitos dinámicos de las cargas de trabajo de IA en tiempo real, como la administración de caché KV.
TurboQuant es un algoritmo que “ajena los datos” y no requiere ajustes ni calibraciones específicas del conjunto de datos. Está diseñado para ser altamente compatible con aceleradores modernos como las GPU al aprovechar operaciones vectorizadas en lugar de búsquedas binarias lentas y no paralelizables.
La mecánica geométrica de TurboQuant
El mecanismo central de TurboQuant implica aplicar una rotación aleatoria Π E Rdxd a los vectores de entrada. Esta rotación induce una distribución Beta concentrada en cada coordenada, independientemente de los datos de entrada originales. En dimensiones altas, estas coordenadas se vuelven casi independientes y distribuidas de manera idéntica (iid).
Esta casi independencia simplifica el diseño de cuantificación, lo que permite a TurboQuant resolver un problema de cuantificación escalar continuo 1D k-medias/Max-Lloyd por coordenada. El cuantificador escalar óptimo para un ancho de bits b dado se encuentra minimizando la siguiente función de costo MSE:
$$\mathcal{C}(f_{X},b):=min_{-1\le c_{1}\le c_{2}\le…\le c_{2^{b}}\le1}\sum_{i=1}^{2^{b}}\int_{\frac{c_{i-1}+c_{i}}{2}}^{\frac{c_{i}+c_{i+1}}{2}}|x-c_{i}|^{2}\cdot f_{X}(x)dx$$
Al resolver esta optimización una vez para anchos de bits relevantes y almacenar los libros de códigos resultantes, TurboQuant puede cuantificar vectores de manera eficiente durante la inferencia en línea.
Eliminación del sesgo interno del producto
Un desafío principal en la cuantificación es que los mapas optimizados estrictamente para MSE a menudo introducen sesgos al estimar los productos internos, que son las operaciones fundamentales en los mecanismos de atención del transformador. Por ejemplo, un cuantificador MSE óptimo de 1 bit en dimensiones altas puede exhibir un sesgo multiplicativo de 2/π.
Para corregir esto, Google Research desarrolló TURBOQUANTprod, un enfoque de dos etapas:
Etapa MSE: Aplica un cuantificador TURBOQUANTmse usando un ancho de bits de b-1 para minimizar la norma L2 del vector residual. Etapa imparcial: aplica una transformada cuantificada de Johnson-Lindenstrauss (QJL) de 1 bit al vector residual.
Esta combinación da como resultado un ancho de bits general de b al tiempo que proporciona un estimador demostrablemente insesgado para los productos internos:
\(\mathbb{E}_{Q}[\langle y,Q^{-1}(Q(x))\rangle ]=\langle y,x\rangle \)
Desempeño teórico y empírico
El equipo de investigación estableció límites inferiores de la teoría de la información utilizando el límite inferior de Shannon (SLB) y el principio minimax de Yao. La distorsión MSE de TurboQuant está demostrablemente dentro de un pequeño factor constante (≈ 2,7) del límite teórico absoluto en todos los anchos de bits. Con un ancho de bits de b=1, está a sólo un factor de aproximadamente 1,45 del óptimo.
En los puntos de referencia de generación de LLM de extremo a extremo utilizando Llama-3.1-8B-Instruct y Ministral-7B-Instruct, TurboQuant demostró una retención de alta calidad. Con una relación de compresión de 4x, el modelo mantuvo una precisión de recuperación del 100% en el punto de referencia Needle-In-A-Haystack. En el punto de referencia Needle-In-A-Haystack, TurboQuant igualó el rendimiento de máxima precisión hasta 104.000 tokens bajo una compresión de 4×.
Para anchos de bits no enteros, el sistema emplea una estrategia de tratamiento de valores atípicos, asignando mayor precisión (por ejemplo, 3 bits) a canales atípicos específicos y menor precisión (por ejemplo, 2 bits) a valores no atípicos, lo que resulta en tasas de bits efectivas como 2,5 o 3,5 bits por canal.
Velocidad y eficiencia de indexación
En las tareas de búsqueda de vecinos más cercanos, TurboQuant superó a la Cuantificación de Producto (PQ) estándar y a RabitQ en recuperación, al tiempo que redujo el tiempo de indexación a prácticamente cero. Debido a que TurboQuant no tiene en cuenta los datos, elimina la necesidad de la fase de entrenamiento de k-medias que requiere mucho tiempo y que requiere PQ, que puede llevar cientos de segundos para grandes conjuntos de datos.
TurboQuant representa un cambio matemáticamente fundamentado hacia una cuantificación vectorial eficiente y compatible con hardware que cierra la brecha entre los límites teóricos de distorsión y la implementación práctica de la IA.
Conclusiones clave
No se requiere preprocesamiento: a diferencia de la cuantificación de productos (PQ) estándar, TurboQuant no tiene en cuenta los datos y funciona instantáneamente sin necesidad de un entrenamiento de k-means que requiere mucho tiempo en su conjunto de datos específico. Perfección casi teórica: logra tasas de distorsión casi óptimas, permaneciendo dentro de un pequeño factor constante de aproximadamente 2,7 del límite inferior teórico de la información establecido por Shannon. Productos internos imparciales: al utilizar un enfoque de dos etapas (aplicando una cuantificación óptima MSE seguida de una transformación QJL de 1 bit en el residual), se proporcionan estimaciones imparciales del producto interno, lo cual es vital para mantener la precisión de los mecanismos de atención del transformador. Ahorros masivos de memoria: en la implementación de LLM, comprime la caché KV más de 5 veces. Logra una neutralidad de calidad absoluta a 3,5 bits por canal y mantiene una recuperación del 100 % en pruebas de “aguja en un pajar” de hasta 104 000 tokens. Indexación instantánea para búsqueda: para bases de datos vectoriales, TurboQuant reduce el tiempo de indexación prácticamente a cero (por ejemplo, 0,0013 s para vectores de 1536 dimensiones) y, al mismo tiempo, supera consistentemente el PQ tradicional en la recuperación de búsqueda.
Consulte el documento y los detalles técnicos. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 120.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.