A Faster Generalized Two-Stage Approximate Top-K

Resumen Técnico: Una Aproximación Generalizada de Dos Etapas Más Rápida para el Top-K

Enunciado del Problema
La identificación de los elementos top-$K$ en un array es una operación fundamental en muchos algoritmos de aprendizaje automático, incluida la Búsqueda del Máximo Producto Interno (MIPS), los Vecinos Más Cercanos (KNN) y la optimización de grandes modelos base (por ejemplo, para la dispersión en logits de clasificación, bloques MLP, mecanismos de atención y compresión de KVCache). En aceleradores modernos como TPUs y GPUs, calcular el Top-K exacto suele ser un cuello de botella de rendimiento. Estos dispositivos están optimizados para multiplicaciones de matrices densas (MXUs/TensorCores), mientras que las operaciones Top-K dependen típicamente de unidades escalares/vectoriales (VPUs/CUDA cores) y pueden ser significativamente más lentas que las multiplicaciones de matrices que generan los datos. Por ejemplo, en TPUv5e, encontrar el 2% superior de las activaciones ocultas en un modelo Gemma 2 9B tarda 27 veces más que la multiplicación de matrices correspondiente.

Aunque los algoritmos aproximados de Top-K ofrecen una solución, el método más avanzado de Chern et al. (2022), implementado como jax.lax.approx_max_k, aún incurre en una latencia significativa. Su enfoque de dos etapas particiona la entrada en cubetas, selecciona el elemento top-1 de cada cubeta y ordena los resultados. Sin embargo, la segunda etapa (ordenar los elementos seleccionados) sigue siendo un cuello de botella, y la primera etapa a menudo subutiliza los recursos de cómputo escalar/vectorial disponibles.

Metodología
Los autores proponen una generalización del algoritmo de Chern et al. (2022) que modifica la primera etapa para seleccionar los top-$K'$ elementos de cada cubeta (donde $1 \le K' \le K$) en lugar de restringir la selección al top-1.

1. Generalización Algorítmica:

   • Etapa 1: El array de entrada de tamaño $N$ se particiona en $B$ cubetas. En lugar de seleccionar solo el máximo, el algoritmo selecciona los top-$K'$ elementos de cada cubeta.
   • Etapa 2: Los $B \times K'$ elementos seleccionados se fusionan y ordenan para devolver los elementos top-$K$ finales.
   • Implementación: La primera etapa se implementa como una subrutina de actualización en línea utilizando Pallas (un lenguaje de kernels para JAX diseñado para TPUs). Mantiene una lista ordenada de los valores e índices top-$K'$ para cada cubeta. Los nuevos elementos se insertan mediante un pase similar a la ordenación por burbuja, optimizado para la vectorización a través de las cubetas. Los autores proporcionan tanto una versión "no fusionada" (dos kernels separados) como una versión "fusionada con multiplicación de matrices" donde la primera etapa se fusiona con la multiplicación de matrices precedente, ocultando efectivamente el costo de la primera etapa.

2. Análisis Teórico:

   • Los autores derivan una expresión exacta para la recuperación esperada de este algoritmo generalizado bajo particionamiento aleatorio, modelando la distribución de los elementos top-$K$ en las cubetas utilizando la distribución hipergeométrica.
   • Demuestran que, para un objetivo de recuperación esperado fijo, aumentar $K'$ permite reducir el número de cubetas $B$ de tal manera que se minimiza el número total de elementos pasados a la segunda etapa ($B \times K'$).
   • Crucialmente, derivan un nuevo límite para el número de cubetas requerido cuando $K'=1$. Este límite es demostrablemente más estricto por un factor de 2 en comparación con el límite en Chern et al. (2022), lo que permite menos cubetas y, por tanto, menos elementos para ordenar incluso en el escenario original.

3. Modelado de Rendimiento:

   • Utilizando el análisis de punto de cumbre (ridge-point), los autores caracterizan los subsistemas del acelerador (HBM, VPU, MXU). Demuestran que la primera etapa suele estar limitada por la memoria hasta que $K'$ alcanza cierto umbral (por ejemplo, $K' \approx 6$ en TPUv5e). Dentro de este régimen limitado por la memoria, aumentar $K'$ no incrementa significativamente la latencia de la primera etapa, pero reduce drásticamente la carga de trabajo de la segunda etapa.

Contribuciones Clave

• Algoritmo Generalizado: Un nuevo algoritmo aproximado de Top-K de dos etapas que selecciona elementos top-$K'$ por cubeta, ofreciendo un compromiso flexible entre el cómputo de la primera etapa y la carga de ordenación de la segunda etapa.
• Límites Teóricos:
  • Una expresión exacta para la recuperación esperada en función de $N, K, B,$ y $K'$.
  • Un límite demostrablemente más estricto (factor de 2) para el número de cubetas requerido en el caso $K'=1$ en comparación con trabajos anteriores, mejorando la selección de parámetros.
• Implementación: Una implementación de alto rendimiento en Cloud TPUv5e utilizando Pallas, que presenta variantes tanto no fusionadas como fusionadas con multiplicación de matrices.
• Selección Automática de Parámetros: Una rutina que selecciona automáticamente $K'$ y $B$ para cumplir con un objetivo de recuperación especificado por el usuario, eliminando la necesidad de ajuste manual.

Resultados

• Eficiencia Teórica: Para un objetivo de recuperación fijo (por ejemplo, 95%), aumentar $K'$ reduce significativamente el número de elementos procesados en la segunda etapa. Por ejemplo, seleccionar los top-1024 de 262,144 elementos requiere 16,384 elementos de salida con $K'=1$, pero solo 2,048 elementos con $K'=4$ (una reducción de 8x).
• Mejoras de Latencia (No Fusionado): En TPUv5e, el algoritmo generalizado con $K'=4$ logra una aceleración de 11x sobre jax.lax.approx_max_k original (27µs vs 305µs) para un objetivo de recuperación del 99% en la misma tarea.
• Mejoras de Latencia (Fusionado): En una carga de trabajo de Búsqueda del Máximo Producto Interno (MIPS) (1M de vectores, 1024 consultas), la implementación fusionada con multiplicación de matrices con $K'=4$ reduce el tiempo de ejecución total de 594ms (Top-K exacto) a 10ms. Esto representa una aceleración general de 6.7x en comparación con la base densa, situando el paso de Top-K aproximado por debajo del costo de la propia multiplicación de matrices.
• Impacto de Extremo a Extremo: En una prueba de referencia de bloque de transformador disperso (Gemma 2 9B), el método propuesto reduce la sobrecarga del bloque MLP disperso de ser 2.7 veces más lento que el denso (usando el algoritmo original) a solo una sobrecarga modesta de +5ms en relación con la base densa.

Significado y Afirmaciones
El artículo afirma que la elección de $K'=1$ en algoritmos previos de Top-K aproximado no utiliza plenamente los recursos de cómputo disponibles en aceleradores modernos. Al generalizar a $K' > 1$, el algoritmo desplaza efectivamente el cuello de botella lejos de la etapa de ordenación sin incurrir en un costo adicional significativo en la etapa de filtrado, siempre que la primera etapa permanezca limitada por la memoria.

Los autores enfatizan que su análisis teórico proporciona una base rigurosa para la selección de parámetros, ofreciendo una interfaz amigable para el usuario (approx_top_k(array, K, recall_target)) que no requiere ajuste manual. Si bien la evaluación empírica se centra en TPUv5e, los autores afirman que el algoritmo, el modelado de rendimiento y el análisis son generales y aplicables a otras plataformas de hardware como GPUs. El trabajo demuestra que una generalización algorítmica basada en principios, combinada con una implementación consciente del hardware (fusión), puede eliminar el Top-K como un cuello de botella en el entrenamiento y la inferencia de modelos a gran escala.