Non-asymptotic estimates for TUSLA algorithm for non-convex learning with applications to neural networks with ReLU activation function

Abstract

We consider non-convex stochastic optimization problems where the objective functions have super-linearly growing and discontinuous stochastic gradients. In such a setting, we provide a nonasymptotic analysis for the tamed unadjusted stochastic Langevin algorithm (TUSLA) introduced in Lovas et al. (2021). In particular, we establish non-asymptotic error bounds for the TUSLA algorithm in Wasserstein-1 and Wasserstein-2 distances. The latter result enables us to further derive non-asymptotic estimates for the expected excess risk. To illustrate the applicability of the main results, we consider an example from transfer learning with ReLU neural networks, which represents a key paradigm in machine learning. Numerical experiments are presented for the aforementioned example which supports our theoretical findings. Hence, in this setting, we demonstrate both theoretically and numerically that the TUSLA algorithm can solve the optimization problem involving neural networks with ReLU activation function. Besides, we provide simulation results for synthetic examples where popular algorithms, e.g. ADAM, AMSGrad, RMSProp, and (vanilla) SGD, may fail to find the minimizer of the objective functions due to the super-linear growth and the discontinuity of the corresponding stochastic gradient, while the TUSLA algorithm converges rapidly to the optimal solution.

Non-asymptotic convergence bounds for modified tamed unadjusted Langevin algorithm in non-convex setting

“…the Hessian of U . This additional smoothness condition is similar to Assumption H4 of [Bro+19] and we impose it in order to obtain improved convergence rates compared to those derived in [Lov+20] and [Lim+21] for TUSLA.…”

“…In [Lim+21], the polynomial Lipschitz condition imposed in Assumptions 2 and 3 of [Lim+21], and the convex at infinity condition imposed in Assumption 4 of [Lim+21] are directly analogous to Assumptions 2 and 3 of this paper in the case of deterministic gradient. With our choice of taming factor as (1+λ|θ| 2r ) 1/2 , see (7), as opposed to (1 + √ λ|θ| r ) in [Lim+21], as well as assuming additional smoothness conditions on U in Assumption 4 of our work, we were able to improve the orders of convergence of 1/2 and 1/4 for Wasserstein distances of orders 1 and 2, respectively, derived in [Lim+21], to 1 and 1/2, respectively, in Theorems 2.7 and 2.8.…”

“…Comparison with Related Works. In this section, we compare our main results and assumptions under which they were obtained with those of [Bro+19] and [Lim+21] where tamed variants of ULA were proposed to deal with the case of super-linearly growing gradient.…”

“…In [Lov+20] and [Lim+21], the Tamed Unadjusted Stochastic Langevin Algorithm (TUSLA), a tamed variant of the SGLD algorithm, was proposed and studied under a polynomial Lipschitz assumption on the gradient and in a non-convex setting. Non-asymptotic upper bounds for Wasserstein distances of orders 1 and 2 with respective convergence orders 1/2 and 1/4 in the step size were obtained for TUSLA.…”

“…In comparison with [Bro+19], we propose a different taming factor for the gradient, and our theoretical guarantees in Wasserstein distances are obtained under an equivalent polynomial Lipschitz assumption on the gradient, but in a non-convex setting, as opposed to the typical strong convexity condition imposed in [Bro+19] in order to obtain the bounds in Wasserstein-2 distance. In addition, we impose a comparable 'convex at infinity' assumption on the potential U as in [Lim+21], but achieve markedly improved convergence orders of 1 and 1/2 in step size for Wasserstein distances of orders 1 and 2, respectively, compared to 1/2 and 1/4, respectively, obtained in [Lim+21].…”