SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Katayoon Goshvadi
Google Deepmind
Xingchao Liu
Azade Nova
Ruqi Zhang
Will Grathwohl
Dale Schuurmans
Hanjun Dai
Sampling in discrete spaces, with critical applications in simulation and optimization, has recently been boosted by significant advances in gradient-based approaches that exploit modern accelerators like GPUs. However, two key challenges are hindering further advancement in research on discrete sampling. First, since there is no consensus on experimental settings and evaluation setups, the empirical results in different research papers are often not comparable. Second, implementing samplers and target distributions often requires a nontrivial amount of effort in terms of calibration and parallelism. To tackle these challenges, we propose DISCS (DISCrete Sampling), a tailored package and benchmark that supports unified and efficient experiment implementation and evaluations for discrete sampling in three types of tasks: sampling from classical graphical models and energy based generative models, and sampling for solving combinatorial optimization. Throughout the comprehensive evaluations in DISCS, we gained new insights into scalability, design principles for proposal distributions, and lessons for adaptive sampling design. DISCS efficiently implements representative discrete samplers in existing research works as baselines and offers a simple interface that researchers can conveniently add new discrete samplers and directly compare their performance with the benchmark result in a calibrated setup.
Sampling in discrete spaces has been an important problem for decades in physics (Edwards & Anderson, 1975; Baumgartner et al., 2012), statistics (Robert & Casella, 2013; Carpenter et al., 2017), and computer science (LeCun et al., 2006; Wang & Cho, 2019). Since sampling from a target probability distribution  $\pi(x) \propto \exp(-f(x))$  defined on a discrete space  $\mathcal{X}$  is typically intractable, one usually resorts to Markov chain Monte Carlo (MCMC) methods (Metropolis et al., 1953; Hastings, 1970). However, except for a few algorithms such as Swendsen-Wang for the Ising model (Swendsen & Wang, 1987) and Hamze-Freitas for hierarchical models (Hamze & de Freitas, 2012), which exploit the special structure of the underlying problem, sampling in a general discrete space has primarily relied on Gibbs sampling, which exhibits notoriously poor efficiency in high dimensional spaces.
Recently, a family of locally balanced MCMC samplers for discrete spaces (Zanella, 2020; Grathwohl et al., 2021; Sun et al., 2021; Zhang et al., 2022), using ratio informed proposal distributions,  $\frac{\pi(y)}{\pi(x)}$ , have significantly improved sampling efficiency by exploiting modern accelerators like GPUs and TPUs. From the perspective of gradient flow on the Wasserstein manifold of distributions, Gibbs sampling is simply a coordinate descent algorithm, whereas locally balanced samplers perform as full gradient descent (Sun et al., 2022a). Despite the advances in locally balanced samplers, a quantitative benchmark is still missing. As a result, the empirical results in different research papers may not be comparable. One important reason is that there is no consensus on the experimental
setting. Particularly, the initialization of energy based generative models, random seeds used in graphical models, and the protocol of hyper-parameter tuning all have a significant impact on samplers performance. Under this circumstance, there is a critical need for a unified benchmark to advance the research in discrete sampling.
There are two key challenges that seriously hinder the appearance of such a benchmark. First, a sampler may perform well in one target distribution while poorly in another one. To thoroughly examine the performance of a sampler, a qualified benchmark needs to collect a set of representative distributions that covers the potential applications of discrete samplers. Second, the evaluation of discrete samplers is complicated. Although the commonly used metric Effective Sample Size (ESS) (Vehtari et al., 2021) can effectively reflect the efficiency of a sampler in Monte Carlo integration or Bayesian inference, it is not very informative in scenarios when the sampler guides the search in combinatorial optimization problems or performs as a decoder in deep generative models.
To address the two challenges, we propose DISCS, a tailored benchmark for discrete sampling. In particular, DISCS consists of three groups of tasks: sampling from classical graphical models, sampling for solving combinatorial optimization problems, and sampling from deep energy based models (EBMs). These tasks cover both the topics of simulation and optimization, and also target distributions, ranging from hand-designed graphical models to learned deep EBMs. For each task, we collect the representative problems from both synthetic and real-world applications, for example, graph partitioning for distributed computing and language model for text generation. We carefully design the evaluation metrics in DISCS. In sampling from classical graphical models tasks, DISCS uses ESS as a standard. In sampling for solving combinatorial optimization tasks, DISCS runs simulated annealing (Kirkpatrick et al., 1983) with multiple chains and reports the average of the best results in each chain. In sampling from energy based generative models, DISCS employs domain specific scores to measure the sample quality.
DISCS offers a convenient interface for researchers to implement new discrete samplers, without worrying about parallelism, experiment loop, and evaluation. DISCS can efficiently sweep over different tasks and experiments configurations in parallel, making it easy to reproduce the benchmark results of this paper. Also, DISCS implements existing discrete samplers including random walk Metropolis (Metropolis et al., 1953), block Gibbs, Hamming Ball sampler (Titsias & Yau, 2017), Locally Balanced (Zanella, 2020), Gibbs with Gradients (Grathwohl et al., 2021), Path Auxiliary Sampler (Sun et al., 2021), Discrete Metropolis Adjusted Langevin Algorithm (Zhang et al., 2022), Discrete Langevin Monte Carlo (Sun et al., 2022a), and is actively maintaining to add new samplers. Researchers can directly compare the results with state-of-the-art methods.
DISCS can also provide insights on the existing open questions in the space of discrete sampling. In our experiments, we observe an interesting phenomenon that the locally balanced weight function  $g(t) = \sqrt{t}$  outperforms  $g(t) = \frac{t}{t + 1}$  when Ising model has a temperature higher than the critical temperature and  $g(t) = \frac{t}{t + 1}$  outperforms  $g(t) = \sqrt{t}$  when the temperature is lower than the critical temperature. We further observe similar phenomenon in our experiments with more complicated deep energy generative model where  $g(t) = \frac{t}{t + 1}$  outperforms  $g(t) = \sqrt{t}$  on overparameterized neural network with low temperature and sharp landscape. There have been numerous studies on how to select the locally balanced function for a locally balanced sampler (Zanella, 2020; Sansone, 2022), but the answer still remains open. We hope the observations in this paper can provide some insight on this question.
We wrap the DISCS package as a JAX library to facilitate the research in discrete sampling. The library is open sourced at DISCS<sup>1</sup>. The dataset used in our benchmark experiments can be accessed at DISCS DATA<sup>2</sup> The paper is organized as follows:
- In section 2, we provide an overview of related work on different tasks for discrete sampling and recent advances in discrete samplers.
- In section 3, we formulate the discrete sampling problem.
- In section 4, we introduce the discrete sampling tasks and evaluation metrics in DISCS. We also present several results with interesting insights and observations.
- In section 5, we discuss the contribution and limitations of DISCS.
$^{1}$ code base found at https://github.com/google-research/discs  
 $^{2}$ data set found at https://drive.google.com/drive/u/1/folders/1nEppxuUJj8bsV9Prc946LN_buo30AnDx
We provide comprehensive results of our benchmark and studies with further details on the experimental setups, evaluation metrics, mathematical formulations, and the data set used in the Appendix A.
Discrete sampling has been widely used to study the physical picture of spin glasses (Hukushima & Nemoto, 1996; Katzgraber et al., 2001), solve combinatorial optimization via simulated annealing (Kirkpatrick et al., 1983), and for training or decoding deep energy based generative models (Wang & Cho, 2019; Du et al., 2020; Dai et al., 2020b). However, these tasks primarily depend on Gibbs sampling, which could be very slow in high dimensional space.
Since the seminal work Zanella (2020), the recent years have witnessed significant progresses for discrete sampling in both theory and practice. Zanella (2020) introduces the locally balanced proposal  $q(x,y)\propto g(\frac{\pi(y)}{\pi(x)})$ , where  $y\in N(X)$  restricted within a small neighborhood of  $x$  and  $g(\cdot):\mathbb{R}_{+}\to \mathbb{R}_{+}$  satisfying  $g(a) = ag(\frac{1}{a})$ , and prove it is asymptotically optimal. In the following works, PAS (Sun et al., 2021) and DMALA (Zhang et al., 2022) generalize locally balanced proposal to large neighborhoods by introducing an auxiliary path and mimicking the diffusion process, respectively. Inspired by these locally balanced samplers, Sun et al. (2022a) generalize the Langevin dynamics in continuous space to discrete Langevin dynamics (DLD) in discrete space as a continuous time Markov chain  $\frac{d}{dh}\mathbb{P}(X^{t + h} = y|X^t = x) = g(\frac{\pi(y)}{\pi(x)})$ , and show that previous locally balanced samplers are simulations of DLD with different discretization strategies. In the view of Wasserstein gradient flow, the Gibbs sampling can be seen as coordinate descent and DLD gives a full gradient descent. Hence, locally balanced samplers induced from DLD provide a principled framework to utilize modern accelerators like GPUs and TPUs to accelerate discrete sampling. Besides the discretization of DLD, another crucial part to design a locally balanced sampler is estimating the probability ratio  $\frac{\pi(y)}{\pi(x)}$ . Grathwohl et al. (2021) proposes to used gradient approximation  $\frac{\pi(y)}{\pi(x)}\approx \exp (-\langle \nabla f(x),y - x\rangle)$  and obtains good performance on various classical models and deep energy based models. When the Hessian is available, Rhodes & Gutmann (2022); Sun et al. (2023a) use second order approximation via Gaussian integral trick (Hubbard, 1959) to further improve the sampling efficiency on skewed target distributions. When the gradient is not available, Xiang et al. (2023) use zero order approximation via Newton's series.
Besides designing the sampler, Sun et al. (2022b) proves that when tuning path length in PAS (Sun et al., 2021), the optimal efficiency is obtained when the average acceptance rate is 0.574, and design an adaptive tuning algorithm for PAS. Sansone (2022) learn locally balanced weight function for locally balanced proposal, but how to select the weight function in a principled manner is still unclear.
3 Formulation for Sampling in Discrete Space
The sampling in discrete space can be formulated as the following problem: in a finite discrete space  $\mathcal{X}$ , we have an energy function  $f(\cdot): \mathcal{X} \to \mathbb{R}$ . We consider a target distribution
$$
\pi (x) = \frac {\exp (- \beta f (x))}{Z}, \quad Z = \sum_ {z \in \mathcal {X}} \exp (- \beta f (z)), \tag {1}
$$
where  $\beta$  is the inverse temperature. When the normalizer  $Z$  is intractable, people usually resort to Markov chain Monte Carlo (MCMC) to sample from the target distribution  $\pi$ . Metropolis-Hastings (M-H) (Metropolis et al., 1953; Hastings, 1970) is a commonly used general purpose MCMC algorithm. Specifically, given a current state  $x^{(t)}$ , the M-H algorithm proposes a candidate state  $y$  from a proposal distribution  $q(x^{(t)}, y)$ . Then, with probability
$$
\min  \left\{1, \frac {\pi (y) q \left(y , x ^ {(t)}\right)}{\pi \left(x ^ {(t)}\right) q \left(x ^ {(t)}, y\right)} \right\}, \tag {2}
$$
the proposed state is accepted and  $x^{(t + 1)} = y$ ; otherwise,  $x^{(t + 1)} = x^{(t)}$ . In this way, the detailed balance condition is satisfied and the M-H sampler generates a Markov chain  $x^{(0)}, x^{(1)}, \ldots$  that has  $\pi$  as its stationary distribution.
4 Benchmark for Sampling in Discrete Space
The recent development of locally balanced samplers that use the ratio  $\frac{\pi(y)}{\pi(x)}$  to guide the proposal distribution  $q(x, \cdot)$  has significantly improved the sampling efficiency in discrete space. However, there is no consensus for many experimental settings. As a result, the empirical results in different research papers may not be directly comparable. Under this circumstance, we propose DISCS as a benchmark for general purpose samplers in discrete space. In section 4.1, we introduce implemented sampling methods as the baselines in DISCS and make some remarks on how we present the results. DISCS implements both the classical discrete samplers and the recent developed locally balanced samplers. In the following sections, we introduce the tasks considered in DISCS as follows: sampling from classical graphical models (section 4.2), sampling for solving combinatorial optimization problems (section 4.3), and sampling from deep energy based generative models (section 4.4). We also describe how the discrete samplers are evaluated on these tasks and report several results of them with some insights. We leave the remaining comprehensive results and experimentation and report them in detail in Appendix A.
We include both classical discrete samplers and state-of-the-art locally balanced samplers in recent research papers as baselines in our benchmark. Specifically, DISCS implements
1. random walk Metropolis (RWM) (Metropolis et al., 1953).
2. block Gibbs (BG), where BG-  $<  a>$  denotes using block Gibbs with block size  $a$
3. Hamming Ball Sampler (HB) (Titsias & Yau, 2017), where HB-  $<a>$ - $<b>$  denotes using block size  $a$  and Hamming ball size  $b$ .
4. Gibbs with Gradients (GWG) (Grathwohl et al., 2021), a locally balanced sampler that uses gradient to approximate the probability ratio. For binary distribution, GWG has a scaling factor  $L$  to determine how many sites to flip per step.
5. Path Auxiliary Sampler (PAS) (Sun et al., 2021), a locally balanced sampler that has a scaling factor  $L$  to determine the path length.
6. Discrete Metropolis Adjusted Langevin Algorithm (DMALA)(Zhang et al., 2022), a locally balanced sampler that has a scaling factor  $\alpha$  to determine the step size.
7. Discrete Langevin Monte Carlo (DLMC) (Sun et al., 2022a), a locally balanced sampler that has a scaling factor  $\tau$  to determine the simulation time of DLD. DLMC has multiple choices for its numerical solver to approximate the transition matrix. DISCS considers the two versions used in the original paper, DLMC that uses an interpolation, and DLMCf that uses Euler's forward method.
Note that in the depicted plots in this paper, each of the previously described samplers is represented with a distinct and unique color of its own. We present samplers RWM, BG, HB, GWG, PAS, DMALA, DLMC, DLMCf as the colors green, yellow, blue, red, brown, purple, pink and grey.
Remark: weight function DISCS offers a range of locally balanced functions, including  $g(t) = \sqrt{t}$ ,  $g(t) = \frac{t}{t + 1}$ ,  $g(t) = 1 \wedge t$ , and  $g(t) = 1 \vee t$ . All the locally balanced samplers have the flexibility to select from these locally balanced functions.  $g(t) = \sqrt{t}$  and  $g(t) = \frac{t}{t + 1}$  are the two most commonly used weight functions which we also rely on for our experiments. We use <samp> to refer to the type of the weight function for the locally balanced sampler. In cases that the weight function is not reported, we use  $\sqrt{t}$  by default.
Remark: scaling Since the scaling of the proposal distribution in RWM, PAS, DMALA, and DLMC are tunable, we consider two versions: one with adaptive tuning and another with binary search tuning to ensure fair comparison. Sun et al. (2022b, 2023b) propose an adaptive tuning algorithm for PAS and DLMC when the target distribution is factorized. In practice, we find that they also apply well to other locally balanced samplers and for more general target distributions. Hence, in this paper, we use the adaptive tuning algorithm by default to tune the scaling for locally balanced samplers. In the several exceptions where the adaptive algorithm does not apply, we will use <sampier-name>-noA to indicate the results from binary search tuning.
4.2 Sampling from Classical Graphical Models
This section covers the classical graphical models frequently employed in physics and statistics, including Bernoulli Models, Ising Models (Ising, 1924), and Factorial Hidden Markov Models (Ghahramani & Jordan, 1995). The graphical models have large configuration flexibility, for example, the number of discrete variables, the number of categories for each discrete variable, and the temperature of the model. The performances of different samplers can heavily depend on these configurations. DISCS provides tools to automatically sweep over hundreds of configurations with one click. Following the common practice in Monte Carlo integration or Bayesian inference, DISCS uses the Effective Sample Size (ESS) to evaluate the efficiency of each sampler and reports the ESS normalized by the number of calling energy functions and by the running time. Below, we present several Ising Model results as illustrative examples. We report more results and details of our experiments and ESS calculation in the Appendix A.1, A.6.1. More specifically, we report the performance results of different samplers on Bernoulli Model 7, Categorical Model 8, Ising Model 9, Potts Model 10 and FHMM 11. We study the effect of number of discrete variables (sample dimension), the number of categories for each discrete variable, weight function for locally balanced samplers, and the temperature of the models (smoothness/sharppness).
The Ising Model is defined on a 2D grid, where the state space  $\mathcal{X} = \{-1,1\}^{p\times p}$  represents the spins on all nodes. For each state  $x\in \mathcal{X}$ , the energy function is defined as:
$$
f (x) = - \sum_ {i, j} J _ {i j} x _ {i} x _ {j} - \sum_ {i} h _ {i} x _ {i} \tag {3}
$$
where  $J_{ij}$  is the internal interaction and the  $h_i$  is the external field. In the main text, we report the results using the configuration from Zanella (2020). Specifically,  $J_{ij} = 0.5$ ,  $h_i = \mu_i + \sigma_i$ , where  $\sigma_i \sim \mathrm{Uniform}(-1.5, 1.5)$  and  $\mu_i = 0.5$  if node  $i$  is located in a circle has the same center as the 2D grid and radius  $\frac{p}{2\sqrt{2}}$ , else  $-0.5$ . We consider the target distribution  $\pi(x) \propto \exp(-\beta f(x))$ , where  $\beta$  is the inverse temperature. Using DISCS, one can easily investigate the influence of the number of discrete variables (sample dimension). In Figure 1, one can see that the classical samplers, RWM, BG, HB, have a significant decrease in ESS when the model dimension increases, while the locally balanced samplers are less affected as the ratio information  $\frac{\pi(y)}{\pi(x)}$  effectively guides the proposal distribution. The overall trends basically follow the prediction from Sun et al. (2022b) that the ESS is  $O(d^{-1})$  for RWM and  $O(d^{-\frac{1}{3}})$  for PAS.
Figure 1: Results of Ising model with different dimensions.
Through DISCS, researchers can easily define various experiments, configure tasks, evaluate sampler performance and gain invalidabale insights on open questions. As an example, we experiment with sampling from Ising Models with a range of temperatures and sample dimensions. In Figure 2, we experiment with sampling from Ising Models with inverse temperatures from 0.1607 to 0.7607 for both sample dimension of  $50 \times 50$  and  $100 \times 100$ . We consider Ising Model without external field:  $h_i \equiv 0$  and  $J_{ij} \equiv 1$  as we know the critical temperature for this configuration is  $\frac{2}{\log(1 + \sqrt{2})}$ . This gives us the critical point for inverse temperature as  $\beta = 0.4407$ . From the results in Figure 2, we can see that
- The Ising model is harder to sample from when the inverse temperature  $\beta$  is closer to the critical point, which is consistent with the theory in statistical physics.
- When the inverse temperature  $\beta$  is lower than the critical point, using weight function  $g(t) = \sqrt{t}$  gives larger ESS; When the inverse temperature is larger than the critical point, using weight function  $g(t) = \frac{t}{t + 1}$  consistently obtains larger ESS.
The second observation implies that one should use ratio function  $\frac{t}{t + 1}$  for target distributions with a sharp landscapes. We will revisit this conclusion in Table 2. We report more results on the effect of sample dimension in Appendix (7, 8, 9, 10).
Figure 2: Performance of locally balanced samplers with different types of weight functions v.s. temperature on: (left)  $50 \times 50$  Ising model, (right)  $100 \times 100$  Ising model.
The categorical version of Ising model is Potts model, where each site of a state  $x_{i}$  has values in a symmetry group, instead of  $\{-1,1\}$ . For simplicity, we denote the symmetry group as a set of one hot vectors  $\mathcal{C} = \{e_1,\dots ,e_c\}$  with  $h_i\in \mathbb{R}^C,J_{ij}\in \mathbb{R}^{C\times C}$ . In this way, the energy function becomes:
$$
f (x) = - \sum_ {i, j} x _ {i} ^ {\top} J _ {i j} x _ {j} - \sum_ {i} \left\langle h _ {i}, x _ {i} \right\rangle \tag {4}
$$
In Figure 3, one can see the sampling efficiency is very robust with respect to the number of categories, with the exception of the classical HB and BG samplers, which demonstrate a decline in their sampling performance. The result for BG-2 on Potts model with 256 categories is omitted as it takes over 100 hours. We report more results on the effect of number of categories in Appendix (8, 10, 11).
Figure 3: Results of Potts models with different number of categories.
4.3 Sampling for Solving Combinatorial Optimization
Combinatorial optimization is a core challenge in domains like logistics, supply chain management, and hardware design, and has been a fundamental problem of study in computer science for decades. Combining with simulated annealing Kirkpatrick et al. (1983), the discrete sampling algorithm is a powerful tool to solve combinatorial optimization problems (Sun et al., 2023b). In expectation, a sampler with a faster mixing rate can find better solutions. Hence, the second type of task is sampling for solving combinatorial optimization problems. Currently, DISCS covers four problems:
Maximum Independent Set (MIS), Max Clique, MaxCut, and Balanced Graph Partition. Without loss of generality, we consider combinatorial optimization that admits the following form:
$$
\min  _ {x \in \mathcal {C} = \{0, 1, \dots , C - 1 \} ^ {d}} a (x), \quad \text {s . t .} \quad b (x) = 0 \tag {5}
$$
For ease of exposition, we assume  $b(x) \geq 0, \forall x \in \mathcal{C}$ , but otherwise do not limit the form of  $a$  and  $b$ . To convert the optimization problem to a sampling problem, we first rewrite the constrained optimization into a penalty form via a penalty coefficient  $\lambda$ , then treat this as an energy function for an EBM. In particular, the energy function takes the form:
$$
f (x) = a (x) + \lambda \cdot b (x) \tag {6}
$$
Then, we define the probability of  $x$  at inverse temperature  $\beta$  by:
$$
p _ {\beta} (x) \propto \exp (- \beta f (x)) \tag {7}
$$
A naive approach to this problem would be directly sampling from  $p_{\beta \to \infty}(x)$ , but such a distribution is highly nonsmooth and unsuitable for MCMC methods. Instead, following classical simulated annealing, we define a sequence of distributions parameterized by a sequence of decaying temperatures:
$$
\mathcal {P} = \left[ p _ {\beta_ {0}} (x), p _ {\beta_ {1}} (x), \dots , p _ {\beta_ {T}} (x) \right] \tag {8}
$$