Nesterov smoothing for sampling without smoothness
We study the problem of sampling from a target distribution in $\mathbb{R}^d$ whose potential is not smooth. Compared with the sampling problem with smooth potentials, this problem is much less well-understood due to the lack of smoothness. In this paper, we propose a novel sampling algorithm for a class of non-smooth potentials by first approximating them by smooth potentials using a technique that is akin to Nesterov smoothing. We then utilize sampling algorithms on the smooth potentials to generate approximate samples from the original non-smooth potentials. We select an appropriate smoothing intensity to ensure that the distance between the smoothed and un-smoothed distributions is minimal, thereby guaranteeing the algorithm's accuracy. Hence we obtain non-asymptotic convergence results based on existing analysis of smooth sampling. We verify our convergence result on a synthetic example and apply our method to improve the worst-case performance of Bayesian inference on a real-world example.
\section{Introduction}
\looseness=-1 Sampling from a target distribution $\pi (x ) \propto \exp(-s(x)) $ known up to a normalization constant is an important problem in many areas such as machine learning, due to its pivotal role in Bayesian statistics and inference~\citep{gelman1995bayesian,durmus2018efficient,krauth2006statistical, fan2021variational}.
Over the past decades, there has been a vast amount of research carried out on sampling problems from smooth distributions~\citep{dalalyan2017theoretical,vempala2019rapid}, i.e., the gradient of the potential $s(x)$ is Lipschitz. In contrast, the understanding of non-smooth sampling is still quite limited
~\citep{chatterji2020langevin,salim2020primal,liang2021proximal,liang2022proximal,chewi2021analysis}.

In this work, we consider the task of sampling from a specific class of non-smooth distributions $\pi(x) \propto \exp(-s(x))$ whose potentials $s(x)$ permit an explicit max-structure as
\begin{align}\label{eq:potential}
  s(x) = f(x) +  \max_{y \in {\mathcal Y}} \{ \<h(x), y\> -g(y) \}
 
 
\end{align}
where $f$ is smooth, $g$ is convex, ${\mathcal Y}$ is convex and bounded, and $h$ is Lipschitz and smooth.
\fan{$s(x)$ can be non-smooth.}
For instance, when $f(x)=0 $ and $h(x) =Ax$, $s(x)$ becomes a piece-wise affine function and is clearly non-smooth. 
\chen{A special instance of particular interest is 
\begin{equation}\label{eq:maxh}
  s(x) = f(x)+ \max_{i} \{h_1(x), h_2(x),\ldots, h_n(x)\}.
\end{equation}
The optimization of minimizing $s(x)$ has been widely used for science and engineering problems to optimize the worst-case performance \citep{ben2009robust,diehl2006approximation}. Its Bayesian counterpart that accounts for uncertainties can be captured by a sampling problem from the potential $s(x)$.}

Compared to other recent work ~\citep{chatterji2020langevin,salim2020primal,liang2021proximal,liang2022proximal,chewi2021analysis} for sampling from black box non-smooth distribution, we assume certain structures in the target distribution and hope to leverage these structures to achieve better performance. Indeed, in practice, it is rare to have a black box model as a target distribution; we almost always know some structures of the problem. It could be advantageous to utilize such structures in the algorithm design.
With this in mind, we take one step forward in the direction of sampling beyond black box models and propose a sampling method for non-smooth distribution that utilizes the specific max-structure \eqref{eq:potential}. We develop a sampling algorithm for non-smooth potentials whose complexity has the same dimensional dependency as its smooth counterparts.


Our strategy is akin to Nesterov smoothing~\citep{nesterov2005smooth} in optimization that constructs a smooth function $s_\beta(x)$ whose difference with $s(x)$ can be controlled through the smoothing intensity $\beta$. In Nesterov smoothing, the idea is to use a fast smooth optimization algorithm to minimize $s_\beta(x)$ for the purpose of optimizing the non-smooth function $s(x)$. In our problem, instead, we sample from $\pi_\beta (x) \propto \exp(-s_\beta(x))$ for the purpose of sampling from $\pi(x) \propto \exp(-s(x))$. Thanks to the smoothness of $s_\beta(x)$, we can take advantage of the many existing sampling algorithms for smooth potentials to sample from $\pi_\beta (x) \propto \exp(-s_\beta(x))$. Our smoothing strategy is compatible with any such algorithms. With state-of-the-art sampling algorithms for smooth potentials and a proper smoothing intensity $\beta$, we obtain a series algorithms with competitive complexity bounds. Surprisingly,
the complexity bounds we obtained with Nesterov smoothing have exactly the same dimension dependency as their smooth counterparts. This is crucial for high dimensional problems.
In a high level, this work is along the direction of the recent line of research that tries to bridge optimization and sampling~\citep{wibisono2018sampling,barp2022geometric,liang2022proximal}. Our results prove that some smoothing techniques used in optimization are equally effective in sampling.


\paragraph{Related Works:}
Over the last few years, several algorithms for sampling from potentials that are not smooth have been developed. A majority of them is devoted to the analysis of original LMC in the non-smooth setting \citep{erdogdu2021convergence,nguyen2021unadjusted,chewi2021analysis}. These theorectical results are challenging and significantly different from those in the smooth setting. In \citet{bernton2018langevin,durmus2019analysis,salim2020primal}, the non-smooth sampling problem is studied and analyzed from an optimization point of view. Another algorithm that is effective for non-smooth sampling is the proximal sampler developed in the sequence of work \citep{lee2021structured,liang2021proximal,liang2022proximal,chen2022improved}. The methods that are most related to ours are \citet{chatterji2020langevin,durmus2018efficient,lau2022bregman}, in which the authors use Gaussian smoothing and Moreau envelope.
\chen{
Gaussian smoothing can be applied to any \fan{convex} potential, however the evaluation of the smoothed potential is difficult. Normally one can only get unbiased estimation of it through Gaussian sampling, inducing additional variance in the algorithm. Moreau envelope can be applied to any \fan{weakly-convex} potential.
To evaluate the smoothed potential or its gradient, one needs to solve an optimization problem (to get the proximal map), inducing additional complexity, \fan{especially when the potential is not convex. We include a more detailed comparison with these two smoothing techniques in the end of Section \ref{sec:error}.}}

\paragraph{Contributions:}
We summarize our contributions as follows. \\
i) ~~~Inspired by Nesterov smoothing in optimization, we develop a smoothing technique, different from existing ones like Gaussian smoothing and Moreau envolep, for sampling problems for a class of potentials that permit the max-structure \eqref{eq:potential}.  \\
ii) ~~Our smoothing technique is compatible with any first-order sampling algorithm for smooth potentials and the performance is guaranteed via a proper smoothing intensity $\beta$. \\
iii) ~Combining this smoothing technique with start-of-the-art sampling algorithms for smooth potentials we obtain competitive complexity bounds for strongly-log-concave distributions, log-concave distributions, and distributions satisfying log-Sobolev inequality (see Table \ref{table}). In particular, we achieve ${d}^{1/3}$ dependence in sampling non-smooth strongly-log-concave distributions, better than all existing methods.
Besides, our analysis covers certain composite potentials with non-smooth (Lipschitz continuous) components and the distribution satisfies LSI. The complexity bound in this setting has not been established before.


\begin{remark}
  \fan{
   
    In Table \ref{table}, our result contains one additional problem-dependent parameter $D$ defined in \eqref{eq:bound_l}. It can hide dimension dependency, like any other parameter $\lambda_h,\lambda_\rho$ or $ L_f $. In Abstract and Contribution, our claim in terms of dimension dependency assumes other parameters are dimension-free. In realistic examples, this is unlikely to be true, so one need to include the dimension dependency hidden in all parameters for a fair comparison. 
An existing caveat is if $D$ hides dimension dependency, 
then our algorithm is clearly not state-of-art. \chen{Interestingly, for the class of problems with potentials of the form \eqref{eq:maxh}, $D = \log (n)$ will be seen as in Section \ref{sec:ex}.}
    }
\end{remark}

\begin{table}[]
  \caption{
   
    Convergence comparison for \fan{state-of-the-art} non-smooth sampling. We highlight the dependency on dimension $d$, accuracy ${\mathcal E}$, \fan{smoothness $L_f$ of $f(x)$,} Lipschitz  $\lambda_h$ of $h(x)$, Lipschitz $\lambda_\rho$ of $\rho(x):= \max_{y \in {\mathcal Y}} \{ \<h(x), y\> -g(y) \}$, \fan{boundness  $D$ defined in \eqref{eq:bound_l},} and log-Sobolev constant $C_\pi$. A trivial upper bound for $\lambda_\rho $ is $\lambda_h \max \|y\|_2 $. \fan{
   
      Other than \citet{durmus2019analysis},
     
      all results are the last-iterate convergence.
   
      } }
  \label{table}
  \centering
  {\renewcommand{\arraystretch}{2.1}%
    \begin{tabular}{|c|c|c|c|c|}
      \hline
      Assumption                                 & Source                           & Complexity                                                                                 & Metric & Note                     \\ \hline
      \multirow{5}{*}{\begin{tabular}[c]{@{}c@{}}$\alpha$-strongly \\ log-concave\end{tabular}} & \citep{liang2022proximal}        & $\widetilde{\mathcal O} \left( \frac{(\lambda_\rho^{2} \vee L_f) d }{\alpha } \right)$              & R\'enyi \xspace   &                          \\\cline{2-5}
                                                 & \citep{salim2020primal}          & $\widetilde{\mathcal O} \left( \frac{\lambda_\rho^{2} + L_f d }{\alpha^2 {\mathcal E}^2 } \right)$           & $W_2$  &                          \\\cline{2-5}
                                                 & \citep{chatterji2020langevin}    & $\widetilde{\mathcal O} \left( \frac{(\lambda_\rho+L_f)^{2}  d^{5/2} }{\alpha^4 {\mathcal E}^4 } \right)$              & $W_2$     & Gaussian smoothing       \\\cline{2-5}
                                                 & Theorem \ref{thm:str_convex} (1) & $\widetilde{\mathcal O} \left(
    \frac{\lambda_h^{2} D d^{1/3} }{\alpha^{11/6} {\mathcal E}^{5/3}} 
    +    \frac{\lambda_h^{7/3} D^{7/6} d^{1/6} }{\alpha^{23/12} {\mathcal E}^{3/2}}
    \right)$ & $W_2$  &                          \\\cline{2-5}
                                                 & Theorem \ref{thm:str_convex} (2) & $ \widetilde{\mathcal O} \left( \frac{ \lambda_h^2 D d^{1/2}}{\alpha {\mathcal E} } \right)$                 & TV     & Assume warm-start        \\\hline
     
      \multirow{5}{*}{log-concave}               & \citep{liang2022proximal}        & ${\mathcal O} \left(\frac{( \lambda_\rho^2 \vee L_f ) d^2 }{ {\mathcal E}} \right)$                          & R\'enyi \xspace   &                          \\\cline{2-5}
                                                 & \citep{lehec2021langevin}        & ${\mathcal O} \left(\frac{\lambda_\rho^2  d^2 }{ {\mathcal E}^{4}} \right)$                                  & $W_2$  & Assume $M_\pi= {\mathcal O}({d})$ \\\cline{2-5}
                                                 & \citep{durmus2019analysis}       & ${\mathcal O} \left(\frac{ L_f d^2 + \lambda_\rho^2  d }{ {\mathcal E}^2} \right)$                           & KL     &                         \\\cline{2-5}
                                                 & \citep{durmus2018efficient}      & ${\mathcal O} \left(\frac{  \lambda_\rho^2  d^5 }{ {\mathcal E}^2} \right)$                                  & TV     & Moreau envelope          \\\cline{2-5}
                                                 & Theorem \ref{thm:convex}         & ${\mathcal O} \left(\frac{\lambda_h^2 D M_\pi d }{ {\mathcal E}^{3}} \right)$                                & TV     &                          \\\hline
      log-Sobolev                                & Theorem \ref{thm:lsi}            & $\widetilde {\mathcal O} \left( \frac{\lambda_h^2 D C_\pi d }{{\mathcal E}}   \right)$                       & TV     &                          \\ \hline
    \end{tabular}
  }
\end{table}





\section{Background}\label{sec:back}
Before we proceed, we briefly introduce the notation basics.
We assume the norm in the space of $x$ and $y$ can be different and denote them as $\|\cdot \|_{\mathcal X}$ and $\|\cdot \|_{\mathcal Y}$ respectively.
We use $J_h(x): {\mathbb R}^d \rightarrow {\mathbb R}^{n \times d}  $ to represent the  Jacobian matrix of the function $h$ at $x$.
The norm of a matrix $A: {\mathbb R}^d \rightarrow {\mathbb R}^{n \times d} $ is defined as $\|A\|_{{\mathcal X},{\mathcal Y}} = \max_{x,y} \{
  \< A x, y\>_{\mathcal Y}
  : \|x\|_{\mathcal X} =1, \|y\|_{\mathcal Y} =1 \}.$
\fan{Given a primal space equipped with norm $\|\cdot\|$, we denote the norm in its dual space as $\|\cdot\|^*$.}
In this article, we assume the probability measure have density with respect to the Lebesgue measure.

\subsection{Nesterov smooothing}

It is well-known that with subgradient methods, the optimal bound of non-smooth convex minimization complexity is ${\mathcal O}( 1/\epsilon^2)$ to achieve $\epsilon$ error in the function value~\citep{nemirovskij1983problem,nesterov2003introductory}. However, beyond the black-box oracle model, there are various structures that can be used to improve this bound. Nesterov smoothing~\citep{nesterov2005smooth} is an interesting technique for functions with explicit max structure \eqref{eq:potential}. In doing so, the complexity order can be improved to ${\mathcal O}(1/{\mathcal E})$. The underlying principle of Nesterov smoothing is that
\[ \textit{The dual of a strongly-convex function is smooth and vice versa~\citep{kakade2009duality}.} \]
Specifically, it adds a $\sigma$-strongly-convex penalty function $\ell(y)$ to the original maximization such that
\begin{align}\label{eq:s_beta}
  s_\beta(x) = f(x) +\max_{y \in {\mathcal Y}} \{\<h(x), y\> -g(y) -
  {\beta} \ell(y) \} ,
\end{align}
where $\beta$ is the smoothing intensity. In the original paper of Nesterov smoothing, they consider the specific task $h(x) =Ax$ and find $s_\beta$ is smooth with constant $L_f + \|A\|^2_{{\mathcal X},{\mathcal Y}}/\beta \sigma. $ Then they apply the Nesterov accelerated gradient descent algorithm to minimize $s_\beta$, an optimal scheme for smooth optimization.
Our idea is very similar on the high level: use smoothing for the non-smooth potential first, and then apply the existing sampling algorithm for the smoothed distribution.

\subsection{Smooth sampling}\label{sec:smooth}

In this section, we summarize several state-of-art sampling algorithms for smooth distributions. In general, there are several dominant popular algorithms: Unadjusted Langevin Monte Carlo (LMC)~\citep{parisi1981correlation}, Underdamped/Kinetic Langevin Monte Carlo (KLMC)~\citep{cheng2018underdamped} Metropolis-Adjusted Langevin Algorithm (MALA)~\citep{roberts1996exponential}, and Hamiltonian Monte Carlo (HMC)~\citep{neal2011mcmc}.
Suppose we are given the target distribution $\pi \propto \exp(-V(x))$ for some smooth potential $V(x)$.
The LMC algorithm runs the update
\begin{align} \label{eq:lmc}
  x_{t+1} = x_{t} - \gamma \nabla V(x) + \sqrt{2\gamma} {\epsilon}_t,
\end{align}
where $\epsilon_t$ is the standard Gaussian random variable in ${\mathbb R}^d$, $x_0$ is a user-defined starting point and $\gamma>0$ is the step size that is small enough. Indeed, \eqref{eq:lmc} is the Euler discretization of Langevin diffusion dynamics,
\begin{align}\label{eq:langevin}
  {\text{d}} x_t = - \nabla V(x_t) {\text{d}} t + \sqrt{2} {\text{d}} B_t.
\end{align}
This dynamics has an invariant density $\pi$ no matter what $x_0$ is. The underdamped Langevin diffusion adds an Hamiltonian ingredient to the standard Langevin diffusion \eqref{eq:langevin}, so that an additional random variable $v_t := {\text{d}} x_t /{\text{d}} t$ that controls the velocity in the dynamics. Similarly, HMC maintains a velocity term, but with a different randomization manner.
MALA is also a variant of LMC. As its name implies, it adds a Metropolis-Hastings~\citep{metropolis1953equation,hastings1970monte} rejection step to LMC, in order to eliminate the bias induced by discretization.
We refer the readers to \citet{mackay2003information} for the an overview of these algorithms.

In spite of the canonical role of LMC in sampling, it is not the fastest algorithm empirically. Theoretically, in the standard strongly-convex potential case, the complexity of LMC can achieve $\widetilde {\mathcal O}(d)$~\citep{dalalyan2017theoretical,vempala2019rapid,erdogdu2022convergence}, whereas KLMC, HMC algorithms can achieve $\widetilde {\mathcal O}(d^{1/2})$~\citep{dalalyan2020sampling,chen2022optimal}. \fan{The complexity $\widetilde {\mathcal O}(d^{1/2})$ also holds for MALA, however, under a warm start~\citep{wu2021minimax}.
 
  Regarding the accuracy dependency,
  MALA is notably the best  with a benefit from Metropolis step.} With a considerable research accumulated for decades, several novel sampling algorithms or the variants are proposed recently, such as the proximal algorithm~\citep{lee2021structured,chen2022improved,liang2022proximal} based on restricted Gaussian oracle, and the randomized midpoint method for KLMC~\citep{shen2019randomized}.  The proximal algorithm  has been proved to have the state-of-art dimension dependency for log-concave distribution~\citep{chen2022improved}.



There exist some accelerating algorithms that can achieve even lower dimension order in complexity~\citep{mangoubi2018dimensionally,sanz2021wasserstein}. However, they require stronger regularity assumptions, e.g. bounded third-order derivative. Although we do not consider these higher-order algorithms in this paper for the sake of generality, one can still apply them in our framework if the higher-order smoothness of $s_\beta(x)$ can be verified in practice.





\section{Nesterov smoothing for sampling}\label{sec:algo}
In this section, we will introduce our Nesterov smoothing sampling algorithm in Section \ref{sec:setup} and bound the error caused by smoothing in Section \ref{sec:error}.
\subsection{Problem setup and algorithm}\label{sec:setup}


Formally, we consider the task to sample from a non-smooth distribution $ \pi(x) \propto \exp(-s(x))$ where
\begin{align}
  s(x) = f(x) +\max_{y \in {\mathcal Y}} \{\<h(x), y\> -g(y) \}.
\end{align}
We will assume $f(x):{\mathcal X} \rightarrow {\mathbb R}$ is $L_f$-smooth; $h(x): {\mathcal X} \rightarrow {\mathcal Y} $ is $\lambda_h$-Lipschitz and $L_h$-smooth; $g(y): {\mathcal Y} \rightarrow {\mathbb R}$ is continuous and convex through out the paper.
A function $f(x)
 
$ is $L$-smooth if \fan{$\|\nabla f(x_1) -\nabla f(x_2)\|^*_{\mathcal X} \leq L \|x_1-x_2\|_{\mathcal X} $} ~for~ $\forall x_1,x_2 \in {\mathbb R}^d.$
Correspondingly, we say a multivariate function
$h(x)
 
$ is $L$-smooth if $\|J_h(x_1) - J_h(x_2) \|_{{\mathcal X},{\mathcal Y}} \le L \|x_1 - x_2\|_{\mathcal X} $ ~for~ $\forall x_1,x_2 \in {\mathcal X}.$
We say a multivariate function $h(x)$ is $\lambda$-Lipschitz if
\fan{$\|h(x_1) -h(x_2) \|^*_{\mathcal Y} \le \lambda \|x_1 -x_2\|_{\mathcal X}$, which is equivalent to $\| J_h(x)  \|_{{\mathcal X},{\mathcal Y}} \le \lambda $ ~for~ $\forall x \in {\mathcal X}$ if $h(x)$ is differentiable.
}
This boils down to that the
\fan{operator norm}
of $A$ is bounded when $h(x) =Ax$.
And this allows us to compare the smoothing effect with the case considered in \citet{nesterov2005smooth}. \fan{Without more specification, in the rest of paper, we assume $\|\cdot\|^*_{\mathcal X} =\|\cdot\|_{\mathcal X} =\|\cdot\|_2 $ is Euclidean norm to be consistent with the sampling literature.}

It is very likely that $s(x)$ is not differentiable since it is a "piece-wise" function segmented by different maximizer $y$. As such, it is not feasible to directly apply the smooth sampling algorithms.
Inspired by the Nesterov smoothing technique, our strategy is to
\[ \textit{1) smooth $s(x)$ as $s_\beta(x)$, 2) apply smooth sampling algorithm to $\pi_\beta(x)$.} \]

\begin{algorithm}[tb]
  \caption{ Nesterov smoothing sampling framework}
  \label{algo}
  \begin{algorithmic}
    \STATE{{\bfseries Input:}
      A first-order based smooth sampling algorithm,
     
      the target distribution $\pi \propto \exp(-s(x))$,
      number of iterations $K$,  initial point $x_0 \sim \mu_0$
     
    }
    \STATE{{\bfseries Initialization:}
      Choose the smoothing intensity $\beta$
     
    }
    \FOR{ $k=0, \ldots, K-1$}
    \STATE{\it{\# Calculate $\nabla s_\beta(x)  $ according to \eqref{eq:grad_s_beta}} }
    \STATE{Iterate $x_{k+1} $ by the smooth sampling algorithm
    }
   
    \ENDFOR
    \STATE {{\bfseries Output:}  $x_K $ }
  \end{algorithmic}
\end{algorithm}


For smoothing, we consider the surrogate function
$$s_\beta(x) = f(x) +\max_{y \in {\mathcal Y}} \{\<h(x), y\> -g(y) -
  {\beta} \ell(y) \} ,
 
$$ where $\ell(y)$ is a $\sigma$-strongly-convex prox-function
$$\ell(y) \ge
 
  \frac{\sigma}{2}  \| y-y_0\|_{\mathcal Y}^2 .$$
The center of ${\mathcal Y}$ is defined as $y_0 = \argmin_y \{ \ell (y): y \in {\mathcal Y}\} $ and we can assume that $\ell(y_0) =0.$
The compactness of ${\mathcal Y}$ also allows us to
define
\begin{align}\label{eq:bound_l}
  D = \max_y\{\ell(y): y \in {\mathcal Y} \}
\end{align}
and $R = \max_y\{\|y\|_{\mathcal Y} : y \in {\mathcal Y} \} $. From the definition, $D$ is the squared diameter of the set ${\mathcal Y}$, and the $R$ is the furthest distance of any point in the set to the origin. By the triangular inequality, it holds that  $R \le \|y_0\| + \sqrt{\frac{ 2D}{\sigma}}$.
In this way, $s_\beta(x)$ is a smooth function, which is justified by Lemma \ref{lem:phi_smooth}.
The proof is postponed to Appendix.
\begin{lemma}\label{lem:phi_smooth}
  Let $\phi_\beta(x,y): = \<h(x), y\> -g(y) -
    {\beta} \ell(y) $, and $y_\beta(x) := \argmax_{y \in {\mathcal Y}} \phi_\beta(x,y)
   
  $, then
  \begin{align}\label{eq:grad_s_beta}
    \nabla s_\beta(x) = \nabla f(x) +J_h^\top (x) y_\beta(x).
  \end{align}
  Moreoever,
  $s_\beta(x)$ is $L_{s_\beta}$-smooth with
  \begin{align}\label{eq:L_phi}
    L_{s_\beta}= L_f + R L_h + \frac{\lambda_h^2 }{\beta \sigma}.
  \end{align}
\end{lemma}
If $h(x) =Ax$, then $L_h=0$ and $\lambda_h = \|A\|_{{\mathcal X},{\mathcal Y}}$, we get $L_{s_\beta} = L_f + \frac{1 }{\beta \sigma} \|A\|^2_{{\mathcal X},{\mathcal Y}} $ in this special case. This recovers the smoothness constant in \citet[Theorem 1]{nesterov2005smooth}.
In many cases, $\lambda_h$ is dimensional-free. For example, when $h(x) =Ax $ and $A = c I_m$, then $\lambda_h=c$ regardless of the dimensionality.

We will also assume that
we can solve $ \max_y \phi_\beta(x,y)$ for any $x \in {\mathcal X}$
in a fast dimension-free manner, which has been verified by the complexity analysis~\citep[Section 3]{bubeck2015convex} and especially several examples in \citet[Section 4]{nesterov2005smooth}.
Thus we can query the gradient of $s_\beta(x)$ according to \eqref{eq:grad_s_beta} and
resort to first-order sampling algorithm to sample from $\pi_\beta \propto  \exp(- s_\beta (x))$.  The pseudo-code of our scheme is presented in Algorithm \ref{algo}.


\subsection{Bounded distance between $\pi$ and $\pi_\beta$} \label{sec:error}


When $\beta \rightarrow 0,$ it holds that $\pi_\beta \rightarrow \pi $.
In this section, we rigorously bound the distance between $\pi$ and $\pi_\beta$. A key observation is that $s_\beta(x)$ is a uniform smooth approximation of $s(x)$, i.e., $ s_\beta(x) \le s(x) \le s_\beta(x)+ \beta D$ (see \eqref{eq:s_diff})~\citep[Eq. 2.7]{nesterov2005smooth}.
To bound the difference between $\pi$ and $\pi_\beta$, we select two widely-used distributional distance as the error criteria.
The first one is the total variation  ${\textup{TV}} (p,q) = \frac{1}{2} \| p-q\|_1 $, which can be bounded by Kullback–Leibler divergence ${\textup{KL}} (p \|q)$ through Pinsker's inequality
\begin{align}\label{eq:pinsker}
  {\textup{TV}} (p,q) \le \sqrt{\frac{1}{2} {\textup{KL}}(p \| q)}.
\end{align}
The second one is the Wasserstein-2 distance~\citep{villani2021topics}
$$ W_2 (p,q) := \sqrt{\inf_{\nu \in \Pi (p,q)} \int_{{\mathbb R}^d \times {\mathbb R}^d } \|a-b\|_2^2 d\nu(a,b) }, $$
where $\Pi (p,q) $ denotes the set of the all joint distributions of $p$ and $q$.
The $W_2$ distance can also be bounded by ${\textup{KL}}(p \| q)$ through
the Talagrand
inequality~\citep{otto2000generalization}
\begin{align} \label{eq:tala}
  \tag{Talagrand}
 
  W_2(p,q) \le \sqrt{2 C_{\textup{LSI}} {\textup{KL}}( p\|q)}.
 
\end{align}
if $q$ satisfies the following logarithmic-Sobolev inequality (LSI) with constant $C_{\textup{LSI}}$.
\begin{definition}[Log-Sobolev inequality] \label{def:lsi}
  A probability distribution $p$ satisfies a logarithmic-Sobolev inequality with constant $C_{\textup{LSI}}$ if, for all smooth functions $u: {\mathbb R}^d \rightarrow {\mathbb R}$, it holds that
  \begin{align}\tag{LSI} \label{eq:lsi}
    {\mathbb E}_p[u^2 \log{u^2} ] -{\mathbb E}_p[u^2] \log {\mathbb E}_p[u^2] \le {2}{C_{\textup{LSI}}} {\mathbb E}_p [ \|\nabla u\|^2].
  \end{align}
\end{definition}
LSI is a powerful tool for the sampling convergence analysis when there is no log-concavity. It has nice properties, for instance, it implies the sub-Gaussian concentration property~\citep{ledoux1999concentration}. Moreoever, if $s(x)$ is $\alpha$-strongly-convex, then  $\pi(x)$ satisfies \eqref{eq:lsi} with constant $1/\alpha$~\citep{bakry1985diffusions}. With the help of the above inequalities, we are able to show that Proposition \ref{prop:tv}, \ref{prop:w} hold.

\begin{proposition}\label{prop:tv}
  ${\textup{TV}}(\pi , \pi_\beta) \le \frac{\beta D}{2}$.
\end{proposition}


\begin{proposition}\label{prop:w}If $\pi$ satisfies \eqref{eq:lsi} with constant $C_\pi$, then
  $W_2(\pi , \pi_\beta) \le \sqrt{C_\pi} {\beta D} $.
\end{proposition}

We next compare the smoothing effect of Nesterov smoothing, Moreau envelope and Gaussian smoothing. 
Recall that we define the non-smooth part in our potential as $\rho(x)$, and Lipschitz constant of $\rho(x)$ to be $\lambda_\rho.$
All of three smoothing methods can deal with composite smooth$+$non-smooth potential, but
\emph{they require the convexity or weakly-convexity of $\rho(x)$,
while Nesterov smoothing does not require so.}
\paragraph{Moreau envelope}
The Moreau-Yosida envelope of a l.s.c. convex function $\rho(x)$ is defined as 
\begin{align*}
\rho_\beta(x) := \min_{y \in {\mathbb R}^d} \left\{ \rho(y) + \frac{1}{2\beta} \| x-y\|_2^2 \right\}.
\end{align*}
$\rho_\beta$ is $L_\rho$-smooth with
$
    L_\rho = \frac{1}{\beta}.
$
The smoothed composite potential is then written as $s_\beta(x) = f(x) + \rho_\beta(x).$
The smoothing effect of Moreau envelope applies to not only l.s.c. convex functions, but also l.s.c. weakly-convex functions with an upper bound constraint on $\beta$~\citep{bohm2021variable,hoheiselproximal}. 
According to Proposition 1 in \citet{durmus2018efficient}, 
there is
$
    {\textup{TV}} (\pi, \pi_\beta) \le \frac{1}{2} \beta \lambda_\rho^2,
$
which is implied by $0 \le s(x) -s_\beta(x) \le {\beta  \|\lambda_\rho\|^2}/{2} $.

\paragraph{Gaussian smoothing}
Unlike the proposed method in \citet{chatterji2020langevin}, where they apply Gaussian smoothing to both the smooth part $f(x)$ and the non-smooth part $\rho(x)$ in $s(x)$, our discussion is based on only smoothing the $\rho_\beta(x)$  here. We find smoothing both parts would introduce a worse dimension dependency in bounding the difference between $\pi(x)$ and $\pi_\beta(x)$.
 The Gaussian smoothing for $\rho(x)$ is defined as 
 $\rho_\beta (x) =  {\mathbb E}_\xi[ \rho ( x + \beta \xi )] ,$ where $\xi \sim \mathcal{N}(0, I_{d\times d}). $
The smoothing effect of Gaussian smoothing only applies to convex functions. 
 Indeed, according to Lemma 2.2 in~\citet{chatterji2020langevin}, if $\rho(x)$ is convex, then
 $ \rho_\beta (x)$ is $L_{ \rho}$-smooth with
$
   L_{\rho} =  {\lambda_\rho d^{1/2}}/{\beta}.$
To bound the difference between $\pi $ and $\pi_\beta  $, one doesn't need a stronger assumption.
By Lemma E.2 in \citet{chatterji2020langevin}, 
$
    {\textup{TV}} (\pi, \pi_\beta) \le \frac{1}{2} \beta \lambda_\rho d^{1/2} .
$
Additionally, if $f(x)$ is $\alpha$-strongly convex, by Lemma 3.3 in \citet{chatterji2020langevin},
$
    W_2 (\pi, \pi_\beta) =
    \widetilde{\mathcal{O}}\left(
   
    \frac{ \beta \lambda_\rho d}{\alpha}
    \right).
$

Now, if we want to achieve ${\textup{TV}}(\pi,\pi_\beta) = {\mathcal E}/2$, then Moreau envelope results in
$L_{\rho} = \frac{\lambda_\rho^2}{{\mathcal E}};$
Gaussian smoothing results in
$L_{\rho} = \frac{\lambda_\rho^2 d}{{\mathcal E}};$
Nesterov smoothing results in
$L_{\rho} =  R L_h + \frac{\lambda_h^2 D }{\mathcal{E} \sigma}.$
Given the same smooth sampling algorithm, the dimension dependency in $L_\rho$ determines the dimension dependency in the final convergence rate.
Normally, $\sigma=1$. Assuming $\lambda_\rho = \lambda_h , D =\widetilde{\mathcal{O}}(1) $, then Nesterov smoothing effect is the same order as Moreau envelope, and better than Gaussian smoothing. If $D$ hides dimension dependency, then Nesterov smoothing is worse than Moreau envelope. However, Moreau envelope requires more strict assumption on $\rho(x)$, i.e. weakly-convexity; and since the potential is not smooth, solving the proximal operator of Moreau envelope requires additional complexity in each sampling step, even for structured potential in \eqref{eq:potential}.

\section{Non-asymptotic convergence analysis}\label{sec:complexity}
In this section, we present our convergence analysis for several fundamental scenarios:
strongly log-concavity, log-concavity, and log-Sobolev inequality.
Denote the underlying distribution of $x_{k}$ as $\mu_{k}$. We will give the sufficient number of gradient queries for $\mu_k$ to reach a given error tolerance. As discussed in Section \ref{sec:smooth}, there are multiple algorithm candidates to use for the smooth sampling in Algorithm \ref{algo}. Our strategy is to test every algorithm among them and select the one with the best complexity. In the main theorems of this section, each complexity is accompanied by a state-of-art smoothing algorithm, which we will specify concretely.




\subsection{Strong log-concavity}\label{sec:strong_concave}
The strong log-concavity condition is mostly-widely studied for sampling convergence since it can be viewed as the strongest assumption.
To our best knowledge, under this condition and without higher-order smoothness,\footnote{Here, higher-order smoothness means the condition that is higher than bounded Hessian.} the randomized midpoint KLMC has the best dimension dependency $\widetilde {\mathcal O}(d^{1/3})$ among the first-order sampling algorithms. Whereas MALA doesn't have the best dependency on dimension, it has best accuracy dependency. So we will adopt these two algorithms as our smooth sampling algorithm.

We also need to point out that, to have such good complexity, MALA needs a warm-start initialization. The initial distribution $\mu_0$ is called $G$-warm if the $\sup_{S \in {\mathcal B}({\mathbb R}^d) } \mu_0(S) / \pi(S) \le G $, where ${\mathcal B}$ means the Borel sigma algebra. If $G$ does not scale with dimension or condition number, then $\mu_0$ is called a warm start. This warm-start may not hold without a delicate initialization. For example, if Gaussian initialization is used, $G={\mathcal O}(\exp(d))$.

\begin{theorem}\label{thm:str_convex}
  Assume   $f$ is $\alpha$-strongly-convex, and $ \< h(\cdot), y\>$ is convex for any $ y \in {\mathcal Y}$.\\
 
  1) In Algorithm \ref{algo}, if we choose the smooth sampling algorithm to be randomized midpoint KLMC~\citep{shen2019randomized} with initial point $x_0 = \argmin s(x)$,
  $\beta =
    \frac{\sqrt{\alpha} {\mathcal E}}{2D}
  $,
  then for any $0<{\mathcal E}< 2\sqrt{\frac{d}{\alpha}}$,
  the iterate $\mu_K$  satisfies $   W_2(\mu_K , \pi) \le {\mathcal E}$ for
  $K = \widetilde{\mathcal O} \left(
    \left( \frac{L_{s_\beta}}{\alpha} \right)^{7/6} ( \frac{2}{{\mathcal E}} \sqrt{\frac{d}{\alpha} } )^{1/3}
    +
    \left( \frac{L_{s_\beta}}{\alpha} \right) ( \frac{2}{{\mathcal E}} \sqrt{\frac{d}{\alpha} } )^{2/3}
    \right)
  $.
 
  \fan{If ${\mathcal E}$ is sufficiently small such that $\frac{L_f + R L_h}{\lambda_h^2 D}=
      {\mathcal O}(1/{\mathcal E})$, }
  we obtain
$$\fan{K= \widetilde{\mathcal O} \left(
    \frac{\lambda_h^{2} D d^{1/3} }{\alpha^{11/6} {\mathcal E}^{5/3}} 
    +    \frac{\lambda_h^{7/3} D^{7/6} d^{1/6} }{\alpha^{23/12} {\mathcal E}^{3/2}}
    \right) .} $$
  Note that $L_f$ and $ L_h$ show up in the coefficient of lower order of $1/ {\mathcal E}$.
  \\
  2) In Algorithm \ref{algo}, if we choose the smooth sampling algorithm to be MALA~\citep{wu2021minimax} with the initial $\mu_0$ being a warm start,
  $\beta =
    \frac{{\mathcal E}}{D}
  $,
  then for any $0<{\mathcal E}< 2$,
  the iterate $\mu_K$ satisfies ${\textup{TV}}(\mu_K , \pi) \le {\mathcal E}$ for some $K = \widetilde{\mathcal O} \left( \frac{L_{s_\beta}}{\alpha} \sqrt{d} \right)
  $.
 
  If ${\mathcal E}$ is sufficiently small
  such that $\frac{L_f + R L_h}{\lambda_h^2 D}=
   
    {\mathcal O}(1/{\mathcal E})$, we get
 
  $$K = \widetilde{\mathcal O} \left( \frac{ \lambda_h^2 D \sqrt{d}}{\alpha {\mathcal E} } \right) . $$
 
 
\end{theorem}


We need the assumption $ \< h(\cdot), y\>$ is convex for any $ y \in {\mathcal Y}$ since it not only promises $s(x)$ is convex, but also $s_\beta(x)$. Thus we can utilize the convergence result for strong log-concave smooth sampling. \fan{However, this assumption does not necessarily restrict $h(x)$ to be linear.} Here, we give several examples where this assumption hold: (i) $h(x) =Ax$, (ii) $h_i(x)$ is convex for $\forall i$, and ${\mathcal Y} \subset {\mathbb R}_{+}^n$, (iii) $h_i(x)$ is concave for $\forall i$, and ${\mathcal Y} \subset {\mathbb R}_{-}^n$, where ${\mathbb R}^n_+$ means the non-negative part of ${\mathbb R}^n$ and ${\mathbb R}^n_-$ means the non-positive part.





\subsection{Log-concavity}


Denote the second moment of $\pi$ is $M_\pi$. Under the log-concavity condition, the smooth sampling algorithm with best dimension dependency~\citep[Corollary 6]{chen2022improved} is the proximal algorithm. Similar to Section \ref{sec:strong_concave}, MALA can achieve the better accuracy order~\citep{chen2020fast} and worse dimension order compared to the proximal algorithm. For the sake of conciseness, we do not list MALA result in this section.
\begin{theorem}\label{thm:convex}
  Assume   $f$ is convex, and $\< h(\cdot), y\>$ is convex for any $ y \in {\mathcal Y}$. In Algorithm \ref{algo}, if we choose the smooth sampling algorithm to be proximal algorithm~\citep{lee2021structured} with
  $\beta =
    \frac{{\mathcal E}}{D}
  $,
  then for any ${\mathcal E}> 0 $,
  the iterate $\mu_K$  satisfies $ {\textup{TV}}(\mu_K , \pi) \le {\mathcal E} $ within  $K = {\mathcal O} \left( \frac{L_{s_\beta} d
       
        W_2^2(\mu_0, \pi_\beta )
      }{{\mathcal E}^2}  \right)
  $ iterations in expectation.
  For simplification, assume that $W_2^2(\mu_0, \pi_\beta ) = {\mathcal O}(M_\pi)$,
 
  and ${\mathcal E}$ is sufficiently small such that $\frac{L_f + R L_h}{\lambda_h^2 D}=
   
    {\mathcal O}(1/{\mathcal E})$,
  then we have
  \begin{align}\label{eq:concave_K}
    K= {\mathcal O} \left(\frac{\lambda_h^2 D M_\pi d }{ {\mathcal E}^{3}} \right)  .
  \end{align}
 
\end{theorem}

One may notice that an additional term $M_\pi$ shows up in Theorem \ref{thm:convex}. This is because in Theorem \ref{thm:str_convex}, this term is absorbed into logarithm. Generally, $M_\pi$ would scale with dimension $d$. For example $M_\pi = {\mathcal O}(d)$ for a product measure. One needs to take this additional dimension dependency into account when applying Theorem \ref{thm:convex}.


The randomness of $K$ in \eqref{eq:concave_K} comes from the rejection sampling step in proximal algorithm~\citet{lee2021structured}.
The bound is in expectation but one can also derive a high probability bound by  Multiplicative Chernoff bound since each iteration in rejection sampling is independent.


\subsection{Log-Sobolev}

A notable condition when there is no log-concavity is the LSI condition. For example, the double-well potential $s(x) = ax^4 -bx^2$ would satisfy the LSI condition but it is not convex.
For the LSI condition, the convergence analysis still concentrate on LMC~\citep{vempala2019rapid,chewi2021analysis}.
The recent study on proximal algorithm also has the convergence guarantee~\citep{chen2022improved}. After the comparison on their convergence, no one is strictly better than another. So we list the complexity for both algorithms here.

In this part, we will use the second order R\'enyi \xspace divergence, defined as ${\mathcal R}_2 (p \| q) : = \log \left\|\frac{{\text{d}} p}{{\text{d}} q} \right\|^2_{L^2(q)} $, to control the initialization error. It is a strong divergence since it can upper bound both ${\textup{KL}}$ divergence, and squared total variation. It has been a popular analysis tool in sampling~\citep{chewi2021analysis}.
\begin{lemma}\label{lem:q_lsi}
  If $\pi(x)$ satisfies \eqref{eq:lsi} with constant $C_\pi$, then
 
 
  $\pi_\beta$ satisfies \eqref{eq:lsi} with constant
 
  $C_{\beta} =C_\pi \exp \left(
   
    4 \beta D
    \right) $.
 
\end{lemma}

\begin{remark}
  \fan{Although $C_\beta$ has exponential dependence on $D$, we will choose $\beta = \frac{{\mathcal E}}{D}$ in the algorithm such that $\exp (4 \beta D) = \exp(4{\mathcal E}) \approx 1 $. In this manner, $C_\beta$ has the same order of $C_\pi$.}
\end{remark}

With Lemma \ref{lem:q_lsi}, we are able to apply the convergence of smooth sampling algorithms in our scheme and reach Theorem \ref{thm:lsi}.


\begin{theorem}\label{thm:lsi}
  Assume $\pi(x)$ satisfies \eqref{eq:lsi} with constant $C_\pi$.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
  In Algorithm \ref{algo}, if we choose the smooth sampling algorithm to be proximal algorithm~\citep{lee2021structured},
  $\beta =
    \frac{ {\mathcal E}}{D}
  $,
  then for any ${\mathcal E} > 0 $,
  the iterate $\mu_K$ satisfies ${\textup{TV}}(\mu_K , \pi) \le {\mathcal E} $ within $K = {\mathcal O} \left( L_{s_\beta} C_\beta d
   
   
   
    \log(2/{\mathcal E}^2)
    \right)
  $ iterations in expectation.
  Assume
 
 
  \fan{${\mathcal E}$ is sufficiently small such that $\frac{L_f + R L_h}{\lambda_h^2 D}=
      {\mathcal O}(1/{\mathcal E})$, }
 
 
  we obtain
  \begin{align}\label{eq:lsi_prox}
    K= \widetilde {\mathcal O} \left( \frac{\lambda_h^2D C_\pi d
     
    }{
     
      {\mathcal E}}
   
    \right).
  \end{align}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
\end{theorem}


\subsection{Example}\label{sec:ex}
In this section, we present case study to illustrate our algorithm.
We consider the case where $s(x) $ is the summation of a quadratic term and the maximum of absolute values
(also see Section 4.4 of \citep{nesterov2005smooth})
\begin{align}\label{eq:ex}
  s(x) = \|x\|_{\mathcal X}^2 +  \max_{1 \le j \le m } \{ |\< a_j , x\> -b^{(j)} | \}, \quad a_j \in {\mathbb R}^d, ~b^{(j)} \in {\mathbb R},
\end{align}
where we choose $\|\cdot\|_{\mathcal X}$ as Euclidean norm $\|\cdot\|_2$.
It is not difficult to find that $\lambda_\rho = \max_{1\le j \le m} \|a_j\|_2$
since it is a piece-wise affine function.
To write \eqref{eq:ex} in the form of \eqref{eq:potential}, we denote the matrix $A  $ with rows $a_j, j=1, \ldots, m$, and $b = (b_1, \ldots, b_m)$. We also choose $n=2m,$ ${\mathcal Y} = \Delta_{n}$, where $\Delta_{n}$ is the probability simplex on ${\mathbb R}^n$. Then with $\widehat A =
  \begin{pmatrix}
    A \\ -A
  \end{pmatrix}
$
and
$\widehat b =
  \begin{pmatrix}
    b \\ -b
  \end{pmatrix},
$
\begin{align}
  s(x) = \|x\|_2^2 +  \max_{y \in \Delta_{n} } \{ \< \widehat A x, y\> - \< \widehat b, y\> \}.
\end{align}
\chen{This potential is associated with a Bayesian inference problem with a Gaussian prior and likelihood consists of (absolute values of) linear functions. The goal is to infer a posterior distribution that would optimize the worst-case performance.}

To construct the smoothed $s_\beta(x)$, we choose $\| \cdot\|_{\mathcal Y} $ as 1-norm $ \|\cdot\|_1
 
$, ~$y_0$ is the uniform distribution,
and $\ell(y) = \log (n) + \sum_{j=1}^n y^{(j)} \log y^{(j)} $. Thus by Lemma 3 in \citet{nesterov2005smooth}, we result in $\sigma=1$ and $D = \log(n)$.
Clearly, $s(x)$ is strongly-convex, and we have strongly-convex constant $\alpha=2 $, smoothness $L_f= 2, L_h=0$, and Lipschitz
$$\lambda_h
  = \|\widehat A \|_{{\mathcal X},{\mathcal Y}}
 
 
 
  = \max_{x} \{
  \max_{1\le j \le m} |\< a_j, x\>|
  : \|x\|_2 =1 \}
  = \max_{1\le j \le m} \|a_j\|_2  =\lambda_\rho
$$
Putting these together in Theorem \ref{thm:str_convex} bound, with \fan{randomized midpoint KLMC,}
it is enough to have $\widetilde {\mathcal O} \left(
  \frac{
    \lambda_h^{7/3}
    d^{1/3} }{ {\mathcal E}^{5/3}} \right)$, however, the best bound up to now for strongly-convex non-smooth sampling is $\widetilde{\mathcal O} \left( {
   
    \lambda_h^2 d } \right)$ \citep{liang2022proximal}. \fan{If we assume $\lambda_h$ is dimensional free,} we have better convergence result in terms of dimension.









\section{Conclusion}
We consider sampling from a non-smooth distribution $\pi(x) \propto \exp(-s(x))$ that permits the explicit max structure~\eqref{eq:potential}. Inspired by Nesterov smoothing, we firstly transfer the non-smooth potential $s(x)$ to a smooth potential $s_\beta(x)$ by adding a strongly convex penalty in the max structure, then we apply the smooth sampling algorithms directly to $\pi_\beta(x) \propto \exp(- s_\beta(x)) $.
Our smoothing technique is compatible with any first-order sampling algorithm for smooth potentials.
We also show that the error introduced by smoothing can be controlled by a careful choice of smoothing intensity $\beta$. Finally, we systematically study the convergence rate under several common conditions, and compare them with the state-of-the-art results in Table \ref{table}. Our convergence rates are better than or on par with existing methods \fan{if no parameter is hiding dimension dependency.}


\paragraph{Limitation and future work}
In this paper, we only consider one class of non-smooth potential with max structure~\eqref{eq:potential} for smoothing.
It is interesting to explore more structures that can benefit from smoothing techniques.
In addition, from Table \ref{table}, one can see that our accuracy dependency does not achieve state-of-art, since we need to compensate the error caused by smoothing.
\fan{If some parameter like $D$ is hiding dimension dependency, then our method may not be state-of-art anymore.}
Finally, in the Nesterov smoothing work~\citep{nesterov2005smooth}, the optimal estimate for smooth convex optimization is proved.
However, the lower bound of smooth sampling is still unclear and seems more challenging. Only a specific class of sampling algorithms has been shown to have lower bound complexity~\citep{wu2021minimax}. It would also be pivotal to derive the lower bound of smooth sampling.

\newpage

2022-08-17T02:03:41.000
2208
2208.07459
en
https://arxiv.org/abs/2208.07459
We study the problem of sampling from a target distribution in $\mathbb{R}^d$ whose potential is not smooth. Compared with the sampling problem with smooth potentials, this problem is much less well-understood due to the lack of smoothness. In this paper, we propose a novel sampling algorithm for a class of non-smooth potentials by first approximating them by smooth potentials using a technique that is akin to Nesterov smoothing. We then utilize sampling algorithms on the smooth potentials to generate approximate samples from the original non-smooth potentials. We select an appropriate smoothing intensity to ensure that the distance between the smoothed and un-smoothed distributions is minimal, thereby guaranteeing the algorithm's accuracy. Hence we obtain non-asymptotic convergence results based on existing analysis of smooth sampling. We verify our convergence result on a synthetic example and apply our method to improve the worst-case performance of Bayesian inference on a real-world example.
Computation (stat.CO)
Nesterov smoothing for sampling without smoothness
Qwen/Qwen-72B
1. YES
2. YES