GAN và diffusion

Sinh mẫu khác với dự đoán nhãn

Classifier ước lượng nhãn theo ảnh. Model sinh cần mô tả sự biến thiên của dữ liệu: những kiểu ảnh nào có thể xuất hiện và xuất hiện với tần suất nào. Nếu luôn trả cùng một ảnh đẹp, model bỏ mất sự đa dạng.

Nói theo xác suất, model phân biệt (discriminative) học $p(y \mid x)$, còn model sinh (generative) học chính phân phối dữ liệu $p(x)$, hoặc $p(x \mid c)$ khi có điều kiện $c$ như một câu mô tả, để lấy được mẫu mới. Gõ “con mèo đội mũ len, vẽ màu nước” rồi nhận một bức ảnh mới, như DALL·E hay Stable Diffusion, là sinh ảnh từ văn bản.

Nguồn noise $z$ không phải lỗi cần xóa: nó cung cấp biến thiên để cùng một prompt cho nhiều mẫu. Seed cố định dãy số ngẫu nhiên, nên hai seed cho hai noise khác nhau, còn cùng seed trong cùng điều kiện cho lại cùng thí nghiệm. “Cùng điều kiện” gồm cả model, sampler, phiên bản thư viện và phần cứng: đổi một thứ, ảnh có thể khác dù seed giữ nguyên.

GAN: hai mạng có hai mục tiêu

Generator $G(z)$ biến noise thành mẫu giả. Discriminator $D(x)$ học phân biệt mẫu thật với mẫu do generator tạo. Generator được cập nhật để discriminator đánh giá mẫu giả như thật. Hai mạng thay đổi cùng quá trình, nên mục tiêu mà một mạng thấy cũng đang thay đổi.

Khi dùng, chỉ còn generator chạy: trang web sinh khuôn mặt giả chỉ cần $G$ biến noise thành ảnh, còn $D$ chỉ có ích lúc train.

Một cách train thông dụng dùng binary cross-entropy. Discriminator muốn nhãn 1 cho thật, 0 cho giả:

Với non-saturating generator loss, generator muốn fake được D chấm 1:

Với $D(G(z))=0.2$, loss này là $-\log0.2\approx1.609$; nếu G đẩy được điểm lên 0.8 thì loss còn khoảng 0.223. Phía D, chấm 0.9 cho ảnh thật và 0.2 cho ảnh giả cho loss $-\ln 0.9 - \ln 0.8 \approx 0.33$.

Cách viết minimax nguyên thủy cho G cực tiểu $\E_z\log(1-D(G(z)))$. Hai cách cùng hướng đẩy $D(G(z))$ lên, nhưng khác nhau ở độ lớn gradient. Viết $D=\sigma(a)$ với $a$ là logit của fake:

Đầu train, D dễ nhận ra fake nên $D(G(z))$ gần 0. Khi đó cách viết minimax cho gradient gần 0 (bão hòa) và G gần như đứng yên, còn non-saturating cho gradient có độ lớn gần 1. Vì vậy code thực hành dùng $L_G=-\E_z\log D(G(z))$, tức BCE với nhãn “thật” cho fake.

Một bước GAN thật, nhưng trên số một chiều

Trước khi train GAN ảnh, nên kiểm đồ thị gradient trên một bài nhỏ đến mức đọc được từng tham số: dữ liệu thật là bốn số quanh 2, G và D mỗi mạng là một lớp tuyến tính. Điều cần nhìn là tham số nào đổi ở bước nào; giá trị loss của một bước chưa nói gì về chất lượng sinh.

import numpy as np
import torch
from torch import nn

torch.set_num_threads(1)
torch.manual_seed(58)
generator = nn.Linear(1, 1)
discriminator = nn.Linear(1, 1)
optimizer_g = torch.optim.SGD(generator.parameters(), lr=0.1)
optimizer_d = torch.optim.SGD(discriminator.parameters(), lr=0.1)
real = torch.tensor([[1.7], [1.9], [2.1], [2.3]])
noise = torch.tensor([[-1.], [0.], [1.], [2.]])
criterion = nn.BCEWithLogitsLoss()

g_before_d = generator.weight.detach().clone()
optimizer_d.zero_grad()
fake = generator(noise)
real_logits = discriminator(real)
fake_logits = discriminator(fake.detach())
loss_d = criterion(real_logits, torch.ones_like(real_logits))
loss_d += criterion(fake_logits, torch.zeros_like(fake_logits))
loss_d.backward()
optimizer_d.step()
print("G unchanged during D step:", torch.equal(g_before_d, generator.weight))
print("D loss:", round(loss_d.item(), 4))

Kết quảG unchanged during D step: True
D loss: 1.4571

D loss 1.4571 gần $2\ln2\approx1.386$, mức của một D chấm 0.5 cho mọi mẫu: D mới khởi tạo chưa phân biệt được thật và giả. detach() ở fake làm gradient của bước D không đi vào G. Bước tiếp theo giữ trọng số D nhưng vẫn cần đạo hàm của D theo đầu vào:

d_before_g = discriminator.weight.detach().clone()
g_before_g = generator.weight.detach().clone()
for parameter in discriminator.parameters():
    parameter.requires_grad_(False)
optimizer_g.zero_grad()
fake = generator(noise)
fake_logits = discriminator(fake)
loss_g = criterion(fake_logits, torch.ones_like(fake_logits))
loss_g.backward()
optimizer_g.step()
for parameter in discriminator.parameters():
    parameter.requires_grad_(True)
print("D unchanged during G step:", torch.equal(d_before_g, discriminator.weight))
print("G changed during G step:", not torch.equal(g_before_g, generator.weight))
print("G loss:", round(loss_g.item(), 4))
print("hand loss for D(fake)=0.2:", round(-np.log(0.2), 4))

Kết quảD unchanged during G step: True
G changed during G step: True
G loss: 1.0656
hand loss for D(fake)=0.2: 1.6094

G loss 1.0656 ứng với trung bình nhân của $D(G(z))$ khoảng $e^{-1.0656}\approx0.34$: D vẫn nghiêng về “giả”, nên G còn tín hiệu để học. Dòng cuối kiểm lại số tính tay ở mục trước.

Không đặt no_grad() quanh D trong bước G: như vậy sẽ cắt đường gradient cần về G. Đóng băng parameter D vẫn giữ phép tính theo fake. Ba dòng kiểm tham số phát hiện loại bug này trước khi chạy một GAN lớn.

GAN có thể dao động, discriminator quá mạnh hoặc generator chỉ sinh vài kiểu (mode collapse). Loss G/D không có quy tắc “càng thấp đồng thời càng tốt” giống supervised loss cố định. Cần xem chất lượng, đa dạng và metric phù hợp trên nhiều seed.

Diffusion: học sửa ảnh ở nhiều mức nhiễu

Thay vì đối đầu hai mạng, diffusion tạo một nhiệm vụ supervised từ dữ liệu: thêm noise đã biết vào mẫu sạch, rồi học đoán noise hoặc đại lượng tương đương. Lúc sinh, bắt đầu từ noise và làm nhiều bước khử nhiễu theo scheduler đã định nghĩa.

Đặt $\beta_t$ là mức nhiễu mỗi bước, $\alpha_t=1-\beta_t$, $\bar\alpha_t=\prod_{s=1}^t\alpha_s$. Có thể lấy trực tiếp mẫu nhiễu ở thời điểm $t$:

$x_0$ là dữ liệu sạch; $t$ là mức noise, không phải epoch. Hệ số căn bậc hai phân phối phần phương sai cho tín hiệu và noise. Khi $\bar\alpha_t$ gần 1, giữ nhiều tín hiệu; gần 0, mẫu gần noise.

Noise predictor cần nhận cả $x_t$ và $t$ vì cùng một giá trị nhiễu ở mức khác có ý nghĩa khác; $t$ thường được đưa vào mạng bằng embedding sin/cos như mã hóa vị trí ở bài 3.9. Khi $\bar\alpha_t \to 0$, $x_t$ tiến về $\mathcal N(0, I)$ bất kể $x_0$; với $\beta = 0.02$ không đổi, sau 50 bước $\bar\alpha = 0.98^{50} \approx 0.36$.

Ví dụ: với $x_0=2$, $\bar\alpha_t=0.64$, $\epsilon=0.5$, ta có $x_t=0.8\cdot2+0.6\cdot0.5=1.9$. Target khi train noise predictor là 0.5, không phải 2. Khi có dự đoán noise $\hat\epsilon$, có thể ước lượng sạch bằng $(x_t-\sqrt{1-\bar\alpha_t}\hat\epsilon)/\sqrt{\bar\alpha_t}$; đây chưa phải toàn bộ sampler.

clean = 2.0
alpha_bar_example = 0.64
epsilon_example = 0.5
noisy = np.sqrt(alpha_bar_example) * clean
noisy += np.sqrt(1 - alpha_bar_example) * epsilon_example
estimated_clean = (noisy - np.sqrt(1 - alpha_bar_example) * epsilon_example)
estimated_clean /= np.sqrt(alpha_bar_example)
print("noisy:", round(noisy, 4))
print("reconstruction with known noise:", round(estimated_clean, 4))

Kết quảnoisy: 1.9
reconstruction with known noise: 2.0

Ta dùng noise thật ở dòng khôi phục, nên kết quả 2 là kiểm đại số, không phải model đã học khử nhiễu. Ở lúc sinh, ảnh sạch chưa biết và noise phải được model ước lượng.

Train noise predictor trên CPU với dữ liệu một chiều

Thí nghiệm sau train thật một noise predictor trên dữ liệu một chiều: mỗi mẫu sạch nằm quanh $-2$ hoặc $+2$, mức nhiễu lấy ngẫu nhiên trong 40 mức (chỉ số 0-39), MLP nhận $x_t$ cùng chỉ số mức chia cho 39 và đoán $\epsilon$. Cơ chế giống hệt khi train trên ảnh; chỉ khác là $x_0$ là một số thay cho tensor ảnh, và mạng là MLP thay cho U-Net.

Loss là MSE giữa $\epsilon$ và dự đoán. DDPM xuất phát từ mục tiêu likelihood; sau khi rút gọn, mỗi số hạng thành bình phương sai số giữa $\epsilon$ và $\epsilon_\theta(x_t,t)$ nhân một trọng số theo $t$, và bài báo bỏ trọng số ấy để được MSE đơn giản trên noise.

torch.manual_seed(580)
steps = 40
betas = torch.linspace(0.01, 0.15, steps)
alpha_bars = torch.cumprod(1 - betas, dim=0)
denoiser = nn.Sequential(nn.Linear(2, 24), nn.Tanh(), nn.Linear(24, 1))
optimizer = torch.optim.Adam(denoiser.parameters(), lr=0.01)

def make_noisy_batch(count):
    clean = (torch.randint(0, 2, (count, 1)) * 4 - 2).float()
    clean += 0.2 * torch.randn(count, 1)
    times = torch.randint(0, steps, (count,))
    eps = torch.randn(count, 1)
    kept = alpha_bars[times, None]
    noisy = kept.sqrt() * clean + (1 - kept).sqrt() * eps
    inputs = torch.cat([noisy, times[:, None] / (steps - 1)], dim=1)
    return inputs, eps

valid_inputs, valid_noise = make_noisy_batch(128)
with torch.no_grad():
    before = nn.functional.mse_loss(denoiser(valid_inputs), valid_noise).item()
for update in range(250):
    inputs, targets = make_noisy_batch(64)
    optimizer.zero_grad()
    loss = nn.functional.mse_loss(denoiser(inputs), targets)
    loss.backward()
    optimizer.step()
with torch.no_grad():
    after = nn.functional.mse_loss(denoiser(valid_inputs), valid_noise).item()
    zero_baseline = valid_noise.square().mean().item()
print("noise MSE before/after:", round(before, 4), round(after, 4))
print("zero predictor MSE:", round(zero_baseline, 4))
print("last alpha_bar:", round(alpha_bars[-1].item(), 4))

Kết quảnoise MSE before/after: 0.9092 0.3434
zero predictor MSE: 0.9815
last alpha_bar: 0.0342

MSE giảm từ 0.9092 xuống 0.3434, thấp hơn nhiều so với baseline luôn đoán 0 (0.9815): model đã học. Ở mức cuối, $\bar\alpha=0.0342$: tín hiệu chỉ còn hệ số $\sqrt{0.0342}\approx0.18$, nên mẫu ở mức này đã gần noise chuẩn.

MSE không thể về 0, vì cùng một $x_t$ có thể đến từ nhiều cặp $(x_0,\epsilon)$. Ở mức nhiễu thấp, noise nhỏ hơn độ lệch tự nhiên của $x_0$ quanh $\pm2$ nên không tách được; ở mức giữa, không biết $x_t$ đến từ cụm $-2$ hay $+2$. Con số này đo việc đoán noise; chất lượng mẫu sinh còn phụ thuộc sampler và cần đánh giá riêng.

Từ noise predictor đến model sinh có điều kiện

Sampler bắt đầu từ $x_T\sim\mathcal N(0,I)$ và dùng dự đoán noise ở từng mức để đi về $x_{T-1},\dots,x_0$; tùy scheduler, mỗi bước có thể thêm một lượng noise mới.

Lặp công thức khôi phục $x_0$ ở trên không phải là DDPM: hệ số và phương sai của bước reverse phải khớp lịch nhiễu đã dùng khi train (DDPM §3.2 viết đủ các hệ số này). Mỗi bước là một lần gọi denoiser, và bước sau cần kết quả bước trước: 50 bước là 50 lần forward nối tiếp cho một ảnh, trong khi generator GAN chỉ cần một.

Diffusion có điều kiện chữ đưa embedding của prompt vào denoiser, thường qua cross-attention.

Latent diffusion (họ Stable Diffusion) chạy cả quá trình trong latent nén của ảnh: encoder của một autoencoder nén ảnh, khử nhiễu diễn ra trên latent, decoder biến latent cuối thành pixel. Stable Diffusion 1.x dùng latent $64\times64\times4$ cho ảnh $512\times512\times3$, nên mỗi lần gọi denoiser rẻ hơn nhiều.

Classifier-free guidance (CFG) điều khiển mức bám prompt. Mỗi bước gọi denoiser hai lần, có prompt $c$ và không prompt $\varnothing$, rồi đẩy dự đoán ra xa bản không điều kiện:

$w=1$ là dự đoán có điều kiện thường. $w$ lớn hơn làm ảnh bám prompt hơn, đổi lại đa dạng giảm và màu, chi tiết có thể quá gắt; mỗi bước cũng tốn gấp đôi. Model học cả hai chế độ nhờ khi train bỏ prompt ở một phần mẫu. Ví dụ $\epsilon$ không prompt 0.1 và có prompt 0.4 với $w = 5$ cho $0.1 + 5 \cdot 0.3 = 1.6$. Số bước, sampler, kích thước, $w$, seed và prompt đều cần ghi khi so kết quả.

Script của chương đưa quy trình này vào một lệnh: local_assets.py --task diffusion kiểm gói diffusers và transformers, thư mục pipeline có model_index.json, prompt và đường dẫn ảnh ra; không có --run thì chỉ kiểm điều kiện rồi dừng.

Với --run, script nạp pipeline (tokenizer, text encoder, U-Net, VAE, scheduler, weights safetensors) ở chế độ offline trên CPU float32 và sinh một ảnh $256\times256$ bằng 10 bước với seed 5, không ghi đè ảnh cũ. Một ảnh như vậy xác nhận đường nạp và inference chạy được; so chất lượng hai cấu hình cần cả bộ prompt và seed, như mục sau.

Đánh giá model sinh

FID so phân phối feature (lấy từ một mạng phân loại pretrained) của ảnh sinh với ảnh thật. FID càng thấp càng tốt, nhưng cần hàng nghìn mẫu và cách tính nhất quán giữa các lần so: 500 ảnh mỗi model là quá ít, và FID tính trên số mẫu khác nhau không so được với nhau.

Sinh dữ liệu để train classifier cần kiểm nhãn, trùng lặp và mức khớp miền ảnh thật. Model sinh học từ dữ liệu nguồn, có thể lặp nội dung hoặc thiên lệch; ảnh sinh không tự là bằng chứng về thế giới thật.