A Faster Generalized Two-Stage Approximate Top-K

Technische Zusammenfassung: Eine schnellere generalisierte zweistufige approximative Top-K-Suche

Problemstellung
Die Identifizierung der Top-$K$-Elemente in einem Array ist eine fundamentale Operation in vielen maschinellen Lernalgorithmen, einschließlich der Suche nach dem Maximum des inneren Produkts (Maximum Inner Product Search, MIPS), der K-Nächsten-Nachbarn-Suche (KNN) und der Optimierung großer Fundamentaler Modelle (z. B. für Sparsität in Klassifikations-Logits, MLP-Blöcken, Aufmerksamkeitsmechanismen und KVCache-Kompression). Auf modernen Beschleunigern wie TPUs und GPUs stellt die Berechnung exakter Top-$K$-Werte oft einen Leistungsengpass dar. Diese Geräte sind für dichte Matrixmultiplikationen (MXUs/TensorCores) optimiert, whereas Top-$K$-Operationen typischerweise auf skalare/Vektor-Einheiten (VPUs/CUDA-Kerne) angewiesen sind und deutlich langsamer sein können als die Matrixmultiplikationen, die die Daten erzeugen. Beispielsweise dauert es auf einem TPUv5e, die Top-2 % der versteckten Aktivierungen in einem Gemma-2-9B-Modell zu finden, 27-mal länger als die entsprechende Matrixmultiplikation.

Während approximative Top-K-Algorithmen eine Umgehungsmöglichkeit bieten, verursacht die fortschrittlichste Methode von Chern et al. (2022), implementiert als jax.lax.approx_max_k, immer noch signifikante Latenz. Ihr zweistufiger Ansatz unterteilt die Eingabe in Buckets, wählt das Top-1-Element aus jedem Bucket aus und sortiert die Ergebnisse. Der zweite Schritt (Sortieren der ausgewählten Elemente) bleibt jedoch ein Engpass, und der erste Schritt nutzt die verfügbaren skalaren/Vektor-Rechenressourcen oft nicht aus.

Methodik
Die Autoren schlagen eine Verallgemeinerung des Algorithmus von Chern et al. (2022) vor, die den ersten Schritt so modifiziert, dass die Top-$K'$-Elemente aus jedem Bucket ausgewählt werden (wobei $1 \le K' \le K$), anstatt die Auswahl auf das Top-1 zu beschränken.

1. Algorithmische Verallgemeinerung:

   • Schritt 1: Das Eingabe-Array der Größe $N$ wird in $B$ Buckets unterteilt. Anstatt nur das Maximum auszuwählen, selektiert der Algorithmus die Top-$K'$-Elemente aus jedem Bucket.
   • Schritt 2: Die $B \times K'$ ausgewählten Elemente werden zusammengeführt und sortiert, um die finalen Top-$K$-Elemente zurückzugeben.
   • Implementierung: Der erste Schritt wird als Online-Aktualisierungs-Unterprogramm unter Verwendung von Pallas (eine JAX-Kernelsprache für TPUs) implementiert. Es wird eine sortierte Liste der Top-$K'$-Werte und Indizes für jeden Bucket verwaltet. Neue Elemente werden mittels eines bubble-sort-ähnlichen Durchlaufs eingefügt, der für die Vektorisierung über die Buckets hinweg optimiert ist. Die Autoren stellen sowohl eine „entkoppelte" Version (zwei separate Kerne) als auch eine „mit Matrixmultiplikation fusionierte" Version bereit, bei der der erste Schritt mit der vorangegangenen Matrixmultiplikation verschmolzen wird, wodurch die Kosten des ersten Schritts effektiv versteckt werden.

2. Theoretische Analyse:

   • Die Autoren leiten einen exakten Ausdruck für die erwartete Recall-Rate dieses verallgemeinerten Algorithmus unter zufälliger Partitionierung ab, wobei sie die Verteilung der Top-$K$-Elemente in die Buckets mit Hilfe der hypergeometrischen Verteilung modellieren.
   • Sie beweisen, dass bei einem festen Zielwert für die erwartete Recall-Rate eine Erhöhung von $K'$ eine Verringerung der Anzahl der Buckets $B$ ermöglicht, sodass die Gesamtzahl der an den zweiten Schritt weitergegebenen Elemente ($B \times K'$) minimiert wird.
   • Entscheidend leiten sie eine neue Schranke für die erforderliche Anzahl von Buckets ab, wenn $K'=1$ gilt. Diese Schranke ist um einen Faktor von 2 nachweislich enger als die Schranke in Chern et al. (2022), was weniger Buckets und somit weniger zu sortierende Elemente selbst im ursprünglichen Setting ermöglicht.

3. Leistungsmodellierung:

   • Mithilfe einer Ridge-Punkt-Analyse charakterisieren die Autoren die Subsysteme des Beschleunigers (HBM, VPU, MXU). Sie zeigen auf, dass der erste Schritt typischerweise speicherbandbreitenbegrenzt ist, bis $K'$ einen bestimmten Schwellenwert erreicht (z. B. $K' \approx 6$ auf TPUv5e). Innerhalb dieses speicherbandbreitenbegrenzten Regimes führt eine Erhöhung von $K'$ nicht zu einer signifikanten Erhöhung der Latenz des ersten Schritts, reduziert jedoch die Arbeitslast des zweiten Schritts drastisch.

Hauptbeiträge

• Verallgemeinerter Algorithmus: Ein neuer zweistufiger approximativer Top-K-Algorithmus, der Top-$K'$-Elemente pro Bucket auswählt und einen flexiblen Kompromiss zwischen der Rechenlast des ersten Schritts und der Sortierlast des zweiten Schritts bietet.
• Theoretische Schranken:
  • Ein exakter Ausdruck für die erwartete Recall-Rate als Funktion von $N, K, B$ und $K'$.
  • Eine nachweislich engere Schranke (Faktor 2) für die erforderliche Anzahl von Buckets im Fall $K'=1$ im Vergleich zu früheren Arbeiten, was die Parametersélection verbessert.
• Implementierung: Eine Hochleistungs-Implementierung auf Cloud TPUv5e unter Verwendung von Pallas, die sowohl entkoppelte als auch mit Matrixmultiplikation fusionierte Varianten bietet.
• Automatische Parametersélection: Eine Routine, die automatisch $K'$ und $B$ auswählt, um ein vom Benutzer spezifiziertes Recall-Ziel zu erreichen, wodurch die Notwendigkeit einer manuellen Feinabstimmung entfällt.

Ergebnisse

• Theoretische Effizienz: Bei einem festen Recall-Ziel (z. B. 95 %) reduziert eine Erhöhung von $K'$ die Anzahl der im zweiten Schritt verarbeiteten Elemente signifikant. Beispielsweise erfordert die Auswahl der Top-1024 aus 262.144 Elementen mit $K'=1$ 16.384 Ausgabeelemente, mit $K'=4$ jedoch nur 2.048 Elemente (eine Reduktion um den Faktor 8).
• Latenzverbesserungen (Entkoppelt): Auf TPUv5e erreicht der verallgemeinerte Algorithmus mit $K'=4$ eine 11-fache Beschleunigung gegenüber dem ursprünglichen jax.lax.approx_max_k (27 µs vs. 305 µs) bei einem Recall-Ziel von 99 % für dieselbe Aufgabe.
• Latenzverbesserungen (Fusioniert): Bei einer Workload für die Suche nach dem Maximum des inneren Produkts (MIPS, 1 Mio. Vektoren, 1024 Abfragen) reduziert die mit Matrixmultiplikation fusionierte Implementierung mit $K'=4$ die Gesamtlauflzeit von 594 ms (exakte Top-K) auf 10 ms. Dies entspricht einer Gesamtbeschleunigung um den Faktor 6,7 im Vergleich zur dichten Basislinie und bringt den Schritt der approximativen Top-K-Suche unter die Kosten der Matrixmultiplikation selbst.
• End-to-End-Auswirkung: In einem Benchmark für einen spärlichen Transformer-Block (Gemma 2 9B) reduziert die vorgeschlagene Methode den Overhead des spärlichen MLP-Blocks von 2,7-fach langsamer als die dichte Variante (unter Verwendung des ursprünglichen Algorithmus) auf lediglich einen moderaten Overhead von +5 ms relativ zur dichten Basislinie.

Bedeutung und Behauptungen
Die Arbeit behauptet, dass die Wahl von $K'=1$ in früheren approximativen Top-K-Algorithmen die verfügbaren Rechenressourcen auf modernen Beschleunigern nicht vollständig ausnutzt. Durch die Verallgemeinerung auf $K' > 1$ verlagert der Algorithmus den Engpass effektiv von der Sortierphase weg, ohne signifikante zusätzliche Kosten in der Filterphase zu verursachen, sofern der erste Schritt speicherbandbreitenbegrenzt bleibt.

Die Autoren betonen, dass ihre theoretische Analyse eine rigorose Grundlage für die Parametersélection bietet und eine benutzerfreundliche Schnittstelle (approx_top_k(array, K, recall_target)) bereitstellt, die keine manuelle Feinabstimmung erfordert. Obwohl die empirische Evaluation auf TPUv5e fokussiert ist, versichern die Autoren, dass der Algorithmus, die Leistungsmodellierung und die Analyse allgemein sind und auf andere Hardware-Plattformen wie GPUs anwendbar sind. Die Arbeit zeigt auf, dass eine prinzipielle algorithmische Verallgemeinerung in Kombination mit hardwarebewusster Implementierung (Fusion) Top-K als Engpass im Training und der Inferenz großer Modelle eliminieren kann.