The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity

近年来，深度学习模型的复杂性不断提升，导致优化潜在函数呈现非凸、非光滑和超线性增长的特性。传统的Langevin算法在凸光滑场景中表现优异，但在非凸、非光滑和高维环境下，稳定性和收敛性受到严重制约。早期工作如Dalalyan [2017]和Durmus-Moulines [2019]提出的平滑和强凸假设，限制了算法在实际深度模型中的应用。近年来，研究者开始关注在非凸、非光滑、超线性增长条件下的采样算法，尝试引入taming技术、Proximal方法和高阶正则化，以应对梯度爆炸和不稳定问题。相关工作如Brosse et al. [2019]、Lovas et al. [2023]等，虽在特定条件下取得一定进展，但仍未解决多重复杂性共存的场景。本文在此背景下，提出了直接在次梯度上操作的SG-TULA算法，突破了现有限制，为高维非凸采样提供了理论基础和实践方案。

核心问题在于如何在潜在函数具有非光滑、非凸且超线性增长的条件下，保证采样算法的稳定性和收敛速度。传统方法依赖平滑或强凸假设，难以应对深度学习中的复杂潜能。具体挑战包括梯度爆炸、非连续性以及高维空间中的收敛速率难以保证。现有的子梯度方法多在单一条件下有效，面对多重复杂性时表现不佳。如何设计一种既能直接操作次梯度，又能确保算法稳定、收敛的方案，成为亟待解决的问题。这不仅关系到采样的理论基础，也直接影响深度模型预训练和优化的实际效果。

本研究的创新点主要体现在以下几个方面：

1) 直接在次梯度上操作，避免平滑带来的计算成本和偏差，保持算法的原始特性；

2) 引入taming技术，有效控制超线性梯度爆炸，确保算法稳定性；

3) 提出非渐近的Wasserstein-2距离收敛界限，明确追踪常数与维度、逆温度的关系，提升理论深度；

4) 结合具体应用验证，展示在大规模深度模型预训练中的优越性能。这些创新突破了传统平滑和强凸假设的限制，为非凸、非光滑、超线性增长潜在函数的采样和优化提供了新思路。

• �� 构建目标潜在函数u，假设其满足半凸、强凸（外部）和多项式增长条件。
• �� 设计次梯度操作h∈∂u，结合taming技术，将h进行尺度调节，避免梯度爆炸。
• �� 采用显式欧拉-马鲁雅马离散方案，更新公式为θ_{n+1} = θ_n - λ h_λ(θ_n) + √(2λ/β) ξ_{n+1}，其中h_λ为taming后的次梯度，ξ_{n+1}为标准高斯噪声。
• �� 通过引入非渐近的Wasserstein-2距离收敛分析，追踪常数与维度、逆温度的关系，确保在高维空间中的稳定性。
• �� 结合具体潜在函数的验证，推导出收敛速率为O(λ^{1/4})，并分析算法在优化中的超额风险。

实验设计包括在GPT-2潜能的正则化预训练任务中验证算法效果，使用大规模文本数据集，比较SG-TULA与AdamW、Muon的性能。调节步长λ和逆温度β，观察采样误差和优化超额风险。还在稀疏相位检索任务中测试算法的鲁棒性，评估其在非凸、非光滑场景中的表现。通过多组超参数和不同模型深度，验证理论推导的收敛速率和常数依赖关系。实验结果显示，SG-TULA在高维复杂潜能中保持稳定，收敛速度优越，优于现有子梯度方法。

在潜在函数满足半凸、强凸（外部）和多项式增长条件下，算法实现了Wasserstein-2距离的收敛速率为O(λ^{1/4})，常数明确，维度依赖多项式，逆温度影响指数级。实验证明，在GPT-2潜能预训练中，SG-TULA以较小的步长实现快速收敛，优于finetuned AdamW和Muon，验证了理论预期。稀疏相位检索中，算法表现出良好的鲁棒性和稳定性，适应不同潜能结构。超额风险分析显示，算法在实际训练中具有较强的泛化能力，特别是在高维非凸场景中表现出优越性能。

该算法适用于深度学习中的大规模预训练、稀疏优化、贝叶斯推断等场景，特别是在潜在函数复杂、非光滑、超线性增长的环境下。其稳定性和收敛性保证，为训练深层神经网络提供了理论基础，有助于提升模型的泛化能力和鲁棒性。未来，结合自适应参数调节和模型结构优化，有望在实际工业界实现更高效、更稳定的深度模型训练流程。

尽管算法在高维非凸场景中表现优越，但其收敛速率在逆温度较高时呈指数级依赖，可能限制在极端高逆温度环境中的应用。此外，潜在函数的半凸性和多项式增长假设对某些极端非凸或超线性潜能仍有局限。实际操作中，参数调节和数值稳定性仍需优化，尤其在极高维模型中可能面临挑战。未来需研究更宽泛的潜能类别和更高效的参数调优策略，以拓展算法的适用范围。