Particle-based Generalised Stochastic Optimisation

技术摘要：基于粒子的广义随机优化

问题陈述
本文旨在解决优化损失函数 $\ell: \Theta \to \mathbb{R}$ 的挑战，其中梯度是一个关于参数依赖分布 $\pi_\theta$ 的难解期望。具体而言，梯度形式如下：
$$\nabla_\theta \ell(\theta) = \mathbb{E}_{\pi_\theta}[F(\theta, X)] = \int_X F(\theta, x)\pi_\theta(dx)$$
这种结构广泛存在于能量模型（EBMs）的训练、隐变量模型、最大边际似然估计（MMLE）以及生成模型的微调中。分布 $\pi_\theta$ 通常具有不可解的归一化常数 $C_\theta$ 的未归一化密度 $e^{-E(\theta, x)}/C_\theta$。

标准方法通常采用嵌套优化-采样方案，即在每一个外部梯度更新步骤中，通过一个内部的马尔可夫链蒙特卡洛（MCMC）循环来近似 $\pi_\theta$ 的样本。这些方法计算成本高昂，且存在与 MCMC 偏差、混合时间（mixing times）以及步长敏感性相关的收敛问题。虽然近期的文献引入了相互作用粒子方法（例如粒子梯度下降、动量粒子梯度下降）以避免嵌套循环，但这些方法通常是针对特定的优化器或扩散过程定制的，缺乏一个能够容纳现代复杂优化器（如 Adam 或基于动量的算法）和高阶采样动力学的统一理论框架。

方法论
作者提出了一个基于平均场动力学（mean-field dynamics）及其相互作用粒子近似的通用连续时间框架。该方法的核心是一个耦合的随机微分方程（SDE）系统，涉及优化器状态 $\vartheta_t$ 和采样器状态 $Z_t$：

1. 优化器动力学：
   $$d\vartheta_t = \Psi_t(\vartheta_t, G_{\nu_t}(\vartheta_t))dt$$
   这里 $\vartheta_t \in \Theta \times M$ 包含参数和辅助变量（例如动量）。$\Psi_t$ 是一个通用的更新函数，允许各种优化器（梯度下降、重球法、类 Adam 算法）。$G_{\nu_t}$ 是由粒子分布导出的经验梯度估计。

2. 采样器动力学：
   $$dZ_t = b_{\Pi_\theta(\vartheta_t)}(Z_t)dt + \sigma_{\Pi_\theta(\vartheta_t)}(Z_t)dB_t$$
   $Z_t \in X \times V$ 代表粒子的状态。漂移项 $b$ 和扩散项 $\sigma$ 被设计为：在参数 $\theta$ 固定时，其 $X$ 分量的边缘分布收敛于 $\pi_\theta$。该框架允许通过选择适当的 $b$ 和 $\sigma$（例如通过涉及对称和反对称矩阵的广义朗之万构造）来实现过阻尼、欠阻尼和高阶朗之万动力学。

3. 粒子近似：
   该平均场系统由 $N$ 个受独立布朗运动驱动的相互作用粒子 $\{Z^i_t\}_{i=1}^N$ 来近似。粒子通过共享的优化器状态 $\vartheta^N_t$ 进行相互作用，该状态根据从粒子计算出的梯度的经验平均值进行更新。

核心贡献

1. 统一框架： 本文建立了一个通用的基于扩散的框架，将现有的粒子方法（如 PGD、MPGD 和隐式扩散）统一为一个整体。至关重要的是，它允许在单一理论结构内集成任意现代优化器和复杂的采样动力学。
2. 理论收敛保证：
   • 指数收缩（Exponential Contraction）： 在满足良定性（well-posedness）和联合强单调性条件（一种针对耦合优化器-采样器系统的扭曲收缩性条件）的假设下，作者证明了平均场系统指数级收敛至平稳分布。
   • 平稳最优性： 在额外的保守性和联合强凸结构下，证明了平稳参数是损失函数的唯一全局极小值。
   • 有限粒子误差界： 文中推导了粒子近似的非渐近误差界。具体而言，它展示了 $O(N^{-1/2})$ 的相互作用误差，并提供了一个参数误差界，该误差随时间呈指数级衰减，前提是粒子数量 $N$ 足够大。
3. 新颖算法设计： 该框架被实例化为两种新算法：
   • HOL-MYPGD： 一种用于 MMLE 的高阶朗之万动力学算法，利用近端技术处理非微分先验（如全变分 TV）。
   • M-MFPO： 一种用于训练 EBM 的基于动量的粒子优化器，将 MPGD 推广到更广泛的具有欠阻尼采样动力学的问题类。

实验结果
作者在两个不同的任务上评估了所提方法：

1. 贝叶斯图像去模糊（MMLE）： 在涉及各向异性全变分先验的设置中，提出的 HOL-MYPGD（高阶朗之万）与基准方法 Moreau-Yosida 粒子梯度下降（MYPGD）进行了对比。结果表明，高阶方案实现了更低的均方误差（MSE）和更高的结构相似性（SSIM），产生了更清晰的图像重建结果。
2. 能量模型（EBM）训练： 将 M-MFPO 算法应用于合成 2D 数据集（环形、珠状、晶格）并与持久对比散度（PCD）进行测试。与 PCD 相比，该动量方法在费舍尔散度（Fisher Divergence）的对数尺度上表现出更快的收敛速度，同时保持了相当的生成样本视觉质量。

意义与主张
本文声称为具有难解梯度的损失函数提供了一种“原则性的、通用的”随机优化方法。其主要意义在于：

• 泛化性： 超越了特定问题或特定优化器的构建，转向一个灵活的框架，可以整合最先进的优化器（如 Adam 或重球法）和高级采样动力学（高阶、非可逆）。
• 理论严谨性： 为这类算法提供了严谨的收敛证明，特别是在联合收缩性假设下建立了指数收敛和有限样本误差界。
• 实用价值： 通过在特定选择（如用于 MMLE 的高阶动力学，或用于 EBM 的动量）下实例化该框架，证明了其性能优于现有的专门化方法。

作者承认了局限性，指出其理论保证依赖于全局收缩性和一致矩假设，这对于大型神经网络中高度非凸的目标函数可能过于严格。此外，目前的结论适用于连续时间系统，建立针对特定离散化方案的有限步误差界仍是未来的研究课题。