Discrete sampling

Many tasks within deep learning involve sampling over discrete spaces. Such sampling tasks are distinct from sampling over continuous spaces: while the primary frameworks of sampling—Markov chain Monte Carlo (MCMC), rejection sampling, and importance sampling—introduced in Chapter Introduction to sampling can be applied to discrete spaces by adjusting the proposal distribution, this adjustment is often highly non-trivial. In this chapter, we focus on high-dimensional targets whose unnormalised log density can be evaluated but whose state spaces are too large to enumerate. Such targets arise in binary energy-based models, categorical latent-variable models, and neural models over combinatorial structures.

Introduction¶

Discrete sampling is notably difficult due to the sparse, discontinuous, and combinatorial nature of discrete spaces. If a state consists of $d$ categorical variables with $K$ possible values each, its state space may contain $K^d$ configurations. In principle, one could evaluate the unnormalised density at every configuration, normalise the resulting weights, and sample from the corresponding categorical distribution. This enumeration-based approach becomes infeasible as the dimension grows. Standard local MCMC methods avoid full enumeration, but target-independent proposals can mix slowly because they do not distinguish promising moves from moves into low-density regions.

As a running example, consider a binary energy-based model with $x\in\{0,1\}^d$ and unnormalised log density $U(x)$. The model assigns a score to each of its $2^d$ possible configurations, making enumeration infeasible for large $d$. Nevertheless, $U(x)$, and often the gradient of a differentiable extension of $U$, can be evaluated at any given configuration. We use this binary model to illustrate the methods below, although the framework also applies to the broader categorical and combinatorial settings described above.

Gradients are a staple of sampling within deep learning, as illustrated by the continuous-space Langevin methods in Section Metropolis-adjusted Langevin Algorithm. Although a function defined only on a discrete space has no intrinsic gradient, many deep-learning targets are represented through a differentiable function on a continuous ambient space and evaluated at binary, categorical, or other discrete states. The gradient of this extension can provide a cheap source of local information about promising discrete moves.

We begin with an uninformed local MCMC proposal and show how pointwise-informed and locally balanced proposals incorporate target-specific information. We then use this framework to develop gradient-based discrete samplers, first for single-coordinate moves and then for parallel multi-coordinate moves. Finally, we discuss methods designed for larger-scale exploration.

Notation¶

We assume a discrete state space $\mathcal{X}$. We denote a state by $x=(x_1,\dots,x_d)\in\mathcal{X}$, where each coordinate can take one of $K$ possible values. Throughout the chapter, $\pi$ denotes the normalised target distribution and $U$ its unnormalised log density. Thus our target is

The normalising constant or partition function $Z$ in Equation (1) generally requires a sum over as many as $K^d$ states. This makes enumeration-based direct sampling computationally infeasible in high dimensions. MCMC can nevertheless operate on differences such as $U(x')-U(x)$, in which $Z$ cancels.

Locally Balanced Proposals¶

A simple Metropolis–Hastings approach proposes a nearby state using a target-independent kernel $K_\sigma(x'\mid x)$. For example, the kernel might choose a coordinate uniformly and replace its value. Such a proposal is inexpensive, but it assigns the same proposal probability to moves towards high- and low-density states. Large moves may explore quickly but often reach low-density states and are rejected; small moves are more likely to be accepted but can produce highly correlated samples and traverse the state space slowly.

For the running binary example, the simplest local kernel selects one of the $d$ bits uniformly and flips it, without using the values of $U$ to prioritise the available moves.

Zanella (2020) introduces a framework that addresses this trade-off by weighting the candidate moves from $K_\sigma$ using target-density information. The resulting pointwise-informed proposal is

The pointwise-informed proposal in Equation (2) consists of the target-density ratio, a symmetric base kernel or weight $K_\sigma$ whose scale controls the size of candidate moves, a balancing function $g$ that weights moves using the density ratio, and the state-dependent normaliser $Z_{g,\sigma}(x)$. The base kernel may be specified only up to a state-independent multiplicative constant, because that constant cancels through $Z_{g,\sigma}(x)$. We assume $K_\sigma(x'\mid x)=K_\sigma(x\mid x')$. The limit $\sigma\to0$ represents a family of kernels that becomes increasingly concentrated at the current state.

The roles of invariance and ergodicity in MCMC are reviewed in Section Markov chain Monte Carlo, while detailed balance and the Metropolis–Hastings construction are discussed in Section Metropolis–Hastings. In the present setting, $P_{g,\sigma}$ is reversible with respect to $\pi_{g,\sigma}$ if it satisfies detailed balance,

Detailed balance is sufficient for $\pi_{g,\sigma}$ to be invariant under the transition. Invariance alone does not ensure convergence from an arbitrary initial state. For a discrete chain, irreducibility means that every state with positive target probability can eventually be reached from every admissible initial state, while aperiodicity means that the chain is not constrained to return according to a fixed cycle. Together with invariance and the relevant regularity conditions, these properties yield an ergodic chain. The general role of ergodicity in MCMC is reviewed in Section Markov chain Monte Carlo.

A pointwise-informed proposal is called locally balanced if the transition kernel $P_{g,\sigma}$ is reversible with respect to some distribution $\pi_{g,\sigma}$ such that $\pi_{g,\sigma}$ converges weakly to $\pi$ as $\sigma \to 0$ Zanella, 2020. In plain terms, as the proposal becomes increasingly local, the invariant distribution of the uncorrected informed proposal approaches the desired target. Zanella (2020) demonstrates that a pointwise-informed proposal is locally balanced if and only if the balancing function satisfies $g(t) = t g(1/t)$.

If this property holds for the function $g$, then the locally balanced proposal will be reversible with respect to the following distribution without any correction:

Here, $Z_{g,\sigma}$ is defined as follows:

For a local kernel satisfying $K_\sigma(x'\mid x)\to\mathbf 1\{x'=x\}$ and the conventional normalisation $g(1)=1$, we have $Z_{g,\sigma}(x)\to1$. Consequently, $\pi_{g,\sigma}$ converges to $\pi$ as $\sigma\to0$ Zanella, 2020.

At a non-zero scale, $\pi_{g,\sigma}$ need not equal $\pi$, so a Metropolis–Hastings correction is still required for exact sampling. Local balance makes the uncorrected proposal close to the target in the small-scale regime, which can improve the acceptance rate, while a non-zero $\sigma$ is required to produce non-degenerate movement across the state space.

Zanella (2020) shows that locally balanced proposals can improve efficiency over uninformed local proposals in the asymptotic regimes considered there. Constructing an exact informed proposal nevertheless requires evaluating the density ratio for every candidate. Full enumeration can involve $K^d$ states; even a Hamming-one neighbourhood contains $d(K-1)$ candidates. The gradient-based methods below replace these repeated neighbouring-state evaluations with one gradient evaluation at the current state.

Gradient-Based Discrete Sampling¶

A function specified only on a discrete state space does not have an intrinsic gradient. Gradient-based discrete samplers therefore assume that the unnormalised log density $U$ is supplied with, or represented through, a differentiable extension to a continuous ambient space. For example, a neural network may accept a continuous vector representation even though the sampler evaluates it only at binary or one-hot encoded states. This extension need not be unique, and its gradient is useful only insofar as it predicts changes between the discrete states of interest.

Given such an extension, the log-density difference for a nearby state can be approximated by

This approximation is most reliable when the extension is sufficiently smooth and $x'$ is close to $x$; differentiability alone does not guarantee accuracy if the function changes sharply between neighbouring states. Since $\log\pi(x)=U(x)-\log Z$, their gradients are identical and the unknown normalising constant is not required. One forward and backward pass at $x$ then supplies approximate scores for all local candidate moves, at the expense of Taylor-approximation error.

Gibbs with Gradient¶

Grathwohl et al. (2021) first introduces this notion of gradient-based discrete sampling through the Gibbs-With-Gradient algorithm. This algorithm can be seen as a version of the standard Gibbs algorithm where the updated coordinate is selected based on the gradient.

Returning to the general locally balanced sampler, we can derive the Gibbs-With-Gradient update by choosing an uninformed kernel that uniformly selects a state differing from the current state in exactly one coordinate. Define the Hamming-one neighbourhood by

In the running binary example, $H(x)$ contains exactly the $d$ states obtained by flipping one bit of $x$.

For this fixed-scale update, the uniform kernel is

When every coordinate has $K$ possible categories, $|H(x)|=d(K-1)$. This factor is constant across candidate moves and is therefore absorbed into the proportionality below.

Next, we choose the balancing function $g(t) = \sqrt{t}$. Before approximation, this yields the exact informed proposal

Using the unnormalised log density $U$ from Equation (1), this proposal can be expressed without evaluating $Z$:

Replacing the exact log-density difference by its first-order Taylor approximation defines the Gibbs-With-Gradient proposal:

Figure 1 illustrates this Gibbs-With-Gradient construction. The gradient at $x$ assigns different approximate weights to the states in $H(x)$, and $x'$ is then sampled from the resulting proposal over those discrete neighbours.

This can be computed quickly by viewing the update as the selection of a coordinate to modify. In the binary case, this update is always a bit flip. Here, $\operatorname{Softmax}$ normalises a vector of real-valued scores into probabilities, and $\operatorname{Cat}$ denotes the categorical distribution with those probabilities. The resulting proposal is

The categorical case requires selecting both a coordinate and its proposed category. Let $c_i$ denote the current category of coordinate $i$. We represent each categorical coordinate by a one-hot vector when evaluating the gradient. For every alternative category $j \neq c_i$, define the approximate log-probability difference

The first term measures the local effect of assigning category $j$ to coordinate $i$, while the second corresponds to its current category.

There are $d(K-1)$ possible coordinate–category updates. We define a categorical proposal over these updates by

After sampling $(i,j)\sim q_{\mathrm{GWG}}(\cdot,\cdot\mid x)$, the proposed state $x'$ is obtained by changing coordinate $i$ from $c_i$ to $j$ and leaving all other coordinates unchanged. The full categorical procedure is given in Algorithm 1.

Because the gradient-based proposal uses a first-order approximation and is generally asymmetric, it does not by itself preserve the target distribution. We therefore apply the standard Metropolis–Hastings acceptance correction, described in Section Metropolis–Hastings, which ensures that the target distribution is invariant. Under the ergodicity conditions stated above, the resulting Markov chain converges to the target distribution.

Discrete Langevin Proposal¶

Gibbs-With-Gradient can outperform uninformed single-coordinate proposals by using the gradient to prioritise promising updates. Its Hamming-one restriction also keeps the Taylor approximation local, but only one coordinate can change per iteration, which may produce strongly correlated samples in high dimensions. Zhang et al. (2022) addresses this limitation with the Discrete Langevin Proposal (DLP), which retains gradient-based directionality while adding a distance penalty that permits multiple coordinates to be proposed simultaneously.

In the running binary example, Gibbs-With-Gradient proposes at most one bit flip, whereas DLP constructs a proposal for every bit and may change several bits in the same iteration.

To derive the proposal, consider the discrete analogue of a Gaussian random-walk kernel,

Combining this kernel with a first-order approximation of the locally balanced weight gives the DLP

The linear term in Equation (17) favours moves predicted to increase $U$, while the quadratic term penalises moves far from the Taylor expansion point. The step size $\sigma$ controls the strength of this penalty.

Because the state space factorises as $\mathcal{X}=\prod_{i=1}^d\mathcal{X}_i$ and the exponent in Equation (17) is a sum of coordinate-wise terms, the proposal factorises as $q(x'\mid x)=\prod_{i=1}^d q_i(x'_i\mid x)$. Each factor is the categorical distribution

Thus, larger values of $\sigma$ permit more distant proposals, whereas smaller values concentrate the proposal near the current state.

While this may look similar to the proposal formula for Gibbs with Gradient, it is important to note that in DLP, the softmax is computed per coordinate over the potential classes that can be realised. In Gibbs with Gradient, the softmax is over the entire sample to pick only one coordinate. This allows for the parallel update of all coordinates.

The factorisation makes it possible to propose all coordinates in parallel after one gradient evaluation. Although the proposal factors coordinate-wise conditional on the current state, the gradient of the joint log density still guides each coordinate using information from the full current configuration. A Metropolis–Hastings correction can then account for the remaining approximation error and dependence structure. Figure 2 contrasts this parallel update with the single-coordinate update used by Gibbs-With-Gradient.

Using DLP yields two samplers: the discrete Metropolis-adjusted Langevin algorithm (DMALA), which includes a Metropolis–Hastings correction, and the discrete unadjusted Langevin algorithm (DULA), which omits it. DULA generally has asymptotic bias. For log-quadratic targets, however, Zhang et al. (2022) show that this bias vanishes as $\sigma\to0$; they also bound the bias for targets whose gradients are close to those of a log-quadratic distribution.

In terms of implementation, the cost difference between DULA and DMALA is negligible: while DMALA incorporates a Metropolis–Hastings correction that requires the computation of both $q(\cdot\mid x)$ and $q(\cdot\mid x')$, caching means that only one new proposal distribution is computed per step after initialisation. The cached version of DMALA is given in Algorithm 2.

The core logic lies in the fact that the $x'$ of the current iteration is either rejected or becomes the $x$ for the next sampling step. In the case of rejection, the proposal function $q(\cdot| x)$ does not change as $x$ remains the same. In the case of acceptance, the proposal step simply gets updated to match the updated $x'$ by using $q(\cdot | x') \to q(\cdot | x)$, since this proposal must be computed anyway for the accept/reject step. Thus at each step, only one proposal function must be computed.

Beyond Locally Balanced Proposals¶

Locally balanced proposals offer a convenient framework for incorporating gradients into discrete samplers, but their defining approximation concerns the small-scale limit $\sigma\to0$. At finite scales, local proposals may take a long time to move between separated modes, necessitating longer chains or multiple chains.

For example, if the binary target places high probability on two configurations separated by many low-probability intermediate states, a sequence of single-bit proposals may struggle to move from one mode to the other.

This motivates proposals designed for non-local movement. Zanella (2020) already considers $g(t)=t$, calling the resulting proposal the naively informed choice because it has the form $P_{g,\sigma}(x'\mid x)\propto K_\sigma(x'\mid x)\pi(x')$. Building on this observation, Sun et al. (2023) formalise the large-scale interpretation: they call a proposal family globally balanced if the invariant distribution of its uncorrected informed proposal converges to $\pi$ as $\sigma\to\infty$. They assume that the base kernel becomes uniform in this limit. Since a state-independent constant cancels in the normalisation of the informed proposal, they write this limit as $K_\sigma(x'\mid x)\to1$; for a normalised kernel on a finite state space, the corresponding limit is $1/|\mathcal X|$. For $g(t)=t$, the limiting proposal is then

Thus the proposal becomes an independent draw from the target itself and is globally balanced in the sense of Sun et al. (2023). Notice that $g(t)=t$ does not satisfy the locally balanced identity $g(t)=t g(1/t)$; local and global balance refer to different limiting regimes.

Any-scale Balanced Sampler¶

The any-scale balanced sampler introduced by Sun et al. (2023) makes two distinct modifications for proposals outside the local regime. First, it adjusts the density-ratio weighting to match the proposal scale. The family $g(t)=t^\beta$, with $\beta\in[1/2,1]$, interpolates between the locally balanced choice $g(t)=\sqrt{t}$ as $\sigma\to0$ and the globally balanced choice $g(t)=t$ as $\sigma\to\infty$. This scale-dependent relationship is illustrated schematically in Figure 3(a).

Adjusting $\beta$ changes how strongly the proposal favours high-density candidates, but it does not reduce the error in a first-order approximation of $U(x')-U(x)$. The method therefore makes a second modification: it introduces a quadratic approximation intended to remain informative over larger moves. The log-density difference is approximated by

where $W$ is a symmetric matrix shared across states. Setting $W=\nabla^2U(x)$ would recover the usual second-order Taylor approximation locally, but repeatedly computing a Hessian is expensive. Instead, the sampler estimates a global $W$ from pairs $(x_n,x'_n)$ collected during burn-in:

Thus $W$ is fitted to explain observed changes in the gradient and acts as a global curvature approximation. Scale-dependent balancing and curvature approximation address different problems: $\beta$ controls the density-ratio weighting, whereas $W$ improves the approximation used to score non-local candidates.

Substituting the quadratic approximation in Equation (21) and $g(t)=t^\beta$ into the informed proposal gives

The quadratic term captures cross-coordinate interactions, but it also prevents direct coordinate-wise sampling. Sun et al. (2023) therefore introduce an augmented proposal with a Gaussian auxiliary variable. The augmented construction is used directly with the Metropolis–Hastings correction below.

Choose a diagonal positive semi-definite matrix $D$ such that $W+D$ is positive semi-definite, and draw

The matrix square root can be computed by an eigendecomposition or Cholesky factorisation. Conditional on $Y$, define

and sample $x'$ from this distribution. Because $D$ is diagonal and the state space factorises, this conditional proposal factorises across coordinates. Equations (23) and the augmented proposal express, respectively, the motivating quadratic approximation and the auxiliary-variable transition implemented by the sampler. The Metropolis–Hastings step accounts for the resulting augmented proposal exactly.

The Metropolis–Hastings acceptance probability for the augmented proposal is

The diagonal matrix $D$ controls the conditional proposal variance. Sun et al. (2023) consider two semidefinite programs for choosing it. The AB-max variant minimises the largest eigenvalue of $W+D$:

This criterion focuses on the worst spectral direction. The AB-trace variant instead solves

The motivation for this objective comes from the continuous quadratic surrogate studied by Sun et al. (2023) in the special case $\sigma=\infty$ and $\beta=1$. In that surrogate, the covariance of the proposal increment is $2D^{-1}$, so minimising the trace of $D$ encourages dispersion across coordinates. This covariance identity does not hold literally for the general discrete categorical proposal, whose dispersion also depends on $\beta$, $\sigma$, and the state space. The authors report that AB-trace performs better empirically and use it as their preferred variant.

The proposal parameters $\beta$ and $\sigma$ are correlated and target-dependent. Sun et al. (2023) therefore adapt them by alternating coordinate-wise updates that seek to maximise the empirical average jump distance. For one parameter type $a\in\{\sigma,\beta\}$, let $\phi$ denote its current value and define the feasible sets $\mathcal I_\sigma=(0,\infty)$ and $\mathcal I_\beta=[1/2,1]$. With the other parameter held fixed, the adaptation block compares $\phi$ with feasible multiplicative perturbations of it over three consecutive chain segments, as shown in Algorithm 3.

The three segments in Algorithm 3 are consecutive: the second begins from $x^{(N)}$ and the third from $x^{(2N)}$. The full adaptation alternates a block for $\phi=\sigma$ and a block for $\phi=\beta$. If neither parameter changes during an outer iteration, the update rate is reduced geometrically, $\gamma\gets\rho\gamma$ for a decay factor $\rho\in(0,1)$; otherwise the selected parameter values are retained for the next iteration. Because adaptation based on short-run jump distances can be noisy during early mixing, this procedure is intended for burn-in rather than continued adaptation after samples are retained for inference.

Through the combination of a quadratic approximation of the log-probability difference $U(x') - U(x)$ and the adaptive scheme for $\beta$ and $\sigma$, the any-scale sampler interpolates between locally and globally balanced proposals.

Automatic Cyclical Sampling¶

One common issue with gradient-based samplers for continuous spaces is exploration: as the gradient points in the direction of the closest mode, it gives no information regarding further, more distant modes. This can hinder the accuracy of the sampler in terms of approximating the target distribution.

The same problem applies to gradient-based discrete samplers. Unlike a continuous state space, a discrete space has no arbitrarily small non-trivial perturbations: even the most local move must change at least one coordinate. This granularity can make the trade-off between conservative, high-acceptance moves and larger exploratory moves especially pronounced. Interpolating between locally and globally balanced proposals is therefore a natural way to vary the proposal scale when large movements may be required to escape a local mode.

In the continuous case, many approaches have been proposed for handling the mode-escape problem. One popular approach is cyclical SG-MCMC, discussed in Chapter Stochastic gradient MCMC. A cyclical schedule gradually reduces the step-size of a gradient-based sampler from very large magnitudes to small magnitudes and repeats for several cycles. The large steps at the beginning allow for more exploratory movements, which enable mode discovery; and the smaller step-sizes at the end of the cycle allow for exploitation of the mode. The repetition of the cycles allows for the discovery and exploitation of many modes, which is required for effective characterisation of multi-modal distributions.

Building on this intuition, Pynadath et al. (2024) introduce a sampler that cycles jointly through step sizes and balancing exponents. Varying $\sigma$ alone can produce poorly balanced proposals and low acceptance probabilities, so the method uses the generalised DLP

Setting $\beta=1/2$ in Equation (29) recovers the standard DLP. Within each cycle, large values of $\sigma$ are paired with values of $\beta$ closer to 1 to encourage global exploration; as $\sigma$ decreases, $\beta$ moves towards $1/2$ and the proposal becomes more local.

The automatic tuning procedure begins with user-specified endpoint exponents $\beta_{\min}$ and $\beta_{\max}$ and a target acceptance rate. Short burn-in runs determine corresponding endpoint step sizes $\sigma_{\min}$ and $\sigma_{\max}$. These endpoints define the cyclical step-size schedule, following the continuous-space construction of Zhang et al. (2020).

For each step size $\sigma_k$, the corresponding $\beta_k$ is selected to maximise an empirical acceptance-rate objective subject to $\beta_k\in[1/2,\beta_{k-1}]$. This monotonic constraint reflects the transition from large moves with more globally balanced weighting at the beginning of a cycle to small, locally balanced moves at its end, while reducing the cost of searching for the full schedule. Figure 3(b) shows the coupled schedules used to move repeatedly from global exploration to local characterisation.

Conclusion¶

Gradient information can make sampling in high-dimensional discrete spaces substantially more efficient, even though the state space itself is not differentiable. Locally balanced proposals provide the organising principle: they combine an uninformed local kernel with target-specific information and use a Metropolis–Hastings correction to retain the desired target distribution. Gibbs-With-Gradient uses a first-order approximation to choose an informed single-coordinate update, whereas DLP factorises a Langevin-like proposal so that many coordinates can be proposed in parallel.

The appropriate proposal depends on the scale of movement required. Small, locally balanced moves can characterise a mode accurately but may mix slowly between separated modes. Any-scale balancing adapts the balancing exponent and adds a global curvature approximation for more effective non-local moves. Automatic cyclical sampling instead repeats a schedule from large exploratory proposals to small exploitative proposals, pairing each step size with a suitable balancing exponent.

In the running binary example, this progression moves from informed single-bit proposals, through parallel multi-bit updates, to methods that adapt the scale and weighting of those updates to improve movement between separated modes.

These methods trade additional structure and tuning for improved exploration. Their practical performance depends on the quality of the differentiable extension of $U$, the accuracy of the local or quadratic approximation, and the cost of evaluating reverse proposals in the Metropolis–Hastings correction. The adjusted variants provide asymptotically exact sampling under the usual irreducibility and aperiodicity conditions; unadjusted variants can be cheaper, but generally introduce stationary bias.