Improved Stochastic Optimization of LogSumExp

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.

Many machine learning methods need to combine a large number of scores, losses, or costs into a single objective. A common way to do this is the LogSumExp function, which first sums many exponentials and then takes a logarithm. Such a structure makes the objective difficult to optimize with random small batches, because a small batch does not faithfully represent the full objective. We propose a new approximation that moves the problem into a form better suited for stochastic gradient methods, which are the standard tools for large-scale machine learning. The approximation preserves important mathematical properties of LogSumExp, such as smoothness and convexity, while allowing unbiased stochastic gradients for the new objective. It is based on a controlled modification of the mathematical formulation behind LogSumExp, which also prevents any single sampled term from receiving an excessively large weight. Our experiments show that this approach is effective for optimal transport and robust learning problems. In difficult regimes, it avoids numerical failures and can converge faster than existing methods designed for these applications.