A Faster Generalized Two-Stage Approximate Top-K

Riepilogo Tecnico: Un Approximate Top-K Generalizzato a Due Stadi Più Veloce

Enunciato del Problema
L'identificazione degli elementi top-$K$ in un array è un'operazione fondamentale in molti algoritmi di machine learning, inclusa la Ricerca del Massimo Prodotto Interno (MIPS), i K-Nearest-Neighbors (KNN) e l'ottimizzazione di grandi modelli fondazionali (ad esempio, per la sparsità nei logit di classificazione, nei blocchi MLP, nei meccanismi di attenzione e per la compressione della KVCache). Sui moderni acceleratori come TPU e GPU, il calcolo esatto del Top-$K$ è spesso un collo di bottiglia delle prestazioni. Questi dispositivi sono ottimizzati per le moltiplicazioni di matrici dense (MXU/TensorCore), mentre le operazioni Top-$K$ fanno tipicamente affidamento su unità scalari/vettoriali (VPU/CUDA cores) e possono essere significativamente più lente rispetto alle moltiplicazioni di matrici che generano i dati. Ad esempio, su TPUv5e, trovare il 2% superiore delle attivazioni nascoste in un modello Gemma 2 9B richiede 27 volte più tempo rispetto alla corrispondente moltiplicazione di matrice.

Sebbene gli algoritmi di approximate Top-K offrano una soluzione, il metodo all'avanguardia di Chern et al. (2022), implementato come jax.lax.approx_max_k, comporta ancora una latenza significativa. Il loro approccio a due stadi partiziona l'input in bucket, seleziona l'elemento top-1 da ciascun bucket e ordina i risultati. Tuttavia, la seconda fase (l'ordinamento degli elementi selezionati) rimane un collo di bottiglia, e la prima fase spesso sottoutilizza le risorse di calcolo scalari/vettoriali disponibili.

Metodologia
Gli autori propongono una generalizzazione dell'algoritmo di Chern et al. (2022) che modifica la prima fase per selezionare gli elementi top-$K'$ da ciascun bucket (dove $1 \le K' \le K$) invece di limitare la selezione al solo top-1.

1. Generalizzazione Algoritmica:

   • Fase 1: L'array di input di dimensione $N$ viene partizionato in $B$ bucket. Invece di selezionare solo il massimo, l'algoritmo seleziona gli elementi top-$K'$ da ciascun bucket.
   • Fase 2: Gli $B \times K'$ elementi selezionati vengono uniti e ordinati per restituire gli elementi finali top-$K$.
   • Implementazione: La prima fase è implementata come una subroutine di aggiornamento online utilizzando Pallas (un linguaggio di kernel JAX per TPU). Mantiene una lista ordinata dei valori e degli indici top-$K'$ per ciascun bucket. I nuovi elementi vengono inseriti utilizzando un passaggio simile a bubble-sort, ottimizzato per la vettorizzazione attraverso i bucket. Gli autori forniscono sia una versione "non fusa" (due kernel separati) sia una versione "fusa con matmul" in cui la prima fase è unita alla moltiplicazione di matrice precedente, nascondendo efficacemente il costo della prima fase.

2. Analisi Teorica:

   • Gli autori derivano un'espressione esatta per il richiamo atteso di questo algoritmo generalizzato sotto partizionamento casuale, modellando la distribuzione degli elementi top-$K$ nei bucket utilizzando la distribuzione ipergeometrica.
   • Dimostrano che, per un target di richiamo atteso fisso, l'aumento di $K'$ consente di ridurre il numero di bucket $B$ in modo che il numero totale di elementi passati alla seconda fase ($B \times K'$) sia minimizzato.
   • Crucialmente, derivano un nuovo limite per il numero di bucket richiesto quando $K'=1$. Questo limite è dimostrabilmente più stretto di un fattore 2 rispetto al limite in Chern et al. (2022), consentendo meno bucket e quindi meno elementi da ordinare anche nell'impostazione originale.

3. Modellazione delle Prestazioni:

   • Utilizzando l'analisi del punto di sella (ridge-point), gli autori caratterizzano i sottosistemi dell'acceleratore (HBM, VPU, MXU). Dimostrano che la prima fase è tipicamente limitata dalla memoria fino a quando $K'$ non raggiunge una certa soglia (ad esempio, $K' \approx 6$ su TPUv5e). All'interno di questo regime limitato dalla memoria, l'aumento di $K'$ non aumenta significativamente la latenza della prima fase, ma riduce drasticamente il carico di lavoro della seconda fase.

Contributi Chiave

• Algoritmo Generalizzato: Un nuovo algoritmo di approximate Top-K a due stadi che seleziona elementi top-$K'$ per bucket, offrendo un compromesso flessibile tra il calcolo della prima fase e il carico di ordinamento della seconda fase.
• Limiti Teorici:
  • Un'espressione esatta per il richiamo atteso in funzione di $N, K, B,$ e $K'$.
  • Un limite dimostrabilmente più stretto (fattore 2) per il numero di bucket richiesti nel caso $K'=1$ rispetto ai lavori precedenti, migliorando la selezione dei parametri.
• Implementazione: Un'implementazione ad alte prestazioni su Cloud TPUv5e utilizzando Pallas, con varianti sia non fuse sia fuse con matmul.
• Selezione Automatica dei Parametri: Una routine che seleziona automaticamente $K'$ e $B$ per soddisfare un target di richiamo specificato dall'utente, eliminando la necessità di sintonizzazione manuale.

Risultati

• Efficienza Teorica: Per un target di richiamo fisso (ad esempio, 95%), l'aumento di $K'$ riduce significativamente il numero di elementi elaborati nella seconda fase. Ad esempio, selezionare il top-1024 da 262.144 elementi richiede 16.384 elementi di output con $K'=1$, ma solo 2.048 elementi con $K'=4$ (una riduzione di 8 volte).
• Miglioramenti della Latenza (Non Fuso): Su TPUv5e, l'algoritmo generalizzato con $K'=4$ raggiunge un speedup di 11 volte rispetto all'originale jax.lax.approx_max_k (27µs contro 305µs) per un target di richiamo del 99% sullo stesso compito.
• Miglioramenti della Latenza (Fuso): In un carico di lavoro di Ricerca del Massimo Prodotto Interno (MIPS) (1 milione di vettori, 1024 query), l'implementazione fusa con matmul con $K'=4$ riduce il tempo di esecuzione totale da 594ms (Top-K esatto) a 10ms. Ciò rappresenta uno speedup complessivo di 6,7 volte rispetto alla baseline densa, portando il passaggio di approximate Top-K al di sotto del costo della moltiplicazione di matrice stessa.
• Impatto End-to-End: In un benchmark di blocco transformer sparso (Gemma 2 9B), il metodo proposto riduce l'overhead del blocco MLP sparso da 2,7 volte più lento rispetto alla versione densa (utilizzando l'algoritmo originale) a un modesto overhead di soli +5ms rispetto alla baseline densa.

Significato e Affermazioni
Il documento afferma che la scelta di $K'=1$ negli algoritmi di approximate Top-K precedenti non sfrutta appieno le risorse di calcolo disponibili sui moderni acceleratori. Generalizzando a $K' > 1$, l'algoritmo sposta efficacemente il collo di bottiglia dalla fase di ordinamento senza sostenere costi aggiuntivi significativi nella fase di filtraggio, a condizione che la prima fase rimanga limitata dalla memoria.

Gli autori sottolineano che la loro analisi teorica fornisce una base rigorosa per la selezione dei parametri, offrendo un'interfaccia user-friendly (approx_top_k(array, K, recall_target)) che non richiede sintonizzazione manuale. Sebbene la valutazione empirica sia focalizzata su TPUv5e, gli autori affermano che l'algoritmo, la modellazione delle prestazioni e l'analisi sono generali e applicabili ad altre piattaforme hardware come le GPU. Il lavoro dimostra che una generalizzazione algoritmica fondata su principi, combinata con un'implementazione consapevole dell'hardware (fusione), può eliminare il Top-K come collo di bottiglia nell'addestramento e nell'inferenza di modelli su larga scala.