Convergence Analysis of Stochastic Gradient Descent with MCMC Estimators

To solve the stochastic optimization problem (1.1), a common approach is to compute the stochastic gradient for iterations. As the Assumption 1 holds, the gradient of objective function $\mathcal{L}(\theta)$ is derived by

The second equality is due to $\mathbb{E}_{x\sim\pi_{\theta}}\left[\nabla_{\theta}f_{\theta}(x)\right]=0$ and $\nabla_{\theta}\log\pi_{\theta}(x)=\nabla_{\theta}\phi_{\theta}(x)-\mathbb{E}_{x\sim\pi_{\theta}}\left[\nabla_{\theta}\phi_{\theta}(x)\right]$, implied by Assumption 1. We notice that the gradient can be expressed in the form of an expectation, thus allowing the approximation of the gradient through the Monte Carlo method. As for the unnormalized distribution $\pi_{\theta}(x)\propto\phi_{\theta}(x)$, it is natural to employ MCMC methods to estimate the expectations presented within gradients.

Let us consider a Markov chain $\{X_{i}\}_{i=1}^{\infty}$ with a stationary distribution $\pi_{\theta}(x)\propto e^{\phi_{\theta}(x)}$. We have obtained a set of samples $S=\{x_{i}\}_{i=n_{0}+1}^{n_{0}+n}$ following $n_{0}$ burn-in steps. These samples are generated to approximate expectations with respect to $\pi_{\theta}(x)$, which is typically intractable. The objective function and its gradient is estimated by the MCMC samples $S$:

While MCMC methods offer a tractable and efficient strategy for approximating expectations, it must be noted that estimators derived from finite MCMC samples are generally biased. Furthermore, the samples are correlated, which affects the variance of the estimator. To mitigate bias, one must ensure a sufficiently large sample size. However, an excessively large sample size may increase computational burden and reduce the stochasticity for escaping local regions. In light of these considerations, the practitioner must carefully balance the sample size to minimize bias while preserving randomness and computational tractability. A rigorous approach may involve adaptive MCMC methods that adjust the burn-in period and sample size based on convergence diagnostics of Markov chains. It suggests that we should pay attention to the influence of sampling methods in optimization.

The SGD algorithm is utilized to update the parameters with MCMC samples:

where $\alpha_{k}$ are chosen stepsizes and $S_{k}$ is the obtained samples with the stationary distribution $\pi_{\theta_{k}}$. MCMC-SGD provides a powerful framework for optimization in settings where the objective function is defined in terms of expectations over complex distributions. A detailed analysis on the interplay between SGD and MCMC sampling will aid us in deepening our understanding of convergence and refining the algorithms.

To lay the groundwork for our subsequent assumptions, we first expound upon the definition of sub-exponential random variables. This class of random variables is characterized by their tail behavior, which, in a certain sense, is less extreme than that of heavy-tailed distributions. A precise mathematical definition is as follows:

We introduce some regularity conditions on $f_{\theta}$. Noticing that $f_{\theta}$ may be unbounded for a general class problems, we give the sub-exponential assumption of $f_{\theta}$ under the distribution $\pi_{\theta}$.

The mentioned two assumptions are both reasonable and encompass a broad class of scenarios. A common example is $f_{\theta}=-r(x)+\beta\log\pi_{\theta}(x)$ with $|r(x)|\leq R$ for all $x\in\mathcal{X}$, which often appears in the optimization of entropy regularization and KL divergence. In this situation, $\left\|f_{\theta}\right\|_{\psi_{1}(\pi_{\theta};\mathcal{X})}\leq t$ is equivalent to

Since $\int_{x\in\mathcal{X}}\pi_{\theta}dx=1$ holds for the distribution $\pi_{\theta}$, there must exist a constant $t_{\theta}$ relying $\theta$ such that $\int_{x\in\mathcal{X}}e^{R/t_{\theta}}\pi_{\theta}(x)^{1-\beta/t_{\theta}}dx\leq 2$ and then Assumption 2 (1) is easily satisfied for a compact parameter space $\Theta$. Meanwhile, $\mathbb{E}_{\pi_{\theta}}\left[\left\|\nabla_{\theta}f_{\theta}(x)\right\|^{2}\right]=\beta^{2}\mathbb{E}_{\pi_{\theta}}\left[\left\|\nabla_{\theta}\phi_{\theta}(x)-\mathbb{E}_{\pi_{\theta}}\left[\nabla_{\theta}\phi_{\theta}(x)\right]\right\|^{2}\right]$, which can be bounded through Assumption 3 (1) on the variational distribution $\pi_{\theta}$. By establishing the reasonableness of the Assumption 3 (1) on the variational distribution, one can deduce the validity of the Assumption 2 (2) in this setting. For another instances, when solving the many-electron Schrödinger equations, people often use the Jastrow factor [18] to enable the network to efficiently capture the cusps and decay of the wavefunction. By exploiting physical knowledge in the construction of the wavefunction ansatz can make the function $f_{\theta}$ smooth with respect to the parameters $\theta$, such that Assumption 2 is satisfied in practical VMC calculations.

There are also some regularity conditions on the variational distribution $\pi_{\theta}(x)$, in order to guarantee Lipschitz continuity of the gradient $g(\theta)$. Typically, it is more practical to apply these smoothing measures to the energy function $\phi_{\theta}(x)$ rather than directly to the distribution $\pi_{\theta}(x)$. This kind of smoothing helps us control how the gradient changes, which is important for making sure our optimization methods work efficiently and reliably.

The reasonableness of these two assumptions is solely based on the choice of parameterization. For instance, $\pi_{\theta}$ is taken as a canonical form of exponential family, i.e., $\phi_{\theta}(x)=\theta^{\mathrm{T}}T(x)$, where $\theta$ represents the natural parameter vector and $T(x)$ is the sufficient statistic. As $\nabla_{\theta}\phi_{\theta}(x)=T(x)$ and $\nabla_{\theta}^{2}\phi_{\theta}(x)=0$, the Assumption 3 holds when $T(x)$ is uniformly bounded for all $x\in\mathcal{X}$. If $\phi_{\theta}$ is parameterized by a complex network, it is equally imperative to ensure that its gradients with respect to any input and parameters are bounded.

We introduce following notations which are used frequently throughout this paper. For any function $f:\mathcal{X}\rightarrow\mathbb{R}$ and any distribution $\pi$ on $\mathcal{X}$, we write its expectation $\mathbb{E}_{\pi}[f]:=\int f(x)\pi(dx)$ and $p$-th central moment $\sigma^{p}_{p}[f]:=\mathbb{E}[(f-\mathbb{E}_{\pi}[f])^{p}]$. The second central moment, called the variance, is denoted by $\mathrm{Var}_{\pi}[f]=\sigma_{2}^{2}[f]:=\mathbb{E}[(f-\mathbb{E}_{\pi}[f])^{2}]$.

Before the analysis of MCMC methods, we provide several general definitions about Markov chains. Within our consideration, the state space $\mathcal{X}$ is Polish and equipped with its $\sigma$-algebra $\mathcal{B}$. Let $\{X_{i}\}_{i=1}^{n}$ be a time-homogeneous Markov chain defined on $\mathcal{X}$. The distribution of the Markov chain is uniquely determined by its initial distribution $\nu$ and its transition kernel $P$. For any Borel set $A\in\mathcal{B}$, let

A distribution $\pi$ is called stationary with respect to a transition kernel $P$ if

When the initial distribution $\nu=\pi$, we call the Markov chain stationary.

Our analysis starts from the perspective of operator theory on Hilbert spaces. Let $\pi$ be the stationary distribution of a Markov chain and $L^{2}(\pi)=\{f:\mathbb{E}_{\pi}[f^{2}]<\infty\}$ be the Hilbert space equipped with the norm $\left\|f\right\|_{\pi}=(\mathbb{E}_{\pi}[f^{2}])^{1/2}$. Each transition kernel can be viewed as a Markov operator on the Hilbert space $L^{2}(\pi)$. The Markov operator $\operatorname{P}:L^{2}(\pi)\rightarrow L^{2}(\pi)$ is defined by

It is easy to show that $\operatorname{P}$ has the largest eigenvalue $1$. Intuitively but not strictly, the gap between $1$ and other eigenvalues matters to the Markov chain from non-stationarity towards stationarity. Hence, we introduce the definition of the absolute spectral gap.

We consider the sample set $S=\{x_{i}\}_{i=n_{0}+1}^{n_{0}+n}$ generated from the desired distribution $\pi_{\theta}$ by the MCMC algorithm . The Markov operator, denoted by $\operatorname{P}_{\theta}$, determines the convergence rate of the Markov chain. We assume that there exists a uniform lower bound of the spectral gap.

This assumption excludes the situation that $\inf_{\theta\in\Theta}\gamma(\operatorname{P}_{\theta})=0$, which ensures the uniform convergence of the Markov chain.. The spectral gap $\gamma$ measures the lower bound on the mixing rate of the Markov chain, thereby guaranteeing the ergodicity of the MCMC algorithm. Compared to the uniformly geometric ergodicity assumptions commonly found in other articles [22, 3], the spectral gap assumption is significantly tighter, intuitively reflecting the factor that influence the convergence rate of a Markov chain. The spectral gap of the MCMC algorithm has been studied for some specific examples. On finite-state spaces, $\operatorname{P}$ becomes a transition matrix while $1-\gamma(\operatorname{P})$ relates to its eigenvalue. A survey of spectrums for Markov chains on discrete state-spaces is in [43]. This develops amounts of analytic techniques and [10] has further applications to the Metropolis-Hastings algorithm. For the continuous spaces, there are few examples of sharp rates of convergence for the Metropolis-Hastings algorithm. In [24, 31], it is claimed that the spectral gap $\gamma\sim O(\epsilon^{2})$ for the Gaussian proposal $Q(x^{\prime}|x)\sim\exp\left(-\tfrac{1}{2\epsilon^{2}}\left\|x^{\prime}-x\right\|^{2}\right)$.

We commence by stating Theorem 1, which articulates the core result that analyzes the error of MCMC gradient estimators. The proof of this theorem will rely on a series of auxiliary propositions, including Lemmas 1, 2 and 3. Lemma 1 introduces a Bernstein inequality for general Markov chains. Using the Bernstein inequality, we obtain the tail bound of the MCMC estimators for sub-exponential functions in Lemma 2. Then, a detailed and intuitive error bound for the MCMC estimators is of consideration in Lemma 3. Finally, through Lemma 3, we succinctly capture the sampling error inherent in MCMC algorithms. In the following theorem, we delineate the error bounds of this stochastic gradient, which is one of main results in this paper.

To refine the proof of the Theorem 1, we first introduce a pivotal tool in this paper, concentration inequalities, to analyze the error of the MCMC estimators in details. The essence of concentration inequalities lies in their capacity to quantify how a random variable can concentrate around its expected value. A Bernstein inequality for general Markov chains is proposed by Jiang and Fan [19] and beneficial to our analysis. The following lemma is a direct corollary of Theorem 2 in [19].

The Bernstein inequality (3.5) shows how the average of MCMC samples concentrates at the expectation. However, the Markov chain is not always non-stationary. We define ${\mathchoice{\raisebox{0.0pt}{$\displaystyle\chi$}}{\raisebox{0.0pt}{$\textstyle\chi$}}{\raisebox{0.0pt}{$\scriptstyle\chi$}}{\raisebox{0.0pt}{$\scriptscriptstyle\chi$}}}^{2}(\nu,\pi):=\left\|\tfrac{d\nu}{d\pi}-1\right\|^{2}_{\pi}$ represents the chi-squared divergence between $\nu$ and $\pi$, by which the Bernstein inequality can be extended into a non-stationary one. We abbreviate ${\mathchoice{\raisebox{0.0pt}{$\displaystyle\chi$}}{\raisebox{0.0pt}{$\textstyle\chi$}}{\raisebox{0.0pt}{$\scriptstyle\chi$}}{\raisebox{0.0pt}{$\scriptscriptstyle\chi$}}}={\mathchoice{\raisebox{0.0pt}{$\displaystyle\chi$}}{\raisebox{0.0pt}{$\textstyle\chi$}}{\raisebox{0.0pt}{$\scriptstyle\chi$}}{\raisebox{0.0pt}{$\scriptscriptstyle\chi$}}}(\nu,\pi)$ and $C=(1+{\mathchoice{\raisebox{0.0pt}{$\displaystyle\chi$}}{\raisebox{0.0pt}{$\textstyle\chi$}}{\raisebox{0.0pt}{$\scriptstyle\chi$}}{\raisebox{0.0pt}{$\scriptscriptstyle\chi$}}}^{2})^{\frac{1}{2}}$. Then, we derive the tail bound of the MCMC estimator for sub-exponential functions in the following lemma.