Tối Ưu Hóa Sở Thích Trực Tiếp (DPO Loss - Rafailov et al. NeurIPS 2023)

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

• Quy trình RLHF truyền thống đòi hỏi 3 giai đoạn phức tạp: (1) SFT, (2) Huấn luyện Reward Model, (3) Tối ưu hóa PPO với 4 mô hình trong VRAM (Actor, Critic, Reference, Reward Model).
• Định lý DPO (Rafailov et al. 2023) chứng minh: Ta có thể trích xuất nghiệm tối ưu đóng của bài toán RLHF có phạt KL, biến điểm thưởng tiềm ẩn thành tỷ số log xác suất giữa mô hình đang huấn luyện π_θ và mô hình tham chiếu đóng băng π_ref:

r(x, y) = β log π_θ(y | x)π_ref(y | x)

• Thay vào hàm mất mát Bradley-Terry ta được DPO Loss:

L_DPO(θ) = - E_{(x, yw, yl)} [ log σ ( β log π_θ(y_w | x)π_ref(y_w | x) - β log π_θ(y_l | x)π_ref(y_l | x) ) ] Biến đổi đại số thành: h = β ( [log π_θ(y_w) - log π_ref(y_w)] - [log π_θ(y_l) - log π_ref(y_l)] ) L_DPO = - 1B ∑_i=1^B log σ(h_i)

• DPO không cần mạng Critic, không cần Reward Model, ổn định và nhanh gấp nhiều lần PPO!

Yêu cầu

Viết hàm:

def compute_dpo_loss(
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    ref_chosen_logps: torch.Tensor,
    ref_rejected_logps: torch.Tensor,
    beta: float = 0.1
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    pass

• Đầu vào: 4 tensor 1D (B,) chứa tổng log-probabilities của chuỗi token.
• Đầu ra trả về tuple (loss, chosen_rewards, rejected_rewards):
• loss: scalar tensor mất mát DPO.
• chosen_rewards: Tensor (B,) ẩn: β (log π_θ(y_w) - log π_ref(y_w)).
• rejected_rewards: Tensor (B,) ẩn: β (log π_θ(y_l) - log π_ref(y_l)).

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