Hàm Mất Mát Khớp Dòng Tuyến Tính (Flow Matching / Rectified Flow - SD3 & Flux.1)

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

• Mô hình Diffusion truyền thống (DDPM) dựa trên quỹ đạo khuếch tán cong (Curved Brownian Motion) khiến phương trình vi phân SDE/ODE bị uốn lượn phức tạp và cần nhiều bước Euler để giải.
• Khớp dòng (Flow Matching - Lipman et al. 2022, Liu et al. 2022) thay thế hoàn toàn đường cong bằng đoạn thẳng trực tiếp (Straight Line Interpolation / Optimal Transport) nối giữa phân phối nhiễu x₀ ∼ N(0, I) và ảnh thật x₁:

x_t = (1 - t) x₀ + t x₁, t ∈ [0, 1]

• Đạo hàm vận tốc tức thời theo thời gian là một hằng số cố định dọc suốt đoạn thẳng:

u_t(x₁ | x₀) = d x_td t = x₁ - x₀

• Hàm mất mát hồi quy trường vector (Conditional Flow Matching Loss):

Mạng nơ-ron v_θ(x_t, t) học cách dự đoán trực tiếp vector vận tốc x₁ - x₀: L_FM(θ) = E_{t, x0, x1} [ | v_θ(x_t, t) - (x₁ - x₀) |² ]

• Nhờ quỹ đạo thẳng tắp (Straight trajectories), Flux.1 và Stable Diffusion 3 có thể sinh ảnh chất lượng đỉnh cao chỉ với 4 đến 10 bước Euler!

Yêu cầu

Viết hàm:

def compute_flow_matching_loss(
    model_predicted_velocity: torch.Tensor,
    x_0: torch.Tensor,
    x_1: torch.Tensor
) -> torch.Tensor:
    pass

• model_predicted_velocity: Tensor v_θ(x_t, t).
• x_0: Tensor nhiễu ban đầu ∼ N(0, I).
• x_1: Tensor ảnh thật mục tiêu.
• Trả về scalar tensor MSE loss giữa v_θ và (x₁ - x₀).

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