Résumé Technique : ProxyKV

Énoncé du Problème

L'inférence efficace en contexte long dans les Modèles de Langage à Grande Échelle (LLM) est actuellement limitée par les exigences mémoire du cache Key-Value (KV). À mesure que la longueur de séquence ($N$) augmente, la croissance linéaire du cache KV dépasse souvent la capacité de la mémoire GPU, entraînant des surcoûts I/O prohibitifs et une latence d'inférence accrue.

Les méthodes existantes de pruning de cache KV font face à un compromis fondamental entre l'efficacité du scoring et la précision du pruning :

• Méthodes Heuristiques (par ex., H2O, SnapKV) : Elles sont rapides et à faible latence, mais reposent sur des fenêtres d'observation locales, manquant souvent les dépendances sémantiques globales requises pour le raisonnement complexe.
• Méthodes Basées sur la Reconstruction (par ex., KVZip) : Elles atteignent une haute précision en effectuant un second passage de préremplissage forcé par un enseignant sur le modèle cible pour reconstruire le contexte global. Cependant, cela engendre une surcharge de préremplissage prohibitive, annulant les bénéfices de latence du pruning.

Le défi central consiste à combler le fossé entre les heuristiques à faible latence et la reconstruction à haute précision sans encourir le coût de calcul lourd d'un second préremplissage sur le modèle cible.

Méthodologie

Les auteurs proposent ProxyKV, un cadre de pruning par proxy inter-modèle conçu pour délester la tâche de scoring de l'importance du modèle cible vers un « Petit-Modèle Proxy » léger et asynchrone.

Architecture Système

Le système fonctionne via un pipeline à trois composants :

1. Modèle Cible Grand ($M_l$) : Le modèle principal responsable du chemin critique de préremplissage et de l'ensemble de la phase de décodage.
2. Petit-Modèle Proxy ($M_s$) : Un modèle plus petit de la même famille (par ex., 1,5 ou 4 milliards de paramètres) qui s'exécute de manière asynchrone sur un flux d'exécution séparé (ou flux CUDA). Il extrait des caractéristiques d'attention inter-têtes ($X$) du contexte d'entrée.
3. HybridAxialMapper ($\Phi_\theta$) : Un module apprenable qui mappe les caractéristiques du proxy ($X$) vers des scores d'importance alignés sur la cible ($\hat{Y}$).

Flux de travail :
Pendant l'inférence, le contexte d'entrée est dispatché simultanément vers les deux modèles. Le proxy génère des caractéristiques d'attention, qui sont traitées par le HybridAxialMapper pour produire un masque de pruning binaire ($M$) par seuillage Top-K. Le modèle cible effectue ensuite le décodage en n'utilisant que le cache KV pruning. Crucialement, le cache KV du proxy est libéré immédiatement après la phase de préremplissage, rendant la surcharge mémoire transitoire.

Composants Clés

1. Architecture HybridAxialMapper

Pour combler l'écart architectural entre des modèles hétérogènes (différents nombres de têtes et profondeurs de couches), le mapper emploie un pipeline à trois étapes qui dissocie l'extraction de caractéristiques temporelles de l'alignement inter-têtes :

• Étape 1 : Extraction de Caractéristiques Temporelles : Un Stem Convolutif 1D projette les caractéristiques brutes du proxy dans un espace latent agnostique au modèle, normalisant la configuration des têtes.
• Étape 2 : Encodeur de Contexte Axe-Temps : Un encodeur Transformer traite la séquence avec des encodages de position sinusoïdaux pour capturer le contexte global, utilisant une approche de fenêtre glissante pour les longues séquences (jusqu'à 170k tokens).
• Étape 3 : Attention Inter-Têtes Axe-Têtes : Des requêtes cibles apprenables agrègent des tokens de tête synthétiques du proxy vers la configuration spécifique des têtes de la cible. Cette étape permet au mapper de s'adapter à des nombres de têtes cibles arbitraires sans réentraîner l'ensemble du réseau.

2. Perte Hybride Multi-Granularité

Reconnaissant que la métrique de déploiement est la qualité de rétention Top-K plutôt que la reconstruction de score point par point, les auteurs conçoivent une fonction de perte composite ($L_{total}$) composée de cinq termes :

• $L_{bin}$ (Binaire Multi-Ratio) : Supervise la frontière de décision binaire à travers plusieurs ratios de rétention, mettant l'accent sur une compression agressive.
• $L_{mse}$ (MSE Pondérée par la Valeur) : Assure l'alignement de la magnitude numérique pour les tokens saillants.
• $L_{fine}$ (Rang Intra-Top-K) : Impose l'ordre relatif au sein de l'ensemble Top-K retenu.
• $L_{global}$ (Rang Global Sensible à la Masse) : Optimise la séparation entre les tokens Top-K et non-Top-K (frontière signal/bruit).
• $L_{cos}$ (Alignement Cosinus) : Aligner l'orientation globale du vecteur de score prédit avec celle de la cible.

Cette conception de perte déplace l'objectif d'une régression rigide vers la cohérence du classement relatif, ce qui est plus directement aligné avec la qualité du pruning.

Contributions Clés

1. Cadre Proxy Asynchrone : Un cadre de pruning novateur qui élimine le goulot d'étranglement du scoring du chemin critique de la cible en le délestant vers un petit modèle asynchrone, réalisant jusqu'à 3,21× d'accélération de préremplissage sur Llama-3.1-8B (double GPU) et maintenant des accélérations jusqu'à 170k tokens.
2. HybridAxialMapper & Perte Hybride : Une architecture et une fonction de perte apprenables qui comblent avec succès les écarts structurels entre des modèles hétérogènes (par ex., cible 8B vs proxy 1,5B), surmontant la rigidité de l'alignement tête-à-tête statique utilisé dans les travaux antérieurs.
3. Validation Empirique : Des expériences extensives sur les familles Llama-3.1, Qwen-2.5 et Qwen-3 (cibles 7B–32B) démontrant que ProxyKV récupère ~98,7% de la précision moyenne de l'oracle KVZip tout en réduisant significativement la latence.

Résultats

• Précision : ProxyKV correspond aux performances de l'oracle KVZip (la référence haute précision) sur les benchmarks LongBench, SCBench et RULER. Il récupère ~98,7% de la précision moyenne de KVZip et suit à moins de 1–2 points sur la plupart des tâches, même avec des ratios de rétention agressifs ($\rho \le 0,2$).
• Efficacité :
  • Sur Llama-3.1-8B, ProxyKV offre une accélération de 3,21× de la latence de préremplissage par rapport à KVZip (double GPU) et ~1,5× sur un GPU partagé unique.
  • Le HybridAxialMapper lui-même ne contribue que pour 4,7%–6,8% du temps mural de préremplissage.
  • La méthode s'adapte efficacement aux cibles 32B (Qwen-3) avec un ratio de taille cible/proxy d'environ 8×.
• Robustesse : Le cadre démontre une généralisation zero-shot robuste vers des jeux de données retenus (par ex., SCBench RepoQA) et maintient les performances sur des tâches complexes de raisonnement multi-sauts où les méthodes heuristiques (SnapKV) échouent.

Importance et Revendications

L'article revendique que ProxyKV comble avec succès le fossé « coût de scoring–précision » qui a contraint l'inférence LLM en contexte long. En exploitant les corrélations d'attention intra-famille et une cartographie inter-axiale apprenable, il atteint la haute précision des méthodes basées sur la reconstruction sans la surcharge de préremplissage prohibitive.

Les auteurs soulignent que tandis que l'alignement statique couche-à-couche (tel qu'utilisé dans IAM ou KVZip+IAM) devient une contrainte contraignante à mesure que les modèles cibles s'agrandissent (par ex., jusqu'à 32B), le HybridAxialMapper apprenable reste efficace. Cela suggère que ProxyKV offre une voie évolutive pour une inférence en contexte long efficace, récupérant une précision proche de l'oracle tout en réduisant significativement la latence et la pression mémoire pendant la phase de préremplissage. La méthode est présentée comme une solution au moment du déploiement qui ne nécessite qu'un proxy léger intra-famille et un petit mapper entraînable, évitant le besoin de réentraînement complet du modèle ou de modifications matérielles complexes.