A Faster Generalized Two-Stage Approximate Top-K

Résumé technique : Une approximation Top-K généralisée à deux étapes plus rapide

Énoncé du problème
L'identification des $K$ premiers éléments d'un tableau est une opération fondamentale dans de nombreux algorithmes d'apprentissage automatique, notamment la recherche du maximum de produit scalaire (MIPS), les $K$ plus proches voisins (KNN) et l'optimisation de grands modèles de base (par exemple, pour la parcimonie dans les logits de classification, les blocs MLP, les mécanismes d'attention et la compression du KVCache). Sur les accélérateurs modernes tels que les TPU et les GPU, le calcul exact du Top-$K$ constitue souvent un goulot d'étranglement de performance. Ces dispositifs sont optimisés pour les multiplications matricielles denses (MXU/TensorCores), tandis que les opérations Top-K reposent généralement sur des unités scalaires/vectorielles (VPU/CUDA cores) et peuvent être considérablement plus lentes que les multiplications matricielles qui génèrent les données. Par exemple, sur un TPUv5e, trouver les 2 % des activations cachées d'un modèle Gemma 2 9B prend 27 fois plus de temps que la multiplication matricielle correspondante.

Bien que les algorithmes de Top-K approximatifs offrent une solution de contournement, la méthode de l'état de l'art proposée par Chern et al. (2022), implémentée sous la forme de jax.lax.approx_max_k, entraîne toujours une latence significative. Leur approche à deux étapes partitionne l'entrée en buckets, sélectionne le premier élément de chaque bucket, puis trie les résultats. Cependant, la deuxième étape (le tri des éléments sélectionnés) reste un goulot d'étranglement, et la première étape sous-utilise souvent les ressources de calcul scalaires/vectorielles disponibles.

Méthodologie
Les auteurs proposent une généralisation de l'algorithme de Chern et al. (2022) qui modifie la première étape pour sélectionner les $K'$ premiers éléments de chaque bucket (où $1 \le K' \le K$) plutôt que de limiter la sélection au premier élément.

1. Généralisation algorithmique :

   • Étape 1 : Le tableau d'entrée de taille $N$ est partitionné en $B$ buckets. Au lieu de sélectionner uniquement le maximum, l'algorithme sélectionne les $K'$ premiers éléments de chaque bucket.
   • Étape 2 : Les $B \times K'$ éléments sélectionnés sont fusionnés et triés pour retourner les $K$ premiers éléments finaux.
   • Implémentation : La première étape est implémentée comme une sous-routine de mise à jour en ligne utilisant Pallas (un langage de noyau JAX pour les TPU). Elle maintient une liste triée des $K'$ premières valeurs et indices pour chaque bucket. Les nouveaux éléments sont insérés à l'aide d'un passage de type tri à bulles, optimisé pour la vectorisation à travers les buckets. Les auteurs fournissent à la fois une version « non fusionnée » (deux noyaux distincts) et une version « fusionnée avec matmul » où la première étape est fusionnée avec la multiplication matricielle précédente, masquant ainsi efficacement le coût de la première étape.

2. Analyse théorique :

   • Les auteurs dérivent une expression exacte pour le rappel attendu de cet algorithme généralisé sous un partitionnement aléatoire, en modélisant la distribution des éléments Top-$K$ dans les buckets à l'aide de la distribution hypergéométrique.
   • Ils démontrent que pour un objectif de rappel attendu fixe, l'augmentation de $K'$ permet de réduire le nombre de buckets $B$ de telle sorte que le nombre total d'éléments transmis à la deuxième étape ($B \times K'$) soit minimisé.
   • Crucialement, ils dérivent une nouvelle borne pour le nombre de buckets requis lorsque $K'=1$. Cette borne est prouvée comme étant plus stricte d'un facteur 2 par rapport à la borne de Chern et al. (2022), permettant d'utiliser moins de buckets et donc moins d'éléments à trier, même dans le cadre original.

3. Modélisation des performances :

   • En utilisant l'analyse du point de crête (ridge-point), les auteurs caractérisent les sous-systèmes de l'accélérateur (HBM, VPU, MXU). Ils démontrent que la première étape est généralement limitée par la mémoire jusqu'à ce que $K'$ atteigne un certain seuil (par exemple, $K' \approx 6$ sur TPUv5e). Dans ce régime limité par la mémoire, l'augmentation de $K'$ n'augmente pas significativement la latence de la première étape, mais réduit drastiquement la charge de travail de la deuxième étape.

Contributions clés

• Algorithme généralisé : Un nouvel algorithme de Top-K approximatif à deux étapes qui sélectionne les $K'$ premiers éléments par bucket, offrant un compromis flexible entre le calcul de la première étape et la charge de tri de la deuxième étape.
• Bornes théoriques :
  • Une expression exacte du rappel attendu en fonction de $N, K, B,$ et $K'$.
  • Une borne prouvée plus stricte (facteur 2) pour le nombre de buckets requis dans le cas $K'=1$ par rapport aux travaux antérieurs, améliorant la sélection des paramètres.
• Implémentation : Une implémentation haute performance sur Cloud TPUv5e utilisant Pallas, présentant des variantes non fusionnées et fusionnées avec matmul.
• Sélection automatique des paramètres : Une routine qui sélectionne automatiquement $K'$ et $B$ pour atteindre un objectif de rappel spécifié par l'utilisateur, éliminant le besoin de réglage manuel.

Résultats

• Efficacité théorique : Pour un objectif de rappel fixe (par exemple 95 %), l'augmentation de $K'$ réduit considérablement le nombre d'éléments traités dans la deuxième étape. Par exemple, sélectionner les 1024 premiers éléments parmi 262 144 éléments nécessite 16 384 éléments de sortie avec $K'=1$, mais seulement 2 048 éléments avec $K'=4$ (une réduction de 8x).
• Améliorations de la latence (non fusionné) : Sur TPUv5e, l'algorithme généralisé avec $K'=4$ réalise une accélération de 11x par rapport à l'original jax.lax.approx_max_k (27 µs contre 305 µs) pour un objectif de rappel de 99 % sur la même tâche.
• Améliorations de la latence (fusionné) : Dans une charge de travail de recherche du maximum de produit scalaire (MIPS) (1 million de vecteurs, 1024 requêtes), l'implémentation fusionnée avec matmul avec $K'=4$ réduit le temps d'exécution total de 594 ms (Top-K exact) à 10 ms. Cela représente une accélération globale de 6,7x par rapport à la base dense, plaçant l'étape de Top-K approximatif en dessous du coût de la multiplication matricielle elle-même.
• Impact de bout en bout : Dans un benchmark de bloc de transformateur parcimonieux (Gemma 2 9B), la méthode proposée réduit la surcharge du bloc MLP parcimonieux de 2,7 fois plus lent que le dense (en utilisant l'algorithme original) à une surcharge modeste de seulement +5 ms par rapport à la base dense.

Signification et affirmations
L'article affirme que le choix de $K'=1$ dans les algorithmes de Top-K approximatifs précédents n'exploite pas pleinement les ressources de calcul disponibles sur les accélérateurs modernes. En généralisant vers $K' > 1$, l'algorithme déplace efficacement le goulot d'étranglement loin de l'étape de tri sans engendrer de coût supplémentaire significatif dans l'étape de filtrage, à condition que la première étape reste limitée par la mémoire.

Les auteurs soulignent que leur analyse théorique fournit une base rigoureuse pour la sélection des paramètres, offrant une interface conviviale (approx_top_k(array, K, recall_target)) qui ne nécessite aucun réglage manuel. Bien que l'évaluation empirique soit centrée sur le TPUv5e, les auteurs affirment que l'algorithme, la modélisation des performances et l'analyse sont généraux et applicables à d'autres plateformes matérielles comme les GPU. Ce travail démontre qu'une généralisation algorithmique raisonnée, combinée à une implémentation consciente du matériel (fusion), peut éliminer le Top-K en tant que goulot d'étranglement dans l'entraînement et l'inférence de modèles à grande échelle.