Latent Diffusion Model (LDM)

Model khuếch tán tiềm ẩn (LDM) là model khuếch tán tạo dữ liệu bằng cách khử nhiễu một biểu diễn nén từ autoencoder pretrained thay vì xử lý pixel thô. Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser và Björn Ommer giới thiệu phương pháp này trong bài Tổng hợp ảnh độ phân giải cao bằng model khuếch tán tiềm ẩn, trình bày tại CVPR 2022. Autoencoder loại bỏ chi tiết pixel nhỏ vốn tốn kém khi mô hình hóa nhưng ít ảnh hưởng đến cảm nhận về ảnh, nhờ đó mạng khuếch tán tập trung năng lực vào nội dung và bố cục. Vì vậy, chi phí huấn luyện và lấy mẫu thấp hơn nhiều so với khuếch tán trong không gian pixel, cho phép Stable Diffusion, LDM nổi tiếng nhất, tạo ảnh trên một GPU phổ thông. Hầu hết các model tạo ảnh, video và âm thanh hiện nay trong AI tạo sinh đều tuân theo cùng thiết kế hai giai đoạn.

Cách diffusion trong không gian latent hoạt động#

Giai đoạn đầu tiên nén ảnh. Encoder ánh xạ ảnh kích thước H×W×3 thành latent kích thước H/f×W/f×c, trong đó f là hệ số giảm mẫu và c là số kênh latent; decoder ánh xạ latent trở lại thành pixel. Rombach và cộng sự huấn luyện autoencoder này bằng perceptual loss và mục tiêu đối kháng dựa trên patch như trong discriminator của GAN để ảnh tái tạo vẫn sắc nét, đồng thời chính quy hóa latent bằng một penalty KL nhỏ hướng đến phân phối chuẩn tắc chuẩn hoặc bằng lượng tử hóa vector. Hệ số giảm mẫu từ 4 đến 16 cân bằng tốc độ và chất lượng: huấn luyện trong không gian pixel (f = 1) chậm, còn f = 32 nén quá nhiều khiến chất lượng mẫu giảm. Autoencoder của Stable Diffusion dùng f = 8 và 4 kênh, vì vậy đoạn mã này mã hóa ảnh, giải mã lại rồi đo lường những gì bị mất qua quá trình khứ hồi:

import urllib.request

import numpy as np
import torch
from diffusers import AutoencoderKL
from PIL import Image

# Stable Diffusion's autoencoder, the first stage of a latent diffusion model
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")

# 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((480, 640))
x = torch.tensor(np.array(image), dtype=torch.float32).permute(2, 0, 1)[None] / 127.5 - 1

# Encode to the latent that diffusion runs on, then decode back to pixels
with torch.no_grad():
    latent = vae.encode(x).latent_dist.mode()
    recon = vae.decode(latent).sample.clamp(-1, 1)

psnr = 10 * torch.log10(4 / torch.mean((recon - x) ** 2))
print(f"image {tuple(x.shape)} -> latent {tuple(latent.shape)}: {x.numel() // latent.numel()}x fewer values")
print(f"reconstruction PSNR: {psnr:.1f} dB")
Image.fromarray(((recon[0] + 1) * 127.5).byte().permute(1, 2, 0).numpy()).save("bus_reconstructed.jpg")

Đoạn mã in ra image (1, 3, 640, 480) -> latent (1, 4, 80, 60): 48x fewer values và reconstruction PSNR: 23.0 dB. Tệp bus_reconstructed.jpg đã lưu vẫn giữ được xe buýt, người đi đường và con phố, nhưng chữ nhỏ trên thân xe buýt bị nhòe khi tái tạo; đó là chi tiết mịn mà latent 4 kênh không thể lưu giữ.

Giai đoạn thứ hai là một model khuếch tán thông thường được huấn luyện trên các latent này. Quá trình huấn luyện mã hóa một ảnh, thêm nhiễu Gaussian tại một bước ngẫu nhiên, rồi huấn luyện mạng khử nhiễu—U-Net trong công trình gốc và Transformer khuếch tán trong hầu hết model mới hơn—để dự đoán nhiễu đó; còn trong các model rectified flow, mạng dự đoán vận tốc đưa nhiễu về phía latent sạch. Các điều kiện được đưa vào thông qua các lớp attention chéo, cho phép latent chú ý đến các embedding của prompt văn bản, nhãn lớp hoặc bản đồ ngữ nghĩa. Quá trình tạo sinh bắt đầu từ nhiễu latent ngẫu nhiên, khử nhiễu latent này từng bước theo điều kiện, rồi chỉ chạy decoder một lần ở cuối, nhờ vậy mọi bước tốn kém đều xử lý trên latent nhỏ.

Từ Stable Diffusion đến các model hiện nay#

• Stable Diffusion: Ra mắt tháng 8 năm 2022, model này mở rộng công thức LDM cho tác vụ chuyển văn bản thành ảnh 512×512 với latent 64×64×4, U-Net 860 triệu tham số và encoder văn bản CLIP; trọng số có thể được tải xuống tự do.
• Stable Diffusion 3: Chuyển sang autoencoder 16 kênh, nâng PSNR tái tạo từ 25.12 dB với 4 kênh lên 28.62 dB, đồng thời thay U-Net bằng Transformer diffusion đa phương thức được huấn luyện bằng dòng hiệu chỉnh. Các model FLUX.1 của Black Forest Labs cũng khử nhiễu latent 16 kênh bằng Transformer dòng hiệu chỉnh.
• Video và âm thanh: Model tạo video nén clip theo không gian và thời gian; OpenAI mô tả Sora là công nghệ nén video vào không gian latent có số chiều thấp hơn rồi chia thành các patch không-thời gian. AudioLDM và Stable Audio của Stability AI áp dụng cùng phương pháp cho latent phổ âm thanh hoặc dạng sóng.
• Model nhất quán latent: Chưng cất LDM để tạo sinh trong 2 đến 4 bước thay vì 20 đến 50 bước, cho phép tạo sinh gần thời gian thực.

LDM và các thuật ngữ liên quan#

• Model diffusion trong không gian pixel: Thực hiện quá trình thêm và khử nhiễu trên mọi pixel. LDM cũng là model diffusion, chỉ khác ở không gian thực hiện quá trình.
• Variational autoencoder: VAE tạo sinh bằng cách giải mã latent được lấy mẫu trong một lượt, và các mẫu thường bị mờ. LDM chỉ dùng autoencoder kiểu VAE để nén, còn việc tạo sinh do model diffusion đảm nhiệm.
• Stable Diffusion: Một họ LDM từ Stability AI, CompVis và Runway; do đó, mọi model Stable Diffusion đều là LDM, trong khi LDM còn bao gồm FLUX, các model tạo video và model âm thanh.

Ứng dụng#

LDM hỗ trợ tạo ảnh từ văn bản, inpainting ảnh—lĩnh vực mà bài báo gốc lập kỷ lục mới vào thời điểm đó—, siêu phân giải và tạo ảnh từ bố cục bằng bản đồ ngữ nghĩa. Trong thị giác máy tính, LDM cũng tạo dữ liệu tổng hợp: lỗi hiếm gặp, thời tiết bất thường hoặc góc nhìn khó chụp. Ảnh được tạo ra không có nhãn; hàm auto_annotate của Ultralytics gán nhãn cả thư mục ảnh bằng cách phát hiện đối tượng với model YOLO26 và phác họa từng box được phát hiện bằng model phân đoạn Meta SAM, rồi ghi nhãn phân đoạn theo định dạng YOLO để huấn luyện.