RotRNN: Modelling Long Sequences with Rotations

Resumen Técnico: RotRNN – Modelado de Secuencias Largas con Rotaciones

Planteamiento del Problema
El modelado de secuencias largas sigue siendo un dominio desafiante en el aprendizaje automático debido a la dificultad de capturar dependencias de largo alcance. Mientras que los Transformers sufren una escalabilidad computacional cuadrática y sesgos inductivos deficientes, y las RNN no lineales luchan contra el desvanecimiento y la explosión de gradientes, los modelos recientes de Modelos de Espacio de Estados (SSM) y Unidades Recurrentes Lineales (LRU) han logrado un rendimiento de vanguardia. Sin embargo, estos modelos presentan desventajas significativas:

1. Inicialización Compleja: Los SSM como S4 dependen de derivaciones teóricas intrincadas (por ejemplo, la inicialización HiPPO basada en proyecciones polinómicas óptimas) que no se comprenden bien durante el proceso de entrenamiento.
2. Desajuste entre Implementación y Teoría: La LRU, aunque es conceptualmente más simple, no refleja plenamente su motivación teórica en la práctica. Específicamente, no logra imponer pares conjugados para los autovalores de matrices reales, carece de restricciones para asegurar salidas de valores reales y depende de un normalizador aprendido que no garantiza una magnitud constante del estado oculto durante todo el entrenamiento.
3. Estabilidad: Los métodos existentes a menudo requieren esquemas de normalización complejos o procedimientos de inicialización que pueden no ser robustos a través de diferentes profundidades y semillas aleatorias.

Metodología
Los autores proponen RotRNN (Red Neuronal Recurrente Rotacional), un modelo recurrente lineal que parametriza la matriz de estado recurrente $A$ como una matriz de rotación. La metodología se basa en tres pilares fundamentales:

1. Parametrización mediante Matrices de Rotación:
   En lugar de aprender directamente una matriz restringida, RotRNN aprende una matriz de pesos general $M$ y la mapea al grupo ortogonal especial $SO(N)$ utilizando la exponencial de matriz de una matriz antisimétrica ($A = \exp(M - M^\top)$). Esto asegura que $A$ permanezca como una matriz de rotación válida durante todo el entrenamiento.
   Para permitir el cálculo eficiente de las potencias de la matriz (requerido para la forma convolucional de la recurrencia lineal), los autores descomponen la matriz de rotación $A$ en una forma de bloque diagonal: $A = P \Theta P^\top$. Aquí, $P$ es una matriz ortogonal y $\Theta$ consiste en bloques de rotación de $2\times2$ definidos por ángulos $\theta_i$. Esta estructura permite el cálculo rápido de $A^k$ sin operaciones de matrices densas costosas.

2. Normalización Basada en Principios:
   Aprovechando la ortogonalidad de las matrices de rotación (donde los autovalores residen en el círculo unitario), los autores derivan un procedimiento de normalización que mantiene una norma constante del estado oculto esperado ($E[\|x_t\|^2] = c$) durante todo el entrenamiento.
   La recurrencia se define como $x_t = \gamma A x_{t-1} + \xi B u_t$, donde $\gamma$ es un factor de decaimiento escalar aprendible y $\xi$ es un coeficiente de reescalado derivado para satisfacer la condición de norma constante. A diferencia de la LRU, que normaliza basándose en límites de secuencias infinitas, la normalización de RotRNN se mantiene para todos los pasos de tiempo, evitando el decaimiento exponencial de las entradas tempranas que suele verse cuando las constantes de normalización se elevan a la potencia $t$.

3. Arquitectura Multi-Cabezal:
   Para abordar la limitación de un único factor de decaimiento escalar $\gamma$ (que puede no generalizar bien a tareas que requieren diferentes horizontes de memoria), RotRNN emplea $H$ cabezales independientes de baja dimensión. Cada cabezal opera con su propia matriz de rotación $A^{(h)}$, matriz de entrada $B^{(h)}$ y factor de decaimiento $\gamma^{(h)}$. Las salidas se concatenan y se mezclan linealmente, permitiendo al modelo capturar dependencias a múltiples escalas.

Contribuciones Clave

• Nueva Capa Recurrente Lineal: La propuesta de una capa recurrente donde la matriz de estado se parametriza como una rotación, asegurando que el algoritmo sea fiel a su motivación teórica durante todo el entrenamiento.
• Cálculo Eficiente: Un método para calcular potencias de matrices eficientes de rotaciones paramétricas mediante la descomposición en formas de bloque diagonal, permitiendo una recurrencia lineal rápida.
• Normalización Robusta: Una derivación de un procedimiento de normalización basado en principios que retiene una magnitud constante del estado oculto esperado, ofreciendo garantías de estabilidad más fuertes que trabajos previos que dependen de la convergencia asintótica.
• Perspectiva Teórica: Una derivación que muestra la equivalencia matemática entre un caso especial de la LRU (donde los autovalores forman pares conjugados) y la RotRNN multi-cabezal, lo que potencialmente arroja luz sobre la mecánica interna de la LRU.

Resultados
Los autores evaluaron RotRNN en varios benchmarks de modelado de secuencias largas:

• Long Range Arena (LRA): RotRNN demostró un rendimiento competitivo frente a modelos recurrentes lineales de vanguardia (S4, S5, LRU). Logró las puntuaciones más altas entre los baselines en tareas de entrada discreta como ListOps, Texto (IMDB) y Recuperación (Retrieval). Sin embargo, se desempeñó ligeramente por debajo de algunos baselines en tareas de imagen a nivel de píxel (Path-X, CIFAR).
• Speech Commands: En la tarea de clasificación de Speech Commands de 35 vías, RotRNN funcionó de manera idéntica a la LRU reproducida (95.2% de precisión) con un conteo de parámetros similar, manteniéndose competitiva con modelos de espacio de estados profundos más complejos.
• Análisis de Estabilidad: El análisis empírico de las normas del estado oculto en el conjunto de datos ListOps reveló que RotRNN mantiene una magnitud casi constante a través del entrenamiento, la profundidad de la capa y las semillas aleatorias. En contraste, la LRU exhibió variaciones salvajes en las normas del estado oculto y tardó significativamente más en converger a una magnitud estable.

Significancia y Reivindicaciones
El artículo afirma que RotRNN ofrece una alternativa conceptualmente simple, eficiente y matemáticamente fundamentada a los modelos recurrentes lineales existentes. Su principal significancia radica en cerrar la brecha entre la motivación teórica y la implementación práctica. Al utilizar matrices de rotación, RotRNN evita los esquemas de inicialización complejos de los SSM y las inconsistencias de implementación de la LRU. Los autores aseguran que su modelo proporciona un procedimiento de normalización robusto que garantiza la estabilidad sin depender de inicializaciones de "caja negra", logrando al mismo tiempo un rendimiento competitivo en benchmarks estándar. Además, el trabajo sugiere que ver la LRU a través del lente de las estructuras de bloques basadas en rotaciones puede ayudar a aclarar la mecánica de la recurrencia lineal en modelos previos.