Virtual Head Attention réduit le cache tout en préservant la qualité du modèle

À mesure que les modèles de langage traitent des contextes plus longs, leur cache d'attention consomme davantage de mémoire, et réduire les têtes clé-valeur partagées peut nuire à la qualité du modèle. Une conception de l'attention peut-elle alléger cette charge mémoire et le coût de calcul à l'entraînement sans perdre en précision ?

L'article teste si des transformations apprises peuvent restaurer un comportement d'attention utile avec moins de têtes physiques.

Virtual Head Attention apprend la diversité des requêtes et le mélange entre têtes pendant l'entraînement, puis intègre ces deux opérations dans une attention groupée standard pour l'inférence.

Les auteurs ont pré-entraîné quatre échelles de modèles d'architecture Qwen3 sur 50 milliards de tokens et ont comparé VHA à l'attention groupée (grouped-query attention) en utilisant les nombres standard de têtes clé-valeur par échelle.

Sur huit tests à choix multiples en zéro-shot, la précision moyenne à chaque échelle était supérieure de 0,05 à 0,81 point de pourcentage avec VHA.

À 1,7 milliard de paramètres, VHA utilisait également 1,00 Ko de cache clé-valeur BF16 par couche par token, contre 4,00 Ko pour GQA-8 ; son coût de calcul à l'entraînement était inférieur de 7,9 % à une longueur de séquence de 4 096.

Ensemble, ces mesures appuient un équilibre entre qualité et efficacité dans ces tests, tandis que les moyennes de précision n'impliquent pas de gains sur chaque test individuellement.

Les chiffres de précision sur les petits modèles sont des moyennes sur huit tests ; ils ne montrent donc pas que VHA a amélioré chaque test individuellement.

Les comparaisons n'établissent pas non plus que tous les paramètres d'entraînement étaient identiques ; les résultats sur les modèles plus grands et en contexte long proviennent d'études distinctes, avec leurs propres configurations de modèles et d'entraînement.