Energy-Based Models: Training and Inference

Module: Generative AI | Difficulty: Advanced

Energy Function

Score Function of EBM

Contrastive Divergence

where is obtained by steps of Gibbs sampling.

Persistent CD (CD-k with memory)

Keep a buffer of negative samples, update with MCMC.

Annealed Importance Sampling

where

import torch, torch.nn as nn

class EBM(nn.Module):
    def __init__(self, dim=784, hidden=512):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, 1))
    def energy(self, x): return self.net(x).squeeze(-1)
    def cd_loss(self, x_data, k=10, buffer=None):
        x = x_data.clone()
        for _ in range(k):
            x = x + 0.01 * torch.randn_like(x)  # simplified
        return self.energy(x_data).mean() - self.energy(x.detach()).mean()

Research Insight: EBMs are the most natural generative model because they directly model the energy landscape. The challenge is partition function estimation, which motivates score-based approaches.