Diffusion Transformer (DiT)

Diffusion Transformer (DiT) là một model diffusion có mạng khử nhiễu là transformer thay vì U-Net tích chập. DiT chia latent ảnh nhiễu thành các patch, xử lý chúng như một chuỗi token bằng self-attention giống cách Vision Transformer (ViT) xử lý ảnh, rồi dự đoán lượng nhiễu cần loại bỏ ở mỗi bước khử nhiễu. William Peebles và Saining Xie giới thiệu mô hình này trong bài báo Scalable Diffusion Models with Transformers (ICCV 2023); trong đó, DiT-XL/2 đạt Khoảng cách Inception Fréchet (FID) 2.27 trên ImageNet 256×256 có điều kiện theo lớp, vượt kết quả tốt nhất trước đó là 3.60 của latent diffusion. Bài báo cho thấy chất lượng ảnh cải thiện ổn định khi mức compute của DiT tăng, và từ đó, bộ khử nhiễu transformer đã thay thế U-Net trong hầu hết trình tạo ảnh và video quy mô lớn, bao gồm Stable Diffusion 3, FLUX.1 và Sora của OpenAI.

Cách Diffusion Transformer hoạt động#

DiT chạy trong latent space của variational autoencoder (VAE) pretrained được Stable Diffusion sử dụng; encoder của VAE giảm kích thước 8 lần, nên ảnh 256×256×3 trở thành latent 32×32×4. Lớp patchify chia latent nhiễu thành các patch p×p và nhúng tuyến tính từng patch thành một token; kích thước patch 2, 4 và 8 lần lượt tạo ra 256, 64 và 16 token. Embedding vị trí sin-cos đánh dấu vị trí ban đầu của mỗi token, một chuỗi các block transformer xử lý chuỗi, rồi lớp tuyến tính cuối cùng chuyển mỗi token trở lại thành một patch gồm nhiễu và phương sai được dự đoán; các patch này được ghép lại thành hình dạng của latent. Timestep diffusion và nhãn lớp, hoặc prompt văn bản trong model text-to-image, điều kiện hóa mọi block. Quá trình tạo ảnh lặp lại việc dự đoán này từ nhiễu thuần cho đến bước cuối, rồi decoder VAE chuyển latent sạch thành pixel:

Bài báo so sánh bốn cách đưa điều kiện vào mỗi block: token bổ sung trong chuỗi, cross-attention, chuẩn hóa lớp thích ứng (adaLN) và adaLN-Zero. adaLN-Zero đạt FID thấp nhất và cũng tiết kiệm compute nhất. Một mạng nhỏ đọc tổng embedding timestep và embedding lớp, rồi hồi quy độ dịch và hệ số tỷ lệ áp dụng sau mỗi layer norm, cùng một gate điều chỉnh tỷ lệ cho mỗi nhánh residual. Các gate khởi tạo bằng 0, nên mỗi block bắt đầu training như hàm đồng nhất. Bản vá code PyTorch này chia latent thành patch và áp dụng nhánh attention của một block:

import torch
from torch import nn

torch.manual_seed(0)
# A 256x256 image becomes a 4x32x32 latent; 2x2 patches turn it into 256 tokens of width 384 (DiT-S/2)
latent = torch.randn(1, 4, 32, 32)
tokens = nn.Conv2d(4, 384, kernel_size=2, stride=2)(latent).flatten(2).transpose(1, 2)

# adaLN-Zero: the timestep and class embedding regresses a shift, scale, and gate for the attention branch
condition = torch.randn(1, 384)
modulation = nn.Linear(384, 3 * 384)
nn.init.zeros_(modulation.weight)
nn.init.zeros_(modulation.bias)
shift, scale, gate = modulation(condition).unsqueeze(1).chunk(3, dim=-1)

# Self-attention branch of a DiT block, modulated by the condition and scaled by the gate
attention = nn.MultiheadAttention(384, num_heads=6, batch_first=True)
h = nn.functional.layer_norm(tokens, (384,)) * (1 + scale) + shift
out = tokens + gate * attention(h, h, h)[0]
print(f"tokens {tuple(tokens.shape)}, block is identity at initialization: {torch.equal(out, tokens)}")

Kết quả in ra là tokens (1, 256, 384), block is identity at initialization: True. 256 token có chiều rộng 384 tương ứng với model nhỏ nhất, DiT-S/2; gate bằng 0 truyền nguyên đầu vào qua cho đến khi training điều chỉnh nó. Một block DiT đầy đủ hồi quy thêm độ dịch, hệ số tỷ lệ và gate thứ hai cho nhánh MLP.

Mở rộng theo số token và compute#

Bài báo training bốn kích thước model, từ DiT-S (12 layer, chiều rộng 384) đến DiT-XL (28 layer, chiều rộng 1.152), mỗi model có kích thước patch 8, 4 và 2. Trên cả 12 model, FID tương quan với compute của lượt forward tính bằng Gflops, dù compute tăng do mạng sâu hơn, rộng hơn hay có nhiều token hơn. Giảm một nửa kích thước patch làm số token và compute tăng gấp bốn, trong khi gần như không thêm tham số. DiT-XL/2 có 675 triệu tham số, chưa tính VAE 84 triệu tham số, và sử dụng 119 Gflops cho mỗi lượt forward ở độ phân giải 256×256. Ở độ phân giải 512×512, model xử lý 1.024 token từ latent 64×64×4 với 525 Gflops và đạt FID 3.04. Vì self-attention so sánh từng token với mọi token khác, độ phân giải cao hơn và video dài hơn làm chi phí tăng theo bình phương. Ultralytics YOLO26 tránh chi phí này bằng cách chỉ đặt attention ở stage tỷ lệ 1/32, với 400 vị trí cho ảnh 640×640, đồng thời dùng phép tích chập ở mọi vị trí khác để phát hiện theo thời gian thực.

DiT so với U-Net và ViT#

• Bộ khử nhiễu U-Net: U-Net trong DDPM và Stable Diffusion 1.x, 2.x là các encoder-decoder tích chập hoạt động ở nhiều độ phân giải, có skip connection giữa các cấp tương ứng và bổ sung attention ở độ phân giải thấp hơn. DiT duy trì một độ phân giải token, dùng attention toàn cục trong mọi block và có ít lựa chọn thiết kế đặc thù cho ảnh hơn, nên mở rộng theo cách của transformer trong các mô hình ngôn ngữ lớn.
• Vision Transformer: ViT mã hóa ảnh cho các tác vụ như phân loại. DiT tái sử dụng embedding patch và các block transformer của ViT, bổ sung điều kiện hóa timestep qua adaLN-Zero, rồi xuất dự đoán cho mọi patch.
• Latent diffusion và flow matching: Model latent diffusion là khuôn khổ tổng thể để khử nhiễu trong latent space của VAE; diffusion hoặc flow matching là mục tiêu training. DiT là mạng khử nhiễu bên trong khuôn khổ đó, và các trình tạo hiện nay kết hợp backbone DiT với mục tiêu flow matching.

DiT cho ảnh và video từ văn bản#

• DiT đa phương thức (MMDiT): Stable Diffusion 3 đưa token ảnh và token văn bản qua các bộ trọng số riêng cho từng phương thức, rồi kết hợp chúng trong một phép toán attention duy nhất để thông tin lưu chuyển hai chiều giữa prompt và ảnh. Thiết kế này cải thiện kiểu chữ và khả năng làm theo prompt; Stable Diffusion 3.5 Large mở rộng thiết kế lên khoảng 8 tỷ tham số.
• FLUX.1: FLUX.1 của Black Forest Labs là transformer rectified flow có 12 tỷ tham số để tạo ảnh từ văn bản.
• Video: Báo cáo kỹ thuật Sora của OpenAI mô tả một diffusion transformer nén video vào latent space rồi chia video thành các patch không-thời gian, nhờ đó các clip có thời lượng, độ phân giải và tỷ lệ khung hình khác nhau trở thành chuỗi token có độ dài khác nhau. Để tạo video, model Movie Gen lớn nhất của Meta là transformer 30 tỷ tham số, có ngữ cảnh 73.000 token video tương ứng với 16 giây video ở tốc độ 16 frame/giây.