A Faster Generalized Two-Stage Approximate Top-K

Resumo Técnico: Uma Aproximação Top-K Generalizada em Duas Etapas Mais Rápida

Declaração do Problema
A identificação dos elementos top-$K$ em um array é uma operação fundamental em muitos algoritmos de aprendizado de máquina, incluindo Busca de Produto Interno Máximo (MIPS), K-Vizinhos Mais Próximos (KNN) e a otimização de grandes modelos de base (por exemplo, para esparsidade em logits de classificação, blocos MLP, mecanismos de atenção e compressão de KVCache). Em aceleradores modernos como TPUs e GPUs, o cálculo exato do Top-$K$ é frequentemente um gargalo de desempenho. Esses dispositivos são otimizados para multiplicações de matrizes densas (MXUs/TensorCores), enquanto as operações Top-$K$ geralmente dependem de unidades escalares/vetoriais (VPUs/CUDA cores) e podem ser significativamente mais lentas do que as multiplicações de matrizes que geram os dados. Por exemplo, na TPUv5e, encontrar os 2% superiores das ativações ocultas em um modelo Gemma 2 9B leva 27 vezes mais tempo do que a correspondente multiplicação de matrizes.

Embora algoritmos de Top-$K$ aproximado ofereçam uma solução alternativa, o método mais avançado de Chern et al. (2022), implementado como jax.lax.approx_max_k, ainda incorre em latência significativa. Sua abordagem em duas etapas particiona a entrada em buckets, seleciona o elemento top-1 de cada bucket e classifica os resultados. No entanto, a segunda etapa (classificar os elementos selecionados) permanece um gargalo, e a primeira etapa frequentemente subutiliza os recursos de computação escalar/vetorial disponíveis.

Metodologia
Os autores propõem uma generalização do algoritmo de Chern et al. (2022) que modifica a primeira etapa para selecionar os elementos top-$K'$ de cada bucket (onde $1 \le K' \le K$), em vez de restringir a seleção ao top-1.

1. Generalização Algorítmica:

   • Etapa 1: O array de entrada de tamanho $N$ é particionado em $B$ buckets. Em vez de selecionar apenas o máximo, o algoritmo seleciona os elementos top-$K'$ de cada bucket.
   • Etapa 2: Os $B \times K'$ elementos selecionados são mesclados e classificados para retornar os elementos top-$K$ finais.
   • Implementação: A primeira etapa é implementada como uma sub-rotina de atualização online usando Pallas (uma linguagem de kernel JAX para TPUs). Ela mantém uma lista classificada dos valores e índices top-$K'$ para cada bucket. Novos elementos são inseridos usando uma passagem semelhante ao bubble-sort, otimizada para vetorização entre buckets. Os autores fornecem tanto uma versão "desfusão" (dois kernels separados) quanto uma versão "fusão de matmul", onde a primeira etapa é mesclada com a multiplicação de matrizes precedente, efetivamente ocultando o custo da primeira etapa.

2. Análise Teórica:

   • Os autores derivam uma expressão exata para o recall esperado deste algoritmo generalizado sob particionamento aleatório, modelando a distribuição dos elementos top-$K$ nos buckets usando a distribuição hipergeométrica.
   • Eles provam que, para um alvo de recall esperado fixo, o aumento de $K'$ permite uma redução no número de buckets $B$ tal que o número total de elementos passados para a segunda etapa ($B \times K'$) seja minimizado.
   • Crucialmente, eles derivam um novo limite para o número de buckets necessários quando $K'=1$. Este limite é comprovadamente mais apertado por um fator de 2 em comparação com o limite em Chern et al. (2022), permitindo menos buckets e, portanto, menos elementos para classificar, mesmo no cenário original.

3. Modelagem de Desempenho:

   • Usando análise de ponto de sela (ridge-point), os autores caracterizam os subsistemas do acelerador (HBM, VPU, MXU). Eles demonstram que a primeira etapa é tipicamente limitada pela memória até que $K'$ atinja um certo limiar (por exemplo, $K' \approx 6$ na TPUv5e). Dentro deste regime limitado pela memória, aumentar $K'$ não aumenta significativamente a latência da primeira etapa, mas reduz drasticamente a carga de trabalho da segunda etapa.

Principais Contribuições

• Algoritmo Generalizado: Um novo algoritmo de Top-$K$ aproximado em duas etapas que seleciona elementos top-$K'$ por bucket, oferecendo um compromisso flexível entre a computação da primeira etapa e a carga de classificação da segunda etapa.
• Limites Teóricos:
  • Uma expressão exata para o recall esperado como função de $N, K, B,$ e $K'$.
  • Um limite comprovadamente mais apertado (fator de 2) para o número de buckets necessários no caso $K'=1$ em comparação com trabalhos anteriores, melhorando a seleção de parâmetros.
• Implementação: Uma implementação de alto desempenho na Cloud TPUv5e usando Pallas, apresentando variantes desfusão e fusão de matmul.
• Seleção Automática de Parâmetros: Uma rotina que seleciona automaticamente $K'$ e $B$ para atender a um alvo de recall especificado pelo usuário, eliminando a necessidade de ajuste manual.

Resultados

• Eficiência Teórica: Para um alvo de recall fixo (por exemplo, 95%), o aumento de $K'$ reduz significativamente o número de elementos processados na segunda etapa. Por exemplo, selecionar os top-1024 de 262.144 elementos requer 16.384 elementos de saída com $K'=1$, mas apenas 2.048 elementos com $K'=4$ (uma redução de 8x).
• Melhorias de Latência (Desfusão): Na TPUv5e, o algoritmo generalizado com $K'=4$ alcança uma aceleração de 11x sobre o jax.lax.approx_max_k original (27µs vs 305µs) para um alvo de recall de 99% na mesma tarefa.
• Melhorias de Latência (Fusão): Em uma carga de trabalho de Busca de Produto Interno Máximo (MIPS) (1M de vetores, 1024 consultas), a implementação com fusão de matmul com $K'=4$ reduz o tempo de execução total de 594ms (Top-K exato) para 10ms. Isso representa uma aceleração geral de 6,7x em comparação com a base densa, trazendo a etapa Top-K aproximada abaixo do custo da própria multiplicação de matrizes.
• Impacto de Ponta a Ponta: Em um benchmark de bloco transformador esparsificado (Gemma 2 9B), o método proposto reduz a sobrecarga do bloco MLP esparsificado de 2,7x mais lento que o denso (usando o algoritmo original) para apenas uma sobrecarga modesta de +5ms em relação à base densa.

Significado e Alegações
O artigo afirma que a escolha de $K'=1$ em algoritmos anteriores de Top-$K$ aproximado não utiliza plenamente os recursos de computação disponíveis em aceleradores modernos. Ao generalizar para $K' > 1$, o algoritmo desloca efetivamente o gargalo da etapa de classificação sem incorrer em custo adicional significativo na etapa de filtragem, desde que a primeira etapa permaneça limitada pela memória.

Os autores enfatizam que sua análise teórica fornece uma base rigorosa para a seleção de parâmetros, oferecendo uma interface amigável ao usuário (approx_top_k(array, K, recall_target)) que não requer ajuste manual. Embora a avaliação empírica seja focada na TPUv5e, os autores afirmam que o algoritmo, a modelagem de desempenho e a análise são gerais e aplicáveis a outras plataformas de hardware, como GPUs. O trabalho demonstra que uma generalização algorítmica principista, combinada com uma implementação consciente do hardware (fusão), pode eliminar o Top-K como um gargalo no treinamento e inferência de modelos em grande escala.