Training Energy-Based Models with Non-MCMC Samplers and Efficient Temperature Estimation

Technical Summary: Training Energy-Based Models with Non-MCMC Samplers and Efficient Temperature Estimation

1. Problem Statement

Efficient sampling from Boltzmann distributions over discrete variables is a fundamental requirement for training Energy-Based Models (EBMs), such as Boltzmann machines, and for applications in statistical physics and quantum many-body systems. Conventional Markov chain Monte Carlo (MCMC) methods, including Gibbs sampling and Metropolis–Hastings, are inherently sequential and often suffer from slow mixing in high-dimensional or rugged energy landscapes, limiting their scalability.

While fast non-Markov chain Monte Carlo (non-MCMC) samplers (e.g., simulated bifurcation, quantum annealers, coherent Ising machines) have emerged as promising alternatives, their practical application to probabilistic learning is hindered by a critical limitation: the effective inverse temperature ($\beta_{\text{eff}}$) of the generated samples is generally unknown and depends on both the problem instance and sampler parameters. Without accurate estimation of $\beta_{\text{eff}}$, sampling does not correspond to the desired distribution, leading to biased gradient estimates and degraded learning performance. Furthermore, standard methods for estimating $\beta_{\text{eff}}$, such as minimizing the Kullback–Leibler (KL) divergence, are computationally prohibitive for large-scale problems.

2. Methodology

The authors propose a three-component framework to address these challenges: a novel sampler, an efficient temperature estimation method, and an adaptive learning framework.

A. Langevin Simulated Bifurcation (LSB)
The authors introduce Langevin Simulated Bifurcation (LSB), a novel Boltzmann sampler inspired by the simulated bifurcation algorithm. LSB is designed to enable fast, parallel sampling while maintaining accuracy comparable to sequential MCMC methods.

• Mechanism: LSB solves equations of motion for positions ($x_i$) and momenta ($y_i$) derived from the global energy function $E(s|u)$ using the symplectic Euler method.
• Discretization and Stochasticity: To approximate the discrete Boltzmann distribution, LSB applies two key steps at the end of each iteration:
  1. Discretization: $x_i \to \text{sgn}(x_i)$.
  2. Stochastic Initialization: $y_i \to \xi_i \sim \mathcal{N}(0, \sigma)$, where $\xi_i$ is drawn from a Gaussian distribution.
• Parallelism: Unlike Gibbs sampling, the update rules in LSB are based on matrix-vector operations, allowing for natural parallel execution.

B. Conditional Expectation Matching (CEM)
To resolve the unknown $\beta_{\text{eff}}$ problem, the authors propose Conditional Expectation Matching (CEM). This method is applicable to EBMs with exploitable conditional independence structures (e.g., Semi-Restricted Boltzmann Machines).

• Principle: In models like Semi-Restricted Boltzmann Machines (SRBMs), hidden variables are conditionally independent given the visible variables. This allows for an analytical expression of the conditional expectation $\langle h_j \rangle_r$ as a function of $\beta_{\text{eff}}$ (specifically involving the $\tanh$ function).
• Procedure: CEM estimates $\beta_{\text{eff}}$ by minimizing the squared difference between the analytically derived conditional expectations and those obtained via conditional LSB sampling.
• Advantage: Unlike KL-divergence minimization, CEM is computationally efficient, parallelizable with the sampling process, and does not scale exponentially with system size.

C. Sampler Adaptive Learning (SAL)
The authors integrate LSB and CEM into a learning framework called Sampler Adaptive Learning (SAL).

• Adaptivity: In conventional learning, the model temperature is fixed (typically $\beta=1$). In SAL, the model temperature is set to the estimated $\beta_{\text{eff}}$ of the sampler's output distribution. As model parameters $u$ are updated, $\beta_{\text{eff}}$ changes, and the model function $Q_{\beta_{\text{eff}}}$ is adaptively adjusted to match the sampler's distribution.
• Optimization: Learning proceeds via gradient descent on the KL divergence $D_{\text{KL}}(P_D || Q_{\beta_{\text{eff}}})$, utilizing the negative phase (expectation under the sampler) and the positive phase (data distribution).

3. Key Contributions and Results

The proposed methods were validated primarily on Semi-Restricted Boltzmann Machines (SRBMs), a class of EBMs that extend Restricted Boltzmann Machines (RBMs) by adding pairwise connections among visible units. SRBMs are more expressive but difficult to train with conventional methods due to the loss of conditional independence in visible units.

• Sampling Performance (LSB vs. Gibbs):

  • On random SRBMs ($N_v=10, N_h=5$), LSB achieved sampling accuracy (measured by KL divergence) comparable to or better than Gibbs sampling.
  • LSB provided orders-of-magnitude acceleration over Gibbs sampling while maintaining similar memory usage, attributed to its parallelizable nature versus the sequential nature of Gibbs updates.
  • LSB demonstrated robust performance on the Sherrington–Kirkpatrick model and larger system sizes.

• Temperature Estimation (CEM):

  • CEM enabled accurate estimation of $\beta_{\text{eff}}$ for LSB outputs with negligible computational overhead.
  • The mean signed relative error between CEM estimates and those obtained via KL minimization was approximately 3.6%.
  • CEM successfully handled the instance-dependent variability of $\beta_{\text{eff}}$ that plagues non-MCMC samplers.

• Learning Performance (SAL):

  • 3-Spin Model: On synthetic spin-glass datasets, SRBMs trained with SAL minimized the training cost function ($D_{\text{KL}}$) significantly better than FBMs, RBMs, and RBMs trained with Contrastive Divergence (CD-100). This highlights the advantage of SAL in modeling rugged energy landscapes.
  • Bars-and-Stripes (BAS) Dataset: SAL-trained SRBMs demonstrated high-fidelity image generation and reconstruction. The models learned the underlying distribution rather than memorizing samples, achieving near-perfect reconstruction (0.5% error) even with 47.6% of pixels masked.
  • OptDigits Dataset: On handwritten digit images ($N_v=74$), SAL-trained SRBMs achieved high-quality unconditional and conditional generation. Furthermore, the models achieved ~90% classification accuracy on the test set, rapidly converging from a random baseline.

4. Significance and Claims

The paper positions these results as a proof of principle for the effectiveness of combining fast non-MCMC sampling with accurate temperature estimation via CEM. The authors explicitly state that the work does not claim immediate competitiveness with state-of-the-art generative AI but rather aims to bridge the gap between fast non-MCMC samplers and their practical applications.

Key claims include:

• LSB is established as a fast and accurate Boltzmann sampler that overcomes the sequential bottlenecks of MCMC.
• CEM provides a scalable solution to the effective temperature estimation problem, enabling the use of non-MCMC samplers in probabilistic learning.
• SAL enables the efficient training of more expressive models (like SRBMs) that are traditionally difficult to train, offering performance advantages on synthetic spin-glass datasets and real-world image tasks.
• The framework is not limited to SRBMs; it is applicable to any EBM with exploitable conditional independence structures, including models with higher-order interactions and potential applications in neural quantum states and energy-efficient hardware implementations.

The authors acknowledge open questions regarding the theoretical basis for LSB's approximation of discrete Boltzmann distributions and the need for systematic hyperparameter selection strategies, identifying these as directions for future research.