Virtual Head Attention reduce la caché sin sacrificar la calidad del modelo

A medida que los modelos de lenguaje procesan contextos más largos, la caché de atención consume más memoria, y reducir las cabezas de clave-valor compartidas puede dañar la calidad del modelo. ¿Puede un diseño de atención reducir esta carga de memoria y el cómputo de entrenamiento sin perder precisión?

El artículo examina si transformaciones aprendidas pueden recuperar un comportamiento de atención útil con menos cabezas físicas.

Virtual Head Attention aprende diversidad de consultas y mezcla entre cabezas durante el entrenamiento, y luego integra ambas operaciones en la atención agrupada estándar para la inferencia.

Los autores preentrenaron cuatro escalas de modelos con arquitectura Qwen3 en 50B tokens y compararon VHA con la atención agrupada (grouped-query attention), usando los recuentos estándar de cabezas de clave-valor por escala.

En ocho pruebas de opción múltiple sin entrenamiento previo para esa tarea (zero-shot), la precisión promedio en cada escala fue de 0,05 a 0,81 puntos porcentuales mayor con VHA.

En 1,7B, VHA también usó 1,00 KB de caché de clave-valor en BF16 por capa por token frente a 4,00 KB de GQA-8; su cómputo de entrenamiento fue un 7,9 % menor en una longitud de secuencia de 4096.

En conjunto, estas mediciones respaldan un equilibrio entre calidad y eficiencia en estas pruebas, aunque los promedios de precisión no implican mejoras en cada prueba individual.

Las cifras de precisión en modelos pequeños son promedios de ocho pruebas, por lo que no muestran que VHA haya mejorado cada prueba individual.

Las comparaciones tampoco establecen que todos los ajustes de entrenamiento fueran idénticos; los resultados en modelos más grandes y en contextos largos provienen de estudios separados con sus propias configuraciones de modelo y entrenamiento.