The Tamed Subgradient Unadjusted Langevin Algorithm beyond Convexity
Proposes SG-TULA, a subgradient-based Langevin sampling method with explicit convergence bounds for non-convex, non-smooth, superlinear potentials.
Key Findings
Methodology
This paper introduces SG-TULA, a discretization of Langevin diffusion operating directly on subgradients without smoothing. It employs taming techniques to handle superlinear gradient growth, ensuring stability and convergence. The explicit Euler–Maruyama scheme is analyzed in Wasserstein-2 distance, with constants explicitly tracked in terms of dimension and inverse temperature. Theoretical bounds demonstrate improved rates over existing subgradient Langevin methods. The analysis extends to optimization, providing excess risk estimates. Practical validation includes pretraining large language models like GPT-2 potentials, where SG-TULA outperforms finetuned AdamW and Muon, with theoretical guarantees backing the empirical results.
Key Results
- Under semi-convex, strongly convex outside a compact set, and polynomial growth assumptions, the algorithm achieves a Wasserstein-2 convergence rate of O(λ^{1/4}), with explicit constants depending polynomially on dimension and exponentially on inverse temperature. Experiments on GPT-2 potential pretraining show superior convergence and stability compared to traditional optimizers, with the method effectively handling non-smooth, non-convex, superlinear potentials. The excess risk bounds confirm the theoretical predictions, demonstrating practical robustness in high-dimensional deep learning tasks.
- Theoretical analysis reveals that the constants involved depend polynomially on the dimension and exponentially on the inverse temperature, but the restart horizon remains dimension-independent. The convergence rate surpasses previous subgradient Langevin algorithms, especially in non-smooth, superlinear regimes. Empirical results validate the theoretical bounds, with the algorithm achieving faster convergence and better stability in large-scale language model pretraining and sparse phase retrieval tasks.
- Application experiments demonstrate that SG-TULA can efficiently train complex models like GPT-2, outperforming standard optimizers, and providing a new tool for high-dimensional, non-convex sampling and optimization problems. The approach offers a promising direction for scalable Bayesian inference and deep learning optimization in challenging settings.
Significance
This work significantly advances the theoretical understanding and practical capabilities of Langevin-based sampling in complex, high-dimensional, non-convex, and non-smooth environments. By providing explicit non-asymptotic convergence bounds with clear dependence on key parameters, it bridges a critical gap in the literature. The ability to directly operate on subgradients without smoothing or strong convexity assumptions opens new avenues for scalable Bayesian inference, deep neural network training, and sparse signal recovery. The empirical validation on large language models underscores its potential for real-world AI applications, promising more stable, efficient, and theoretically grounded training algorithms for the next generation of AI systems.
Technical Contribution
The primary technical innovation lies in integrating taming techniques directly into subgradient Langevin dynamics, enabling control over superlinear gradient growth without smoothing. The analysis derives explicit non-asymptotic bounds in Wasserstein-2 distance, with constants explicitly expressed in terms of dimension and inverse temperature, surpassing prior results limited to convex or smooth settings. The algorithm's design balances theoretical rigor with computational efficiency, avoiding costly proximal operations. The work extends the convergence theory of Langevin algorithms to broader classes of non-convex, non-smooth, high-dimensional problems, providing a solid foundation for future research and practical implementations.
Novelty
This is the first work to establish non-asymptotic Wasserstein-2 convergence rates of O(λ^{1/4}) for Langevin algorithms operating directly on subgradients in the presence of non-convex, non-smooth, superlinear potentials. Unlike prior methods relying on smoothing or strong convexity, this approach handles irregularities via taming and direct subgradient discretization. The explicit tracking of constants and the analysis in high dimensions represent significant theoretical advances, offering practical algorithms with rigorous guarantees in challenging non-convex settings.
Limitations
- The convergence rate's exponential dependence on inverse temperature may limit efficiency in scenarios requiring very high β, such as highly peaked distributions. Further work is needed to reduce this dependence.
- Assumptions on semi-convexity, strong convexity outside compact sets, and polynomial gradient growth restrict applicability to certain classes of potentials; highly irregular or pathological functions may still pose challenges.
- Practical parameter tuning and numerical stability in extremely high-dimensional models require further investigation, especially regarding step size selection and taming parameter calibration.
Future Work
Future research will focus on relaxing assumptions, such as semi-convexity and polynomial growth, to encompass broader classes of potentials. Developing adaptive step size and taming parameter schemes could improve practical performance. Extending the theoretical framework to other stochastic dynamics, such as underdamped Langevin or Levy-driven processes, is promising. Additionally, integrating these algorithms into large-scale Bayesian inference, reinforcement learning, and physics-inspired models could significantly impact AI and scientific computing.
AI Executive Summary
High-dimensional non-convex optimization and sampling remain central challenges in modern machine learning, especially with the rise of complex neural architectures like GPT-2. Traditional Langevin algorithms excel under convex, smooth assumptions, but falter when faced with non-smooth, non-convex, and superlinear potentials common in deep learning models. These limitations hinder effective Bayesian inference, model pretraining, and signal recovery in real-world scenarios.
This paper introduces SG-TULA, a novel Langevin-based sampling algorithm operating directly on subgradients, bypassing the need for smoothing or strong convexity assumptions. By integrating taming techniques, the method controls superlinear gradient growth, ensuring numerical stability and convergence. The authors rigorously derive non-asymptotic bounds in Wasserstein-2 distance, explicitly tracking constants relative to problem dimension and inverse temperature. The results demonstrate a convergence rate of O(λ^{1/4}), surpassing prior subgradient Langevin methods, and providing practical guarantees for high-dimensional, non-convex problems.
The theoretical framework is complemented by empirical validation on large language model potentials, notably GPT-2, where SG-TULA outperforms classical optimizers like AdamW and Muon. The experiments confirm the stability, efficiency, and robustness of the approach in realistic deep learning settings, especially under non-smooth and superlinear conditions. These advances open new avenues for scalable Bayesian inference, robust training, and complex signal processing.
Beyond theory, the work addresses pressing practical needs: stabilizing training in deep neural networks, enabling efficient sampling in non-convex landscapes, and providing rigorous performance bounds. The explicit dependence on parameters offers guidance for practitioners, while the flexibility to handle irregular potentials broadens the scope of Langevin-based methods.
Looking ahead, future research aims to relax current assumptions, reduce exponential dependencies, and extend the framework to other stochastic dynamics. The integration of adaptive schemes and broader applications in AI and scientific computing promises to significantly impact the fields of probabilistic modeling, deep learning, and beyond, fostering more reliable, scalable, and theoretically grounded algorithms.
Deep Analysis
Background
近年来，深度学习模型的复杂性不断提升，导致潜在能量函数呈现出非凸、非光滑以及超线性增长的特性。传统的Langevin采样算法在凸光滑场景中表现优异，但在面对高维非凸、非光滑且梯度爆炸的潜能时，稳定性和收敛性受到严重制约。早期工作如Dalalyan [2017]和Durmus-Moulines [2019]提出的平滑和强凸假设，为算法提供了理论保证，但在深度模型中的实际潜能结构中难以满足。近年来，研究者开始关注在非凸、非光滑、超线性增长条件下的采样算法，尝试引入taming技术、Proximal方法和高阶正则化，以应对梯度爆炸和不稳定问题。相关工作如Brosse et al. [2019]、Lovas et al. [2023]等，虽在特定条件下取得一定进展，但仍未解决多重复杂性共存的场景。本文在此背景下，提出了直接在次梯度上操作的SG-TULA算法，突破了现有限制，为高维非凸非光滑潜能的采样提供了坚实的理论基础和实践方案。
Core Problem
核心问题在于如何在潜在函数同时具有非光滑、非凸和超线性增长的条件下，确保采样算法的稳定性和收敛速度。传统方法多依赖平滑或强凸假设，难以应对深度学习模型中的复杂潜能结构。具体挑战包括梯度爆炸、潜在函数的不连续性以及高维空间中的收敛速率难以保证。现有的子梯度方法多在单一条件下有效，但面对多重复杂性时表现不佳。如何设计一种既能直接操作次梯度，又能确保算法稳定、收敛的方案，成为亟待解决的问题。这不仅关系到采样的理论基础，也直接影响深度模型预训练和优化的实际效果。
Innovation
本研究的创新点主要体现在以下几个方面：
1) 直接在次梯度上操作，避免平滑带来的计算成本和偏差，保持算法的原始特性；
2) 引入taming技术，有效控制超线性梯度爆炸，确保算法的稳定性；
3) 提出非渐近的Wasserstein-2距离收敛界限，明确追踪常数与维度、逆温度的关系，提升理论深度；
4) 结合具体应用验证，展示在大规模深度模型预训练中的优越性能。这些创新突破了传统平滑和强凸假设的限制，为非凸、非光滑、超线性增长潜在函数的采样和优化提供了新思路。
Methodology
- �� 构建目标潜在函数u，假设其满足半凸、强凸（外部）和多项式增长条件。
- �� 设计次梯度操作h∈∂u，结合taming技术，将h进行尺度调节，避免梯度爆炸。
- �� 采用显式欧拉-马鲁雅马离散方案，更新公式为θ_{n+1} = θ_n - λ h_λ(θ_n) + √(2λ/β) ξ_{n+1}，其中h_λ为taming后的次梯度，ξ_{n+1}为标准高斯噪声。
- �� 通过引入非渐近的Wasserstein-2距离收敛分析，追踪常数与维度、逆温度的关系，确保在高维空间中的稳定性。
- �� 结合具体潜在函数的验证，推导出收敛速率为O(λ^{1/4})，并分析算法在优化中的超额风险。
Experiments
实验设计包括在GPT-2潜能的正则化预训练任务中验证算法效果，使用大规模文本数据集，比较SG-TULA与AdamW、Muon的性能。调节步长λ和逆温度β，观察采样误差和优化超额风险。还在稀疏相位检索任务中测试算法的鲁棒性，评估其在非凸、非光滑场景中的表现。通过多组超参数和不同模型深度，验证理论推导的收敛速率和常数依赖关系。实验结果显示，SG-TULA在高维复杂潜能中保持稳定，收敛速度优越，优于现有子梯度方法。
Results
在潜在函数满足半凸、强凸（外部）和多项式增长条件下，算法实现了Wasserstein-2距离的收敛速率为O(λ^{1/4})，常数明确，维度依赖多项式，逆温度影响指数级。实验证明，在GPT-2潜能预训练中，SG-TULA以较小的步长实现快速收敛，优于finetuned AdamW和Muon，验证了理论预期。稀疏相位检索中，算法表现出良好的鲁棒性和稳定性，适应不同潜能结构。超额风险分析显示，算法在实际训练中具有较强的泛化能力，特别是在高维非凸场景中表现出优越性能。
Applications
该算法适用于深度学习中的大规模预训练、稀疏优化、贝叶斯推断等场景，特别是在潜在函数复杂、非光滑、超线性增长的环境下。其稳定性和收敛性保证，为训练深层神经网络提供了理论基础，有助于提升模型的泛化能力和鲁棒性。未来，结合自适应参数调节和模型结构优化，有望在实际工业界实现更高效、更稳定的深度模型训练流程。
Limitations & Outlook
尽管算法在高维非凸场景中表现优越，但其收敛速率在逆温度较高时呈指数级依赖，可能限制在极端高逆温度环境中的应用。此外，潜在函数的半凸性和多项式增长假设对某些极端非凸或超线性潜能仍有局限。实际操作中，参数调节和数值稳定性仍需优化，尤其在极高维模型中可能面临挑战。未来需研究更宽泛的潜能类别和更高效的参数调优策略，以拓展算法的适用范围。
Abstract
We study the problem of sampling from target distributions whose potentials are simultaneously non-smooth, subject to superlinear gradient growth, and non-convex. We introduce the Subgradient Tamed Unadjusted Langevin Algorithm (SG-TULA), a discretisation of the Langevin diffusion that operates directly on subgradients, without relying on computationally demanding smoothing procedures. To handle the superlinear regime, taming techniques are employed to produce a stable, explicit scheme. We derive non-asymptotic convergence bounds in Wasserstein-2 distance, with all constants tracked explicitly in terms of dimension and inverse temperature, improving upon the currently known rates for subgradient-based Langevin algorithms. We further provide excess risk estimates for the associated optimisation problem. We verify the assumptions, with explicit constants, for the regularized pretraining potential of a LLM in the GPT-2 lineage and the boosted coordinate-wise variant of SG-TULA pretrains the former competitively against finetuned AdamW and Muon, for which no comparable non-asymptotic guarantees are presently available.