PaperlayerView exact revision on arXiv

Source-linked AI summary

Oops I Took A Gradient: Scalable Sampling for Discrete Distributions

Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, Chris J. Maddison

arXiv:2102.04509v2cs.LG

TL;DR

Discrete sampling is increasingly needed for complex models, yet useful structure can be difficult to identify and hard-code. This paper uses gradients to guide Metropolis-Hastings proposals and reports strong sampling and deep-model results, while noting substantial compute costs for Gibbs in the deep-EBM setting.

  • Problem

    As discrete-data models grow more complex, sampling methods are needed, but useful block structure can be difficult to detect automatically and may require users to hard-code it.

    Source: §1 · p. 1

  • Method

    The paper uses gradients of continuous functions underlying discrete distributions to construct Metropolis-Hastings proposals, termed Gibbs-With-Gradients.

    Source: §1 · p. 1, §9 · p. 9

  • Results

    Gibbs-With-Gradients outperforms diverse baseline and gradient-based discrete samplers, and supports unconstrained deep EBMs on discrete data that outperform other deep generative models.

    Source: §9 · p. 9

  • Takeaways & Limitations

    The reported results show that gradient-informed sampling can support training deep energy-based models on discrete data.

    Source: §9 · p. 9

    For the reported deep-EBM setting, Gibbs requires 128x more compute per iteration than GWG and roughly 870x the runtime, making Gibbs infeasible for this training task.

    Source: §8 · p. 8

Abstract

from arXiv · show

We propose a general and scalable approximate sampling strategy for probabilistic models with discrete variables. Our approach uses gradients of the likelihood function with respect to its discrete inputs to propose updates in a Metropolis-Hastings sampler. We show empirically that this approach outperforms generic samplers in a number of difficult settings including Ising models, Potts models, restricted Boltzmann machines, and factorial hidden Markov models. We also demonstrate the use of our improved sampler for training deep energy-based models on high dimensional discrete data. This approach outperforms variational auto-encoders and existing energy-based models. Finally, we give bounds showing that our approach is near-optimal in the class of samplers which propose local updates.

1. Introduction

Discrete models are increasingly important, but exploiting their structure can require users to identify and hard-code it. The paper uses gradients of continuous functions underlying discrete distributions to guide MCMC proposals, reporting large sampling gains and improved deep energy-based models.

1. Introduction

  • The need for discrete sampling methods grows as scientific models for discrete data become more complex.The passage contrasts discrete outcomes with continuous distributions, whose differentiable densities can often be exploited for sampling and inference.Source: §1 · p. 1
  • Users may need to identify block-independence structure in combinatorial distributions and hard-code it into an algorithm to speed sampling.Discrete distributions over combinatorial spaces often have block independence, but detecting it automatically may be difficult.Source: §1 · p. 1
  • Gradients of continuous functions underlying discrete distributions inform Metropolis-Hastings proposals for large discrete probabilistic models.The resulting proposal family may be viewed as adaptive Gibbs samplers or as a fast approximation to locally informed proposals.Source: §1 · p. 1
  • Gradient-informed proposals can improve sampling efficiency by orders of magnitude and sometimes outperform samplers using hard-coded independence structure.The authors also report that the approach supports training deep energy-based models that outperform baseline variational auto-encoders and existing energy-based models.Source: §1 · pp. 1–2

2. Background

The paper frames discrete sampling through generic Gibbs updates and locally informed proposals, highlighting the costs of unproductive updates and exhaustive local scoring. Its approach motivates input-dependent proposals that use gradients to target likely changes more cheaply.

2. Background

  • The study considers unnormalized distributions with unknown normalizing constants, using finite-dimensional binary or one-hot categorical data representations.The paper states that all finite-dimensional discrete distributions can be embedded in these forms.Source: §2 · p. 2

2.1. Gibbs Sampling

  • One-dimension-at-a-time Gibbs updates can waste computation on pixels unlikely to change, while gradient-informed, input-dependent proposals target dimensions more likely to change.For MNIST-like binary data, background and digit-interior pixels often remain unchanged; the likely-to-change edge depends on the entire input.Source: §2.1 · p. 2

2.2. Locally-Informed Proposals

  • Locally informed proposals balance likelihood gains against reverse-proposal probability, but scoring every neighbor costs O(D); the gradient-based approach reduces this to O(1) evaluations.The cited locally balanced proposal uses temperature τ = 2 to balance the terms.Source: §2.2 · pp. 2–3

3. Searching for Structure

The paper seeks a broadly reusable structure for estimating local likelihood differences in discrete distributions. It observes that common distributions are expressed as differentiable continuous functions restricted to discrete inputs, allowing gradients and Taylor approximations to guide those estimates.

3. Searching for Structure

  • The desired structure should be broadly available and generic, since bespoke derivations and implementations are needed to exploit local differences efficiently in some distributions.The paper presents gradients as a way to estimate those differences across diverse discrete distribution families.Source: §3 · p. 3
  • Many distributions are implemented as differentiable functions of real-valued inputs, with discrete structure imposed by restricting inputs to a discrete subset.Table 1 presents common distributions' unnormalized log-likelihoods, which the caption says are differentiable with respect to x.Source: §3 · p. 3, Table 1 · p. 3
  • Taylor-series approximations using gradients estimate likelihood ratios between an input and nearby discrete states.For binary data, the approximation concerns flipping a bit; for categorical data, it concerns changing a dimension to another value, and linear operators can extend it to larger windows.Source: §3 · p. 3
  • A given discrete distribution may have multiple underlying continuous representations, and the paper presents theory to guide the choice when that ambiguity arises.The paper notes that its results cover finite-dimensional discrete distributions and that the method can also be applied to integer-valued distributions with a {x − 1, x + 1} proposal window.Source: §3 · p. 3

4. Gibbs With Gradients

Gibbs-With-Gradients uses Taylor-series estimates of local likelihood ratios to form Metropolis-Hastings proposals. Under smoothness assumptions, its approximation is bounded in efficiency relative to locally balanced proposals, with an Ising example illustrating the speed trade-off.

4. Gibbs With Gradients

  • The sampler approximates locally balanced proposals by using Taylor-series estimates of local likelihood ratios in a standard Metropolis-Hastings algorithm.The experiments use local moves within a Hamming window of size 1; for binary data, the proposal selects a dimension to change.Source: §4 · p. 3
  • For categorical data, the proposal chooses both a dimension and its replacement value using a D(K − 1)-way Softmax.The sampler is named Gibbs-With-Gradients because of its relationship to Adaptive Gibbs.Source: §4 · pp. 3–4

4.1. Analyzing Approximations

  • Under Lipschitz continuity of the log-probability gradient, the approximation's efficiency loss is bounded by a constant factor depending on the gradient's Lipschitz constant and proposal window size.The comparison is with the locally balanced proposal; when the Taylor approximation is accurate, the method matches its efficiency.Source: §4.1 · p. 4
  • For θ = .1 on a cyclic 2D Ising lattice with a size-1 Hamming window, Gibbs-With-Gradients has at most a constant efficiency loss per iteration and an O(D) speed increase over the exact proposal.The passage gives c = .67 regardless of D; computing the exact proposal costs O(D3), while Gibbs-With-Gradients incurs at most a constant per-iteration efficiency loss.Source: §4.1 · p. 4

5. Relationship to Continuous Relaxations

The paper contrasts direct gradient-informed proposals in discrete space with methods that sample through continuous relaxations. On RBMs, the relaxation baselines fail to scale to higher dimensions, while the proposed approach does not degrade with dimension.

5. Relationship to Continuous Relaxations

  • Computing exact local differences can be efficient for Ising models, but requires bespoke derivations and does not generalize to neural-network-parameterized distributions.The paper contrasts this with estimating local differences using gradients.Source: §5 · p. 4
  • Relaxation-based methods construct a continuous distribution, sample from it, then transform samples back to the discrete target.Examples include Discrete-SVGD and Discontinuous HMC, which use gradient-based methods such as SVGD or HMC in the relaxed space.Source: §5 · pp. 4–5
  • Relaxed distributions can be highly multimodal and require small step sizes; mixing in relaxed space may not predict mixing in discrete space.These methods also require tuning additional hyperparameters, whereas the paper’s approach operates directly in discrete space and has none.Source: §5 · p. 5
  • Beyond 200 dimensions, relaxation baselines do not scale, while GWG does not degrade with dimension on RBMs up to 1,000 dimensions.Figure 2 compares log-MMD against Block-Gibbs samples across increasing RBM dimensions.Source: §5 · p. 5

6. Sampling From EBMs

The section evaluates Gibbs-With-Gradients (GWG) on RBMs, lattice Ising models, and factorial hidden Markov models, comparing it with generic samplers and, when available, samplers that exploit model structure.

Sampling comparisons

  • 1.6, 6.6, and 2.1 times Gibbs-1 are the per-step costs of Gibbs-2, HB-10-1, and GWG, respectively.GWG’s per-step runtime is therefore most comparable to Gibbs-2 in this implementation.Source: §6 · p. 5

Restricted Boltzmann Machines

  • GWG matches Block-Gibbs in RBM MMD and has higher ESS than unstructured baselines, reaching about halfway between those baselines and Block-Gibbs in log-space.Block-Gibbs updates 784 dimensions per iteration, while GWG updates one.Source: §6 · p. 5, Figure 3 · p. 6

Lattice Ising Models

  • GWG provides a notable ESS increase for Ising models with higher connectivity, where dimensions are more correlated, and outperforms in most tested settings.Figure 4 reports log-ESS over 100,000 sampling steps for 10x10 and 40x40 lattices.Source: Figure 4 · p. 6, §6 · p. 6

Factorial Hidden Markov Models

  • GWG notably outperforms baseline samplers on FHMM posterior sampling, including a Hamming-Ball sampler that exploits the model’s structure.The evaluations use joint likelihood and reconstruction error; the structured sampler updates each 10-dimensional latent state as a block.Source: Figure 5 · p. 6, §6 · p. 6

7. Training EBMs

Slow-mixing MCMC can bias the gradient estimates used to train energy-based models, making sampler efficiency important for parameter inference. The section applies Gibbs-With-Gradients (GWG) to Ising and Potts models, including large proteins.

7. Training EBMs

  • Slow-mixing MCMC yields biased gradient estimates and sub-par learning, so improving sampling can improve parameter inference for unnormalized models.The Ising experiments use Persistent Contrastive Divergence with an ℓ1 penalty and assess connectivity recovery by RMSE.Source: §7 · pp. 6–7, §7.1 · p. 7

7.1. Training Ising models on generated data

  • GWG outperforms Gibbs sampling across the Ising settings, enabling faster training and higher-quality recovery of the true connectivity matrix.Figure 6 reports log(RMSE) between learned and true J for lattice and Erdos-Renyi Ising models.Source: §7.1 · p. 7, Figure 6 · p. 7

7.2. Protein Coupling Prediction with Potts Models

  • On large proteins, GWG performs best; Gibbs sampling falls below PLM, while it outperforms PLM on the smaller protein.The proteins have 276 and 526 amino acids, respectively; Figure 7 reports contact-prediction recall curves, with GWG achieving higher recall.Source: §7.2 · p. 7, Figure 7 · p. 7

8. Deep EBMs for Discrete Data

Recent deep energy-based models rely on gradient-based sampling that requires continuous data, leaving high-dimensional discrete data outside the scope of those scalable methods. This section uses GWG to train deep EBMs on binary and categorical data, and reports that the resulting models outperform competing methods on most datasets.

8. Deep EBMs for Discrete Data

  • GWG enables PCD training of ResNet-parameterized deep EBMs on high-dimensional binary and categorical data, including continuous images represented as 1-of-256 categorical data.The paper states that, to its knowledge, such models had not previously been successfully trained on this data.Source: §8 · p. 8
  • Binary-data training required 800 Gibbs steps per iteration versus 40 with GWG, making Gibbs 9.52x slower; the comparison used equal wall-clock time.Source: §8 · p. 8
  • For categorical data, Gibbs requires 256 energy evaluations per step versus 2 for GWG, making its per-iteration compute 128x greater and estimated runtime roughly 870x higher.The authors report that training this model form with Gibbs is not feasible.Source: §8 · p. 8
  • On most datasets, the ResNet EBM outperforms the VAEs and the RBM and DBN baselines on test-set likelihoods.The paper presents long-run samples from trained models in Figure 8, generated using 300,000 GWG steps.Source: §8 · p. 8, Figure 8 · p. 8

9. Future Directions and Conclusion

Gibbs-With-Gradients (GWG) uses gradients to construct Metropolis–Hastings proposals and outperforms many baseline and gradient-relaxation samplers across diverse distributions. The appendices give supporting analyses and experiments, including a multi-sample variant and results for RBMs and Ising models.

9. Future Directions and Conclusion

  • GWG considerably outperforms baseline samplers, including many that exploit known target structure, and prior discrete samplers using continuous relaxations.For deep energy-based models, it enables training on discrete data and outperforms other deep generative models.Source: §9 · p. 9
  • The experiments used window size 1; larger windows may improve the method but require more efficient proposal-sampling algorithms.The authors leave exploration of larger windows for future work.Source: §9 · p. 9

A. Balancing Locally-Informed Proposals

  • Setting the locally informed proposal temperature to τ = 2 cancels the most variable acceptance-rate terms, leaving a ratio of proposal normalizing constants.With a small Hamming ball and a Lipschitz function, nearby states have similar normalizing constants and acceptance is high.Source: A. Balancing Locally-Informed Proposals · p. 12

B. Proof of Theorem 1

  • The proof establishes a lower bound on the gradient-informed transition kernel relative to the locally informed kernel, then applies Zanella’s (2020) theorem.The bound is shown by separately bounding the corresponding minimum terms.Source: B. Proof of Theorem 1 · p. 12, §B.4 · p. 13

C. Relationship to Relaxations

  • Relaxed MALA, relaxed HMC, and D-SVGD use continuous surrogates to propose discrete updates, and their performance is sensitive to hyperparameters, especially relaxation temperature.Reported baseline results are the best across tested hyperparameters; further tuning might improve them.Source: §C.1 · p. 14, §C.2 · p. 14

D.2. Multi-Sample Variant

  • In some RBM settings, the multi-sample proposal converges faster and achieves greater effective sample size than the single-sample variant.The reported comparison tests N = 3 and N = 5; the approach uses additional approximations and assumptions.Source: §D.2 · p. 15, Figure 9 · p. 15

E.2. Additional Results

F.1. Experimental Details: Sampling

  • On both 10×10 and 40×40 lattice Ising models, GWG outperforms Gibbs and the Hamming Ball sampler.The comparison evaluates effective sample size for a test statistic.Source: §F.1 · p. 16

F.2. Experimental Details: Training

Additional experiments report better Ising parameter recovery with GWG, compare protein-model results, and describe the FHMM and neural-network settings. The results include a tuned protein dataset and visualized coupling estimates, alongside architecture choices for discrete-data energy models.

F.2. Experimental Details: Training

  • The Ising training experiments vary lattice size and connectivity and compare models trained with persistent contrastive divergence using Gibbs or GWG updates.They assess recovery by RMSE between true and inferred connectivity matrices.Source: §F.2 · p. 16

F.3. Additional Results: Training Ising Models

  • GWG considerably outperforms Gibbs sampling when training Ising models.Figure 11 expands the Figure 6 results; its plotted values are log(RMSE) between learned and true coupling matrices.Source: §F.3 · p. 17, Figure 11 · p. 17

G. Factorial Hidden Markov Models

  • An FHMM’s posterior p(x|y) has no closed form, so sampling methods are used to obtain posterior samples.Its latent state consists of a product of independent Markov processes.Source: G. Factorial Hidden Markov Models · p. 17

G.1. Experimental Details

  • The FHMM comparison uses a Hamming Ball sampler with blocks of 10 latent dimensions at a randomly selected time step, exploiting known posterior block structure.That sampler uses more hard-coded model structure than the Gibbs and GWG baselines.Source: §G.1 · p. 17

H. Potts Models for Proteins

  • Potts-model training uses persistent contrastive divergence and block-ℓ1 regularization; protein contacts are defined by amino-acid pairs within 5 angstroms.The regularization strength tested as best for PLM, Gibbs, and GWG was .01.Source: H. Potts Models for Proteins · pp. 17–18

H.1. Recall on PF10018

  • On PF10018, MCMC training outperforms PLM by a smaller margin, while GWG and Gibbs perform comparably.PF10018 was used for hyperparameter tuning, so its results are reported outside the main body.Source: §H.1 · p. 18, Figure 12 · p. 18

H.2. Visualized Contacts

  • For CADH1 HUMAN, GWG’s inferred couplings most accurately match the known structure; Gibbs infers spurious couplings and PLM misses many near the diagonal.The figure contrasts the known-distance ground truth with norms of the Potts-model coupling parameters.Source: §H.2 · p. 18, Figure 13 · p. 18

I.1. Architectures

  • Categorical image pixels are represented as one-hot vectors and mapped to learned four-dimensional embeddings before entering the network.The authors chose this straightforward parameterization to demonstrate high-dimensional discrete-data EBM training, not to achieve state-of-the-art dataset results.Source: §I.1 · p. 18

I.2. Experimental Details

The deep energy-based models use replay buffers and model selection based on annealed importance sampling (AIS). The authors caution that the resulting likelihood estimates are upper bounds, and provide additional long-run samples from several datasets.

I.2. Experimental Details

  • The EBM training setup uses replay buffers of 10,000 samples for large datasets and 1,000 for smaller datasets, without reinitializing buffer samples to noise.Reinitialization was found to cause unstable training and lower likelihoods.Source: §I.2 · p. 18

I.2.1. PARTITION FUNCTION ESTIMATION WITH AIS

  • AIS underestimates the log-partition function and therefore overestimates likelihood, so the reported likelihoods are upper bounds rather than definitive evidence that one model is better.The reported estimates use 300,000 intermediate distributions; the authors say the Caltech Silhouettes estimates appear converged between 30,000 and 300,000.Source: §I.2.1 · p. 19

I.3. Additional Results

  • Additional long-run samples from the convolutional EBM are generated using an annealed Markov chain with GWG as the base MCMC sampler.The appendix shows samples for static and dynamic MNIST, Omniglot, Caltech Silhouettes, Freyfaces, and Histopathology.Source: §I.3 · p. 19, Figure 15 · p. 19, Figure 16 · p. 20

J. Preliminary Text EBM Results

The section tests whether Gibbs-With-Gradients can train an MCMC text EBM despite the cost of sampling over a 10,000-word vocabulary. The preliminary model achieves much more efficient sampling than Gibbs, though its test-set log-likelihood is slightly below an autoregressive model.

J. Preliminary Text EBM Results

  • Text EBMs offer flexible, potentially nonautoregressive modeling, but Gibbs sampling would require over 10,000 energy evaluations per step and prior NCE approaches rely on pretrained autoregressive noise models.Those NCE approaches are limited by the quality of their autoregressive models.Source: J. Preliminary Text EBM Results · p. 22
  • The proof-of-concept uses shortened Penn Tree Bank sentences with a 10,000-word vocabulary, PCD with a 1,000-sample buffer, and 20 GWG MCMC steps per iteration.Sentences longer than 20 words are removed and shorter ones are padded with an end-of-sequence token; the setting is chosen because Gibbs sampling is not feasible.Source: J. Preliminary Text EBM Results · p. 22
  • −77.16 EBM versus −74.0 autoregressive-model test-set log-likelihood; the autoregressive model slightly outperforms the EBM.Uniform-token and training-fit categorical distributions obtain −184.21 and −100.05, respectively.Source: J. Preliminary Text EBM Results · p. 22
  • Just above 70% acceptance accompanies sampling that is at least 3500 times more efficient than Gibbs, despite 200,000 possible updates per step.The 200,000 choices arise from 10,000 words across 20 positions.Source: J. Preliminary Text EBM Results · p. 22
  • The authors note that these language-modeling results are far from state of the art and identify efficient, numerically stable sampling from very large categorical distributions as future work.They suggest larger models and more tuning could improve results.Source: J. Preliminary Text EBM Results · p. 22

Related papers

Loading 2102.04509v2…