Diffusion model hoạt động thế nào? Tự cài đặt DDPM bằng NumPy

Stable Diffusion, các mô hình tạo ảnh từ văn bản, và nhiều công cụ sinh ảnh lỗi giả lập cho Machine Vision đều dựa trên một ý tưởng: học cách khử nhiễu. Bài Generative AI là gì? đã giới thiệu khái niệm. Bài này đi sâu hơn: ta tự cài đặt quá trình thêm nhiễu và khử nhiễu của DDPM (Denoising Diffusion Probabilistic Models, Ho và cộng sự, 2020) bằng NumPy, để thấy tận mắt từng con số.

Ý tưởng trong ba câu

1. Quá trình thuận (không cần học): lấy ảnh thật, thêm từng chút nhiễu Gauss qua T bước, cho tới khi chỉ còn nhiễu thuần.
2. Huấn luyện: cho mạng nơ-ron xem ảnh đã bị nhiễu ở bước t, bắt nó đoán xem nhiễu đã thêm vào là gì.
3. Sinh ảnh: bắt đầu từ nhiễu thuần, dùng mạng để trừ dần nhiễu qua T bước, thu được ảnh mới.

1. Lịch nhiễu (noise schedule)

Ở mỗi bước t, ta thêm nhiễu với "liều lượng" β_t. Bài báo DDPM dùng T = 1000 bước và β tăng tuyến tính từ 10⁻⁴ đến 0,02. Đặt α_t = 1 − β_t và ᾱ_t là tích của mọi α từ bước 1 đến t. Đại lượng ᾱ_t cho biết còn lại bao nhiêu "tín hiệu" sau t bước.

import numpy as np

T = 1000
beta = np.linspace(1e-4, 0.02, T)          # lịch nhiễu tuyến tính của DDPM
alpha = 1.0 - beta
alpha_bar = np.cumprod(alpha)              # ᾱ_t = α_1 · α_2 · ... · α_t

for t in [1, 10, 100, 250, 500, 1000]:
    ab = alpha_bar[t - 1]
    print(f"t = {t:4d}: giữ lại tín hiệu sqrt(ᾱ) = {np.sqrt(ab):.3f}, nhiễu sqrt(1-ᾱ) = {np.sqrt(1 - ab):.3f}")

t =    1: giữ lại tín hiệu sqrt(ᾱ) = 1.000, nhiễu sqrt(1-ᾱ) = 0.010
t =   10: giữ lại tín hiệu sqrt(ᾱ) = 0.999, nhiễu sqrt(1-ᾱ) = 0.044
t =  100: giữ lại tín hiệu sqrt(ᾱ) = 0.947, nhiễu sqrt(1-ᾱ) = 0.321
t =  250: giữ lại tín hiệu sqrt(ᾱ) = 0.724, nhiễu sqrt(1-ᾱ) = 0.690
t =  500: giữ lại tín hiệu sqrt(ᾱ) = 0.280, nhiễu sqrt(1-ᾱ) = 0.960
t = 1000: giữ lại tín hiệu sqrt(ᾱ) = 0.006, nhiễu sqrt(1-ᾱ) = 1.000

Ở bước 1000, tín hiệu gốc chỉ còn hệ số 0,006: ảnh gần như đã biến thành nhiễu Gauss chuẩn. Đó chính là điểm xuất phát khi sinh ảnh.

2. Nhảy thẳng tới bước t

Một tính chất rất tiện của nhiễu Gauss: không cần thêm nhiễu 250 lần để có ảnh ở bước 250. Có công thức đóng:

x_t = √ᾱ_t · x₀ + √(1 − ᾱ_t) · ε, với ε là nhiễu Gauss chuẩn.

Thử với một "ảnh" 64×64 có vết xước sáng chạy chéo trên nền tối (giá trị chuẩn hóa về khoảng [−1, 1] như trong các mô hình diffusion):

rng = np.random.default_rng(0)

# "Ảnh" 64x64: nền tối, một vết xước sáng chéo, giá trị chuẩn hóa về [-1, 1]
x0 = np.full((64, 64), -0.8)
for i in range(10, 54):
    x0[i, i - 2:i + 2] = 0.9


def them_nhieu(x0, t, rng):
    """Nhảy thẳng tới bước t: x_t = sqrt(ᾱ_t)·x0 + sqrt(1-ᾱ_t)·ε"""
    eps = rng.standard_normal(x0.shape)
    ab = alpha_bar[t - 1]
    return np.sqrt(ab) * x0 + np.sqrt(1 - ab) * eps, eps


for t in [10, 100, 250, 500, 1000]:
    xt, eps = them_nhieu(x0, t, rng)
    tuong_quan = np.corrcoef(x0.ravel(), xt.ravel())[0, 1]
    print(f"t = {t:4d}: tương quan với ảnh gốc = {tuong_quan:.3f}")

t =   10: tương quan với ảnh gốc = 0.992
t =  100: tương quan với ảnh gốc = 0.710
t =  250: tương quan với ảnh gốc = 0.329
t =  500: tương quan với ảnh gốc = 0.102
t = 1000: tương quan với ảnh gốc = 0.024

Hệ số tương quan đo "ảnh nhiễu còn giống ảnh gốc bao nhiêu": từ 0,99 ở bước 10 giảm về gần 0 ở bước 1000. Công thức này giúp huấn luyện rất nhanh: mỗi lần lấy một ảnh, chọn ngẫu nhiên một t, tạo x_t ngay lập tức.

3. Mạng học gì? Học đoán nhiễu

Trong DDPM, mạng (thường là một U-Net, xem U-Net) nhận x_t và t, xuất ra dự đoán ε̂ của nhiễu ε. Hàm mất mát chỉ là sai số bình phương trung bình giữa ε̂ và ε. Vì sao đoán nhiễu lại đủ? Vì biết nhiễu thì đảo ngược được công thức ở trên:

t = 250
xt, eps = them_nhieu(x0, t, rng)
ab = alpha_bar[t - 1]

# Giả sử mạng dự đoán nhiễu hoàn hảo: khôi phục x0 chính xác
x0_uoc = (xt - np.sqrt(1 - ab) * eps) / np.sqrt(ab)
print("sai số khi biết đúng nhiễu:", np.abs(x0_uoc - x0).max().round(6))

# Mạng dự đoán nhiễu sai 10%: ảnh khôi phục có sai số
eps_sai = eps + 0.1 * rng.standard_normal(eps.shape)
x0_uoc = (xt - np.sqrt(1 - ab) * eps_sai) / np.sqrt(ab)
print("sai số trung bình khi nhiễu dự đoán lệch:", np.abs(x0_uoc - x0).mean().round(3))

sai số khi biết đúng nhiễu: 0.0
sai số trung bình khi nhiễu dự đoán lệch: 0.075

Trên thực tế, mạng không bao giờ đoán đúng hoàn toàn, nhất là ở bước t lớn khi ảnh gần như chỉ còn nhiễu. Vì vậy thay vì đảo ngược một phát, DDPM đi lùi từng bước nhỏ, mỗi bước sửa sai một ít.

4. Sinh ảnh: đi lùi 1000 bước

Công thức một bước lùi của DDPM:

x_t−1 = (x_t − β_t / √(1 − ᾱ_t) · ε̂) / √α_t + √β_t · z, với z là nhiễu mới (bỏ ở bước cuối).

Để chạy được mà không cần huấn luyện, ta dùng một "mạng lý tưởng" cho trường hợp đơn giản nhất: tập dữ liệu chỉ có đúng một ảnh x0. Khi đó nhiễu tốt nhất để đoán chính là phần x_t lệch khỏi √ᾱ_t·x₀:

def mang_ly_tuong(xt, t):
    """'Mạng' lý tưởng khi dữ liệu chỉ có đúng một ảnh x0: nhiễu = phần lệch khỏi sqrt(ᾱ)·x0."""
    ab = alpha_bar[t - 1]
    return (xt - np.sqrt(ab) * x0) / np.sqrt(1 - ab)


x = rng.standard_normal(x0.shape)                  # bắt đầu từ nhiễu thuần
for t in range(T, 0, -1):                          # t = 1000, 999, ..., 1
    eps_du_doan = mang_ly_tuong(x, t)
    he_so = beta[t - 1] / np.sqrt(1 - alpha_bar[t - 1])
    x = (x - he_so * eps_du_doan) / np.sqrt(alpha[t - 1])
    if t > 1:                                      # bước cuối không thêm nhiễu
        x = x + np.sqrt(beta[t - 1]) * rng.standard_normal(x.shape)
    if t in (1000, 500, 100, 1):
        print(f"t = {t:4d}: tương quan với ảnh gốc = {np.corrcoef(x0.ravel(), x.ravel())[0, 1]:.3f}")

t = 1000: tương quan với ảnh gốc = 0.008
t =  500: tương quan với ảnh gốc = 0.078
t =  100: tương quan với ảnh gốc = 0.706
t =    1: tương quan với ảnh gốc = 1.000

Từ nhiễu thuần, vòng lặp dần "hiện" ra vết xước. Để ý: ở nửa đầu (t từ 1000 về 500) ảnh vẫn gần như nhiễu; hình dạng chỉ rõ dần ở các bước cuối. Với mô hình thật, tập dữ liệu có hàng triệu ảnh, nên mạng không "nhớ" một ảnh mà học được phân bố chung; mỗi lần bắt đầu từ một nhiễu khác nhau sẽ cho ra một ảnh mới khác nhau.

5. Liên hệ với Machine Vision

• Sinh ảnh lỗi khi thiếu dữ liệu NG: kỹ thuật inpainting dùng diffusion để "vẽ" vết lỗi vào vùng chọn trên ảnh OK. Ảnh sinh ra phải được kỹ sư kiểm tra lại và không nên thay thế hoàn toàn ảnh lỗi thật; xem dữ liệu tổng hợp khi thiếu ảnh lỗi.
• Anomaly detection: một số phương pháp dùng diffusion để "sửa" ảnh có lỗi thành ảnh trông bình thường, rồi so sánh với ảnh gốc để tìm vùng khác biệt.
• Khử nhiễu ảnh: bài toán gốc của diffusion, hữu ích với ảnh chụp thiếu sáng.

Điểm cần nhớ khi đánh giá: mô hình sinh ảnh chỉ tái tạo những gì nó đã thấy. Một loại lỗi chưa từng xuất hiện trong dữ liệu huấn luyện thì mô hình cũng không tự "nghĩ ra" đúng hình dạng vật lý của nó.

Nguồn tham khảo: Ho, Jain, Abbeel, Denoising Diffusion Probabilistic Models (arXiv:2006.11239), Nichol và Dhariwal, Improved DDPM (arXiv:2102.09672), Hugging Face Diffusers: DDPMScheduler.