Property-Based Testing for Robust Softmax and Cross-Entropy

Muhammad Maaz

Home/Archive/arXiv

arXiv · 2610.36427

Property-Based Testing for Robust Softmax and Cross-Entropy

Muhammad Maaz

View original source ↗Read source document ↗Export metadata ↓

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.

Explore related subjects

cs.LG

Keep this discovery

Explore connections, maps & timelines

BibTeX RIS

Muhammad Maaz. 2026-10-27. Property-Based Testing for Robust Softmax and Cross-Entropy. https://arxiv.org/abs/2610.36427

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

A Comparative Study of Convolutional and Recurrent Neural Networks for Storm Surge Prediction in Tampa Bay

In this paper, we compare the performance of three popular deep learning architectures: Long Short-Term Memory (LSTM), Convolutional Neural Networks (CNN), and CNN-LSTM in water level prediction over different time window sizes. We used wind, pressure, and tide data from National Oceanic and Atmospheric Administration (NOAA) stations. Our results show that the LSTM model generally outperformed the other models in forecasting water level trends. By applying LSTM, we achieved a mean squared error (MSE) of 0.006, making it the most reliable option for real-world use. We predicted water levels using wind, pressure, and tides. Wind, pressure, and tide data are time-based features that lack spatial patterns. Consequently, LSTM performed better than spatial methods, such as CNN. We also investigated Hurricane Ian, which occurred in Tampa Bay in September 2022 and caused a significant storm surge. Our findings show that CNN and CNN-LSTM performed slightly better than LSTMs during powerful storms, such as hurricanes. Hurricanes are seldom observed in training data, and CNN-based models are more reliable when predicting data that were not included in the training. To explain how environmental factors affect the predicted water level, we used the perturbation method. The perturbation shows that the tide is more effective at predicting water levels than wind and pressure. This study demonstrates the potential of deep learning to improve storm surge prediction accuracy.

cs.LG↗

arXiv2026

Anomaly Detection via Autoencoder Composite Features and NCE

Unsupervised anomaly detection is a challenging task. Autoencoders (AEs) or generative models are often employed to model the data distribution of normal inputs and subsequently identify anomalous, out-of-distribution inputs by high reconstruction error or low likelihood, respectively. However, AEs may generalize and achieve small reconstruction errors on abnormal inputs. We propose a decoupled training approach for anomaly detection that both an AE and a likelihood model trained with noise contrastive estimation (NCE). After training the AE, NCE estimates a probability density function, to serve as the anomaly score, on the joint space of the AE's latent representation combined with features of the reconstruction quality. To further reduce the false negative rate in NCE we systematically varying the reconstruction features to augment the training and optimize the contrastive Gaussian noise distribution. Experimental assessments on multiple benchmark datasets demonstrate that the proposed approach matches the performance of prevalent state-of-the-art anomaly detection algorithms.

cs.LG↗

arXiv2026

Accuracy vs. Accuracy: Computational Tradeoffs Between Classification Rates and Utility

We revisit the foundations of fairness and its interplay with utility and efficiency in settings where the training data contain richer labels, such as individual types, rankings, or risk estimates, rather than just binary outcomes. In this context, we propose algorithms that achieve stronger notions of evidence-based fairness than are possible in standard supervised learning. Our methods support classification and ranking techniques that preserve accurate subpopulation classification rates, as suggested by the underlying data distributions, across a broad class of classification rules and downstream applications. Furthermore, our predictors enable loss minimization, whether aimed at maximizing utility or in the service of fair treatment. Complementing our algorithmic contributions, we present impossibility results demonstrating that simultaneously achieving accurate classification rates and optimal loss minimization is, in some cases, computationally infeasible. Unlike prior impossibility results, our notions are not inherently in conflict and are simultaneously satisfied by the Bayes-optimal predictor. Furthermore, we show that each notion can be satisfied individually via efficient learning. Our separation thus stems from the computational hardness of learning a sufficiently good approximation of the Bayes-optimal predictor. These computational impossibilities present a choice between two natural and attainable notions of accuracy that could both be motivated by fairness.

cs.LG↗