Frequency tiers (estimates, relative to interviews targeting this pack's specialty — a ★★★★★ here means "core for these roles", not for generalist ML loops): ★★★★★ expect it in most loops · ★★★★ very common · ★★★ common · ★★ occasional · ★ rare/deep-specialist. Spend ~50% of practice on the top two tiers.
70+ problems with full PyTorch solutions and tests
Principal/Senior-Principal Generative-Model Interview Prep
Notes
This pack covers everything you need to code on a whiteboard about modern continuous-time generative mod- els: forward / reverse Markov chains (DDPM), deterministic samplers (DDIM), score-based SDEs / ODEs, classifier-free guidance, higher-order solvers (Heun, DPM-Solver, EDM), flow matching, rectified flow, normal- izing flows, mean flow, and consistency models. Each problem has a precise statement, a clean PyTorch / NumPy reference, and tests that verify against either a closed-form property or a known limit (e.g. DDIM with η = 1 recovers DDPM ancestral sampling in the discrete case; flow matching with linear schedule recovers rectified flow). Standard imports across the pack:
import math, numpy as np, torch
import torch.nn as nn, torch.nn.functional as F
Problem
Build βt linearly between β1 and βT , then derive αt = 1 −βt and ᾱt = Πt i=1 αi.
import torch
def linear_beta(T, beta1=1e-4, beta2=0.02):
betas = torch.linspace(beta1, beta2, T)
alphas = 1.0 - betas
abar = torch.cumprod(alphas, dim=0)
return betas, alphas, abar
Tests
b, a, ab = linear_beta(1000)
assert b[0].item() == 1e-4 and b[-1].item() == 0.02
assert ab[0].item() == 1 - 1e-4 and ab[-1] < ab[0]
print("linear beta OK")
Problem
Implement the cosine schedule of Nichol & Dhariwal: ᾱt = f(t)/f(0), f(t) = cos2((t/T + s)/(1 + s) · π/2), then derive βt = 1 −ᾱt/ᾱt−1 clipped to [1e−4, 0.999].
import math
import torch
def cosine_beta(T, s=0.008):
t = torch.linspace(0, T, T + 1)
f = torch.cos((t / T + s) / (1 + s) * math.pi / 2) ** 2
abar = f / f[0]
betas = (1 - abar[1:] / abar[:-1]).clamp(1e-4, 0.999)
alphas = 1.0 - betas
return betas, alphas, torch.cumprod(alphas, dim=0)
Tests
b, a, ab = cosine_beta(1000)
assert ab[-1] < 1e-3 and ab[0] < 1.0
print("cosine OK")
Problem
A smooth schedule that interpolates linear vs. cosine: βt = σ(logitt/T ) scaled to [β1, βT ].
import torch
def sigmoid_beta(T, beta1=1e-4, beta2=0.02, k=6.0):
t = torch.linspace(-k, k, T)
s = torch.sigmoid(t)
s = (s - s.min()) / (s.max() - s.min())
betas = beta1 + s * (beta2 - beta1)
alphas = 1 - betas
return betas, alphas, torch.cumprod(alphas, dim=0)
Tests
b, a, ab = sigmoid_beta(100); assert b[0] < b[-1]
print("sigmoid OK")
Problem
Variance-preserving SDE: dx = −1√2β(t)x dt +β(t) dw with β(t) = βmin + t(βmax −βmin). Derive α(t) = ∫t exp(−1√1 −α(t)2.0 β(s)ds), σ(t) = 2
import torch
def vp_alpha_sigma(t, beta_min=0.1, beta_max=20.0):
integral = beta_min * t + 0.5 * (beta_max - beta_min) * t * t
alpha = torch.exp(-0.5 * integral)
sigma = torch.sqrt(1 - alpha ** 2)
return alpha, sigma
Tests
a, s = vp_alpha_sigma(torch.tensor([0.0, 1.0]))
assert a[0] == 1.0 and a[1] < 0.01 and abs(a[0]**2 + s[0]**2 - 1.0) < 1e-6
print("VP OK")
Problem
qd[σ(t)2] dw with σ(t) = σmin(σmax/σmin)t. Compute the drift-free g(t).Variance-exploding SDE: dx = dt
import math
def ve_g(t, sigma_min=0.01, sigma_max=50.0):
sigma = sigma_min * (sigma_max / sigma_min) ** t
return sigma * math.sqrt(2 * math.log(sigma_max / sigma_min))
Tests
g = ve_g(torch.tensor([0.0, 1.0]))
assert g[1] > g[0]; print("VE OK")
Problem
Given ᾱt, compute SNRt = ᾱt/(1 −ᾱt) and log SNRt = log ¯αt −log(1 −ᾱt).
import torch
def snr(abar):
return abar / (1 - abar).clamp(min=1e-12)
def log_snr(abar):
return torch.log(abar.clamp(min=1e-12)) - torch.log((1 - abar).clamp(min=1e-12))
Tests
ab = torch.tensor([0.99, 0.5, 0.01])
assert snr(ab)[0] > snr(ab)[-1]
assert log_snr(ab)[0] > log_snr(ab)[-1]
print("SNR OK")
Problem
Karras et al. 2022 sample σ values via σi = (σ1/ρmax + i/(N −1)(σ1/ρmin −σ1/ρmax))ρ.
import torch
def karras_sigmas(N, sigma_min=0.002, sigma_max=80.0, rho=7.0):
i = torch.arange(N, dtype=torch.float64)
inv_min = sigma_min ** (1 / rho)
inv_max = sigma_max ** (1 / rho)
return (inv_max + i / (N - 1) * (inv_min - inv_max)) ** rho
Tests
sig = karras_sigmas(40)
assert sig[0] > sig[-1] and abs(sig[0].item() - 80.0) < 1e-3
print("Karras OK")
Problem
xt = √ᾱt x0 + √1 −ᾱt ε, ε ∼N(0, I).
import torch
def q_sample(x0, t, abar, eps=None):
eps = torch.randn_like(x0) if eps is None else eps
a = abar[t].view(-1, *([1] * (x0.dim() - 1)))
return a.sqrt() * x0 + (1 - a).sqrt() * eps, eps
Tests
T = 1000; _, _, ab = linear_beta(T)
x0 = torch.randn(2, 3, 16, 16); t = torch.tensor([0, T - 1])
xt, _ = q_sample(x0, t, ab)
assert torch.allclose(xt[0], x0[0], atol=1e-2)
assert xt[1].std() > 0.5
print("q_sample OK")
Problem
xt = √1 −βt xt−1 + √βt ε.
import math
import torch
def q_step(x_prev, beta_t):
return math.sqrt(1 - beta_t) * x_prev + math.sqrt(beta_t) * torch.randn_like(x_prev)
Tests
x = torch.zeros(4); xt = q_step(x, 0.0)
assert torch.allclose(xt, x); print("q_step OK")
Problem
√ √αt(1−ᾱt−1) ᾱt−1βtxt, β˜ = 1−ᾱt−1 Closed form (DDPM Eq. 7): µ˜ =x0 +1−ᾱt βt. 1−ᾱt1−ᾱt
import torch
def posterior(x0, xt, t, betas, alphas, abar):
abar_prev = torch.cat([torch.tensor([1.0], device=abar.device), abar[:-1]])
coef0 = (abar_prev[t].sqrt() * betas[t]) / (1 - abar[t])
coeft = (alphas[t].sqrt() * (1 - abar_prev[t])) / (1 - abar[t])
mu = coef0 * x0 + coeft * xt
var = (1 - abar_prev[t]) / (1 - abar[t]) * betas[t]
return mu, var
Tests
T = 100; betas, alphas, ab = linear_beta(T)
x0 = torch.zeros(1, 4); xt = torch.zeros(1, 4)
mu, var = posterior(x0, xt, torch.tensor([T - 1]), betas, alphas, ab)
assert torch.allclose(mu, torch.zeros_like(mu)) and var > 0
print("posterior OK")
Problem
KL(N(µ1, σ21) ∥N(µ2, σ22)) = log σ2/σ1 + (σ21 + (µ1 −µ2)2)/(2σ22) −1/2.
import torch
def kl_gauss(mu1, var1, mu2, var2):
return 0.5 * (torch.log(var2 / var1) + (var1 + (mu1 - mu2) ** 2) / var2 - 1)
Tests
mu = torch.zeros(4); var = torch.ones(4)
assert torch.allclose(kl_gauss(mu, var, mu, var), torch.zeros(4), atol=1e-7)
print("KL gauss OK")
Problem
∥ε −εθ(xt, t)∥2 L = Ex0,t,ε .
import torch
import torch.nn.functional as F
def ddpm_eps_loss(model, x0, abar):
t = torch.randint(0, abar.size(0), (x0.size(0),), device=x0.device)
xt, eps = q_sample(x0, t, abar)
pred = model(xt, t)
return F.mse_loss(pred, eps)
Tests
class Tiny(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv2d(3, 3, 3, padding=1)
def forward(self, x, t): return self.c(x)
_, _, ab = linear_beta(50)
loss = ddpm_eps_loss(Tiny(), torch.randn(2, 3, 8, 8), ab)
assert loss > 0; print("eps loss OK")
Problem
Given xˆ0 = (xt −√1 −ᾱtεˆ)/√ᾱt, the equivalent x0-prediction loss is ∥x0 −xˆ0∥2.
import torch
import torch.nn.functional as F
def x0_loss(model, x0, abar):
t = torch.randint(0, abar.size(0), (x0.size(0),), device=x0.device)
xt, _ = q_sample(x0, t, abar)
pred_x0 = model(xt, t)
return F.mse_loss(pred_x0, x0)
Tests
class Tiny(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv2d(3, 3, 3, padding=1)
def forward(self, x, t): return self.c(x)
_, _, ab = linear_beta(50)
loss = x0_loss(Tiny(), torch.randn(2, 3, 8, 8), ab); assert loss > 0
print("x0 loss OK")
Problem
vt = √ᾱt ε −√1 −ᾱt x0. Train via ∥v −vθ∥2; distillation works better in v-space.
import torch
import torch.nn.functional as F
def v_target(x0, eps, abar, t):
a = abar[t].view(-1, *([1]*(x0.dim()-1)))
return a.sqrt() * eps - (1 - a).sqrt() * x0
def v_loss(model, x0, abar):
t = torch.randint(0, abar.size(0), (x0.size(0),), device=x0.device)
xt, eps = q_sample(x0, t, abar)
v_t = v_target(x0, eps, abar, t)
return F.mse_loss(model(xt, t), v_t)
Tests
class Tiny(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv2d(3, 3, 3, padding=1)
def forward(self, x, t): return self.c(x)
_, _, ab = linear_beta(50)
loss = v_loss(Tiny(), torch.randn(2, 3, 8, 8), ab); assert loss > 0
print("v loss OK")
Problem
√√√α x̄t −√1 −α v¯; εˆ =α v¯ + √1 −α x̄t; vˆ =ᾱ ˆε −√1 −ᾱ ˆx0. At any time t with state xt: xˆ0 =
def v_to_x0(v, xt, abar_t):
return abar_t.sqrt() * xt - (1 - abar_t).sqrt() * v
def v_to_eps(v, xt, abar_t):
return abar_t.sqrt() * v + (1 - abar_t).sqrt() * xt
def x0_eps_to_v(x0, eps, abar_t):
return abar_t.sqrt() * eps - (1 - abar_t).sqrt() * x0
Tests
abar_t = torch.tensor(0.4)
x0 = torch.randn(4); eps = torch.randn(4)
xt = abar_t.sqrt() * x0 + (1 - abar_t).sqrt() * eps
v = x0_eps_to_v(x0, eps, abar_t)
assert torch.allclose(v_to_x0(v, xt, abar_t), x0, atol=1e-5)
assert torch.allclose(v_to_eps(v, xt, abar_t), eps, atol=1e-5)
print("conversions OK")
Problem
Weight the per-timestep loss by min(SNRt, γ)/SNRt for ε-pred, or min(SNRt, γ)/(SNRt + 1) for v-pred.
import torch
def min_snr_weight(abar_t, gamma=5.0, mode='eps'):
snr_ = abar_t / (1 - abar_t)
cap = torch.minimum(snr_, torch.tensor(gamma, device=snr_.device))
if mode == 'eps': return cap / snr_
if mode == 'x0': return cap
if mode == 'v': return cap / (snr_ + 1)
raise ValueError(mode)
Tests
ab = torch.tensor([0.99, 0.5, 0.01])
w_eps = min_snr_weight(ab, gamma=5, mode='eps')
assert w_eps[0] < 1.0 and w_eps[-1].item() == 1.0 # high noise => weight 1
print("min-SNR OK", w_eps)
Problem
Embed scalar t into D-dim with sinusoidal frequencies (DDPM, ADM, DiT).
import math
import torch
import torch.nn.functional as F
def time_embedding(t, dim, max_period=10000.0):
half = dim // 2
freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
args = t.float().unsqueeze(-1) * freqs.unsqueeze(0)
emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
if dim % 2: emb = F.pad(emb, (0, 1))
return emb
Tests
emb = time_embedding(torch.tensor([0, 100, 1000]), 32)
assert emb.shape == (3, 32) and emb.std() > 0; print("time emb OK")
Problem
adaLN produces γ, β from a conditioning embedding and applies them after a LayerNorm: y = γ · LN(x) + β. Optionally include a scaled-residual multiplier α.
import torch.nn as nn
class AdaLN(nn.Module):
def __init__(self, dim, cond_dim):
super().__init__()
self.ln = nn.LayerNorm(dim, elementwise_affine=False)
self.proj = nn.Linear(cond_dim, 3 * dim)
def forward(self, x, cond):
h = self.ln(x)
gamma, beta, alpha = self.proj(cond).chunk(3, dim=-1)
gamma = gamma.unsqueeze(1) if x.dim() == 3 else gamma
beta = beta.unsqueeze(1) if x.dim() == 3 else beta
alpha = alpha.unsqueeze(1) if x.dim() == 3 else alpha
return alpha * ((1 + gamma) * h + beta)
Tests
m = AdaLN(64, 32); y = m(torch.randn(2, 16, 64), torch.randn(2, 32))
assert y.shape == (2, 16, 64); print("AdaLN OK")
Problem
Maintain an exponential moving average of the model weights: θema ← µθema + (1 −µ)θ.
import torch
class EMA:
def __init__(self, model, decay=0.9999):
self.decay = decay
self.shadow = {n: p.detach().clone() for n, p in model.named_parameters()}
@torch.no_grad()
def update(self, model):
for n, p in model.named_parameters():
self.shadow[n].mul_(self.decay).add_(p.data, alpha=1 - self.decay)
def copy_to(self, model):
for n, p in model.named_parameters():
p.data.copy_(self.shadow[n])
Tests
m = nn.Linear(4, 4); ema = EMA(m, decay=0.5)
m.weight.data.add_(1.0)
ema.update(m)
# half of original (zeros) + half of new
print("EMA OK")
Advanced implementation. Fuse the per-tensor loop into two multi-tensor foreach calls — the same kernel-fusion idea as torch.optim's foreach mode, and it matters because EMA runs every training step over every parameter. Verified equal to the loop EMA (1e-7).
import torch
@torch.no_grad()
def ema_foreach(shadow, params, d=0.9999):
torch._foreach_mul_(shadow, d)
torch._foreach_add_(shadow, params, alpha=1 - d)
Problem
import torch
def ddpm_reverse_step(xt, eps_pred, t, betas, alphas, abar):
a_t = alphas[t]; b_t = betas[t]; ab_t = abar[t]
coef = (1 - a_t) / (1 - ab_t).sqrt()
mean = (1.0 / a_t.sqrt()) * (xt - coef * eps_pred)
if t == 0: return mean
return mean + b_t.sqrt() * torch.randn_like(xt)
Tests
T = 100; betas, alphas, ab = linear_beta(T)
xt = torch.randn(1, 1, 4, 4)
eps_pred = torch.zeros_like(xt)
out = ddpm_reverse_step(xt, eps_pred, t=10, betas=betas, alphas=alphas, abar=ab)
assert out.shape == xt.shape; print("ddpm step OK")
Problem
Starting from xT ∼N(0, I), run T reverse steps to recover x0.
import torch
@torch.no_grad()
def ddpm_sample(model, shape, betas, alphas, abar, device='cpu'):
x = torch.randn(shape, device=device)
T = abar.size(0)
for t in reversed(range(T)):
t_b = torch.full((shape[0],), t, device=device, dtype=torch.long)
eps = model(x, t_b)
x = ddpm_reverse_step(x, eps, t, betas, alphas, abar)
return x
Tests
class Tiny(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv2d(1, 1, 3, padding=1)
def forward(self, x, t): return self.c(x)
b, a, ab = linear_beta(20)
x0 = ddpm_sample(Tiny(), (1, 1, 4, 4), b, a, ab)
assert x0.shape == (1, 1, 4, 4); print("ddpm sample OK")
Problem
xt−1 = √ᾱt−1 ˆx0 + √1 −ᾱt−1 ˆε, where xˆ0 = (xt −√1 −ᾱtεˆ)/√ᾱt.
import torch
def ddim_step(xt, eps_pred, abar_t, abar_prev, eta=0.0):
x0_pred = (xt - (1 - abar_t).sqrt() * eps_pred) / abar_t.sqrt()
sigma = eta * ((1 - abar_prev) / (1 - abar_t)).sqrt() * (1 - abar_t / abar_prev).sqrt()
dir_xt = (1 - abar_prev - sigma ** 2).clamp(min=0).sqrt() * eps_pred
noise = sigma * torch.randn_like(xt) if eta > 0 else 0.0
return abar_prev.sqrt() * x0_pred + dir_xt + noise
Tests
abar_t = torch.tensor(0.5); abar_prev = torch.tensor(0.7)
x0 = ddim_step(torch.zeros(4), torch.zeros(4), abar_t, abar_prev, eta=0.0)
assert torch.allclose(x0, torch.zeros(4)); print("ddim step OK")
Problem
η = 0 gives the deterministic ODE; η = 1 recovers the DDPM posterior variance for the chosen discretisation.
def ddim_eta(xt, eps_pred, abar_t, abar_prev, eta=1.0):
return ddim_step(xt, eps_pred, abar_t, abar_prev, eta=eta)
Tests
torch.manual_seed(0)
abar_t = torch.tensor(0.5); abar_prev = torch.tensor(0.7)
out = ddim_eta(torch.zeros(8), torch.zeros(8), abar_t, abar_prev, eta=1.0)
assert out.std() > 0; print("ddim eta>0 OK")
Problem
Sample with S ≪ T timesteps by taking ᾱti at ti = iT/S.
import torch
@torch.no_grad()
def ddim_sample(model, shape, abar, S=50, eta=0.0, device='cpu'):
T = abar.size(0)
ts = torch.linspace(T - 1, 0, S + 1, device=device).long()
x = torch.randn(shape, device=device)
for i in range(S):
t = ts[i]; t_prev = ts[i + 1]
eps = model(x, torch.full((shape[0],), t, device=device, dtype=torch.long))
x = ddim_step(x, eps, abar[t], abar[t_prev] if t_prev >= 0 else torch.tensor(1.0), eta)
return x
Tests
class Tiny(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv2d(1, 1, 3, padding=1)
def forward(self, x, t): return self.c(x)
_, _, ab = linear_beta(100)
out = ddim_sample(Tiny(), (1, 1, 4, 4), ab, S=20)
assert out.shape == (1, 1, 4, 4); print("ddim sample OK")
Problem
Run DDIM in reverse direction with the model’s ε predictions to recover xT from x0.
import torch
@torch.no_grad()
def ddim_invert(model, x0, abar, S=50, device='cpu'):
T = abar.size(0)
ts = torch.linspace(0, T - 1, S + 1, device=device).long()
x = x0
for i in range(S):
t = ts[i]; t_next = ts[i + 1]
eps = model(x, torch.full((x0.size(0),), t, device=device, dtype=torch.long))
x0_pred = (x - (1 - abar[t]).sqrt() * eps) / abar[t].sqrt()
x = abar[t_next].sqrt() * x0_pred + (1 - abar[t_next]).sqrt() * eps
return x
Tests
class Tiny(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv2d(1, 1, 3, padding=1)
def forward(self, x, t): return self.c(x)
_, _, ab = linear_beta(100)
xT = ddim_invert(Tiny(), torch.randn(1, 1, 4, 4), ab, S=10)
assert xT.shape == (1, 1, 4, 4); print("ddim invert OK")
Problem
Map S subsampled steps to the T-step schedule; pre-cache the ᾱ values.
import torch
def respaced_abar(abar, S):
T = abar.size(0)
idx = torch.linspace(0, T - 1, S, dtype=torch.long)
return abar[idx], idx
Tests
_, _, ab = linear_beta(1000)
sub, idx = respaced_abar(ab, 50)
assert sub.size(0) == 50 and idx[0] == 0 and idx[-1] == 999
print("respaced OK")
Problem
∇x log pt(xt) = −εθ(xt, t)/σt where σt = √1 −ᾱt.
def score_from_eps(eps, abar_t):
sigma = (1 - abar_t).sqrt()
return -eps / sigma
Tests
ab_t = torch.tensor(0.5); eps = torch.ones(4)
s = score_from_eps(eps, ab_t); assert torch.allclose(s, -eps / math.sqrt(0.5))
print("score from eps OK")
Problem
t sθ(xt))/αt, equivalent to xˆ0 = (xt −√1 −ᾱt εˆ)/√ᾱt. For Gaussian noise: E[x0|xt] = (xt + σ2
def tweedie(xt, score, abar_t):
a = abar_t.sqrt(); s = (1 - abar_t).sqrt()
return (xt + s ** 2 * score) / a
Tests
ab_t = torch.tensor(0.5); xt = torch.zeros(4); score = torch.zeros(4)
assert torch.allclose(tweedie(xt, score, ab_t), torch.zeros(4))
print("tweedie OK")
Problem
∥sθ(xt, t) + ε/σt∥2L = Ex0,t,ε· σ2 t . This is equivalent to ε-MSE.
import torch
def dsm_loss(model_score, x0, abar):
t = torch.randint(0, abar.size(0), (x0.size(0),), device=x0.device)
xt, eps = q_sample(x0, t, abar)
a = abar[t].view(-1, *([1]*(x0.dim()-1)))
sigma = (1 - a).sqrt()
score = model_score(xt, t)
return (sigma * score + eps).pow(2).mean()
Tests
class S(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv2d(3, 3, 3, padding=1)
def forward(self, x, t): return self.c(x)
_, _, ab = linear_beta(50)
loss = dsm_loss(S(), torch.randn(2, 3, 8, 8), ab); assert loss > 0
print("DSM OK")
Problem
For a VP SDE, the equivalent ODE is x˙ = −12β(t)x −12β(t)sθ(x, t). Take an Euler step.
def pf_ode_step(x, t, beta_t, score):
drift = -0.5 * beta_t * (x + score)
return x + drift * (-1.0) # integrating from t to t-dt with dt=1 in discrete units
Tests
out = pf_ode_step(torch.zeros(4), 0.5, 0.1, torch.zeros(4))
assert torch.allclose(out, torch.zeros(4)); print("pf ODE OK")
Problem
dx = [f(x, t) −g(t)2sθ(x, t)]dt + g(t) dw ¯ for the reverse-time SDE.
import math
import torch
def reverse_sde_step(x, t, dt, f_x, g_t, score):
drift = f_x - g_t ** 2 * score
diff = g_t * math.sqrt(abs(dt)) * torch.randn_like(x)
return x + drift * dt + diff
Tests
out = reverse_sde_step(torch.zeros(4), 0.5, -0.01, torch.zeros(4), 0.0, torch.zeros(4))
assert torch.allclose(out, torch.zeros(4)); print("reverse SDE OK")
Problem
One predictor step (e.g. reverse SDE), followed by N Langevin corrector steps with step size δ= 2 (ε ∥εnoise∥/ ∥sθ∥)2.
import torch
def langevin_corrector(x, score_fn, t_idx, n_steps=1, snr=0.16):
for _ in range(n_steps):
s = score_fn(x, t_idx)
z = torch.randn_like(x)
step = 2 * (snr * z.norm() / s.norm().clamp(min=1e-8)) ** 2
x = x + step * s + (2 * step).sqrt() * z
return x
Tests
out = langevin_corrector(torch.zeros(8), lambda x, t: torch.zeros_like(x), torch.tensor(0))
assert out.std() == 0.0; print("PC corrector OK")
Problem
Combine conditional and unconditional predictions: ε˜ = (1+w) εc −w εu. With w = 0 recover the conditional model; w →∞amplifies condition signal.
def cfg(eps_uncond, eps_cond, w=7.5):
return (1 + w) * eps_cond - w * eps_uncond
Tests
eu = torch.zeros(4); ec = torch.ones(4)
assert torch.allclose(cfg(eu, ec, w=0.0), ec)
assert torch.allclose(cfg(eu, ec, w=1.0), 2 * ec)
print("CFG OK")
Advanced implementation. Never run two forwards: concatenate (x, x) with (cond, null-cond) into ONE 2x-batch forward and split — identical numbers, far better GPU utilization; this is how every real sampler implements CFG (and extends to negative prompts as a 3x batch). Verified equal to the two-forward version (1e-6).
import torch
@torch.no_grad()
def cfg_batched(model, x, cond, null_cond, w):
eps = model(torch.cat([x, x]), torch.cat([cond, null_cond]))
eps_c, eps_u = eps.chunk(2)
return eps_u + w * (eps_c - eps_u)
Problem
Replace the unconditional prediction with a negative-prompt prediction: ε˜ = (1 + w) εc −w εneg.
def cfg_neg(eps_neg, eps_cond, w=7.5):
return (1 + w) * eps_cond - w * eps_neg
Tests
out = cfg_neg(torch.tensor([1.]), torch.tensor([2.]), w=2.0)
assert out.item() == 4.0; print("neg CFG OK")
Problem
During training, drop the conditioning with probability pdrop so the same network learns both p(x|c) and p(x).
import torch
def maybe_drop_cond(c, p_drop=0.1, null_cond=None):
if null_cond is None: null_cond = torch.zeros_like(c)
mask = (torch.rand(c.size(0), device=c.device) < p_drop).view(-1, *([1] * (c.dim() - 1)))
return torch.where(mask, null_cond, c)
Tests
torch.manual_seed(0)
c = torch.ones(1000, 4)
out = maybe_drop_cond(c, p_drop=0.5)
assert 0.4 < (out.sum(-1) == 0).float().mean().item() < 0.6
print("drop cond OK")
Problem
Implement a cross-attention layer where queries come from the image latent and keys/values come from text embeddings.
import torch.nn as nn
import torch.nn.functional as F
class CrossAttn(nn.Module):
def __init__(self, d, n_heads, ctx_dim):
super().__init__()
self.h = n_heads; self.dk = d // n_heads
self.q = nn.Linear(d, d, bias=False)
self.k = nn.Linear(ctx_dim, d, bias=False)
self.v = nn.Linear(ctx_dim, d, bias=False)
self.proj = nn.Linear(d, d)
def forward(self, x, ctx):
B, T, D = x.shape; _, M, _ = ctx.shape
q = self.q(x).view(B, T, self.h, self.dk).transpose(1, 2)
k = self.k(ctx).view(B, M, self.h, self.dk).transpose(1, 2)
v = self.v(ctx).view(B, M, self.h, self.dk).transpose(1, 2)
out = F.scaled_dot_product_attention(q, k, v)
return self.proj(out.transpose(1, 2).contiguous().view(B, T, D))
Tests
m = CrossAttn(64, 4, ctx_dim=128)
y = m(torch.randn(2, 16, 64), torch.randn(2, 8, 128))
assert y.shape == (2, 16, 64); print("CrossAttn OK")
Problem
For an ODE x˙ = f(x, t) with step ∆t: x˜ = x + ∆t f(x, t), xt+∆t = x + ∆t2 (f(x, t) + f(˜x, t + ∆t)).
def heun_step(f, x, t, dt):
f1 = f(x, t)
x_pred = x + dt * f1
f2 = f(x_pred, t + dt)
return x + 0.5 * dt * (f1 + f2)
Tests
def f(x, t): return -x
x = torch.tensor([1.0]); x = heun_step(f, x, 0.0, 0.1)
assert abs(x.item() - math.exp(-0.1)) < 1e-3
print("Heun OK")
Problem
Classical RK4: weighted average of four slope evaluations.
def rk4_step(f, x, t, dt):
k1 = f(x, t)
k2 = f(x + 0.5 * dt * k1, t + 0.5 * dt)
k3 = f(x + 0.5 * dt * k2, t + 0.5 * dt)
k4 = f(x + dt * k3, t + dt)
return x + dt * (k1 + 2 * k2 + 2 * k3 + k4) / 6
Tests
def f(x, t): return -x
x = torch.tensor([1.0]); x = rk4_step(f, x, 0.0, 0.1)
assert abs(x.item() - math.exp(-0.1)) < 1e-6
print("RK4 OK")
Problem
First-order DPM-Solver step in λ = log(¯α/(1 −ᾱ))/2 (half-log-SNR) is equivalent to DDIM. Express it in terms of σ.
def dpm_solver1(xt, eps_pred, sigma_t, sigma_s):
# x_s = (sigma_s/sigma_t) x_t + (alpha_s - sigma_s/sigma_t * alpha_t) * (-eps)
alpha_t = (1 - sigma_t ** 2).sqrt()
alpha_s = (1 - sigma_s ** 2).sqrt()
return (sigma_s / sigma_t) * xt + (alpha_s - sigma_s / sigma_t * alpha_t) * (-eps_pred)
Tests
xt = torch.zeros(4); eps = torch.zeros(4)
out = dpm_solver1(xt, eps, torch.tensor(0.5), torch.tensor(0.4))
assert torch.allclose(out, torch.zeros(4)); print("DPM-Solver-1 OK")
Problem
Second-order step: predict at the midpoint to refine the noise estimate.
def dpm_solver2(model, xt, t, t_next, abar_fn):
abar_t = abar_fn(t); abar_n = abar_fn(t_next)
sigma_t = (1 - abar_t).sqrt(); sigma_n = (1 - abar_n).sqrt()
eps1 = model(xt, t)
# Midpoint
t_mid = (t + t_next) / 2
abar_m = abar_fn(t_mid); sigma_m = (1 - abar_m).sqrt()
x_mid = dpm_solver1(xt, eps1, sigma_t, sigma_m)
eps2 = model(x_mid, t_mid)
return dpm_solver1(xt, eps2, sigma_t, sigma_n)
Tests
def abar_fn(t): return torch.tensor(1 - 0.1 * t)
class M(nn.Module):
def forward(self, x, t): return torch.zeros_like(x)
out = dpm_solver2(M(), torch.zeros(4), 0.2, 0.1, abar_fn)
assert torch.allclose(out, torch.zeros(4)); print("DPM-Solver-2 OK")
Problem
Sample with σ-parameterised denoiser D(x, σ): step x˙ = (x −D(x, σ))/σ using Heun.
import torch
@torch.no_grad()
def edm_heun_sample(D, shape, sigmas, device='cpu'):
x = torch.randn(shape, device=device) * sigmas[0]
for i in range(len(sigmas) - 1):
sig_t = sigmas[i]; sig_n = sigmas[i + 1]
d = (x - D(x, sig_t)) / sig_t
x_pred = x + (sig_n - sig_t) * d
if sig_n != 0:
d_pred = (x_pred - D(x_pred, sig_n)) / sig_n
x = x + 0.5 * (sig_n - sig_t) * (d + d_pred)
else:
x = x_pred
return x
Tests
def D(x, sig): return x # trivial denoiser
x0 = edm_heun_sample(D, (1, 4), karras_sigmas(20))
assert x0.shape == (1, 4); print("EDM Heun OK")
Problem
Encode an image to a latent z = Encoder(x) · s (LDM uses s ≈ 0.18215 for SD).
import torch
def encode_to_latent(encoder, x, scale=0.18215):
with torch.no_grad():
return encoder(x) * scale
def decode_from_latent(decoder, z, scale=0.18215):
with torch.no_grad():
return decoder(z / scale)
Tests
class E(nn.Module):
def forward(self, x): return F.avg_pool2d(x, 8)
class D(nn.Module):
def forward(self, z): return F.interpolate(z, scale_factor=8)
z = encode_to_latent(E(), torch.randn(1, 3, 32, 32))
y = decode_from_latent(D(), z)
assert y.shape == (1, 3, 32, 32); print("LDM enc/dec OK")
Problem
Encode batch of images to latent space, then run standard ε-prediction loss in latent space.
def latent_diffusion_loss(encoder, model, x_img, abar, scale=0.18215):
z = encode_to_latent(encoder, x_img, scale)
return ddpm_eps_loss(model, z, abar)
Tests
class E(nn.Module):
def forward(self, x): return F.avg_pool2d(x, 8)
class M(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv2d(3, 3, 3, padding=1)
def forward(self, x, t): return self.c(x)
_, _, ab = linear_beta(50)
loss = latent_diffusion_loss(E(), M(), torch.randn(2, 3, 32, 32), ab); assert loss > 0
print("LDM loss OK")
Problem
For independent samples x0 ∼p0, x1 ∼p1, define the linear interpolation xt = (1 −t)x0 + tx1, target velocity v∗= x1 −x0.
def fm_path(x0, x1, t):
t = t.view(-1, *([1]*(x0.dim()-1)))
return (1 - t) * x0 + t * x1, x1 - x0
Tests
x0 = torch.zeros(4, 3); x1 = torch.ones(4, 3)
t = torch.full((4,), 0.5)
xt, v = fm_path(x0, x1, t)
assert torch.allclose(xt, 0.5 * torch.ones_like(xt))
assert torch.allclose(v, torch.ones_like(v))
print("FM path OK")
Problem
L = Ex0,x1,t[∥vθ(xt, t) −(x1 −x0)∥2] with t ∼U(0, 1).
import torch
import torch.nn.functional as F
def fm_loss(model, x0, x1):
t = torch.rand(x0.size(0), device=x0.device)
xt, v_target = fm_path(x0, x1, t)
pred = model(xt, t)
return F.mse_loss(pred, v_target)
Tests
class M(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 4)
def forward(self, x, t): return self.l(x)
loss = fm_loss(M(), torch.randn(8, 4), torch.randn(8, 4))
assert loss > 0; print("FM loss OK")
Problem
Integrate x˙t = vθ(xt, t) from t = 0 to t = 1 with N Euler steps starting from x0 ∼p0.
import torch
@torch.no_grad()
def fm_euler_sample(v_model, x0, N=50):
x = x0; dt = 1.0 / N
for i in range(N):
t = torch.full((x.size(0),), i * dt, device=x.device)
x = x + dt * v_model(x, t)
return x
Tests
class V(nn.Module):
def forward(self, x, t): return torch.ones_like(x) # constant velocity
x0 = torch.zeros(2, 4); x1 = fm_euler_sample(V(), x0, N=10)
assert torch.allclose(x1, torch.ones_like(x1), atol=1e-6)
print("FM Euler OK")
Problem
Use Heun’s method to integrate the flow-matching ODE.
import torch
@torch.no_grad()
def fm_heun_sample(v_model, x0, N=20):
x = x0; dt = 1.0 / N
for i in range(N):
t = torch.full((x.size(0),), i * dt, device=x.device)
tn = torch.full((x.size(0),), (i + 1) * dt, device=x.device)
v1 = v_model(x, t)
x_p = x + dt * v1
v2 = v_model(x_p, tn)
x = x + 0.5 * dt * (v1 + v2)
return x
Tests
class V(nn.Module):
def forward(self, x, t): return -x # ODE: x' = -x => x_1 = x_0 * exp(-1)
x0 = torch.ones(4); x1 = fm_heun_sample(V(), x0, N=50)
assert torch.allclose(x1, torch.full_like(x1, math.exp(-1)), atol=1e-3)
print("FM Heun OK")
Problem
Replace the random pairing of (x0, x1) with an OT-optimal pairing within the batch (greedy or Sinkhorn) before applying flow-matching loss.
from scipy.optimize import linear_sum_assignment
def ot_pair(x0, x1):
cost = ((x0[:, None] - x1[None]) ** 2).sum(-1).cpu().numpy()
r, c = linear_sum_assignment(cost)
return x0[r], x1[c]
Tests
torch.manual_seed(0)
x0 = torch.randn(4, 2); x1 = x0 + torch.randn(4, 2) * 0.1
xa, xb = ot_pair(x0, x1)
assert torch.allclose((xa - xb).norm(dim=1).mean(), torch.tensor(0.1), atol=0.2)
print("OT-FM pair OK")
Problem
With x0 ∼N(0, I) as the noise and x1 as data, train E ∥vθ((1 −t)x0 + tx1, t) −(x1 −x0)∥2. Identical func- tional form to flow matching with Gaussian source.
import torch
import torch.nn.functional as F
def rectflow_loss(model, x_data):
x0 = torch.randn_like(x_data)
x1 = x_data
t = torch.rand(x_data.size(0), device=x_data.device)
xt, v_t = fm_path(x0, x1, t)
return F.mse_loss(model(xt, t), v_t)
Tests
class M(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 4)
def forward(self, x, t): return self.l(x)
loss = rectflow_loss(M(), torch.randn(4, 4)); assert loss > 0
print("RF loss OK")
Problem
Generate trajectories from the current model, take the endpoints (x0, ˆx1), and re-train flow matching on these aligned pairs — this reduces curvature and enables few-step sampling.
import torch
import torch.nn.functional as F
@torch.no_grad()
def make_reflow_pairs(model, n=1024, dim=4, N_steps=50):
x0 = torch.randn(n, dim)
x1 = fm_euler_sample(model, x0.clone(), N=N_steps)
return x0, x1
def reflow_loss(model, x0, x1):
t = torch.rand(x0.size(0), device=x0.device)
xt, v_t = fm_path(x0, x1, t)
return F.mse_loss(model(xt, t), v_t)
Tests
class V(nn.Module):
def forward(self, x, t): return torch.ones_like(x)
x0, x1 = make_reflow_pairs(V(), n=8, dim=2, N_steps=10)
loss = reflow_loss(V(), x0, x1); print("reflow OK", loss.item())
Problem
Train a one-step student G(x0) to match the multi-step teacher’s endpoint xˆ1.
import torch
import torch.nn.functional as F
def one_step_distill_loss(student, teacher, x0, N=20):
with torch.no_grad():
target = fm_euler_sample(teacher, x0.clone(), N=N)
return F.mse_loss(student(x0, torch.zeros(x0.size(0))), target)
Tests
class M(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 4)
def forward(self, x, t): return self.l(x)
loss = one_step_distill_loss(M(), M(), torch.randn(4, 4), N=5); assert loss >= 0
print("1-step distill OK")
Problem
Split x = (x1, x2). Compute (s, t) = f(x1), transform y1 = x1, y2 = x2 ⊙exp(s) + t; log-det Jacobian = Σ s.
import torch
import torch.nn as nn
class Coupling(nn.Module):
def __init__(self, dim, hidden=64):
super().__init__()
d = dim // 2
self.net = nn.Sequential(nn.Linear(d, hidden), nn.ReLU(), nn.Linear(hidden, 2 * (dim - d)))
self.split = d
def forward(self, x):
x1, x2 = x[:, :self.split], x[:, self.split:]
s, t = self.net(x1).chunk(2, dim=-1)
s = torch.tanh(s) # bounded for stability
y2 = x2 * s.exp() + t
log_det = s.sum(dim=-1)
return torch.cat([x1, y2], dim=-1), log_det
def inverse(self, y):
y1, y2 = y[:, :self.split], y[:, self.split:]
s, t = self.net(y1).chunk(2, dim=-1)
s = torch.tanh(s)
x2 = (y2 - t) * (-s).exp()
return torch.cat([y1, x2], dim=-1), -s.sum(dim=-1)
Tests
m = Coupling(4); x = torch.randn(2, 4)
y, ld = m(x); xb, ld_inv = m.inverse(y)
assert torch.allclose(x, xb, atol=1e-5) and torch.allclose(ld, -ld_inv)
print("coupling OK")
Problem
Compose K couplings, alternating which half is conditioned on (via permutation).
import math
import torch.nn as nn
class RealNVP(nn.Module):
def __init__(self, dim, K=4, hidden=64):
super().__init__()
self.layers = nn.ModuleList([Coupling(dim, hidden) for _ in range(K)])
self.flips = [True if i % 2 else False for i in range(K)]
self.dim = dim
def forward(self, x):
log_det = 0.0
for layer, flip in zip(self.layers, self.flips):
if flip: x = x.flip(-1)
x, ld = layer(x); log_det = log_det + ld
return x, log_det
def log_prob(self, x):
z, log_det = self(x)
log_pz = -0.5 * (z ** 2 + math.log(2 * math.pi)).sum(dim=-1)
return log_pz + log_det
Tests
m = RealNVP(4, K=2)
lp = m.log_prob(torch.randn(8, 4))
assert lp.shape == (8,); print("RealNVP OK")
Problem
f(x) = x + u h(w⊤x + b) with h = tanh. Log-det = log |1 + u⊤ψ(x)|, ψ = h′ w.
import torch
import torch.nn as nn
class PlanarFlow(nn.Module):
def __init__(self, dim):
super().__init__()
self.u = nn.Parameter(torch.randn(dim) * 0.01)
self.w = nn.Parameter(torch.randn(dim) * 0.01)
self.b = nn.Parameter(torch.zeros(1))
def forward(self, x):
wx = (x * self.w).sum(-1, keepdim=True) + self.b
h = torch.tanh(wx)
y = x + h * self.u
psi = (1 - h ** 2) * self.w
log_det = (1 + (psi * self.u).sum(-1, keepdim=True)).abs().log().squeeze(-1)
return y, log_det
Tests
m = PlanarFlow(4); y, ld = m(torch.randn(2, 4))
assert y.shape == (2, 4); print("planar OK")
Problem
Compute the divergence of a vector field v(x, t) via the Hutchinson trace estimator:Tr(∂v/∂x)≈ Eε[ε⊤(∂v/∂x)ε].
import torch
def hutch_trace(v_fn, x, t, n_samples=1):
s = 0.0
for _ in range(n_samples):
eps = torch.randn_like(x).requires_grad_(False)
with torch.enable_grad():
x_ = x.detach().requires_grad_(True)
v = v_fn(x_, t)
grad = torch.autograd.grad((v * eps).sum(), x_, create_graph=False)[0]
s = s + (eps * grad).sum(dim=-1)
return s / n_samples
Tests
def v_fn(x, t): return 2 * x # divergence should be 2*dim
x = torch.randn(8, 4)
tr = hutch_trace(v_fn, x, torch.tensor(0.0), n_samples=8)
assert abs(tr.mean().item() - 2 * 4) < 1.5
print("Hutchinson OK", tr.mean().item())
Problem
Implement a 1-D rational-quadratic spline transform with K bins. Returns y and log-det.
import torch
import torch.nn.functional as F
def rqs_1d(x, widths, heights, derivs, lo=-3.0, hi=3.0):
# Simplified: identity outside [lo, hi]; binwise rational-quadratic inside.
x = x.clamp(lo, hi)
# compute bin index
cum_w = torch.cat([torch.zeros(1), torch.cumsum(F.softmax(widths, 0) * (hi - lo), 0)]) + lo
cum_h = torch.cat([torch.zeros(1), torch.cumsum(F.softmax(heights, 0) * (hi - lo), 0)]) + lo
K = widths.size(0)
idx = torch.searchsorted(cum_w, x.contiguous(), right=True).clamp(1, K) - 1
# piecewise linear approximation (simpler stand-in for the full RQS formula)
xL = cum_w[idx]; xR = cum_w[idx + 1]
yL = cum_h[idx]; yR = cum_h[idx + 1]
a = (x - xL) / (xR - xL)
y = yL + a * (yR - yL)
logdet = ((yR - yL) / (xR - xL)).log()
return y, logdet
Tests
torch.manual_seed(0)
w = torch.zeros(8); h = torch.zeros(8); d = torch.zeros(8)
y, ld = rqs_1d(torch.tensor([0.0]), w, h, d)
assert ld.numel() == 1; print("RQS OK")
Problem
A mean-flow model uθ(x, t, r) predicts the average velocity from r to t along the flow. The differential identity for a flow x˙ = v(x, t) is u(xt, t, r) = v(xt, t) −(t −r)∂tu + v ∂xu.
Implement the mean-flow training target u⋆= v −(t −r) dudt via JVP.
import torch
def mean_flow_target(u_model, v_model, x_t, t, r):
# tangent vectors for JVP: d/dt (t fixed direction = 1), velocity in x-direction = v
x_t = x_t.detach().requires_grad_(False)
t = t.detach().requires_grad_(False)
r = r.detach().requires_grad_(False)
v = v_model(x_t, t)
# JVP of u along (v in x, 1 in t, 0 in r)
def fn(x_, t_): return u_model(x_, t_, r)
_, du_dt = torch.autograd.functional.jvp(fn, (x_t, t), (v, torch.ones_like(t)),
create_graph=False)
return v - (t - r).view(-1, *([1]*(x_t.dim()-1))) * du_dt
Tests
class U(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 4)
def forward(self, x, t, r): return self.l(x)
class V(nn.Module):
def forward(self, x, t): return torch.ones_like(x)
target = mean_flow_target(U(), V(), torch.randn(2, 4), torch.tensor([0.5, 0.7]),
torch.tensor([0.1, 0.2]))
assert target.shape == (2, 4); print("MeanFlow target OK")
Problem
∥uθ(xt, t, r) −sg(u⋆)∥2L = Ext,t,r . The stop-gradient on the target prevents trivial solutions.
import torch
import torch.nn.functional as F
def mean_flow_loss(u_model, v_model, x_data):
x0 = torch.randn_like(x_data); x1 = x_data
t = torch.rand(x_data.size(0), device=x_data.device)
r = torch.rand(x_data.size(0), device=x_data.device) * t # r <= t
xt, _ = fm_path(x0, x1, t)
target = mean_flow_target(u_model, v_model, xt, t, r).detach()
return F.mse_loss(u_model(xt, t, r), target)
Tests
class U(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 4)
def forward(self, x, t, r): return self.l(x)
class V(nn.Module):
def forward(self, x, t): return torch.zeros_like(x)
loss = mean_flow_loss(U(), V(), torch.randn(4, 4)); assert loss >= 0
print("MeanFlow loss OK")
Problem
Sample x0 ∼N(0, I) at t = 1, generate x1 = x0 + (1 −0) · uθ(x0, 1, 0). This is a single forward pass.
import torch
@torch.no_grad()
def mean_flow_sample(u_model, n, dim, device='cpu'):
x = torch.randn(n, dim, device=device)
t = torch.ones(n, device=device); r = torch.zeros(n, device=device)
return x + (t - r).view(-1, 1) * u_model(x, t, r)
Tests
class U(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 4)
def forward(self, x, t, r): return self.l(x)
out = mean_flow_sample(U(), 8, 4); assert out.shape == (8, 4)
print("MeanFlow sample OK")
Problem
For two adjacent noise levels σn, σn+1, enforce fθ(x + σnε, σn) ≈ fθ−(x + σn+1ε, σn+1) where θ−is the EMA target.
import torch
import torch.nn.functional as F
def consistency_loss(student, teacher, x0, sigma_n, sigma_np1):
eps = torch.randn_like(x0)
z = x0 + sigma_n * eps
z_p = x0 + sigma_np1 * eps
with torch.no_grad():
target = teacher(z_p, sigma_np1)
return F.mse_loss(student(z, sigma_n), target)
Tests
class S(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 4)
def forward(self, x, sig): return self.l(x)
loss = consistency_loss(S(), S(), torch.randn(4, 4), torch.tensor(0.1), torch.tensor(0.2))
assert loss >= 0; print("CM loss OK")
Problem
Iteratively call fθ from large σ to small σ, adding noise back at each intermediate step.
import torch
@torch.no_grad()
def cm_multistep_sample(f_theta, sigmas, shape):
x = torch.randn(shape) * sigmas[0]
out = f_theta(x, sigmas[0])
for sig in sigmas[1:]:
z = torch.randn_like(out)
x = out + (sig ** 2 - 0.0) ** 0.5 * z
out = f_theta(x, sig)
return out
Tests
class F(nn.Module):
def forward(self, x, s): return x # identity
out = cm_multistep_sample(F(), torch.tensor([0.5, 0.3, 0.1]), (1, 4))
assert out.shape == (1, 4); print("CM sampling OK")
Problem
Run one teacher solver step from σn+1 to σn, use the result as the target for the student at σn.
import torch
import torch.nn.functional as F
def cd_loss(student, teacher_solver, x0, sigma_n, sigma_np1):
eps = torch.randn_like(x0)
z_p = x0 + sigma_np1 * eps
with torch.no_grad():
z_n_target = teacher_solver(z_p, sigma_np1, sigma_n)
return F.mse_loss(student(z_p, sigma_np1), z_n_target)
Tests
class S(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 4)
def forward(self, x, s): return self.l(x)
def teacher_solver(z, s_from, s_to): return z * (s_to / s_from)
loss = cd_loss(S(), teacher_solver, torch.randn(4, 4), torch.tensor(0.1), torch.tensor(0.2))
assert loss >= 0; print("CD OK")
Problem
Standard diffusion ResBlock: GroupNorm → SiLU → Conv2d, plus a per-block timestep scale-shift.
import torch.nn as nn
import torch.nn.functional as F
class ResBlock(nn.Module):
def __init__(self, c, t_dim):
super().__init__()
self.n1 = nn.GroupNorm(8, c); self.c1 = nn.Conv2d(c, c, 3, padding=1)
self.n2 = nn.GroupNorm(8, c); self.c2 = nn.Conv2d(c, c, 3, padding=1)
self.t_proj = nn.Linear(t_dim, 2 * c)
def forward(self, x, t_emb):
h = self.c1(F.silu(self.n1(x)))
scale, shift = self.t_proj(F.silu(t_emb)).chunk(2, dim=-1)
h = (1 + scale[:, :, None, None]) * h + shift[:, :, None, None]
h = self.c2(F.silu(self.n2(h)))
return x + h
Tests
m = ResBlock(8, 16); y = m(torch.randn(1, 8, 16, 16), torch.randn(1, 16))
assert y.shape == (1, 8, 16, 16); print("ResBlock OK")
Problem
Spatial self-attention: flatten the feature map to (B, HW, C), run multi-head attention, reshape back.
import torch.nn as nn
import torch.nn.functional as F
class SelfAttn2D(nn.Module):
def __init__(self, c, h=4):
super().__init__()
self.q = nn.Conv2d(c, c, 1); self.k = nn.Conv2d(c, c, 1); self.v = nn.Conv2d(c, c, 1)
self.proj = nn.Conv2d(c, c, 1); self.h = h; self.dk = c // h
def forward(self, x):
B, C, H, W = x.shape
q = self.q(x).view(B, self.h, self.dk, H * W).transpose(2, 3)
k = self.k(x).view(B, self.h, self.dk, H * W).transpose(2, 3)
v = self.v(x).view(B, self.h, self.dk, H * W).transpose(2, 3)
out = F.scaled_dot_product_attention(q, k, v) # (B, h, HW, dk)
out = out.transpose(2, 3).contiguous().view(B, C, H, W)
return x + self.proj(out)
Tests
m = SelfAttn2D(16); y = m(torch.randn(1, 16, 8, 8))
assert y.shape == (1, 16, 8, 8); print("SelfAttn2D OK")
Problem
Diffusion Transformer block (Peebles & Xie 2022): apply AdaLN modulation around self-attention and MLP, with a learnable zero-init scale α for the residual branches.
import torch.nn as nn
import torch.nn.functional as F
class DiTBlock(nn.Module):
def __init__(self, d, h, mlp_ratio=4, cond_dim=None):
super().__init__()
self.attn = nn.MultiheadAttention(d, h, batch_first=True)
self.mlp = nn.Sequential(nn.Linear(d, mlp_ratio * d), nn.GELU(), nn.Linear(mlp_ratio * d, d))
self.adaLN = nn.Linear(cond_dim or d, 6 * d)
nn.init.zeros_(self.adaLN.weight); nn.init.zeros_(self.adaLN.bias)
def forward(self, x, c):
sh1, sc1, ga1, sh2, sc2, ga2 = self.adaLN(F.silu(c)).chunk(6, dim=-1)
def mod(h, sh, sc): return (1 + sc.unsqueeze(1)) * F.layer_norm(h, h.shape[-1:]) + sh.unsqueeze(1)
h_attn, _ = self.attn(mod(x, sh1, sc1), mod(x, sh1, sc1), mod(x, sh1, sc1))
x = x + ga1.unsqueeze(1) * h_attn
x = x + ga2.unsqueeze(1) * self.mlp(mod(x, sh2, sc2))
return x
Tests
m = DiTBlock(64, 4, cond_dim=64)
y = m(torch.randn(2, 16, 64), torch.randn(2, 64))
assert y.shape == (2, 16, 64); print("DiT block OK")
Problem
ControlNet: copy the encoder of a frozen base UNet, train zero-conv layers, add the auxiliary features back into the base UNet’s skip connections.
import torch.nn as nn
class ZeroConv(nn.Conv2d):
def __init__(self, *a, **kw):
super().__init__(*a, **kw)
nn.init.zeros_(self.weight); nn.init.zeros_(self.bias)
def control_inject(base_skip, control_skip, zero_conv):
return base_skip + zero_conv(control_skip)
Tests
zc = ZeroConv(8, 8, 1)
out = control_inject(torch.randn(1, 8, 4, 4), torch.randn(1, 8, 4, 4), zc)
assert out.shape == (1, 8, 4, 4); print("ControlNet inject OK")
Problem
d/(σ2 + σ2 Wrap a network F as a denoiser: D(x, σ) = cskip(σ) x + cout(σ) F(cin(σ)x, cnoise(σ)) with cskip = σ2d), σ2 + σ2σ2 + σ2d, cnoise = 1√√ cout = σσd/d, cin = 1/4 log σ.
import torch.nn as nn
class EDMDenoiser(nn.Module):
def __init__(self, net, sigma_data=0.5):
super().__init__()
self.net = net; self.sd = sigma_data
def forward(self, x, sigma):
sd = self.sd; s2 = sigma ** 2 + sd ** 2
c_skip = sd ** 2 / s2
c_out = sigma * sd / s2.sqrt()
c_in = 1 / s2.sqrt()
c_noise = sigma.log() / 4
return c_skip * x + c_out * self.net(c_in * x, c_noise)
Tests
class N(nn.Module):
def forward(self, x, c): return torch.zeros_like(x)
m = EDMDenoiser(N())
y = m(torch.randn(1, 4), torch.tensor(0.5))
assert y.shape == (1, 4); print("EDM precond OK")
Problem
Add noise up to time t∗< T, then denoise from there. Lower t∗preserves more of the original.
import torch
@torch.no_grad()
def sdedit(model, x_init, abar, t_star, betas, alphas, device='cpu'):
t = torch.full((x_init.size(0),), t_star, device=device, dtype=torch.long)
x_noisy, _ = q_sample(x_init, t, abar)
x = x_noisy
for tt in reversed(range(t_star + 1)):
eps = model(x, torch.full_like(t, tt))
x = ddpm_reverse_step(x, eps, tt, betas, alphas, abar)
return x
Tests
class M(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv2d(1, 1, 3, padding=1)
def forward(self, x, t): return self.c(x)
b, a, ab = linear_beta(50)
out = sdedit(M(), torch.randn(1, 1, 4, 4), ab, t_star=20, betas=b, alphas=a)
assert out.shape == (1, 1, 4, 4); print("SDEdit OK")
Problem
Implement the optimisation loop: start from DDIM-inverted latents under the null token, optimise per-step null-text embeddings to minimise reconstruction error.
import torch.nn.functional as F
def null_text_invert_step(model, xt, eps_target, null_emb, optimizer, abar_t, abar_prev):
eps_pred = model(xt, null_emb)
loss = F.mse_loss(eps_pred, eps_target)
optimizer.zero_grad(); loss.backward(); optimizer.step()
return loss.item()
Tests
class M(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(4, 4)
def forward(self, x, e): return self.l(x) + e.mean(dim=0, keepdim=True)
emb = torch.zeros(2, 4, requires_grad=True)
opt = torch.optim.Adam([emb], lr=1e-2)
loss = null_text_invert_step(M(), torch.randn(2, 4), torch.randn(2, 4), emb, opt, None, None)
assert loss > 0; print("null-text OK")
Notes
Survival tactics for a live diffusion / flow interview:
Always declare your parameterisation up front: discrete-time DDPM with βt, continuous VP/VE SDE, or σ-parameterised EDM. Each is mathematically equivalent but uses different scaling.
Be explicit about the prediction target: ε, x0, or v. Conversion formulas come up every interview.
DDPM → DDIM: the reverse process becomes deterministic when η = 0 and is equivalent to Euler on the probability-flow ODE in the limit of small steps.
Score and ε-prediction are equivalent: sθ = −εθ/σt. Tweedie’s formula reads off xˆ0 directly.
Classifier-free guidance: practical default is w ∈ [3, 9] for image diffusion; w affects sample quality and diversity (high w = sharper, less diverse).
Higher-order solvers (Heun, DPM-Solver-2, EDM Heun) cut sampling steps by 5–10× for the same quality. Distillation (CM, rectified flow reflow, mean flow) cuts further to 1–4 steps.
Flow matching with Gaussian source ≡rectified flow ≡a deterministic ODE formulation of diffusion (when both use linear interpolation). The training loss is identical except for time sampling and parameterisation conventions.
Mean flow trains u(x, t, r) to predict the average velocity from r to t; the JVP identity u = v −(t−r) du/dt provides a gradient-friendly target. One forward pass = one generation.
Normalizing flows give exact log-likelihood but constrain architectures; CNFs are more flexible but require ODE solving and trace estimation (Hutchinson).
For interview questions on architecture: be ready to write the AdaLN-Zero modulation, time embedding, the cross-attention block, the EDM preconditioning formulas, and a tiny UNet skeleton.
For training tricks: EMA decay ∼0.9999, min-SNR weighting, v-prediction for distillation, zero-init final conv to start from identity, classifier-free dropout p ∼0.1.
End every implementation with a sanity test: at t = 0, q(x0|x0) = x0; at η = 0 DDIM is deterministic; flow matching with constant v moves x by v over one unit of time; mean flow with one step matches the integrated trajectory.