The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity

技术摘要：超越凸性的受控次梯度未调整朗之万算法 (The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity)

1. 问题陈述

本文旨在解决在目标分布 $\pi_\beta(dx) \propto e^{-\beta u(x)} dx$ 下进行采样，并解决相关的优化问题 $\theta^\star \in \arg\min u(\theta)$ 的挑战。在这些场景中，势函数 $u: \mathbb{R}^d \to \mathbb{R}$ 同时具备三种困难属性：

1. 非光滑性 (Non-smoothness)： 势函数并非处处可微（例如，由于 ReLU 激活或 $\ell_1$ 惩罚项），因此需要使用次梯度 $h \in \partial u$。
2. 超线性梯度增长 (Superlinear Gradient Growth)： 次梯度的模长呈多项式增长（$|h(x)| \leq m + K|x|^p$，其中 $p \geq 1$），这会导致标准的 Euler–Maruyama 朗之万扩散离散化在数值上变得不稳定（矩在有限时间内发散）。
3. 非凸性 (Non-convexity)： 势函数并非全局凸，尽管它满足特定的结构条件（半凸性以及无穷远处的强凸性）。

现有文献通常孤立地处理这些问题（例如，通过平滑处理非光滑性、通过 taming 处理超线性增长，或通过泛函不等式处理非凸性）。本文认为，针对这三种性质同时存在的非渐近保证在以往的研究中是缺失的。

2. 方法论

作者提出了次梯度受控未调整朗之万算法 (SG-TULA)。这是对过阻尼朗之万扩散的一种显式 Euler–Maruyama 离散化，它直接作用于可测的次梯度选择 $h \in \partial u$，而不依赖于计算昂贵的平滑程序（如 Moreau–Yosida 正则化）。

算法实现：
迭代公式定义为：
$$\theta^\lambda_{n+1} = \theta^\lambda_n - \lambda h^\lambda(\theta^\lambda_n) + \sqrt{2\lambda \beta^{-1}} \xi_{n+1}$$
其中 $\xi_{n+1}$ 是独立同分布的标准高斯向量。其核心创新在于控制函数 (taming function) $h^\lambda$，它能稳定漂移项：
$$h^\lambda(x) = \frac{h(x) - \mu x}{1 + \lambda^{1/2}|h(x) - \mu x|} + \mu x$$
该机制根据次梯度的量级连续地重新缩放次梯度，从而防止爆炸，而无需设定固定阈值（不同于梯度裁剪）。论文分析了统一形式变体 (SG-TULA$_u$) 和坐标形式变体 (SG-TULA$_c$)。

理论框架：
其分析依赖于对 $u$ 的三个主要假设：

• A1 (半凸性)： $u + \frac{L}{2}|\cdot|^2$ 是凸的。这允许在测度为零的集合上存在梯度不连续性（例如 ReLU），同时提供单侧 Lipschitz 界。
• A2 (无穷远处的强凸性)： 在紧集之外，势函数是强凸的，确保了扩散过程的耗散性和收缩性。
• A3 (多项式增长)： 次梯度满足 $|h(x)| \leq m + K|x|^p$。

证明策略包括：

1. 建立不变测度的存在性与唯一性。
2. 在这些假设下，证明目标测度 $\pi_\beta$ 满足对数索博列夫不等式 (LSI)。
3. 推导连续 SDE 向 $\pi_\beta$ 的指数收缩率。
4. 通过控制机制建立离散算法的一致矩界。
5. 将离散化误差估计与收缩率结合，以界定 Wasserstein-2 距离。

3. 核心贡献与结果

非渐近采样界限：
论文推导了在 Wasserstein-2 距离 ($W_2$) 下显式的非渐近收敛界限。对于步长 $\lambda < 1/(4\mu)$ 和 $n$ 次迭代：
$$W_2(\pi_\beta, \mathcal{L}(\theta^\lambda_n)) \leq 2 C_E \lambda^{1/4} + e C_D e^{-n\lambda/T_0}$$

• 速率： 离散化误差以 $O(\lambda^{1/4})$ 的速率收敛。
• 常数： 所有常数 ($C_E, C_D, T_0$) 都根据维度 $d$ 和逆温度 $\beta$ 进行了显式追踪。维度依赖关系是多项式的 ($O((d/\beta)^{p/2})$)。
• 对比： 这改进了在类似的非光滑、超线性及非凸条件下，基于次梯度的朗之万算法的现有最优速率 $O(\lambda^{1/8})$（具体改进了 Johnston 等人 [2025a] 的结果）。它达到了平滑非凸设置下的 $O(\lambda^{1/4})$ 速率，并将此扩展到了非光滑、超线性的情形。

优化保证：
采样界限被转移到优化问题中，用以界定超额风险 (excess risk)：
$$\mathbb{E}[u(\theta^\lambda_n)] - u(\theta^\star) \leq C_{T1} W_2(\pi_\beta, \mathcal{L}(\theta^\lambda_n)) + C_{T2}$$
其中 $C_{T2}$ 代表“温度间隙”（由有限 $\beta$ 导致的偏差）。达到 $\epsilon$ 精度所需的复杂度为 $O(\epsilon^{-4} \log(1/\epsilon))$ 次迭代，步长 $\lambda = O(\epsilon^4)$。

实验验证：
作者在两个特定问题上验证了其假设和算法：

1. LLM 预训练 (Nanochat)： 他们应用了一种增强的坐标形式变体 (SG-TheoPouLA) 来预训练一个 GPT-2 系列模型。论文验证了该势函数满足所需的曲率界限（具体为 $\nabla^2 u(\theta) \succeq -c_\star(1+|\theta|)^{12L+2} I$）。实验表明，在深度为 12 时，SG-TheoPouLA 在特定配置下，在测试的优化器中取得了最低的验证比特/字节 (bits-per-byte) 和最高的 CORE 分数，表现出与经过微调的 AdamW 和 Muon 相当的竞争力。在深度为 24 时，Muon 在验证比特/字节方面领先，而 SG-TheoPouLA 紧随其后，并在 CORE 指标上与 AdamW 保持竞争态势，且均在评估误差范围内。
2. 稀疏相位检索 (Sparse Phase Retrieval)： 论文展示了稀疏相位检索问题（四次能量 + $\ell_1$ 惩罚）满足所有这三种具有挑战性的条件（$p=3$ 的超线性增长、非凸性以及坐标超平面上的梯度不连续性）。

4. 重要性与主张

本文声称，提供了第一个针对在不借助平滑近似的情况下，能同时处理非光滑性、超线性梯度增长和非凸性的基于朗之万算法的非渐近收敛保证。

• 理论影响： 它填补了文献中的空白，即以往的工作要么需要光滑性（以使用标准 taming），要么需要凸性（以确保稳定性）。$O(\lambda^{1/4})$ 的速率被视为在此特定情形下对现有次梯度方法的改进。
• 实践影响： 研究表明，taming 技术在计算上是非常高效的（开销显著小于矩阵乘法），并且在深度学习语境下行之有效。在 LLM 预训练上的实验结果表明，该理论框架适用于现代深度架构，其中梯度相对于权重呈多项式增长且具有非凸性。
• 局限性： 论文指出，常数中对逆温度 $\beta$ 的指数依赖很大程度上是证明策略（通过采样路径传递优化保证）产生的产物。作者建议，在满足 Polyak–Łojasiewicz 条件的情况下，这种依赖关系可以变为多项式。此外，注意到 $O(\epsilon^{-4})$ 的复杂度是由于特定的漂移正则性假设所导致的，若具备更高阶的光滑性，理论上可以获得更好的速率（例如 $O(\lambda^{1/2})$）。

总之，这项工作为在深度学习的“野外”条件（非光滑、非凸、超线性）下使用次梯度朗之万动力学建立了严谨的理论基础，并提供了实验证据，证明所提出的控制方案是标准优化器（如 AdamW 和 Muon）的一个可行且具竞争力的替代方案。