Mẹo Tái Tham Số Hóa Và Hàm Mất Mát KL Cho VAE (Kingma & Welling 2013)

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

• Trong Variational Autoencoder (VAE), mạng Encoder xuất ra kỳ vọng mu và log phương sai log sigma² của phân phối tiên nghiệm ẩn q_phi(z | x) = N(mu, diag(sigma²)).
• Mẹo tái tham số hóa (Reparameterization Trick):

Thao tác lấy mẫu ngẫu nhiên z sim N(mu, sigma²) là một nút ngẫu nhiên không thể truyền ngược đạo hàm (non-differentiable). Kingma & Welling tách sự ngẫu nhiên ra khỏi đồ thị tính toán: z = mu + sigma odot epsilon, với epsilon sim N(0, I) Trong đó sigma = exp(0.5 · log sigma²).

• Hàm mất mát phân kỳ KL dạng giải tích (Closed-form KL Divergence) giữa q(z | x) và phân phối chuẩn tắc N(0, I):

D_KL(q(z | x) parallel N(0, I)) = -12 Σ_j=1^D ( 1 + log sigma_j² - mu_j² - sigma_j² )

Yêu cầu

Viết 2 hàm:

def reparameterize(mu: torch.Tensor, logvar: torch.Tensor, eps: torch.Tensor | None = None) -> torch.Tensor:
    # Lấy mẫu z = mu + std * eps
    pass

def compute_vae_loss(
    recon_x: torch.Tensor,
    x: torch.Tensor,
    mu: torch.Tensor,
    logvar: torch.Tensor,
    beta: float = 1.0
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    # Trả về: (total_loss, recon_loss, kl_loss)
    pass

• reparameterize: Nếu eps là None, tạo eps = torch.randn_like(mu).
• compute_vae_loss:
• recon_loss = F.mse_loss(recon_x, x, reduction='mean').
• kl_loss = -0.5 * torch.mean(torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=-1)).
• total_loss = recon_loss + beta * kl_loss.

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ả.