Résumé Technique : StiefAttention

Énoncé du Problème

L'inférence en contexte long pour les grands modèles de langage (LLM) est de plus en plus limitée par la capacité de mémoire et les exigences de bande passante du cache Clé-Valeur (KV). Lors du décodage autorégressif, le cache KV croît linéairement avec la longueur de la séquence et la taille du lot (batch size), consommant souvent une mémoire comparable à celle des poids du modèle (par exemple, environ 16 Go pour Llama3-8B avec un contexte de 32k).

Les méthodes existantes de compression post-entraînement emploient généralement la décomposition de matrice de rang faible (telle que la décomposition en valeurs singulières, SVD) pour projeter les clés et les valeurs dans des dimensions inférieures. Cependant, ces méthodes optimisent des objectifs de substitution (proxy objectives), tels que la minimisation de l'erreur de reconstruction de Frobenius des vecteurs clés/valeurs bruts ou des scores d'attention $QK^\top$. Les auteurs soutiennent que ces objectifs de substitution ne parviennent pas à rendre compte de la manière dont les erreurs de reconstruction interagissent avec le softmax de l'attention, le mélange des valeurs (value mixing) et les transformations ultérieures des couches du décodeur (normalisation, voies résiduelles, et MLP). Par conséquent, une projection qui est optimale pour reconstruire les tenseurs intermédiaires peut néanmoins induire des erreurs significatives sur la sortie finale de la couche du décodeur, dégradant ainsi la qualité de la génération de bout en bout.

Méthodologie : StiefAttention

Le document présente StiefAttention, une méthode de compression du cache KV post-entraînement qui minimise directement l'erreur de reconstruction au niveau de la sortie de la couche du décodeur plutôt qu'au niveau des quantités d'attention intermédiaires.

Mécanisme Central

1. Fonction d'Objectif : Au lieu de minimiser $\|K - \tilde{K}\|_F$, StiefAttention minimise l'erreur relative entre la sortie originale de la couche du décodeur $f_\ell(x)$ et la sortie $\tilde{f}_\ell(x)$ produite par le mécanisme d'attention compressé. Cet objectif capture l'effet cumulatif du softmax, du mélange des valeurs et des non-linéarités en aval.
2. Projection Orthonormale sur la Variété de Stiefel : La méthode apprend des bases de projection orthonormées ($P_K, P_V$) situées sur la variété de Stiefel (matrices avec des colonnes orthonormées).
3. Prédiction de Base par Gradient :
   • Plutôt que de calculer la SVD sur l'ensemble complet du jeu de calibration (ce qui est gourmand en mémoire), StiefAttention utilise un prédicteur léger et entraînable (un MLP) pour générer des bases de projection à partir de statistiques d'activation simples.
   • Caractéristiques d'Entrée : Le prédicteur prend comme caractéristiques d'entrée la moyenne ($\mu$) et la variance ($\sigma^2$) par dimension des vecteurs clés et valeurs.
   • Entraînement : Le MLP est entraîné via la descente de gradient pour mapper ces statistiques vers des matrices $A_K$ et $A_V$.
   • Orthonormalisation : Les matrices de sortie sont orthonormalisées via la décomposition QR ($A = QR$). Le facteur orthogonal $Q$ sert de base de projection, car le facteur triangulaire $R$ ne représente qu'un changement de coordonnées au sein du sous-espace.
4. Allocation de Rang : La méthode construit des profils de rang par couche. Elle emploie une stratégie d'allocation de budget séquentielle où les rangs sont sélectionnés couche par couche afin de minimiser l'erreur cumulée de la sortie de la couche du décodeur sous un budget total de cache KV spécifié par l'utilisateur. Cela tient compte de la propagation de l'erreur entre les couches.
5. Inférence : Lors de l'inférence, les clés et les valeurs sont stockées sous forme compressée ($K\downarrow = K P_K$). Les dimensions complètes sont reconstruites à la volée ($\tilde{K} = K\downarrow P_K^\top$) avant le calcul de l'attention. Les bases de projection peuvent être intégrées (folded) dans les poids linéaires de l'attention, n'entraînant aucun surcoût de temps d'exécution par rapport aux méthodes standards de rang faible.

Contributions Clés

1. Optimisation Directe de la Sortie de la Couche du Décodeur : Les auteurs proposent d'optimiser la compression du cache KV directement contre l'erreur de sortie de la couche du décodeur, allant au-delà des objectifs de substitution intermédiaires comme la reconstruction de tenseur ou l'approximation des scores d'attention.
2. Cadre StiefAttention : Une méthode post-entraînement qui apprend des bases orthonormées à partir de statistiques d'activation en utilisant un prédicteur léger, permettant une compression efficace sans réentraîner l'ensemble du modèle.
3. Allocation de Rang Séquentielle : Une stratégie pour allouer les rangs à travers les couches en se basant sur des surfaces d'erreur locales évaluées sur des entrées qui incluent déjà les effets des compressions des couches précédentes.
4. Validation Empirique : Une évaluation complète sur Llama3-8B et Qwen3-8B montrant que StiefAttention surpasse les meilleures méthodes de référence (notamment EigenAttention) pour des budgets de cache KV équivalents.

Résultats Expérimentaux

Les auteurs ont évalué StiefAttention sur Llama3-8B et Qwen3-8B, en le comparant à la ligne de base non compressée FP16 et à EigenAttention (une méthode basée sur la SVD qui incorpore les requêtes dans la projection).

• Gains de Performance : Dans des conditions d'iso-compression (même budget de cache KV), StiefAttention a nettement surpassé EigenAttention :
  • Llama3-8B : A obtenu une amélioration de 4,2 points de la perplexité C4 et une augmentation de 8,9 points de la précision MMLU zero-shot par rapport à EigenAttention.
  • Qwen3-8B : A démontré des tendances similaires, StiefAttention maintenant une perplexité plus faible et une précision plus élevée, particulièrement lorsque les ratios de compression augmentent (compression plus forte).
• Robustesse : La méthode a montré une dégradation fluide des performances à mesure que le ratio de cache KV diminuait, sans effondrement brutal.
• Compatibilité avec la Quantification : StiefAttention s'est avéré substantiellement plus robuste à la quantification du cache KV en bas-bit (ex: 4-bit et 8-bit) que EigenAttention, maintenant une perplexité plus faible pour des tailles de cache comparables.

Analyse de la Signification

Le document fournit une analyse diagnostique expliquant pourquoi StiefAttention améliore la performance de bout en bout :

• Consistance Directionnelle : Bien qu'EigenAttention reconstruise mieux les sorties d'attention intermédiaires (réduisant l'erreur de sortie d'attention d'environ 29,6 % par rapport à StiefAttention), StiefAttention parvient à une meilleure préservation de la direction de la sortie de la couche du décodeur (similitude cosinus plus élevée, +3,3 %) et à une erreur de sortie de couche globale plus faible (amélioration de 5,2 %).
• Sensibilité des Couches Précoces : L'analyse révèle que l'écart de similitude cosinus est plus prononcé dans les premières couches. Les auteurs avancent que les perturbations dans les couches précoces sont amplifiées par les blocs non linéaires suivants (softmax, voies résiduelles, MLP). Les méthodes de type SVD, qui optimisent la reconstruction intermédiaire, échouent souvent à préserver l'information directionnelle spécifique requise par ces blocs en aval. En optimisant directement pour la sortie du décodeur, StiefAttention aligne les bases de projection avec le "sous-espace pertinent pour la sortie", conduisant à une meilleure qualité de génération finale.

Limites et Travaux Futurs

Les auteurs reconnaissent plusieurs limites :

• Portée : Les expériences ont été principalement menées à une longueur de séquence de 2048 en raison des contraintes de mémoire GPU ; une évaluation plus large sur des contextes plus longs est nécessaire.
• Simplification de l'Optimisation : La sélection du rang utilise une trajectoire séquentielle gloutonne plutôt qu'une optimisation globale exhaustive.
• Représentation des Caractéristiques : Le prédicteur n'utilise actuellement que la moyenne et la variance diagonale ; des résumés d'activation plus riches (ex: esquisses de covariance) sont laissés pour des travaux futurs.
• Évaluation Système : Une évaluation complète au niveau système, incluant la latence de prefill/decode et le débit dans des conditions de service réalistes, est en attente.

Le document conclut qu'optimiser directement pour la sortie finale de la couche du décodeur est une stratégie plus efficace pour la compression du cache KV que l'optimisation d'objectifs de substitution intermédiaires, offrant une voie pratique pour réduire l'empreinte mémoire sans sacrifier la qualité de la génération.