🤖 AI

Fast LapSum: Exact Differentiable Top-k at Million Scale

L'article introduit Fast LapSum, une primitive de top- $k$ doux, exacte et différentiable, qui préserve une masse de sélection précise de $k$ tout en s'exécutant en temps linéaire sur GPU, permettant un calcul creux efficace à l'échelle du million pour des applications telles que la génération d'exemples adverses et le codage d'images différentiable.

Auteurs originaux : Łukasz Struski, Joanna Wojciechowicz, Jakub Antczak, Marcin Mazur, Kamil Książek, Jacek Tabor

Publié 2026-08-10

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Łukasz Struski, Joanna Wojciechowicz, Jakub Antczak, Marcin Mazur, Kamil Książek, Jacek Tabor

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous dirigez une bibliothèque numérique massive où des millions de livres sont numérisés chaque seconde. Pour donner un sens à ce déluge d'informations, l'IA de la bibliothèque doit décider quels quelques livres sont les plus importants à lire en ce moment. Dans le monde de l'intelligence artificielle, cela s'appelle la « sélection top-k » : choisir les k meilleurs éléments parmi une liste immense. Habituellement, l'IA fait cela en étant un bibliothécaire strict qui choisit les meilleurs livres et ignore complètement les autres. C'est excellent pour la vitesse, mais c'est terrible pour l'apprentissage car l'IA ne peut pas comprendre comment s'améliorer ; c'est comme essayer d'apprendre à conduire en ne regardant la route que lorsque l'on est déjà dans la bonne voie, sans aucun moyen d'ajuster le volant.

Pour corriger cela, les scientifiques ont inventé des versions « douces » de cette sélection. Au lieu d'un « oui ou non » tranché, l'IA attribue un score de « peut-être » à chaque livre, permettant ainsi d'apprendre de ses erreurs. Mais il y a un piège : ces versions douces sont souvent si lentes et coûteuses en calcul qu'elles font planter le système quand la bibliothèque devient trop grande. Elles sont comme si l'on essayait de trier un million de livres à la main pendant que la bibliothèque est en feu. La grande question pour les chercheurs a été : pouvons-nous avoir un bibliothécaire qui soit à la fois assez délicat pour apprendre (différentiable) et assez rapide pour gérer des millions de livres sans transpirer ?

C'est là qu'intervient le nouvel article, « Fast LapSum ». Les auteurs, une équipe de Pologne, ont construit un nouvel outil qui agit comme un bibliothécaire super efficace et mathématiquement parfait. Ils ont créé une méthode appelée Fast LapSum qui permet à une IA de choisir les meilleurs éléments d'une liste de millions d'éléments tout en étant capable d'apprendre du processus. Contrairement aux méthodes précédentes qui renonçaient soit à la précision parfaite pour gagner en vitesse, soit étaient trop lentes pour être utiles, Fast LapSum réussit à faire les deux. Il trouve exactement le bon nombre d'éléments à choisir (le « budget ») et calcule les scores de « peut-être » parfaits pour eux en un clin d'œil.

Le ingrédient secret est une astuce ingénieuse impliquant une vue « floue » des scores. Imaginez que les scores ne soient pas des points nets mais des nuages vaporeux. L'IA doit tracer une ligne à travers ces nuages de sorte que la quantité totale de « nuage » au-dessus de la ligne soit exactement égale au nombre de livres qu'elle est autorisée à choisir. Les anciennes méthodes tentaient de trouver cette ligne en devinant et en vérifiant encore et encore, ce qui prenait une éternité. Fast LapSum, cependant, utilise une formule mathématique spéciale (basée sur ce qu'on appelle la distribution de Laplace) qui lui permet de calculer la ligne instantanément après un seul tri.

Pour les listes vraiment gigantesques — comme un million ou même cent millions de scores — les auteurs ont ajouté un second tour de magie appelé « encadrement probabiliste ». Au lieu de trier la liste entière d'un million d'éléments, ce qui revient à essayer d'organiser un stade rempli de gens, le système prend un échantillon rapide pour deviner où la ligne se trouve probablement. Il ne trie ensuite que le petit groupe de personnes se tenant juste près de cette ligne. Cela rend le processus incroyablement rapide, ne prenant que quelques millisecondes même pour des ensembles de données massifs.

L'article prouve que cela fonctionne en testant la méthode sur deux tâches difficiles. Premièrement, ils l'ont utilisée pour créer des « exemples adverses », qui sont des images qui paraissent normales pour les humains mais trompent les classificateurs d'IA. Ils ont réussi à modifier une image de manière si infime — en altérant seulement environ 0,02 % des pixels (environ 600 pixels sur 3,3 millions) — que l'IA a identifié à tort une image de tigre. Cela a été fait beaucoup plus rapidement et avec moins de « dommages » à l'image que les méthodes précédentes. Deuxièmement, ils ont construit un codeur d'image différentiable de toutes pièces, un système qui compresse les images en sélectionnant uniquement les parties les plus importantes à conserver. Dans les deux cas, Fast LapSum a agi comme le moteur, gérant des millions de décisions par seconde sans ralentir le processus d'apprentissage.

Les auteurs démontrent que cette méthode n'est pas seulement une idée théorique, mais un outil pratique qui s'exécute en millisecondes sur des puces informatiques standards. Ils ont comparé leur travail à d'autres tentatives récentes, comme une méthode appelée DFTopK, et ont constaté que bien que ces méthodes soient rapides, elles sacrifient l'exactitude de la sélection (le nombre total d'éléments choisis s'éloigne de la cible). Fast Lapsum, affirment-ils, est le premier à maintenir une sélection parfaitement exacte tout en restant assez rapide pour les systèmes d'IA réels à grande échelle. Il transforme un goulot d'étranglement lent et coûteux en une opération fluide et rapide, permettant à l'IA d'être à la fois intelligente et efficace.

Résumé Technique : Fast LapSum : Top-k Différentiable Exact à l'Échelle du Million

1. Énoncé du Problème

L'opération top-k est une primitive fondamentale dans le calcul parcimonieux moderne, essentielle pour le routage de jetons dans les modèles Mixture-of-Experts (MoE), la sélection de mémoire dans les architectures à contexte long et les systèmes de recherche. Cependant, la sélection hard top-k standard est non-différentiable, bloquant le flux de gradient à travers la frontière de décision. Inversement, les relaxations continues (soft) existantes (par exemple, NeuralSort, SoftSort, Sinkhorn-OT) introduisent souvent des coûts de calcul prohibitifs ( $O(n^2)$ ou une complexité itérative) ou ne parviennent pas à préserver le budget de sélection exact (la somme des probabilités sélectionnées est égale à $k$ ).

Les méthodes actuelles de top-k différentiable font face à un compromis : elles sont soit trop coûteuses pour être exécutées dans la boucle d'entraînement de modèles à grande échelle, soit elles relaxent la contrainte de normalisation ( $\sum p_i = k$ ), entraînant une parcimonie incontrôlée. L'article soutient qu'un parcimoniseur différentiable doit être une « primitive système » qui maintient le budget exact du hard top-k tout en restant suffisamment peu coûteux pour des vecteurs à l'échelle du million.

2. Méthodologie

2.1 L'Opérateur LapSum

Les auteurs s'appuient sur l'opérateur LapSum, qui modélise la sélection top-k de manière probabiliste. Au lieu d'un masque binaire déterministe, LapSum traite les scores $r_i$ comme étant perturbés par un bruit de Laplace $\epsilon \sim \text{Laplace}(0, t)$ . La probabilité de sélection $p_i$ est définie comme la probabilité qu'un score perturbé dépasse un seuil scalaire $b$ :
$p_i = \sigma\left(\frac{r_i - b}{t}\right)$
où $\sigma$ est la fonction de répartition (CDF) de Laplace. Le seuil $b$ est déterminé de manière unique par la contrainte que la masse totale est égale au budget $k$ :
$\sum_{i=1}^n p_i = k$
Contrairement aux méthodes antérieures qui approchent ce seuil, LapSum résout $b$ exactement. La passe arrière est analytique, utilisant un produit vecteur-Jacobien (VJP) qui évite les substituts biaisés comme les estimateurs straight-through.

2.2 Fast LapSum : Le Solveur GPU

La contribution principale est Fast LapSum, un solveur natif pour GPU qui calcule le seuil exact $b$ en temps linéaire après le tri.

Solveur par Tri Complet (Full-Sort Solver) : Pour les vecteurs où le tri est réalisable, la méthode calcule une « courbe de seuil » en utilisant un scan unique de préfixe/suffixe stable sur les scores triés. Le seuil $b$ est trouvé via une recherche d'index et un calcul de racine en forme close dans l'intervalle spécifique où l'équation du budget s'équilibre. Cela remplace la bisection itérative utilisée dans les implémentations précédentes, réduisant le travail post-tri à $O(n)$ .
Encadrement Probabiliste (Échelle du Million) : Pour les échelles extrêmes (par exemple, $n > 10^6$ $n > 1 0^{6}$ ) où le tri de l'ensemble du vecteur devient le goulot d'étranglement, les auteurs introduisent un pipeline d'encadrement probabiliste :
1. Localisation : Un petit échantillon de scores bruités par noyau est prélevé pour estimer le quantile correspondant au budget. En utilisant les statistiques d'ordre binomiales, un intervalle étroit $[a_L, a_R)$ est identifié, contenant le véritable seuil $b$ avec une haute probabilité.
2. Streaming et Filtrage : Le vecteur complet est parcouru une seule fois (streaming). Les scores en dehors de l'encadrement sont agrégés en sommes de queues stables ( $\tilde{S}_L, \tilde{S}_R$ ) basées sur les propriétés de la CDF de Laplace. Seule la « bande intermédiaire incertaine » de scores située dans $[a_L, a_R)$ est conservée.
3. Solution Exacte : Le solveur LapSum exact est appliqué uniquement à cette bande intermédiaire réduite.
4. Vérification : Une étape de certification vérifie si l'encadrement tient ; sinon, il est élargi de manière itérative. Cela garantit que le résultat final satisfait la contrainte de budget exacte $\sum p_i = k$ .

2.3 Relation avec DFTopK

L'article distingue Fast LapSum de la méthode concurrente DFTopK. DFTopK atteint un temps linéaire en fixant le seuil au point médian entre les statistiques d'ordre $k$ et $k+1$ , relaxant ainsi la contrainte de normalisation. Les auteurs prouvent que DFTopK est la limite de température zéro de LapSum. Cependant, à des températures finies, Fast LapSum préserve le budget exact, tandis que DFTopK permet à la masse active de dériver (par exemple, résultant en une masse d'environ $3.46k$ au lieu de $k$ ).

3. Contributions Clés

Opérateur Fast LapSum : Le premier solveur GPU qui atteint une complexité de temps linéaire après le tri tout en préservant strictement la contrainte non relaxée $\sum p_i = k$ . Il élimine le besoin de substituts biaisés.
Pipeline d'Encadrement Probabiliste : Une méthode pour les vecteurs à l'échelle du million qui localise le seuil via l'échantillonnage, trie uniquement la bande intermédiaire incertaine, et vérifie la solution exacte, maintenant une complexité temporelle attendue linéaire.
Scalabilité et Efficacité : Capacité démontrée à traiter des vecteurs allant jusqu'à $10^8$ scores en quelques millisecondes (par exemple, 5,23 ms pour $10^8$ scores), soit des ordres de grandeur plus rapide que les implémentations exactes précédentes.

4. Résultats Expérimentaux

4.1 Benchmarks de Performance

Vitesse : La branche de tri complet de Fast LapSum est presque deux ordres de grandeur plus rapide que la ligne de base par bisection par requête. La branche d'encadrement atteint des accélérations allant jusqu'à 19 $\times$ par rapport au tri complet à $10^8$ scores et 1200 $\times$ par rapport à la version officielle de LapSum à $3,2 \times 10^7$ scores.
Mémoire : Contrairement aux relaxations de tri dense (NeuralSort, SoftSort) qui saturent la mémoire (OOM) autour de $n=10^5$ , Fast LapSum fonctionne efficacement à $n=10^8$ .
Précision : Fast LapSum maintient une normalisation de budget exacte (erreur relative $< 10^{-5}$ ) à toutes les échelles. En revanche, DFTopK présente une erreur relative de $\approx 2,46$ (masse $\approx 3,46k$ ) en raison de sa contrainte relaxée.

4.2 Applications

Les auteurs démontrent l'opérateur dans deux applications exigeantes à l'échelle du million :

Codeur d'Image Mégapixel Différentiable : Un codeur d'image parcimonieux entièrement différentiable entraîné à partir de zéro. L'opérateur gère la sélection de coordonnées sur des millions de scores au sein d'une boucle taux-distorsion, assemblant avec succès une image à partir d'atomes parcimonieux sans dominer le coût de calcul.
Attaques Adversaires Parcimonieuses : L'opérateur est utilisé pour générer des exemples adverses parcimonieux sur des images à résolution native (DIV2K) avec un classifieur ConvNeXt-V2.
- Efficacité : L'attaque atteint un taux de succès de 100 % avec un budget médian de 199,8 équivalents de pixels entièrement modifiés.
- Comparaison : C'est nettement plus efficace que les bases : $\sigma$ -zero nécessite 3 101, Sparse-PGD nécessite 16 134, et PGD-L0 nécessite 74 557.
- Précision : L'attaque concentre les perturbations sur les textures définissant la classe (par exemple, l'œil d'un tigre) plutôt que sur du bruit dispersé, n'utilisant que $\sim 0,02 \%$ des pixels de l'image.

5. Signification et Revendications

L'article affirme que Fast LapSum résout le compromis de longue date entre différentiabilité, contraintes de budget exact et efficacité de calcul. En traitant le top-k différentiable comme une primitive système, la méthode permet :

Entraînement de Bout en Bout : L'opérateur est assez rapide pour être intégré directement dans la boucle d'entraînement de grands modèles parcimonieux (routage, recherche, élagage) sans devenir le goulot d'étranglement computationnel.
Contrôle de la Parcimonie Exacte : Contrairement aux relaxations douces qui font dériver la masse, Fast Lapsum garantit que le nombre d'experts actifs, d'entrées de mémoire ou de candidats de recherche respecte strictement le budget spécifié $k$ .
Praticité à Grande Échelle : La méthode rend le top-k doux exact pratique pour des tâches réelles impliquant des millions de coordonnées, telles que le traitement d'images haute résolution et la recherche à grande échelle, là où les méthodes exactes précédentes étaient infaisables.

Les auteurs notent des limites : la solution en forme close est spécifique au noyau de Laplace, et la mise à l'échelle linéaire est attendue plutôt que garantie dans le pire des cas (des distributions de scores pathologiques pourraient déclencher des passes de vérification supplémentaires). Cependant, les résultats empiriques suggèrent que l'opérateur est un coût négligeable à l'échelle, permettant une nouvelle classe de systèmes à budget exact.