🤖 machine learning

Improved Stochastic Optimization of LogSumExp

This paper proposes a novel convexity- and smoothness-preserving approximation for the LogSumExp function, grounded in a new "Safe KL" divergence, which enables efficient stochastic optimization for large-scale problems like distributionally robust optimization and entropy-regularized optimal transport.

Original authors: Egor Gladin, Alexey Kroshnin, Jia-Jie Zhu, Pavel Dvurechensky

Published 2026-06-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Egor Gladin, Alexey Kroshnin, Jia-Jie Zhu, Pavel Dvurechensky

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the "average" height of a crowd, but instead of just adding up heights and dividing by the number of people, you have to calculate a special kind of average where the tallest people count way more than everyone else. In the world of math and machine learning, this is called the LogSumExp function. It's a crucial tool used in everything from teaching AI to recognize images to making sure self-driving cars don't crash when the weather gets weird.

However, there's a big problem with this tool: It's a numerical nightmare.

The Problem: The "Explosion"

Think of the LogSumExp function like a very sensitive scale. If you put a heavy weight on it, the scale doesn't just tip; it explodes. In computer terms, when the numbers inside the calculation get too big, the computer's memory "overflows." It's like trying to pour a gallon of water into a thimble; the water spills everywhere, and the calculation crashes.

This happens often when:

There are too many people: The crowd (data) is massive or infinite.
The weights are extreme: The "tallest" people are so tall that their numbers become impossible for a standard computer to handle.

To fix this, traditional methods try to be very careful, using tiny steps to avoid the explosion. But this makes the process incredibly slow, like trying to walk across a room by taking baby steps to avoid tripping.

The Solution: The "Safe KL" Shield

The authors of this paper propose a clever new way to look at the problem. Instead of trying to calculate the "explosive" average directly, they build a shield around it.

They introduce a new concept called the Safe KL Divergence. Imagine you are trying to measure the distance between two groups of people. The old way (standard KL divergence) is like measuring the distance with a ruler that stretches infinitely if the groups get too far apart. The new "Safe" way uses a ruler that has a hard stop; it can't stretch beyond a certain point.

By using this "Safe" ruler, they create a new version of the LogSumExp function that:

Doesn't explode: It has a built-in safety valve that prevents numbers from getting too big.
Is still accurate: It stays very close to the original, difficult-to-calculate function.
Is smooth: It allows the computer to take big, confident steps instead of tiny, cautious ones.

The Analogy: The "SoftPlus" Bridge

The paper uses a mathematical trick called SoftPlus. Imagine you are trying to cross a river.

The Old Way: You try to jump across the whole river at once. If the river is wide (large data), you might fall in (overflow). If you try to jump in tiny hops, it takes forever.
The New Way: You build a bridge that gently slopes up and then levels off. You can walk across it quickly and safely. The bridge doesn't go exactly where the river is deepest (it's an approximation), but it gets you to the other side efficiently without falling in.

Why This Matters

The authors tested this new "Safe" method in two main areas:

Optimal Transport (Moving Data): Imagine you have a pile of sand in one spot and want to move it to another spot with the least amount of effort. This is a common problem in AI. The old methods often crash when the "sand" is very spread out or the "effort" calculation gets too intense. The new method handles these messy, complex situations without crashing, allowing the AI to learn faster.
Robust Optimization (Preparing for the Worst): Imagine you are planning a picnic. You want to prepare for the worst weather possible. The old way of calculating the "worst-case scenario" often leads to computer errors when the weather data is extreme. The new method calculates this worst-case scenario smoothly, ensuring the plan is robust without breaking the computer.

The Bottom Line

The paper claims that by swapping out the old, explosive math for this new "Safe" version, we can solve complex machine learning problems much faster and more reliably. It's like replacing a fragile glass ladder with a sturdy steel one: you can climb higher (solve harder problems) without the fear of it shattering under pressure.

The authors show that this method works better than existing techniques, especially when the data is messy or the numbers get huge, and it does so without needing massive amounts of computing power.

Technical Summary: Improved Stochastic Optimization of LogSumExp

1. Problem Statement

The LogSumExp function, defined as $F(\phi; \mu) = \log \int e^{\phi(x)} d\mu(x)$ , is central to numerous optimization problems in machine learning, including entropy-regularized optimal transport (OT), distributionally robust optimization (DRO), multiclass classification, and variational inference.

The primary challenge in optimizing this function arises when the number of exponential terms is large or the measure $\mu$ is continuous. Computing the exact gradient requires differentiating every term in the sum or integral, which is computationally prohibitive. Standard stochastic gradient descent (SGD) approaches that sample a subset of terms generally yield biased gradient estimators. Furthermore, existing methods often suffer from numerical instability (overflow) when the regularization parameter is small, as the exponential function grows rapidly, leading to floating-point exceptions.

2. Methodology

The authors propose a novel convexity- and smoothness-preserving approximation to the LogSumExp function. This approach modifies the variational formulation of the problem by replacing the standard Kullback-Leibler (KL) divergence in the dual with a new $f$ -divergence termed the Safe KL divergence.

2.1 Safe KL Divergence

The method introduces a parameter $\rho \in (0, 1)$ and defines a generator function $f_\rho(t)$ :
$f_\rho(t) = \begin{cases} t \log t + \frac{1-\rho t}{\rho} \log(1 - \rho t) & \text{if } 0 \le t \le \frac{1}{\rho} \\ +\infty & \text{otherwise} \end{cases}$
This generator induces the Safe KL divergence $D_\rho(\nu, \mu)$ . Unlike the standard KL divergence, $f_\rho$ imposes a hard constraint on the density ratio $d\nu/d\mu$ , preventing it from exceeding $1/\rho$ . This property bounds the reweighting in DRO and OT, offering stability against outliers and numerical overflow.

2.2 Variational Approximation

Using convex duality, the authors derive an equivalent variational representation for the approximated log-partition function $F_\rho(\phi; \mu)$ :
$F_\rho(\phi; \mu) = \inf_{\alpha \in \mathbb{R}} \left( \alpha - 1 + \frac{1}{\rho} \int \log(1 + \rho e^{\phi(x) - \alpha}) d\mu(x) \right)$
This formulation replaces the raw exponential $e^{\phi(x)}$ with a rescaled SoftPlus function, $\frac{1}{\rho}\log(1 + \rho e^{\phi(x) - \alpha})$ .

2.3 Optimization Properties

Unbiased Gradients: The stochastic gradients computed on sampled batches for the surrogate objective $G_\rho(\theta, \alpha)$ are unbiased estimators for the approximated objective. This eliminates the batch-size-dependent bias inherent in standard sampled LogSumExp approximations.
Convexity and Smoothness: If the underlying loss function is convex, the surrogate objective remains jointly convex in the parameters. The conjugate function $f_\rho^*$ is shown to be $1/\rho$ -smooth, ensuring controlled second moments for the stochastic subgradients.
Connection to CVaR: The authors demonstrate that $F_\rho$ is closely related to the Conditional Value at Risk (CVaR) functional, providing a smooth approximation that bridges LogSumExp minimization and CVaR optimization.

3. Key Contributions

The paper outlines four main contributions:

General-Purpose Approximation: A computationally efficient relaxation of the LogSumExp function that preserves convexity and smoothness, transforming the problem into an expectation minimization amenable to standard SGD. The approximation error is controlled by a single tunable scalar parameter $\rho$ .
Theoretical Foundation: The derivation of the Safe KL divergence as the dual counterpart to this approximation, establishing its theoretical properties and its relationship to existing $f$ -divergences.
Empirical Validation: Evaluation on continuous entropy-regularized OT and various DRO formulations (KL-divergence and Unbalanced OT). The method demonstrates competitive performance against application-specific baselines while avoiding numerical overflow issues.
Theoretical Insights: Clarification of the connections between the proposed approximation, CVaR, and the smooth approximation of risk functionals.

4. Experimental Results

The authors validate the method on three specific tasks:

Continuous Entropy-Regularized OT: Using MNIST and EMNIST datasets, the proposed method (parameterized by neural networks) was compared against the LSOT baseline.
- Result: The baseline (LSOT) suffered from numerical overflow for small regularization parameters ( $\epsilon$ ), requiring prohibitively small learning rates or terminating early. The proposed method remained stable and efficient across all tested $\epsilon$ values, achieving faster convergence with larger learning rates.
KL-Divergence DRO: Applied to a regression task on the California housing dataset.
- Result: For small batch sizes, the proposed estimator outperformed the baseline (which relies on large batches to reduce bias). The intermediate parameter $\rho = 10^{-3}$ offered the best trade-off between approximation accuracy and optimization smoothness.
Unbalanced OT DRO: Applied to MNIST classification with corrupted labels.
- Result: The baseline method (optimizing the sum of exponentials) frequently diverged or required extremely small step sizes due to overflow. The proposed approach achieved lower objective values and demonstrated robustness even with high label noise (up to 85% in supplementary experiments).

5. Significance and Limitations

Significance:
The paper claims that the proposed approach serves as a versatile tool for large-scale optimization problems involving LogSumExp. Its primary advantage is the ability to handle weak regularization and large-scale data without the numerical instability (overflow) that plagues standard exponential formulations. By providing unbiased stochastic gradients for the surrogate objective, it circumvents the need for massive batch sizes to reduce bias.

Limitations:
The authors modestly acknowledge several limitations:

Structural Properties: Since Safe KL modifies the original KL divergence, specific structural properties inherent to KL-based formulations may require separate analysis in specific applications.
Parameter Tuning: The parameter $\rho$ controls the bias-stability trade-off and currently requires manual tuning; an adaptive selection procedure is suggested as future work.
Scope: The experiments cover a limited set of downstream tasks, and further study is needed to extend the Safe KL formulation to other settings where LogSumExp and KL duality are central.