Resumo Técnico: GeoNorm: Unificando Pre-Norm e Post-Norm com Otimização Geodésica

Declaração do Problema

A colocação de camadas de normalização dentro de arquiteturas Transformer — especificamente a escolha entre Pre-Norm e Post-Norm — continua sendo uma decisão de design crítica que influencia a estabilidade do treinamento, o desempenho do modelo e o fluxo de gradiente. Embora o Pre-Norm tenha se tornado a escolha predominante em modelos de Linguagem de Grande Escala (LLMs) modernos devido à sua estabilidade, ele introduz gradientes desequilibrados onde as camadas inferiores recebem atualizações maiores do que as superiores. Por outro outro lado, o Post-Norm, originalmente projetado para facilitar o treinamento de redes profundas, é propenso à instabilidade de treinamento e picos de perda (loss spikes). Soluções existentes como DeepNorm e SandwichNorm melhoram a estabilidade por meio de inicialização de parâmetros ou camadas de normalização adicionais, mas carecem de um arcabouço teórico unificado para analisar esses mecanismos.

Metodologia

Os autores propõem o GeoNorm, um novo método de normalização fundamentado na otimização de variedades (manifold optimization). A percepção teórica central é reinterpretar as operações padrão de atualização do Transformer ($x + \text{FFN}(x)$ e $x + \text{Attention}(x)$) como passos de otimização em uma esfera. Nesta visão:

1. Embeddings de tokens atuam como o estado inicial em uma variedade.
2. Módulos de Atenção e FFN geram direções de atualização (gradientes estimados).
3. Normalização é vista como uma projeção de volta para um conjunto de restrição (a esfera).

A normalização padrão (Post-Norm) realiza uma atualização não restringida seguida por uma projeção euclidiana de volta à esfera. Os autores argumentam que essa projeção descarta informações geométricas e pode distorcer as direções de atualização. Em vez disso, o GeoNorm substitui o passo de projeção por uma atualização geodésica usando o mapa exponencial.

Para um ponto $x$ em uma esfera e um vetor tangente $v$ (a direção de atualização projetada), a atualização é realizada intrinsecamente ao longo da geodésica:
$$\text{exp}_x(v) = \cos\left(\frac{\|v\|}{\|x\|}\right)x + \|x\|\sin\left(\frac{\|v\|}{\|x\|}\right)\frac{v}{\|v\|}$$

Esta formulação garante que a representação se mova suavemente ao longo da variedade, preservando a geometria intrínseca do espaço de características. Além disso, os autores introduzem um mecanismo de decaimento de atualização camada a camada para o tamanho do passo ($\alpha_k$), análogo a cronogramas de taxa de aprendizado (learning rate schedules). Eles avaliam três estratégias de decaimento:

• Harmônica: $\alpha_k = \alpha / k$
• Raiz Quadrada: $\alpha_k = \alpha / \sqrt{k}$
• Linear: $\alpha_k = \alpha (T - k) / T$

O artigo também demonstra teoricamente que o Pre-Norm pode ser interpretado como um caso especial do GeoNorm com ângulos dependentes da profundidade e não adaptativos, unificando assim as duas abordagens sob um único arcabouço geométrico.

Principais Contribuições

1. Arcabouço Teórico: Os autores fornecem uma perspectiva teórica unificada interpretando as camadas Transformer como passos de otimização iterativos em uma variedade esférica, estabelecendo uma ponte entre Pre-Norm e Post-Norm.
2. Algoritmo GeoNorm: Eles introduzem o GeoNorm, que substitui a normalização padrão baseada em projeção por atualizações geodésicas via mapa exponencial. Este método não requer parâmetros treináveis adicionais e adiciona um overhead computacional negligenciável.
3. Decaimento por Camada: Uma proposta inovadora de decaimento de atualização camada a camada para componentes de atenção e FFN para estabilizar a dinâmica de treinamento.
4. Validação Empírica: Experimentos extensos em vários tamanhos de modelo (125M a 1.3B), datasets (Arxiv, Books3, FinWeb-Edu) e comprimentos de treinamento demonstram a eficácia do método.

Resultados Experimentais

O artigo relata que o GeoNorm supera consistentemente os métodos de normalização existentes (Pre-Norm, Post-Norm, DeepNorm, SandwichNorm) em todas as configurações testadas:

• Desempenho: Nos datasets Arxiv e Books3 com um modelo de 125M, o GeoNorm alcançou a menor perda de validação (1.8792 no Arxiv; 3.4040 no Books3) em comparação com todos os baselines.
• Escalabilidade: A vantagem de desempenho persiste e aumenta à medida que o tamanho do modelo cresce. Para um modelo de 1.3B, o GeoNorm alcançou uma perda de 3.1759, superando o melhor baseline (Pre-Norm com 3.2067).
• Estabilidade: O GeoNorm demonstrou superior estabilidade de treinamento, particularmente em execuções de treinamento longas (até 4096 de comprimento de sequência). Enquanto o Post-Norm exibiu picos de perda e colapso de treinamento em comprimentos maiores, o GeoNorm manteve uma convergência estável.
• Mecanismo de Decaimento: Entre as estratégias de decaimento, o decaimento Harmônico ($\alpha/k$) produziu o melhor desempenho, estabelecendo novos marcos nos datasets testados.

Significância e Alegações

O artigo afirma que o GeoNorm oferece uma integração contínua em arquiteturas Transformer padrão, não exigindo modificações nos blocos principais (Atenção ou FFN) e não introduzindo parâmetros treináveis adicionais. Os autores afirmam que, ao alavancar a otimização geodésica, o GeoNorm permite que as representações se movam de forma mais coerente ao longo da variedade esférica, levando a uma melhor propagação de informação e estabilidade de treinamento.

O trabalho posiciona-se como um passo fundamental para entender mecanismos de normalização através da lente da geometria de Riemann. Sugere que o sucesso empírico do Pre-Norm pode ser explicado por seu decaimento de passo implícito, um mecanismo que o GeoNorm torna explícito e adaptativo. Os autores concluem que esta visão unificada fornece uma base robusta para futuras inovações no design de arquiteturas Transformer, particularmente para escalar modelos para maiores profundidades e complexidades.