Kronecker Product Attention compresse le cache KV jusqu'à 10 fois et rattrape MLA en qualité

La mémoire consommée par le cache KV est depuis longtemps le principal goulot d'étranglement de la génération autorégressive : elle croît linéairement avec la longueur du contexte, la taille du batch et la profondeur du modèle. Un nouveau préprint sur arXiv propose un mécanisme, Kronecker Product Attention (KPA), qui, selon ses auteurs, compresse ce cache jusqu'à 10 fois sans perdre en qualité face à MHA, MQA, GQA, MLA et TPA. La comparaison porte sur quatre tailles de modèles et neuf benchmarks downstream, avec des conditions d'entraînement identiques.

Comment fonctionne KPA

L'idée centrale : les clés et les valeurs ne sont pas stockées en entier, mais décomposées en une somme de produits de Kronecker de petits facteurs matriciels. Ces facteurs dépendent du token : ils sont calculés à partir de l'entrée et non figés dans les poids. Seuls ces facteurs restent dans le cache, et les tenseurs complets sont reconstruits au moment du calcul de l'attention.

Les approches existantes font chacune des compromis différents. MQA et GQA économisent de la mémoire en regroupant les têtes, au prix d'une perte d'expressivité. MLA compresse le cache via une représentation latente, mais exige une branche « découplée » distincte pour l'encodage positionnel RoPE. TPA recourt à une factorisation par produit extérieur de vecteurs, c'est-à-dire à une structure rigide de rang un. KPA, affirment les auteurs, généralise strictement TPA : le produit extérieur de vecteurs en est un cas dégénéré de la décomposition de Kronecker, tandis que les facteurs matriciels non dégénérés peuvent décrire des motifs par blocs plus riches.

Un autre résultat technique concerne les positions. Les auteurs introduisent un opérateur de rotation structuré, compatible avec la représentation de Kronecker. Il conserve la propriété de position relative de RoPE, mais ne fait tourner que les facteurs compacts. Une branche supplémentaire pour l'encodage positionnel, comme dans MLA, n'est donc pas nécessaire.

Ce que revendique l'article

• Compression du cache KV : jusqu'à 10x.
• Comparaison avec MHA, MQA, GQA, MLA et TPA : résultats équivalents ou supérieurs.
• Échelle : quatre tailles de modèles, neuf benchmarks downstream.
• Conditions : entraînement identique pour toutes les architectures.
• Encodage positionnel : opérateur compatible RoPE appliqué aux facteurs compacts, sans branche séparée.
• Réglage : la « partition de Kronecker des têtes » (Kronecker head split) est définie comme un hyperparamètre.

Ce dernier point est peut-être le plus pratique. La partition détermine le compromis entre taux de compression et précision, et on peut la modifier à la conception sans changer d'architecture. L'équipe n'a pas à réécrire le modèle pour passer d'une compression plus douce à une compression plus agressive. Il suffit de réentraîner avec une autre valeur.

Ce que doivent faire les ingénieurs, et où rester prudent

Pour l'instant, il s'agit d'un préprint v1, sans vérification indépendante. L'abstract n'indique ni les tailles précises des modèles, ni les longueurs de contexte, ni la vitesse de décodage sur du matériel réel. Un gain de mémoire n'équivaut pas toujours à un gain de latence : la reconstruction des clés et des valeurs à partir des facteurs ajoute du calcul, dont le coût dépend des kernels et du GPU. Le « jusqu'à 10x » annoncé correspond au meilleur régime ; la configuration qui permet de l'atteindre, et le prix à payer en qualité, ne ressortiront que du texte complet et du code, si les auteurs le publient.

Les équipes qui construisent de l'inférence à long contexte ou servent de gros batchs devraient garder KPA dans leur liste de candidats, aux côtés de MLA. Trois points sont à vérifier : combien de mémoire est libérée à vos longueurs de contexte, comment la qualité évolue selon les partitions, et quel débit on obtient avec vos kernels d'attention. Pour ceux qui ont déjà entraîné un modèle avec GQA ou MLA, la transition exigera un réentraînement : mieux vaut donc intégrer KPA au prochain cycle de pre-training que tenter de remplacer le mécanisme dans des poids déjà entraînés.