NVIDIA et l’Université d’Édimbourg Révolutionnent le Cache KV des Modèles de Langue avec la Dynamique de Sparsification Mémoire (DMS)

Résumé du papier sur la Dynamique de Compression en Mémoire (DMS) pour une densification de 8x des caches KV dans les grands modèles linguistiques à base de Transformer (LLMs) Contexte et Problématique Les grands modèles linguistiques (LLMs) comme GPT, LLaMA, et Qwen sont de plus en plus utilisés pour des tâches requérant une forte capacité de raisonnement, ce qui implique la génération de séquences textuelles plus longues ou des chaînes de raisonnement parallèles. Toutefois, la performance pendant l'inférence est fortement limitée par la taille du cache key-value (KV). Ce cache stocke les représentations des tokens passés pour la génération auto-régressive, et sa taille croît linéairement avec la longueur de la séquence et la largeur des threads parallèles, consommant ainsi beaucoup de mémoire GPU. Cette consommation entraîne des ralentissements de l'inférence du fait des accès fréquents à la mémoire. Approches Existantes et leurs Limites Les techniques actuelles de compression du cache KV utilisent soit des heuristiques sans apprentissage, comme l'élimination de tokens basée sur les poids d'attention, soit des méthodes nécessitant une refonte lourde après l'apprentissage, comme la compression en mémoire dynamique (DMC). Les premières peuvent altérer la précision du modèle, tandis que les secondes sont coûteuses en termes de calcul. Introduction de DMS : Compression Avec Flexibilité Les chercheurs de NVIDIA et de l'Université d'Édimbourg ont proposé la Dynamique de Compression en Mémoire (DMS), une méthode robuste permettant de comprimer le cache KV tout en préservant la précision du modèle. DMS combine les avantages de la raréfaction traditionnelle avec un entrainement minimaliste et un système d'élimination retardée des tokens. Concept principal de DMS : - Pronostic d'élimination différentiable : Utilisation d'un mécanisme de prédiction basé sur la sigmoïde de Gumbel pour rendre la décision d'élimination des tokens différentiable lors de l'entrainement. - Élimination retardée des tokens : Les tokens marqués pour élimination demeurent utilisables durant une fenêtre de temps glissante avant d'être définitivement supprimés, permettant au modèle d'absorber mieux leur valeur informationnelle. Avantages de DMS Contrairement à DMC, DMS n'introduit pas de paramètres additionnels par tête d'attention et réutilise une petite partie du mécanisme d'attention (un seul neurone) pour prédire l'élimination. Cela facilite le réaménagement de modèles existants sans nécessiter de modifications architecturales importantes. Résultats Empiriques Les expérimentations sur différents modèles et tâches de raisonnement montrent que DMS peut atteindre une compression du cache KV de 8x tout en préservant voire améliorant les performances du modèle. Par exemple : Tâches de raisonnement : Amélioration de 9.1 points sur AIME, 7.6 sur GPQA, et 9.6 sur LiveCodeBench pour les modèles Qwen-R1 1.5B, 7B et 32B. Comparaisons avec les baselines : Meilleure efficacité de lecture de cache KV et utilisation optimale de la mémoire, surpassant des modèles tels que Quest et TOVA. Tâches de contexte court : Préservation des performances avec des ratios de compression jusqu'à 4x et une dégradation minimale (~3.5 points), testés sur MMLU, GSM8K et HellaSwag. Tâches de contexte long : Surpassement des modèles de base, notamment sur Needle-in-a-Haystack et Variable Tracking, suggérant une capacité à réduire les problèmes de trop forte compression de l'information dans les séquences longues. Conclusion La Dynamique de Compression en Mémoire (DMS) offre une solution pratique et évolutive pour améliorer l'efficacité de l'inférence des modèles linguistiques basés sur Transformer. En comprimant intelligentement le cache KV avec un entrainement minimum, DMS permet aux modèles de traiter des séquences plus longues ou en parallèle sans augmenter le temps de calcul ou la demande en mémoire. Ses performances constantes dans une variété de tâches, tant en raison de raisonnement qu'en contexte général, témoignent de sa polyvalence et de son efficacité. DMS pourrait donc jouer un rôle crucial dans l'optimisation des LLMs pour les environnements à ressources limitées, en équilibrant compression, précision et facilité d'intégration. Évaluation des Professionnels et Profil de NVIDIA Les experts de l'industrie saluent DMS pour sa capacité à améliorer significativement les performances inférentielles tout en minimisant l'impact sur la précision. NVIDIA, leader mondial en solutions de calcul parallèle et en technologies de l'intelligence artificielle, continue d'innover pour rendre les LLMs plus accessibles et efficaces. Ce travail reflète l'engagement de NVIDIA dans la recherche et le développement de solutions de pointe pour l'IA, soutenant ainsi les exigences croissantes des applications modernes.