Organizations: School of Data, Mathematical, and Statistical Sciences, University of Central Florida, Orlando, USA · Department of CIS, University of Macau, Macao, China
Recent advances in machine learning have emphasized the integration of structured optimization components into end-to-end differentiable models, enabling richer inductive biases and tighter alignment with task-specific objectives. In this work, we introduce a novel differentiable approximation to the zero-one loss-long considered the gold standard for classification performance, yet incompatible with gradient-based optimization due to its non-differentiability. Our method constructs a smooth, order-preserving projection onto the n,k-dimensional hypersimplex through a constrained optimization framework, leading to a new operator we term Soft-Binary-Argmax. After deriving its mathematical properties, we show how its Jacobian can be efficiently computed and integrated into binary and multiclass learning systems. Empirically, our approach achieves significant improvements in generalization under large-batch training by imposing geometric consistency constraints on the output logits, thereby narrowing the performance gap traditionally observed in large-batch training.
Figures & tables
Figure 1: Binary-argmax@k of a point x=(0.1,1.6,1) into the exterior of the Hypersimplex (left). At k=1 , the solution is (0,1,0) . Introducing temperature to the program yields an interior solution (right), i.e., the soft-binary-argmax@1. In R3 different k values yield points on a standard simplex, but in higher dimensions yields a point on the hypersimplex. In R4 with k=2 , the solution lies on an octahedron [ 2 ] .
CIFAR-10
FashionMNIST
Batch
CE
HS (ours)
Δ
t-stat
p-val
CE
HS (ours)
Δ
t-stat
p-val
128
0.8885
0.8917
0.0032
2.29
0.084
0.9456
0.9467
0.0011
1.31
0.262
256
0.8843
0.8874
0.0030
3.47
0.026
0.9434
0.9473
0.0040
4.34
0.012
512
0.8807
0.8857
0.0050
6.08
<0.01
0.9436
0.9469
0.0033
5.51
<0.01
1024
0.8776
0.8821
0.0045
9.78
<0.01
0.9413
0.9453
0.0040
6.67
<0.01
2048
0.8725
0.8791
0.0066
7.81
<0.01
0.9388
0.9446
0.0058
12.97
<0.01
Table 1: Batch-wise Comparison of Cross-Entropy (CE) vs. HyperSimplex (HS) Losses on CIFAR-10 and Fashion-MNIST. The highlighted scores are statistically significant at a 10% level of significance.
Automatic differentiation (AD) frameworks such as JAX and PyTorch have enabled gradient-based optimization for a wide range of scientific fields. Yet, many "hard" primitives in these libraries such as thresholding, Boolean logic, discrete indexing, and sorting operations yield zero or undefined gradients that are not useful for optimization. While numerous "soft" relaxations have been proposed that provide informative gradients, the respective implementations are fragmented across projects, making them difficult to combine and compare. This work introduces SoftJAX and SoftTorch, open-source, feature-complete libraries for soft differentiable programming. These libraries provide a variety of soft functions as drop-in replacements for their hard JAX and PyTorch counterparts. This includes (i) elementwise operators such as clip or abs, (ii) utility methods for manipulating Booleans and indices via fuzzy logic, (iii) axiswise operators such as sort or rank -- based on optimal transport or permutahedron projections, and (iv) offer full support for straight-through gradient estimation. Overall, SoftJAX and SoftTorch make the toolbox of soft relaxations easily accessible to differentiable programming, as demonstrated through benchmarking and a practical case study. Code is available at github.com/a-paulus/softjax and github.com/a-paulus/softtorch.
Anselm Paulus, A. René Geist, Vít Musil +3
University of Tübingen, Germany · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Masaryk University, Czechia +1
Solving massive-scale optimization problems requires scalable first-order methods with low per-iteration cost. This tutorial highlights a shift in optimization: using differentiable programming not only to execute algorithms but to learn how to design them. Modern frameworks such as PyTorch, TensorFlow, and JAX enable this paradigm through efficient automatic differentiation. Embedding first-order methods within these systems allows end-to-end training that improves convergence and solution quality. Guided by Fenchel-Rockafellar duality, the tutorial demonstrates how duality-informed iterative schemes such as the alternating direction method of multipliers, and the primal-dual hybrid gradient can be learned and adapted through representative case studies.
The choice of loss function in classification involves a fundamental trade-off: smooth losses (like Cross-Entropy) enable fast optimization rates but yield slow square-root consistency bounds, while piecewise-linear losses (like Hinge) offer fast linear consistency rates but suffer from non-differentiability. We propose Linear-Core (LC) Surrogates, a new family of convex loss functions that resolve this tension by stitching a linear core to a smooth tail. We prove that these surrogates are differentiable everywhere while retaining strict linear H-consistency bounds, effectively combining the optimization benefits of smoothness with the statistical efficiency of margin-based losses. In the structured prediction setting, we show that this smoothness unlocks a massive computational and energy advantage: it allows for an unbiased stochastic gradient estimator that bypasses the quadratic complexity O(∣Y∣2) of exact inference (e.g., Viterbi). Empirically, our method achieves a 23× speedup over Structured SVMs on large-vocabulary sequence tagging tasks and demonstrates superior robustness to instance-dependent label noise, outperforming Cross-Entropy by 2.6% on corrupted CIFAR-10.
Mehryar Mohri, Yutao Zhong
Google Research, New York, NY · Courant Institute of Mathematical Sciences, New York, NY