Kronecker Product Attention comprime la KV-cache hasta 10 veces y alcanza a MLA en calidad

La memoria de la KV-cache es desde hace tiempo el principal cuello de botella de la generación autorregresiva: crece de forma lineal con la longitud del contexto, el tamaño del batch y la profundidad del modelo. Un nuevo preprint en arXiv propone el mecanismo Kronecker Product Attention (KPA), que, según sus autores, comprime esa caché hasta 10 veces sin perder calidad frente a MHA, MQA, GQA, MLA y TPA. La comparación se realizó en cuatro escalas de modelo y nueve benchmarks downstream, con las mismas condiciones de entrenamiento.

Cómo funciona KPA

La idea central: las claves y los valores no se almacenan completos, sino que se descomponen en una suma de productos de Kronecker de factores matriciales compactos. Los factores dependen del token, es decir, se calculan a partir de la entrada y no están fijados en los pesos. En la caché solo quedan esos factores, y los tensores completos se reconstruyen al calcular la atención.

Los enfoques existentes hacen concesiones distintas. MQA y GQA ahorran memoria al agrupar cabezas, a costa de perder expresividad. MLA comprime la caché mediante una representación latente, pero requiere una rama «desacoplada» aparte para la codificación posicional RoPE. TPA usa una factorización basada en el producto externo de vectores, es decir, una estructura rígida de rango uno. KPA, según los autores, generaliza TPA de forma estricta: el producto externo de vectores es un caso degenerado de la partición de Kronecker, mientras que los factores matriciales no degenerados pueden describir patrones de bloques más ricos.

Otro resultado técnico afecta a las posiciones. Los autores introducen un operador rotacional estructurado, coherente con la representación de Kronecker. Conserva la propiedad de posición relativa de RoPE, pero rota solo los factores compactos. Por eso no se necesita una rama adicional para la codificación posicional, como en MLA.

Qué afirma el artículo

• Compresión de la KV-cache: hasta 10x.
• Comparación con MHA, MQA, GQA, MLA y TPA: resultado igual o superior.
• Escalas: cuatro tamaños de modelo, nueve benchmarks downstream.
• Condiciones: entrenamiento idéntico para todas las arquitecturas.
• Codificación posicional: operador compatible con RoPE que actúa sobre los factores compactos, sin rama separada.
• Ajuste: la «partición de Kronecker de cabezas» (Kronecker head split) se define como hiperparámetro.

Este último punto es quizá el más práctico. La partición determina el equilibrio entre grado de compresión y precisión, y se puede modificar en la fase de diseño sin cambiar la arquitectura. El equipo no tiene que reescribir el modelo para pasar de una compresión más suave a otra más agresiva. Basta con reentrenar con otro valor.

Qué deben hacer los ingenieros y dónde conviene ser prudentes

Por ahora se trata de un preprint v1 sin verificación independiente. El resumen no indica ni los tamaños concretos de los modelos, ni las longitudes de contexto, ni la velocidad de decodificación en hardware real. Una mejora en memoria no siempre equivale a una mejora en latencia: reconstruir claves y valores a partir de los factores añade cómputo, y su coste depende de los kernels y de la GPU. El «hasta 10x» declarado corresponde al mejor régimen. En qué configuración se alcanza y a costa de cuánta calidad solo se sabrá con el texto completo y el código, si los autores lo publican.

Los equipos que construyen inferencia de contexto largo o atienden batches grandes deberían tener KPA en la lista de candidatos junto a MLA. Hay tres cosas que comprobar: cuánta memoria se libera con vuestras longitudes de contexto, cómo se comporta la calidad con distintas particiones y qué throughput se obtiene con vuestros kernels de atención. Para quienes ya han entrenado un modelo con GQA o MLA, la migración exigirá reentrenar, así que tiene más sentido incorporar KPA en el siguiente ciclo de pre-training que intentar sustituir el mecanismo en pesos ya entrenados.