Resumen Técnico: FourierQK

Planteamiento del Problema

Los mecanismos de atención estándar en los Transformers computan puntuaciones de similitud por pares mediante productos escalares de las proyecciones aprendidas de Consulta ($Q$) y Clave ($K$) ($e_{ij} = q_i \cdot k_j$). Esto opera enteramente dentro del espacio de incrustación (embedding) aprendido, comparando representaciones de un solo token. El artículo investiga si transformar $Q$ y $K$ hacia un espacio de representación diferente —específicamente, aplicando un preprocesamiento espectral en el dominio de la frecuencia— antes de computar estas puntuaciones de similitud puede mejorar el rendimiento del modelado de lenguaje.

Un desafío crítico en este dominio es distinguir el aprendizaje de secuencias genuino del "filtración posicional" (positional leakage). Muchos métodos espectrales (particularmente aquellos que utilizan la FFT inversa) corren el riesgo de sufrir artefactos de frontera circulares donde la información de tokens futuros se filtra hacia posiciones pasadas debido a la naturaleza periódica de la Transformada Discreta de Fourier (DFT). El artículo busca determinar si el preprocesamiento espectral ofrece un beneficio genuino sobre la atención estándar y verificar rigurosamente que las mejoras no sean artefactos de filtración de información.

Metodología

Arquitectura de Atención Espectral

El enfoque propuesto FourierQK aplica un filtro de dominio de frecuencia $\phi(\omega)$ a las proyecciones $Q$ y $K$ antes de la computación de la puntuación de atención.

1. Transformación: Las proyecciones aprendidas $q, k \in \mathbb{R}^{T \times d}$ se transforman mediante DFT: $\hat{q}(\omega) = \mathcal{F}[q](\omega)$.
2. Filtrado: Se aplica un filtro en el dominio de la frecuencia: $\hat{\tilde{q}}(\omega) = \hat{q}(\omega) \cdot \phi(\omega)$.
3. Reconstrucción: Las representaciones filtradas se transforman de nuevo al dominio del tiempo mediante la DFT inversa (IDFT): $\tilde{q} = \mathcal{F}^{-1}[\hat{\tilde{q}}]$.
4. Puntuación: Las puntuaciones de atención se computan utilizando las representaciones filtradas: $e_{ij} = \tilde{q}_i \cdot \tilde{k}_j / \sqrt{d}$.

El artículo pone a prueba varias variantes de filtros:

• Random-QK: Filtros espectrales fijos y no aprendidos (tipo Morlet) para probar si cualquier mezcla espectral ayuda.
• Fourier-QK: Un selector gaussiano suave con una única frecuencia dominante aprendida $f^*$ por cabeza.
• MultiFourier-QK: Múltiples frecuencias aprendidas por cabeza para capturar estructuras de multiescala.
• Filtros Causales: Convoluciones causales en el dominio del tiempo (Gaussian, Mexican Hat, Causal Morlet) para probar si un filtrado local y sin filtraciones proporciona beneficios similares.

Diagnósticos de Verificación

Para asegurar que los resultados no se deban a la filtración posicional (donde el modelo "hace trampa" viendo tokens futuros), los autores emplean un diagnóstico de validación con permutación (shuffled validation diagnostic). Los modelos se evalúan en un conjunto de validación donde el orden de los tokens ha sido permutado.

• Aprendizaje Genuino: Una gran brecha entre la pérdida de validación ordenada y la pérdida de validación permutada indica que el modelo depende del orden temporal.
• Filtración (Leakage): Una brecha pequeña (pérdida similar en datos ordenados y permutados) indica que el modelo está explotando información futura filtrada.

Además, el artículo distingue su enfoque de FNet (Lee-Thorp et al., 2021). Mientras que FNet reemplaza todo el mecanismo de atención con una mezcla de Fourier global de las incrustaciones de los tokens (eliminando las proyecciones $Q/K$ y la matriz de puntuación), FourierQK preserva la estructura de atención estándar (proyecciones $Q/K$ aprendidas, matriz de puntuación, máscara causal) y aplica el filtrado espectral solo a las proyecciones.

Resultados Clave

1. El Preprocesamiento Espectral Mejora el Rendimiento

En el conjunto de datos de nivel de carácter TinyShakespeare ($T=256$), el preprocesamiento espectral produce ganancias sustanciales sobre la atención de producto punto estándar (BASE-DOT, val $\approx$ 1.474):

• Filtro Espectral Aleatorio: Un filtro aleatorio fijo y no aprendido logra val = 1.031 ($\Delta = +0.443$).
• Frecuencia Única Aprendida: Un modelo con una frecuencia aprendida inicializada a escala de párrafo logra val = 0.608 ($\Delta = +0.867$).
• Atención de Multi-Frecuencia: Cuatro frecuencias aprendidas que abarcan escalas de párrafo a palabra logran val = 0.309 ($\Delta = +1.166$), una reducción del 79% en la pérdida de validación.

Estos resultados se confirman a través de múltiples semillas aleatorias (media val = 0.236 para frecuencia única, std = 0.019), estableciendo la reproducibilidad.

2. El Mecanismo: Coherencia de Fase y Mezcla Global

• La Fase es Esencial: Los estudios de ablación muestran que preservar el espectro de amplitud mientras se randomiza la fase resulta en un rendimiento peor que la atención estándar. La ganancia es impulsada por la estructura temporal coherente de fase, no solo por la distribución de energía.
• Global vs. Local: Las ganancias son específicas de la FFT bilateral (mezcla global de la secuencia). Los filtros causales en el dominio del tiempo (Gaussian, Mexican Hat, Causal Morlet) fallan en mejorar la atención estándar a escala de caracteres.
  • Razón: Los filtros causales con campos receptivos limitados ($K \le T/2$) no pueden capturar la coherencia de fase global necesaria para los patrones a escala de párrafo. La FFT bilateral es estructuralmente no causal, pero proporciona el contexto global necesario.
• No es Distorsión de Métrica: Las rotaciones ortogonales aleatorias y las proyecciones no ortogonales de $Q/K$ (sin filtrado espectral) no producen ganancias, confirmando que el beneficio proviene de la mezcla en el dominio de la frecuencia, no de un simple remapeo de la representación.

3. Ordenamiento Multi-Escala

Las frecuencias aprendidas en el modelo MultiFourier convergen hacia una jerarquía casi geométrica:

• Frecuencias: ~49, 27, 10 y 6 tokens por ciclo.
• Interpretación: Estas corresponden a las escalas de párrafo, sub-párrafo, frase y palabra en textos dramáticos.
• Validación: La Descomposición de Modos Empíricos (EMD) aplicada post-entrenamiento recupera de forma independiente estas mismas escalas, confirmando que el modelo descubre una estructura genuina de multiresolución en el espacio $Q/K$.

4. Filtración e Implementación de Restricciones

• Filtración por Irfft: El artículo reporta un "resultado negativo" respecto a la correlación cruzada de la wavelet de Morlet usando irfft. La naturaleza bilateral del núcleo ($\kappa(-\tau) = \kappa(\tau)$) acopla intrínsecamente cada posición con tokens futuros, causando una filtración catastrófica (val $\approx$ 0.018) que desaparece solo cuando se aplica el diagnóstico de brecha de permutación.
• Solución FourierQK: La implementación exitosa de FourierQK evita la reconstrucción mediante irfft de una secuencia indexada por posición. En su lugar, colapsa la dimensión de frecuencia a un escalar ponderado por posición, evitando la filtración inherente a la reconstrucción completa de la señal.

Significancia y Reivindicaciones

El artículo sostiene que el preprocesamiento espectral de las proyecciones de Consulta y Clave es una modificación arquitectónica distinta y efectiva que mejora la atención de los Transformers al permitir que el modelo detecte estructuras temporales coherentes de fase en múltiples escalas lingüísticas.

Las reivindicaciones clave incluyen:

1. Especificidad Espectral: La mejora se debe estrictamente a la mezcla global en el dominio de la frecuencia, no a transformaciones lineales aleatorias o distorsión de métrica.
2. Límite Arquitectónico: Existe un límite claro entre la atención espectral de mezcla global (que funciona pero es no causal debido a la FFT bilateral) y la atención espectral genuinamente causal (que falla a escala de caracteres debido al aliasing y los campos receptivos limitados, pero podría funcionar a escala de palabra).
3. Distinción de FNet: A diferencia de FNet, que reemplaza la atención, este método potencia el mecanismo de atención existente, permitiendo que la jerarquía de frecuencias emerja de las propias puntuaciones de atención.
4. Dominancia de la Fase: El mecanismo depende de la coherencia de fase (alineación temporal de los componentes de frecuencia) más que de la amplitud (energía) por sí sola.

El trabajo identifica una limitación precisa: aunque el preprocesamiento por FFT bilateral es altamente efectivo, su naturaleza no causal impide su aplicación directa en modelos de lenguaje causales sin filtración. Los autores sugieren que el trabajo futuro sobre la atención espectral causal debe abordar las limitaciones de aliasing y del campo receptivo observadas en los filtros de tiempo, posiblemente moviéndose hacia una tokenización a escala de palabra donde las restricciones de Nyquist sean menos severas.