🤖 machine learning

Output-Aware Rotation for INT2 KV-Cache Quantization

Este artículo propone OptR, un método de rotación consciente de la salida que minimiza los errores de atención-salida post-proyección mediante correcciones ortogonales por cabeza y reparametrización de claves, mejorando así significativamente el rendimiento de la cuantización del caché KV en INT2 para modelos de lenguaje de gran tamaño mientras mantiene un sobrecoste de inferencia insignificante.

Autores originales: Vincent-Daniel Yun, Woosang Lim, Minsoo Cheong, Sunwoo Lee, Murali Annavaram, Sai Praneeth Karimireddy, Sungjoo Yoo

Publicado 2026-08-05

📖 8 min de lectura🧠 Análisis profundo

CC0 1.0

Autores originales: Vincent-Daniel Yun, Woosang Lim, Minsoo Cheong, Sunwoo Lee, Murali Annavaram, Sai Praneeth Karimireddy, Sungjoo Yoo

Artículo original dedicado al dominio público bajo CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás tratando de recordar una historia masiva para contársela a un amigo. Tienes un cuaderno donde anotas cada detalle importante que has escuchado hasta ahora. Cuanto más grande es la historia, más páginas necesitas. Ahora, imagina que tu cuaderno se está quedando sin espacio y no puedes permitirte comprar uno más grande. Para solucionar esto, decides reducir el tamaño de tu letra. En lugar de escribir letras completas y claras, las garabateas usando solo cuatro símbolos diminutos: un punto, un guion, un círculo y una cruz. Esto es como comprimir una enorme cantidad de información en un espacio muy pequeño.

En el mundo de la inteligencia artificial, específicamente en los Grandes Modelos de Lenguaje (LLMs), estos cuadernos se llaman "KV caches" (cachés de clave-valor). Almacenan el contexto de una conversación para que la IA pueda recordar lo que se dijo anteriormente. A medida que las conversaciones se alargan, estos cachés se vuelven enormes, consumiendo memoria y ralentizando todo. Para solucionarlo, los científicos intentan encoger los datos a solo 2 bits (usando solo cuatro niveles, como nuestros cuatro símbolos). Pero aquí está el truco: cuando aplastas los datos con tanta fuerza, los "outliers" —los detalles inusuales y realmente importantes— se aplastan en la forma incorrecta, y la IA empieza a cometer errores. Es como intentar meter una papa gigante y deforme en una caja diminuta; la papa se golpea y, cuando la sacas, no se ve como la papa que pusiste dentro.

Durante un tiempo, los investigadores intentaron arreglar esto rotando la papa antes de aplastarla, con la esperanza de repartir los bultos de manera uniforme, pero estaban midiendo el éxito por qué tan bien se veía la papa dentro de la caja, no por qué tan bien la IA podía usarla para contar la historia después. Este artículo introduce un nuevo método llamado OptR (Rotación Consciente de la Salida) que cambia las reglas del juego. En lugar de solo intentar que los datos se vean bien al ser aplastados, OptR comprueba si la IA aún puede entender la historia después de que los datos han sido aplastados y des-aplastados. Resulta que, al prestar atención al resultado final, la IA puede recordar mucho más con precisión, incluso con ese cuaderno de cuatro símbolos diminutos.

El Problema: El Efecto de la "Papa Aplastada"

Cuando una IA lee un texto largo, construye un mapa mental de las palabras que ha visto. Este mapa se almacend en el KV cache. Para ahorrar espacio, los investigadores utilizan una técnica llamada cuantización para encoger este mapa. La versión más extrema es la cuantización INT2, que reduce los datos a solo cuatro valores posibles. Es increíblemente eficiente —utilizando solo 1/8 de la memoria de los formatos estándar— pero es arriesgado.

Piensa en los datos del caché como un grupo de estudiantes parados en una fila. La mayoría de los estudiantes tienen una estatura promedio, pero algunos son gigantes. Si intentas meterlos a todos en una habitación pequeña (el rango INT2), los gigantes terminan aplastados, y los estudiantes promedio también se ven aplastados porque la habitación es demasiado pequeña para los gigantes. Esto causa el "error de cuantización", donde la IA recuerda erróneamente tanto a los gigantes como a los estudiantes promedio.

Para solucionar esto, los métodos anteriores utilizaban la rotación. Imagina que haces girar la fila de estudiantes para que los gigantes ya no estén parados rectos, sino inclinados hacia un lado. Esto distribuye su altura a través de la habitación, haciendo que sea más fácil meter a todos sin aplastarlos. Sin embargo, el artículo argumenta que los métodos existentes estaban cometiendo un error. Estaban haciendo girar a los estudiantes solo para que encajaran bien en la habitación (minimizando el error de los datos almacenados), pero no estaban comprobando si los estudiantes aún podrían correr una carrera correctamente una vez que salieran de la habitación.

Los autores señalan un desajuste: la rotación que hace que los datos se vean mejor dentro de la caja no es necesariamente la rotación que ayuda a la IA a desempeñarse mejor después de usar esos datos. A la IA no le importa si los datos se ven perfectos en el caché; le importa si la respuesta final es correcta.

La Solución: OptR (Rotación Consciente de la Salida)

El artículo propone OptR, un método que optimiza la rotación basándose en el resultado final de la IA, no solo en el almacenamiento.

Así es como funciona OptR, paso a paso:

Centrar los Datos (El Truco de "Nivelación"): Antes de aplastar los datos, OptR resta el valor promedio de las claves (la "Reparametrización de Claves"). Imagina que los gigantes en nuestra fila en realidad estaban parados sobre una plataforma alta. OptR baja la plataforma para que todos estén parados en el mismo nivel del suelo. Esto no cambia quién es alto o bajo en relación con los demás, pero evita que los gigantes golpeen el techo de la pequeña habitación. Crucialmente, este paso es "equivalente en atención", lo que significa que cambia los números pero mantiene el enfoque de atención de la IA exactamente igual. Esto estrecha el rango de números, haciendo que el aplastamiento INT2 sea mucho menos doloroso.
la Aprendizaje de la Rotación Perfecta (El Paso "Consciente de la Salida"): En lugar de usar un giro fijo (como una rotación matemática estándar), OptR aprende un giro personalizado para cada "cabeza" (una parte específica del cerebro de la IA) en el modelo. Lo hace simulando todo el proceso: aplastar los datos, des-aplastarlos, ejecutar el mecanismo de atención de la IA y luego proyectarlo a la respuesta final.
- Observa el error en la respuesta final (el "error de salida de atención post-WO").
- Descompone este error en dos partes: errores causados por las claves (que determinan a qué presta atención la IA) y errores causados por los valores (que determinan qué información se recupera realmente).
- Luego, ajusta ligeramente los ángulos de rotación para minimizar el error en la respuesta final, no solo el error en el almacenamiento.

Piensa en ello como sintonizar una radio. Los métodos antiguos intentaban que el estática sonara lo más silenciosa posible dentro de la caja del altavoz. OptR escucha la música que sale del altavoz y ajusta las perillas hasta que la canción suena perfecta, incluso si la estática dentro de la caja no es perfectamente silenciosa.

Lo Que Encontraron

Los investigadores probaron OptR en tres modelos de IA diferentes (Qwen3-4B, Qwen3-8B y Phi4-14B) y cinco benchmarks desafiantes, incluyendo problemas matemáticos (AIME25), tareas de programación (LiveCodeBench) y recuperación de contexto largo (encontrar una aguja en un pajar de 64,000 palabras).

Los resultados fueron impactantes:

Ganancias Masivas de Precisión: En el modelo Qwen3-8B, usar la cuantización INT2 estándar con el mejor método anterior (QuaRot) resultó en una precisión de solo 17.33% en una prueba matemática difícil. Cuando añadieron OptR, la precisión se disparó al 66.67%. Eso es casi una mejora de cuatro veces.
Superando la Línea Base: Incluso comparado con el método de vanguardia (OSCAR), que ya era bastante bueno, OptR elevó la precisión del 54.67% al 66.00%.
Superpoder de Contexto Largo: El hallazgo más impresionante fue en las tareas de contexto largo. A medida que la historia se hacía más larga (hasta 64,000 tokens), los métodos estándar de INT2 fallaban estrepitosamente, cayendo a una precisión cercana a cero. OptR mantuvo la precisión de recuperación alta, cayendo solo ligeramente del 99.83% (a 4k tokens) al 70.02% (a 64k tokens), mientras que el método estándar se desplomó al 0.04%.
Sin Penalización de Velocidad: Los autores confirmaron que esta magia no viene con un costo pesado. Integraron OptR en el sistema y descubrieron que añade una sobrecarga insignificante. La velocidad de la IA (latencia) y la cantidad de datos que puede procesar (throughput) permanecieron casi idénticas a los métodos estándar. Es como obtener un motor supercargado sin añadir peso extra al coche.

Por Qué Esto Importa

El artículo sugiere que para que la cuantización de muy bajo bit (como INT2) funcione, no podemos mirar los datos de forma aislada. Tenemos que ver cómo fluyen esos datos a través de todo el sistema hasta la respuesta final. Al optimizar para la salida en lugar del almacenamiento, OptR cierra la brecha entre la compresión extrema y el alto rendimiento.

Los autores enfatizan que esto no es solo una victoria teórica; es una victoria práctica. Permite que los modelos de IA manejen conversaciones mucho más largas y grupos de usuarios más grandes sin quedarse sin memoria, todo esto manteniendo al modelo lo suficientemente inteligente como para resolver problemas matemáticos difíciles y escribir código. No solo sugirieron que esto podría funcionar; lo midieron en varios modelos y encontraron mejoras consistentes y significativas, demostando que la optimización "consciente de la salida" es la clave para desbloquear el potencial total de la IA de 2 bits.

Resumen Técnico: Rotación Consciente de la Salida para la Cuantización del KV-Cache en INT2

Declaración del Problema

A medida que los Modelos de Lenguaje de Gran Escala (LLMs) escalan en tamaño y longitud de contexto, el caché de Clave-Valor (KV) se ha convertido en un cuello de botella primario para la memoria y el ancho de banda durante la inferencia de contexto largo. Si bien la cuantización de ultra bajo bit (específicamente INT2) ofrece un camino para reducir el almacenamiento a 1/8 de los requisitos de BF16, introduce errores de cuantización significativos. En INT2, donde solo existen cuatro niveles, unos pocos valores atípicos extremos pueden expandir el rango de cuantización, causando grandes errores de redondeo para la mayoría de los valores.

Los métodos de rotación existentes (por ejemplo, QuaRot, OSCAR) intentan mitigar esto rotando los vectores del caché hacia una base favorable para la cuantización. Sin embargo, estos métodos optimizan las matrices de rotación basándose en objetivos indirectos (proxies) definidos antes de la lectura completa de la atención, como minimizar el error de reconstrucción del caché o igualar las estadísticas previas a la atención. El artículo argumenta que esto crea un desajuste de objetivos: la rotación que mejor reconstruye las claves y valores almacenados en caché no necesariamente minimiza el error propagado a través del mecanismo de atención y la matriz de proyección de salida ( $W_O$ ) hacia el estado oculto del modelo. En consecuencia, la predicción final del modelo puede sufrir a pesar de tener un caché bien reconstruido.

Metodología: OptR

Los autores proponen OptR, un método de rotación consciente de la salida diseñado para minimizar directamente el error de la salida de la atención post- $W_O$ . El método opera mediante un proceso de calibración offline de una sola vez donde los pesos del modelo permanecen congelados y solo se aprenden las correcciones de rotación.

1. Descomposición del Error en el Espacio de Salida

OptR formula el problema de cuantización en el espacio de salida. Descompone el error entre la salida de atención en BF16 y la salida cuantizada en INT2 en dos términos distintos:

Error inducido por la clave: Derivado de los errores de cuantización en las claves ( $\Delta k$ ) que perturban los logits de atención y la distribución softmax resultante ( $\Delta p$ ).
Error inducido por el valor: Derivado de los errores de cuantización en los valores ( $\Delta v$ ) agregados por los pesos de atención (potencialmente perturbados).

El objetivo de optimización es el error cuadrático de la contribución final de la salida ( $\Delta y$ ) tras la proyección por $W_O$ , en lugar del error de reconstrucción bruto del caché.

2. Reparametrización de Claves Equivalentes a la Atención

Antes de la rotación, OptR aplica un desplazamiento de media por canal (centrado) a las claves. Al restar la media de calibración $\mu$ de todas las claves, el método desplaza los logits de atención por una constante dependiente de la consulta (query). Crucialmente, esta operación deja la distribución de atención softmax sin cambios (equivalente a la atención) pero estrecha significamente el rango dinámico de las claves. Esto reduce el impacto de los valores atípicos sobre los limitados niveles de cuantización INT2 sin alterar el comportamiento del mecanismo de atención.

3. Correcciones Ortogonales por Cabeza

OptR aprende correcciones ortogonales específicas para cada cabeza KV para refinar una rotación base (por ejemplo, de Hadamard o basada en covarianza).

Parametrización: Las correcciones se aprenden como matrices no restringidas $A_K, A_V$ , que se convierten en generadores antisimétricos ( $S = A - A^\top$ ) y se exponentian para asegurar que las rotaciones resultantes sigan siendo ortogonales: $R = R_0 \exp(S)$ .
Calibración en Dos Etapas:
1. Calibración de Claves: La rotación de la clave se optimiza manteniendo los valores en BF16. La función de pérdida combina un término de divergencia KL (para preservar la distribución de atención) y un término que mide el error de la salida post- $W_O$ inducido por la cuantización de la clave.
2. Calibración de Valores: Una vez fijada la rotación de la clave, se optimiza la rotación del valor para minimizar el error de salida resultante de la cuantización del valor bajo la distribución de atención seleccionada.
Diferenciabilidad: Dado que el redondeo INT2 no es diferenciable, los gradientes se estiman utilizando un Estimador de Paso Directo (STE, Straight-Through Estimator).

4. Integración con el Sistema

OptR se integra en un pipeline de caché KV en INT2 basado en SGLang. Mantiene la compatibilidad con sistemas de caché KV paginados y kernels de decodificación fusionados al preservar la forma del tensor y el diseño regular de la memoria. La reparametrización de la clave y la rotación se implementan mediante kernels de Triton optimizados que fusionan operaciones (rotación, centrado, recorte, cuantización) para minimizar la sobrecarga.

Principales Contribuciones

Formulación de Optimización en el Espacio de Salida: El artículo redefine la cuantización del caché KV en INT2 como un problema de optimización dirigido al error de la salida de la atención post- $W_O$ , descomponiéndolo en términos inducidos por la clave y el valor.
Algoritmo OptR: Un nuevo método que combina la reparametrización de claves equivalente a la atención (para reducir los desplazos por canal) con correcciones ortogonales por cabeza aprendidas a través de la ruta completa de cuantización INT2 y de atención.
Eficiencia y Compatibilidad: El método mantiene el formato de caché KV paginado e introduce una sobrecarga de inferencia insignificante, lo que lo hace práctico para su despliegue en sistemas de servicio existentes.

Resultados Experimentales

Los autores evaluaron OptR en tres modelos (Qwen3-4B, Qwen3-8B y Phi4-14B) en cinco benchmarks (AIME24/25, GPQA-Diamond, MBPP+, LiveCodeBench v6) y recuperación de contexto largo (RULER-NIAH).

Precisión en Razonamiento y Codificación: OptR mejora significativamente el rendimiento respecto a los métodos de rotación base. Para Qwen3-8B en AIME25, la precisión aumentó de 17.33% (QuaRot-INT2) a 66.67% (QuaRot-INT2 + OptR), y de 54.67% (OSCAR) a 66.00% (OSCAR + OptR). Estos resultados se acercan a la línea base de BF16 de 68.00%.
Recuperación de Contexto Largo: OptR demuestra una robustez superior a medida que aumenta la longitud del contexto. En RULER-NIAH con un contexto de 64K, QuaRot-INT2 cayó a una precisión cercana a cero, mientras que QuaRot-INT2 + OptR mantuvo una precisión del 70.02%.
Análisis de Error: El análisis por capa muestra que OptR reduce tanto el error inmediato post- $W_O$ como el error residual propagado en el flujo residual en la mayoría de las capas y cabezas KV, no solo en un subconjunto.
Eficiencia: La integración de OptR añade una latencia insignificante (dentro del 2% del pipeline base) y una sobrecarga de memoria mínima (por ejemplo, +18 MiB para un lote de 128 con contexto de 2K/4K).

Significado

El artículo afirma que OptR aborda un desajuste fundamental en las estrategias actuales de cuantización de bajo bit: optimizar para la reconstrucción del caché no garantiza un rendimiento óptimo del modelo. Al optimizar directamente para el error observado por las capas subsiguientes (post- $W_O$ ), OptR permite que el caché KV de ultra bajo bit (INT2) logre un rendimiento comparable a BF16 en tareas complejas de razonamiento y de contexto largo. El método se presenta como una mejora práctica y de fácil implementación para los pipelines de rotación existentes que preserva la eficiencia del sistema mientras reduce drásticamente la degradación inducida por la cuantización.