14 (Bonus). Energy-Based Models and Diffusion LLMs

14 (Bonus). Energy-Based Models and Diffusion LLMs#

⚡Compute Note: You can run this notebook on CPU.

Large language models (LLMs) today are dominated by autoregressive Transformers, but the landscape is rapidly expanding. Two families of models are becoming increasingly important in modern AI research:

Diffusion Models
Energy-Based Models (EBMs)

Although they emerged from different motivations, they share a deep conceptual connection: both model data through iterative refinement rather than one-shot prediction.

In this notebook, we will build both of them from scratch using simple 2D toy data and then extend the same principles toward language modeling.

🔄 Diffusion Models: Learning to Reverse Noise#

A diffusion model learns how to undo a progressive noising process.

The forward process gradually destroys structure by adding Gaussian noise: \(q(x_t \mid x_{t-1}) = \mathcal{N}\!\big(\sqrt{1 - \beta_t}\,x_{t-1},\,\beta_t I\big)\).
The reverse process is learned: \(p_\theta(x_{t-1} \mid x_t)\) so the model can denoise step-by-step and recover the original data.

Why diffusion is popular:

✔ State-of-the-art generative modeling for images, audio, 3D, video
✔ Stable training compared to GANs
✔ Supports iterative reasoning, making it attractive for text and LLM research
✔ Enables parallel generation in language models (non-autoregressive)

Diffusion models are now used in:

Image models like Stable Diffusion, DALL·E 3, Imagen
Speech synthesis (e.g., Grad-TTS)
Emerging diffusion LLMs (EDLMs, DLMs, Plaid, D3PM)

To know more, read the paper by Nie et al. (2025) on diffusion LLMs.

⚡ Energy-Based Models (EBMs): Learning an Energy Landscape#

Energy-based models do not predict pixels or tokens directly. Instead, they learn a scalar energy function \(E_\theta(x)\) such that:

Low energy → likely / desirable
High energy → unlikely / undesirable

They define unnormalized probability models:

\[ p_\theta(x) \propto e^{-E_\theta(x)}. \]

Sampling is done by optimizing the input using gradient descent:

\[ x \leftarrow x + \eta \nabla_x E_\theta(x), \]

which pulls samples toward regions of high energy.

EBMs are popular because:

✔ They can model global constraints (e.g., logical consistency in text)
✔ Easy to combine with diffusion models — diffusion can be viewed as a special EBM
✔ Great for reasoning, control, and structured generation
✔ Used in recent research to improve LLM consistency and correct errors

I highly recommend reading the great paper by Gladstone et al. (2025) on energy based language models.

🌉 Why Study Both Together?#

Recent breakthroughs show that diffusion models can be interpreted as EBMs, and EBMs can refine or guide diffusion sampling.

This leads to powerful hybrid approaches:

Energy-based Diffusion LMs (EDLM): Improve logical correctness and global consistency in text generation. Read more on this in the paper by Xu et al. (2025).
Diffusion of Thoughts: Treat diffusion steps as internal reasoning iterations
Non-autoregressive parallel text generation: Faster generation by replacing sequential decoding with iterative refinement

Understanding both gives insight into:

How modern generative models work beneath the surface
How LLMs may evolve beyond purely autoregressive Transformers
How iterative refinement can improve reasoning and reduce hallucinations

We’ll keep everything tiny and visual:

Start in a 2D toy world with only three points:
- Build a diffusion model that starts from a random 2D vector and iteratively moves towards one of the three points.
- Build an energy-based model that learns an energy landscape and generate samples by optimizing the input.
Scale the same ideas to language:
- Build a tiny autoregressive LLM baseline.
- Turn it into a discrete diffusion language model (DLM) that denoises corrupted text over multiple steps.
- Add a sequence-level energy head to get a tiny energy-based diffusion LLM.
- Connect diffusion iterations to thinking / reasoning steps.

The goal is not state-of-the-art performance, but a hands-on understanding of the mechanics of diffusion and energy-based language models.

import math
import time
from dataclasses import dataclass
from typing import Optional, List, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

import matplotlib.pyplot as plt
import numpy as np

torch.manual_seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device

device(type='cpu')

1️⃣ A Tiny 2D Dataset with Three Points#

We start with the simplest possible dataset:

Only three points in 2D – think of them as three “modes” or “clusters” of our data distribution. Consider these as representatives of the sparse points in the space of natural language texts (sequenes of tokens) or images (sequences of pixels) that make sense out of the all possible token or pixel sequences, majority of which would be gibberish or noisy images.

I highly recommend watching the video by Julia Turc on diffusion models.

We’ll:

Add Gaussian noise in a forward diffusion process (data → noise).
Train a tiny MLP to predict the noise at each timestep.
Run the learned reverse process (noise → data) to sample points that land near these three modes.

# Three data points in 2D space
X_data = torch.tensor([
    [0.0, 0.0],
    [2.0, 2.0],
    [-2.0, 2.0],
], device=device)

def sample_data(batch_size: int) -> torch.Tensor:
    idx = torch.randint(0, X_data.shape[0], (batch_size,), device=device)
    return X_data[idx]

X_data

tensor([[ 0.,  0.],
        [ 2.,  2.],
        [-2.,  2.]])

2️⃣ Forward Diffusion: From Clean Points to Noise#

We’ll use a simple variance schedule where we keep adding noise to the modes in the dataset:

At timestep \(t\), we apply some noise controlled by \(\beta_t\).
We define:

\[ q(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}) = \mathcal{N}\!\left( \mathbf{x}_{t}; \sqrt{1 - \beta_t}\,\mathbf{x}_{t-1}, \beta_t\,\mathbf{I} \right). \]

This means:

Each diffusion step only depends on the previous state \(x_{t-1}\).
Noise is added incrementally, not all at once.
To sample \(x_t\), we must iterate the transition kernel from step 0 → t.

Below we implement the exact step-wise transition
(and a helper that iterates it from \(x_0\) to \(x_t\)), fully matching the Markov definition of DDPM forward diffusion.

T = 100  # number of diffusion steps

# Simple linear beta schedule
beta_start, beta_end = 1e-8, 0.05
betas = torch.linspace(beta_start, beta_end, T, device=device)
alphas = 1.0 - betas
alpha_bars = torch.cumprod(alphas, dim=0)


def q_step(x_prev: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None):
    """
    Single diffusion step:
        q(x_t | x_{t-1}) = N(x_t; sqrt(1 - beta_t) * x_{t-1}, beta_t * I)

    x_prev: (B, D) = x_{t-1}
    t:      (B,)   = integer timestep in [0, T-1] (same for all in batch is typical)
    """
    if noise is None:
        noise = torch.randn_like(x_prev)

    # beta_t for each element in the batch
    beta_t = betas[t].view(-1, 1)          # (B, 1)
    mean = torch.sqrt(1.0 - beta_t) * x_prev
    std = torch.sqrt(beta_t)

    return mean + std * noise              # x_t


def q_sample(x0: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None):
    """
    Sample x_t by *iterating* the Markov chain from x_0 up to time t,
    using q_step at each intermediate step.

    This matches the step-wise definition:
        q(x_t | x_{t-1}) = N( sqrt(1 - beta_t) x_{t-1}, beta_t I )
    """
    # Ensure we have one trajectory per element in the batch
    x = x0
    B = x0.size(0)

    # If a single t is passed, broadcast it; otherwise assume per-sample timesteps
    if t.dim() == 0:
        t = t.repeat(B)
    assert t.shape[0] == B

    # For efficiency, we could still use the closed-form, but here we explicitly follow
    # the step-wise process to mirror the equation you gave.
    max_t = int(t.max().item())
    for step in range(max_t + 1):
        # Only update those samples whose target t >= current step
        mask = (t >= step)
        if not mask.any():
            break

        x_prev = x[mask]
        t_batch = torch.full(
            (x_prev.size(0),), step, device=x0.device, dtype=torch.long
        )
        x_new = q_step(x_prev, t_batch)

        x = x.clone()
        x[mask] = x_new

    return x

Visualizing Forward Diffusion Rollouts#