Source-linked AI summary
Improved Contrastive Divergence Training of Energy Based Models
Yilun Du, Shuang Li, Joshua Tenenbaum, Igor Mordatch
TL;DR
Contrastive divergence training for energy-based models is difficult to stabilize, limiting architecture compatibility and scalability. The paper estimates a usually omitted gradient term and combines it with data augmentation and multi-scale processing. The resulting framework stabilizes training and improves performance across image generation, OOD detection, and compositional generation.
Problem
Contrastive divergence training of energy-based models suffers from instabilities that limit modeling power, modern-architecture compatibility, and scalable long-running training.
Method
The paper efficiently estimates the neglected gradient term and supplements contrastive divergence with data augmentation and multi-scale processing.
Results
The framework improves performance across image generation, OOD detection, and compositional generation while enabling more stable training.
Takeaways & Limitations
The approach provides a simple framework for high-resolution compositional generation and out-of-distribution robustness.
Takeaways & Limitations
Entropy estimation becomes difficult in high dimensions, so it is used as a fast regularizer rather than estimated accurately.
Abstract
from arXiv · showhide
Contrastive divergence is a popular method of training energy-based models, but is known to have difficulties with training stability. We propose an adaptation to improve contrastive divergence training by scrutinizing a gradient term that is difficult to calculate and is often left out for convenience. We show that this gradient term is numerically significant and in practice is important to avoid training instabilities, while being tractable to estimate. We further highlight how data augmentation and multi-scale processing can be used to improve model robustness and generation quality. Finally, we empirically evaluate stability of model architectures and show improved performance on a host of benchmarks and use cases,such as image generation, OOD detection, and compositional generation.
1 Introduction
Energy-based models support diverse applications, but contrastive divergence training is unstable and often incompatible with modern architectures. The paper addresses this by retaining the method’s simplicity while estimating a neglected gradient term and using augmentation and multi-scale processing to improve stability, mixing, and generation.
- Energy-based models have been applied to image generation, robustness, compositional generation, memory, language, video, reinforcement learning, continual learning, and protein design.
- Contrastive divergence trains EBMs by lowering energy on training data and raising it on MCMC-generated confabulations.Its appeal comes from simplicity, extensibility, and avoiding auxiliary networks.
- Training instability has required gradient clipping, spectral normalization, parameter tuning, early stopping, or avoiding self-attention and layer normalization.These requirements limit modeling power, architecture compatibility, and long-running training.
- The neglected gradient term is tractable to estimate, contributes significantly to the training gradient, and stabilizes training.The procedure enables self-attention, removes capacity-limiting spectral normalization, and supports longer training without introducing a new objective.
- Data augmentation encourages MCMC mixing and diversity, while multi-scale energy evaluation improves spatial coherence in generated samples.
- The paper evaluates architecture stability and reports improved performance across image generation, OOD detection, and compositional generation.
2 An Improved Contrastive Divergence Framework for Energy-Based Models
The framework corrects contrastive divergence by retaining and estimating the often-ignored KL divergence gradient term, then improves sampling through entropy regularization, data augmentation, and multi-scale energies.
- EBM maximum-likelihood training is difficult because computing the partition function requires integration over the entire input domain.
- Contrastive divergence replaces the intractable model expectation with finite MCMC transitions, but those transitions may mix slowly or fail to reach the model distribution.
- A Missing Term in Contrastive Divergence: Changes to the energy function alter MCMC samples, introducing a KL divergence gradient term that prior methods commonly ignored or estimated with high variance.
- Estimating the missing gradient term: The full objective combines traditional contrastive loss with KL loss, whose energy and entropy components encourage low-energy, diverse generated samples and improve training stability.The energy component is optimized by differentiating through Langevin sampling, while entropy is estimated with nearest-neighbor distances.
- Estimating the missing gradient term: Nearest-neighbor entropy estimation scales poorly in high dimensions, but here it serves as a fast regularizer intended to prevent sampling collapse rather than accurately estimate entropy.
- Data Augmentation Transitions: Periodic data-augmentation transitions improve exploration by moving chains between related inputs, while multi-scale energies sum functions across image resolutions.The training procedure applies augmentation to replay-buffer samples and propagates gradients through sampling to compute the KL loss.
3 Experiments
Experiments test the proposed EBM components across generation, sampling, training stability, OOD detection, and compositionality. The results show that KL loss improves stability, data augmentation improves mixing and diversity, and the framework performs competitively across several benchmarks.
- Experimental Setup: The experiments evaluate component effects and downstream performance on image generation, OOD detection, and compositionality across CIFAR-10, ImageNet32x32, and CelebA-HQ.Models were trained with Adam on Volta GPUs, with detailed configurations in the appendix.
- Image Quality: The approach outperforms many prior EBMs on CIFAR-10, is close to SSGAN on CelebA-HQ, and matches PixelIQN on ImageNet32x32 with about one tenth the parameters.The authors note that overall generative performance is not the best reported, while improving existing EBM performance.
- Data Augmentation: Data augmentation transitions allow Langevin chains to mix across image modes and produce diverse samples from the same initialization, unlike unaugmented chains that converge to one face.The sampling chains use augmentation transitions every 60 steps in the illustrated setup.
- Mode Convergence: With KL loss and data augmentation, very low-temperature samples retain naturalistic modes on CIFAR-10 and CelebA-HQ, unlike IGEBM samples that develop odd hues or textures.The comparison follows 1000 gradient-descent steps on the energy function.
- Stability/KL Loss: The KL loss significantly improves training stability and enables stable use of self-attention, whereas spectral normalization stabilizes training but reduces Inception-score improvement.Stable training is associated with energy differences near zero between real and generated images.
- Downstream Applications: The framework significantly outperforms most OOD baselines, generates conditioned attributes with higher resolution, and achieves higher attribute recovery than IGEBM and JVAE across compositions.The OOD exception is CIFAR-10 interpolations; compositional attributes include face and object factors.
4 Related Work
Related work situates the approach within contrastive-divergence EBM training, alternative negative-sample strategies, and alternative objectives. The paper distinguishes its KL-gradient estimator from prior high-variance estimation methods.
- Contrastive Divergence: Contrastive divergence trains EBMs by contrasting negative samples from a model distribution with real data, but negative-sample generation is a central difficulty.Prior work also includes image-domain applications and multiscale approaches.
- Alternative Training Strategies: Auxiliary networks have been used to amortize the negative-sampling procedure, while other work replaces contrastive divergence with score matching, denoising, noise contrastive estimation, Stein discrepancy, or learned f-divergences.These approaches represent distinct strategies for training EBMs.
- KL-Term Estimation: Most prior contrastive-divergence work ignores the KL term; Ruiz and Titsias estimate a related term with high-variance REINFORCE, whereas this approach uses autodifferentiation and nearest-neighbor entropy estimation.The comparison concerns how the KL-gradient contribution is estimated.
5 Conclusion
The paper presents a framework intended to make EBM training easier while improving generation, OOD robustness, and compositional image synthesis. It reports high-resolution compositional generation and distinct samples from the CelebA-HQ training set.
- Conclusion: The framework is presented as a simple, general approach for improving EBM generation and training ease, including high-resolution compositional generation and OOD robustness.The authors identify text, video, and reasoning as future application domains.
- Conclusion: Unconditional CelebA-HQ samples generated by the trained EBM are reported to be distinct from the training set under an L2 nearest-neighbor comparison.The comparison uses nearest neighbors in the CelebA-HQ training dataset.
A.2 Additional Quantitative Results
The paper reports stronger unconditional image quality than SNGAN on LSUN bedrooms and globally coherent generations across LSUN, CIFAR-10, and ImageNet.
- FID is 33.46 for the proposed approach versus 64.05 for SNGAN on LSUN 128x128 bedroom scenes.SNGAN was re-implemented with default hyperparameters for the comparison.
- Unconditional generations are described as relatively globally coherent on LSUN bedrooms, CIFAR-10, and ImageNet 32x32.
B.1 Model Architectures
The experiments use multiscale energy functions whose outputs combine full-, half-, and quarter-resolution processing, with specified architectural reductions at lower resolutions.
- The final multiscale energy is the sum of energy functions applied at full, half, and quarter image resolutions.
- Half-resolution models remove the first down-sampled residual block, while quarter-resolution models remove the first two.
- Each residual block uses group normalization and the Swish nonlinearity.
B.2 Experiment Configurations For Different Datasets
Dataset-specific configurations use 40-step Langevin sampling, with step sizes and augmentation resize ranges adjusted between CIFAR-10/ImageNet and CelebA-HQ/LSUN.
- CIFAR-10/ImageNet 32x32: CIFAR-10 and ImageNet 32x32 use 40 Langevin steps with step size 500 and Gaussian noise magnitude 0.001.
- Both configurations use color augmentation, random horizontal flips, and image resizing, with resize ranges of 0.02–1.0 or 0.08–1.0.
- CelebA/LSUN Bedroom: CelebA-HQ and LSUN use 40 Langevin steps with step size 1000 and Gaussian noise magnitude 0.001.
C Loss Gradient Derivation
The derivation shows that the full contrastive-divergence gradient includes the additional LKL contribution, whose use improves training stability, including with random-noise initialization.
- The gradient of the full contrastive-divergence objective is equivalent to the gradient of LFull = LKL + LCD.
- The full gradient combines the derivatives of the first and second KL terms in the contrastive-divergence objective.
- The derived gradient is identified as equal to the full contrastive-divergence gradient LCD Full.
- With Langevin dynamics initialized from random noise, the approach again improves training stability.
D.2 Analysis of Truncated Langevin Backpropagation
Truncating backpropagation through Langevin sampling produces similar MNIST performance while substantially reducing training time. Backpropagation through all sampling steps is over three times slower to train.
- FID is 90.54 with backpropagation through only the final Langevin step, versus 94.85 with backpropagation through all steps.
- Truncating backpropagation to one Langevin step causes little degradation in performance.
- Backpropagation through all Langevin steps is over 3 times slower to train than backpropagation through only the last step.
D.3 Analysis of Effect of KL Loss on Mode Sampling
The KL loss acts as a repelling regularizer that encourages low-energy, diverse MCMC samples and prevents the sampling collapse observed without it. The paper illustrates this contrast across generated-image examples and collapse visualizations.
- LKL encourages EBM MCMC samples to have both low energy and diversity by acting as a repelling term.
- Without LKL, EBM sampling always collapses and eventually produces significantly less diverse samples.
- With LKL, training does not suffer from sampling collapse, as illustrated by the generated samples and collapse comparison.