Resumo Técnico: Estabilidade de Transformers sob Normalização de Camada

Declaração do Problema

Apesar de seu domínio no aprendizado profundo moderno, o treinamento de Transformers profundos permanece propenso à instabilidade, particularmente à medida que a profundidade do modelo aumenta. Essas instabilidades manifestam-se em duas áreas críticas: instabilidade direta (crescimento descontrolado dos estados ocultos levando a problemas numéricos) e instabilidade reversa (explosão de gradientes durante a retropropagação).

Um componente arquitetônico primário usado para mitigar esses problemas é a Normalização de Camada (LN). No entanto, a colocação da LN dentro dos blocos Transformer evoluiu amplamente através de heurísticas empíricas em vez de uma justificativa teórica rigorosa. O artigo identifica três configurações principais:

• Post-LN: LN aplicada após as conexões residuais. Frequentemente requer um escalonamento de otimização delicado e exibe desempenho subótimo.
• Pre-LN: LN aplicada antes dos módulos de atenção e feedforward. Embora seja o padrão para sua melhoria na estabilidade de treinamento em relação ao Post-LN, observa-se empiricamente que produz estados ocultos excessivamente grandes, levando a potenciais instabilidades numéricas e explosão de gradientes em redes profundas.
• Peri-LN: Uma configuração mais recente que aplica LN tanto na entrada quanto na saída dos módulos. Estudos empíricos sugerem que ela oferece estabilidade superior e estados ocultos mais regulares em comparação ao Pre-LN, mas suas propriedades teóricas permanecem pouco compreendidas.

O problema central abordado é a falta de um arcabouço teórico principado para explicar por que essas diferentes colocações levam a dinâmicas de treinamento e perfis de estabilidade distintos, e se o treinamento conduz os modelos para soluções regulares ou comportamentos patológicos.

Metodologia

Os autores propõem um novo arcabouço teórico fundamentado na teoria de controle ótimo de tempo discreto para analisar a dinâmica de treinamento de Transformers.

1. Formulação como Controle Ótimo: O problema de treinamento do Transformer é formulado como um problema de controle de campo médio em tempo discreto. O índice da camada $i$ representa o tempo discreto, e o estado oculto $X_i$ representa o estado do sistema. As conexões residuais atuam como a velocidade das dinâmicas.
2. Análise de Hamilton-Jacobi-Bellman (HJB): Os autores analisam as condições de otimalidade do problema de treinamento usando a equação diferencial parcial (PDE) HJB acoplada com equações de transporte de massa.
   • Eles examinam o Hamiltoniano, que atua como um mecanismo de controle de feedback definindo o campo de velocidade ótimo para os estados ocultos.
   • Para o Pre-LN, eles demonstram que o Hamiltoniano é ilimitado, levando a condições de otimalidade mal postas onde os estados ocultos podem crescer arbitrariamente grandes.
   • Para o Peri-LN, a normalização de saída restringe o campo de velocidade a um elipsoide (baseado no Lema 1), garantindo que o Hamiltoniano seja limitado e o problema seja bem posto.
3. Análise de Gradiente: O artigo analisa a estabilidade reversa examinando a propagação de gradientes através do produto de sensibilidades locais (matrizes Jacobianas) através das camadas.
4. Escalonamento de Resíduo: O arcabouço é estendido para incluir um parâmetro de incremento de tempo $\Delta t < 1$, escalonando efetivamente a magnitude dos passos residuais para controlar a propagação do sinal.

Principais Contribuições

1. Arcabouço Teórico para Modelos Treinados

Ao contrário de trabalhos anteriores que frequentemente focam na inicialização do modelo ou dependem de evidências empíricas, este arcabouço analisa modelos treinados na convergência. Ele fornece uma avaliação sistemática de se as dinâmicas de treinamento levam a soluções regulares ou patologias.

2. Análise Rigorosa de Colocações de LN

• Instabilidade do Pre-LN: O artigo prova que, para o Pre-LN, a solução ótima para o problema de treinamento é ilimitada em magnitude. Mesmo com weight decay, os estados ocultos podem crescer exponencialmente ($O(e^D)$), e os gradientes crescem proporcionalmente às ativações, levando à instabilidade de treinamento.
• Estabilidade do Peri-LN: Os autores derivam limites explícitos mostrando que o Peri-LN garante crescimento linear para o valor absoluto médio dos estados ocultos ($O(D)$) e crescimento quadrático para a variância ($O(D^2)$). Crucialmente, a sensibilidade local do Peri-LN é mostrada como invariante à magnitude da ativação, prevenindo a explosão de gradientes mesmo quando as ativações são grandes.

3. Quantificação de Incerteza

A análise estabelece que o Peri-LN fornece uma propagação controlada de incerteza (medida via distância de Wasserstein) da entrada para os estados ocultos terminais, enquanto o Pre-LN pode amplificar a incerteza da entrada para níveis ilimitados.

4. Escalonamento de Passo Residual

Guiados pela análise de estabilidade, os autores introduzem um mecanismo de escalonamento de passo residual ($\Delta t < 1$). Eles demonstram teoricamente que este escalonamento:

• Afunila os limites de crescimento do estado oculto e da incerteza de saída.
• Mitiga a explosão de gradientes ao escalar a sensibilidade local.
• Melhora tanto a estabilidade direta quanto a reversa sem custo computacional adicional.

Resultados Experimentais

Os autores validam suas descobertas teóricas usando variantes do GPT-2 (124M a 1.5B de parâmetros) treinadas no dataset OpenWebText.

• Estabilidade de Treinamento: Em experimentos com 5 tentativas, os modelos Pre-LN divergiram em 1 execução quando o weight decay foi habilitado e em 3 execuções quando o weight decay foi desabilitado. Em contraste, os modelos Peri-LN permaneceram estáveis em todas as 5 tentativas sob ambas as configurações de weight decay.
• Desempenho: O Peri-LN alcançou desempenho comparável ao Pre-LN mesmo quando os hiperparâmetros foram ajustados para o Pre-LN.
• Impacto do Escalonamento: Combinar o Peri-LN com o escalonamento de passo residual ($\Delta t = 0.1$) consistentemente rendeu o melhor desempenho em todas as métricas (Val Loss, Perplexity, ROUGE, scores BERT) e controlou efetivamente o crescimento dos momentos dos estados ocultos (valor absoluto médio e variância) através das camadas, alinhando-se com as previsões teóricas.

Significância e Alegações

O artigo afirma fornecer a primeira justificativa teórica rigorosa para as observações empíricas referentes às colocações de Normalização de Camada em Transformers.

• Explicação da Instabilidade: Explica por que o Pre-LN leva a estados ocultos ilimitados e explosão de gradientes via a má definição da PDE HJB subjacente.
• Justificativa para o Peri-LN: Fornece uma base teórica para a estabilidade do Peri-LN, mostrando que ele mantém o crescimento controlado e a sensibilidade invariante.
• Orientação de Design: O arcabouço oferece um método principado para realizar o "teste de sanidade" (sanity-check) de novas modificações arquitetônicas antes de treinamentos empíricos caros.
• Melhoria Prática: A introdução do escalonamento de passo residual é apresentada como uma modificação simples e sem custo que melhora significativamente a estabilidade e o desempenho, validada por experimentos mostrando que o Peri-LN com escalonamento supera as configurações padrão.

Os autores enfatizam que seu trabalho preenche a lacuna entre heurísticas empíricas e compreensão teórica, oferecendo uma perspectiva de tempo contínuo que gera tanto insights teóricos quanto melhorias arquitetônicas práticas.