ByteDance Seed propõe o cronograma de passo de treinamento AIR: erro final 0,01 menor que o do cosine

Pesquisadores da ByteDance Seed (Qifan Yu, Xinyu Ma, Xingyan Bin, Di He) propuseram o AIR, um cronograma de taxa de aprendizado. No pré-treinamento de um modelo Mixture-of-Experts (MoE) com 128 bilhões de tokens, ele reduziu o erro final nos dados de treinamento (training loss) em cerca de 0,01 em comparação com o decaimento cossenoidal do passo (cosine). A comparação foi feita com a mesma soma de passos ao longo de todo o treinamento. Aqui, taxa de aprendizado (learning rate, LR) é o tamanho do passo com que o otimizador altera os pesos do modelo. Em geral, ela diminui ao longo do treinamento, na maioria das vezes seguindo uma curva cossenoidal.

Por que a normalização de pesos entrou em discussão

Para controlar o tamanho dos pesos, normalmente se usa o weight decay (WD): a cada passo, ele puxa levemente os pesos em direção a zero. Recentemente, passou-se a considerar como alternativa a weight normalization (WN), o treinamento com pesos normalizados. Alguns trabalhos relataram um ganho expressivo de loss no pré-treinamento.

Os autores verificaram de onde vem esse ganho e chegaram a outra explicação. Com WN, a LR nominal se converte em um tamanho real de atualização diferente do obtido com WD. Isso gera uma trajetória própria de taxa de aprendizado efetiva (effective learning rate, ELR). Os pesquisadores ajustaram a ELR dos dois métodos para que as trajetórias coincidissem. Com isso, as curvas de loss de WN e WD ficaram praticamente sobrepostas. Logo, o ganho vinha do formato da ELR, e não da normalização em si.

A conclusão prática é que o cronograma deve ser construído para a ELR, e não apenas para a LR nominal.

Como o AIR funciona

O nome é a sigla de Affine Inverse-square-Root decay, ou seja, decaimento afim pelo inverso da raiz quadrada. O cronograma tem duas partes.

• Base. O passo cai proporcionalmente a 1/√t, em que t é o número do passo de treinamento. Os autores derivaram essa forma teoricamente. • Correção. Uma transformação linear (afim). O treinamento dura um tempo finito, e o decaimento puro de 1/√t não fornece os valores desejados no início e no fim. A correção permite definir de antemão a LR de pico e a LR final.

Como no cosine, o usuário especifica dois pontos extremos: o pico e o valor final.

Há ainda uma observação sobre a fase final. Se várias execuções têm o mesmo trecho final do cronograma de ELR, esse trecho reduz a influência de escolhas de LR feitas antes. Os autores consideram o planejamento tardio da ELR uma das principais fontes do ganho.

O que o experimento mostrou

O teste foi feito no pré-treinamento de um modelo MoE com 128 bilhões de tokens. Em um MoE, apenas parte dos blocos de especialistas trabalha em cada token. O AIR foi comparado ao cosine e a outros cronogramas candidatos. Os resultados:

• O training loss é menor com o AIR. O loss final ficou cerca de 0,01 abaixo do obtido com cosine, com a mesma soma de LR ao longo de todo o treinamento. • O AIR também tem melhor qualidade nas tarefas aplicadas (downstream). • O cronograma é mais robusto à escolha de hiperparâmetros.

Em um experimento separado, a mesma trajetória de ELR foi aplicada ao weight decay comum, e a dinâmica de treinamento se reproduziu. Os autores concluem que o benefício não se limita ao treinamento com norma de pesos fixa e apontam a ELR como o objeto que vale a pena planejar.

O que isso significa na prática

Uma diferença de 0,01 no loss é pequena, mas foi obtida apenas com a troca do cronograma de passo. O trabalho oferece três orientações práticas:

• verificar como o seu esquema de pesos e o seu otimizador convertem a LR nominal em efetiva; • testar o cronograma AIR, em que só é preciso definir a LR de pico e a LR final; • dar mais atenção ao fim do treinamento, porque um trecho final comum do cronograma reduz a influência de configurações anteriores.

O resumo do artigo descreve os experimentos em pré-treinamento de MoE com 128 bilhões de tokens. Ele não informa nada sobre outros tamanhos de modelo e volumes de dados.

O artigo foi publicado no arXiv com o número 2610.40362 e entrou na seleção do Hugging Face Daily Papers.