技术摘要：改进的 LogSumExp 随机优化方法

1. 问题陈述

LogSumExp 函数定义为 $F(\phi; \mu) = \log \int e^{\phi(x)} d\mu(x)$，它是众多机器学习优化问题中的核心，包括熵正则化最优传输（OT）、分布鲁棒优化（DRO）、多分类问题以及变分推理。

优化该函数的主要挑战在于当指数项的数量很大或测度 $\mu$ 是连续时。计算精确梯度需要对所有项或积分进行微分，这在计算上是难以承受的。标准的随机梯度下降（SGD）方法通过采样一部分项来近似，通常会产生有偏的梯度估计器。此外，现有的方法在正则化参数较小时常会出现数值不稳定（溢出）问题，因为指数函数增长极快，会导致浮点异常。

2. 方法论

作者提出了一种针对 LogSumExp 函数的新型凸性与平滑性保持近似方法。该方法通过引入一种被称为 Safe KL 散度 的新 $f$-散度来替换对偶中的标准 Kullback-Leibler (KL) 散度，从而修改了问题的变分表述。

2.1 Safe KL 散度

该方法引入了一个参数 $\rho \in (0, 1)$，并定义了一个生成函数 $f_\rho(t)$：
$$f_\rho(t) = \begin{cases} t \log t + \frac{1-\rho t}{\rho} \log(1 - \rho t) & \text{若 } 0 \le t \le \frac{1}{\rho} \\ +\infty & \text{否则} \end{cases}$$
该生成函数诱导了 Safe KL 散度 $D_\rho(\nu, \mu)$。与标准 KL 散度不同，$f_\rho$ 对密度比 $d\nu/d\mu$ 施加了硬约束，防止其超过 $1/\rho$。这一特性限制了 DRO 和 OT 中的重加权，提供了针对离群值和数值溢出的稳定性。

2.2 变分近似

利用凸对偶性，作者推导出了近似对数配分函数 $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)$$
该公式使用缩放后的 SoftPlus 函数 $\frac{1}{\rho}\log(1 + \rho e^{\phi(x) - \alpha})$ 取代了原始的指数 $e^{\phi(x)}$。

2.3 优化性质

• 无偏梯度： 在代理目标 $G_\rho(\theta, \alpha)$ 上通过采样批次计算的随机梯度是近似目标的无偏估计量。这消除了标准采样 LogSumExp 近似中固有的依赖于批次大小的偏差。
• 凸性与平滑性： 如果底层损失函数是凸的，则代理目标在参数上保持联合凸性。其共轭函数 $f_\rho^*$ 被证明是 $1/\rho$-平滑的，确保了随机次梯度的二阶矩受控。
• 与 CVaR 的联系： 作者展示了 $F_\rho$ 与条件风险价值（CVaR）泛函密切相关，提供了一种连接 LogSumExp 最小化与 CVaR 优化的平滑近似。

3. 主要贡献

本文概述了四个主要贡献：

1. 通用近似： 一种计算高效的 LogSumExp 函数的松弛方法，它保持了凸性和平滑性，将问题转化为适用于标准 SGD 的期望最小化问题。近似误差由单个可调标量参数 $\rho$ 控制。
2. 理论基础： 推导了作为该近似对偶部分的 Safe KL 散度，确立了其理论属性及其与现有 $f$-散度的关系。
3. 实证验证： 在连续熵正则化 OT 和各种 DRO 形式（KL 散度及非平衡 OT）上进行了评估。该方法在竞争性能上达到了应用特定基准线的高度，同时避免了数值溢出问题。
4. 理论见解： 阐明了所提近似、CVaR 以及风险泛函平滑近似之间的联系。

4. 实验结果

作者在三个特定任务上验证了该方法：

• 连续熵正则化 OT： 使用 MNIST 和 EMNIST 数据集，将所提方法（由神经网络参数化）与 LSOT 基准进行了比较。
  • 结果： 基准方法（LSOT）在正则化参数 ($\epsilon$) 较小时会出现数值溢出，需要极小的学习率或提前终止。所提方法在所有测试的 $\epsilon$ 值下均保持稳定且高效，能够以较大的学习率实现更快的收敛。
• KL-散度 DRO： 应用于加州住房数据集的回归任务。
  • 结果： 在小批次规模下，所提估计器优于基准方法（后者依赖大批次来减少偏差）。中间参数 $\rho = 10^{-3}$ 在近似精度与优化平滑性之间提供了最佳权衡。
• 非平衡 OT DRO： 应用于带有标签噪声的 MNIST 分类任务。
  • 结果： 基准方法（优化指数之和）由于溢出经常发生发散或需要极小的步长。所提方法实现了更低的客观值，并表现出即使在高度标签噪声（补充实验中高达 85%）下的鲁棒性。

5. 意义与局限性

意义：
本文声称所提方法是处理涉及 LogSumExp 的大规模优化问题的通用工具。其主要优势在于能够处理弱正则化和大尺度数据，而不会出现困扰标准指数形式的数值不稳定性（溢出）。通过为代理目标提供无偏随机梯度，它规避了需要大规模批次来降低偏差的需求。

局限性：
作者谦逊地承认了以下几点局限性：

• 结构属性： 由于 Safe KL 修改了原始的 KL 散度，在特定应用中，原本基于 KL 的公式所具有的特定结构属性可能需要单独分析。
• 参数调优： 参数 $\rho$ 控制着偏差与稳定性的权衡，目前仍需要手动调优；将自适应选择程序作为未来的研究方向。
• 范围： 实验涵盖了一定的下游任务，但仍需进一步研究以将 Safe KL 形式扩展到其他 LogSumExp 与 KL 对偶处于核心地位的场景中。