Property-Based Testing for Robust Softmax and Cross-Entropy

Muhammad Maaz

cs.LGOct 27, 2026

Property-Based Testing for Robust Softmax and Cross-Entropy

Authors: Muhammad Maaz

Abstract

A recent paper of Barron et. al. (2026) identified two numerical issues with the softmax and cross-entropy functions in virtually all standard machine learning libraries (e.g., PyTorch and JAX). We show that these issues can be recovered through simple property-based tests in Hypothesis. However, our tests identify an additional edge case where their implementation fails. We then propose an even more robust softmax implementation which passes these tests and also supports use in causal attention masking. We compare our implementation with theirs and discuss some of the differences. Overall, this work shows that property-based testing can be a powerful tool for discovering and fixing numerical issues in machine learning libraries. Our implementations and property-based tests can be found at https://github.com/mmaaz-git/softmax-fix.

Figures & tables

Implementation	Forward probabilities	Computed Jacobian
Standard (PyTorch)	$[1,0]$	All zeros
Barron et al. (2026) (PyTorch translation)	$[1,0]$	All NaNs
Ours	$[1,0]$	All zeros

Table 1: Forward softmax probabilities and computed Jacobians for

[10^{308},-10^{308}]

, evaluated in float64 on CPU using PyTorch 2.7.1.

Robust Loss Functions

Robust loss functions limit the influence of outliers, noisy labels, and corrupted observations during model training.

See all 50 papers

A Fix for Softmax and Cross-Entropy
Oct 22, 2026Jonathan T. Barron, Dan Gnanapragasam, Michael TschannenNeural Network Activation Functions
Google DeepMind
SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients
Mar 9, 2026Anselm Paulus, A. René Geist, Vít Musil +3Straight-Through Estimator Differentiable Programming
University of Tübingen, Germany · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Masaryk University, Czechia +1
Ghosts of Softmax: Complex Singularities That Limit Safe Step Sizes in Cross-Entropy
Mar 13, 2026Piyush SaoDeep Learning Optimization Convergence Analysis
Oak Ridge National Laboratory, Oak Ridge, TN, USA

Oct 22, 2026J/K move · Enter open · S save

A Fix for Softmax and Cross-Entropy

Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen

Google DeepMind

Neural Network Activation Functions

Open paper AskPDF arXiv

Standard implementations of softmax, log-softmax, and cross-entropy have two numerical bugs. The first bug, Softmax Collapse, rounds the target logit's gradient to zero if that one logit is much larger than the others, resulting in gradients that no longer sum to zero. The second bug, Softmax Cancellation, occurs when student logits are close to teacher logits, and produces non-zero gradients even when the student and teacher logits are equal. We present simple implementations of log-softmax, softmax, and cross-entropy that fix both bugs. Our fix eliminates the previously observed "slingshot" phenomenon, speeds up grokking with weight decay, prevents accuracy collapse without weight decay, stops student networks from drifting away from the teacher during distillation, and shows that after memorization, logit growth is bounded by an equilibrium between weight decay and Adam's

ε

hyperparameter.

Abstract

Robust Loss Functions

Explore similar work

A Fix for Softmax and Cross-Entropy

SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

Ghosts of Softmax: Complex Singularities That Limit Safe Step Sizes in Cross-Entropy