Cascade Token Selection for Transformer Attention Acceleration

Resumen Técnico: Selección de Tokens en Cascada para la Aceleración de la Atención en Transformadores

Enunciado del Problema
Los modelos transformadores enfrentan un cuello de botella computacional significativo en el mecanismo de autoatención, donde el costo de calcular el producto $QK^\top$ y la suma ponderada posterior por softmax escala cuadráticamente con la longitud de la secuencia ($O(T^2 d_h)$). Aunque métodos recientes como FlashAttention han optimizado el ancho de banda de memoria, la complejidad algorítmica sigue siendo una barrera para secuencias largas.

El artículo complementario, Atención de Descorrelación de Activaciones (ADA), propuso una solución para reducir esta complejidad de $O(T^2)$ a $O(r^2)$ seleccionando un pequeño conjunto de $r \ll T$ tokens "representativos" en cada capa. ADA identifica estos tokens utilizando un umbral de Gram: un token se considera representativo si su máxima correlación coseno con todos los tokens precedentes está por debajo de un umbral $1 - \tau^2$. Sin embargo, el propio proceso de selección requiere calcular una matriz de Gram completa de $T \times T$ en cada capa, incurriendo en un costo de $O(T^2 d)$. Para longitudes de secuencia moderadas (por ejemplo, $T=2048$), esta sobrecarga de selección es comparable al costo de la atención que pretende reducir, anulando la posible aceleración. El problema central abordado por este artículo es el alto costo computacional del paso de selección de tokens representativos en ADA.

Metodología: El Mecanismo en Cascada
El artículo introduce un mecanismo de "cascada" que explota la coherencia en profundidad del conjunto de tokens representativos. Los autores observan que el conjunto de tokens informativos en la capa $l$ es fuertemente predictivo del conjunto en la capa $l+1$. Esto se fundamenta teóricamente en la continuidad Lipschitz de los bloques transformadores: los tokens que son casi-duplicados (redundantes) o descorrelacionados (representativos) en una capa tienden a mantener esa relación en la siguiente capa, con solo perturbaciones menores a medida que se añade información contextual.

En lugar de recalcular la matriz de Gram completa de $T \times T$ en cada capa, el mecanismo en cascada:

1. Hereda el conjunto representativo $R^{(l)}$ de la capa anterior.
2. Valida el conjunto heredado calculando dos productos de Gram más pequeños:
   • Un producto de $r \times r$ entre los representantes heredados para verificar si alguno se ha vuelto redundante.
   • Un producto cruzado de $(T-r) \times r$ entre los no representativos y los representantes heredados válidos para verificar si algún no representativo se ha vuelto descorrelacionado (es decir, nuevos representantes).
3. Actualiza el conjunto con un pequeño número de adiciones y eliminaciones basadas en estas verificaciones.

Este enfoque reduce la complejidad de la selección de $O(T^2 d)$ a $O(Trd)$ por capa, un factor de aceleración de $T/r$.

Contribuciones Clave

• Innovación Algorítmica: La propuesta de un algoritmo de selección de tokens en cascada que reemplaza el cálculo completo de la matriz de Gram con pasos de validación y actualización incrementales.
• Justificación Teórica: Una demostración (Proposición 3) que prueba que la coherencia en profundidad del conjunto representativo es una consecuencia de la continuidad Lipschitz de los bloques transformadores y la naturaleza aditiva de las conexiones residuales.
• Garantía Conservadora: Una demostración (Proposición 4) que muestra que el método en cascada es conservador; produce un superconjunto del conjunto representativo independiente. Nunca omite un verdadero representante, asegurando que la calidad de la atención no se degrade, aunque puede aumentar ligeramente el número de representantes ($r$).
• Aceleración Ortogonal: El artículo posiciona este método como el análogo en la dimensión de tokens a la Inferencia en Subespacio Puertado (GSI), que acelera las capas lineales mediante la herencia de bases coherentes en profundidad. Juntos, cubren el paso forward completo del transformador (tanto el eje $d$ como el eje $T$).

Resultados Experimentales
El método fue validado en tres familias de modelos (GPT-2 124M, GPT-J 6B, OPT 6.7B) utilizando hardware AMD MI300X con una longitud de secuencia de $T=512$ y un umbral de Gram $\tau=0.30$.

• Coherencia en Profundidad: El índice de Jaccard entre conjuntos representativos en capas consecutivas es alto (0.83 a 0.94) en capas profundas, confirmando la estabilidad del conjunto de tokens.
• Ahorro en Operaciones de Gram: El mecanismo en cascada redujo las operaciones totales de Gram en:
  • 22% para GPT-2 (12 capas)
  • 52% para GPT-J (28 capas)
  • 63% para OPT (32 capas)
    Los ahorros aumentan con la profundidad del modelo porque el costo inicial completo de Gram en la capa 0 se amortiza sobre más capas.
• Tasa de Rotación: En capas profundas, la tasa de rotación (fracción del conjunto modificada) desciende al 3–18%, lo que indica que el conjunto heredado requiere un ajuste mínimo.
• Predicciones de Escalado: Los autores proyectan que para secuencias más largas (por ejemplo, $T=8192$), donde $r$ crece sublinealmente, la aceleración de la selección podría alcanzar $18\times$, reduciendo los costos de selección de un término dominante a uno insignificante.

Significado y Afirmaciones
El artículo afirma que el mecanismo en cascada revela una propiedad estructural de los transformadores entrenados: el conjunto de tokens informativos es una propiedad estructural coherente que se propaga a través de la profundidad de la red. Al explotar esta coherencia, el método hace que la sobrecarga de selección sea práctica para ADA, permitiendo que la compresión de la atención logre su aceleración teórica sin verse obstaculizada por el paso de selección.

Los autores enfatizan que esto es una aceleración complementaria a métodos existentes como FlashAttention (que optimiza el ancho de banda de memoria) y la atención dispersa (que utiliza patrones fijos). A diferencia de los métodos de poda o fusión de tokens que alteran permanentemente las representaciones, la ADA en cascada selecciona dinámicamente representantes basándose en la estructura de entrada sin cambios arquitectónicos ni reentrenamiento. La combinación de GSI y ADA en cascada ofrece una reducción multiplicativa en el costo de inferencia, reduciendo efectivamente la complejidad del paso forward a la "complejidad efectiva" de la entrada en lugar del tamaño nominal del modelo.