Weight Decay Improves Language Model Plasticity

Resumo Técnico: O Decaimento de Peso Melhora a Plasticidade de Modelos de Linguagem

Declaração do Problema
As práticas atuais para o treinamento de Grandes Modelos de Linguagem (LLMs) tipicamente desacoplam o pré-treinamento do pós-treinamento (ajuste fino, alinhamento). A otimização de hiperparâmetros e as leis de escala durante o pré-treinamento são predominantemente avaliadas com base na perda de validação do modelo base (entropia cruzada), sob a suposição de que uma perda de pré-treinamento menor correlaciona-se com um modelo downstream mais capaz. No entanto, essa abordagem negligencia a plasticidade do modelo — a capacidade de um modelo treinado de se adaptar com sucesso a novos dados e tarefas após o treinamento adicional sem reinicialização. O artigo investiga se a otimização dos hiperparâmetros de pré-treinamento visando exclusivamente o desempenho de pré-treinamento compromete inadvertidamente a adaptabilidade do modelo durante o ajuste fino subsequente. Especificamente, os autores focam no decaimento de peso (weight decay), um hiperparâmetro de regularização padrão, para determinar se seu valor ideal para minimizar a perda de pré-treinamento difere de seu valor ideal para maximizar a adaptabilidade downstream.

Metodologia
Os autores conduzem um estudo empírico sistemático através de duas famílias de modelos (Llama-2 e OLMo-2) e múltiplos tamanhos de modelo (até 4 bilhões de parâmetros). O design experimental varia o hiperparâmetro de decaimento de peso ($\lambda$) durante o pré-treinamento, controlando outros fatores.

1. Regimes de Treinamento: Os modelos são treinados em dois regimes distintos:
   • Computacionalmente otimizado: 20 tokens por parâmetro (TPP).
   • Sobretreinado (Overtrained): 140 TPP.
2. Avaliação de Ajuste Fino (Fine-tuning): Modelos pré-treinados são submetidos a ajuste fino em um conjunto diversificado de tarefas downstream, incluindo:
   • Seis conjuntos de dados de Cadeia de Pensamento (CoT) (ex: MetaMathQA, MMLUProCoT).
   • Cinco benchmarks de compreensão de linguagem e raciocínio comum (ex: HellaSwag, WinoGrande).
   • Uma tarefa de alinhamento de segurança.
3. Métricas: O desempenho downstream é avaliado usando múltiplas métricas, incluindo Greedy (Pass@1), Pass@16, Voto Majoritário (Maj@16), Pontuação de Modelo de Recompensa (RM@16) e Razão de Correção (Correct Ratio).
4. Análise Mecanística: Para entender por que o decaimento de peso afeta a plasticidade, os autores analisam:
   • Sondagem Linear (Linear Probing): Medindo a separabilidade linear das representações internas (classificação de sentimento e tópico).
   • Rank da Atenção: Analisando o pseudo-rank das matrizes de projeção de consulta-chave ($W_{QK}$) e de valor ($W_{VP}$).
   • Overfitting: Calculando a lacuna entre a perda de treinamento e a perda de validação (Train-Val Gap).

Principais Contribuições

1. Decaimento de Peso como Impulsionador de Plasticidade: O artigo demonstra que o decaimento de peso é um fator crítico na moldagem da plasticidade do modelo. Contrariando os padrões usuais, valores maiores de decaimento de peso durante o pré-treinamento levam consistentemente a um melhor desempenho downstream após o ajuste fino.
2. Desacoplamento de Objetivos: O estudo fornece evidências de que o valor do hiperparâmetro que minimiza a perda de validação do pré-treinamento não é necessariamente o valor que maximiza o desempenho downstream. Em vários regimes, modelos com maior perda de pré-treinamento (devido a um maior decaimento de peso) superam modelos com menor perda de pré-treinamento após o ajuste fino.
3. Insights Mecanísticos: Os autores identificam três mecanismos específicos através dos quais o decaimento de peso aumenta a plasticidade:
   • Ele incentiva representações linearmente separáveis, tornando as características internas mais acessíveis para tarefas downstream.
   • Ele regulariza as matrizes de atenção, especificamente reduzindo o rank do produto consulta-chave ($W_{QK}$), levando a estruturas mais comprimidas e robustas.
   • Ele reduz o overfitting nos dados de pré-treinamento, efetivamente "esquecendo" o ruído específico do treinamento e preservando a capacidade do modelo de aprender novas informações.

Principais Resultados

• Mudança do Decaimento de Peso Ótimo: Para o treinamento computacionalmente otimizado (20 TPP), o decaimento de peso ideal para o desempenho downstream é significativamente maior do que o padrão de 0,1. Por exemplo, valores ótimos foram encontrados em 0,5–1,0 para os modelos Llama-2 e OLMo-2 neste regime. No regime sobretreinado (140 TPP), o valor ideal para o desempenho downstream (0,3) permanece superior ao valor que minimiza a perda de pré-treinamento (0,1).
• Trade-offs Contraintuitivos: O artigo destaca um trade-off onde modelos base "piores" (maior perda de pré-treinamento) podem se tornar modelos finais "melhores" após o ajuste fino. Por exemplo, um modelo OLMo-2-1B-140x pré-treinado com $\lambda=0,3$ ou $1,0$ apresentou maior perda de pré-treinamento do que a linha de base $\lambda=0,1$, mas alcançou uma precisão downstream visivelmente superior.
• Robustez: Essas descobertas sustentam-se em diferentes famílias de modelos, tamanhos e mesmo quando outros hiperparâmetros (taxa de aprendizado, tamanho do lote) são variados durante o pré-treinamento e o ajuste fino.
• Correlações Mecanísticas:
  • Um maior decaimento de peso correlaciona-se com maior precisão de sondagem linear, indicando representações mais estruturadas.
  • Um maior decaimento de peso correlaciona-se com um menor pseudo-rank nas matrizes de atenção (particularmente $W_{QK}$), sugerindo uma transição para hipóteses de baixo rank e robustas.
  • Um maior decaimento de peso correlaciona-se com uma menor lacuna entre Treino e Validação (Train-Val Gap), indicando redução de overfitting.

Significância e Alegações
O artigo argumenta que o campo deve reconsiderar a dependência exclusiva da perda de validação de pré-treinamento como a única métrica para otimização de hiperparâmetros. Os autores afirmam que a plasticidade do modelo é uma propriedade distinta e crucial que é frequentemente sacrificada quando se otimiza puramente pela perda de pré-treinamento.

A significância deste trabalho reside em:

1. Reavaliação de Padrões: Sugere que escolhas padrão de hiperparâmetros (ex: $\lambda=0,1$) podem ser subótimas para todo o ciclo de vida do desenvolvimento de LLMs, particularmente quando a adaptabilidade downstream é o objetivo final.
2. Limitações da Entropia Cruzada: Destaca que a perda de entropia cruzada é um proxy imperfeito para a capacidade downstream, pois modelos com perdas semelhantes podem exibir plasticidades vastamente diferentes.
3. Papel Multifacetado dos Hiperparâmetros: Ressalta que um único hiperparâmetro como o decaimento de peso desempenha um papel complexo e multifacetado na moldagem do comportamento do modelo, influenciando não apenas a generalização, mas também as propriedades estruturais (linearidade, rank) que facilitam o aprendizado futuro.

Os autores mantêm a modéstia, observando que, embora seus resultados sejam robustos através das escalas testadas, os valores ótimos específicos podem mudar em cenários de escala extremamente grande ou fortemente sobretreinados além de seu escopo experimental atual. Eles chamam para trabalhos futuros que investiguem mais profundamente o trade-off estabilidade-plasticidade e a aplicabilidade dessas descobertas a outras modalidades de modelos de fundação.