AI Research Portal

This paper introduces DREXEL (Discrete Replica EXchangE Langevin) and DREAM (Discrete Replica Exchange with Adjusted Metropolis), two novel samplers designed for efficient exploration of non-convex discrete energy landscapes. The core idea is to combine Gradient-based Discrete Samplers (GDSs) with replica exchange, employing two GDSs at different temperatures and step sizes—one for local exploitation and another for broader exploration. A tailored swap mechanism ensures detailed balance in discrete sampling. The paper claims that these samplers satisfy detailed balance and converge to the target distribution under mild conditions, demonstrating superior performance across 2D synthetic simulations, Ising models, Restricted Boltzmann Machines, and deep energy-based models.

Stagnation in complex, non-convex landscapes: Gradient-based Discrete Samplers (GDSs) are effective for discrete energy landscapes but often get trapped in local minima in complex, non-convex settings, limiting their exploration capabilities.
: Despite improvements in Locally Balanced Proposals (LBPs), achieving a balance between exploring the broader energy landscape and exploiting local modes remains a challenge.

A novel integration of GDS with replica exchange to improve non-convex exploration.
A swap mechanism tailored for detailed balance and sample efficiency in discrete sampling.
A novel theoretical convergence analysis based on a spectral framework without smoothness assumptions.
Superior performance in synthetic tasks, Ising models, restricted Boltzmann machines, and deep energy-based models for exploring non-convex discrete energy landscapes.

The proposed samplers, DREXEL and DREAM, combine Gradient-based Discrete Samplers (DLS) with replica exchange Markov Chain Monte Carlo (reMCMC) to efficiently explore non-convex discrete spaces.

Sampler Design: Two DLSs are employed in parallel, each with distinct temperatures and step sizes:

A low-temperature sampler (τ1, α1) focuses on local exploitation.
A high-temperature sampler (τ2, α2) explores broader energy landscapes, with τ1 < τ2 and α1 < α2. Larger step sizes and higher temperatures encourage more exploratory moves.

Discrete Langevin Sampler (DLS) Update: The transition probability for updating a coordinate d for sampler k is given by: θ(k)i+1,d ~ Cat [Softmax (∇U(θ(k)i,d)(θ(k)i+1,d−θ(k)i,d)/(2τk) - (θ(k)i+1,d−θ(k)i,d)²/(2αk))] (6) This update guides the sampler towards lower-energy states while penalizing large jumps.

Metropolis-Hastings Correction (for DREAM): To correct for bias and ensure convergence, DREAM applies a Metropolis-Hastings (MH) correction. After generating a proposal state ω from θi, it is accepted with probability A(ω, θi): A(ω, θi) = min {1, π(ω)q(θi|ω) / (π(θi)q(ω|θi))} (4) where q(·) is the transition probability from (2). These MH steps are performed independently for each sampler before any swap.

Sample Swap Mechanism: Periodically, a swap between the states of the two chains is proposed. The swap function tailored for discrete sampling is: S(θ(1)i+1, θ(2)i+1 | θ(1)i, θ(2)i) = exp (−1/τ1 − 1/τ2) [ΔUi+1−ΔUi] (8) where ΔUi := U(θ(1)i) – U(θ(2)i). This mechanism ensures detailed balance with respect to the joint distribution over replica states by incorporating energy estimates at previous samples. The swap occurs if a uniformly drawn random number u is less than p min {1, S(θ(1)i+1, θ(2)i+1 | θ(1)i, θ(2)i)}, where p > 0 is the swap intensity.

Algorithm Flow:

Initialize two samples θ(1) and θ(2) with distinct step sizes α1, α2 and temperatures τ1, τ2.
For a specified number of iterations I:
- For each sampler k (low-temp and high-temp):
  - Construct proposal distributions qk(θ(k)i+1 | θ(k)i) using equation (6).
  - Sample θ(k)i+1 from qk(· | θ(k)i).
- (For DREAM only) For each sampler k:
  - Compute acceptance probability A(θ(k)i+1, θ(k)i) using equation (4).
  - Accept or reject the proposed sample based on A.
- Compute the swap function S(θ(1)i+1, θ(2)i+1 | θ(1)i, θ(2)i) using equation (8).
- Swap θ(1)i+1 and θ(2)i+1 with probability p min {1, S}.
Output the samples from the low-temperature chain.

Target Distribution: π(θ) = exp [U(θ)/τ]/Z, ∀θ∈ Θ (1), where Θ is a high-dimensional, finite, discrete domain, e.g., binary {0,1}d or categorical {0,..., N-1}d. U is the energy function, τ is temperature, and Z is the intractable normalization constant.
DLS Update Equation: θ(k)i+1,d ~ Cat [Softmax (∇U(θ(k)i,d)(θ(k)i+1,d−θ(k)i,d)/(2τk) - (θ(k)i+1,d−θ(k)i,d)²/(2αk))] (6).
Metropolis-Hastings Acceptance Probability: A(ω, θi) = min {1, π(ω)q(θi|ω) / (π(θi)q(ω|θi))} (4).
Discrete Swap Function: S(θ(1)i+1, θ(2)i+1 | θ(1)i, θ(2)i) = exp (−1/τ1 − 1/τ2) [ΔUi+1−ΔUi] (8).
Bias Correction (bDREXEL and bDREAM variants): A bias correction term S(θ(1)i+1, θ(2)i+1) = exp (−1/τ1 − 1/τ2) [ΔUi+1+(σi)²] (7) is considered, where ΔUi := U(θ(1)i) – U(θ(2)i) and σi² compensates for stochastic gradient due to mini-batch data.
Optimization (EBMs): Adam optimizer with a learning rate = 0.001 for 20,000 iterations, batch size of 256.
Energy-Based Model Architecture: ResNet-64 backbone.
EBM Training: Intractable likelihood gradient approximated through Persistent Contrastive Divergence. A replay buffer containing 1,000 past samples is implemented.
EBM Inference: Annealed Importance Sampling is conducted with DULA to estimate test log-likelihoods.

2D Synthetic Problems
- Datasets: wave, 8gaussians, 16gaussians, moon, 2moon, twist, flower. Domain: Θ = {1, 2, ..., N}d, where N = 256 and d = 51 × 51.
- Metrics: Kullback-Leibler (KL) divergence (10⁻²), Maximum Mean Discrepancy (MMD) (10⁻⁴), Negative Log-Likelihood (NLL) (10⁻³).
- Baselines: DLS (DULA and DMALA), Any-scale Balanced sampling (AB), Automatic Cyclical Sampler (ACS).
- Numeric Results:
  - KL (10⁻²): DREAM consistently achieves the lowest KL divergence across all tasks. E.g., for 'wave', DREAM: 4.975±0.222 vs. DMALA: 10.891±2.408. For 'moon', DREAM: 0.500±0.119 vs. DMALA: 2.921±0.531.
  - MMD (10⁻⁴): DREAM consistently achieves the lowest MMD across all tasks. E.g., for 'wave', DREAM: 6.211±0.235 vs. DMALA: 20.409±1.444. For 'moon', DREAM: 1.394±0.387 vs. DMALA: 4.309±0.840.
  - NLL (10⁻³): DREAM consistently achieves the lowest NLL across all tasks. E.g., for 'wave', DREAM: 8.419±0.010 vs. DMALA: 7.937±0.013. For 'flower', DREAM: 6.499±0.004 vs. DMALA: 6.801±0.015.
- Qualitative Results: DREAM effectively captures all modes of complex distributions (wave, 8gaussians, 16gaussians), while DMALA captures only 50% of modes and gets stuck in local minima.
Sampling from Ising Models
- Dataset: 2D Ising model (θ∈ {−1,1}d) with energy U(θ) = wθT Jθ + bTθ. Connectivity strength w = 0.15.
- Metric: log Root Mean Square Error (log RMSE) of marginals.
- Baselines: DULA, ACS, DREXEL, DMALA, mh-ACS.
- Numeric Results: DREAM achieves the lowest log-RMSE: -4.884±0.204. DMALA: -4.595±0.133. DREXEL: -2.544±0.033. DULA: -1.769±0.022.
- Runtime/Training Cost: Performance quantified across 50,000 sample sizes and computational budgets (t = 50.0 seconds).
Learning Energy-Based Models
- Datasets: Static MNIST, Dynamic MNIST, EMNIST, kMNIST, Fashion, Omniglot, Caltech Silhouettes.
- Metric: MMD (for RBMs), Test log-likelihoods (for EBMs).
- Baselines: DMALA, ACS, AB (for RBMs); DULA, DMALA, bDREXEL, bDREAM (for EBMs).
- Numeric Results (RBMs - MMD): DREAM consistently achieves the lowest MMD values across all datasets. E.g., for MNIST, DREAM: -6.349±0.061 vs. DMALA: -6.248±0.059. For Omniglot, DREAM: -6.626±0.057 vs. DMALA: -6.536±0.099. DREAM reduces MMD by 9.7% for MNIST, and 8.2% (eMNIST) to 12.3% (Omniglot) compared to DMALA.
- Numeric Results (EBMs - Test log-likelihoods): DREAM consistently achieves the highest log-likelihoods across all datasets. E.g., for Static MNIST, DREAM: -83.711±0.930 vs. DULA: -84.641±0.407. For Omniglot, DREAM: -98.525±1.782 vs. DULA: -118.627±5.402. DREAM outperforms baselines by 11.9–14.7% on Omniglot and Caltech Silhouettes.

Challenges with decaying step sizes in discrete sampling: In discrete spaces, small steps do not equate to gradual movements and can lead to repeated proposals of nearly identical samples, causing the sampler to become trapped in local regions.
Difficulty for high-temperature samplers in exploiting regions: High-temperature samplers may struggle to exploit certain regions due to abrupt exploration, requiring excessive time to fully characterize local modes and achieve mixing.
Suboptimal performance of bias-corrected variants: The bias-corrected variants, bDREXEL and bDREAM, generally performed worse across datasets compared to DREXEL and DREAM, indicating that MH corrections are essential for improving sampling fidelity in discrete tasks.

Method evidence: "To bridge this gap, we propose two samplers: Discrete Replica EXchangE Langevin (DREXEL) and Discrete Replica Exchange with Adjusted Metropolis (DREAM)." (section: Introduction)
Result evidence: "As shown in Table 1, DREAM consistently outperforms the baselines across most metrics and distributions." (section: Sampling from 2D Synthetic Problems)
Limitation evidence: "It should be noted that while decaying step sizes are commonly advantageous in Langevin MCMC for handling big data (Teh, Thiéry, and Vollmer 2016), they present potential challenges in discrete sampling. In discrete spaces, small steps do not equate to gradual movements as they do in continuous spaces. Instead, they tend to repeatedly propose nearly identical samples, which causes the sampler to become trapped in local regions." (section: Discrete Samplers with Different Temperatures and Step Sizes)

Exploring Non-Convex Discrete Energy Landscapes: An Efficient Langevin-Like Sampler with Replica Exchange

Summary

Problem

Key Contributions

Method

Technical Details

Results

Weaknesses / Limitations

Evidence

Additional Figures