Logo.png

Razonar de manera eficiente a través de secuencias extendidas es una dificultad importante en el aprendizaje automático. Recientemente, las convoluciones han surgido como una primitiva crítica para el modelado de secuencias, respaldando el rendimiento más moderno en modelado de lenguaje, análisis de series temporales, visión por computadora, modelado de ADN y más. A pesar de estos impresionantes hallazgos de calidad y ventajas adicionales, como una estabilidad mejorada y una mejor escalabilidad a medida que aumenta la longitud de la secuencia, los modelos de secuencia convolucional siguen siendo significativamente más lentos que los Transformers.

Una causa principal es el soporte de hardware poco confiable. Las convoluciones para el modelado de secuencias emplean frecuentemente filtros tan largos como la secuencia de entrada, en contraste con los filtros cortos utilizados en las convoluciones clásicas para aplicaciones visuales. El algoritmo de convolución de Transformada Rápida de Fourier (FFT) calcula la convolución entre una entrada u y un núcleo de convolución k mapeando las frecuencias de entrada y salida.

A pesar de ser asintóticamente eficiente, el algoritmo de convolución FFT tiene un tiempo de reloj bajo en los aceleradores contemporáneos. Sin embargo, el progreso tecnológico en los sistemas ha permitido a Transformers alcanzar los límites de los aceleradores actuales, con un uso de FLOP de extremo a extremo superior al 72% cuando se utiliza FlashAttention-v2.

Para ofrecer capacidades de contexto más largo, una nueva investigación de la Universidad de Stanford investiga cómo optimizar el método de convolución FFT en aceleradores contemporáneos. Los investigadores creen que, a medida que los avances en sistemas como FlashAttention condujeron a mejores modelos y nuevos algoritmos de atención, la optimización de la convolución FFT conducirá a nuevos y mejores algoritmos, aumentando la calidad de los modelos de secuencia convolucional.

La convolución FFT se puede optimizar fácilmente para secuencias cortas. Es una práctica común reutilizar los filtros del kernel en varios lotes, lo que permite calcular previamente la FFT del filtro antes de reutilizarlo. Por lo tanto, la convolución FFT es paralela entre lotes y filtros, y la fusión del núcleo permite almacenar en caché las salidas de convolución intermedias en SRAM o registros.

  1. Sin embargo, el equipo destaca que aparecen dos obstáculos importantes a medida que crece la duración de la secuencia. Con respecto a los aceleradores actuales, las convoluciones FFT no utilizan de manera óptima las unidades de multiplicación matriz-matriz especializadas.
  2. En segundo lugar, la fusión del kernel falla porque las secuencias crecen demasiado para caber en la SRAM y se requieren costosas operaciones de E/S. Las operaciones de relleno para la causalidad y las conversiones de entradas/salidas de valor real a intermediarios FFT de valor complejo podrían aumentar aún más estos costos de E/S.

En respuesta, los investigadores ofrecen FlashFFTConv, un algoritmo novedoso que emplea una descomposición Monarch de la FFT para optimizar la convolución de la FFT para secuencias extendidas. La FFT se puede transferir eficazmente al hardware gracias a una descomposición Monarch de orden p, que reescribe la FFT como una serie de p operaciones de multiplicación matriz-matriz. Los valores de p más altos generan menos costos de FLOP debido a matrices más pequeñas, pero requieren más E/S para transmitir resultados intermedios. Por lo tanto, hay una compensación involucrada.

El estudio demuestra cómo optimizar p para el costo FLOP y el costo de E/S en una GPU utilizando un modelo de costo sencillo basado en la longitud de la secuencia. Además de facilitar la fusión del núcleo en secuencias de mayor longitud, esta descomposición reduce la cantidad de secuencia que debe mantenerse en SRAM. Por lo tanto, FlashFFTConv puede manejar fácilmente secuencias de entre 256 y 4 millones de caracteres. Al utilizar un algoritmo FFT de valor real y omitir partes de las operaciones de multiplicación de matrices cuando la entrada tiene relleno con ceros, FlashFFTConv puede reducir la duración de la operación FFT hasta a la mitad. Por último, pero no menos importante, la vista matricial de la convolución FFT proporciona una interfaz simple para implementar dos modificaciones arquitectónicas: convoluciones parciales, que aprenden con un núcleo de convolución que es más corto que la secuencia de entrada, y convoluciones de frecuencia dispersa, que ponen a cero secciones de la convolución FFT. núcleo en el espacio de frecuencias. Ambos enfoques se pueden implementar simplemente omitiendo secciones de la descomposición de la matriz, reduciendo el uso de memoria y el tiempo de ejecución del reloj de pared, y pueden considerarse como paralelos convolucionales de atención escasa/aproximada en Transformers.

Los investigadores demuestran que FlashFFTConv acelera la convolución FFT, lo que da como resultado modelos de secuencia más largos, más eficientes y de mejor calidad.

  • FlashFFTConv mejora la calidad de los modelos de secuencia convolucional a través de una mayor eficiencia: para el mismo presupuesto de cómputo, FlashFFTConv permite a Hyena-GPT-s lograr 2,3 puntos más de perplejidad y permite a M2-BERT-base lograr una puntuación GLUE promedio hasta 3,3 más alta: una ganancia en rendimiento equivalente a duplicar los parámetros del modelo.
  • FlashFFTConv mejora la eficiencia de las convoluciones hasta en 7,93 y hasta en 5,60 en ahorro de memoria en comparación con PyTorch, y esta eficiencia se mantiene en más de cuatro órdenes de magnitud en la longitud de la secuencia. FlashFFTConv es más rápido en tiempo de reloj de pared que FlashAttention-v2 de extremo a extremo para longitudes de secuencia de 2K y más debido a menores costos de FLOP y logra hasta un 62,3% de uso de FLOP de extremo a extremo, que es solo un 10% menos que FlashAttention. -v2.
  • Los modelos de secuencias más largas son posibles con FlashFFTConv. FlashFFTConv ha producido el único modelo capaz de completar el largo trabajo Path-512 del benchmark de arena (longitud de secuencia 256K) para la clasificación de imágenes de alta resolución. FlashFFTConv es el primer modelo que incorpora los genes humanos más largos (hasta 2,3 millones de pares de bases) con una resolución de un solo nucleótido; extiende HyenaDNA a una longitud de secuencia de 4 M mediante convoluciones parciales.

El equipo espera que FlashFFTConv allane el camino para un uso más amplio de modelos de secuencia convolucional y que las lecciones aprendidas conduzcan a arquitecturas informáticas más eficientes en cuanto a recursos.


Revisar la Papel, GitHuby Artículo de blog. Todo el crédito por esta investigación va a los investigadores de este proyecto. Además, no olvides unirte. nuestro SubReddit de 33k+ ML, 41k+ comunidad de Facebook, Canal de discordia, y Boletín electrónicodonde compartimos las últimas noticias sobre investigaciones de IA, interesantes proyectos de IA y más.

Si te gusta nuestro trabajo, te encantará nuestra newsletter.


Dhanshree Shenwai es ingeniero en informática y tiene una buena experiencia en empresas de tecnología financiera que cubren el ámbito financiero, tarjetas y pagos y banca con un gran interés en las aplicaciones de IA. Le entusiasma explorar nuevas tecnologías y avances en el mundo en evolución de hoy que facilita la vida de todos.