A Langevin-like Sampler for Discrete Distributions

Introduction

Discrete variables are ubiquitous in machine learning problems ranging from discrete data such as text (Wang & Cho 2019; Gu et al. 2018) and genome (Wang et al. 2010), to discrete models such as low-precision neural networks (Courbariaux et al. 2016; Peters & Welling 2018). As data and models become large-scale and complicated, there is an urgent need for efficient sampling algorithms on complex high-dimensional discrete distributions.

Markov Chain Monte Carlo (MCMC) methods are typically used to perform sampling, of which the efficiency is largely affected by the proposal distribution (Brooks et al. 2011). For general discrete distributions, Gibbs sampling is broadly applied, which resamples a variable from its conditional distribution with the remaining variables fixed. Recently, gradient information has been incorporated in the proposal of Gibbs sampling, leading to a substantial boost to the convergence speed in discrete spaces (Grathwohl et al. 2021). However, Gibbs-like proposals often suffer from high-dimensional and highly correlated distributions due to conducting a small update per step. In contrast, proposals in continuous spaces that leverage gradients can usually make large effective moves. One of the most popular methods is the Langevin algorithm (Grenander & Miller 1994; Roberts & Tweedie 1996; Roberts & Stramer 2002), which drives the sampler towards high probability regions following a Langevin diffusion. Due to its simplicity and efficiency, the Langevin algorithm has been widely used for sampling from complicated high-dimensional continuous distributions in machine learning and deep learning tasks (Welling & Teh 2011; Li et al. 2016; Grathwohl et al. 2019; Song & Ermon 2019). Its great success makes us ask: what is the simplest and most natural analogue of the Langevin algorithm in discrete domains?

In this paper, we develop such a Langevin-like proposal for discrete distributions, which can update many coordinates of the variable based on one gradient computation. By reforming the proposal from the standard Langevin algorithm, we find that it can be easily adapted to discrete spaces and can be cheaply computed in parallel due to coordinatewise factorization. We call this proposal discrete Langevin proposal (DLP). Inheriting from the Langevin algorithm, DLP is able to update all coordinates in a single step in parallel and the magnitude of changes is controlled by a stepsize. Using this proposal, we are able to obtain high-quality samples conveniently on a variety of tasks. We summarize our contributions as the following:

We propose discrete Langevin proposal (DLP), a gradient-based proposal for sampling discrete distributions. DLP is able to update many coordinates in a single step with only one gradient computation.

We theoretically prove the efficiency of DLP by showing that without a Metropolis-Hastings correction, the asymptotic bias of DLP is zero for log-quadratic distributions, and is small for distributions that are close to being log-quadratic.

With DLP, we develop several variants of sampling algorithms, including unadjusted, Metropolis-adjusted, stochastic and preconditioned versions, indicating the general applicability of DLP for different scenarios.

We provide extensive experimental results, including Ising models, restricted Boltzmann machines, deep energy-based models, binary Bayesian neural networks and text generation, to demonstrate the superiority of DLP in general settings.

Related Work

Gibbs sampling is perhaps the de facto method for sampling from general discrete distributions. In each step, it iteratively updates one variable leaving the others unchanged. Updating a block of variables is possible, but typically with an increasing cost along with the increase of the block size. To speed up the convergence of Gibbs sampling in high dimensions, Grathwohl et al. 2021 uses gradient information to choose which coordinate to update and Titsias & Yau 2017 introduces auxiliary variables to trade off the number of updated variables in a block for less computation. However, inheriting from Gibbs sampling, these methods still require a large overhead to make significant changes (e.g. $>5$ coordinates) to the configuration in one step.

Based on the information of a local neighborhood of the current position, locally-balanced proposals have been developed for sampling from discrete distributions (Zanella 2020). Later they have been extended to continuous-time Markov processes (Power & Goldman 2019) and have been tuned via mutual information (Sansone 2021). Similar to Gibbs sampling-based methods, this type of proposals is very expensive to construct when the local neighborhood is large, preventing them from making large moves in discrete spaces. A concurrent work (Sun et al. 2022) explores a larger neighborhood by making a sequence of small movements. However, it still only updates one coordinate per gradient computation and the update has to be done in sequence, while on the contrary, our method can update many coordinates based on one gradient computation in parallel.

Incorporating gradients in the proposal has been a great success in continuous spaces, such as the Langevin algorithm, Hamiltonian Monte Carlo (HMC) (Duane et al. 1987; Neal et al. 2011) and their variants. To take advantage of this success, continuous relaxation is applied which performs sampling in a continuous space by gradient-based methods and then transforms the collected samples to the original discrete space (Pakman & Paninski 2013; Nishimura et al. 2020; Han et al. 2020; Zhou 2020; Jaini et al. 2021; Zhang et al. 2022). The efficiency of continuous relaxation highly depends on the properties of the extended continuous distributions which may be difficult to sample from. As shown in previous work, this type of methods usually does not scale to high dimensional discrete distributions (Grathwohl et al. 2021).

Preliminaries

We consider sampling from a target distribution

where $\theta$ is a $d$-dimensional variable, $\Theta$ is a finite We consider finite discrete distributions in the paper. However, our algorithms can be easily extended to infinite distributions. See Appendix C for a discussion. variable domain, $U$ is the energy function, and $Z$ is the normalizing constant for $\pi$ to be a distribution. In this paper, we restrict to a factorized domain, that is $\Theta=\prod_{i=1}^{d}\Theta_{i}$, and mainly consider $\Theta$ to be $\{0,1\}^{d}$ or $\{0,1,\ldots,S-1\}^{d}$. Additionally, we assume that $U$ can be extended to a differentiable function in $\RR^{d}$. Many popular models have such natural extensions such as Ising models, Potts models, restricted Boltzmann machines, and (deep) energy-based models.

In continuous spaces, one of the most powerful sampling methods is the Langevin algorithm, which follows a Langevin diffusion to update variables:

where $\alpha$ is a stepsize. The gradient helps the sampler to explore high probability regions efficiently. Generally, computing the gradient and sampling a Gaussian variable can be done cheaply in parallel on CPUs and GPUs. As a result, the Langevin algorithm is especially compelling for complex high-dimensional distributions, and has been extensively used in machine learning and deep learning.

Discrete Langevin Proposal

In this section, we propose discrete Langevin proposal (DLP), a simple counterpart of the Langevin algorithm in discrete domains.

At the current position $\theta$, the proposal distribution $q(\cdot|\theta)$ produces the next position to move to. As introduced in Section 3, $q(\cdot|\theta)$ of the Langevin algorithm in continuous spaces can be viewed as a Gaussian distribution with mean $\theta+\alpha/2\nabla U(\theta)$ and covariance $\alpha I_{d\times d}$. Obviously we could not use this Gaussian proposal in discrete spaces. However, we notice that by explicitly indicating the spaces where the normalizing constant is computed over, this proposal is essentially applicable to any kind of spaces. Specifically, we write out the variable domain $\Theta$ explicitly in the proposal distribution,

where the normalizing constant is integrated (continuous) or summed (discrete) over $\Theta$ (we use sum below)

Here, $\Theta$ can be any space without affecting $q$ being a valid proposal. As a special case, when $\Theta=\RR^{d}$, it follows that $Z_{\RR^{d}}(\theta)=(2\pi\alpha)^{d/2}$ and recovers the Gaussian proposal in the standard Langevin algorithm. When $\Theta$ is a discrete space, we naturally obtain a gradient-based proposal for discrete variables.

Computing the sum over the full space in $Z_{\Theta}(\theta)$ is generally very expensive, for example, the cost is $\mathcal{O}(S^{d})$ for $\Theta=\{0,1,\ldots,S-1\}^{d}$. This is why previous methods often restrict their proposals to a small neighborhood. A key feature of the proposal in Equation (1) is that it can be factorized coordinatewisely. To see this, we write Equation (1) as $q(\theta^{\prime}|\theta)=\prod_{i=1}^{d}q_{i}(\theta_{i}^{\prime}|\theta)$, where $q_{i}(\theta_{i}^{\prime}|\theta)$ is a simple categorical distribution of form:

with $\theta^{\prime}_{i}\in\Theta_{i}$ (note that the equation does not contain the term $(\alpha/2\nabla U(\theta)_{i})^{2}$ because it is independent of $\theta^{\prime}$ and will not affect the softmax result). Combining with coordinatewise factorized domain $\Theta$, the above proposal enables us to update each coordinate in parallel after computing the gradient $\nabla U(\theta)$. The cost of gradient computation is also $\mathcal{O}(d)$, therefore, the overall cost of constructing this proposal depends linearly rather than exponentially on $d$. This allows the sampler to explore the full space with the gradient information without paying a prohibitive cost.

We denote the proposal in Equation (2) as Discrete Langevin Proposal (DLP). DLP can be used with or without a Metropolis-Hastings (MH) step (Metropolis et al. 1953; Hastings 1970), which is usually combined with proposals to make the Markov chain reversible. Specifically, after generating the next position $\theta^{\prime}$ from a distribution $q(\cdot|\theta)$, the MH step accepts it with probability

By rejecting some of the proposed positions, the Markov chain is guaranteed to converge asymptotically to the target distribution.

We outline the sampling algorithms using DLP in Algorithm 1. We call DLP without the MH step as discrete unadjusted Langevin algorithm (DULA) and DLP with the MH step as discrete Metropolis-adjusted Langevin algorithm (DMALA). Similar to MALA and ULA in continuous spaces (Grenander & Miller 1994; Roberts & Stramer 2002), DMALA contains two gradient computations and two function evaluations and is guaranteed to converge to the target distribution, while DULA may have asymptotic bias, but only requires one gradient computation, which is especially valuable when performing the MH step is expensive such as in large-scale Bayesian inference (Welling & Teh 2011; Durmus & Moulines 2019).

Zanella 2020 has developed a class of locally-balanced proposals that can be used in both discrete and continuous spaces. One of the locally-balanced proposals is defined as

where $\theta^{\prime}\in\Theta$. DLP can be viewed as a first-order Taylor series approximation to $r(\theta^{\prime}|\theta)$ using

Zanella 2020 discussed the connection between their proposals and Metropolis-adjusted Langevin algorithm (MALA) in continuous spaces but did not explore it in discrete spaces. Grathwohl et al. 2021 uses a similar Taylor series approximation for another locally-balanced proposal,

where $\theta^{\prime}$ belongs to a hamming ball centered at $\theta$ with window size 1. Like Gibbs sampling, their proposal only updates one coordinate per step. They also propose an extension of their method to update $X$ coordinates per step, but with $X$ times gradient computations. See Appendix D for a discussion on Taylor series approximation for Equation (4) without window sizes.

Beyond previous works, we carefully investigate the properties of the Langevin-like proposal in Equation (2) in discrete spaces, providing both convergence analysis and extensive empirical demonstration. We find this simple approach explores the discrete structure surprisingly well, leading to a substantial improvement on a range of tasks.

Convergence Analysis for DULA

In the previous section, we showed that DLP is a convenient gradient-based proposal for discrete distributions. However, the effectiveness of a proposal also depends on how close its underlying stationary distribution is to the target distribution. Because if it is far, even if using the MH step to correct the bias, the acceptance probability will be very low. In this section, we provide an asymptotic convergence analysis for DULA (i.e. the sampler using DLP without the MH step). Specifically, we first prove in Section 5.1 that when the stepsize $\alpha\rightarrow 0$, the asymptotic bias of DULA is zero for log-quadratic distributions, which are defined as

Later in Section 5.2, we extend the result to general distributions where we show the asymptotic bias of DULA is small for distributions that are close to being log-quadratic.

We consider a log-quadratic distribution $\pi(\theta)$ as defined in Equation (5). This type of distributions appears in common tasks such as Ising models. The following theorem summarizes DULA’s asymptotic accuracy for such $\pi$.

If the target distribution $\pi$ is log-quadratic as defined in Equation (5). Then the Markov chain following transition $q(\cdot|\theta)$ in Equation (2) (i.e. DULA) is reversible with respect to some distribution $\pi_{\alpha}$ and $\pi_{\alpha}$ converges weakly to $\pi$ as $\alpha\rightarrow 0$. In particular, let $\lambda_{\text{min}}$ be the smallest eigenvalue of $W$, then for any $\alpha>0$,

where $Z$ is the normalizing constant of $\pi$.

Theorem 5.1 shows that the asymptotic bias of DULA decreases at a $\mathcal{O}(\exp(-1/(2\alpha))$ rate which vanishes to zero as the stepsize $\alpha\rightarrow 0$. This is similar to the case of the Langevin algorithm in continuous spaces, where it converges asymptotically when the stepsize goes to zero.

We empirically verify this theorem in Figure 1a. We run DULA with varying stepsizes on a $2$ by $2$ Ising model. For each stepsize, we run the chain long enough to make sure it converged. The results clearly show that the distance between the stationary distribution of DULA and the target distribution decreases as the stepsize decreases. Moreover, the decreasing speed roughly aligns with a function containing $\exp(-1/(2\alpha))$, which demonstrates the convergence rate with respect to $\alpha$ in Theorem 5.1.

2 Convergence on General Distributions

Then we have the following theorem for the asymptotic bias of DULA on general distributions.

Let $\pi$ be the target distribution and $\pi^{\prime}(\theta)=\exp\left(\theta^{\intercal}W\theta+b\theta\right)/Z^{\prime}$ be the log-quadratic distribution satisfying the assumption in Equation (6), then the stationary distribution of DULA satisfies

where $c_{1}$ is a constant depending on $\pi^{\prime}$ and $\alpha$; $c_{2}$ is a constant depending on $\Theta$ and $\max_{\theta,\theta^{\prime}\in\Theta}\left\|\theta^{\prime}-\theta\right\|_{\infty}$.

The first term in the bound captures the bias induced by the deviation of $\pi$ from being log-quadratic, which decreases in a $\mathcal{O}(\exp(\epsilon))$ rate. The second term captures the bias by using a non-zero stepsize which directly follows from Theorem 5.1. To ensure a satisfying convergence, Theorem 5.2 suggests that we should choose a continuous extension for $U$ of which the gradient is close to a linear function.

In fact, the example of $\pi(\theta)\propto\exp\left(2\epsilon\sin(\theta\pi/2)\right)$ with $\theta\in\{-1,1\}$ can be considered as an counterexample for $any$ existing gradient-based proposals in discrete domains. The gradient on $\{-1,1\}$ is always zero regardless the value of $\epsilon$ whereas the target distribution clearly depends on $\epsilon$. This suggests that we should be careful with the choice of the continuous extension of $U$ since some extensions will not provide useful gradient information to guide the exploration. Our Theorem 5.2 provides a guide about how to choose such an extension.

Other Variants

Thanks to the similarity with the standard Langevin algorithm, discrete Langevin proposal can be extended to different usage scenarios following the rich literature of the standard Langevin algorithm. We briefly discuss two such variants in this section.

Similar to Stochastic gradient Langevin dynamics (SGLD) (Welling & Teh 2011), we can replace the full-batch gradient with an unbiased stochastic estimation $\hat{\nabla}U$ in DLP, which will further reduce the cost of our method on large-scale problems. To show the influence of stochastic estimation, we consider a binary domain for simplicity. We further assume that the stochastic gradient has a bounded variance and the norm of the true gradient and the stochastic gradient are bounded. Then we have the following theorem.

Let $\Theta=\{0,1\}^{d}$. We assume that the true gradient $\nabla U$ and the stochastic gradient $\hat{\nabla}U$ satisfy: $\mathbf{E}[\hat{\nabla}U_{i}]=\nabla U_{i}$, $\mathbf{E}[\hat{\nabla}U_{i}]\leq\sigma^{2}$ for some constant $\sigma$; $|\hat{\nabla}U_{i}|,\left|\nabla U_{i}\right|\leq L$ for some constant $L$. Let $q_{i}$ and $\hat{q}_{i}$ be the discrete Langevin proposal for the coordinate $i$ using the full-batch gradient and the stochastic gradient respectively, then

This suggests that when the variance of the stochastic gradient or the stepsize decreases, the stochastic DLP in expectation will be closer to the full-batch DLP. We test DULA with the stochastic gradient on a binary Bayesian neural network and empirically verify it works well in practice in Appendix I.5.

When $\pi$ alters more quickly in some coordinates than others, a single stepsize may result in slow mixing. Under this situation, a preconditiner that adapts the stepsize for different coordinates can help alleviate this problem. We show that it is easy for DLP to incorporate diagonal preconditioners, as long as the coordinatewise factorization still holds. For example, when the preconditioner is constant, that is, we scale each coordinate by a number $g_{i}$, discrete Langevin proposal becomes

Similar to preconditioners in continuous spaces, $g_{i}$ adjusts the stepsize for different coordinates considering their variation speed. The above proposal is obtained by applying a coordinate transformation to $\theta$ and then transforming the DLP update back to the original space. In this way, the theoretical results in Section 5 directly apply to it. We put more details in Appendix H.

Experiments

We conduct a thorough empirical evaluation of discrete Langevin proposal (DLP), comparing to a range of popular baselines, including Gibbs sampling, Gibbs with Gradient (GWG) (Grathwohl et al. 2021), Hamming ball (HB) (Titsias & Yau 2017)—three Gibbs-based approaches; discrete Stein Variational Gradient Descent (DSVGD) (Han et al. 2020) and relaxed MALA (RMALA) (Grathwohl et al. 2021)—two continuous relaxation methods; and a locally balanced sampler (LB-1) (Zanella 2020) which uses the locally-balanced proposal in Equation (4) with window size $1$. We denote Gibbs-$X$ for Gibbs sampling with a block-size of $X$, GWG-$X$ for GWG with $X$ indices being modified per step (see D.2 in Grathwohl et al. 2021 for more details of GWG-$X$), HB-$X$-$Y$ for HB with a block size of $X$ and a hamming ball size of $Y$. All methods are implemented in Pytorch and we use the official release of code from previous papers when possible. In our implementation, DMALA (i.e. discrete Langevin proposal with an MH step) has a similar cost per step with GWG-$1$ (the main costs for them are one gradient computation and two function evaluations), which is roughly 2.5x of Gibbs-$1$. DULA (i.e. discrete Langevin proposal without an MH step) has a similar cost per step with Gibbs-$1$ (the main cost for DULA is one gradient computation, and for Gibbs-$1$ is $d$ function evaluations). We released the code at https://github.com/ruqizhang/discrete-langevin.

We consider a $5$ by $5$ lattice Ising model with random variable $\theta\in\{-1,1\}^{d}$, and $d=5\times 5=25$. The energy function is

where $W$ is a binary adjacency matrix, $a=0.1$ is the connectivity strength and $b=0.2$ is the bias. We first show that DLP can change many coordinates in one iteration while still maintaining a high acceptance rate in Figure 2 (Left). When the stepsize $\alpha=0.6$, on average DMALA can change 6 coordinates in one iteration with an acceptance rate 52% in the MH step. In comparison, GWG-6 (which at most changes 6 coordinates) only has 43% acceptance rate, not to mention it requires 6x cost of DMALA. This demonstrates that DLP can make large and effective moves in discrete spaces. We compare the root-mean-square error (RMSE) between the estimated mean and the true mean in Figure 3. DMALA is the fastest to converge in terms of both runtime and iterations. This demonstrates the importance of (1) using gradient information to explore the space compared to Gibbs and HB; (2) sampling in the original discrete space compared to DSVGD and RMALA; and (3) changing many coordinates in one step compared to LB-1 and GWG-1. GWG-4 underperforms because of a lower acceptance rate than DMALA. DULA can achieve a similar result as LB-1 and GWG-1 but worse than DMALA, indicating that the MH step accelerates the convergence on this task. In Figure 2 (Right), we compare the effective sample size (ESS) per second for exact samplers (i.e. having the target distribution as its stationary distribution). DMALA significantly outperforms other methods, indicating the correlation among its samples is low due to making significant updates in each step. We additionally present results on Ising models with different connectivity strength $a$ in Appendix I.2.

In what follows, we mainly compare our method with GWG-1 and Gibbs-1, as other methods either could not give reasonable results or are too costly to run.

2 Sampling From Restricted Boltzmann Machines

Restricted Boltzmann Machines (RBMs) are generative artificial neural networks, which learn an unnormalized probability over inputs,

where $\{W,a,b\}$ are parameters and $\theta\in\{0,1\}^{d}$. Following Grathwohl et al. 2021, we train $\{W,a,b\}$ with contrastive divergence (Hinton 2002) on the MNIST dataset for one epoch. We measure the Maximum Mean Discrepancy (MMD) between the obtained samples and those from Block-Gibbs sampling, which utilizes the known structure.

Results The results are shown in Figure 4. Since DULA and DMALA can update multiple coordinates in a single iteration, they converge remarkably faster than baselines in both iterations and runtime. Furthermore, DMALA reaches the lowest MMD ($\approx-6.5$) among all the methods after 5,000 iterations, demonstrating the importance of the MH step. We leave the sampled images in Appendix I.3.

3 Learning Energy-based Models

Though the first term is easy to estimate from the data, the second term requires samples from $p_{\theta}$. Better samplers can improve the training process of $E_{\theta}$, leading to EBMs with higher performance.

As in Grathwohl et al. 2021, we generate a 25 by 25 Ising model and generate training data by running a Gibbs sampler. In this experiment, $E_{\theta}$ is an Ising model with learnable parameter $\hat{W}$. We evaluate the samplers by computing RMSE between the estimated $\hat{W}$ and the true $W$.

Results Our results are summarized in Figure 5. In (a), DMALA and DULA always have smaller RMSE than baselines given the same number of iterations. In (b), DMALA and DULA get a log-RMSE of $-5.0$ in $800$s, while the baseline methods fail to reach $-5.0$ in $1,400$s. In (c), we vary the number of sampling steps per iteration from $5$ to $100$ (we omit the results of Gibbs-1 since it diverges with less than 100 steps) and report the RMSE after 10,000 iterations. DMALA and DULA outperform GWG consistently and the improvement becomes larger when the number of sampling steps becomes smaller, demonstrating the fast mixing of our discrete Langevin proposal.

3.2 Deep EBMs

We train deep EBMs where $E_{\theta}$ is a ResNet (He et al. 2016) with Persistent Contrastive Divergence (Tieleman 2008; Tieleman & Hinton 2009) and a replay buffer (Du & Mordatch 2019) following Grathwohl et al. 2021. We run DMALA and DULA for 40 steps per iteration. After training, we adopt Annealed Importance Sampling (Neal 2001) to estimate the likelihood. The results for GWG and Gibbs are taken from Grathwohl et al. 2021, and for VAE are taken from Tomczak & Welling 2018.

Results In Table 1, we see that DMALA yields the highest log-likelihood among all methods and its generated images in Figure 9 in Appendix I.4 are very close to the true images. DULA runs the same number of steps as DMALA and GWG but only with half of the cost. We hypothesis that running DULA for more steps or with an adaptive stepsize schedule (Song & Ermon 2019) will improve its performance.

4 Binary Bayesian Neural Networks

Bayesian neural networks have been shown to provide strong predictions and uncertainty estimation in deep learning (Hernández-Lobato & Adams 2015; Zhang et al. 2020; Liu et al. 2021a). In the meanwhile, binary neural networks (Courbariaux et al. 2016; Rastegari et al. 2016; Liu et al. 2021b), i.e. the weight is in $\{-1,1\}$, accelerate the learning and significantly reduce computational and memory costs. To combine the benefits of both worlds, we consider training a binary Bayesian neural network with discrete sampling. We conduct regression on four UCI datasets (Dua & Graff 2017), and the energy function is defined as,

where $D=\{x_{i},y_{i}\}_{i=1}^{N}$ is the training dataset, and $f_{\theta}$ represents a two-layer neural network with Tanh activation and 500 hidden neurons. The dimension of the weight $d$ varies on different datasets ranging from $7,500$ to $45,000$. We report the log-likelihood on the training set together with root mean-square-error (RMSE) on the test set.

Results From Table 2, we observe that DMALA and DULA outperform other methods significantly on all datasets except test RMSE on COMPAS, which we hypothesize is because of overfitting (the training set only has $4,900$ data points). These results demonstrate that our methods converge fast for high dimensional distributions, due to the ability to make large moves per iteration, and suggest that our methods are compelling for training low-precision Bayesian neural networks of which the weight is discrete.

5 Text Infilling

Text infilling is an important and intriguing task where the goal is to fill in the blanks given the context (Zhu et al. 2019; Donahue et al. 2020). Prior work has realized it by sampling from a categorical distribution produced by BERT (Devlin et al. 2019; Wang & Cho 2019). However, there are a huge number of word combinations, which makes sampling from the categorical distribution difficult. Therefore, an efficient sampler is needed to producing high-quality text.

We randomly sample 20 sentences from TBC (Zhu et al. 2015) and WiKiText-103 (Merity et al. 2017), mask 25% of the words in the sentence (Zhu et al. 2019; Donahue et al. 2020), and sample 25 sentences from the probability distribution given by BERT. We run all samplers for 50 steps based on two models, BERT-base and BERT-large. As a common practice in non-autoregressive text generation, we select the top-5 sentences with the highest likelihood out of the 25 sentences to avoid low-quality generation (Gu et al. 2018; Zhou et al. 2019). We evaluate the methods from two perspectives, diversity and quality. For diversity, we use self-BLEU (Zhu et al. 2018) and the number of unique n-grams (Wang & Cho 2019) to measure the difference between the generated sentences. For quality, we measure the BLEU score (Papineni et al. 2002) between the generated texts and the original dataset (TBC+WikiText-103) (Wang & Cho 2019; Yu et al. 2017). Note that the BLEU score can only approximately represent the quality of the generation since it cannot handle out-of-distribution generations. We use one-hot vectors to represent categorical variables. We put the discussion of DLP with categorical variables in practice in Appendix B.

Results The quantitative and qualitative results are shown in Table 3 and Figure 6. We find that DULA and DMALA can fill in the blanks with similar quality as Gibbs and GWG but with much higher diversity. Due to the nature of languages, there exist strong correlations among words. It is generally difficult to change one word given the others fixed while still fulfilling the context. However, it is likely to have another combination of words, which are all different from the current ones, to satisfy the infilling. Because of this, the ability to update all coordinates in one step makes our methods especially suitable for this task, as reflected in the evaluation metrics and generated sentences.

Conclusion

We propose a Langevin-like proposal for efficiently sampling from complex high-dimensional discrete distributions. Our method, discrete Langevin proposal (DLP), is able to explore discrete structure effectively based on the gradient information. For different usage scenarios, we have developed several variants with DLP, including unadjusted, Metropolis-adjusted, stochastic, and preconditioned versions. We prove the asymptotic convergence of DLP without the MH step under log-quadratic and general distributions. Empirical results on many different problems demonstrate the superiority of our method over baselines in general settings.

While the Langevin algorithm has achieved great success in continuous spaces, there has always lacked a counterpart of such simple, effective and general-purpose samplers in discrete spaces. We hope our method sheds light on building practical and accurate samplers for discrete distributions.

Acknowledgements

We would like to thank Yingzhen Li and the anonymous reviewers for their thoughtful comments on the manuscript. This research is supported by CAREER-1846421, SenSE2037267, EAGER-2041327, Office of Navy Research, and NSF AI Institute for Foundations of Machine Learning (IFML).

References

Appendix A Algorithm with Binary Variables

When the variable domain $\Theta$ is binary $\{0,1\}^{d}$, we could simplify Algorithm 1 further and obtain Algorithm 2, which clearly shows that our method can be cheaply computed in parallel on CPUs and GPUs.

Appendix B Algorithm with Categorical Variables

When using one-hot vectors to represent categorical variables, our discrete Langevin proposal becomes

where $\theta_{i},\theta^{\prime}_{i}$ are one-hot vectors.

However, if the variables are ordinal with clear ordering information, we can also use integer representation $\theta\in\{0,1,\ldots,S-1\}^{d}$ and compute DLP as in Equation(2).

Appendix C Extension to Infinite Domains

We only consider finite discrete distributions in the paper. However, our algorithms can be easily extended to infinite distributions. One way to do so is to add a window size for the proposal such that the proposal only considers a local region in each step. For example, if the domain is infinite integer-valued: $\Theta=\{0,1,2,\ldots\}^{d}$, then our DLP can use a proposal window of $\theta^{\prime}_{i}\in\{\theta_{i}-1,\theta_{i},\theta_{i}+1\}$. It is also possible to use other window sizes. We leave the comprehensive study of our method on infinite discrete distributions for future work.

Appendix D Importance of the Term ‖θ′−θ‖2/(2​α)\left\|\theta^{\prime}-\theta\right\|^{2}/(2\alpha)

Similar to DLP, we can also use a first-order Taylor series approximation to Equation (4) and get the following proposal

The above proposal does not contain the stepsize term as in DLP, which will result in a very low acceptance probability in practice. For example, in the RBM experiment (Section 7.2), its acceptance probability is $0.005\%$ while DMALA’s acceptance probability is $51\%$ on average. This demonstrates the importance of the term $\left\|\theta^{\prime}-\theta\right\|^{2}/(2\alpha)$ in DLP where the stepsize $\alpha$ can control the closeness of $\theta^{\prime}$ to the current $\theta$. The effect of $\alpha$ is similar to the stepsize in the standard Langevin algorithm, and by tuning it our method can achieve a desirable acceptance probability leading to efficient sampling.

Appendix E Proof of Theorem 5.1

We divide the proof into two parts. In the fisrt part, we will prove the weak convergence and in the second part we will prove the convergence rate with respect to the stepsize $\alpha$.

The main idea of the proof is to replace the gradient term in the proposal by the energy difference $U(\theta^{\prime})-U(\theta)$ using Taylor series approximation, and then show the reversibility of the chain based on the proof of Theorem 1 in Zanella 2020 .

Recall that the target distribution is $\pi(\theta)=\exp\left(\theta^{\intercal}W\theta+b^{\intercal}\theta\right)/Z$. We have that $\nabla U(\theta)=2W^{\intercal}\theta+b$, $\nabla^{2}U(\theta)=2W$. Since $\nabla^{2}U$ is a constant, we can rewrite the proposal distribution as the following

where the last equation is because $U(\theta^{\prime})-U(\theta)=\nabla U(\theta)^{\intercal}(\theta^{\prime}-\theta)+\frac{1}{2}(\theta^{\prime}-\theta)^{\intercal}2W(\theta^{\prime}-\theta)$ by Taylor series approximation.

Let $Z_{\alpha}(\theta)=\sum_{x}\exp\left(\frac{1}{2}\left(U(x)-U(\theta)\right)-(x-\theta)^{\intercal}(\frac{1}{2\alpha}I+\frac{1}{2}W)(x-\theta)\right)$, and $\pi_{\alpha}=\frac{Z_{\alpha}(\theta)\pi(\theta)}{\sum_{x}Z_{\alpha}(x)\pi(x)}$, now we will show that $q_{\alpha}$ is reversible w.r.t. $\pi_{\alpha}$.

It is clear that this expression is symmetric in $\theta$ and $\theta^{\prime}$. Therefore $q_{\alpha}$ is reversible and the stationary distribution is $\pi_{\alpha}$.

Now we will prove that $\pi_{\alpha}$ converges weakly to $\pi$ as $\alpha\rightarrow 0$. Notice that for any $\theta$,

where $\delta_{\theta}(x)$ is a Dirac delta. It follows that $\pi_{\alpha}$ converges pointwisely to $\pi(\theta)$. By Scheffé’s Lemma, we attain that $\pi_{\alpha}$ converges weakly to $\pi$.

Let us consider the convergence rate in terms of the $L_{1}$-norm

Since $\lambda_{\text{min}}(W)\left\|x\right\|^{2}\leq x^{\intercal}Wx,\forall x$, it follows that

We also notice that $\min_{x\neq\theta}\left\|x-\theta\right\|^{2}=1$, thus when $\frac{Z_{\alpha}(\theta)}{\sum_{x}Z_{\alpha}(x)\pi(x)}-1>0$, we get

Similarly, when $\frac{Z_{\alpha}(\theta)}{\sum_{x}Z_{\alpha}(x)\pi(x)}-1<0$, we have,

Therefore, the difference between $\pi_{\alpha}$ and $\pi$ can be bounded as follows

Appendix F Proof of Theorem 5.2

We use a log-quadratic distribution that is close to $\pi$ as an intermediate term to bound the bias of DULA. Recall that $\pi$ is the target distribution, $\pi^{\prime}$ is the log-quadratic distribution that is close to $\pi$ and $\pi_{\alpha}$ is the stationary distribution of DULA. We let $\pi^{\prime}_{\alpha}$ be the stationary distributions of DULA targeting $\pi^{\prime}$, then by triangle inequality,

Let the energy function of $\pi^{\prime}$ be $V(\theta)=\theta^{\intercal}W\theta+b\theta$. Since $\Theta$ is a discrete space, there exists a bounded subset $\Omega\in\RR^{d}$ such that $\Theta$ is a subset of $\Omega$. By Poincaré inequality, we get

where the constant $C_{1}$ depends on $\Omega$.

Recall that $\pi(\theta)=\frac{\exp(U(\theta))}{Z}$. Let $\pi^{\prime}(\theta)=\frac{\exp(V(\theta))}{Z^{\prime}}$ where $Z^{\prime}$ is the normalizing constant to make $\pi^{\prime}$ a distribution. Then

Plugging the above in Equation (7), we obtain

We denote $D=\max_{\theta,\theta^{\prime}\in\Theta}\left\|\theta^{\prime}-\theta\right\|_{\infty}$. By the assumption $\left\|\nabla U(\theta)-\nabla V(\theta)\right\|_{1}\leq\epsilon$, similar to Equation (8), we have

By the perturbation bound in Schweitzer 1968,

where $C_{2}$ is a constant depending on $\pi^{\prime}$ and $\alpha$. Please note that it is also possible to use other perturbation bounds (Cho & Meyer 2001).

We define $c_{1}:=2\max(2,2C_{2})$ and $c_{2}:=\max(2C_{1},D)$, then we reach the the final result

Appendix G Proof of Theorem 6.1

The proposal distribution for the coordinate $i$ with the stochastic gradient is

We consider the binary case i.e. $\Theta_{i}=\{0,1\}$. When the current position $\theta_{i}=0$ the probability for $\theta^{\prime}_{i}$ is

The difference between $\mathbf{E}\left[\hat{q}_{i}\right]$ and $q_{i}$ is

We consider each absolute value term. When $\theta^{\prime}_{i}=0$,

Since the absolute value of the derivative for $f(x)=\frac{1}{\exp(x)+1}$ is $\frac{\exp(x)}{(\exp(x)+1)^{2}}$, which is monotonically increasing for $x\in(-\infty,0]$ and monotonically decreasing for $x\in(0,\infty)$. Then by the assumption $\left|\hat{\nabla}U(\theta)_{i}\right|\leq L$, $\left|\nabla U(\theta)_{i}\right|\leq L$ and Lipschitz continuity,

where the last equation is because of our assumption on the variance of the stochastic gradient $\operatorname{Var}\left[\hat{\nabla}U(\theta)_{i}\right]\leq\sigma^{2}$, which leads to

We substitute the above two bounds to the $L_{1}$ distance

The same analysis applies to the case when the current position $\theta_{i}=1$. To see this, we write out the proposal distribution

The remaining follows the same as in $\theta_{i}=0$.

Appendix H Preconditioned Discrete Langevin Proposal

Our discrete Langevin proposal can be easily combined with diagonal preconditioners as in the standard Langevin proposal. We consider the case when the preconditioner is constant.

Assume that the constant diagonal preconditioner is $G=\text{diag}(g)$ where $g\in\RR^{d}$, then we can derive the preconditioned discrete Langevin proposal by a simple coordinate transformation. Let the new coordinates be $\eta_{i}:=g_{i}^{-1}\theta_{i}$, then the domain for $\eta_{i}$ is $H_{i}=\{g_{i}^{-1}\theta_{i}:\theta_{i}\in\Theta_{i}\}$. The discrete Langevin proposal for $\eta_{i}$ is

Let $\eta^{\prime}_{i}=g_{i}^{-1}\theta^{\prime}_{i},x_{i}=g_{i}^{-1}y_{i}$, then the above proposal distribution is equivalent to

To demonstrate the use of the preconditioned discrete Langevin proposal, we consider an Ising model $U(\theta)=\theta^{\intercal}W\theta$ where $W=\text{diag}(-0.001,-1000)$ and $\theta\in\{-1,1\}$. Due to the drastically different scale for the derivative of different coordinates, if using the same stepsize for both coordinates, no values will work. However if we use the preconditioned proposal as shown above and let $\alpha=1,g_{1}=1000,g_{2}=0.001$. Then the Markov chain will mix quickly. For example, in terms of the log RMSE between the estimated mean and the true mean, preconditioned DLP gives $-6.2$ whereas the standard DLP with any $\alpha$ can only achieve around $-0.34$.

Appendix I Additional Experiments Results and Setting Details

To verify Theorem 5.1, we use a 2 by 2 Ising model $U(\theta)=\theta^{\intercal}cW\theta+b\theta$ where $W$ is the binary adjacency matrix and $c=0.1,b=0.2$. To verify Theorem 5.2, we set $a=1$ and $b=0.1$.

I.2 Sampling from Ising Models

We adopt the experiment setting in Section F.1 of Grathwohl et al. 2021. The energy of the ising model is defined as,

where $a$ controls the connectivity strength and $J$ is an adjacency matrix whose elements are either $0$ or $1$. In our experiments, $J$ is the adjacency matrix of a 2D lattice graph. DMALA and DULA use a stepsize $\alpha$ of $0.4$ and $0.2$ respectively. We show the results with varying connectivity strength $a$. We run all methods for the same amount of time (GWG-1 and DMALA for $50,000$ iterations; Gibbs-1 and DULA for $125,000$ iterations). The results are shown in Figure 7. We can clearly observe that DMALA outperforms other methods on all connectivity strength.

I.3 Sampling from RBMs

The RBM used in our experiments has 500 hidden units and 784 visible units. The model is trained using contrastive divergence (Hinton 2002). We set the batch size to 100. The stepsize $\alpha$ is set to be $0.1$ for DULA and $0.2$ for DMALA, respectively. The model is optimized by an Adam optimizer with a learning rate of $0.001$. The Block-Gibbs sampler is described in details by Grathwohl et al. 2021. Here, we give a brief introduction. RBM defines a distribution over data $x$ and hidden variables $h$ by,

The joint and marginal distribution are both unnormalized and hard to sample from. However, the special structure of the problem leads to easy conditional distributions,

Therefore, we can perform Block-Gibbs sampling by updating $x$ and $h$ alternatively. This Block-Gibbs sampler can update all the $784$ coordinates at the same time, and has much higher efficiency. However, this is due to the special structure of the problem. We provide the generated images in Figure 8, showing that the results of DULA and DMALA are closer to that of Block-Gibbs.

I.4 Learning EBMs

For all the experiments in this section, we set the stepsize $\alpha$ to be $0.1$ for DULA and $0.15$ for DMALA.

We adopt the same ResNet structure and experiment protocol as in Grathwohl et al. 2021, where the network has 8 residual blocks with 64 feature maps. There are 2 convolutional layers for each residual block. The network uses Swish activation function (Ramachandran et al. 2017). For static/dynamic MNIST and Omniglot, we use a replay buffer with $10,000$ samples. For Caltech, we use a replay buffer with $1,000$ samples. We evaluate the models every 5,000 iterations by running AIS for $10,000$ steps. The reported results are from the model which performs the best on the validation set. The final reported numbers are generated by running $300,000$ iterations of AIS. All the models are trained with Adam (Kingma & Ba 2015) with a learning rate of $0.0001$ for $50,000$ iterations.

We show the generated images with DULA and DMALA in Figure 9. We see that the generated images from DMALA are very close to the true images on all four datasets. DULA can generate high-quality images on static and dynamic MNIST, but mediocre images on Omniglot and Caltech Silhouette. We hypothesis that we need to run DULA for more steps per iteration to get better results.

I.5 Binary Bayesian Neural Networks

Details of the Datasets (1) COMPAS: COMPAS (J. Angwin & Kirchner 2016) is a dataset containing the criminal records of 6,172 individuals arrested in Florida. The task is to predict whether the individual will commit a crime again in 2 years. We use 13 attributes for prediction. (2) News: Online News Popularity Data Set https://archive.ics.uci.edu/ml/datasets/online+news+popularity contains 39,797 instances of the statistics on the articles published by a website. The task is to predict the number of clicks in several hours after the articles published. (3) Adult: Adult Income Dataset https://archive.ics.uci.edu/ml/datasets/adult is a dataset containing the information of US individuals from 1994 census. The prediction task is to predict whether an individual makes more than 50K dollars per year. The dataset contains 44,920 data points. (4) Blog: Blog Feedback (Buza 2014) is a dataset containing 54,270 data points from blog posts. The raw HTML-documents of the blog posts were crawled and processed. The prediction task associated with the data is the prediction of the number of comments in the upcoming 24 hours. The feature of the dataset has 276 dimensions.

More Details on Training We run 50 chains in parallel and collect the samples at the end of training. All datasets are randomly partitioned into 80% for training and 20% for testing. The features and the predictive targets are normalized to $(0,1)$. We set $\alpha=0.1$ for all datasets. We use a uniform prior over the weights. We train the Bayesian neural network for $1,000$ steps. We use full-batch training for the results in Section 7 so that Gibbs and GWG are also applicable.

Experimental Results of Stochastic DLP As mentioned in Section 6, we empirically evaluate the performance of stochastic DLP here. We use a batch size of 10, 100, 500 and report the performance of the obtained binary Bayesian neural networks in Table 4. Note that Gibbs-based sampling techniques are not suitable for mini-batches. We find that our method works quite well with mini-batches, and the performance increases as the batch size becomes larger.

I.6 Text Infilling

We use the BERT model pre-trained in PyTorch-Pretrained-BERT https://github.com/maknotavailable/pytorch-pretrained-BERT. We use the same TBC and WikiText-103 corpus as in Wang & Cho 2019 by randomly sampling $5,000$ sentences from the complete TBC and WikiText-103 corpus. BERT is a masked language model, which means that given a sentence with MASK tokens and contexts, the model is able to provide a probability distribution over all the candidate words to indicate which word the MASK position is likely to be. Supposing we have a sentence $X=\left(x_{1},x_{2},\dots,x_{n}\right)$ which has $n$ words, and $m$ of the $n$ words are MASKs. For each MASK[i], the BERT model can output a vector $f(\cdot|X_{\text{MASK}[i]})\in\RR^{N}$ which is a normalized probability distribution over the candidate words representing the probability of a certain word appears at MASK[i] when the context is given. The BERT model we used has $N=30,522$ candidate words. Hence, $\theta\in\{1,2,\dots,N\}^{m}$. Given the user-defined context and MASK positions, we set MASK[i] to $\theta_{i}$, and the energy function is defined as,

There are $N^{m}$ combinations in total, so directly sampling from the categorical distribution is intractable for $m\geq 2$. More generated texts are displayed in Figure 10.