From Sparsity to Simplicity: Enabling Simpler Sequential Replacements via Sparse Attention Distillation

Resumen Técnico: De la Dispersión a la Simplicidad (S2S)

Enunciado del Problema
Los mecanismos de autoatención son la base del preentrenamiento de transformadores a gran escala, permitiendo una mezcla flexible de tokens y un modelado de dependencias de largo alcance. Sin embargo, su complejidad cuadrática con respecto a la longitud de la secuencia resulta en una alta latencia de inferencia y sobrecarga de memoria. Aunque reemplazar la atención con módulos secuenciales más simples (por ejemplo, RNN, Modelos de Espacio de Estados) es una vía atractiva para la eficiencia, la sustitución ingenua a menudo conduce a una degradación significativa de la precisión, particularmente a escalas mayores. Los enfoques de eficiencia existentes suelen modificar el propio cálculo de la atención o podar tokens, pero mantienen la atención como el operador central. Este artículo aborda el desafío de reemplazar completamente las capas de atención con módulos secuenciales sin una pérdida sustancial en el rendimiento.

Metodología
El artículo propone un marco llamado De la Dispersión a la Simplicidad (S2S), que postula que los transformadores preentrenados descomponen las dependencias complejas de tokens en mapeos por bloques de complejidad efectiva variable. Los autores hipotetizan que las capas que exhiben patrones de atención dispersos son inherentemente más fáciles de aproximar con módulos secuenciales más simples.

La metodología consta de tres componentes centrales:

1. Reemplazo Unificado de Capas de Atención:

   • Los autores introducen un marco plug-and-play donde capas de atención específicas en un Vision Transformer (ViT) preentrenado se intercambian por mezcladores de tokens secuenciales (específicamente BiMamba y BiLSTM), preservando la interfaz de la capa original (LayerNorm, conexiones residuales y MLP).
   • Los módulos secuenciales adoptan una estructura bidireccional de múltiples cabezas para igualar la naturaleza no causal y de contexto completo de la atención ViT.
   • Se utiliza una supervisión alineada por capas durante la destilación: la salida del mezclador secuencial reemplazado se compara directamente con la salida de la capa de atención original a la misma profundidad utilizando una pérdida de error cuadrático medio (MSE).

2. Destilación Guiada por Dispersión:

   • Para controlar y estudiar el impacto de la dispersión, los autores integran A-ViT (ViT Adaptativo), que utiliza puntuaciones de parada para generar máscaras binarias de retención de tokens.
   • Etapa 1 (Destilación): El modelo estudiante (con reemplazos secuenciales) se entrena bajo las mismas máscaras de tokens por capa generadas por el profesor (A-ViT). Esto asegura que el estudiante y el profesor operen sobre los mismos tokens activos, aislando la dificultad de aproximar la función de mezcla de tokens de las diferencias en el enrutamiento de tokens. La pérdida incluye pérdida de similitud, un término de alineación de parada y pérdida de clasificación.
   • Etapa 2 (Ajuste Fino): El estudiante se ajusta finamente utilizando sus propios módulos de parada aprendidos para generar máscaras, permitiéndole operar de forma autónoma mientras mantiene el rendimiento en la tarea.

3. Estrategia de Reemplazo Disperso Consciente de Capas:

   • En lugar de un reemplazo uniforme, la estrategia se dirige a grupos de capas específicos basándose en sus perfiles de dispersión natural. Los autores reemplazan solo las capas más profundas (que exhiben la mayor dispersión natural) y aplican el marco de destilación guiada por dispersión a estos bloques específicos.

Contribuciones Clave

• La Hipótesis S2S: El artículo establece que la dispersión de la atención es una firma de la complejidad del mapeo. Los bloques dispersos en transformadores preentrenados pueden ser reemplazados por módulos secuenciales más simples con pérdida mínima, mientras que los bloques densos son más difíciles de aproximar.
• Marco de Reemplazo Controlado: Un método unificado para intercambiar atención con módulos secuenciales (Mamba/LSTM) manteniendo la compatibilidad de la interfaz y habilitando el análisis selectivo por capas.
• Control Explícito de Dispersión: La introducción de un pipeline de destilación donde la dispersión inducida por el profesor (a través de máscaras A-ViT) se alinea explícitamente con el estudiante, demostrando que aumentar la dispersión del profesor reduce consistentemente la brecha de rendimiento entre estudiante y profesor.
• Compensación Optimizada: Una estrategia que combina la dispersión natural por capas con la dispersión explícita para lograr mejores compensaciones precisión-eficiencia que las líneas base de atención dispersa.

Resultados Experimentales
Los experimentos se realizaron en ImageNet utilizando backbones DeiT (Tiny, Small, Base) con BiMamba y BiLSTM como sustitutos.

• Dispersión Natural y Reemplazabilidad: Las visualizaciones de los mapas de interacción de tokens muestran que, mientras que las capas superficiales exhiben una mezcla global densa, las capas más profundas se vuelven notablemente más dispersas y selectivas. Los experimentos controlados que reemplazan diferentes grupos de capas revelan que reemplazar capas más profundas (más dispersas) incurre en caídas de precisión significativamente menores que reemplazar capas más superficiales (más densas).
• Efecto de la Dispersión Explícita: Cuando la relación de retención de tokens se reduce mediante A-ViT, la brecha de precisión entre el profesor basado en atención y el estudiante secuencial se reduce consistentemente. Esto confirma que la dispersión explícita simplifica el mapeo objetivo, facilitando que los módulos secuenciales aprendan.
• Compensación Precisión-Eficiencia: El reemplazo disperso consciente de capas propuesto (reemplazando solo las 3 capas más profundas con BiMamba bajo presupuestos de dispersión coincidentes) logra:
  • Precisión comparable o ligeramente superior a las líneas base de atención dispersa (A-ViT) en tasas de retención de tokens similares.
  • Ganancias significativas de eficiencia: Hasta 1.71× de aceleración estimada a una resolución de 384² en comparación con la línea base densa, con un mayor rendimiento de mezcla de tokens que los modelos de atención dispersa en conteos de tokens más grandes.

Significado y Afirmaciones
El artículo afirma que el reemplazo de atención debe tratarse como un problema estructurado y dependiente de la capa en lugar de un intercambio de arquitectura uniforme. El significado del trabajo radica en demostrar que:

1. La dispersión es una señal de reemplazabilidad: Los modelos preentrenados descomponen naturalmente la complejidad de tal manera que las capas dispersas son las más propicias para la simplificación.
2. La dispersión explícita ayuda a la destilación: Forzar la dispersión en el profesor hace que la tarea de aprendizaje para el estudiante secuencial sea más fácil, reduciendo la brecha de rendimiento.
3. El reemplazo estratégico supera a los enfoques uniformes: Al centrar el reemplazo en capas naturalmente dispersas y guiar el proceso con dispersión controlada, se puede lograr un equilibrio superior de precisión y latencia en comparación con simplemente podar la atención o reemplazar todo el modelo.

Los autores concluyen que esta perspectiva fomenta futuras investigaciones sobre modelos híbridos que aprovechen el conocimiento de transformadores preentrenados mientras integran módulos de secuencia más simples y eficientes donde la estructura interna del modelo lo permite.