Google DeepMind encontra falha no softmax: gradiente zera quando a margem chega a 16,6

Em 22 de outubro de 2026, Jonathan Barron, Dan Gnanapragasam e Michael Tschannen, do Google DeepMind, divulgaram o preprint A Fix for Softmax and Cross-Entropy. Os autores analisaram dois defeitos numéricos nas implementações padrão de softmax e cross-entropy usadas no PyTorch e no JAX. O primeiro, Softmax Collapse, zera o gradiente da classe correta em float32 quando a margem em relação aos concorrentes passa de cerca de 16,6. O segundo, Softmax Cancellation, introduz um gradiente diferente de zero na destilação, mesmo quando o modelo aluno reproduz exatamente o professor. Os dois defeitos se resolvem reordenando os termos da fórmula, sem alterar a matemática.

Onde o float32 perde o número pequeno

O esquema padrão já é protegido contra overflow: subtrai-se o máximo dos logits, e a exponencial não explode. Mas a proteção contra overflow não garante precisão nas diferenças pequenas. Tome os logits [20, 0] em float32. A probabilidade da segunda classe sai como 2,06·10⁻⁹, e esse número é armazenado sem problemas. A probabilidade da primeira é arredondada para exatamente 1, porque um acréscimo de 10⁻⁹ a 1 não cabe em sete dígitos significativos. Depois da subtração de 1, o gradiente da classe correta fica em zero, quando deveria ser cerca de −2,06·10⁻⁹. Os gradientes dos concorrentes, porém, permanecem, e a soma das derivadas deixa de ser zero.

Se chamarmos de m a margem da resposta correta e de r = e^(−m), a fórmula usual calcula 1/(1+r) − 1. A corrigida calcula −r/(1+r). É a mesma função, mas o valor pequeno fica no numerador e não é arredondado junto com o 1. Verificação em float32: com m = 17, o esquema usual retorna 0, e o corrigido retorna −4,14·10⁻⁸. Com m = 20, o resultado é −2,06·10⁻⁹; com m = 30, é −9,36·10⁻¹⁴.

Os limiares dependem do número de bits da parte fracionária do formato de cálculo:

• bfloat16, 7 bits: limiar de margem ≈ 5,54;
• float16, 10 bits: ≈ 7,62;
• float32, 23 bits: ≈ 16,63;
• float64, 52 bits: ≈ 36,74.

São limiares da aritmética de cada operação específica. Ter pesos em bfloat16 não significa que softmax e loss também sejam calculados em bfloat16: na precisão mista, as operações sensíveis costumam rodar em float32.

O aluno igualou o professor, mas o gradiente continua lá

O segundo defeito aparece na destilação. O modelo aluno aprende a reproduzir as probabilidades de um professor fixo e, quando há coincidência total, o gradiente em relação aos logits deveria ser zero. Nesse caso, a cross-entropy permanece igual à entropia do professor, enquanto o gradiente e a divergência KL chegam a zero.

Na prática, as probabilidades do professor e do aluno muitas vezes são calculadas por caminhos diferentes: uma pela divisão de exponenciais, a outra pela exponencial do log_softmax. Para os logits [2, 1, 0] em float32, a diferença é de cerca de 1,49·10⁻⁸ em dois componentes, com norma de ruído de 2,1·10⁻⁸. Os logits são idênticos, mas o gradiente não é.

Os autores propõem comparar primeiro os logits e só depois passar às probabilidades. Calcula-se a diferença entre os logits do aluno e do professor, subtrai-se dela a média ponderada pelas probabilidades do professor e, em seguida, calcula-se a pequena correção com expm1 e log1p. Essas funções preservam a precisão para argumentos pequenos. Com logits coincidentes, o resultado é um zero exato: expm1(0) retorna 0.

Por que 10⁻⁸ chega ao treinamento

Um gradiente da ordem de 10⁻⁹ parece inofensivo, mas o Adam normaliza o passo pela estimativa da escala do gradiente, de modo que a pequenez absoluta do sinal não protege contra as consequências. No Softmax Collapse, o gradiente do logit-alvo desapareceu, mas o dos concorrentes ficou. O sinal distorcido atravessa a rede e chega ao otimizador adaptativo. Os autores associam essa assimetria ao efeito Slingshot: depois de um período de loss baixo, surgem picos abruptos.

Os experimentos do preprint foram feitos em float32, com cinco execuções usando seeds diferentes. No primeiro, a adição módulo 31 com weight decay λ = 1 e Adam ε = 10⁻⁸. No segundo, (a² + b²) mod 97 com λ = 0 e ε = 10⁻¹⁶. Nos dois casos, a correção do esquema numérico elimina os picos e muda a dinâmica posterior do treinamento. É um resultado obtido em modelos e tarefas específicos. Ele não prova que qualquer salto de loss em LLMs seja causado pelo softmax.

Para calcular de forma estável o loss com hard labels, os autores usam logsumexp sobre os concorrentes e softplus, em vez da soma direta com 1. O exemplo didático em NumPy apresentado na análise funciona apenas quando a classe correta lidera e não substitui o autograd.

O que verificar no seu código

A correção muda a ordem das operações, não a arquitetura: não exige novos parâmetros nem memória adicional. Para equipes que treinam modelos com loss baixo, é uma verificação barata.

• Passe ao kernel de loss logits com margem de 17 a 30 em float32 e veja se o gradiente da classe correta zera.
• Confirme que, no seu modo de precisão mista, softmax e loss são de fato calculados em float32, e não em bfloat16 ou float16.
• Na destilação, passe ao aluno os logits do professor. O gradiente deve ser exatamente zero, e não da ordem de 10⁻⁸.
• Compare o comportamento com diferentes valores de ε no Adam: deste parâmetro depende o quanto o ruído se converte em passo.
• Se o treinamento mostra picos de loss depois de um longo platô, verifique o esquema numérico antes de mexer nos hiperparâmetros.

A implementação completa para grandes diferenças de logits deve ser tirada do preprint dos autores: as demonstrações curtas da análise servem para diferenças pequenas e não substituem uma função pronta.