Attention à produit de Kronecker (KPA) : clés et valeurs en sommes de produits de Kronecker de facteurs compacts, cache KV compressé jusqu'à 10 fois, sans branche rotative séparée

À l'inférence, un goulot d'étranglement de plus en plus visible pour les grands modèles est le cache KV. Pour éviter de recalculer, le modèle stocke les clés et valeurs de chaque jeton à chaque couche, et ce cache croît linéairement avec la longueur de séquence, la taille de lot et la profondeur. Quand les contextes atteignent des centaines de milliers de jetons et que les lots grandissent pour augmenter le débit, le cache KV peut occuper l'essentiel de la mémoire du GPU, laissant peu de place aux poids et aux activations et limitant la longueur de contexte, le lot et l'échelle de déploiement que l'on peut se permettre. Les remèdes existants font chacun un compromis. L'attention multi-requête (MQA) et l'attention par requêtes groupées (GQA) partagent clés et valeurs entre têtes au prix de la diversité de représentation. L'attention latente multi-têtes (MLA) met en cache des états latents compressés mais exige un mécanisme d'encodage de position découplé supplémentaire. L'attention à produit tensoriel (TPA) factorise par des produits extérieurs de vecteurs de faible rang et peut être limitée par une structure de rang trop stricte. L'article 2610.12657 propose l'attention à produit de Kronecker (KPA) dans ce contexte.

L'idée centrale de KPA est de ne pas mettre en cache les matrices complètes de clés et de valeurs, mais d'écrire les clés et valeurs de chaque jeton comme des sommes de produits de Kronecker de facteurs matriciels compacts qui dépendent du jeton. Un produit de Kronecker construit une grande matrice à partir de deux petites, il peut donc placer de la structure à la fois au niveau des têtes et au niveau des caractéristiques à l'intérieur d'une tête, alors que le produit extérieur de vecteurs de TPA ne couvre que l'un des deux. L'article note que des études empiriques montrent que les têtes d'attention se spécialisent dans des rôles différents, comme des motifs syntaxiques, sémantiques et positionnels, si bien qu'une bonne factorisation du cache KV devrait travailler aux deux niveaux à la fois. Il montre que la décomposition en produit extérieur de rang un de TPA apparaît comme une « division de Kronecker dégénérée » de KPA, donc que KPA généralise strictement TPA, et que des facteurs matriciels non dégénérés peuvent représenter des motifs en blocs plus riches. Seuls les facteurs sont mis en cache au décodage, et le taux de compression dépend de la façon dont on découpe la dimension de tête : plus les facteurs sont asymétriques, plus la compression est forte. Cette « division de têtes de Kronecker » est un hyperparamètre que l'on règle à la conception pour se déplacer le long du compromis compression-précision, sans autre changement d'architecture.

L'encodage de position est l'autre point clé. RoPE agit sur la clé complète, mais KPA ne met en cache que des facteurs, donc l'article doit répondre à la question de savoir où va la rotation. La réponse de MLA est une branche RoPE découplée séparée. L'article donne un opérateur de rotation structuré adapté à la représentation de Kronecker. Il préserve la propriété de position relative de RoPE et ne fait tourner que les facteurs compacts, si bien que la branche découplée n'est plus nécessaire. L'article dit que KPA dépasse MLA sans exiger une telle branche.

Côté expériences, les auteurs entraînent des modèles de style Llama à quatre tailles, environ 125 M, 353 M, 773 M et 1,5 Md de paramètres, avec SwiGLU et RoPE au-dessus de nanoGPT, sur environ 50 Md de jetons de FineWeb-Edu, avec un rang RK=RV=2 pour clés et valeurs et RQ=6 pour les requêtes. Ils évaluent neuf benchmarks zéro-coup avec lm-evaluation-harness : ARC-Easy, ARC-Challenge, BoolQ, HellaSwag, OpenbookQA, PIQA, WinoGrande, MMLU et SciQ. L'article dit qu'à conditions d'entraînement identiques, KPA égale ou dépasse TPA, MHA, MQA, GQA et MLA tout en atteignant jusqu'à 10 fois de compression du cache KV. À une taille, KPA obtient la meilleure moyenne à 6,4 fois de compression, le même rapport que TPA, avec 1,49 point de précision en plus, et le plus gros gain vient d'ARC-Easy. Une autre configuration atteint 8 fois. À environ 6,7 fois de compression, il dépasse les autres bases, avec des gains nets sur ARC-Easy, BoolQ et SciQ. À 1,5 Md, il mène sur huit des neuf benchmarks, et à 125 M toutes les configurations KPA dépassent les 46,21 % de TPA. Les divisions asymétriques avec un facteur plus grand au niveau des têtes battent à la fois les divisions symétriques et TPA.

Comment le juger ? Ses forces sont une idée claire, la preuve d'une généralisation stricte, une réponse à la coexistence de la compression et de l'encodage de position rotatif, et un taux de compression réglable. Les réserves sont tout aussi claires. D'abord, tous les modèles sont petits et entraînés de zéro par les auteurs, jusqu'à 1,5 Md de paramètres et environ 50 Md de jetons, et ne sont pas appliqués à un grand modèle existant. Ensuite, l'évaluation porte surtout sur la précision zéro-coup. L'article dit lui-même en conclusion que construire un noyau de décodage factorisé efficace et tester KPA sur des contextes plus longs sont les prochaines étapes naturelles, ce qui revient à dire que la vitesse de décodage, la latence et le comportement à long contexte n'ont pas encore de données. Enfin, les résultats sont auto-déclarés et nous n'avons vu aucune reproduction indépendante. Si votre souci est le coût du cache KV, voyez-y une conception à suivre et jugez s'il peut remplacer GQA ou MLA quand un noyau et des résultats à long contexte existeront.