RotRNN: Modelling Long Sequences with Rotations

Resumo Técnico: RotRNN – Modelagem de Sequências Longas com Rotações

Declaração do Problema
A modelagem de sequências longas continua sendo um domínio desafiador no aprendizado de máquina devido à dificuldade de capturar dependências de longo alcance. Enquanto os Transformers sofrem com o escalonamento computacional quadrático e vieses indutivos pobres, e as RNNs não lineares lutam contra o desaparecimento e a explosão de gradientes, modelos recentes de Espaços de Estados (SSMs) e Unidades Recorrentes Lineares (LRUs) alcançaram um desempenho de estado da arte. No entanto, esses modelos exibem desvantagens significativas:

1. Inicialização Complexa: Os SSMs como o S4 dependem de derivações teóricas intrincadas (ex: inicialização HiPPO baseada em projeções polinomiais ótimas) que não são bem compreendidas durante o processo de treinamento.
2. Descompasso entre Implementação e Teoria: O LRU, embora conceitualmente mais simples, não reflete totalmente sua motivação teórica na prática. Especificamente, ele falha em impor pares conjugados para autovalores de matrizes reais, carece de restrições para garantir saídas de valores reais e depende de um normalizador aprendido que não garante magnitudes consistentes do estado oculto ao longo do treinamento.
3. Estabilidade: Métodos existentes frequentemente exigem esquemas de normalização complexos ou procedimentos de inicialização que podem não ser robustos através de diferentes profundidades e sementes aleatórias.

Metodologia
Os autores propõem o RotRNN (Rede Neural Recorrente Rotacional), um modelo recorrente linear que parametriza a matriz de estado recorrente $A$ como uma matriz de rotação. A metodologia é construída sobre três pilares centrais:

1. Parametrização via Matrizes de Rotação:
   Em vez de aprender uma matriz restrita diretamente, o RotRNN aprende uma matriz de peso geral $M$ e a mapeia para o grupo ortogonal especial $SO(N)$ usando a exponencial de matriz de uma matriz antissimétrica ($A = \exp(M - M^\top)$). Isso garante que $A$ permaneça uma matriz de rotação válida durante todo o treinamento.
   Para permitir a computação eficiente de potências de matrizes (necessária para a forma convolucional da recorrência linear), os autores decompõem a matriz de rotação $A$ em uma forma de bloco-diagonal: $A = P \Theta P^\top$. Aqui, $P$ é uma matriz ortogonal e $\Theta$ consiste em blocos de rotação $2\times2$ definidos por ângulos $\theta_i$. Essa estrutura permite o cálculo rápido de $A^k$ sem operações densas dispendiosas de matrizes.

2. Normalização Principiada:
   Aproveitando a ortogonalidade das matrizes de rotação (onde os autovalores residem no círculo unitário), os autores derivam um procedimento de normalização que mantém uma norma constante do estado oculto esperado ($E[\|x_t\|^2] = c$) ao longo do treinamento.
   A recorrência é definida como $x_t = \gamma A x_{t-1} + \xi B u_t$, onde $\gamma$ é um fator de decaimento escalar aprendível e $\xi$ é um coeficiente de reescalonamento derivado para satisfazer a condição de norma constante. Diferente do LRU, que normaliza com base em limites de sequências infinitas, a normalização do RotRNN mantém-se para todos os passos de tempo, evitando o decaimento exponencial de entradas iniciais frequentemente observado quando constantes de normalização são elevadas à potência de $t$.

3. Arquitetura Multi-Head (Multi-Cabeça):
   Para abordar a limitação de um único fator de decaimento escalar $\gamma$ (que pode não generalizar bem para tarefas que exigem diferentes horizontes de memória), o RotRNN emprega $H$ cabeças de baixa dimensão independentes. Cada cabeça opera com sua própria matriz de rotação $A^{(h)}$, matriz de entrada $B^{(h)}$ e fator de decaimento $\gamma^{(h)}$. As saídas são concatenadas e misturadas linearmente, permitindo que o modelo capture dependências em múltiplas escalas.

Principais Contribuições

• Nova Camada de RNN Linear: A proposta de uma camada recorrente onde a matriz de estado é parametrizada como uma rotação, garantindo que o algoritmo permaneça fiel à sua motivação teórica durante todo o treinamento.
• Computação Eficiente: Um método para computar potências de matrizes paramétricas de forma eficiente através da decomposição em formas bloco-diagonais, permitindo uma recorrência linear rápida.
• Normalização Robusta: Uma derivação principiada de um procedimento de normalização que retém uma magnitude constante do estado oculto esperado, oferecendo garantias de estabilidade mais fortes do que trabalhos anteriores que dependem de convergência assintótica.
• Insight Teórico: Uma derivação mostrando a equivalência matemática entre um caso especial do LRU (onde os autovalores formam pares conjugados) e o RotRNN multi-head, potencialmente lançando luz sobre a mecânica interna do LRU.

Resultados
Os autores avaliaram o RotRNN em vários benchmarks de modelagem de sequências longas:

• Long Range Arena (LRA): O RotRNN demonstrou um desempenho competitivo contra modelos recorrentes lineares de estado da arte (S4, S5, LRU). Ele alcançou as pontuações mais altas entre as bases de comparação em tarefas de entrada discreta, como ListOps, Texto (IMDB) e Recuperação (Retrieval). No entanto, teve um desempenho ligeiramente inferior a algumas bases em tarefas de imagem de nível de pixel (Path-X, CIFAR).
• Speech Commands: Na tarefa de classificação de comandos de voz de 35 vias, o RotRNN teve um desempenho idêntico ao LRU reproduzido (95,2% de precisão) com uma contagem de parâmetros semelhante, permanecendo competitivo com modelos de espaço de estados mais complexos.
• Análise de Estabilidade: A análise empírica das normas do estado oculto no conjunto de dados ListOps revelou que o RotRNN mantém uma magnitude quase constante através do treinamento, profundidade de camada e sementes aleatórias. Em contraste, o LRU exibiu variações selvagens nas normas do estado oculto e levou significativamente mais tempo para convergir para uma magnitude estável.

Significância e Alegações
O artigo afirma que o RotRNN oferece uma alternativa conceitualmente simples, eficiente e matematicamente fundamentada aos modelos recorrentes lineares existentes. Sua principal significância reside em preencher a lacuna entre a motivação teórica e a implementação prática. Ao usar matrizes de rotação, o RotRNN evita os esquemas de inicialização complexos dos SSMs e as inconsistências de implementação do LRU. Os autores afirmam que seu modelo fornece um procedimento de normalização robusto que garante estabilidade sem depender de inicializações de "caixa preta", alcançando simultaneamente um desempenho competitivo em benchmarks padrão. Além disso, o trabalho sugere que visualizar o LRU através da lente de estruturas de blocos baseadas em rotação pode ajudar a esclarecer a mecânica da recorrência linear em modelos anteriores.