Lesson 62's diffusion models learn to reverse a noising process one still image at a time. This lesson extends that same denoising recipe to video: generate a whole clip at once while keeping it temporally consistent — the new failure mode a video generator has that an image generator never faces is flicker, where independently-denoised frames disagree slightly from one frame to the next, so a shape's position or appearance jitters instead of moving smoothly. The fix has two parts: architecture (a temporal-attention layer, plus a frame-position embedding so the network knows which frame of the clip it is denoising) and sampling (correlating, rather than fully independently drawing, the random noise injected into each frame at every step). Both parts turn out to matter, and this lesson shows why with a concrete, measurable failure mode along the way.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
Reuse Lesson 64's shape generator, but now give the shape a constant velocity: each clip is a short sequence of frames in which a plus or circle translates by a fixed (vx, vy) offset per frame. "Smooth, consistent motion" is exactly the property a good video generator needs to reproduce, and exactly what flicker destroys — so a dataset where motion is literally constant velocity by construction makes it easy to check, quantitatively, whether a generated clip moved smoothly or not.
SIZE = 16
T = 60
NF = 4 # frames per clip
betas = torch.linspace(1e-4, 0.02, T)
alphas = 1.0 - betas
alpha_bars = torch.cumprod(alphas, dim=0)
def make_image(shape_type, cx, cy, size=SIZE):
img = np.zeros((size, size), dtype=np.float32)
if shape_type == 'plus':
img[cy-1:cy+2, cx-3:cx+4] = 1.0
img[cy-3:cy+4, cx-1:cx+2] = 1.0
else:
yy, xx = np.mgrid[0:size, 0:size]
img[((xx-cx)**2 + (yy-cy)**2) <= 9] = 1.0
return img * 2 - 1 # scale to [-1, 1], same convention as Lesson 64
def make_clip(rng, n_frames=NF, size=SIZE):
shape_type = rng.choice(['plus', 'circle'])
vx, vy = 0, 0
while vx == 0 and vy == 0: # exclude the degenerate "doesn't move" case
vx, vy = rng.integers(-1, 2), rng.integers(-1, 2)
cx0, cy0 = rng.integers(7, 9), rng.integers(7, 9)
frames = [make_image(shape_type, cx0 + vx * i, cy0 + vy * i, size) for i in range(n_frames)]
return np.stack(frames, axis=0) # (NF, H, W)
def make_clip_dataset(rng, n, n_frames=NF, size=SIZE):
return np.array([make_clip(rng, n_frames, size) for _ in range(n)], dtype=np.float32)
rng = np.random.default_rng(3)
clips = make_clip_dataset(rng, 500)
clips_t = torch.tensor(clips) # (N, NF, H, W)
fig, axes = plt.subplots(2, NF, figsize=(7, 3.5))
for row in range(2):
for i in range(NF):
axes[row, i].imshow(clips[row, i], cmap='gray'); axes[row, i].axis('off')
if row == 0:
axes[row, i].set_title(f't={i}', fontsize=8)
plt.suptitle('Two training clips: a shape moving at constant velocity')
plt.show()
DenoiseNet below is Lesson 64's own per-frame denoiser, unchanged — it will serve as the frame-independent baseline. VideoDenoiseNet adds two things on top of the same spatial backbone: a temporal-attention layer (Lesson 45's attention mechanism, repurposed so that instead of mixing information across positions in a sequence, it mixes information across frames at the same pixel location — every spatial location's feature vector becomes a query that attends to that same location's feature vectors in the other NF frames), and a frame-position embedding (the same sinusoidal-embedding idea as the diffusion timestep t, but indexing "which of the NF frames is this" instead of "how noisy is this"). The embedding matters more than it looks: without it, every frame gets an identical input representation whenever the noise happens to be correlated across frames, so a purely symmetric network has no way to tell frame 0 from frame 3 and nothing stops it from collapsing to the trivial, degenerate "solution" of generating the exact same frame four times (zero flicker, but also zero motion). The frame-position embedding breaks that symmetry, giving the temporal-attention layers an actual position to reason about.
Real video diffusion models usually factorize spatial and temporal mixing into two separate passes, rather than one joint attention over every pixel of every frame at once, since joint attention's cost grows with the product of spatial and temporal length; this lesson's architecture already reflects that factorization, with ordinary spatial convolutions handling within-frame structure and TemporalAttention handling across-frame structure separately.
class TimeEmbedding(nn.Module):
def __init__(self, dim=32):
super().__init__()
self.dim = dim
def forward(self, t):
half = self.dim // 2
freqs = torch.exp(-np.log(10000) * torch.arange(half).float() / half)
args = t[:, None].float() * freqs[None, :]
return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
class DenoiseNet(nn.Module):
"""Lesson 64's per-frame denoiser, unmodified -- the frame-independent baseline."""
def __init__(self, ch=48, time_dim=32):
super().__init__()
self.time_embed = TimeEmbedding(time_dim)
self.time_mlp = nn.Linear(time_dim, ch)
self.conv1 = nn.Conv2d(1, ch, 3, padding=1)
self.conv2 = nn.Conv2d(ch, ch, 3, padding=1)
self.conv3 = nn.Conv2d(ch, 1, 3, padding=1)
def forward(self, x, t):
temb = self.time_mlp(self.time_embed(t))[:, :, None, None]
h = F.relu(self.conv1(x) + temb)
h = F.relu(self.conv2(h) + temb)
return self.conv3(h)
class TemporalAttention(nn.Module):
"""Lesson 45's attention, applied along the frame axis: each spatial location's
feature vector attends across that same location in all NF frames."""
def __init__(self, ch, n_frames):
super().__init__()
self.n_frames = n_frames
self.q = nn.Linear(ch, ch)
self.k = nn.Linear(ch, ch)
self.v = nn.Linear(ch, ch)
self.scale = ch ** -0.5
def forward(self, h):
BF, ch, Hh, Ww = h.shape # BF = batch_size * n_frames
B = BF // self.n_frames
x = h.view(B, self.n_frames, ch, Hh, Ww).permute(0, 3, 4, 1, 2).reshape(B * Hh * Ww, self.n_frames, ch)
q, k, v = self.q(x), self.k(x), self.v(x)
attn = torch.softmax((q @ k.transpose(-1, -2)) * self.scale, dim=-1)
out = attn @ v
out = out.reshape(B, Hh, Ww, self.n_frames, ch).permute(0, 3, 4, 1, 2).reshape(BF, ch, Hh, Ww)
return h + out # residual, same pattern as Lesson 46's encoder block
class VideoDenoiseNet(nn.Module):
def __init__(self, ch=48, time_dim=32, n_frames=NF):
super().__init__()
self.n_frames = n_frames
self.time_embed = TimeEmbedding(time_dim)
self.time_mlp = nn.Linear(time_dim, ch)
self.frame_embed = nn.Embedding(n_frames, ch) # tells the network WHICH frame of the clip this is
self.conv1 = nn.Conv2d(1, ch, 3, padding=1)
self.temporal_attn1 = TemporalAttention(ch, n_frames)
self.conv2 = nn.Conv2d(ch, ch, 3, padding=1)
self.temporal_attn2 = TemporalAttention(ch, n_frames)
self.conv3 = nn.Conv2d(ch, 1, 3, padding=1)
def forward(self, x, t, frame_idx):
# x: (B*NF, 1, H, W), t: (B*NF,), frame_idx: (B*NF,) with values in [0, NF)
temb = self.time_mlp(self.time_embed(t))[:, :, None, None]
femb = self.frame_embed(frame_idx)[:, :, None, None]
h = F.relu(self.conv1(x) + temb + femb)
h = self.temporal_attn1(h)
h = F.relu(self.conv2(h) + temb + femb)
h = self.temporal_attn2(h)
return self.conv3(h)
Both models are trained with the identical noise-prediction objective as Lesson 64. The baseline is trained on every frame of every clip, pooled together as if they were unrelated still images — it never sees two frames of the same clip together, so it has no way to learn any relationship between them. The video model is trained clip-by-clip: all NF frames of a clip are denoised together, at the same randomly-sampled noise level t, with each frame also given its own frame_idx so VideoDenoiseNet can learn frame-position-dependent content.
def train_baseline(clips_t, epochs=800, bs=64, seed=0):
torch.manual_seed(seed)
model = DenoiseNet()
opt = torch.optim.Adam(model.parameters(), lr=0.001)
frames_flat = clips_t.reshape(-1, 1, SIZE, SIZE) # every frame of every clip, treated as an independent image
n = frames_flat.shape[0]
for epoch in range(epochs):
idx = np.random.default_rng(epoch).permutation(n)[:bs]
x0 = frames_flat[idx]
t = torch.randint(0, T, (x0.shape[0],))
noise = torch.randn_like(x0)
ab = alpha_bars[t][:, None, None, None]
x_t = torch.sqrt(ab) * x0 + torch.sqrt(1 - ab) * noise
pred = model(x_t, t)
loss = F.mse_loss(pred, noise)
opt.zero_grad(); loss.backward(); opt.step()
return model, loss.item()
def train_video(clips_t, epochs=2500, bs_clips=16, seed=0):
torch.manual_seed(seed)
model = VideoDenoiseNet()
opt = torch.optim.Adam(model.parameters(), lr=0.001)
n = clips_t.shape[0]
frame_idx_single = torch.arange(NF)
for epoch in range(epochs):
idx = np.random.default_rng(epoch).permutation(n)[:bs_clips]
B = idx.shape[0]
x0 = clips_t[idx].reshape(-1, 1, SIZE, SIZE) # (B*NF, 1, H, W)
frame_idx = frame_idx_single.repeat(B) # [0,1,2,3, 0,1,2,3, ...]
t_clip = torch.randint(0, T, (B,))
t = t_clip.repeat_interleave(NF) # same noise LEVEL across a clip's frames
noise = torch.randn_like(x0) # still an independent noise DRAW per frame
ab = alpha_bars[t][:, None, None, None]
x_t = torch.sqrt(ab) * x0 + torch.sqrt(1 - ab) * noise
pred = model(x_t, t, frame_idx)
loss = F.mse_loss(pred, noise)
opt.zero_grad(); loss.backward(); opt.step()
return model, loss.item()
baseline_model, bl_loss = train_baseline(clips_t)
video_model, vid_loss = train_video(clips_t)
print(f'baseline (frame-independent) final loss: {bl_loss:.3f}')
print(f'video (temporal-attention) final loss: {vid_loss:.3f}')
Sampling mirrors training's asymmetry. The baseline reverses the diffusion process for each frame completely independently — its own starting noise, its own fresh random noise injected at every one of the T reverse steps — so nothing connects what frame 0 looks like to what frame 1 looks like, beyond both being some plausible shape. The video model's sampler instead draws mostly-shared noise: at every step, a single shared noise tensor is mixed with a small independent-per-frame component (SHARE = 0.9 below), rather than being either fully independent (no correlation at all) or fully shared (which, perhaps surprisingly, is also wrong — with literally identical noise and no frame-position embedding, every frame receives an identical input at every step and the reverse process provably produces four pixel-identical frames: zero flicker, but also zero motion, a degenerate "solution" that technically minimizes flicker metrics while completely failing at the actual task. Try it: remove the frame-position embedding and set SHARE = 1.0, and check that every generated frame is bit-for-bit identical). Mostly-shared noise keeps the frames' trajectories correlated enough to suppress flicker while leaving enough independent signal, combined with the frame-position embedding, for the network to still generate genuinely different (and smoothly evolving) content per frame.
To measure flicker numerically, track each frame's shape centroid and compute the jerk: the mean magnitude of the frame-to-frame change in velocity (the second difference of centroid position). Real clips move at constant velocity by construction, so their jerk is exactly 0; a generator that produces smooth, physically plausible motion should score close to 0 too, while one that effectively redraws an unrelated shape each frame will have large, essentially random jerk. (Jerk alone can't distinguish "smooth motion" from "no motion at all" — a frozen clip also scores 0 — so the code below also reports each clip's mean frame-to-frame speed, to confirm the video model is actually moving the shape, not just holding it still.)
@torch.no_grad()
def sample_baseline_clip(model, seed, n_frames=NF):
"""Each frame is its own completely independent reverse-diffusion run."""
torch.manual_seed(seed)
x = torch.randn(n_frames, 1, SIZE, SIZE)
for t in reversed(range(T)):
t_batch = torch.full((n_frames,), t, dtype=torch.long)
pred_noise = model(x, t_batch)
alpha_t, alpha_bar_t, beta_t = alphas[t], alpha_bars[t], betas[t]
mean = (1 / torch.sqrt(alpha_t)) * (x - (beta_t / torch.sqrt(1 - alpha_bar_t)) * pred_noise)
x = mean + torch.sqrt(beta_t) * torch.randn_like(x) if t > 0 else mean
return x.squeeze(1)
SHARE = 0.9 # fraction of each noise draw's variance shared across frames; the rest is independent per frame
def mixed_noise(n_frames):
shared = torch.randn(1, 1, SIZE, SIZE)
indep = torch.randn(n_frames, 1, SIZE, SIZE)
return np.sqrt(SHARE) * shared + np.sqrt(1 - SHARE) * indep
@torch.no_grad()
def sample_video_clip(model, seed, n_frames=NF):
"""Temporal-attention + frame-position-aware architecture, PLUS mostly-shared noise."""
torch.manual_seed(seed)
frame_idx = torch.arange(n_frames)
x = mixed_noise(n_frames)
for t in reversed(range(T)):
t_batch = torch.full((n_frames,), t, dtype=torch.long)
pred_noise = model(x, t_batch, frame_idx)
alpha_t, alpha_bar_t, beta_t = alphas[t], alpha_bars[t], betas[t]
mean = (1 / torch.sqrt(alpha_t)) * (x - (beta_t / torch.sqrt(1 - alpha_bar_t)) * pred_noise)
x = mean + torch.sqrt(beta_t) * mixed_noise(n_frames) if t > 0 else mean
return x.squeeze(1)
def frame_centroid(frame):
ys, xs = np.where(frame.numpy() > 0)
return None if len(xs) == 0 else np.array([xs.mean(), ys.mean()])
def clip_jerk(clip):
centroids = [frame_centroid(f) for f in clip]
if any(c is None for c in centroids):
return None
vels = [centroids[i + 1] - centroids[i] for i in range(len(centroids) - 1)]
jerks = [np.linalg.norm(vels[i + 1] - vels[i]) for i in range(len(vels) - 1)]
return np.mean(jerks)
def clip_speed(clip):
"""mean distance moved between consecutive frames -- near 0 would mean a frozen clip."""
centroids = [frame_centroid(f) for f in clip]
if any(c is None for c in centroids):
return None
return np.mean([np.linalg.norm(centroids[i + 1] - centroids[i]) for i in range(len(centroids) - 1)])
real_jerks = [clip_jerk(torch.tensor(c)) for c in clips[:30] if clip_jerk(torch.tensor(c)) is not None]
n_eval = 30
baseline_jerks, video_jerks, baseline_speeds, video_speeds = [], [], [], []
for seed in range(n_eval):
bclip, vclip = sample_baseline_clip(baseline_model, seed=1000 + seed), sample_video_clip(video_model, seed=2000 + seed)
bj, vj = clip_jerk(bclip), clip_jerk(vclip)
bs, vs = clip_speed(bclip), clip_speed(vclip)
if bj is not None: baseline_jerks.append(bj)
if vj is not None: video_jerks.append(vj)
if bs is not None: baseline_speeds.append(bs)
if vs is not None: video_speeds.append(vs)
print(f'real training clips: jerk = {np.mean(real_jerks):.3f}')
print(f'baseline (frame-independent sampling): jerk = {np.mean(baseline_jerks):.3f}, speed = {np.mean(baseline_speeds):.3f} (n={len(baseline_jerks)}/{n_eval})')
print(f'video model (temporal attn + mostly-shared noise): jerk = {np.mean(video_jerks):.3f}, speed = {np.mean(video_speeds):.3f} (n={len(video_jerks)}/{n_eval})')
bclip_show = sample_baseline_clip(baseline_model, seed=1000)
vclip_show = sample_video_clip(video_model, seed=2000)
fig, axes = plt.subplots(2, NF, figsize=(7, 3.5))
for i in range(NF):
axes[0, i].imshow(bclip_show[i], cmap='gray'); axes[0, i].axis('off')
axes[1, i].imshow(vclip_show[i], cmap='gray'); axes[1, i].axis('off')
axes[0, i].set_title(f't={i}', fontsize=8)
plt.suptitle('Baseline (top, flickers) vs. temporal video model (bottom, moves smoothly)')
plt.show()
The video model's jerk is a clear fraction of the baseline's, while its speed confirms the shape is genuinely moving (not frozen) — both halves of the fix are doing real work, but a small tiny-model-on-tiny-data gap remains between the video model's jerk and the real clips' exact 0, since 2,500 toy-scale training steps on 500 clips can't perfectly learn "constant velocity." (Exercise 2 asks you to probe how much each half of the fix is pulling its own weight.)
A model that can only generate an unconditional 4-frame clip of a moving blob isn't yet useful on its own. Two conditioning recipes turn it into something closer to a real product:
Both conditioning recipes generate a clip of a fixed length. Real video models need to go well beyond NF frames, and no single joint denoising pass scales to minutes of footage. The standard solution is autoregressive rollout: generate a short chunk of frames (as this lesson does), then generate the next chunk conditioned on the last few already-generated frames (image-to-video conditioning, applied to the model's own output) instead of starting fresh — repeating this produces an arbitrarily long video, at the cost of slowly accumulating drift from one chunk to the next, since every chunk's conditioning is itself a generated (not real) clip.
SHARE = 1.0 in mixed_noise and re-sample from the already-trained video_model (which already includes frame_embed). Is every frame identical, or does the model still generate genuinely different content per frame? Now build and train a second copy of VideoDenoiseNet with the + femb terms removed from forward (so it never learns to use frame position at all), and sample from that model, also with SHARE = 1.0. Which single piece — removing the frame-position embedding, or using fully shared noise — is actually responsible for the degenerate frozen-clip solution described above, and which one turns out not to matter much on its own?TemporalAttention + frame_embed) and the sampling trick (mostly-shared noise, SHARE = 0.9). Sample from the trained video_model using fully independent per-frame noise (SHARE = 0.0) instead, and separately, sample from the trained baseline_model using mixed_noise instead of independent noise (ignoring that it was never trained with frame indices, so just drop the frame_idx argument it doesn't accept). Which change moves the jerk metric more on its own?vx, vy's range in make_clip to allow speeds of 2 pixels/frame instead of 1, and retrain both models. Does the video model's jerk advantage over the baseline hold up at the faster, more dramatic motion, or does it degrade — and if it degrades, why might faster motion be a harder case for a fixed NF-frame temporal window and a fixed SHARE to track well?