Subgroups of $U(d)$ Induce Natural RNN and Transformer Architectures

Resumo Técnico: Subgrupos de U(d) Induzem Arquiteturas Naturais de RNN e Transformer

1. O Problema e Motivação

Os modelos de sequência modernos (como RNNs e Transformers) são construídos a partir de uma complexa interação de escolhas arquitetônicas (parametrização de atualização, camadas de estabilização, mecanismos de memória). A maioria desses modelos opera em espaços vetoriais euclidianos ($\mathbb{R}^n$), onde a estabilidade dinâmica é frequentemente garantida por restrições pós-hoc ou normalizações.

O autor propõe uma mudança de paradigma fundamental: em vez de manter os estados ocultos em $\mathbb{R}^n$ e restringir os operadores, os próprios estados ocultos e as representações de tokens devem ser elementos de um grupo. Especificamente, o trabalho assume que os estados residem em um subgrupo fechado de $U(d)$ (o grupo unitário de dimensão $d$).

A motivação é metodológica:

• Dinâmica Limitada: Grupos matriciais compactos garantem dinâmicas de estado limitadas por construção.
• Geometria Natural: A estrutura de grupo fornece uma álgebra de Lie para atualizações locais e uma medida de similaridade natural baseada em produtos internos matriciais.
• Unificação: O objetivo é derivar arquiteturas completas (RNN e Transformer) diretamente dessa suposição, sem introduzir muitos componentes adicionais.

2. Metodologia e Framework

O artigo apresenta um framework minimalista baseado em axiomas para modelos de sequência com estados ocultos em subgrupos fechados de $U(d)$.

Axiomas Fundamentais:

1. Causalidade: O estado $H_{t+1}$ depende apenas de $(H_t, x_t)$.
2. Fechamento do Grupo: Se $H_t \in G$, então $H_{t+1} \in G$.
3. Atualização Local: Cada passo é gerado por um elemento tangente na álgebra de Lie $\mathfrak{g}$ do grupo.

O Esqueleto do Modelo:
A atualização do estado segue a regra de construção exponencial:
$$H_{t+1} = H_t \exp(\Phi_\theta(H_t, x_t))$$
Onde $\Phi_\theta$ mapeia para a álgebra de Lie $\mathfrak{g}$. Isso garante que o estado permaneça no grupo $G$ devido ao fechamento sob a exponenciação e multiplicação de grupos.

Leitura (Readout):
A saída é calculada usando matrizes protótipo $P_v \in G$ e logits definidos pela similaridade intrínseca do grupo:
$$\ell_v(H_t) = \tau \Re\text{tr}(H_t^* P_v) + b_v$$
Esta é uma leitura nativa do subgrupo, baseada no traço do produto matricial.

Templates Compartilhados (RNN e Transformer):
O autor define templates agnósticos ao subgrupo, onde a escolha do subgrupo atua como um componente "plug-and-play" (drop-in):

• RNN: Atualização recursiva direta no espaço tangente, seguida por exponenciação.
• Transformer: Utiliza atenção baseada em similaridade de grupo ($\Re\text{tr}(H_i^* H_j)$). O modelo realiza dois passos no espaço tangente por camada: um passo de atenção (agregação ponderada) e um passo de ancoragem (ajuste em direção ao embedding do token atual).
• Mistura Linear (Linear Mixing): Uma extensão opcional que aplica um mapa linear aprendido no espaço tangente antes da exponenciação. Isso permite reponderar e rotear atualizações entre direções "visíveis à leitura" e direções de "memória".

Subgrupos Específicos:
O framework é genérico para $U(d)$, $SU(d)$, $O(d)$, $SO(d)$ e toros $T^k$. O papel do subgrupo define:

• O espaço de estado ($G$).
• O espaço tangente ($\mathfrak{g}$).
• A projeção tangente (como os outputs brutos do modelo são mapeados para $\mathfrak{g}$).

3. Contribuições Principais

1. Framework Unificado: Estabelece um esqueleto comum para RNNs e Transformers onde a escolha do subgrupo substitui diretamente a definição do espaço de estado, projeção tangente e mapa de atualização, sem alterar a interface externa do modelo.
2. Mudança de Paradigma: Move a restrição do operador (como em RNNs ortogonais tradicionais) para o próprio estado oculto, que vive na variedade do grupo.
3. Extensão de Mistura Linear: Propõe uma camada de mistura linear no espaço tangente que melhora o desempenho em orçamentos de parâmetros finitos e é aplicável a qualquer escolha de subgrupo.
4. Validação Empírica: Implementa e avalia especificamente a família de subgrupos Ortogonais ($O(d)$), criando modelos de "estado ortogonal" (OSM).

4. Resultados Experimentais

Os experimentos focaram em modelagem de linguagem em nível de caractere nos conjuntos de dados Tiny Shakespeare e Penn Treebank (PTB), comparando os modelos OSM (Orthogonal State Models) com baselines padrão (Transformers e LSTMs) sob orçamentos de parâmetros iguais.

Desempenho (Bits por Caractere - BPC, menor é melhor):

• Tiny Shakespeare (~100K parâmetros):
  • O OSMFormer (com mistura linear e leitura USIM) alcançou 2.464 BPC no teste, superando o Transformer baseline (2.583 BPC).
  • O OSM-RNN (~55K parâmetros) alcançou 2.539 BPC, superando o LSTM baseline (2.594 BPC).
• Penn Treebank (~100K parâmetros):
  • O OSMFormer alcançou 1.614 BPC no teste, superando o Transformer baseline (1.664 BPC).
• Escalabilidade:
  • Em escalas de 100K a 500K parâmetros, os modelos OSMFormer consistentemente superaram os baselines de Transformer em ambos os conjuntos de dados, especialmente em orçamentos menores.
  • A variante com mistura linear mostrou-se a opção de melhor desempenho, mesmo com menos parâmetros que algumas configurações sem mistura.
• Robustez:
  • Os modelos OSM demonstraram maior estabilidade em varreduras de hiperparâmetros (taxas de aprendizado, clipping, weight decay) em comparação aos baselines, mantendo-se treináveis em configurações mais agressivas.

5. Significado e Implicações

• Eficiência e Estabilidade: A abordagem baseada em subgrupos compactos oferece uma via natural para estabilizar a dinâmica de treinamento sem depender excessivamente de técnicas de normalização externas, resultando em modelos mais robustos a variações de hiperparâmetros.
• Interpretabilidade Geométrica: Ao operar no espaço tangente e mapear de volta via exponencial, o modelo oferece uma alça direta sobre a geometria das atualizações. Isso abre caminho para diagnósticos de "atividade de atualização por plano" e análise de estabilidade local.
• Novo Eixo de Design: O trabalho introduz a escolha do subgrupo como um novo eixo de design arquitetônico. Subgrupos abelianos (como toros) versus não-abelianos (como $O(d)$ ou $SU(d)$) podem induzir padrões de interação e memória distintos, sugerindo futuras linhas de pesquisa sobre memória de longo prazo baseada em álgebra de Lie.
• Generalidade: O framework é agnóstico ao subgrupo, permitindo que futuros trabalhos explorem outras variedades (como $SU(d)$ para fases complexas ou $T^k$ para dinâmicas de canal abeliano) usando a mesma infraestrutura de treinamento.

Em conclusão, o artigo demonstra que assumir estados ocultos em subgrupos de $U(d)$ não é apenas uma restrição teórica, mas uma construção prática que induz arquiteturas de RNN e Transformer competitivas, estáveis e geometricamente fundamentadas, com a variante ortogonal ($O(d)$) já superando baselines estabelecidas em tarefas de linguagem padrão.