Google DeepMind encuentra un fallo en softmax: el gradiente se anula con un margen de 16,6

El 22 de octubre de 2026, Jonathan Barron, Dan Gnanapragasam y Michael Tschannen, de Google DeepMind, publicaron el preprint A Fix for Softmax and Cross-Entropy. Los autores analizan dos defectos numéricos en las implementaciones estándar de softmax y cross-entropy que usan PyTorch y JAX. El primero, Softmax Collapse, anula en float32 el gradiente de la clase correcta cuando su margen sobre los competidores supera aproximadamente 16,6. El segundo, Softmax Cancellation, introduce un gradiente distinto de cero en la destilación, incluso cuando el estudiante ha replicado exactamente al profesor. Ambos defectos se corrigen reordenando los términos de la fórmula, sin alterar las matemáticas.

Dónde float32 pierde el número pequeño

El esquema estándar ya está protegido contra el desbordamiento: se resta el máximo de los logits y la exponencial no explota. Pero esa protección frente al overflow no garantiza la precisión de las diferencias pequeñas. Tomemos los logits [20, 0] en float32. La probabilidad de la segunda clase resulta ser 2,06·10⁻⁹, y ese número se almacena sin problema. La probabilidad de la primera se redondea a exactamente 1, porque un incremento de 10⁻⁹ sobre la unidad no cabe en siete cifras significativas. Tras restar la unidad, el gradiente de la clase correcta es cero, cuando debería rondar −2,06·10⁻⁹. Los gradientes de los competidores, en cambio, se conservan, y la suma de las derivadas deja de ser cero.

Si se denomina m al margen de la respuesta correcta y r = e^(−m), la fórmula habitual calcula 1/(1+r) − 1. La corregida calcula −r/(1+r). Es la misma función, pero la cantidad pequeña queda en el numerador y no se redondea junto con la unidad. Comprobación en float32: con m = 17, el esquema habitual da 0 y el corregido da −4,14·10⁻⁸. Con m = 20 se obtiene −2,06·10⁻⁹, y con m = 30, −9,36·10⁻¹⁴.

Los umbrales dependen del número de bits de la parte fraccionaria del formato de cálculo:

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

Son umbrales de la aritmética de una operación concreta. Que los pesos estén en bfloat16 no significa que softmax y la loss también se calculen en bfloat16: con precisión mixta, las operaciones sensibles suelen ejecutarse en float32.

El estudiante coincide con el profesor y el gradiente sigue ahí

El segundo defecto aparece en la destilación. El estudiante aprende a reproducir las probabilidades de un profesor fijo y, si coinciden por completo, el gradiente respecto a los logits debería ser cero. En ese caso, la cross-entropy permanece igual a la entropía del profesor, mientras que el gradiente y la divergencia KL se anulan.

En la práctica, las probabilidades del profesor y del estudiante a menudo se calculan por vías distintas: una mediante la división de exponenciales y otra mediante la exponencial de log_softmax. Para los logits [2, 1, 0] en float32, la diferencia es de unos 1,49·10⁻⁸ en dos componentes, con una norma de ruido de 2,1·10⁻⁸. Los logits son idénticos, pero el gradiente no lo es.

Los autores proponen comparar primero los logits y pasar a las probabilidades después. Se calcula la diferencia entre los logits del estudiante y del profesor, se le resta la media ponderada por las probabilidades del profesor y, a continuación, se obtiene la pequeña corrección con expm1 y log1p. Estas funciones conservan la precisión con argumentos pequeños. Cuando los logits coinciden, el resultado es un cero estricto: expm1(0) devuelve 0.

Por qué 10⁻⁸ llega al entrenamiento

Un gradiente del orden de 10⁻⁹ parece inofensivo, pero Adam normaliza el paso según una estimación de la escala del gradiente, de modo que la pequeñez absoluta de la señal no protege de las consecuencias. Con Softmax Collapse, el gradiente del logit objetivo desaparece mientras el de los competidores permanece. La señal distorsionada atraviesa la red y llega al optimizador adaptativo. Los autores relacionan esta asimetría con el efecto Slingshot: tras un periodo de loss baja aparecen picos bruscos.

Los experimentos del preprint se realizaron en float32, con cinco ejecuciones con distintas seeds. En el primero, una suma módulo 31 con weight decay λ = 1 y Adam ε = 10⁻⁸. En el segundo, (a² + b²) mod 97 con λ = 0 y ε = 10⁻¹⁶. En ambos casos, corregir el esquema numérico elimina los picos y cambia la dinámica posterior del entrenamiento. Es un resultado obtenido en modelos y tareas concretos. No demuestra que cualquier salto de la loss en un LLM esté causado por softmax.

Para calcular de forma estable la loss con hard labels, los autores usan logsumexp sobre los competidores y softplus en lugar de la suma directa con la unidad. El ejemplo didáctico en NumPy incluido en el análisis solo funciona cuando la clase correcta va en cabeza y no sustituye a autograd.

Qué comprobar en tu código

La corrección cambia el orden de las operaciones, no la arquitectura: no requiere parámetros nuevos ni memoria adicional. Para los equipos que entrenan modelos con loss baja, es una comprobación barata.

• Pasa al kernel de la loss logits con un margen de 17–30 en float32 y comprueba si se anula el gradiente de la clase correcta.
• Verifica que, en tu modo de precisión mixta, softmax y la loss se calculan realmente en float32 y no en bfloat16 o float16.
• En destilación, pasa al estudiante los logits del profesor. El gradiente debe ser exactamente cero, no del orden de 10⁻⁸.
• Compara el comportamiento con distintos valores de ε en Adam: de ese parámetro depende cuánto del ruido se convierte en paso.
• Si durante el entrenamiento ves picos de loss tras una meseta prolongada, revisa el esquema numérico antes de cambiar los hiperparámetros.

La implementación completa para grandes diferencias de logits debe tomarse del preprint de los autores: las demostraciones breves del análisis sirven para diferencias pequeñas y no sustituyen a una función ya preparada.