Quá Trình Khuếch Tán Thuận Thêm Nhiễu Closed-Form Cho DDPM (Ho et al. 2020)

Mục tiêu kiến thức

• Nhờ tính chất của phân phối chuẩn Gauss (tổng của hai biến ngẫu nhiên Gauss độc lập vẫn là một biến ngẫu nhiên Gauss), ta không cần mô phỏng tuần tự t bước thêm nhiễu từng bước x₀ → x₁ → … → x_t.
• Định lý Lấy Mẫu Đóng (Closed-Form Marginal Distribution):

q(x_t | x₀) = N(x_t; ; √(α_t) x₀, ; (1 - α_t) I)

• Dùng mẹo tái tham số hóa:

x_t = √(α_t) · x₀ + √(1 - α_t) · ε, với ε ∼ N(0, I)

• Tính chất then chốt này cho phép huấn luyện mạng nơ-ron dự đoán nhiễu ε_θ(x_t, t) tại bước thời gian t bất kỳ được lấy mẫu ngẫu nhiên đồng đều t ∼ U(1, T) mà không tốn chi phí duyệt vòng lặp!

Yêu cầu

Viết hàm:

def q_sample(
    x_0: torch.Tensor,
    t: torch.Tensor,
    alphas_cumprod: torch.Tensor,
    noise: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor]:
    pass

• x_0: Tensor ảnh gốc shape (B, C, H, W) hoặc (B, D).
• t: Tensor 1D (B,) kiểu nguyên chứa chỉ số bước thời gian t ∈ [0, T-1].
• alphas_cumprod: Tensor 1D (T,) chứa α_t.
• noise: Tensor cùng shape với x_0 hoặc None (nếu None tự sinh torch.randn_like(x_0)).
• Trích xuất √(α_t) và √(1 - α_t) định hình lại (reshape) để broadcast được với x_0.
• Trả về tuple (x_t, noise).

Ràng buộc

• Thời gian chạy tối đa: 2000ms.
• Giới hạn bộ nhớ: 256MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.