# Improved Stochastic Optimization of LogSumExp

arXiv: https://arxiv.org/abs/2509.24894 · HTML: https://mlsanity.com/paper/2509.24894/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 66 references that are in the map.
- First posted on arXiv 29 Sep 2025.

## Builds on (most cited first, 30 of 32)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Retiring Adult: New Datasets for Fair Machine Learning](https://mlsanity.com/paper/2108.04884/) (2021)
- [Certifying Some Distributional Robustness with Principled Adversarial Training](https://mlsanity.com/paper/1710.10571/) (2017)
- [Data-driven Distributionally Robust Optimization Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations](https://mlsanity.com/paper/1505.05116/) (2015)
- [Stochastic Optimization for Large-scale Optimal Transport](https://mlsanity.com/paper/1605.08527/) (2016)
- [Large-Scale Methods for Distributionally Robust Optimization](https://mlsanity.com/paper/2010.05893/) (2020)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [Neural Optimal Transport](https://mlsanity.com/paper/2201.12220/) (2022)
- [Large-Scale Optimal Transport and Mapping Estimation](https://mlsanity.com/paper/1711.02283/) (2017)
- [The Bayesian Learning Rule](https://mlsanity.com/paper/2107.04562/) (2021)
- [Score-based Generative Neural Networks for Large-Scale Optimal Transport](https://mlsanity.com/paper/2110.03237/) (2021)
- [Entropic Neural Optimal Transport via Diffusion Processes](https://mlsanity.com/paper/2211.01156/) (2022)
- [On the Fenchel Duality between Strong Convexity and Lipschitz Continuous Gradient](https://mlsanity.com/paper/1803.06573/) (2018)
- [Provable Stochastic Optimization for Global Contrastive Learning: Small Batch Does Not Harm Performance](https://mlsanity.com/paper/2202.12387/) (2022)
- [An Online Method for A Class of Distributionally Robust Optimization with Non-Convex Objectives](https://mlsanity.com/paper/2006.10138/) (2020)
- [Fast yet Simple Natural-Gradient Descent for Variational Inference in Complex Models](https://mlsanity.com/paper/1807.04489/) (2018)
- [Sinkhorn Distributionally Robust Optimization](https://mlsanity.com/paper/2109.11926/) (2021)
- [Neural Optimal Transport with General Cost Functionals](https://mlsanity.com/paper/2205.15403/) (2022)
- [Label Noise Types and Their Effects on Deep Learning](https://mlsanity.com/paper/2003.10471/) (2020)
- [Finite-Sum Coupled Compositional Stochastic Optimization: Theory and Applications](https://mlsanity.com/paper/2202.12396/) (2022)
- [Energy-guided Entropic Neural Optimal Transport](https://mlsanity.com/paper/2304.06094/) (2023)
- [Statistical Learning with Conditional Value at Risk](https://mlsanity.com/paper/2002.05826/) (2020)
- [On Russian Roulette Estimates for Bayesian Inference with Doubly-Intractable Likelihoods](https://mlsanity.com/paper/1306.4032/) (2013)
- [Guaranteed bounds on the Kullback-Leibler divergence of univariate mixtures using piecewise log-sum-exp inequalities](https://mlsanity.com/paper/1606.05850/) (2016)
- [A New Unbiased and Efficient Class of LSH-Based Samplers and Estimators for Partition Function Computation in Log-Linear Models](https://mlsanity.com/paper/1703.05160/) (2017)
- [Stochastic Constrained DRO with a Complexity Independent of Sample Size](https://mlsanity.com/paper/2210.05740/) (2022)
- [Outlier-Robust Wasserstein DRO](https://mlsanity.com/paper/2311.05573/) (2023)
- [One-vs-Each Approximation to Softmax for Scalable Estimation of Probabilities](https://mlsanity.com/paper/1609.07410/) (2016)
- [SUMO: Unbiased Estimation of Log Marginal Probability for Latent Variable Models](https://mlsanity.com/paper/2004.00353/) (2020)
- [(f,Gamma)-Divergences: Interpolating between f-Divergences and Integral Probability Metrics](https://mlsanity.com/paper/2011.05953/) (2020)

## Built on by (most cited first, 4 of 4)

- [Hyperspherical Forward-Forward with Prototypical Representations](https://mlsanity.com/paper/2605.00082/) (2026)
- [A Geometry-Aware Efficient Algorithm for Compositional Entropic Risk Minimization](https://mlsanity.com/paper/2602.02877/) (2026)
- [PERO: Efficient Robust Post-Training Foundation Models for Encrypted Traffic Classification](https://mlsanity.com/paper/2608.15504/) (2026)
- [Scalable Cox Regression via Grouped Risk Sets and Sharper LogSumExp Rates](https://mlsanity.com/paper/2609.40120/) (2026)

## Abstract (arXiv)

The LogSumExp function, dual to the Kullback-Leibler (KL) divergence, plays a central role in many important optimization problems, including entropy-regularized optimal transport (OT) and distributionally robust optimization (DRO). In practice, when the number of exponential terms inside the logarithm is large or infinite, optimization becomes challenging since computing the gradient requires differentiating every term. We propose a novel convexity- and smoothness-preserving approximation to LogSumExp that can be efficiently optimized using stochastic gradient methods. This approximation is rooted in a sound modification of the KL divergence in the dual, resulting in a new f-divergence called the Safe KL divergence. Our experiments and theoretical analysis of the LogSumExp-based stochastic optimization, arising in DRO and continuous OT, demonstrate the advantages of our approach over existing baselines.