Noise prediction in Diffusion

Analogy

Imagine you are an AI who is trained by a group of researcher. They have a bunch of images. They say, those images will be added with noise gradually with some sort of scheduling techniques you don’t know. The images will be added by noise gradually in 1000 timesteps max. But, when training, they will give you images that has been added with random noise in random timestep. For example, in the 1st training, they give you images that has been added with random noise in the 67th timestep. In the 2nd training, they can give you set of noised images from same original images in the same timestep as before (the 67th timestep), but the noise that was added is different. Therefore, you can’t memorize what kind of noise that is added in that image. But the noised images are still have the same underlying structure (the original images). Your job is to predict the noise that was added to the image. After that, the researcher will give you the true noise and let you learn from your mistake to make your prediction better. After 100 of training, you start to understand the pattern. They will tell you to generate a meaningful image that resembles the structures of original images from a purely noised image.

Different noise in the same timestep (t)

Let’s simulate simple forward diffusion:

import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw

# --- create a synthetic clean image x0 ---
W, H = 128, 128
img = Image.new("RGB", (W, H), color=(30, 30, 30))
draw = ImageDraw.Draw(img)
draw.rectangle([16, 16, 48, 48], fill=(220, 50, 50))   # red square
draw.ellipse([80, 20, 116, 56], fill=(50, 200, 50))    # green circle
draw.polygon([(64,80),(100,110),(28,110)], fill=(50,50,220))  # blue triangle

# small horizontal line for extra structure
draw.line([0, H//2, W, H//2], fill=(200,200,50), width=2)

x0 = np.array(img).astype(np.float32) / 255.0            # H,W,C in [0,1]
x0_t = torch.from_numpy(x0).permute(2,0,1).unsqueeze(0)   # (1,3,H,W) in [0,1]
# scale to typical diffusion range [-1,1]
x0_scaled = x0_t * 2.0 - 1.0                             # (1,3,H,W) in [-1,1]

# --- choose a single timestep t represented by alpha_bar ---
# (in practice alpha_bar comes from the beta schedule; here pick a scalar for illustration)
alpha_bar = 0.4
sqrt_a = np.sqrt(alpha_bar)
sqrt_1_minus_a = np.sqrt(1.0 - alpha_bar)

# generate three different eps and x_t
noisy_results = []
for i in range(3):
    eps = torch.randn_like(x0_scaled)    # a fresh noise draw each time
    xt = sqrt_a * x0_scaled + sqrt_1_minus_a * eps
    # convert to HWC numpy for plotting (xt in [-1,1] -> to [0,1])
    noisy_results.append((eps.squeeze(0).permute(1,2,0).numpy(),
                          xt.squeeze(0).permute(1,2,0).numpy()))

# --- plotting ---
fig, axs = plt.subplots(4, 3, figsize=(10,12))

# top row: x_t images
for i in range(3):
    xt_rgb = noisy_results[i][1]
    xt_disp = (xt_rgb + 1.0) / 2.0                  # [-1,1] -> [0,1]
    xt_disp = np.clip(xt_disp, 0.0, 1.0)
    axs[0,i].imshow(xt_disp)
    axs[0,i].axis("off")
    axs[0,i].set_title(f"x_t sample {i+1}")

# second row: epsilon (noise) visualized per sample
for i in range(3):
    eps_rgb = noisy_results[i][0]
    # normalize noise for display (per-sample)
    eps_min, eps_max = eps_rgb.min(), eps_rgb.max()
    eps_disp = (eps_rgb - eps_min) / (eps_max - eps_min + 1e-8)
    axs[1,i].imshow(eps_disp)
    axs[1,i].axis("off")
    axs[1,i].set_title("ε (noise)")

# third row: the clean x0 repeated as reference
x0_disp = x0  # HWC [0,1]
for i in range(3):
    axs[2,i].imshow(x0_disp)
    axs[2,i].axis("off")
    axs[2,i].set_title("x0 (clean)")

# fourth row: residuals xt - scaled_x0 (what model would need to explain)
for i in range(3):
    xt_rgb = noisy_results[i][1]
    residual = xt_rgb - (x0 * 2.0 - 1.0)   # xt minus scaled x0
    rmin, rmax = residual.min(), residual.max()
    residual_disp = (residual - rmin) / (rmax - rmin + 1e-8)
    axs[3,i].imshow(residual_disp)
    axs[3,i].axis("off")
    axs[3,i].set_title("residual (xt - scaled_x0)")

plt.suptitle(f"Same timestep (alpha_bar={alpha_bar}) with 3 different ε samples\nTop: x_t | 2nd: ε | 3rd: x0 | Bottom: residuals")
plt.tight_layout(rect=[0,0,1,0.96])
plt.show()

The 1st row is represent noised image. They are from the original image from the 3rd row that has been added with the noise in the 2nd row. The last row (4th row), shows the total change apply to the original image.

As you can see, even though they are from the same timestep, each image can has different noise. By subtracting those noised image with the original image, you can get how much their differences. Note that noised image — original image is not the same with the added noise.

After we done in forward process, we need to denoise it. It is the process of subtracting the noise gradually until it generates image that resembles the structure from the original image. This process usually use U-Net. We can set simple toy U-Net to simulate it.

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw
import random

torch.manual_seed(0)
np.random.seed(0)
random.seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# -----------------------------
# Create a simple synthetic x0
# -----------------------------
W, H = 128, 128
img = Image.new("RGB", (W, H), color=(30, 30, 30))
draw = ImageDraw.Draw(img)
draw.rectangle([16, 16, 48, 48], fill=(220, 50, 50))   # red square
draw.ellipse([80, 20, 116, 56], fill=(50, 200, 50))    # green circle
draw.polygon([(64,80),(100,110),(28,110)], fill=(50,50,220))  # blue triangle
draw.line([0, H//2, W, H//2], fill=(200,200,50), width=2)
x0 = np.array(img).astype(np.float32) / 255.0  # HWC in [0,1]
x0_t = torch.from_numpy(x0).permute(2,0,1).unsqueeze(0).to(device)  # (1,3,H,W) in [0,1]
x0_scaled = x0_t * 2.0 - 1.0  # scale to [-1,1], typical diffusion input range

# -----------------------------
# Make three xt samples at same t
# -----------------------------
alpha_bar = 0.4
sqrt_a = np.sqrt(alpha_bar)
sqrt_1_a = np.sqrt(1.0 - alpha_bar)

num_samples = 3
eps_list = []
xt_list = []
for i in range(num_samples):
    eps = torch.randn_like(x0_scaled).to(device)
    xt = sqrt_a * x0_scaled + sqrt_1_a * eps
    eps_list.append(eps)
    xt_list.append(xt)

# Stack them as a tiny training set (B, C, H, W)
X = torch.cat(xt_list, dim=0)   # shape (3,3,H,W)
E_true = torch.cat(eps_list, dim=0)

# Also create target t embedding vector (same t for all)
T_scalar = torch.tensor([56], dtype=torch.long)  # any integer timestep
T = torch.full((num_samples,), 56, dtype=torch.long).to(device)

# --------------------------------
# Tiny U-Net-like model for demo
# --------------------------------
class TinyUNet(nn.Module):
    def __init__(self, in_ch=3, base_ch=32, time_emb_dim=32):
        super().__init__()
        # time MLP
        self.time_mlp = nn.Sequential(
            nn.Embedding(1000, time_emb_dim),  # assume T <= 1000
            nn.Linear(time_emb_dim, base_ch * 2),
            nn.ReLU()
        )
        # encoder
        self.enc1 = nn.Sequential(nn.Conv2d(in_ch, base_ch, 3, padding=1), nn.ReLU())
        self.enc2 = nn.Sequential(nn.Conv2d(base_ch, base_ch*2, 3, padding=1), nn.ReLU())
        # bottleneck
        self.bot = nn.Sequential(nn.Conv2d(base_ch*2, base_ch*2, 3, padding=1), nn.ReLU())
        # decoder
        self.up1 = nn.ConvTranspose2d(base_ch*2, base_ch, kernel_size=2, stride=2)
        self.dec1 = nn.Sequential(nn.Conv2d(base_ch*2, base_ch, 3, padding=1), nn.ReLU())
        self.out = nn.Conv2d(base_ch, in_ch, 1)
        self.pool = nn.AvgPool2d(2)

    def forward(self, x, t):
        # x: (B,3,H,W), t: (B,) integer timesteps
        te = self.time_mlp(t)                # (B, time_emb_dim)
        # broadcast time embedding spatially and add to features later
        e = te[:,:,None,None]                # (B, dim, 1, 1)

        e1 = self.enc1(x)                    # (B, base, H, W)
        p1 = self.pool(e1)                   # down
        e2 = self.enc2(p1)                   # (B, base*2, H/2, W/2)

        b = self.bot(e2)
        # add time embedding to bottleneck channels
        b = b + e.expand_as(b)

        u = self.up1(b)                      # (B, base, H, W)
        # concat skip connection
        u = torch.cat([u, e1], dim=1)        # (B, base*2, H, W)
        d = self.dec1(u)
        out = self.out(d)
        return out

model = TinyUNet().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()

# --------------------------------
# Train the tiny UNet to predict eps
# We'll train for a small number of iterations on this tiny dataset.
# --------------------------------
epochs = 800  # small toy training, increase if you want better fit
batch = X  # tiny dataset, we will just re-use it
t_tensor = T

model.train()
for ep in range(epochs):
    optimizer.zero_grad()
    pred = model(batch, t_tensor)   # predict epsilon
    loss = loss_fn(pred, E_true)
    loss.backward()
    optimizer.step()
    if (ep+1) % 200 == 0:
        print(f"Epoch {ep+1}/{epochs}  loss={loss.item():.6f}")

# --------------------------------
# Evaluate: show results for each sample
# --------------------------------
model.eval()
with torch.no_grad():
    pred_E = model(X, t_tensor)    # predicted eps for the 3 samples
    # reconstruct x0_hat from predicted eps: x0_hat = (xt - sqrt(1 - a) * eps_hat) / sqrt(a)
    x0_hat = (X - sqrt_1_a * pred_E) / sqrt_a

# Helper for plotting tensors (CHW -> HWC in [0,1])
def to_img(t):
    t = t.cpu().numpy().transpose(1,2,0)
    # clamp from [-1,1] to [0,1] if needed
    t = (t + 1.0) / 2.0
    t = np.clip(t, 0.0, 1.0)
    return t

n = num_samples
fig, axs = plt.subplots(n, 4, figsize=(12, 4*n))

for i in range(n):
    axs[i,0].imshow(to_img(X[i]))                      # x_t
    axs[i,0].set_title("x_t (noisy)")
    axs[i,1].imshow(to_img(E_true[i]))                 # true eps
    axs[i,1].set_title("true ε")
    axs[i,2].imshow(to_img(pred_E[i]))                 # predicted eps
    axs[i,2].set_title("pred ε̂")
    axs[i,3].imshow(to_img(x0_hat[i]))                 # reconstructed x0_hat
    axs[i,3].set_title("reconstructed x0_hat")
    for j in range(4):
        axs[i,j].axis("off")
plt.tight_layout()
plt.show()

# Print per-sample MSE of predicted epsilon
mse_per_sample = ((pred_E - E_true)**2).reshape(n, -1).mean(dim=1).cpu().numpy()

for i, m in enumerate(mse_per_sample):
    print(f"Sample {i+1} MSE(eps): {m:.6f}")

It still use the same timestep but use different noise

The result from this tiny demonstration is: Sample 1 MSE(eps): 0.010022 Sample 2 MSE(eps): 0.009744 Sample 3 MSE(eps): 0.010310

This Mean Squared Error calculates the difference between the predicted noise with the actual noise. Due to very little dataset, the training is likely overfitting. We can try with using another dataset, MNIST.

Trained on MNIST

You can see full implementation in here:

[embed]Noise diff in Diffusion Explore and run machine learning code with Kaggle Notebooks | Using data from No attached data sourceswww.kaggle.com

Simple U-Net

U-Net simple architecture

Let’s analyze the few result from Simple U-Net. You can see the full result in the link above.

Step 10000 Loss 0,0801

Step 50000 Loss 0,0693

Step 90000 Loss 0,0593

There are 4 images in a row, but we just going to focus more on pred — actual noise. Red pixel indicate that prediction value is higher than the actual value, blue pixel indicate prediction value is lower than the actual value, white pixel indicate prediction is match (or close) with the actual value.

There’s still much red and blue than the white pixel, which mean the model is struggle to predict the noise. After 100 training, we can see the sampling or inference result based on DDPM sampling.

Sampling from pure noise. It doesn’t generates anything meaningful.

The same goes with this. It use x_t rather from pure noise to see whether the underlying structure from ground truth can help guide the sampling step.

There is nothing good generates from this model. Therefore, we can modify the model to improve its performance.

Robust U-Net

Let’s add Sinusoidal embedding, ResBlock, and AttentionBlock within our previous simple U-Net. Let’s see what it learn from robust U-Net.

Step 10000 loss 0,0149

Step 50000 loss 0,0222

Step 90000 loss 0,0143

As you can see, the loss is heavily reduced. Which mean it is a good sign. But, the pred — actual shows that the model still struggle to predict the exact value of the noise around the number structures. The black area are mostly white-ish. It indicates that the model is pretty much can predict the noise that have underlying blank area but have a hard time predict noise that have underlying complex structure. But, the sampling result is actually improve from the previous test.

The structure like number 9 is emerge from pure noise. The line look solid and show clear pattern.

The result from generating determined x_t from actual x0 is different. To be clear, it is expected. Because when sampling happens, it adding random noise every time.

Here is another from pure noise. It uses DDPM sampling like before.

This one not really look clear, but at least the line is solid.

This one is not using attention block and the result still looks good.

This one also not using attention block. Still looks good.

Let’s try sampling using DDIM, where it can sampling only in a very few steps compare to DDPM.

From pure noise result looks good. It can give us a solid ‘8’

Number 7 is emerge only within 50 steps

We can compare on DDPM vs DDIM quality.

DDPM and DDIM use the same pure noise as starting point. The result is different and it is expected, because when sampling they add with random noise.

Same goes with using x_t from actual x0. The result is different. Again, this is expected!

In here we can conclude: for MNIST dataset, attention block is not really necessary because without it, it still have deliver good result. But of course, more complex U-Net is needed for this task by utilizing ResBlock and Sinusoidal embedding.