Proxies em pequena escala para instabilidades no treinamento de Transformers em grande escala

Resumo

Equipes que treinaram grandes modelos baseados em Transformers relataram instabilidades de treinamento em grande escala que não apareciam ao treinar com os mesmos hiperparâmetros em escalas menores. Embora as causas dessas instabilidades sejam de interesse científico, a quantidade de recursos necessários para reproduzi-las tem dificultado a investigação. Neste trabalho, buscamos maneiras de reproduzir e estudar a estabilidade e a instabilidade do treinamento em escalas menores. Primeiro, focamos em duas fontes de instabilidade de treinamento descritas em trabalhos anteriores: o crescimento dos logits nas camadas de atenção (Dehghani et al., 2023) e a divergência dos logits de saída das probabilidades logarítmicas (Chowdhery et al., 2022). Ao medir a relação entre a taxa de aprendizado e a perda em diferentes escalas, mostramos que essas instabilidades também aparecem em modelos pequenos quando treinados com altas taxas de aprendizado, e que as mitigações anteriormente empregadas em grandes escalas são igualmente eficazes nesse regime. Isso nos leva a investigar até que ponto outras intervenções conhecidas no otimizador e no modelo influenciam a sensibilidade da perda final a mudanças na taxa de aprendizado. Para isso, estudamos métodos como warm-up, decaimento de peso e o muParam (Yang et al., 2022), e combinamos técnicas para treinar modelos pequenos que alcançam perdas semelhantes em ordens de magnitude de variação da taxa de aprendizado. Por fim, para concluir nossa exploração, estudamos dois casos em que as instabilidades podem ser previstas antes de surgirem, examinando o comportamento de escalonamento das normas de ativação e gradiente do modelo.

English

Teams that have trained large Transformer-based models have reported training instabilities at large scale that did not appear when training with the same hyperparameters at smaller scales. Although the causes of such instabilities are of scientific interest, the amount of resources required to reproduce them has made investigation difficult. In this work, we seek ways to reproduce and study training stability and instability at smaller scales. First, we focus on two sources of training instability described in previous work: the growth of logits in attention layers (Dehghani et al., 2023) and divergence of the output logits from the log probabilities (Chowdhery et al., 2022). By measuring the relationship between learning rate and loss across scales, we show that these instabilities also appear in small models when training at high learning rates, and that mitigations previously employed at large scales are equally effective in this regime. This prompts us to investigate the extent to which other known optimizer and model interventions influence the sensitivity of the final loss to changes in the learning rate. To this end, we study methods such as warm-up, weight decay, and the muParam (Yang et al., 2022), and combine techniques to train small models that achieve similar losses across orders of magnitude of learning rate variation. Finally, to conclude our exploration we study two cases where instabilities can be predicted before they emerge by examining the scaling behavior of model activation and gradient norms.

Proxies em pequena escala para instabilidades no treinamento de Transformers em grande escala

Small-scale proxies for large-scale Transformer training instabilities

Resumo

Support