PACE compresse l'état des LLM linéaires à 4 bits et économise un tiers de la mémoire

Une équipe de chercheurs propose PACE, un quantificateur post-entraînement qui stocke l'état récurrent des modèles à attention linéaire sur quatre bits et ne nécessite pas de fine-tuning. Sur quatre modèles hybrides de 4B à 35B paramètres, la variante mixte, avec des têtes en 4 et 8 bits, se situe au niveau du DAMP 8 bits sur les trois plus grands modèles, tout en occupant environ deux tiers de sa mémoire. La version où toutes les têtes sont en quatre bits obtient de meilleurs résultats que Q-Mamba à mémoire égale.

Pourquoi l'état devient un goulot d'étranglement

Dans les modèles à attention linéaire, le contexte ne vit pas dans un KV-cache qui grossit, mais dans un état récurrent de taille fixe. Sur le papier, c'est pratique. Mais le système d'inférence doit conserver une copie distincte de cet état pour chaque séquence décodée en parallèle. Plus le nombre de copies qui tiennent dans la mémoire de l'accélérateur est élevé, plus le nombre de requêtes servies simultanément augmente, et plus chaque token coûte moins cher.

D'où l'intérêt de quantifier l'état lui-même. Les méthodes existantes l'arrondissent après chaque token généré. L'erreur d'arrondi s'accumule alors pas à pas, et à quatre bits la précision chute nettement. Les travaux antérieurs aboutissaient tous à la même conclusion : soit on augmente le nombre de bits, soit on réentraîne le modèle en vue de la quantification. PACE montre qu'il existe une troisième voie.

Deux observations à la base de PACE

La première concerne la fréquence d'arrondi. PACE n'arrondit l'état que tous les quelques tokens et conserve les tokens intermédiaires dans un petit buffer. Les noyaux chunked fonctionnent déjà de manière similaire lors du prefill : ils mettent à jour l'état une fois par chunk, et non à chaque token. Moins d'arrondis, c'est moins d'erreur accumulée.

La seconde observation porte sur la structure de l'état dans les couches gated-delta. Les grandes valeurs n'y sont pas réparties uniformément, mais regroupées dans des colonnes. La raison : certains canaux de la mise à jour écrite à chaque token restent élevés pendant de nombreux tokens. PACE utilise donc une échelle de quantification commune pour toute une colonne. Une seule grande valeur ne dégrade plus la précision de toutes les autres, qui devaient auparavant partager son échelle.

Troisième élément : la précision mixte. Certaines têtes conservent l'erreur d'arrondi plus longtemps que d'autres. PACE garde précisément celles-ci en huit bits et fait passer les autres à quatre.

• Méthode : post-entraînement, sans fine-tuning du modèle
• Format cible : quatre bits par élément de l'état, huit pour les têtes dont l'erreur persiste le plus longtemps
• Arrondi : tous les quelques tokens, les tokens intermédiaires étant stockés dans un buffer
• Échelle de quantification : commune à chaque colonne de l'état
• Validation : quatre modèles hybrides de 4B à 35B paramètres

Ce que montrent les mesures

Le PACE mixte avec des têtes en 4 et 8 bits se situe au niveau du DAMP 8 bits sur trois des quatre modèles, les plus grands. La mémoire occupée représente alors environ deux tiers de celle du DAMP. La variante purement 4 bits l'emporte sur Q-Mamba à mémoire égale. Les auteurs soulignent par ailleurs que l'arrondi tous les quelques tokens préserve l'accélération du décodage apportée par le buffer : l'économie de mémoire ne se paie pas par une perte de vitesse.

Ce dernier point compte en pratique. Une quantification de l'état qui ajoute une opération d'arrondi à chaque étape peut annuler le gain obtenu en densifiant le batch. Ici, l'arrondi sort de la boucle critique, ce qui réduit ce risque.

Où cela peut servir côté infrastructure

Les architectures hybrides, qui alternent couches linéaires et attention classique, apparaissent de plus en plus souvent dans les modèles ouverts et commerciaux. Pour ceux qui les servent, PACE offre un levier clair : réduire d'environ un tiers la mémoire consacrée à l'état par rapport à un schéma 8 bits, sans réentraînement ni fine-tuning. La place libérée peut être allouée à un batch plus grand ou à une file plus longue de sessions simultanées.

Il y a aussi des limites. Le résultat est validé sur quatre modèles hybrides, et sur le plus petit d'entre eux, le PACE mixte n'a pas atteint le niveau du DAMP. Les équipes qui envisagent un déploiement devraient vérifier le comportement de la méthode sur leur propre modèle et leur charge avant de remanier la production pour un état en quatre bits. Mais l'idée elle-même, arrondir moins souvent et ajuster l'échelle à la structure réelle des données, semble transposable et peu coûteuse à mettre en œuvre.