Visualizing Forward Diffusion on Images

The previous lesson treated a diffusion schedule as a map from timestep to signal-to-noise ratio. We now turn that description into something concrete: a small PyTorch component that takes clean, normalized images and produces noisy images at any requested diffusion time.

By the end, you will be able to (1) preprocess an image batch on the correct scale, (2) precompute a linear or cosine schedule, (3) sample ${\mathbf{x}}_t$ directly from ${\mathbf{x}}_0$, and (4) visualize both a set of noisy marginals and a genuine sequential forward trajectory.

---

The computational object we are building

The forward process has two equivalent views.

The one-step Markov transition is

$$q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})=\mathcal{N}\left(\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I}\right).$$

A literal trajectory repeatedly draws a fresh Gaussian noise tensor at every step.

For implementation and training, the more useful form is the closed-form marginal:

$$q({\mathbf{x}}_t\mid {\mathbf{x}}_0)=\mathcal{N}\left(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,(1-{\bar{\alpha }}_t)\mathbf{I}\right),$$

so that one can sample directly:

$${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ},\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I}).$$

This direct form is central to DDPM training: the model does not need to simulate all earlier corruptions before receiving a sample at timestep $t$.

Two implementation conventions are common:

• Training-oriented indexing: store $T$ values indexed from $0$ to $T-1$, where index $0$ already means “after one small noising step.”
• Visualization-oriented indexing: explicitly include $t=0$ as the clean image, so that $t$ ranges from $0$ to $T$.

We will use the second convention because it makes plots easier to interpret. Thus, ${\bar{\alpha }}_0=1$, and q_sample(x0, t=0) returns x0 exactly.

---

Study the reference implementation

The following reading supplies a useful implementation pattern: compute one-dimensional schedule tensors once, then use a batch-aware extract function to reshape the coefficients so they multiply a batch of images correctly.

---

Data scaling is part of the model definition

Diffusion noise is normally sampled from a standard normal distribution. Therefore, an image tensor with values in $[0,1]$ is poorly matched to the added noise: its mean is positive, and its dynamic range is asymmetric around zero.

For a standard pixel-space DDPM, map RGB images from $[0,1]$ to $[-1,1]$:

$${\mathbf{x}}_0=2\mathbf{u}-1,$$

where $\mathbf{u}$ is the image returned by ToTensor().

from torchvision import transforms

image_size = 64

transform = transforms.Compose([
    transforms.Resize(image_size),
    transforms.CenterCrop(image_size),
    transforms.ToTensor(),                 # [0, 1], shape [C, H, W]
    transforms.Lambda(lambda x: 2.0 * x - 1.0),  # [-1, 1]
])

Suppose loader is a DataLoader whose batches are image tensors. Select one image for visualization:

import torch

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

x0 = next(iter(loader))[0:1].to(device)  # shape [1, C, H, W]

assert x0.ndim == 4
assert x0.shape[0] == 1
print(x0.shape, x0.min().item(), x0.max().item())

For a dataset with different batch formats, you may need to extract the image field first, such as batch["image"]. The essential contract is that x0 has shape

$$[B,C,H,W].$$

Do not clamp x0 before noising. The clean normalized image should generally lie near $[-1,1]$, but the noised variable ${\mathbf{x}}_t$ has Gaussian support and can naturally exceed that range.

---

A minimal, inspectable forward diffuser

Here is a compact implementation that supports both the original linear beta schedule and the cosine schedule from the preceding lesson.

import math
import torch


def linear_betas(num_steps, beta_start=1e-4, beta_end=2e-2):
    return torch.linspace(beta_start, beta_end, num_steps)


def cosine_betas(num_steps, s=0.008):
    # Define cumulative signal retention first, then recover betas.
    t = torch.linspace(0, num_steps, num_steps + 1) / num_steps

    alpha_bar_target = torch.cos(
        ((t + s) / (1.0 + s)) * math.pi / 2.0
    ) ** 2
    alpha_bar_target = alpha_bar_target / alpha_bar_target[0]

    betas = 1.0 - alpha_bar_target[1:] / alpha_bar_target[:-1]
    return betas.clamp(max=0.999)


class ForwardDiffuser:
    def __init__(self, num_steps=1000, schedule="cosine", device="cpu"):
        if schedule == "linear":
            betas = linear_betas(num_steps)
        elif schedule == "cosine":
            betas = cosine_betas(num_steps)
        else:
            raise ValueError(f"Unknown schedule: {schedule}")

        self.num_steps = num_steps
        self.betas = betas.to(device)

        # One-step signal retention.
        self.alphas = 1.0 - self.betas

        # alpha_bar[0] corresponds to the clean image x_0.
        self.alpha_bar = torch.cat([
            torch.ones(1, device=device),
            torch.cumprod(self.alphas, dim=0),
        ])

        self.sqrt_alpha_bar = torch.sqrt(self.alpha_bar)
        self.sqrt_one_minus_alpha_bar = torch.sqrt(1.0 - self.alpha_bar)

    @staticmethod
    def _extract(values, t, x):
        """
        values: shape [T + 1]
        t:      shape [B], integer timesteps in [0, T]
        x:      shape [B, C, H, W]
        """
        coefficients = values.gather(0, t.long())
        return coefficients.reshape(
            x.shape[0],
            *([1] * (x.ndim - 1))
        )

    def q_sample(self, x0, t, noise=None):
        """
        Directly sample x_t given x_0.

        x0:    [B, C, H, W]
        t:     [B], with values from 0 through T
        noise: optional [B, C, H, W] Gaussian tensor
        """
        if noise is None:
            noise = torch.randn_like(x0)

        signal_scale = self._extract(self.sqrt_alpha_bar, t, x0)
        noise_scale = self._extract(self.sqrt_one_minus_alpha_bar, t, x0)

        return signal_scale * x0 + noise_scale * noise

    def q_step(self, x_prev, t, noise=None):
        """
        Sample one actual Markov transition from x_{t-1} to x_t.

        t must contain integers from 1 through T.
        """
        if torch.any(t < 1) or torch.any(t > self.num_steps):
            raise ValueError("q_step expects timesteps in [1, T].")

        if noise is None:
            noise = torch.randn_like(x_prev)

        beta_t = self._extract(self.betas, t - 1, x_prev)
        return (
            torch.sqrt(1.0 - beta_t) * x_prev
            + torch.sqrt(beta_t) * noise
        )

Instantiate it with a schedule:

diffuser = ForwardDiffuser(
    num_steps=1000,
    schedule="cosine",
    device=device,
)

A few simple checks catch common implementation errors:

torch.manual_seed(7)

zero_t = torch.zeros(x0.shape[0], dtype=torch.long, device=device)
test_noise = torch.randn_like(x0)

x_at_zero = diffuser.q_sample(x0, zero_t, noise=test_noise)

assert torch.allclose(x_at_zero, x0)
assert diffuser.alpha_bar[0].item() == 1.0
assert diffuser.alpha_bar[-1].item() < diffuser.alpha_bar[1].item()

The first assertion is particularly valuable. If it fails, your indexing convention and cumulative schedule are inconsistent.

---

Broadcasting: the small detail that makes batched diffusion work

For a batch of $B$ images, timesteps have shape $[B]$. But a coefficient such as $\sqrt{{\bar{\alpha }}_t}$ must multiply every channel and pixel within the corresponding image.

The reshaping in _extract turns a vector of coefficients into shape

$$[B,1,1,1].$$

PyTorch then broadcasts that coefficient across the image dimensions:

Tensor	Example shape
Batch of images ${\mathbf{x}}_0$	$[B,C,H,W]$
Timesteps $t$	$[B]$
Gathered signal scales	$[B,1,1,1]$
Gaussian noise $\mathbf{ϵ}$	$[B,C,H,W]$

This lets one batch contain images at different noise levels. That is precisely how DDPM training will later work: each image receives its own independently sampled timestep.

---

Visualizing the marginals $q({\mathbf{x}}_t\mid {\mathbf{x}}_0)$

For an explanatory grid, it is helpful to use the same noise tensor at every selected timestep. This keeps a visually coherent noise pattern while the schedule gradually reduces the signal scale and raises the noise scale.

Each column remains a valid sample from the correct marginal $q({\mathbf{x}}_t\mid {\mathbf{x}}_0)$. However, the resulting grid is not a literal sequential Markov path, because a real path draws fresh noise at every transition.

import matplotlib.pyplot as plt


def display_image(x):
    """
    Convert a single CHW tensor from model space to a displayable HWC image.

    Clipping is for visualization only. It must not be used inside q_sample.
    """
    x = x.detach().cpu()
    x = ((x.clamp(-1, 1) + 1.0) / 2.0)
    return x.permute(1, 2, 0).numpy()


def plot_marginal_grid(diffuser, x0, timesteps, seed=0, title=None):
    generator = torch.Generator(device=x0.device).manual_seed(seed)
    fixed_noise = torch.randn(
        x0.shape,
        device=x0.device,
        dtype=x0.dtype,
        generator=generator,
    )

    images = []
    for step in timesteps:
        t = torch.full(
            (x0.shape[0],),
            step,
            device=x0.device,
            dtype=torch.long,
        )
        xt = diffuser.q_sample(x0, t, noise=fixed_noise)
        images.append(xt[0])

    fig, axes = plt.subplots(1, len(timesteps), figsize=(18, 3))

    for ax, image, step in zip(axes, images, timesteps):
        ax.imshow(display_image(image))
        ax.set_title(
            f"t = {step}\n"
            f"alpha_bar = {diffuser.alpha_bar[step].item():.3g}"
        )
        ax.axis("off")

    if title is not None:
        fig.suptitle(title, y=1.05)

    plt.tight_layout()
    plt.show()

Run it on a wide spread of noise levels:

timesteps = [0, 25, 100, 250, 500, 750, 1000]

plot_marginal_grid(
    diffuser=diffuser,
    x0=x0,
    timesteps=timesteps,
    seed=7,
    title="Cosine-schedule forward marginals",
)

When inspecting the output, distinguish three facts:

1. Visual recognizability is not the same as SNR. Images may retain broad color regions after semantic detail is gone.
2. The terminal sample should have no reliable connection to the source image. It need not look like grayscale static because each RGB channel receives independent Gaussian noise.
3. Out-of-range values are expected. The display function clips only because monitors cannot show values outside the standard image range.

---

Visualizing a true sequential forward trajectory

A true Markov trajectory begins at the clean image and calls q_step once per timestep. Each step has a new independent noise tensor.

@torch.no_grad()
def make_markov_trajectory(diffuser, x0, timesteps):
    requested_steps = set(timesteps)
    states = {0: x0.clone()}

    x = x0.clone()
    final_step = max(timesteps)

    for step in range(1, final_step + 1):
        t = torch.full(
            (x.shape[0],),
            step,
            device=x.device,
            dtype=torch.long,
        )
        x = diffuser.q_step(x, t)

        if step in requested_steps:
            states[step] = x.clone()

    return [states[step][0] for step in timesteps]


def plot_markov_grid(diffuser, x0, timesteps, seed=0):
    torch.manual_seed(seed)
    images = make_markov_trajectory(diffuser, x0, timesteps)

    fig, axes = plt.subplots(1, len(timesteps), figsize=(18, 3))

    for ax, image, step in zip(axes, images, timesteps):
        ax.imshow(display_image(image))
        ax.set_title(f"t = {step}")
        ax.axis("off")

    fig.suptitle("One sequential Markov forward trajectory", y=1.05)
    plt.tight_layout()
    plt.show()

plot_markov_grid(
    diffuser=diffuser,
    x0=x0,
    timesteps=[0, 25, 100, 250, 500, 750, 1000],
    seed=7,
)

The marginal grid and Markov grid will not look identical, even with the same random seed. They answer different questions:

Visualization	Noise usage	What it is best for
Direct marginals	One chosen noise tensor reused across displayed times	Explaining how the schedule scales signal and noise
Sequential path	Fresh noise at every transition	Showing the actual forward Markov process
Training sample	One random timestep and one fresh noise tensor per image	Producing examples for a noise-prediction network

For training, direct sampling is the preferred implementation because it has the correct distribution and costs one formula evaluation rather than $t$ sequential transitions.

---

Make the previous schedule comparison visible in image space

The same plotting function can compare linear and cosine schedules at matching nominal timesteps.

linear_diffuser = ForwardDiffuser(
    num_steps=1000,
    schedule="linear",
    device=device,
)

cosine_diffuser = ForwardDiffuser(
    num_steps=1000,
    schedule="cosine",
    device=device,
)

timesteps = [0, 100, 250, 500, 750, 1000]

plot_marginal_grid(
    linear_diffuser,
    x0,
    timesteps,
    seed=11,
    title="Linear beta schedule",
)

plot_marginal_grid(
    cosine_diffuser,
    x0,
    timesteps,
    seed=11,
    title="Cosine schedule",
)

Using the same image, selected timesteps, and visualization seed makes this a controlled comparison. The two grids differ only because their ${\bar{\alpha }}_t$ values differ. In particular, a common timestep such as $t=500$ does not represent a common noise level across schedules.

---

Relating your implementation to a library scheduler

Once the tensor-level implementation is clear, a library can take care of schedule storage and noising. The abstraction is still the same: a scheduler receives clean samples, Gaussian noise, and timestep indices, then returns noisy samples.

A library scheduler is convenient, but your implementation should remain intelligible in terms of the same three cached objects:

$${\beta }_t,{\alpha }_t=1-{\beta }_t,{\bar{\alpha }}_t=\prod _{s=1}^t{\alpha }_s.$$

---

Key takeaways

A practical forward diffuser needs only a schedule and the direct-sampling identity:

$${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ}.$$

The implementation details that matter most are:

• Normalize clean image tensors to $[-1,1]$.
• Precompute schedule quantities once rather than recomputing them per batch.
• Gather timestep-dependent scalars and reshape them to $[B,1,1,1]$.
• Use direct sampling for training and marginal visualization.
• Use repeated one-step sampling only when you specifically want a genuine forward Markov trajectory.
• Clip only for display, never as part of the diffusion operation.

The next module moves from forward corruption to the reverse learning problem. We will first derive the tractable posterior $q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)$, which is the bridge between this known forward process and a learned denoising model.