El contraataque del FFT: una alternativa eficiente a Self-Attention

• El mecanismo tradicional de Self-Attention tiene una complejidad de O(n²), por lo que su escalabilidad para secuencias largas es limitada
• Este artículo propone FFTNet, que aprovecha la Fast Fourier Transform (FFT)
• FFTNet realiza mezcla global de tokens con complejidad temporal de O(n log n)
• Introduce filtros espectrales aprendibles y la función de activación modReLU en el dominio de la frecuencia para resaltar componentes de frecuencia importantes
• En experimentos de benchmark sobre Long Range Arena (LRA) e ImageNet, muestra un rendimiento superior frente a los modelos tradicionales de Self-Attention y a los modelos con transformadas de Fourier fijas

Investigación relacionada

• Complejidad de Self-Attention: los modelos Transformer requieren O(n²) de cómputo, lo que los vuelve ineficientes para procesar secuencias largas
• Enfoques basados en Fourier: modelos como FNet redujeron el costo computacional usando transformadas de Fourier fijas, pero carecen de adaptabilidad a la entrada
• Técnicas de aproximación lineal, dispersa y de baja dimensión: trabajos como Performer, Linformer y BigBird propusieron métodos para aproximar el cómputo de Self-Attention
• Técnicas de descomposición con matrices ortogonales: el uso de transformaciones ortogonales (incluyendo DFT) mejora la estabilidad del entrenamiento del modelo
• Filtrado espectral adaptativo: al agregar filtros aprendibles a transformaciones basadas en FFT, el método resulta más flexible y con mayor capacidad de representación que los enfoques previos

FFTNet: técnica de filtrado espectral adaptativo

Motivación

• Self-Attention tiene complejidad O(n²) y es ineficiente en secuencias largas
• FFT opera en O(n log n) y puede codificar interacciones globales de manera eficiente

Metodología

• Transformada de Fourier (aplicación de FFT)
  • Convierte la secuencia de entrada al dominio de la frecuencia para capturar de forma eficiente dependencias globales
• Aplicación de filtro espectral adaptativo
  • Usa un vector de contexto global para generar filtros aprendibles y resaltar dinámicamente bandas de frecuencia importantes
• Activación no lineal modReLU
  • Aplica una activación basada en ReLU en el dominio de frecuencia complejo para aumentar la capacidad de representación
• Transformada inversa de Fourier (IFFT)
  • Después de aplicar filtrado y activación a los datos transformados, los convierte de nuevo al dominio temporal

Fundamento teórico de FFTNet

• Permite mezcla global de tokens con un costo computacional de O(n log n)
• Attention adaptativa: los filtros aprendibles en el dominio de la frecuencia ajustan las frecuencias según la entrada dada
• Mayor capacidad expresiva gracias a la activación no lineal: con modReLU, puede aprender patrones de orden superior más allá de transformaciones lineales simples
• Estabilidad garantizada basada en el teorema de Parseval: conserva la energía de la señal y minimiza la pérdida de información

Resultados experimentales

Benchmark Long Range Arena (LRA)

• FFTNet registra una precisión general más alta que Transformer y FNet
• En particular, muestra mejor rendimiento en tareas como ListOps, Text, Retrieval, Image y Pathfinder, y obtiene en promedio la puntuación más alta
• Transformer mostró alto rendimiento en algunas tareas, pero tiene limitaciones para manejar dependencias de largo plazo
• Aunque FNet usa FFT, su esquema de transformación fija carece de adaptabilidad y muestra un rendimiento general inferior
• En particular, en la tarea Path-X, Transformer falló por exceso de memoria (OOM), mientras que FFTNet mostró un rendimiento estable

Experimento de clasificación en ImageNet

• El Vision Transformer basado en FFTNet (FFTNetViT) logra mantener una precisión similar a la de ViT convencional, reduciendo al mismo tiempo de forma importante el costo computacional (FLOPs)
• En el caso del modelo Base, FFTNetViT usa aproximadamente 38% menos FLOPs que ViT y aun así aumenta ligeramente la precisión
• En los modelos Large y Huge, FFTNetViT también mantiene un rendimiento similar al de ViT con menor costo computacional
• Esto confirma que FFTNetViT ofrece una alta eficiencia computacional

Ablation Study (análisis de importancia por componente)

• Se analiza el impacto de distintos elementos de FFTNet sobre el rendimiento del modelo al eliminarlos uno por uno
• Cuantos más componentes principales de FFTNet se eliminan, más tiende a disminuir la precisión
  • Eliminación del spectral gating: al desaparecer la función de resaltar frecuencias específicas, la precisión baja ligeramente
  • Eliminación del módulo adaptativo: al desaparecer la capacidad de ajustar dinámicamente los filtros según la entrada, la precisión cae aún más
  • Uso de convolución en lugar de FFT: al perderse la capacidad de mezclar información global de forma eficiente, se produce la mayor caída de rendimiento
• Esto muestra que cada componente de FFTNet cumple un papel importante en la mejora del rendimiento

Conclusión

• FFTNet es una alternativa con mejor eficiencia computacional que Self-Attention
• Al combinar filtros espectrales adaptativos y modReLU en el dominio de la frecuencia, ofrece una fuerte capacidad de representación
• Los resultados experimentales muestran que supera a los modelos tradicionales de Self-Attention en rendimiento y eficiencia en LRA e ImageNet
• Mantiene una complejidad O(n log n) y aun así ofrece un rendimiento de nivel Self-Attention, lo que lo hace favorable para el procesamiento de secuencias largas
• El Vision Transformer basado en FFTNet (FFTNetViT) también alcanza un rendimiento comparable al de ViT con menos FLOPs

• Repositorio de código: enlace de GitHub

[GN⁺]

Comentarios de Hacker News

• Básicamente usa el teorema de convolución: una convolución costosa en el espacio directo se convierte en una multiplicación simple en el espacio recíproco

  • Cuando hay una operación de convolución en los datos, se transforman al dominio conjugado para convertirla en multiplicación
  • Es decir, se trabaja con los datos en un dominio que les resulta natural

• Google presentó en 2022 la idea llamada "FNet: Mixing Tokens with Fourier Transforms"

  • Después descubrieron que sus TPU eran más rápidas que la FFT para multiplicaciones de matrices en la mayoría de los escenarios

• La transformada de Fourier se realiza sobre la dimensión de los "tokens". Pero en muchas aplicaciones, esa dimensión no tiene significado

  • Por eso, los transformadores son una excelente opción para procesar datos invariantes a la permutación
  • Me gustaría ver más experimentos usando transformadas de Fourier sobre grupos finitos menos conocidos
  • Si esto se convierte en lo próximo grande para los LLM, me pregunto qué tan fácil sería integrarlo en motores de inferencia como vLLM, llama.cpp, etc.

• Las matemáticas son demasiado difíciles y cuesta entenderlas. Me pregunto si alguien podría explicar en inglés sencillo cómo esto es equivalente al mecanismo de atención, de qué frecuencias está hablando y cómo codifica las relaciones posicionales entre tokens

• No veo cómo se podría encajar el enmascaramiento causal en este framework. Tampoco hay mención de embeddings posicionales, así que parece que la implementación de self-attention con la que se compara es un NoPE no causal

  • Si los resultados hubieran estado cerca del estado del arte, probablemente el autor lo habría mencionado

• No hay ninguna mención del Hyena Operator, que ya demostró hace unos años una mezcla de contexto completo O(n log n)

• En la era de la telemetría, me parece un gran error no aplicar FFT a la telemetría en la nube para detectar epiciclos y sistemas cuasiestables antes de que provoquen drama

  • "Es más probable que los SLA se incumplan entre 23 y 25 minutos después de desplegar el servicio. Me pregunto por qué... ah, no."

• Me pregunto si alguien tiene una intuición de por qué ayuda ver las cosas en el dominio de la frecuencia

  • Entiendo el término DC, pero no esperaría que los datos de entrada fueran tan periódicos como para que otras frecuencias tengan significado

• Entiendo la notación Big O hasta cierto punto, pero como la mayoría de las cosas relacionadas con computación o ingeniería eléctrica, esto también me cuesta entenderlo

  • Como alguien muy débil en matemáticas, envidio a quienes pueden entender o aprender estas cosas
  • Lo que sé sobre FFT es que transforma señales, que se usa en cierto procesamiento de señales y que en el pasado tuvo un papel importante para detectar explosiones nucleares

• No entiendo por qué se necesita atención. Una capa totalmente conectada también puede "prestar atención" a todas las entradas

  • En conjuntos de datos muy pequeños (0 - 500 tokens), la atención hace que el entrenamiento tarde más y empeora los resultados
  • Parece que las ventajas aparecen en conjuntos de datos más grandes
  • Soy principiante en IA y estoy haciendo un proyecto personal de IA, así que no es una referencia exacta