Resumo Técnico: A Importância da Inicialização Pequena para Grandes Modelos de Linguagem

Declaração do Problema

Embora o progresso nos Grandes Modelos de Linguagem (LLMs) seja tipicamente atribuído à escala, dados e arquitetura, a inicialização de parâmetros permanece uma escolha de design crítica que ressurgiu como um problema complexo na era dos modelos de grande escala. Heurísticas padrão (ex: Xavier, Kaiming) frequentemente exigem ajustes de tentativa e erro dispendiosos. Uma visão teórica prevalecente sugere que a inicialização pequena conduz as redes para um regime "condensado", onde os pesos se alinham ao longo de direções compartilhadas antes de se diversificarem, potencialmente melhorando o raciocínio e a generalização. No entanto, permanece incerto se esses benefícios persistem no pré-treinamento realista de LLMs em escala, ou se componentes arquiteturais inadvertidamente suprimem as vantagens da inicialização pequena.

Metodologia

Os autores investigam sistematicamente o impacto da escala de inicialização na dinâmica de treinamento e no desempenho de LLMs.

• Parametrização de Inicialização: Matrizes de peso $W$ são inicializadas como $W_{ij} \sim \mathcal{N}(0, \sigma^2)$, onde $\sigma = d_{in}^{-\gamma}$. O parâmetro $\gamma$ controla a escala: $\gamma = 0.5$ recupera a escala padrão tipo Xavier, enquanto $\gamma > 0.5$ (especificamente $\gamma = 1$) resulta em uma inicialização menor.
• Arquiteturas de Modelos: Experimentos foram conduzidos tanto em Transformers decodificadores densos quanto em modelos de Mistura de Especialistas (MoE). As variações arquiteturais incluíram LayerNorm/RMSNorm padrão e Atenção Multi-Cabeça (Multi-Head Attention).
• Intervenções Arquiteturais: Para abordar as limitações observadas, duas intervenções específicas foram testadas:
  1. Redução da constante do RMSNorm ($\epsilon$): Reduzir $\epsilon$ do padrão $10^{-5}$ para $10^{-12}$ para evitar que o fator de normalização seja dominado pela constante quando a variância do estado oculto é pequena.
  2. Atenção com Portão (Gated Attention): Introdução de portões dependentes de query nas saídas das cabeças de atenção para mitigar o "atrator de atenção" (atenção desproporcional no primeiro token), que é exacerbado pela inicialização pequena.
• Configuração de Treinamento: Os modelos foram treinados em um corpus bilíngue de alta qualidade (Chinês-Inglês) (1 trilhão de tokens) usando predição do próximo token com AdamW. Os experimentos abrangeram tamanhos de modelo de 0.1B a 1.5B de parâmetros.
• Métricas de Análise:
  • Nível de Token: Gap de probabilidade simétrica ($\Delta_{sym}$) e dificuldade do token (perda média) para localizar ganhos de desempenho.
  • Mecanística: Similaridade de cosseno por linha e Posto Estável (Stable Rank $\|W\|_F^2 / \|W\|_2^2$) para rastrear a evolução da estrutura de parâmetros (condensação vs. expansão).

Principais Resultados

1. A Limitação de Escala e Sua Correção

Sob configurações padrão ($\gamma=0.5$ vs. $\gamma=1$), a inicialização pequena consistentemente reduz a perda de validação, mas o benefício diminui conforme o tamanho do modelo aumenta. Os autores identificam duas barreiras arquiteturais que causam isso:

• Saturação do RMSNorm: Com o $\epsilon=10^{-5}$ padrão, o fator de normalização torna-se insensível à redução da variância da inicialização pequena uma vez que $\sigma^2(h) \lesssim \epsilon$. Reduzir $\epsilon$ para $10^{-12}$ restaura a sensibilidade.
• Atrator de Atenção (Attention Sink): A inicialização pequena intensifica a tendência de atender ao primeiro token. A atenção com portão mitiga isso efetivamente, reduzindo significativamente a perda quando combinada com a inicialização pequena.
• Resultado: Quando ambas as barreiras são removidas, a lacuna de desempenho aumenta. A inicialização pequena ($\gamma=1$) com esses ajustes melhora o tamanho efetivo do modelo em aproximadamente 44% e produz ganhos consistentes em todas as escalas de modelo.

2. Escala de Inicialização Ótima ($\gamma=1$)

O estudo determina que a inicialização não deve ser arbitrariamente pequena.

• Equilíbrio do Fluxo Residual: Em um Transformer, a atualização residual deve permanecer comparável ao caminho de identidade. A análise teórica mostra que o fluxo residual escala como $d^{1-\gamma}$.
  • Se $\gamma < 1$, o residual domina.
  • Se $\gamma > 1$, o residual desaparece, deixando a rede próxima de um mapeamento de identidade.
  • $\gamma = 1$ representa o ponto crítico de equilíbrio onde a atualização residual corresponde à escala de embedding. Empiricamente, o desempenho atinge o pico em $\gamma=1$ e degrada para $\gamma > 1$.

3. Ganhos de Nível de Token e Específicos de Tarefa

• Ganhos Não Uniformes: As melhorias não são uniformes em todos os tokens. Os ganhos concentram-se em tokens de "moderada a alta dificuldade" que requerem integração contextual e raciocínio, em vez de tokens fáceis ou extremamente ruidosos/raros.
• Desempenho em Benchmarks: Em modelos de 1.5B, a inicialização pequena otimizada ($\gamma=1$ + ajustes) supera significativamente a inicialização padrão ($\gamma=0.5$) em benchmarks focados em raciocínio (ex: GSM8K, MATH500, HellaSwag), com ganhos absolutos superiores a 4% em diversas tarefas.

4. Trajetória Mecanística: Condensação

A inicialização pequena induz uma trajetória de desenvolvimento distinta:

• Fase 1 (Condensação): No início do treinamento, as matrizes de peso condensam-se rapidamente em estruturas de baixo posto e baixa complexidade (as linhas se alinham em blocos coerentes).
• Fase 2 (Expansão): À medida que o treinamento progride, essas estruturas expandem-se para representações mais ricas e de maior complexidade.
• Contraste: A inicialização padrão ($\gamma=0.5$) começa com uma estrutura difusa e de alto posto que evolui lentamente, carecendo da fase inicial de baixa complexidade. Essa trajetória de complexidade "baixa para alta" é hipotetizada como um fator que direciona o modelo para regras de raciocínio parcimoniosas em vez de memorização superficial.

Significância e Alegações

O artigo argumenta que a inicialização de parâmetros é um determinante "do tipo gene" da capacidade do modelo e da dinâmica de treinamento, não meramente um detalhe de implementação de baixo nível.

• Eixo de Design: Os autores propõem tratar a escala de inicialização $\gamma$ como um hiperparâmetro explícito no design de LLMs.
• Recomendação Prática: Eles defendem uma regra de inicialização $\gamma$: expor a escala de inicialização como um controle e adotar $\gamma=1$ como o padrão. Esta intervenção é descrita como "quase sem custo", mas capaz de melhorar o pré-treinamento e fortalecer o raciocínio através das escalas.
• Insight Teórico: O trabalho fornece evidência concreta apoiando a ideia de que "compressão é inteligência", demonstrando que um regime específico de inicialização força os modelos através de uma fase de condensação que se correlaciona com melhores capacidades de raciocínio.
• Interdependência Arquitetural: As descobertas destacam que os benefícios da inicialização são contingentes aos componentes arquiteturais (especificamente constantes de normalização e mecanismos de atenção), sugerindo que estes devem ser ajustados conjuntamente e não independentemente.