Diffusion Models

Model khuếch tán là các model AI tạo sinh tạo dữ liệu bằng cách học đảo ngược một quá trình thêm nhiễu dần dần. Trong quá trình huấn luyện, ảnh hoặc dữ liệu khác được làm nhiễu bằng nhiễu Gaussian qua nhiều bước nhỏ, và một mạng nơ-ron học cách dự đoán nhiễu ở từng bước. Khi tạo dữ liệu, model bắt đầu từ nhiễu ngẫu nhiên thuần túy rồi khử nhiễu từng bước cho đến khi xuất hiện một mẫu mới. Năm 2015, Jascha Sohl-Dickstein và các cộng sự đề xuất phương pháp này trong Deep Unsupervised Learning using Nonequilibrium Thermodynamics; năm 2020, Jonathan Ho, Ajay Jain và Pieter Abbeel giúp phương pháp đạt hiệu quả cạnh tranh với Denoising Diffusion Probabilistic Models (DDPM), thiết lập FID tiên tiến nhất ở mức 3.17 trên CIFAR-10. Năm 2021, các nhà nghiên cứu của OpenAI cho thấy model khuếch tán vượt GAN trong tổng hợp ảnh ImageNet; model khuếch tán và các phương pháp liên quan như khớp dòng hiện là nền tảng cho hệ thống tạo ảnh từ văn bản như Stable Diffusion, công cụ tạo video, trình chỉnh sửa ảnh và model khoa học.

Cách hoạt động của model khuếch tán#

Quá trình xuôi được cố định thay vì học. Quá trình này thêm một lượng nhỏ nhiễu Gaussian ở mỗi bước trong T bước; DDPM dùng T = 1,000, với phương sai nhiễu tăng tuyến tính từ 0.0001 đến 0.02 mỗi bước. Vì tổng của các nhiễu Gaussian vẫn là Gaussian, ảnh huấn luyện có thể nhảy thẳng đến bất kỳ bước t nào bằng công thức dạng đóng, xt = sqrt(alpha_bar_t) * x0 + sqrt(1 - alpha_bar_t) * noise, trong đó alpha_bar_t là tích của 1 - beta trong t bước đầu tiên. Quá trình huấn luyện liên tục chọn một ảnh, một bước ngẫu nhiên và nhiễu mới, tạo xt, rồi dạy mạng dự đoán nhiễu đó bằng sai số bình phương trung bình, hàm mục tiêu đơn giản hóa của DDPM. Đoạn mã này áp dụng quá trình xuôi lên một bức ảnh bằng lịch nhiễu của DDPM:

import urllib.request

import numpy as np
import torch
from PIL import Image

# DDPM's linear schedule: 1,000 steps whose noise variance rises from 0.0001 to 0.02
betas = torch.linspace(1e-4, 0.02, 1000)
alpha_bar = torch.cumprod(1 - betas, dim=0)

# Load a photo and scale its pixels to [-1, 1]
image = Image.open(urllib.request.urlopen("https://ultralytics.com/images/bus.jpg")).convert("RGB").resize((256, 256))
x0 = torch.tensor(np.array(image), dtype=torch.float32).permute(2, 0, 1) / 127.5 - 1

# Jump straight to step t in closed form: x_t = sqrt(alpha_bar) * x0 + sqrt(1 - alpha_bar) * noise
frames = []
for t in (0, 99, 299, 599, 999):
    xt = alpha_bar[t].sqrt() * x0 + (1 - alpha_bar[t]).sqrt() * torch.randn_like(x0)
    frames.append(((xt.clamp(-1, 1) + 1) * 127.5).byte().permute(1, 2, 0).numpy())
    print(f"step {t + 1}: image weight {alpha_bar[t].sqrt():.3f}, noise weight {(1 - alpha_bar[t]).sqrt():.3f}")
Image.fromarray(np.hstack(frames)).save("forward_diffusion.jpg")

Trọng số ảnh giảm từ 1.000 ở bước 1 xuống 0.947 ở bước 100, 0.630 ở bước 300, 0.161 ở bước 600 và 0.006 ở bước 1,000, trong khi trọng số nhiễu tăng lên 1.000. forward_diffusion.jpg đã lưu cho thấy chiếc xe buýt ban đầu sắc nét, sau đó nhiễu hạt, rồi mờ dần và cuối cùng biến mất.

Quá trình tạo chạy ngược lại. Quá trình bắt đầu từ nhiễu Gaussian thuần túy; ở mỗi bước, mạng dự đoán nhiễu trong mẫu hiện tại, một bộ giải loại bỏ một phần nhiễu, rồi kết quả sạch hơn một chút được đưa vào bước tiếp theo. Prompt văn bản hoặc nhãn lớp được đưa vào mạng làm điều kiện, thường thông qua chú ý chéo. Hướng dẫn không cần bộ phân loại, do Jonathan Ho và Tim Salimans giới thiệu, huấn luyện một mạng cả khi có lẫn khi không có điều kiện, rồi tại thời điểm lấy mẫu ngoại suy từ dự đoán vô điều kiện theo hướng dự đoán có điều kiện; hệ số hướng dẫn cao hơn, chẳng hạn giá trị mặc định 7.5 trong các pipeline Stable Diffusion 1.x, giúp kết quả bám sát prompt hơn nhưng làm giảm tính đa dạng.

Tăng tốc quá trình lấy mẫu#

Quá trình lấy mẫu DDPM cần 1,000 lượt đánh giá mạng cho mỗi ảnh, vì vậy phần lớn tiến bộ của lĩnh vực này tập trung vào việc giảm con số đó:

• Bộ lấy mẫu tốt hơn: DDIM tái sử dụng DDPM đã huấn luyện với bộ lấy mẫu xác định và tạo mẫu chất lượng cao nhanh hơn 10 đến 50 lần theo thời gian thực. Các bộ giải ODE ra đời sau đó đưa quá trình tạo thông thường xuống còn khoảng 20 đến 50 bước.
• Khuếch tán tiềm ẩn: Model khuếch tán tiềm ẩn khử nhiễu một biểu diễn nén từ autoencoder thay vì các pixel. Autoencoder của Stable Diffusion chuyển ảnh 512×512×3 thành một biểu diễn tiềm ẩn 64×64×4, ít hơn 48 lần số giá trị.
• Chưng cất: Model nhất quán học cách ánh xạ trực tiếp mọi mẫu nhiễu thành mẫu sạch, cho phép tạo nội dung chỉ trong một bước; model nhất quán tiềm ẩn áp dụng cùng ý tưởng cho Stable Diffusion với 2 đến 4 bước.
• Đường đi thẳng hơn: Khớp dòng và dòng được hiệu chỉnh huấn luyện dọc theo các đường thẳng từ nhiễu đến dữ liệu mà bộ giải có thể theo dõi trong ít bước hơn.

Bản thân mạng khử nhiễu đã chuyển từ U-Net tích chập của DDPM và Stable Diffusion 1.x sang Transformer khuếch tán (DiT) được hầu hết trình tạo ảnh và video hiện nay sử dụng.

Model khuếch tán và các model tạo sinh khác#

• Mạng đối sinh tạo (GAN): GAN tạo một mẫu chỉ trong một lượt truyền xuôi, nhưng quá trình huấn luyện đối kháng giữa generator và discriminator dễ mất ổn định và sụp đổ mode, khiến generator chỉ bao phủ một phần dữ liệu. Quá trình huấn luyện khuếch tán là một bài toán hồi quy ổn định và bao phủ phân phối dữ liệu đầy đủ hơn, đổi lại cần nhiều bước lấy mẫu.
• Autoencoder biến phân (VAE): VAE giải mã một vector tiềm ẩn thành mẫu chỉ trong một lượt, nhưng các mẫu thường bị mờ hơn. Model khuếch tán tiềm ẩn chỉ dùng autoencoder kiểu VAE để nén ảnh, còn việc tạo ảnh do model khuếch tán đảm nhiệm.
• Model tự hồi quy: Model ngôn ngữ và model tự hồi quy thị giác tạo từng token hoặc từng tỷ lệ một, dựa trên những gì đã xuất hiện trước đó. Model khuếch tán tinh chỉnh song song mọi phần của mẫu ở mỗi bước, còn model ngôn ngữ khuếch tán áp dụng cùng phương pháp khử nhiễu cho văn bản.

Ứng dụng thực tế#

• Tạo và chỉnh sửa ảnh: Model chuyển văn bản thành ảnh tạo ảnh từ prompt; tô vùng ảnh tạo lại một vùng được che để khớp với phần xung quanh; ControlNet ràng buộc bố cục bằng bản đồ cạnh, độ sâu hoặc tư thế; còn các model khuếch tán tăng độ phân giải thực hiện siêu phân giải.
• Tạo video: Các model tạo video khử nhiễu các biểu diễn tiềm ẩn nén bao phủ nhiều khung hình cùng lúc, nhờ đó mỗi khung hình được tạo đồng thời với các khung hình lân cận.
• Dữ liệu huấn luyện tổng hợp: Ảnh được tạo ra bổ sung những phần còn thiếu trong dataset thị giác, chẳng hạn lỗi hiếm gặp, thời tiết bất thường hoặc góc nhìn ít phổ biến. Trong dataset trên Nền tảng Ultralytics, thao tác Tạo ảnh tương tự tạo từ 1 đến 16 ảnh mới mỗi batch, lấy cảm hứng từ một ảnh hiện có; những ảnh được giữ lại sẽ được thêm vào dataset mà chưa gán nhãn, sẵn sàng để chú giải.
• Robot học: Chính sách khuếch tán tạo chuỗi hành động của robot bằng cách khử nhiễu, với ảnh camera làm điều kiện.
• Khoa học: AlphaFold 3 dự đoán cấu trúc của các phức hợp protein, axit nucleic và phân tử nhỏ bằng một module khuếch tán khử nhiễu các tọa độ nguyên tử thô, thay thế module cấu trúc của AlphaFold 2.