A profundidade se tornou treinável por meio de poucos ajustes pouco glamorosos

O artigo que descreve o PaLM, o modelo de linguagem de 540 bilhões de parâmetros do Google, tem uma seção que se lê mais como um diário de operações do que como um resultado de pesquisa. Durante o treino, a perda disparou cerca de vinte vezes, mesmo com o corte de gradiente ligado, em intervalos irregulares. Os modelos menores da mesma família não apresentaram as disparadas. O remédio da equipe foi procedimental. Reiniciaram a partir de um checkpoint de cerca de 100 passos antes de cada disparada e pularam entre 200 e 500 lotes de dados que cobriam o episódio, e a perda não disparou de novo naquele ponto. Eles não achavam que dados ruins fossem a causa, já que reexecutar os mesmos lotes a partir de um checkpoint anterior não produziu disparada alguma. O problema surgia de uma combinação particular de um lote com o estado dos parâmetros.

Em 2022, portanto, uma equipe com parte da melhor infraestrutura de treino do mundo administrava a estabilidade de uma rede profunda olhando um gráfico e voltando atrás. Isso diz algo sobre o quanto o problema de treinar redes profundas está resolvido. A profundidade é a característica que define o aprendizado profundo e, durante a maior parte da história do campo, foi também a característica que fazia as redes falharem. As razões são matemáticas e simples. As correções que as superaram são modestas. Como as mesmas correções estão dentro de todo grande modelo, quem treina ou ajusta um deles, ou paga alguém para fazê-lo, ganha em saber quais são e qual delas tem mais chance de ser a culpada quando uma execução dá errado.

O ganho se compõe, e se compõe nos dois sentidos

Uma rede com L camadas passa um sinal adiante por L transformações e passa um gradiente para trás pelas mesmas L. Se cada camada altera a magnitude típica do sinal por um fator g, a saída depois de L camadas fica multiplicada por g^L. Para uma rede de cem camadas a aritmética é implacável:

Ganho por camada	Após 50 camadas	Após 100 camadas
0,90	0,005	0,00003
0,95	0,08	0,006
1,00	1	1
1,05	11	130
1,10	117	13.800

Uma camada que reduz os sinais em apenas 10% deixa cerca de 0,003% do sinal no fundo de uma pilha de cem camadas. Uma camada que amplifica em 10% produz um sinal quatorze mil vezes grande demais. O treino ajusta os pesos em proporção aos gradientes, então gradientes que desaparecem impedem as primeiras camadas de aprender e gradientes que explodem lançam a rede inteira num precipício. O fio da navalha é estreito: só ganhos muito próximos de 1 funcionam em profundidade.

É o velho problema do gradiente que desaparece, analisado por Sepp Hochreiter em 1991 e por Yoshua Bengio, Patrice Simard e Paolo Frasconi em 1994, no contexto de redes recorrentes, em que a profundidade vem do tempo. A mesma composição vale para uma pilha feed-forward, e por boa parte de duas décadas ela limitou a profundidade útil de uma rede. A solução, quando veio, não foi uma ideia só. Foram quatro, cada uma tratando de um lugar diferente onde o ganho pode se desviar.

A inicialização define o ganho de partida

A primeira alavanca é a escala dos pesos aleatórios no início. Se os pesos são pequenos demais, cada camada encolhe o sinal, e se são grandes demais, ela o amplia. Xavier Glorot e Yoshua Bengio derivaram em 2010 uma escala que mantém aproximadamente constante a variância das ativações e dos gradientes para funções de ativação simétricas. Unidades lineares retificadas, que zeram metade das entradas, precisam de outra. Kaiming He e colegas a derivaram em 2015 para retificadores: pesos com variância igual a 2 dividido pelo número de entradas, de modo que o fator por camada seja 1 em média. Com essa inicialização treinaram do zero redes retificadas extremamente profundas, e uma rede construída sobre ela chegou a 4,94% de erro top-5 no ImageNet, resultado que apresentaram como o primeiro a superar a referência humana de 5,1%.

Compare com a tabela acima. Uma inicialização cuja variância erra em 10% em cada camada é um ganho de cerca de 1,05 em amplitude por camada, e na camada 100 o sinal está 130 vezes grande demais. A constante certa, 2 e não 1,8 ou 2,2, é a diferença entre uma rede que treina e uma que devolve NaN já no primeiro passo.

A normalização impede que o ganho se desvie

A inicialização correta resolve o começo, mas os pesos mudam durante o treino e o ganho se desvia. A normalização em lote (batch normalization), de Sergey Ioffe e Christian Szegedy (2015), renormaliza as entradas de cada camada para média zero e variância um ao longo do minilote, com escala e deslocamento aprendidos. O resultado de destaque foi prático: um classificador de imagens no estado da arte atingiu a mesma acurácia com 14 vezes menos passos de treino, com taxas de aprendizado muito maiores toleradas e menos sensibilidade à inicialização. Um conjunto de redes com normalização em lote chegou a 4,9% de erro top-5 de validação no ImageNet.

A explicação original era que a normalização reduz o "deslocamento interno de covariáveis", a deriva na distribuição de entrada de cada camada. Esse relato foi contestado. Shibani Santurkar e colegas argumentaram em 2018 que o benefício vem de tornar mais suave a superfície da função de perda, e que a história do deslocamento de distribuição não se sustenta em experimentos. Seja qual for o mecanismo, o efeito é real, e a lição é que resultados de engenharia podem correr anos à frente da teoria.

A normalização em lote depende de estatísticas do lote, o que é incômodo para sequências de comprimento variável e para lotes pequenos. A normalização de camada (layer normalization), de Jimmy Ba, Jamie Kiros e Geoffrey Hinton (2016), calcula as estatísticas sobre os atributos de um único exemplo, e é o que os Transformers usam. Sua forma simplificada, a normalização RMS (Zhang e Sennrich, 2019), elimina a subtração da média e é comum em modelos de linguagem recentes.

O lugar onde a normalização fica importa tanto quanto ela existir. O Transformer original a punha depois de cada soma residual ("post-LN"). Ruibin Xiong e colegas mostraram no ICML 2020 que esse arranjo deixa gradientes esperados grandes perto da saída na inicialização, o que torna o treino instável a menos que a taxa de aprendizado seja aquecida lentamente a partir de quase zero. Mover a normalização para dentro do ramo residual ("pre-LN") dá gradientes bem-comportados na inicialização, e nos experimentos deles a etapa de aquecimento pôde ser descartada com resultados comparáveis e muito menos ajuste. O GPT-2 já havia posto a normalização na entrada de cada subbloco, e a análise de Xiong explicou por que isso ajuda. A maioria dos grandes modelos de linguagem desde então usa pre-LN ou uma variante próxima, e esse pequeno movimento de uma operação é a "normalização movida para a frente" que aparece em qualquer lista de refinamentos do Transformer.

Conexões de atalho transformam um produto numa soma

A correção que tornou rotineiras as redes de cem camadas foi a conexão residual. Kaiming He e colegas observaram em 2015 que redes simples mais profundas davam resultados piores que as mais rasas, mesmo no conjunto de treino, o que significava que o problema era de otimização e não de sobreajuste. O remédio deles foi deixar cada bloco aprender uma correção para a própria entrada, de modo que a camada calcule x + f(x) em vez de f(x). O resumo do artigo relata redes residuais de até 152 camadas no ImageNet, oito vezes mais profundas que as redes VGG da época e com menor complexidade, um erro de conjunto de 3,57% que venceu a tarefa de classificação do ILSVRC de 2015 e uma análise no CIFAR-10 com 100 e 1000 camadas.

A razão pela qual funciona é visível na aritmética da tabela. Sem o atalho, o Jacobiano de cada camada é uma matriz cujo tamanho pode ser qualquer um, e o produto de cem delas é o problema. Com o atalho, o Jacobiano de cada camada é a identidade mais uma correção, e um gradiente pode viajar da saída até qualquer camada anterior pelo caminho da identidade com ganho exato de 1. A rede profunda se comporta, para efeito de fluxo de gradiente, como uma soma de caminhos de muitos comprimentos diferentes, inclusive curtíssimos. Andreas Veit e colegas mostraram que isso tem uma consequência notável: remover um único bloco de uma rede residual treinada quase não a prejudica, ao contrário de apagar uma camada de uma rede simples, e as redes residuais se comportam como conjuntos de redes relativamente rasas.

Um truque relacionado deixa o ponto bem claro. Priya Goyal e colegas, em seu trabalho de 2017 sobre treinar modelos do ImageNet com lotes muito grandes, inicializaram em zero a última escala de normalização de cada bloco residual. No começo do treino, cada bloco então não produz saída alguma, e a rede inteira de cem camadas é exatamente a função identidade, com ganho 1 por construção. Os blocos aprendem a contribuir à medida que o treino avança, e a profundidade é usada apenas onde os dados pedem. Começar raso, na prática, e crescer em profundidade é um bom resumo do que toda a família de correções consegue.

Um fluxo residual também tem um problema mais discreto. Se cada um dos L blocos soma uma contribuição de variância 1, a variância do fluxo depois de 100 blocos é cerca de 100, de modo que o desvio padrão é dez vezes maior que no início. O artigo do GPT-2 tratou disso escalando os pesos das camadas residuais na inicialização por 1 sobre a raiz quadrada do número de camadas, o que mantém a magnitude do fluxo aproximadamente constante. Pesquisadores da Anthropic popularizaram depois uma forma de ver o Transformer inteiro em torno desse objeto: todo bloco de atenção e feed-forward lê do fluxo residual e escreve de volta nele, e o fluxo é o canal pelo qual as camadas se comunicam.

O otimizador e o formato numérico

Duas peças a mais completam o quadro. O Adam, de Diederik Kingma e Jimmy Ba (2014), mantém uma estimativa corrente da média e da variância do gradiente de cada parâmetro e escala o passo de acordo, de modo que parâmetros cujos gradientes são consistentemente minúsculos não ficam para trás em relação àqueles cujos gradientes são grandes. Isso torna o treino bem menos sensível às diferenças de escala de gradiente entre camadas que sobraram, e o AdamW, que separa o decaimento de pesos da atualização por gradiente, virou o padrão para Transformers. O formato numérico bfloat16 mantém os 8 bits de expoente do ponto flutuante de 32 bits e encurta a mantissa, o que lhe dá a mesma faixa dinâmica com metade da memória e evita a maior parte dos problemas de estouro dos formatos de 16 bits.

O que ainda é difícil

As disparadas de perda do PaLM mostram que a pilha de correções estreita o problema e deixa um resto. Em escala muito grande, com taxas de aprendizado enormes, lotes gigantes e aritmética de baixa precisão, eventos raros ocorrem com frequência suficiente para importar, e uma execução que custa milhões de dólares não pode esperar para descobrir o que foram. A resposta prática tem sido operacional: checkpoints frequentes, monitoramento da perda e das normas de gradiente, a capacidade de voltar atrás e pular dados e uma equipe de plantão.

Dois instrumentos baratos pegam a maior parte dos problemas cedo. O corte de gradiente, introduzido por Razvan Pascanu, Tomas Mikolov e Yoshua Bengio em 2013 para redes recorrentes, reescala o gradiente inteiro sempre que sua norma excede um limite, o que transforma uma explosão rara em um passo comum. E um gráfico da norma do gradiente por camada, tirado a cada algumas centenas de passos, mostra a composição diretamente: como regra prática, se a norma do gradiente da primeira camada é mil vezes menor que a da última, ou maior, o ganho por camada está errado e nenhum ajuste de taxa de aprendizado salvará a execução.

As mesmas correções explicam boa parte do que dá errado em treinos pequenos. Quando um job de ajuste fino diverge, a lista de suspeitos decorre das seções acima: uma taxa de aprendizado alta demais para o tamanho do lote, falta de aquecimento, falta de corte de gradiente, uma configuração de precisão que estoura, um estado do otimizador restaurado de forma incorreta, uma ordem de dados que pôs um lote patológico no momento errado. Equipes que ajustam modelos internamente podem montar uma rotina curta em torno disso: registrar a perda e a norma do gradiente a cada passo, gravar checkpoints com frequência suficiente para que uma volta atrás custe minutos e iniciar toda configuração nova com uma taxa de aprendizado menor que a da última que funcionou.

O que a história sugere para líderes

A história traz duas lições para um líder que não treina modelos. Uma é sobre de onde vem a capacidade. O salto das redes de 19 camadas para as de 152 camadas deveu pouco a uma nova teoria da inteligência e muito à aritmética do ganho e a algumas linhas de código. O progresso nesse campo muitas vezes parece encanamento até ser visto de fora. As pessoas que tornam o treino estável valem, portanto, mais do que seus cargos sugerem, e um fornecedor que sabe descrever como suas execuções se mantêm estáveis, inclusive as falhas, dá uma evidência melhor do que o que mostra só uma tabela de benchmarks.

A outra é sobre o ritmo do entendimento. A normalização em lote funcionou por anos antes de alguém conseguir explicar por quê, e a primeira explicação se mostrou, na melhor das hipóteses, incompleta. Uma boa parte do que está em produção repousa em técnicas em que os profissionais confiam porque funcionam, não porque são plenamente compreendidas. Isso não é motivo para desconfiar delas, já que o histórico empírico é forte, mas é motivo para medir os resultados diretamente em vez de raciocinar de princípios básicos sobre o que um modelo deve estar fazendo.

Fontes

1. Chowdhery, A., Narang, S., Devlin, J., et al. (2022). PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311, seção sobre instabilidade de treino.
2. Bengio, Y., Simard, P., Frasconi, P. (1994). Learning long-term dependencies with gradient descent is difficult. IEEE Transactions on Neural Networks 5(2).
3. Glorot, X., Bengio, Y. (2010). Understanding the difficulty of training deep feedforward neural networks. AISTATS 2010.
4. He, K., Zhang, X., Ren, S., Sun, J. (2015). Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. ICCV 2015. arXiv:1502.01852.
5. Ioffe, S., Szegedy, C. (2015). Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. ICML 2015. arXiv:1502.03167.
6. Santurkar, S., Tsipras, D., Ilyas, A., Madry, A. (2018). How Does Batch Normalization Help Optimization? NeurIPS 2018. arXiv:1805.11604.
7. Ba, J. L., Kiros, J. R., Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450.
8. Xiong, R., Yang, Y., He, D., et al. (2020). On Layer Normalization in the Transformer Architecture. ICML 2020. arXiv:2002.04745.
9. He, K., Zhang, X., Ren, S., Sun, J. (2016). Deep Residual Learning for Image Recognition. CVPR 2016. arXiv:1512.03385.
10. Veit, A., Wilber, M., Belongie, S. (2016). Residual Networks Behave Like Ensembles of Relatively Shallow Networks. NeurIPS 2016. arXiv:1605.06431.
11. Radford, A., Wu, J., Child, R., et al. (2019). Language Models are Unsupervised Multitask Learners (GPT-2). Relatório técnico da OpenAI.
12. Elhage, N., Nanda, N., Olsson, C., et al. (2021). A Mathematical Framework for Transformer Circuits. Anthropic, transformer-circuits.pub.
13. Kingma, D. P., Ba, J. (2015). Adam: A Method for Stochastic Optimization. ICLR 2015. arXiv:1412.6980.
14. Hochreiter, S. (1991). Untersuchungen zu dynamischen neuronalen Netzen. Tese de diploma, Technische Universität München.
15. Zhang, B., Sennrich, R. (2019). Root Mean Square Layer Normalization. NeurIPS 2019. arXiv:1910.07467.
16. Loshchilov, I., Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR 2019. arXiv:1711.05101.
17. Pascanu, R., Mikolov, T., Bengio, Y. (2013). On the difficulty of training recurrent neural networks. ICML 2013. arXiv:1211.5063.
18. Goyal, P., Dollár, P., Girshick, R., et al. (2017). Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. arXiv:1706.02677.