Resumo Técnico: Diferentes Camadas, Diferentes Variedades

Declaração do Problema

A geometria do espaço de pesos é um fator crítico na otimização de redes neurais, mas os atuais métodos de treinamento com restrição de variedade tipicamente aplicam uma única restrição geométrica uniforme (por exemplo, ortogonalidade ou restrições de Stiefel) em todas as matrizes de pesos de um Transformer. Este trabalho investiga se diferentes módulos Transformer — especificamente camadas de atenção e blocos MLP (Perceptron Multicamadas) — preferem geometrias de espaço de pesos distintas. Os autores postulam que os papéis computacionais e os mecanismos não lineares desses módulos diferem significamente, sugerindo que uma atribuição de variedade uniforme pode ser subótima ou até mesmo desestabilizadora.

Metodologia

O estudo utiliza o Manifold Muon, um otimizador que aplica atualizações normalizadas por matriz aos pesos das camadas ocultas enquanto os restringe a variedades de matrizes específicas. Os autores comparam duas famílias primárias de restrições:

1. Restrição de Stiefel: Impõe ortonormalidade de colunas ($W^\top W = I$), controlando estritamente os valores singulares e fornecendo um forte controle espectral.
2. Restrição de DGram: Uma restrição de espaço Gram mais fraca que exige que as entradas fora da diagonal da matriz Gram sejam zero ($\text{Off}(W^\top W) = 0$). Isso preserva a ortogonalidade das colunas, mas permite que as normas das colunas (e, portanto, os valores singulares) variem, oferecendo graus de liberdade de escala adicionais.

Os autores avaliam cinco estratégias distintas de atribuição em um modelo GPT-2 small (124M de parâmetros) pré-treinado em OpenWebText:

• UNCONSTRAINED (Sem Restrição): Sem restrições de variedade.
• ALL-STIEFEL: Restrições de Stiefel tanto nos pesos de Atenção quanto nos de MLP.
• ALL-DGRAM: Restrições de DGram tanto nos pesos de Atenção quanto nos de MLP.
• HETERO: Restrições de Stiefel para Atenção, DGram para MLP.
• HETERO-INV: Restrições de DGram para Atenção, Stiefel para MLP.

Todas as configurações compartilham as mesmas arquiteturas de modelo, cronogramas de treinamento e hiperparâmetros do otimizador para isolar o efeito da atribuição de variedade.

Resultados Principais

Os experimentos revelam uma clara assimetria na forma como diferentes módulos respondem às restrições de variedade:

• Configuração Ótima: A atribuição HETERO (Stiefel para Atenção, DGram para MLP) alcançou a menor perda de validação estável (3.3544), superando tanto o baseline sem restrições quanto a configuração uniforme ALL-STIEFEL.
• Instabilidade no DGram-Atenção: Configurações que atribuem restrições de DGram às camadas de atenção (ALL-DGRAM e HETERO-INV) tornaram-se instáveis sob a configuração de hiperparâmetros compartilhados. Essas execuções exibiram crescimento explosivo na perda de validação e não conseguiram se recuperar.
• Estabilidade no DGram-MLP: Por outro lado, atribuir restrições de DGram às camadas MLP (quando pareadas com atenção Stiefel) permaneceu estável e melhorou o desempenho em comparação com a restrição uniforme de Stiefel.

Análise Mecanística

O artigo fornece uma explicação mecanística para a instabilidade observada na atenção com restrição de DGram:

1. Crescimento de Valores Singulares: Ao contrário das restrições de Stiefel, o DGram permite que os valores singulares das matrizes de peso cresçam arbitrariamente grandes. Nas camadas de atenção, as projeções de consulta ($W_Q$) e chave ($W_K$) multiplicam-se para formar os logits de atenção ($Z = X W_Q W_K^\top X^\top / \sqrt{d}$).
2. Amplificação de Logits: O crescimento descontrolado nas normas espectrais de $W_Q$ e $W_K$ amplifica diretamente a magnitude dos logits de atenção.
3. Saturação de Softmax: Grandes magnitudes de logit empurram a função softmax para um regime de saturação, onde a saída se aproxima de um vetor one-hot. Isso faz com que o Jacobiano da softmax desapareça exponencialmente, levando à degeneração do gradiente e à instabilidade do treinamento.
4. Tolerância do MLP: As camadas MLP toleram restrições de DGram porque suas funções de ativação (ex: GELU) operam elemento a elemento e não acoplam pontuações de tokens através de um simplex de probabilidade normalizado. Consequentemente, a amplificação de escala nos pesos do MLP não força todo o módulo a um regime de roteamento saturado.

Contribuições

Os autores identificam três contribuições principais:

1. Avaliação Empírica: Eles avaliam atribuições de variedade por módulo para o Manifold Muon no pré-treinamento do GPT-2, demonstrando que restrições uniformes são insuficientes.
2. Descoberta de Assimetria: Eles mostram que camadas de atenção e MLP preferem geometrias de espaço de pesos diferentes, com a atenção exigindo controle espectral forte (Stiefel) e o MLP beneficiando-se da liberdade de preservação de escala (DGram).
3. Mecanismo de Instabilidade: Eles identificam o crescimento de valores singulares e a subsequente saturação de softmax como um mecanismo plausível para a instabilidade da atenção com DGram.

Significância e Limitações

O artigo argumenta que a otimização consciente de simetria e geometria para Transformers deve ser específica por módulo, e não uniforme. Isso sugere que a geometria do espaço de pesos é um eixo de design prático para melhorar a estabilidade e a dinâmica de treinamento em redes modernas.

Os autores mantêm uma postura modesta em relação às suas afirmações:

• Eles não afirmam que o DGram de atenção é inerentemente instável sob todas as condições, mas sim que ele expõe um caminho instável sob a configuração de hiperparâmetros compartilhados testada sem controle de escala adicional.
• Os experimentos são limitados ao GPT-2 small e a um conjunto específico de sementes (seeds); a diferença de perda de validação entre a melhor configuração heterogênea e a uniforme é modesta.
• A análise mecanística serve como uma explicação para a instabilidade observada, em vez de uma prova completa de que o DGram de atenção não pode ser estabilizado (por exemplo, via ajuste de taxas de aprendizado ou regularização de escala explícita).

Sugere-se que trabalhos futuros desenvolvam métodos automáticos para atribuir restrições de variedade entre diferentes módulos, arquiteturas e escalas.