Improved Stochastic Optimization of LogSumExp

## Abstract

Abstract The LogSumExp function, also known as the free energy, plays a central role in many important optimization problems,
including entropy-regularized optimal transport and distributionally robust optimization (DRO).
It is also the dual to the Kullback-Leibler (KL) divergence, which is widely used in machine learning.
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. Previous approaches that replace the full sum with a small batch introduce significant bias.
We propose a novel 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 f -divergence called the safe KL divergence .
The accuracy of the approximation is controlled by a tunable parameter and can be made arbitrarily small.
Like the LogSumExp, our approximation preserves convexity.
Moreover, when applied to an L L -smooth function bounded from below, the smoothness constant of the resulting objective scales linearly with L L .
Experiments in DRO and continuous optimal transport demonstrate the advantages of our approach over state-of-the-art baselines and the effective treatment of numerical issues associated with the standard LogSumExp and KL.

## 1 Introduction

Optimization problems arising in various fields involve the LogSumExp function, or, more generally, the log-partition functional

$$ $F(\varphi;\mu)\coloneqq\ln\int e^{\varphi(x)}\,\mathrm{d}\mu(x)\in(-\infty,\infty]$ (1) $$

mapping a measurable function $\varphi$ to $(-\infty,\infty]$ based on a probability measure $\mu$. The goal in such optimization problems is to minimize an objective involving $F$ w.r.t. $\varphi$ over some class.

LogSumExp function appears commonly in optimization objectives, e.g., multiclass classification with softmax probabilities (Bishop & Nasrabadi, 2006),
semi-dual formulation of entropy-regularized optimal transport (OT) (Peyré & Cuturi, 2019; Genevay et al., 2016),
minimax problems (Pee & Royset, 2011),
distributionally robust optimization (DRO) (Hu & Hong,; Ben-Tal et al., 2013; Kuhn et al., 2024), maximum likelihood estimation (MLE) for exponential families and graphical models (Wainwright et al., 2008),
variational Bayesian methods (Khan & Nielsen, 2018; Khan & Rue, 2023),
information geometry (Amari & Nagaoka, 2000),
KL-regularized Markov decision processes (Tiapkin et al., 2024).
Such optimization problems are characterized by two challenges. First, the decision variable $\varphi$ (or a parameter $\theta$ defining $\varphi$) often has large or infinite dimension. Second, the support of the measure $\mu$ can be large or even infinite. The first challenge is usually addressed by the use of first-order methods, especially stochastic gradient descent (SGD), with cheap iterations. Unfortunately, there is no universal way to tackle the second challenge. If the number of exponential terms under the logarithm is large or infinite, i.e., the Monte Carlo estimate $\log\sum_{i=1}^{N}e^{\varphi(x)}$ for a large $N$, the gradient computation requires differentiating each term, and, to the best of our knowledge, no cheap unbiased stochastic gradient have been proposed.

In the current work, we propose a general-purpose approach to tackle both mentioned challenges. To that end, we use a SoftPlus approximation of $F(\varphi;\mu)$ that allows using stochastic gradient methods while remaining close to the original objective.
We start with a variational formulation analogous to the one in the Gibbs principle, but with the KL-divergence replaced with another $f$-divergence – the *safe KL* divergence.
The corresponding variational problem can be of interest itself, as it possesses some properties which can be beneficial compared to the KL penalty — e.g., uniform density bound.
Moreover, it can be also viewed as an approximation of a conditional value at risk functional (CVaR).
In fact, the same functional (with different parameters) appeared in Soma & Yoshida (2020) in the context of smooth CVaR approximation.
Thus, we demonstrate that it generates a family of problems including CVaR and LogSumExp minimization as limit cases.

Related works. We are not aware of any universal way to efficiently treat optimization objectives involving the LogSumExp functional ([1](https://arxiv.org/html/2509.24894v1#S1.E1)), especially in infinite-dimensional settings. This functional appeared in different applications and was treated on a case-by-case basis.
Bouchard (2007) studies three upper
bounds on LogSumExp for approximate Bayesian inference. One of them is a particular case of the approximation proposed in the present work.
Titsias (2016) constructs a bound on softmax probabilities and shows that it leads to a bound on LogSumExp in the context of multiclass classification.
Nielsen & Sun (2016) approximate LogSumExp in the context of estimating divergences between mixture models. Their approximation combines LogSumExp bounds based on $\min$ and $\max$.
Hu & Hong and, susequently,
Levy et al. (2020) study DRO problems with $f$-divergences. They propose a batch-based approximation. When the ambiguity set is the unit simplex, and KL divergence penalty is used, the original objective is the LogSumExp of the losses over the entire dataset, while the approximation replaces it with the average of LogSumExp terms computed on individual batches. This approximation introduces bias, which can only be reduced by using large batch sizes.

Contributions. Our main contributions are as follows:

- 1.
We introduce a general-purpose and computationally efficient approach for handling the LogSumExp function in large-scale optimization problems by proposing a novel relaxation of this function. The proposed relaxation preserves key properties of the original LogSumExp function, such as convexity and smoothness, and enables the use of stochastic gradient methods for machine learning tasks.
Furthermore, our method only requires a simple and tunable scalar parameter, allowing the relaxation to be made arbitrarily close to the original LogSumExp objective as desired.
- 2.
We provide the theoretical backbone of this approximation, demonstrating that it is due to a modified version of the KL-divergence in the dual formulation. We termed the resulting $f$-divergence the *(Overflow-)Safe KL* divergence.
It can be applied to various applications where KL-divergence is used.
- 3.
We empirically demonstrate the effectiveness of our approach on tasks, including computing continuous entropy-regularized OT and various DRO formulations(^1^11The source code is available at [https://github.com/egorgladin/logsumexp-approx](https://github.com/egorgladin/logsumexp-approx).).
Our method outperforms existing state-of-the-art baselines in these applications and circumvents the overflow issue (Remark [3.1](https://arxiv.org/html/2509.24894v1#S3.Thmtheorem1)). It can also be combined with existing techniques. Therefore, it serves as a versatile tool for solving large-scale optimization problems.
- 4.
Additionally, we provide insights into a few remarkable connections between the proposed approximation and existing notions such as the conditional value-at-risk.

Notation.
Given $a,a_{1},\dots,a_{n}\in\mathbb{R}$, we define $\operatorname{LogSumExp}(a_{1},\dots,a_{n})\coloneqq\log(\sum_{i=1}^{n}e^{a_{i}})$ and $\operatorname{SoftPlus}(a)\coloneqq\log(1+e^{a})$.
Given a measurable space $\mathcal{X}$, by $\mathcal{P}(\mathcal{X})$ we denote the space of probability measures on $\mathcal{X}$, and by $\mathcal{C}(\mathcal{X})$ the space of continuous functions on $\mathcal{X}$. Let $\mu,\nu\in\mathcal{P}(\mathcal{X})$.
Define Kullback–Leibler (KL) divergence as

$$ $D_{KL}(\mu,\nu)\coloneqq\begin{cases}\int_{\mathcal{X}}\log\frac{\mathrm{d}\mu}{\mathrm{d}\nu}(x)\,\mathrm{d}\mu(x)&\mu\ll\nu,\\ +\infty&\text{otherwise},\end{cases}$ $$

where $\log$ is the natural logarithm and $\mu\ll\nu$ denotes that $\mu$ is absolutely continuous w.r.t. $\nu$.

## 2 SoftPlus approximation of log-partition function

In this section, we present our approximation to the log-partition function ([1](https://arxiv.org/html/2509.24894v1#S1.E1)) and describe its theoretical properties.
Recall that by the Gibbs variational principle

$$ $F(\varphi;\mu)=\sup\left\{\int\varphi(x)\,\mathrm{d}\nu(x)-D_{KL}(\nu,\mu):\nu\in\mathcal{P}(\mathcal{X}),\,\int|\varphi(x)|\,\mathrm{d}\nu(x)<\infty\right\}$ (2) $$

with the minimum attained at the Gibbs measure
$d\nu^{*}(x)=e^{\varphi(x)-F(\varphi;\mu)}d\mu(x),$
once $F(\varphi;\mu)<\infty$ (Polyanskiy & Wu, 2025, Proposition 4.7).

We are going to construct an approximation of $F$ with better regularity properties by changing $D_{KL}$ to another $f$-divergence.
Specifically, for any $0<\rho<1$, let us define the following.

###### Definition 2.1 (Safe KL entropy) .

We define the safe KL entropy generator
$f_{\rho}\colon[0,\infty)\to\mathbb{R}$ by

$$ $f_{\rho}(t)\coloneqq\begin{cases}t\log t+1+\frac{1-\rho t}{\rho}\log(1-\rho t),&0\leq t\leq\frac{1}{\rho},\\ +\infty,&\text{otherwise}.\end{cases}$ (3) $$

The resulting $f_{\rho}$-divergence, which we refer to as the *safe KL divergence*, is given by

$$ $D_{\rho}(\nu,\mu)\coloneqq\begin{cases}\int f_{\rho}\left(\frac{d\nu}{d\mu}(x)\right)\,\mathrm{d}\mu(x),&\nu\ll\mu,\\ +\infty,&\text{otherwise}.\end{cases}$ (4) $$

Figure: Figure 1: $f_{\rho}(t)$ for different values of $\rho$.
Refer to caption: fig_new/f_rho.png

It is easy to see that $f_{\rho}(t)\to f_{0}(t)\coloneqq t\log t+1-t$ as $\rho\to 0$.
Since $f_{0}$ induces the standard KL-divergence, $D_{\rho}$ is its approximation with accuracy regulated by the parameter $\rho$.

Using the variational representation, we define

$$ $\displaystyle F_{\rho}(\varphi;\mu)\coloneqq\sup\left\{\int\varphi(x)\,\mathrm{d}\nu(x)-D_{\rho}(\nu,\mu):\right.$ $\displaystyle\left.\nu\in\mathcal{P}(\mathcal{X}),\,\int|\varphi(x)|\,\mathrm{d}\nu(x)<\infty\right\}.$ (5) $$

(i.e., $F_{\rho}(\cdot;\mu)$ is the convex conjugate of $D_{\rho}(\cdot,\mu)$).
Note that the last term in $f_{\rho}$ prevents the density $\frac{d\nu}{d\mu}$ from being too large.
In particular, it can not be greater than $\frac{1}{\rho}$.
This can make the safe KL divergence a reasonable choice for unbalanced OT or DRO, as it imposes a hard constraint on the reweighting unlike the standard $D_{KL}$.
Moreover, it can also be used instead of the entropy penalization in regularized OT (cf. capacity constrained transport in (Benamou et al., 2015, section 5.2)).

Again, by the convex duality and the variational principle (see Birrell et al., 2022, Theorem 6), we state the following properties.

###### Lemma 2.2 .

The functional $F_{\rho}$ defined by ([5](https://arxiv.org/html/2509.24894v1#S2.E5)) has an equivalent variational representation

$$ $F_{\rho}(\varphi;\mu)=\inf_{\alpha\in\mathbb{R}}\alpha+\int f_{\rho}^{*}(\varphi(x)-\alpha)\,\mathrm{d}\mu(x).$ $$

It is straightforward to check the following.

###### Lemma 2.3 .

The conjugate function to $f_{\rho}$ is a rescaled SoftPlus, specifically,

$$ $f_{\rho}^{*}(s)\coloneqq\sup_{t\in\mathbb{R}_{+}}st-f_{\rho}(t)=\frac{1}{\rho}\log\left(1+\rho e^{s}\right)-1.$ $$

Therefore, we obtain

$$ $F_{\rho}(\varphi;\mu)=\inf_{\alpha\in\mathbb{R}}\alpha-1+\frac{1}{\rho}\int\log\left(1+\rho e^{\varphi(x)-\alpha}\right)\,\mathrm{d}\mu(x).$ (6) $$

In essence, we have replaced the exponential function with a rescaled SoftPlus.
Furthermore, it is easy to see that the optimal $\alpha^{*}$ satisfies

$$ $\int\frac{e^{\varphi(x)-\alpha^{*}}}{1+\rho e^{\varphi(x)-\alpha^{*}}}\,\mathrm{d}\mu(x)=1,$ (7) $$

in particular, $\alpha^{*}<F(\varphi;\mu)$.
Moreover, the maximum in ([5](https://arxiv.org/html/2509.24894v1#S2.E5)) is attained at
$\displaystyle d\nu^{*}_{\rho}(s)=\frac{e^{\varphi(x)-\alpha^{*}}}{1+\rho e^{\varphi(x)-\alpha^{*}}}\,\mathrm{d}\mu(x)$.
Note that $0<\frac{d\nu^{*}_{\rho}(x)}{d\mu(x)}<\frac{1}{\rho}$, which is due to the fact that the derivative of $t\log t$ explodes at 0, preventing reaching the constraint.

The next proposition (proved in Appendix [A](https://arxiv.org/html/2509.24894v1#A1)) ensures that $F_{\rho}$ is a valid approximation of $F$.

###### Proposition 2.4 .

Let $\mu\in\mathcal{P}(\mathcal{X})$ and $\varphi$ be a measurable function on $\mathcal{X}$.

- (i)
For all $0<\rho\leq\rho^{\prime}<1$, it holds $F_{\rho^{\prime}}(\varphi;\mu)\leq F_{\rho}(\varphi;\mu)$.
- (ii)
As $\rho\to 0+,\;F_{\rho}(\varphi;\mu)\to F_{0}(\varphi;\mu)\coloneqq F(\varphi;\mu)$.
- (iii)
If $F(2\varphi;\mu)<\infty$, then for all $0<\rho\leq\frac{1}{4}e^{2F(\varphi;\mu)-F(2\varphi;\mu)}$
$F_{\rho}(\varphi;\mu)\geq F(\varphi;\mu)+\frac{\rho}{2}-4\rho e^{F(2\varphi;\mu)-2F(\varphi;\mu)}.$
(8)
- (iv)
If $\varphi(x)\leq M$ for all $x\in\mathcal{X}$, then $F_{\rho}(\varphi;\mu)\geq F(\varphi;\mu)-\rho e^{M-F(\varphi;\mu)}$ for $\rho\in\left(0,e^{F(\varphi;\mu)-M}\right)$.

In particular, (i) and (iii) show that $F_{\rho}-O(\rho)\leq F\leq F_{\rho}$, and thus the parameter $\rho$ allows one to control the approximation accuracy.
In the case of LogSumExp, the above proposition yields the following simple bounds.

###### Corollary 2.5 .

Let $a_{1},\dots,a_{n}\in\mathbb{R}$.
Then for any $0<\rho<1$

$$ $\operatorname{LogSumExp}(a_{1},\dots,a_{n})-\rho\leq\inf_{\alpha\in\mathbb{R}}\alpha-1+\frac{1}{\rho}\sum_{i=1}^{n}\log(1+\rho e^{a_{i}-\alpha})\leq\operatorname{LogSumExp}(a_{1},\dots,a_{n}).$ $$

For $\rho=1$ our approximation coincides with Bouchard’s bound for LogSumExp (Bouchard, 2007).

###### Definition 2.1 (Safe KL entropy) .

###### Lemma 2.2 .

###### Lemma 2.3 .

###### Proposition 2.4 .

###### Corollary 2.5 .

### 2.1 Links to CVaR

Recall that the conditional value at risk (CVaR) w.r.t. a probability measure $\mu\in\mathcal{P}(\mathcal{X})$ at level $\rho\in(0,1)$, associated with a function $\varphi$, can be defined (in the case of continuous distribution) as

$$ $\operatorname{CVaR}_{\rho}(\varphi;\mu)\coloneqq\mathbb{E}_{X\sim\mu}\left[\varphi(X)\middle|\varphi(X)\geq Q_{1-\rho}\right]=\frac{1}{\rho}\int_{\varphi(x)\geq Q_{1-\rho}}\varphi(x)\,\mathrm{d}\mu(x),$ $$

where $Q_{1-\rho}$ is the $(1-\rho)$-quantile of $\varphi(X)$, $X\sim\mu$ (Rockafellar et al., 2000).
Moreover, by Theorem 1 in Rockafellar et al. (2000) CVaR also has the following variational formulation:

$$ $\operatorname{CVaR}_{\rho}(\varphi;\mu)=\inf_{\alpha\in\mathbb{R}}\alpha+\frac{1}{\rho}\int(\varphi(x)-\alpha)_{+}\,\mathrm{d}\mu(x).$ (9) $$

Remarkably, in Soma & Yoshida (2020) the authors obtained a smooth approximation to CVaR which, up to an additive constant, has the same form as $F_{\rho}$. However, they considered the approximation w.r.t. a different parameter—a ”temperature” inside SoftPlus.
Finally, Levy et al. (2020) proposed another similar smoothed version of CVaR (KL-regularized CVaR) in the context of DRO.
For our approximation, we obtain the following bounds.

###### Proposition 2.6 .

For all $0<\rho<1$ and $\lambda>0$

$$ $\operatorname{CVaR}_{\rho}(\varphi;\mu)+\lambda(\log\rho-1)\leq\lambda F_{\rho}(\varphi/\lambda;\mu)\leq\operatorname{CVaR}_{\rho}(\varphi;\mu)+\lambda\left(\log\rho-1+\frac{1}{\rho}\right).$ (10) $$

###### Proposition 2.6 .

### 2.2 The case of parametric models

In some applications, the function $\varphi$ is defined as the parametric loss function $L(x,\theta)$ and the goal is to minimize objective involving ([1](https://arxiv.org/html/2509.24894v1#S1.E1)) w.r.t. parameter $\theta$ to find the best model from the parametric family. In this section, we study our approximation to ([1](https://arxiv.org/html/2509.24894v1#S1.E1)) in this parametric setting.
Fix some closed parameter set $\Theta\subset\mathbb{R}^{d}$ and a loss function $L\colon\mathcal{X}\times\Theta\to\mathbb{R}$.
Combining our approximation ([6](https://arxiv.org/html/2509.24894v1#S2.E6)) and the minimization w.r.t. parameter $\theta$, we obtain the following minimization problem (note that we shifted $\alpha$ by $\log\rho$ compared to ([6](https://arxiv.org/html/2509.24894v1#S2.E6)))

$$ $\min_{\theta\in\Theta,\alpha\in\mathbb{R}}G_{\rho}(\theta,\alpha)\coloneqq\alpha+\log\rho-1+\frac{1}{\rho}\int\log\left(1+e^{L(x,\theta)-\alpha}\right)\,\mathrm{d}\mu(x).$ $$

Clearly, $G_{\rho}$ is convex in $\alpha$.
Moreover, if $L$ is convex in $\theta$ for $\mu$-a.e. $x$, then $G_{\rho}$ is jointly convex, meaning that our approximation preserves convexity.

Note that
$f_{\rho}(t)=\frac{1}{\rho}\left((\rho t)\log(\rho t)+(1-\rho t)\log(1-\rho t)\right)+1-t\log\rho.$
Thus, unlike the KL entropy function $t\log t+1-t$, $f_{\rho}$
possesses the following favorable properties:

###### Lemma 2.7 .

The entropy function $f_{\rho}$ is $\rho$-strongly convex.
Its conjugate function $f_{\rho}^{*}$ is $\frac{1}{\rho}$-smooth.

The above properties are important from the computational optimization point of view.
Recall that
$\frac{d}{dt}\log(1+e^{t})=\frac{e^{t}}{1+e^{t}}\eqqcolon\sigma(t)$.
Thus, we immediately obtain the following formulas for the gradient:

$$ $\displaystyle\nabla_{\theta}G_{\rho}(\theta,\alpha)=\frac{1}{\rho}\int\sigma\left(L(x,\theta)-\alpha\right)\nabla_{\theta}L(x,\theta)\,\mathrm{d}\mu(x),$ $\displaystyle\partial_{\alpha}G_{\rho}(\theta,\alpha)=1-\frac{1}{\rho}\int\sigma\left(L(x,\theta)-\alpha\right)\,\mathrm{d}\mu(x).$ $$

This yields, in particular, that the variance of the (naïve) stochastic gradient is bounded by $\frac{1}{\rho}$ and the second moment of $\nabla_{\theta}L(X,\theta)$, $X\sim\mu$.
In the same way one can calculate the Hessian of $G_{\rho}$, see Appendix [C](https://arxiv.org/html/2509.24894v1#A3).
By Proposition [C.2](https://arxiv.org/html/2509.24894v1#A3.Thmtheorem2), if $L(x,\theta)$ is bounded from below, then $G_{\rho}$ is smooth on $\Theta\times(-\infty,a]$ for any $a\in\mathbb{R}$ meaning that our approximation preserves smoothness of the loss $L$.

###### Lemma 2.7 .

## 3 Applications

In this section we consider several particular applications involving the objective ([1](https://arxiv.org/html/2509.24894v1#S1.E1)) and show numerically, that our general-purpose approach based on approximation ([6](https://arxiv.org/html/2509.24894v1#S2.E6)) leads to better performance of SGD-type algorithms than the baseline algorithms designed specifically for these applications. The source code for all experiments is available at
[https://github.com/egorgladin/logsumexp-approx](https://github.com/egorgladin/logsumexp-approx).

### 3.1 Continuous entropy-regularized OT

The classical optimal transport (Monge–Kantorovich) problem consists in finding a coupling of two probablility measures $\mu,\nu\in\mathcal{P}(\mathcal{X})$ which minimizes the integral of a given measurable cost function $c\colon\mathcal{X}\times\mathcal{X}\to\mathbb{R}_{+}$ (e.g., a distance), i.e., $W(\mu,\nu)\coloneqq\inf_{\pi\in\Pi(\nu,\mu)}\int c(x,z)\,\mathrm{d}\pi(x,z)$,
where $\Pi(\mu,\nu)\subset\mathcal{P}(\mathcal{X}\times\mathcal{X})$ is the set of couplings (transport plans) of $\mu$ and $\nu$ (see Kantorovich, 1942; Villani, 2008; Santambrogio, 2015).
For simplicity of demonstration, we assume that the measures are defined on the same space $\mathcal{X}$, but the results extend trivially to the case of two different spaces.
Following Cuturi (2013), we consider entropy-regularized optimal transport (eOT) problem:

$$ $\min_{\pi\in\Pi(\nu,\mu)}\int c(x,z)\,\mathrm{d}\pi(x,z)+\varepsilon D_{KL}(\pi,\nu\otimes\mu)$ (11) $$

where $\nu\otimes\mu$ is the product measure.
It is known that eOT admits the following dual and semi-dual formulations (see, e.g., Genevay et al. (2016)):

$$ $W_{\varepsilon}(\mu,\nu)=\underbrace{\max_{u,v\in\mathcal{C}(\mathcal{X})}\iint f_{\varepsilon}(x,y,u,v)\,\mathrm{d}\mu(x)\,\mathrm{d}\nu(y)}_{\text{dual}}=\underbrace{\max_{v\in\mathcal{C}(\mathcal{X})}\int h_{\varepsilon}(x,v)\,\mathrm{d}\mu(x)}_{\text{semi-dual}},$ $$

where

$$ $\displaystyle f_{\varepsilon}(x,y,u,v)$ $\displaystyle\coloneqq u(x)+v(y)-\varepsilon\exp\left(\frac{u(x)+v(y)-c(x,y)}{\varepsilon}\right),$ (12) $\displaystyle h_{\varepsilon}(x,v)$ $\displaystyle\coloneqq\int v(y)\mathrm{d}\nu(y)-\varepsilon\log\left(\int\exp\left(\frac{v(y)-c(x,y)}{\varepsilon}\right)\mathrm{d}\nu(y)\right)-\varepsilon,$ (13) $$

and $\varepsilon>0$ is the regularization coefficient. Genevay et al. (2016) consider a reproducing kernel Hilbert space (RKHS) $\mathcal{H}$ defined on $\mathcal{X}$, with a kernel $\kappa$, and apply SGD to solve the dual problem.

###### Remark 3.1 (The overflow issue) .

The main drawback of this approach is the presence of the exponent in the dual objective (and consequently in the SGD updates, see Appendix [B.1](https://arxiv.org/html/2509.24894v1#A2.SS1)).
Specifically, exponents are prone to floating-point exceptions (Goldberg, 1991), especially if the regularization parameter $\varepsilon$ is relatively small, which is often the case. For example, if $\varepsilon=0.01$ and $z\geq 7.1$, then $e^{z/\varepsilon}$ exceeds the representable range of a double-precision (float64) floating-point number — an overflow happens. When single precision (float32) is used, an overflow happens even for $z\geq 0.89$.

###### Remark 3.1 (The overflow issue) .

#### Our approach.

If we consider instead the semi-dual formulation and use the approximation ([6](https://arxiv.org/html/2509.24894v1#S2.E6)), we arrive at the problem

$$ $\displaystyle\max_{\begin{subarray}{c}v\in\mathcal{C}(\mathcal{X})\\ \alpha\in\mathbb{R}\end{subarray}}\,\iint\,$ $\displaystyle\tilde{h}_{\varepsilon}(x,y,v,\alpha)\,\mathrm{d}\mu(x)\,\mathrm{d}\nu(y)$ (14) with $\displaystyle\tilde{h}_{\varepsilon}(x,y,v,\alpha)\coloneqq v(y)-\alpha-\frac{\varepsilon}{\rho}\log\bigl(1+\rho e^{(v(y)-c(x,y)-\alpha)/\varepsilon}\bigr)-\varepsilon,$ (15) $$

which can also be solved by SGD.
Analytic form of SGD iterates for objectives ([12](https://arxiv.org/html/2509.24894v1#S3.E12)) and ([15](https://arxiv.org/html/2509.24894v1#S3.E15)) can be found in Appendix [B.1](https://arxiv.org/html/2509.24894v1#A2.SS1).
One can show, in the same way as in Genevay et al. (2016), that this corresponds to the regularized OT problem ([11](https://arxiv.org/html/2509.24894v1#S3.E11)) with Safe KL divergence $D_{\rho}$ rather than the usual KL, i.e.

$$ $\min_{\pi\in\Pi(\nu,\mu)}\int c(x,z)\,\mathrm{d}\pi(x,z)+\varepsilon D_{\rho}(\pi,\nu\otimes\mu).$ $$

Note that this problem, in turn, can be viewed as a combination of the entropy-regularized and the capacity-constrained optimal transport.
For $\rho>0$, this approach is much more stable than the previous one when used in SGD. We illustrate this in the following experiments.

Figure: Figure 2: Left: convergence of kernel SGD applied to the dual objective ([12](https://arxiv.org/html/2509.24894v1#S3.E12)) (blue and orange) and approximate semi-dual problem ([14](https://arxiv.org/html/2509.24894v1#S3.E14)) (green, red and purple). Solid lines show average optimality gap across 20 runs, shaded regions indicate $\pm$ one standard deviation. Y-axis uses logarithmic scale. Middle: a zoomed-in view of blue and orange curves from the plot on the left. Right: examples of divergent optimality gap curves obtained by running the baseline approach with the stepsize parameter $C=10^{-2}$.
Refer to caption: fig_new/dual_vs_semidual.png

#### Experiments.

Consider a setup analogous to the one described in Section 5 of Genevay et al. (2016). Specifically, $\mu$ is a 1D Gaussian, and $\nu$ is a
mixture of two Gaussians (see Appendix [B](https://arxiv.org/html/2509.24894v1#A2) for a plot of densities).
Gaussian kernel $\kappa(x,x^{\prime})=\exp\left(-\frac{\|x-x^{\prime}\|^{2}}{\sigma^{2}}\right)$ with a bandwidth hyperparameter $\sigma^{2}>0$ is used.
The regularization coefficient is set to $\varepsilon=0.01$.
We consider kernel SGD applied to the dual objective ([12](https://arxiv.org/html/2509.24894v1#S3.E12)) as a baseline approach (Genevay et al., 2016).
We compare it to the proposed approach, namely, kernel SGD applied to the approximate semi-dual problem ([14](https://arxiv.org/html/2509.24894v1#S3.E14)). For details on how the optimality gap is estimated, see Appendix [B](https://arxiv.org/html/2509.24894v1#A2).

When applying kernel SGD to the dual and approximate semi-dual formulations, we consider hyperparameters $\sigma^{2}\in\{0.1,1,10\}$ (kernel bandwidth), $C\in\{10^{-4},10^{-3},\ldots,10\}$ (stepsize parameter), and $\rho\in\{0.03,0.1,0.3\}$ (approximation accuracy).
Double floating-point precision is used.
In the experiment, the proposed approach works best with $\sigma^{2}=10$, and $C=1$ for $\rho\in\{0.03,0.1\}$, $C=10$ for $\rho=0.3$.
Baseline works best with $\sigma^{2}\in\{0.1,1\}$ and $C=10^{-3}$.
Figure [2](https://arxiv.org/html/2509.24894v1#S3.F2) (left) shows performance of the two approaches. For clarity, we provide a zoomed-in view of the curves generated by the baseline in the middle.
As seen from the figures, the baseline is extremely slow, which happens due to the small stepsize.
Larger values of $C$ lead to numerical instabilities as illustrated by the plot on the right.
Apparently, exponent causes a large magnitude of the gradient at a certain step, which brings an iterate to a region where it stagnates.
On the contrary, our approximate semi-dual formulation permits larger stepsizes, which results in faster convergence. Indeed, the method usually achieves a relatively low optimality gap in about $2\cdot 10^{4}$ iterations, and plateaus after that.

### 3.2 Distributionally robust optimization with KL divergence

One of the approaches to training a model that is robust to data distribution shifts and noisy observations
is called Distributionally Robust Optimization (DRO) (Kuhn et al., 2024).
In contrast to the standard Empirical Risk Minimization (ERM) approach, which minimizes the average loss on the training sample, DRO minimizes the risk for the worst-case distribution among those close to a reference measure (e.g., empirical distribution). A prominent example is KL divergence DRO (Hu & Hong,), which is formulated as the saddle-point problem

$$ $\min_{\theta\in\Theta}\,\max_{p\in\Delta^{n}}\;\sum_{i=1}^{n}p_{i}\ell_{i}(\theta)-\lambda D_{KL}(p,\hat{p}),$ (16) $$

where $\theta\in\Theta$ is the model parameters, $\ell_{i}(\theta)$ is the respective loss on the $i$-th training example, $\Delta^{n}$ is the unit simplex in $\mathbb{R}^{n}$, $\hat{p}\in\Delta^{n}$ is the weight vector defining the empirical distribution (typically $\hat{p}=\frac{1}{n}\mathbf{1}$), and $D_{KL}$ is the Kullback–Leibler divergence
which discourages distributions that are too far from the empirical one,
$\lambda>0$ is the penalty coefficient.
For fixed $\theta$, the solution of the maximization problem is given by
$p_{i}^{*}(\theta)\coloneqq\frac{e^{\ell_{i}(\theta)/\lambda}}{\sum_{j}e^{\ell_{j}(\theta)/\lambda}}$,
which reduces the problem to

$$ $\min_{\theta\in\Theta}\mathcal{L}(\theta)\coloneqq\lambda\log\Bigl(\frac{1}{n}\sum_{i=1}^{n}e^{\ell_{i}(\theta)/\lambda}\Bigr).$ (17) $$

However, when $n$ is large, computing the full gradient $\nabla\mathcal{L}(\theta)=\sum_{i=1}^{n}p_{i}^{*}(\theta)\nabla\ell_{i}(\theta)$ becomes costly.
A straightforward approach (Levy et al., 2020) is to sample a batch $D$, compute the respective softmax weights $p_{i}^{D}(\theta)\coloneqq\frac{e^{\ell_{i}(\theta)/\lambda}}{\sum_{j\in D}e^{\ell_{j}(\theta)/\lambda}}$, and define a gradient estimator by

$$ $\tilde{\nabla}_{D}\mathcal{L}(\theta)=\sum_{i\in D}p_{i}^{D}(\theta)\nabla\ell_{i}(\theta).$ (18) $$

However, this introduces a bias and requires using large batch sizes to keep it sufficiently small.

#### Our approach.

Instead, we propose to use the approximation ([6](https://arxiv.org/html/2509.24894v1#S2.E6)), which results in the problem

$$ $\min_{\begin{subarray}{c}\theta\in\Theta\\ \alpha\in\mathbb{R}\end{subarray}}\,G(\theta,\alpha)\coloneqq\frac{1}{n}\sum_{i=1}^{n}\Bigl\{\alpha+\frac{\lambda}{\rho}\log\bigl(1+\rho e^{(\ell_{i}(\theta)-\alpha)/\lambda}\bigr)\Bigr\}.$ $$

Like in the previous subsection, this can be interpreted as switching from $D_{KL}$ penalty in ([16](https://arxiv.org/html/2509.24894v1#S3.E16)) to Safe KL $D_{\rho}$.
The respective gradient estimators are

$$ $\displaystyle\tilde{\nabla}^{D}_{\theta}G(\theta,\alpha)\coloneqq\frac{1}{|D|}\sum_{i\in D}\sigma_{\rho}\left(\textstyle\frac{\ell_{i}(\theta)-\alpha}{\lambda}\right)\nabla\ell_{i}(\theta),\quad\tilde{\nabla}^{D}_{\alpha}G(\theta,\alpha)\coloneqq 1-\frac{1}{|D|}\sum_{i\in D}\sigma_{\rho}\left(\textstyle\frac{\ell_{i}(\theta)-\alpha}{\lambda}\right).$ (19) $$

#### Experiments.

Consider California housing dataset (Pace & Barry, 1997) consisting of 20640 objects represented by 8 features. Let $\ell_{i}$ be the squared error of a linear model, $\ell_{i}(\theta)=(y_{i}-\theta^{\top}x_{i})^{2}$. We use SGD with Nesterov acceleration and the gradient estimator ([18](https://arxiv.org/html/2509.24894v1#S3.E18)) (Levy et al., 2020) as the baseline approach for solving ([17](https://arxiv.org/html/2509.24894v1#S3.E17)). As an alternative, accelerated SGD is used with our proposed gradient estimator ([19](https://arxiv.org/html/2509.24894v1#S3.E19)).
We considered values of the stepsize $\eta\in\{10^{-8},10^{-7},\ldots,10^{-4}\}$ and approximation accuracy parameter $\rho\in\{10^{-4},10^{-3},10^{-2}\}$. Momentum was set to 0.9 (without tuning) and entropy penalty $\lambda$ to 1. Least squares solution was used as a starting point for optimization.

Figure: (a) Batch size 10
Refer to caption: fig_new/bar_batch10.png

Convergence curves for both approaches corresponding to the best hyperparameters (i.e., resulting in smallest objective values) are displayed in Figure [3](https://arxiv.org/html/2509.24894v1#S3.F3). Y-axis depicts the value of $\mathcal{L}(\theta)$, X-axis displays the number of epochs (i.e., full passes over the dataset), error bars denote $\pm$ one standard deviation across 20 runs. We omit first few epochs to zoom in on smaller objective values. As seen from the figure, convergence of the proposed approach is consistently good across all batch sizes. As for the baseline, $|D|=10$ requires small stepsize $\eta=10^{-6}$ to avoid divergence, and thus results in slow convergence. Larger batches enable the increase of $\eta$ up to $10^{-5}$. The baseline still converges slower than the proposed approach with $|D|=100$, but almost matches its performance for $|D|=1000$.

### 3.3 Distributionally robust optimization with unbalanced OT

In the KL divergence DRO described in the previous subsection,
uncertainty set is limited to distributions with the same support as the empirical measure $\mu=\frac{1}{n}\sum_{i}\delta_{x_{i}}$.
Another popular approach, Wasserstein DRO (WDRO) (Mohajerin Esfahani & Kuhn, 2018; Sinha et al., 2020),
considers the worst-case risk over shifts within a Wasserstein (OT) ball around a reference measure $\mu$
instead of the KL-ball in ([16](https://arxiv.org/html/2509.24894v1#S3.E16)),
thus including continuous probability measures.
Unfortunately, this approach is not resilient to outliers that are geometrically far from the clean distribution since
OT metric is sensitive to them (Nietert et al., 2023).
A natural generalization is to switch to semi-balanced OT (Liero et al., 2018; Chizat et al., 2019; Kondratyev et al., 2016), which replaces a hard constraint on one of the marginals with a mismatch penalty function, e.g.,

$$ $W_{\beta}(\nu,\mu)=\inf_{\begin{subarray}{c}\pi\in\mathcal{P}(\mathcal{X}\times\mathcal{X})\\ \pi_{1}=\nu\end{subarray}}\int c(x,z)\,\mathrm{d}\pi(x,z)+\beta\,D_{KL}(\pi_{2},\mu),$ $$

where $\pi_{1}$ and $\pi_{2}$ are first and second marginals of $\pi$, respectively,
$\beta>0$ is the marginal penalty parameter. Intuitively, this discrepancy measure allows to ignore some points (e.g., outliers) by paying a small price for mismatch in marginals.
The (penalty-form) DRO problem can be written as

$$ $\min_{\theta\in\Theta}\,\max_{\nu\in\mathcal{P}(\mathcal{X})}\;\int\ell(\theta,x)\,\mathrm{d}\nu(x)-\lambda W_{\beta}(\nu,\mu),$ $$

where $\lambda>0$ is the Lagrangian penalty parameter.
Using standard duality,
Wang et al. (2024) showed that when
$\mu=\frac{1}{n}\sum_{i}\delta_{x_{i}}$ is the empirical distribution,
this is equivalent to

$$ $\min_{\theta\in\Theta}F(\theta)\coloneqq\lambda\beta\,\log\Bigl(\frac{1}{n}\sum_{i=1}^{n}e^{\hat{\ell}_{i}(\theta)/(\lambda\beta)}\Bigr)\quad\text{with }\;\hat{\ell}_{i}(\theta)\coloneqq\sup_{z\in\mathcal{X}}\{\ell(\theta;z)-\lambda c(z,x_{i})\},$ (20) $$

To avoid the costly gradient computation of LogSumExp, the authors drop the logarithm and use SGD to optimize the sum of exponents,

$$ $\min_{\theta\in\Theta}\frac{1}{n}\sum_{i=1}^{n}e^{\hat{\ell}_{i}(\theta)/(\lambda\beta)}.$ (21) $$

The major downside of this approach is that the exponent terms have a large variance, and SGD is prone to floating-point exceptions (overflow)
unless a very small stepsize is tuned, which slows down the convergence and can be time-consuming and unstable in practice.

#### Our approach.

To overcome this issue, we propose leveraging the approximation ([6](https://arxiv.org/html/2509.24894v1#S2.E6)), which leads to the problem

$$ $\min_{\begin{subarray}{c}\theta\in\Theta\\ \alpha\in\mathbb{R}\end{subarray}}\,\frac{1}{n}\sum_{i=1}^{n}\Bigl\{\alpha+\frac{\lambda\beta}{\rho}\log\bigl(1+\rho e^{(\hat{\ell}_{i}(\theta)-\alpha)/(\lambda\beta)}\bigr)\Bigr\},$ (22) $$

where $\rho>0$ is a parameter controlling the accuracy of the approximation. This approximation can be efficiently optimized with SGD.
Note that our method can also be applied to other DRO algorithms such as Sinkhorn DRO (Wang et al., 2021), which we omit to avoid redundancy.

#### Experiments.

We consider MNIST dataset (Deng, 2012) with train and validation labels corrupted by feature-dependent noise (see Algan & Ulusoy, 2020) (noise ratio 25%), and original (clean) test labels.
Let $\theta$ denote weights of a CNN with two convolutional layers (32 and 64 channels, kernel size 3, ReLU activations, and 2$\times$2 max pooling), followed by a fully connected classifier with one hidden layer of 128 units, and let $\ell(\theta;z)$ be its cross entropy loss on object $z$.
In the experiment,
SGD (with batch size 1) is applied to problems ([21](https://arxiv.org/html/2509.24894v1#S3.E21)) (baseline) and ([22](https://arxiv.org/html/2509.24894v1#S3.E22)) (proposed approach).
We consider values of the stepsize $\eta\in\{10^{-6},10^{-5},10^{-4},10^{-3}\}$.
Parameters $\lambda$ and $\beta$ were set to 1 since smaller values required a smaller stepsize and resulted in a slow convergence, while larger values pushed the model towards fitting the noisy distribution instead of the true one.
Approximation accuracy parameter $\rho$ in ([22](https://arxiv.org/html/2509.24894v1#S3.E22)) was set to 0.1.
For the inner maximization problem in ([20](https://arxiv.org/html/2509.24894v1#S3.E20)), just 5 iterations of Nesterov’s accelerated gradient method were sufficient to reach plateau in terms of the objective value.
Additionally, we used SGD for the usual empirical risk minimization (ERM) to observe the effects of conventional (non-robust) training on noisy data.

Figure: Figure 4: Performance of ERM and two DRO approaches on MNIST with noisy labels. Left: ERM accuracy on the noisy validation set vs. clean test set. Middle: validation vs. test accuracy for DRO approaches. Right: training loss $F(\theta)$ from ([20](https://arxiv.org/html/2509.24894v1#S3.E20)).
Refer to caption: fig_new/accuracy_erm.png

Figure [4](https://arxiv.org/html/2509.24894v1#S3.F4) demonstrates the performance of different approaches with best hyperparameter $\eta$ ($10^{-3}$ for ERM, $10^{-4}$ for the proposed approach, and $10^{-5}$ for the baseline). Shaded regions indicate $\pm$ one standard deviation across 10 runs, except that for the baseline approach ([21](https://arxiv.org/html/2509.24894v1#S3.E21)) we excluded a single run that caused a floating-point exception.
The plot on the left illustrates that ERM fits to the corrupted data well (accuracy on the noisy validation set is increasing) which results in decreasing accuracy on the clean test set.
In contrast, the plot in the middle shows that an increase in validation accuracy results in the increase in test accuracy for both DRO approaches, which indicates that they are more capable at learning the underlying clean distribution. The plot on the right shows the train loss $F(\theta)$ from ([20](https://arxiv.org/html/2509.24894v1#S3.E20)). As seen from the figure, the proposed approach converges faster than the baseline. This is caused by the fact that the baseline requires a small stepsize to avoid overflows.

## 4 Conclusion

We introduce a novel approximation to the log partition function (and in particular, to LogSumExp), which arises in numerous applications across machine learning and optimization.
In the dual formulation, it corresponds to the safe KL divergence.
Our LogSumExp approximation preserves convexity and smoothness, and can be efficiently minimized using stochastic gradient methods.
Importantly, the respective gradient estimator has controllable bias independent of batch size, in contrast to prior approaches.
Our empirical results highlight the practical advantages of the proposed approximation across tasks in continuous entropy-regularized OT and DRO.
An important direction for future work is to leverage the approximation for other applications, where the LogSumExp function and duality of the KL divergence play a role.

## References

- Algan & Ulusoy (2020)
Görkem Algan and Ilkay Ulusoy.
Label noise types and their effects on deep learning.
*arXiv preprint arXiv:2003.10471*, 2020.
- Amari & Nagaoka (2000)
Shun-ichi Amari and Hiroshi Nagaoka.
*Methods of Information Geometry*.
American Mathematical Soc., 2000.
ISBN 978-0-8218-4302-4.
- Ben-Tal et al. (2013)
Aharon Ben-Tal, Dick den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen.
Robust solutions of optimization problems affected by uncertain probabilities.
*Management Science*, 59(2):341–357, February 2013.
ISSN 0025-1909, 1526-5501.
doi: 10.1287/mnsc.1120.1641.
URL [http://pubsonline.informs.org/doi/abs/10.1287/mnsc.1120.1641](http://pubsonline.informs.org/doi/abs/10.1287/mnsc.1120.1641).
- Benamou et al. (2015)
Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré.
Iterative bregman projections for regularized transportation problems.
*SIAM Journal on Scientific Computing*, 37(2):A1111–A1138, 2015.
- Birrell et al. (2022)
Jeremiah Birrell, Paul Dupuis, Markos A. Katsoulakis, Yannis Pantazis, and Luc Rey-Bellet.
(f,gamma)-divergences: Interpolating between f-divergences and integral probability metrics.
*Journal of Machine Learning Research*, 23(39):1–70, 2022.
URL [http://jmlr.org/papers/v23/21-0100.html](http://jmlr.org/papers/v23/21-0100.html).
- Bishop & Nasrabadi (2006)
Christopher M Bishop and Nasser M Nasrabadi.
*Pattern recognition and machine learning*, volume 4.
Springer, 2006.
- Bouchard (2007)
Guillaume Bouchard.
Efficient bounds for the softmax function and applications to approximate inference in hybrid models.
In *NIPS 2007 workshop for approximate Bayesian inference in continuous/hybrid systems*, volume 6, 2007.
- Chizat et al. (2019)
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard.
Unbalanced optimal transport: Dynamic and Kantorovich formulation.
*arXiv:1508.05216 [math]*, February 2019.
URL [http://arxiv.org/abs/1508.05216](http://arxiv.org/abs/1508.05216).
arXiv: 1508.05216.
- Cuturi (2013)
Marco Cuturi.
Sinkhorn distances: Lightspeed computation of optimal transport.
*Advances in neural information processing systems*, 26, 2013.
- Deng (2012)
Li Deng.
The mnist database of handwritten digit images for machine learning research.
*IEEE Signal Processing Magazine*, 29(6):141–142, 2012.
- Genevay et al. (2016)
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach.
Stochastic optimization for large-scale optimal transport.
In D. D. Lee, M. Sugiyama, U. V. Luxburg, I. Guyon, and R. Garnett (eds.), *Advances in Neural Information Processing Systems 29*, pp. 3440–3448. Curran Associates, Inc., 2016.
- Goldberg (1991)
David Goldberg.
What every computer scientist should know about floating-point arithmetic.
*ACM computing surveys (CSUR)*, 23(1):5–48, 1991.
- (13)
Zhaolin Hu and L Jeff Hong.
Kullback-Leibler Divergence Constrained Distributionally Robust Optimization.
- Kantorovich (1942)
Leonid Kantorovich.
On the translocation of masses.
*(Doklady) Acad. Sci. URSS (N.S.)*, 37:199–201, 1942.
- Khan & Nielsen (2018)
Mohammad Emtiyaz Khan and Didrik Nielsen.
Fast yet simple natural-gradient descent for variational inference in complex models.
In *2018 International Symposium on Information Theory and Its Applications (ISITA)*, pp. 31–35. IEEE, 2018.
- Khan & Rue (2023)
Mohammad Emtiyaz Khan and Håvard Rue.
The Bayesian learning rule.
(arXiv:2107.04562), June 2023.
- Kondratyev et al. (2016)
Stanislav Kondratyev, Léonard Monsaingeon, and Dmitry Vorotnikov.
A new optimal transport distance on the space of finite Radon measures.
*Advances in Differential Equations*, 21(11/12):1117 – 1164, 2016.
doi: 10.57262/ade/1476369298.
URL [https://doi.org/10.57262/ade/1476369298](https://doi.org/10.57262/ade/1476369298).
- Kuhn et al. (2024)
Daniel Kuhn, Soroosh Shafiee, and Wolfram Wiesemann.
Distributionally Robust Optimization.
(arXiv:2411.02549), November 2024.
doi: 10.48550/arXiv.2411.02549.
- Levy et al. (2020)
Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford.
Large-scale methods for distributionally robust optimization.
*Advances in neural information processing systems*, 33:8847–8860, 2020.
- Liero et al. (2018)
Matthias Liero, Alexander Mielke, and Giuseppe Savaré.
Optimal entropy-transport problems and a new Hellinger–Kantorovich distance between positive measures.
*Inventiones mathematicae*, 211(3):969–1117, March 2018.
ISSN 0020-9910, 1432-1297.
doi: 10.1007/s00222-017-0759-8.
URL [http://link.springer.com/10.1007/s00222-017-0759-8](http://link.springer.com/10.1007/s00222-017-0759-8).
- Mohajerin Esfahani & Kuhn (2018)
Peyman Mohajerin Esfahani and Daniel Kuhn.
Data-driven distributionally robust optimization using the Wasserstein metric: Performance guarantees and tractable reformulations.
*Mathematical Programming*, 171(1-2):115–166, September 2018.
ISSN 0025-5610, 1436-4646.
doi: 10.1007/s10107-017-1172-1.
- Nielsen & Sun (2016)
Frank Nielsen and Ke Sun.
Guaranteed bounds on information-theoretic measures of univariate mixtures using piecewise log-sum-exp inequalities.
*Entropy*, 18(12):442, 2016.
- Nietert et al. (2023)
Sloan Nietert, Ziv Goldfeld, and Soroosh Shafiee.
Outlier-robust wasserstein dro.
*Advances in Neural Information Processing Systems*, 36:62792–62820, 2023.
- Pace & Barry (1997)
R Kelley Pace and Ronald Barry.
Sparse spatial autoregressions.
*Statistics & Probability Letters*, 33(3):291–297, 1997.
- Pee & Royset (2011)
EY Pee and Johannes O Royset.
On solving large-scale finite minimax problems using exponential smoothing.
*Journal of optimization theory and applications*, 148(2):390–421, 2011.
- Peyré & Cuturi (2019)
Gabriel Peyré and Marco Cuturi.
Computational optimal transport: With applications to data science.
*Foundations and Trends® in Machine Learning*, 11(5-6):355–607, 2019.
ISSN 1935-8237.
doi: 10.1561/2200000073.
URL [http://dx.doi.org/10.1561/2200000073](http://dx.doi.org/10.1561/2200000073).
arXiv:1803.00567.
- Polyanskiy & Wu (2025)
Yury Polyanskiy and Yihong Wu.
*Information theory: From coding to learning*.
Cambridge university press, 2025.
- Rockafellar et al. (2000)
R Tyrrell Rockafellar, Stanislav Uryasev, et al.
Optimization of conditional value-at-risk.
*Journal of risk*, 2:21–42, 2000.
- Santambrogio (2015)
Filippo Santambrogio.
*Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs, and Modeling*.
Springer International Publishing, 2015.
ISBN 9783319208282.
doi: 10.1007/978-3-319-20828-2.
URL [http://dx.doi.org/10.1007/978-3-319-20828-2](http://dx.doi.org/10.1007/978-3-319-20828-2).
- Sinha et al. (2020)
Aman Sinha, Hongseok Namkoong, Riccardo Volpi, and John Duchi.
Certifying some distributional robustness with principled adversarial training.
*arXiv:1710.10571 [cs, stat]*, May 2020.
- Soma & Yoshida (2020)
Tasuku Soma and Yuichi Yoshida.
Statistical learning with conditional value at risk.
*arXiv preprint arXiv:2002.05826*, 2020.
- Tiapkin et al. (2024)
Daniil Tiapkin, Denis Belomestny, Daniele Calandriello, Eric Moulines, Alexey Naumov, Pierre Perrault, Michal Valko, and Pierre Menard.
Demonstration-regularized RL.
In *The Twelfth International Conference on Learning Representations*, 2024.
URL [https://openreview.net/forum?id=lF2aip4Scn](https://openreview.net/forum?id=lF2aip4Scn).
- Titsias (2016)
Michalis K. Titsias.
One-vs-each approximation to softmax for scalable estimation of probabilities.
*Advances in Neural Information Processing Systems*, 29, 2016.
- Villani (2008)
Cédric Villani.
*Optimal transport: old and new*, volume 338.
Springer Science & Business Media, 2008.
- Wainwright et al. (2008)
Martin J Wainwright, Michael I Jordan, et al.
Graphical models, exponential families, and variational inference.
*Foundations and Trends® in Machine Learning*, 1(1–2):1–305, 2008.
- Wang et al. (2021)
Jie Wang, Rui Gao, and Yao Xie.
Sinkhorn distributionally robust optimization.
*arXiv preprint arXiv:2109.11926*, 2021.
- Wang et al. (2024)
Zifan Wang, Yi Shen, Michael Zavlanos, and Karl H Johansson.
Outlier-robust distributionally robust optimization via unbalanced optimal transport.
*Advances in Neural Information Processing Systems*, 37:52189–52214, 2024.

## Appendix A Proofs for Section 2

###### Proof of Proposition 2.4 .

(i,ii) Consider the function $g(t)\coloneqq\frac{\ln(1+t)}{t}$.
It is decreasing and convex on $(0,\infty)$, $g(t)\to 1$ and $g^{\prime}(t)\to-\frac{1}{2}$ as $t\to 0+$.
Note that

$$ $F_{\rho}(\varphi;\mu)=\inf_{\alpha\in\mathbb{R}}\alpha-1+\int e^{\varphi(x)-\alpha}g\left(\rho e^{\varphi(x)-\alpha}\right)\,\mathrm{d}\mu(x).$ $$

Then (i) follows immediately from ([6](https://arxiv.org/html/2509.24894v1#S2.E6)) and the monotonicity of $g$.
The monotone convergence theorem yields (ii) since

$$ $F(\varphi;\mu)=\inf_{\alpha\in\mathbb{R}}\alpha-1+\int e^{\varphi(x)-\alpha}\,\mathrm{d}\mu(x).$ $$

Now, let us prove (iii). Consider the optimal $\alpha_{\rho}$, satisfying ([7](https://arxiv.org/html/2509.24894v1#S2.E7)).
By Jensen’s inequality

$$ $\displaystyle\int\ln\left(1+\rho e^{\varphi(x)-\alpha_{\rho}}\right)\,\mathrm{d}\mu(x)$ $\displaystyle=-\int\ln\left(1-\frac{\rho e^{\varphi(x)-\alpha_{\rho}}}{1+\rho e^{\varphi(x)-\alpha_{\rho}}}\right)\,\mathrm{d}\mu(x)$ $\displaystyle\geq-\ln\left(1-\int\frac{\rho e^{\varphi(x)-\alpha_{\rho}}}{1+\rho e^{\varphi(x)-\alpha_{\rho}}}\,\mathrm{d}\mu(x)\right)=-\ln(1-\rho),$ $$

thus

$$ $F_{\rho}(\varphi;\mu)=\alpha_{\rho}-1+\frac{1}{\rho}\int\ln\left(1+\rho e^{\varphi(x)-\alpha_{\rho}}\right)\,\mathrm{d}\mu(x)\geq\alpha_{\rho}-1-\frac{\ln(1-\rho)}{\rho}\geq\alpha_{\rho}+\frac{\rho}{2}.$ (23) $$

It remains to get a lower bound on $\alpha_{\rho}$.
By the monotonicity of $\frac{t}{1+t}$ we deduce that $\alpha_{\rho}\geq\alpha$ for any $\alpha$ such that

$$ $\int\frac{\rho e^{\varphi(x)-\alpha}}{1+\rho e^{\varphi(x)-\alpha}}\,\mathrm{d}\mu(x)\geq\rho.$ $$

Since $\frac{t}{1+t}\geq t-t^{2}$,

$$ $\int\frac{\rho e^{\varphi(x)-\alpha}}{1+\rho e^{\varphi(x)-\alpha}}\,\mathrm{d}\mu(x)\geq\int\left(\rho e^{\varphi(x)-\alpha}-\rho^{2}e^{2\varphi(x)-2\alpha}\right)\,\mathrm{d}\mu(x)=\rho e^{F(\varphi;\mu)-\alpha}-\rho^{2}e^{F(2\varphi;\mu)-2\alpha}.$ $$

Denoting $u\coloneqq e^{F(\varphi;\mu)-\alpha}$, it is enough to find $u$ such that

$$ $u-au^{2}\geq 1,~~\text{where}~~a\coloneqq\rho e^{F(2\varphi;\mu)-2F(\varphi;\mu)}\leq\frac{1}{4}.$ $$

Thus, taking

$$ $u\coloneqq\frac{1}{2a}\left(1-\sqrt{1-4a}\right)\leq 1+4a,$ $$

we obtain

$$ $\alpha_{\rho}\geq F(\varphi;\mu)-\ln u\geq F(\varphi;\mu)-\ln\left(1+4a\right)\geq F(\varphi;\mu)-4a.$ $$

Combining this with ([23](https://arxiv.org/html/2509.24894v1#A1.E23)), we get ([8](https://arxiv.org/html/2509.24894v1#S2.E8)).

(iv) Finally, let $\varphi(x)\leq M$ for all $x\in\mathcal{X}$. Then by concavity

$$ $\int\ln\left(1+\rho e^{\varphi(x)-\alpha}\right)\,\mathrm{d}\mu(x)\geq\int e^{\varphi(x)-M}\ln\left(1+\rho e^{M-\alpha}\right)\,\mathrm{d}\mu(x)=e^{F(\varphi;\mu)-M}\ln\left(1+\rho e^{M-\alpha}\right)$ $$

for all $\alpha\in\mathbb{R}$. Therefore,

$$ $\displaystyle F_{\rho}(\varphi;\mu)$ $\displaystyle\geq\min_{\alpha}\alpha-1+\frac{e^{F(\varphi;\mu)-M}}{\rho}\ln\left(1+\rho e^{M-\alpha}\right)$ $\displaystyle=F(\varphi;\mu)-1-\frac{1-\rho e^{M-F(\varphi;\mu)}}{\rho e^{M-F(\varphi;\mu)}}\ln\left(1-\rho e^{M-F(\varphi;\mu)}\right)$ $\displaystyle\geq F(\varphi;\mu)-\rho e^{M-F(\varphi;\mu)}.$ $$

Here we used the inequality

$$ $\frac{1-t}{t}\ln(1-t)\leq t-1,\quad 0<t<1.$ $$

∎

###### Proof of Corollary 2.5 .

Set $\mu_{n}\coloneqq\frac{1}{n}\sum_{i=1}^{n}\delta_{a_{i}}\in\mathcal{P}(\mathbb{R})$.
Then

$$ $\operatorname{LogSumExp}(a_{1},\dots,a_{n})=\ln n+\ln\left(\int x\,\mathrm{d}\mu_{n}(x)\right)=\ln n+F(id;\mu_{n})$ $$

and

$$ $\displaystyle\inf_{\alpha\in\mathbb{R}}\alpha-1+\frac{1}{\rho}\sum_{i=1}^{n}\ln(1+\rho e^{a_{i}-\alpha})$ $\displaystyle=\inf_{\alpha\in\mathbb{R}}\alpha-1+\frac{n}{\rho}\int\ln(1+\rho e^{x-\alpha})\,\mathrm{d}\mu_{n}(x)$ $\displaystyle=\inf_{\alpha\in\mathbb{R}}\alpha-1+\frac{n}{\rho}\int\ln\left(1+\frac{\rho}{n}e^{x-\alpha+\ln n}\right)\,\mathrm{d}\mu_{n}(x)$ $\displaystyle=\ln n+F_{\rho/n}(id;\mu_{n}).$ $$

Since

$$ $e^{F(id;\mu_{n})-\max_{i}a_{i}}=\frac{\sum_{i=1}^{n}e^{a_{i}}}{n\max_{i}e^{a_{i}}}\geq\frac{1}{n}>\frac{\rho}{n},$ $$

Proposition [2.4](https://arxiv.org/html/2509.24894v1#S2.Thmtheorem4)(i,iv) yields

$$ $F(id;\mu_{n})-\rho\leq F_{\rho/n}(id;\mu_{n})\leq F(id;\mu_{n}).$ $$

The claim follows.
∎

###### Proof of Proposition 2.6 .

As $\lambda\ln(1+e^{t/\lambda})>t_{+}:=\max\{0,t\}$, we get

$$ $\displaystyle\lambda F_{\rho}(\varphi/\lambda;\mu)$ $\displaystyle\geq\lambda\inf_{\alpha\in\mathbb{R}}\alpha-1+\frac{1}{\rho}\int\left(\ln\rho+\frac{\varphi(x)}{\lambda}-\alpha\right)_{+}\,\mathrm{d}\mu(x)$ $\displaystyle=\lambda(\ln\rho-1)+\inf_{\alpha\in\mathbb{R}}\alpha+\frac{1}{\rho}\int(\varphi(x)-\alpha)_{+}\,\mathrm{d}\mu(x).$ $$

The infimum in the r.h.s. is the variational formula for CVaR ([9](https://arxiv.org/html/2509.24894v1#S2.E9)), thus we get the first inequality in ([10](https://arxiv.org/html/2509.24894v1#S2.E10)).
The second inequality can be obtained in a similar way using that $\lambda\ln(1+e^{t/\lambda})<t_{+}+\lambda$.
∎

###### Proof of Proposition 2.4 .

###### Proof of Corollary 2.5 .

###### Proof of Proposition 2.6 .

## Appendix B Additional details on the eOT experiment

### B.1 SGD iterates for the dual and approximate semi-dual formulations

By the property of RKHS, if $u\in\mathcal{H}$, then $u(x)=\langle u,\kappa(\cdot,x)\rangle_{\mathcal{H}}$, and the derivatives of $f_{\varepsilon}$ are as follows:

$$ $\displaystyle\nabla_{u}f_{\varepsilon}(x,y,u,v)$ $\displaystyle=\kappa(\cdot,x)-\exp\left(\frac{u(x)+v(y)-c(x,y)}{\varepsilon}\right)\kappa(\cdot,x),$ $\displaystyle\nabla_{v}f_{\varepsilon}(x,y,u,v)$ $\displaystyle=\kappa(\cdot,y)-\exp\left(\frac{u(x)+v(y)-c(x,y)}{\varepsilon}\right)\kappa(\cdot,y).$ $$

Thus, the SGD iterates for the dual objective ([12](https://arxiv.org/html/2509.24894v1#S3.E12)) take the simple form

$$ $\displaystyle\left(u_{k},v_{k}\right)$ $\displaystyle=\left(u_{0},v_{0}\right)+\sum_{i=1}^{k}\beta_{i}\left(\kappa\left(\cdot,x_{i}\right),\kappa\left(\cdot,y_{i}\right)\right)$ (24) $\displaystyle\text{ with }\,\beta_{i}$ $\displaystyle\coloneqq\frac{C}{\sqrt{i}}\Bigl(1-e^{\frac{u_{i-1}\left(x_{i}\right)+v_{i-1}\left(y_{i}\right)-c\left(x_{i},y_{i}\right)}{\varepsilon}}\Bigr),$ (25) $$

where $(x_{i},y_{i})$ are i.i.d. samples from $\mu\otimes\nu$, and
$C>0$ is the initial stepsize.
Similarly, SGD iterates for ([15](https://arxiv.org/html/2509.24894v1#S3.E15)) are computed as follows:

$$ $\displaystyle v_{k}$ $\displaystyle=v_{0}+\sum_{i=1}^{k}\tilde{\beta}_{i}\kappa\left(\cdot,y_{i}\right),$ $\displaystyle\alpha_{k}$ $\displaystyle=\alpha_{0}-\sum_{i=1}^{k}\tilde{\beta}_{i}\quad\text{with }\;\beta_{i}\coloneqq\frac{C}{\sqrt{i}}\Bigl(1-\sigma_{\rho}\left(\textstyle\frac{u_{i-1}\left(x_{i}\right)+v_{i-1}\left(y_{i}\right)-c\left(x_{i},y_{i}\right)}{\varepsilon}\right)\Bigr),$ $$

where $\sigma_{\rho}(t)\coloneqq\frac{e^{t}}{1+\rho e^{t}}$.

### B.2 Densities of source and target distributions

Figure [5](https://arxiv.org/html/2509.24894v1#A2.F5) shows the densities of source and target distributions ($\mu$ and $\nu$, respectively) used in the experiment in Section [3.1](https://arxiv.org/html/2509.24894v1#S3.SS1).
Specifically, $\mu$ is a 1D Gaussian, and $\nu$ is a mixture of two Gaussians.

Figure: Figure 5: Densities of source and target distributions in the eOT experiment.
Refer to caption: fig_new/1d_distr.png

### B.3 Computing a proxy for optimality gap

Optimality gap in the experiment is estimated as follows:

- 1.
Test sets $\{x_{i}\}_{i=1}^{N}$ and $\{y_{i}\}_{i=1}^{N}$ of size $N=10^{4}$ are sampled from $\mu$ and $\nu$. The corresponding empirical distributions are denoted $\hat{\mu}$ and $\hat{\nu}$, respectively.
- 2.
Similarly to Genevay et al. (2016), we obtain a proxy $\hat{W}$ for $W(\mu,\nu)$ by solving the semi-discrete eOT problem
$\displaystyle\max_{\mathbf{v}\in\mathbb{R}^{N}}\,\mathbb{E}_{X\sim\mu}\,$
$\displaystyle\hat{h}_{\varepsilon}(X,\mathbf{v})$
with
$\displaystyle\hat{h}_{\varepsilon}(x,\mathbf{v})\coloneqq\frac{1}{N}\sum_{i=1}^{N}v_{i}-\varepsilon\log\Bigl(\frac{1}{N}\sum_{i=1}^{N}e^{\frac{v_{i}-c(x,y_{i})}{\varepsilon}}\Bigr)-\varepsilon,$
which corresponds to replacing the expectation $\mathbb{E}_{Y\sim\nu}$ in ([13](https://arxiv.org/html/2509.24894v1#S3.E13)) with the average over the test set $\mathbb{E}_{Y\sim\hat{\nu}}$.
We preform 10 runs of SGD, each consisting of $2\cdot 10^{5}$ iterations, and define $\hat{W}$ as largest achieved value on the test set, i.e., the largest $\mathbb{E}_{X\sim\hat{\mu}}\,\hat{h}_{\varepsilon}(X,\mathbf{v})$.
- 3.
Finally, given a potential $v\in\mathcal{C}(\mathcal{X})$, we estimate the optimality gap as
$\hat{W}-\mathbb{E}_{X\sim\hat{\mu}}\hat{h}_{\varepsilon}(X,\mathbf{v})$, where $\mathbf{v}=(v(y_{1}),\ldots,v(y_{N}))^{\top}$ is the evaluation of $v$ on the test set.

## Appendix C Properties of SoftPlus

Let $F(x)=\log(1+e^{f(x)})$, then

$$ $\displaystyle\nabla F(x)$ $\displaystyle=\sigma(f(x))\nabla f(x),$ (26) $\displaystyle\nabla^{2}F(x)$ $\displaystyle=\sigma(f(x))\nabla^{2}f(x)+\sigma(f(x))\bigl(1-\sigma(f(x))\bigr)\nabla f(x)\nabla f(x)^{\top}.$ (27) $$

Suppose $f(x)$ is $L$-smooth (possibly non-convex). Let us derive smoothness constant of $F$. We will use the following

###### Lemma C.1 .

Consider function $f_{a}(x)=\sigma(x)+2\sigma^{\prime}(x)(x-a),\,x\geq a$ with parameter $a\leq 0$. It holds $f_{a}(x)\leq 2-\frac{a}{2}$.

###### Proof.

By the properties of the sigmoid function $\sigma(x)$, $\sigma^{\prime}(x)\leq\frac{1}{4}$ and $\sigma(x)\leq 1$. Therefore, $f_{a}(x)\leq 1+\frac{x-a}{2}$. If $x\leq 2$, the result follows. Let us now show that the derivative

$$ $\frac{d}{dx}f_{a}(x)=\sigma^{\prime}(x)[3+2(1-2\sigma(x))(x-a)]$ $$

is negative if $x>2$. Indeed, due to monotonicity of the sigmoid function $\sigma(x)$,

$$ $\sigma(x)>\sigma(2)>0.88\Rightarrow 2(1-2\sigma(x))<-\frac{3}{2}.$ $$

Moreover, $x-a>2$, so $3+2(1-2\sigma(x))(x-a)<0$ and $\frac{d}{dx}f_{a}(x)<0$. Therefore, if $x>2$, then $f_{a}(x)<f_{a}(2)\leq 2-\frac{a}{2}$.
∎

###### Proposition C.2 .

Let $f\in C^{1}(\mathbb{R}^{d})$ be $L$-smooth and bounded from below by $f_{*}\in\mathbb{R}$,
then $F(x)=\log(1+e^{f(x)})$ is smooth with parameter

$$ $\begin{cases}\frac{4}{3}L&\text{if }f_{*}\geq 0,\\ \left(\frac{4}{3}-\frac{f_{*}}{2}\right)L&\text{if }f_{*}<0.\end{cases}$ (28) $$

###### Proof.

W.l.o.g., we can assume that $f\in C^{2}$.
From ([27](https://arxiv.org/html/2509.24894v1#A3.E27)) and Lemma [C.1](https://arxiv.org/html/2509.24894v1#A3.Thmtheorem1) we get

$$ $\displaystyle\lVert\nabla^{2}F(x)\rVert$ $\displaystyle\leq\sigma(f(x))\lVert\nabla^{2}f(x)\rVert+\sigma^{\prime}(f(x))\lVert\nabla f(x)\rVert^{2}$ $\displaystyle\leq L\sigma(f(x))+2L\sigma^{\prime}(f(x))(f(x)-f_{*})$ $\displaystyle=L\bigl(\sigma(f(x))+2\sigma^{\prime}(f(x))f(x)\bigr)-2L\sigma^{\prime}(f(x))f_{*}.$ $$

Analyzing the function $h(t)\coloneqq\bigl(\sigma(t)+2t\sigma^{\prime}(t)\bigr)$, one can show that $\max_{t}h(t)<\frac{4}{3}$.
Thus, in the case $f_{*}\geq 0$, using the fact that $\sigma^{\prime}(t)>0$ we obtain

$$ $\lVert\nabla^{2}F(x)\rVert\leq Lh(f(x))\leq\frac{4}{3}L.$ $$

Now, consider the case $f_{*}<0$. Since $\sigma^{\prime}(t)=\sigma(t)(1-\sigma(t))\leq\frac{1}{4}$,

$$ $\lVert\nabla^{2}F(x)\rVert\leq Lh(f(x))-2L\sigma^{\prime}(f(x))f_{*}\leq\frac{4}{3}L-\frac{L}{2}f_{*}.$ $$

The claim follows.

∎

###### Remark C.3 .

The factor $\frac{1}{2}$ in front of $-f_{*}$ in ([28](https://arxiv.org/html/2509.24894v1#A3.E28)) can’t be improved. Indeed, consider $f(x)=\frac{1}{2}(x-a)^{2}-\frac{1}{2}a^{2}$ with $f_{*}=-\frac{1}{2}a^{2}$. The second derivative of $F(x)=\log(1+e^{f(x)})$ is

$$ $\displaystyle F^{\prime\prime}(x)$ $\displaystyle=\sigma(f(x))+\sigma(f(x))\left(1-\sigma(f(x))\right)(x-a)^{2},$ $\displaystyle F^{\prime\prime}(0)$ $\displaystyle=\sigma(0)+\sigma(0)\left(1-\sigma(0)\right)a^{2}=\frac{1}{2}+\frac{a^{2}}{4}=\frac{1}{2}-\frac{f_{*}}{2}.$ $$

###### Proposition C.4 .

If $f$ is convex, then $F(x)=\log(1+e^{f(x)})$ is also convex.

###### Proof.

Trivially follows from ([27](https://arxiv.org/html/2509.24894v1#A3.E27)).
∎

###### Lemma C.1 .

###### Proof.

###### Proposition C.2 .

###### Proof.

###### Remark C.3 .

###### Proposition C.4 .

###### Proof.

## Appendix D LLM Usage Disclosure

In the preparation of this manuscript, large language models (LLMs) were used to improve the readability.
All substantive contributions are solely by the authors.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 4
presentation: 3
questions: 1. Could the authors compare with existing unbiased estimators like REBAR or those based on Russian Roulette truncation?
2. Any other baselines Authors would like to add?
3. Why fρ(t) is strongly convex (lines 218, 219)?
soundness: 3
strengths: 1. A key strength is that the bias introduced by the approximation is not arbitrary. The paper provides clear theoretical bounds showing the error is of order O(ρ), allowing practitioners to explicitly trade off accuracy for computational stability (Proposition 2.4 and Corollary 2.5).
2. The approximation preserves essential properties like convexity and smoothness, which are crucial for guaranteeing the performance of gradient-based optimizers.
3. The resulting gradient estimator is simple to implement and does not require complex machinery like control variates or sophisticated sampling schemes, making it an attractive and easy-to-adopt tool.
summary: The paper introduces a novel and mathematically sound approximation for the LogSumExp function, termed "safe KL divergence," which is designed for efficient and stable stochastic optimization. The method preserves desirable properties like convexity, and its approximation error is explicitly controlled by a tunable parameter. The empirical results in continuous optimal transport and distributionally robust optimization are strong, demonstrating clear practical benefits in terms of convergence speed and numerical stability over the **chosen** baselines. However, the paper's central claim regarding its novelty is significantly overstated due to a major omission of relevant literature on existing unbiased estimators for the log-partition function, which directly contradicts the assertion that "no cheap unbiased stochastic gradient have been proposed."
weaknesses: 1. The paper does not do justice to prior work on approximating the log partition function. As of now, the prior work is largely incomplete. Few examples:
https://arxiv.org/pdf/2004.00353
https://arxiv.org/pdf/1703.05160
https://arxiv.org/pdf/1306.4032
https://projecteuclid.org/journals/statistical-science/volume-30/issue-4/On-Russian-Roulette-Estimates-for-Bayesian-Inference-with-Doubly-Intractable/10.1214/15-STS523.full
https://arxiv.org/pdf/1703.07370
https://arxiv.org/abs/1901.10517
2. Please add a proof for Lemma 2.7 into the Appendix.
3. Empirical baselines are weak (see item 1 above).

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
questions: - Could you perform an ablation study of the parameters $\epsilon$ in section 3.1 and parameters $\lambda$ and $\beta$ in section 3.2 and 3.3, respectively?
- Could you test your approximation with some of the more recent EOT methods?

*Typos:*

- in Eq. (2): ‘sup’ -> ‘sup_{\nu}’
- in line 108: minimum -> maximum?
- in Eq. (2), (5), (11)-(13) and almost everywhere in the text: \int -> \int_{X}
- how do you calculate the ‘optimality gap’ visualized in Fig. 2?

**References.** 

Seguy, V., Damodaran, B. B., Flamary, R., Courty, N., Rolet, A., & Blondel, M. (2018, April). Large-Scale Optimal Transport and Mapping Estimation. In ICLR 2018-International Conference on Learning Representations (pp. 1-15).

Gushchin, N., Kolesov, A., Korotin, A., Vetrov, D. P., & Burnaev, E. (2023). Entropic neural optimal transport via diffusion processes. Advances in Neural Information Processing Systems, 36, 75517-75544.

Daniels, M., Maunu, T. & Hand, P. Score-based generative neural networks for large-scale optimal transport. Advances in neural information processing systems, 34:12955–12965, 2021.

Korotin, A., Gushchin, N., & Burnaev, E. Light Schrödinger Bridge (2024). In The Twelfth International Conference on Learning Representations.
soundness: 2
strengths: The idea of the paper is rather interesting and supported by the conducted theoretical analysis. The authors show several use-cases of the proposed approximation.
summary: In this paper,  the authors propose a new strategy of approximating the LogSumExp function which is based on the usage of a modified KL divergence (safe KL) in the dual formulation. This approximation is applied in relevant tasks, including the computation of entropy-regularized optimal transport (EOT) problem. It is shown that the approximation helps to alleviate the overflow issue in these applications.
weaknesses: My major concerns correspond to the overall contribution of the proposed approximation of LogSumExp and its practical significance. While the derived theoretical results show some nice properties of the derived approximation, its experimental validation has some white spots. To begin with, in section 3.1, the authors tested their approximation by introducing it in the method for entropic optimal transport (Genevay et al., 2016) which is rather old while being classic. For example, this method does not allow for the direct computation of EOT plans which is a more important problem than the computation of EOT cost itself. In recent years, many other methods for computing EOT plans between the continuous distributions have been derived, see, e.g., (Seguy et al., 2018; Daniels et al., 2021, Gushchin et al., 2023, Korotin et al., 2023). The paper would benefit from the testing of your approximation in these more recent approaches.

Besides, while in remark 3.1, the authors note that for small value of parameter $\epsilon$  the method of (Genevay et al., 2016) exhibits the overflow issue, the conducted experiment reveals the behaviour of this approach and its modified version (with the proposed regularization) only for one parameter $\epsilon$. Here I kindly suggest the authors perform an ablation study of this parameter (for bigger and smaller values) to show how it affects the performance of the method with and without regularization. Similarly, I kindly suggest the authors perform an ablation study on the parameters $\lambda$ and $\beta$ appearing in denominator of exponent in the distributionally robust optimization problems – balanced (16) and unbalanced (20) –  in experiments from sections 3.2, 3.3.

Moreover, there are several aspects of the paper which could be written more clearly. First, In equation (2), the authors introduce an alternative form of the functional $F(\phi,\mu)$ using the Gibbs variational principle. I kindly suggest including reference to one of the papers where this principle is introduced. Many formulas should be written in a more rigorous manner, i.e, integrals should include the spaces over which they are taken and sup/inf should be taken over some optimization variable (see Eq. (2)). The typos are summarized below. 

**Overall**, I have doubts regarding the contribution of the approximation of the LogSumExp function derived in the paper. While the theoretical results support this approximation, its usage should offer considerable benefits over the usage of the initial function. To prove it, some additional experimental evaluations with the recent entropy-regularized OT methods and ablation study on the parameters should be considered.

---

confidence: 2
contribution: 4
initial_rating: 8
presentation: 2
questions: How to pick the parameter $\rho$ controlling the accuracy of the approximation?
soundness: 3
strengths: The paper studies a fundamental question: the optimization of LogSumExp. The contribution is clear in both theory and the provided empirical results, which demonstrate improvements across various tasks. As far as I can follow, I did not find inconsistencies or technical flaws, and thus the contribution appears significant.
summary: The paper considers the minimization of the functional $f \mapsto \log \int e^f d\mu$, whose specific forms appear in many applications, such as computing softmax probabilities and maximum-likelihood estimation for exponential families. The functional  can be viewed as LogSumExp, whose computation is known to suffer from numerical-stability issues, which the proposed approximation can mitigate substantially as shown emprically in the paper.
weaknesses: Quantitive summary of the results of the experimental section is missing, especially in Section 3.1. Is it Figure 2? This figure is not referenced anywhere in the text. Due to the unclear presentation, it is difficult to obtain an accurate view of how well the proposed approach performs overall compared to the baseline.

Writing is not always smooth. Some examples:

“First, the decision variable φ(or a parameter θdefining φ) often has large or infinite dimension... The first challenge is usually addressed by the use of first-order methods, especially stochastic gradient descent (SGD), with cheap iterations.” What this means? Is a reader supposed to follow? 

Lines 55 and 57: incosistent usage of dashes

---

confidence: 5
contribution: 2
initial_rating: 2
presentation: 3
questions: (1) The manuscript claims novelty for the “rescaled SoftPlus approximation” of the LogSumExp (LSE) function. Please clarify precisely how your method differs in formulation, analysis, or practical behavior from existing works, for instance listed above.
(2) Can you include additional experiments and comparison with another optimization methods of LSE, such as piecewise-linear approximations, Newton-type solvers, or surrogate bounds, and provide wall-clock/iteration numbers for a range of dimensions and scaling parameters?
(3) Could you provide a sensitivity/ablation study of \rho in both training stage and final objective quality, along with guidance on how to choose \rho in practice?
(4) How does your method scale to higher feature dimensions (e.g., dim > 10, or dim > 100), larger term-counts in LogSumExp, and different \rho values?
soundness: 2
strengths: (1)The theoretical reformulation through an alternative f-divergence is clear and well motivated.
(2)The approach preserves convexity and smoothness, with rigorous proofs of approximation bounds.
(3)Applications to EOT and DRO are reasonable and illustrate numerical stability over the standard LogSumExp formulation with larger stepping.
summary: This paper introduces a stochastic optimization framework for objectives involving the LogSumExp function (or log-partition functional), proposing a new “Safe KL” divergence formulation that yields a rescaled SoftPlus approximation of LogSumExp. The authors claim this relaxation preserves convexity, smoothness, and enables efficient stochastic gradient methods applicable to entropy-regularized optimal transport (EOT) and distributionally robust optimization (DRO). The theoretical analysis shows convergence properties and connections to Conditional Value-at-Risk (CVaR). Empirical demonstrations include low-dimensional EOT and DRO on the MNIST and California Housing datasets.

The theoretical formulation is sound and mathematically elegant, but the novelty and experimental validation are limited. The approximation technique and results overlap with prior SoftPlus-based LogSumExp relaxations, and the numerical demonstrations are too small-scale to convincingly support the claimed generality.
weaknesses: (1)The proposed rescaled SoftPlus approximation is not entirely novel. Similar ideas have appeared—for instance, in “Safety Alignment Should Be Made More Than Just a Few Tokens Deep” (arXiv:2406.05946), where the authors also analyze the effect of the scaling parameter \rho. 
(2) The paper lacks comparison with other established techniques for optimizing or bounding LogSumExp, such as those discussed in 
- Convex Maximization via Adjustable Robust Optimization (https://optimization-online.org/wp-content/uploads/2020/07/7881.pdf)
- LSEMINK: A Modified Newton-Krylov Method for Log-Sum-Exp Minimization” (Kan et al. 2023) deals with optimization of LogSumExp via Newton-Krylov methods (https://arxiv.org/abs/2307.04871)
- Accurate Computation of the Log-Sum-Exp and Softmax Functions (https://arxiv.org/abs/1909.03469)
- The broader review “What is the log-sum-exp function?” (Higham, 2021 https://nhigham.com/2021/01/05/what-is-the-log-sum-exp-function/)
(3) The experimental design is too simple for a top ML conference. For instance, the “semi-continuous” EOT test uses 1D feature inputs, which is insufficient to demonstrate scalability or real-world relevance. Stronger baselines (e.g., multi-dimensional or high-resolution OT problems) should be included.
(4) The empirical section does not explore how the rescaling parameter \rho affects performance and accuracy, especially in the MNIST experiments (Section 3.3). Sensitivity or ablation tests are needed to justify the parameter’s influence on convergence and bias.


Decision: Reject