A Faster Generalized Two-Stage Approximate Top-K

Technische Samenvatting: Een Snellere Generaliseerde Tweestaps Benadering voor Top-K

Probleemstelling
Het identificeren van de top-$K$ elementen in een array is een fundamentele bewerking in veel machine learning-algoritmen, waaronder Maximum Inner Product Search (MIPS), K-Nearest-Neighbors (KNN) en de optimalisatie van grote foundation-modellen (bijvoorbeeld voor sparsiteit in classificatie-logits, MLP-blokken, attention-mechanismen en KVCache-compressie). Op moderne versnellers zoals TPUs en GPUs is het berekenen van exacte Top-$K$ vaak een prestatieknelpunt. Deze apparaten zijn geoptimaliseerd voor dichte matrixvermenigvuldigingen (MXUs/TensorCores), terwijl Top-$K$-bewerkingen doorgaans afhankelijk zijn van scalar/vector-eenheden (VPUs/CUDA cores) en aanzienlijk trager kunnen zijn dan de matrixvermenigvuldigingen die de data genereren. Bijvoorbeeld, op TPUv5e duurt het vinden van de top-2% van verborgen activaties in een Gemma 2 9B-model 27 keer langer dan de bijbehorende matrixvermenigvuldiging.

Hoewel benaderende Top-K-algoritmen een oplossing bieden, vertoont de state-of-the-art-methode van Chern et al. (2022), geïmplementeerd als jax.lax.approx_max_k, nog steeds aanzienlijke latentie. Hun tweestapsbenadering partitioneert de invoer in bakken, selecteert het top-1 element uit elke bak en sorteert de resultaten. De tweede fase (het sorteren van de geselecteerde elementen) blijft echter een knelpunt, en de eerste fase maakt vaak onvoldoende gebruik van de beschikbare scalar/vector-berekeningsbronnen.

Methodologie
De auteurs stellen een generalisatie voor van het algoritme van Chern et al. (2022) die de eerste fase aanpast om de top-$K'$ elementen uit elke bak te selecteren (waarbij $1 \le K' \le K$), in plaats van de selectie te beperken tot het top-1 element.

1. Algoritmische Generalisatie:

   • Fase 1: Het invoerarray van grootte $N$ wordt gepartitioneerd in $B$ bakken. In plaats van alleen het maximum te selecteren, selecteert het algoritme de top-$K'$ elementen uit elke bak.
   • Fase 2: De $B \times K'$ geselecteerde elementen worden samengevoegd en gesorteerd om de uiteindelijke top-$K$ elementen terug te geven.
   • Implementatie: De eerste fase is geïmplementeerd als een online update-routine met behulp van Pallas (een JAX-kern taal voor TPUs). Het onderhoudt een gesorteerde lijst van top-$K'$ waarden en indices voor elke bak. Nieuwe elementen worden ingevoegd met een bubbel-sort-achtige doorloop, geoptimaliseerd voor vectorisatie over bakken heen. De auteurs bieden zowel een "niet-gefuseerde" versie (twee aparte kernels) als een "met-matrixvermenigvuldiging-gefuseerde" versie, waarbij de eerste fase wordt samengevoegd met de voorgaande matrixvermenigvuldiging, waardoor de kosten van de eerste fase effectief worden verborgen.

2. Theoretische Analyse:

   • De auteurs leiden een exacte uitdrukking af voor de verwachte recall van dit generaliseerde algoritme onder willekeurige partitionering, waarbij ze de verdeling van top-$K$ elementen over bakken modelleren met de hypergeometrische verdeling.
   • Ze bewijzen dat voor een vastgestelde verwachte recall-doelstelling, het verhogen van $K'$ een vermindering van het aantal bakken $B$ mogelijk maakt, zodat het totale aantal elementen dat naar de tweede fase wordt doorgegeven ($B \times K'$) wordt geminimaliseerd.
   • Cruciaal leiden ze een nieuwe bovengrens af voor het benodigde aantal bakken wanneer $K'=1$. Deze bovengrens is bewezen tweemaal zo strak als de bovengrens in Chern et al. (2022), waardoor minder bakken nodig zijn en dus minder elementen om te sorteren, zelfs in de oorspronkelijke setting.

3. Prestatiemodellering:

   • Met behulp van ridge-point-analyse karakteriseren de auteurs de subsystemen van de versneller (HBM, VPU, MXU). Ze tonen aan dat de eerste fase doorgaans geheugenbeperkt is totdat $K'$ een bepaalde drempel bereikt (bijvoorbeeld $K' \approx 6$ op TPUv5e). Binnen dit geheugenbeperkte regime leidt het verhogen van $K'$ niet tot een significante toename van de latentie van de eerste fase, maar vermindert het de werklast van de tweede fase drastisch.

Belangrijkste Bijdragen

• Generaliseerd Algoritme: Een nieuw tweestaps benaderend Top-K-algoritme dat top-$K'$ elementen per bak selecteert, en een flexibele afweging biedt tussen berekening in de eerste fase en sorteerlast in de tweede fase.
• Theoretische Grenzen:
  • Een exacte uitdrukking voor verwachte recall als functie van $N, K, B,$ en $K'$.
  • Een bewezen strakkere grens (factor 2) voor het aantal benodigde bakken in het geval $K'=1$ vergeleken met eerdere werken, wat de parameterselectie verbetert.
• Implementatie: Een hoogpresterende implementatie op Cloud TPUv5e met behulp van Pallas, met zowel niet-gefuseerde als met-matrixvermenigvuldiging-gefuseerde varianten.
• Automatische Parameterselectie: Een routine die automatisch $K'$ en $B$ selecteert om een door de gebruiker opgegeven recall-doelstelling te halen, waardoor handmatige afstemming overbodig wordt.

Resultaten

• Theoretische Efficiëntie: Voor een vaste recall-doelstelling (bijvoorbeeld 95%) vermindert het verhogen van $K'$ het aantal elementen dat in de tweede fase wordt verwerkt aanzienlijk. Bijvoorbeeld: het selecteren van top-1024 uit 262.144 elementen vereist 16.384 uitvoerelementen met $K'=1$, maar slechts 2.048 elementen met $K'=4$ (een reductie van 8x).
• Latentieverbeteringen (Niet-gefuseerd): Op TPUv5e bereikt het generaliseerde algoritme met $K'=4$ een snelheidswinst van 11x ten opzichte van de originele jax.lax.approx_max_k (27µs versus 305µs) voor een 99% recall-doelstelling op dezelfde taak.
• Latentieverbeteringen (Ge fuseerd): In een Maximum Inner Product Search (MIPS)-werklast (1 miljoen vectoren, 1024 queries) reduceert de met-matrixvermenigvuldiging-gefuseerde implementatie met $K'=4$ de totale runtime van 594ms (exacte Top-K) naar 10ms. Dit vertegenwoordigt een algehele snelheidswinst van 6,7x ten opzichte van de dichte baseline, waardoor de benaderende Top-K-stap goedkoper wordt dan de matrixvermenigvuldiging zelf.
• End-to-End Impact: In een benchmark voor een spars transformer-blok (Gemma 2 9B) reduceert de voorgestelde methode de overhead van het spars MLP-blok van 2,7x trager dan dicht (met het originele algoritme) naar slechts een bescheiden +5ms overhead ten opzichte van de dichte baseline.

Betekenis en Beweringen
Het artikel beweert dat de keuze voor $K'=1$ in eerdere benaderende Top-K-algoritmen niet volledig gebruikmaakt van de beschikbare berekeningsbronnen op moderne versnellers. Door te generaliseren naar $K' > 1$, verplaatst het algoritme het knelpunt effectief weg van de sorteerstap zonder aanzienlijke extra kosten in de filterstap, mits de eerste fase geheugenbeperkt blijft.

De auteurs benadrukken dat hun theoretische analyse een rigoureuze basis biedt voor parameterselectie, en een gebruiksvriendelijke interface biedt (approx_top_k(array, K, recall_target)) die geen handmatige afstemming vereist. Hoewel de empirische evaluatie zich richt op TPUv5e, stellen de auteurs dat het algoritme, het prestatiemodellering en de analyse algemeen zijn en toepasbaar op andere hardwareplatforms zoals GPUs. Het werk toont aan dat een principiële algoritmische generalisatie, gecombineerd met hardware-bewuste implementatie (fusie), Top-K kan elimineren als knelpunt in training en inferentie van modellen op grote schaal.