Uniform Scaling Limits in AdamW-Trained Transformers

Resumo Técnico: Limites de Escalonamento Uniforme em Transformers Treinados com AdamW

Declaração do Problema
O artigo aborda o desafio teórico de compreender a dinâmica de treinamento de Modelos de Linguagem de Grande Escala (LLMs), especificamente Transformers, à medida que eles escalonam em profundidade ($L$) e largura (número de cabeças de atenção, $H$). Embora leis de escalonamento empíricas impulsionem a indústria a aumentar o tamanho dos modelos, uma compreensão teórica rigorosa de como as escolhas de design afetam a generalização e a robustez no limite de profundidade infinita e largura infinita permanece incompleta.

A literatura existente sobre a convergência do gradiente descendente para otimalidade global frequentemente depende de modelos limite (por exemplo, EDOs Neurais ou limites de Campo Médio). No entanto, transferir resultados desses limites para modelos práticos requer limites de convergência que se mantenham uniformemente sobre a entrada do modelo. Existe uma lacuna significativa na literatura em relação a Transformers treinados com o otimizador AdamW. Trabalhos anteriores frequentemente assumem gradiente descendente ou exigem que os parâmetros sejam compartilhados ou dependentes entre camadas, o que não reflete a prática padrão de inicialização independente e identicamente distribuída (i.i.d.). Além disso, o crescimento cúbico dos gradientes na auto-atenção softmax (devido aos produtos de matrizes $Q, K, O, V$) complica as garantias teóricas, frequentemente levando a limites que crescem exponencialmente com os passos de treinamento. Finalmente, os limites de convergência padrão frequentemente dependem do número de tokens ($N$) por meio de argumentos de cobertura (por exemplo, Teorema do Integral de Entropia de Dudley), o que limita sua aplicabilidade à convergência uniforme sobre distribuições de entrada arbitrárias.

Metodologia
Os autores modelam a dinâmica do estado oculto de um Transformer sem máscara, apenas com atenção, treinado com AdamW, como um sistema de partículas interagentes (IPS) acoplado através do mecanismo de atenção. A análise prossegue através das seguintes etapas:

1. Formulação de Campo Médio: A dinâmica discreta dos $N$ tokens e $H$ cabeças é aproximada por um limite de tempo contínuo. Os estados ocultos e as variáveis retropropagadas (adjuntas) são mostrados a convergir para a solução de um sistema forward-backward de Equações Diferenciais Ordinárias (EDOs).
2. EDOs de McKean–Vlasov (MVODEs): Na ausência de mascaramento causal, o sistema limite é identificado como uma EDO de McKean–Vlasov. O campo de velocidade desta EDO é definido pela expectativa sobre a distribuição de parâmetros.
3. Tratamento de AdamW e Não Linearidade: Um componente crítico da metodologia é o tratamento do otimizador AdamW. Diferentemente do SGD padrão, o AdamW desacopla o decaimento de peso. Os autores aproveitam essa propriedade para provar que o suporte das medidas de parâmetros permanece dentro de um conjunto compacto uniformemente em todos os passos de treinamento. Isso permite constantes de Lipschitz globais para o mecanismo de atenção e suas derivadas, prevenindo o crescimento exponencial de limites tipicamente associado ao crescimento cúbico do gradiente na auto-atenção.
4. Convergência Uniforme sem Cobertura: Para alcançar limites independentes do número de tokens ($N$), os autores evitam argumentos tradicionais de cobertura. Em vez disso, eles utilizam:
   • Concentração de Medida: Aplicando técnicas para limitar a norma $L^1$ das diferenças de martingale associadas à inicialização dos parâmetros.
   • Aproximação de Taylor em Blocos: Dividindo as camadas em blocos e aproximando os mapas de fluxo dentro de cada bloco usando expansões de Taylor.
   • Estrutura Tensorial: Explorando a estrutura algébrica específica do mecanismo de atenção (expresso como razões de funções lineares de medidas) para mostrar que o supremo sobre condições iniciais pode ser reduzido a um espaço de dimensão finita, eliminando assim a dependência de $N$.
5. AdamW em Blocos: Os autores introduzem uma variante do AdamW (AdamW em Blocos) que regulariza blocos de parâmetros separadamente. Essa adaptação permite que os limites derivados sejam independentes da dimensão de incorporação de tokens ($d$) também.

Principais Contribuições

• Teorema de Convergência Uniforme: O resultado principal (Teorema 7) estabelece que o Transformer sem máscara, apenas com atenção, com parâmetros inicializados i.i.d. a partir de uma distribuição com suporte compacto, converge em $L^2$ uniformemente sobre condições iniciais para um sistema de partículas interagentes de EDOs médias. A taxa de convergência é $O(L^{-1} + L^{-1/3}H^{-1/2})$.
• Independência da Contagem de Tokens: Diferentemente de limites anteriores que escalonam com $\sqrt{\log(N)}$ ou dependem de $N$ por meio de argumentos de cobertura, os limites derivados são independentes do número de tokens $N$. Isso implica que os resultados valem para condições iniciais não presentes nos dados de treinamento.
• Adaptação do AdamW: O trabalho estende resultados de convergência de grande profundidade para Transformers treinados com AdamW. Demonstra que o decaimento de peso desacoplado do AdamW garante que os parâmetros permaneçam em um conjunto compacto, fornecendo continuidade de Lipschitz global no treinamento para o mecanismo de atenção.
• Independência de Dimensão: Ao utilizar AdamW em Blocos, os limites são mostrados como independentes da dimensão de incorporação de tokens $d$.
• Derivadas Explícitas: O artigo calcula explicitamente as derivadas de Lions para o mecanismo de atenção no componente de medida, estabelecendo continuidade local de Lipschitz em relação à distância de Wasserstein $p$-Wasserstein para todo $p \in [1, \infty]$.

Resultados
O artigo prova que, sob escalonamento apropriado das cabeças de atenção e com parâmetros inicializados i.i.d., a dinâmica conjunta dos estados ocultos e variáveis retropropagadas converge para a solução de um sistema forward-backward de EDOs.

• Taxa de Convergência: O erro entre o modelo Transformer discreto e o limite de campo médio contínuo é limitado por $C(L^{-1} + L^{-1/3}H^{-1/2})$, onde $C$ é uma constante independente de $N$ e, sob AdamW em Blocos, independente de $d$.
• Uniformidade: A convergência é uniforme sobre conjuntos compactos de condições iniciais.
• Limites de Parâmetros: Os autores mostram que as atualizações do AdamW confinam os parâmetros a um conjunto compacto $S(R_\theta)$ uniformemente sobre todos os passos de treinamento, uma propriedade crucial para estabelecer constantes de Lipschitz globais.

Significado e Alegações
O artigo alega fornecer uma base teórica rigorosa para compreender o comportamento de escalonamento de Transformers treinados com o otimizador dominante, AdamW. O significado reside em:

1. Preenchendo a Lacuna: Conecta a dinâmica de treinamento discreta de Transformers práticos e de grande escala a limites de campo médio contínuos, oferecendo um cenário tratável para análise.
2. Robustez ao Tamanho de Entrada: Ao estabelecer limites independentes do número de tokens, os resultados sugerem que o comportamento teórico do modelo é intrínseco à sua arquitetura e inicialização, e não artefatos do tamanho específico do conjunto de dados. Isso é particularmente relevante para a generalização a dados não vistos.
3. Estabilidade de Otimização: A análise destaca o papel estabilizador do decaimento de peso desacoplado do AdamW no controle do crescimento de parâmetros, o que é essencial para as garantias teóricas de continuidade de Lipschitz na presença de mecanismos de atenção não lineares.

Os autores reconhecem uma compensação: a taxa de convergência $O(L^{-1} + L^{-1/3}H^{-1/2})$ é ligeiramente mais lenta que a taxa $O(L^{-1} + (LH)^{-1/2})$ observada em alguns estudos de ResNet, um sacrifício feito para alcançar limites uniformes independentes de $N$. Eles observam que estender esses resultados para mascaramento causal (padrão em modelagem de linguagem) e regimes de escalonamento de EDEs permanece uma área aberta para pesquisa futura.