TurboQuant: compresión 6x de LLM KV-cache

TurboQuant: Caché KV de LLM 6 veces más pequeño sin pérdida de precisión

En la generación de modelos lingüísticos autoregresivos, el mecanismo de atención depende de los vectores clave y valor de tokens anteriores. Recalcular todo el contexto en cada paso es ineficiente, por lo que se almacenan representaciones intermedias. El caché KV crece linealmente con la longitud del contexto, limitando el rendimiento y el tamaño máximo de contexto.

En secuencias largas, la inferencia alcanza límites de ancho de banda de memoria: las GPU pasan más tiempo leyendo datos que calculando. Esto aumenta la latencia y el costo. TurboQuant aborda esto comprimiendo el caché manteniendo el significado semántico para la atención.

Cómo funciona PolarQuant

La primera etapa es PolarQuant. Los vectores se transforman a coordenadas polares tras una rotación aleatoria. El radio codifica la magnitud del vector, el ángulo captura su dirección.

# Pseudocódigo: PolarQuant
rotated = rotate_random(vector)
rho = norm(rotated)    # radio
phi = angle(rotated)   # ángulo
quantized = quantize(rho, phi)

La rotación aleatoria simplifica la distribución de los datos, haciéndola predecible para la cuantización. No se necesitan parámetros adicionales ni diccionarios — la compresión alcanza entre 3 y 4 bits por valor sin ajustar el modelo.

Corrección de errores mediante QJL

La cuantización introduce ruido. La segunda etapa es QJL (Johnson-Lindenstrauss cuantizado), que codifica el error residual en solo 1 bit por valor. La transformación preserva las distancias euclidianas entre vectores —crucial para la atención de producto punto.

QJL aproxima el lema de Johnson-Lindenstrauss en forma cuantizada:

• Entrada: residuo tras PolarQuant
• Salida: corrección de 1 bit
• Propiedad: ||Qx - Qy|| ≈ ||x - y|| para x, y ∈ caché KV

Juntos logran entre 1,25 y 2 bits por valor con mínima degradación de perplejidad.

Pruebas y rendimiento

Probado en Llama-3.1-8B-Instruct (LongBench):

• Compresión: 6–8x frente al modelo base
• Aceleración en cálculo de logits: hasta 8x en H100 (base JAX)
• Calidad: menos del 1% de caída en tareas de "aguja en un pajar" y contexto largo

En GloVe (d=200), TurboQuant lidera en el equilibrio entre recuperación y relación de compresión (1:k óptimo).

| Método | Bits/Valor | Compresión | Caída de perplejidad |

|--------|------------|-------------|----------------------|

| Base | 16 | 1x | 0% |

| INT8 | 8 | 2x | 0,5% |

| TurboQuant (3-bit) | 3 | 6x | 0,2% |

Alcanza hasta 3 bits sin reentrenamiento —un avance clave para la cuantización zero-shot.

Conclusiones clave

• Compresión extrema: El caché KV se reduce 6 veces a 3 bits por valor.
• Inferencia más rápida: Hasta 8x de aceleración en cargas limitadas por memoria.
• Preservación de distancias: QJL asegura compatibilidad con atención.
• Aplicabilidad universal: Funciona para búsqueda de vectores y RAG.
• Sin ajuste fino necesario: Reemplazo listo para usar en LLMs existentes.

Más allá de los LLM

TurboQuant trasciende arquitecturas de transformadores. En bases de datos de vectores, la compresión de embeddings mantiene el rendimiento de búsquedas ANN (HNSW, FAISS). En dispositivos de borde, reduce la demanda de DRAM, permitiendo inferencia local.

Para ingenieros senior: intégralo en motores de inferencia personalizados. El soporte para JAX/Flax simplifica la prototipación. Prueba en contextos largos (>128k tokens) para obtener ganancias reales.

— Editorial Team