Phạt Độ Lệch KL (Kullback-Leibler Penalty) Từng Token Cho PPO RLHF

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

• Trong vòng lặp PPO-RLHF, nếu chỉ tối đa hóa điểm thưởng R, mô hình sẽ bị hiện tượng Hack Điểm Thưởng (Reward Hacking): sản sinh ra các câu lặp vô nghĩa hoặc nịnh hót để đánh lừa Reward Model.
• Để giữ câu từ tự nhiên và bảo toàn năng lực gốc, ta thêm số hạng phạt KL Divergence giữa chính sách π_θ và chính sách tham chiếu ban đầu π_ref:

R_penalized(x, y) = R(x, y) - β · D_KL(π_θ ∥ π_ref)

• Các ước lượng KL từng mẫu (Schulman 2020):

Đặt r_t = log π_θ(y_t | x, y_<t) - log π_ref(y_t | x, y_<t).

• Ước lượng chuẩn k₁:

D_KL⁽¹⁾ = r_t (Có kỳ vọng đúng bằng KL, nhưng phương sai lớn và có thể nhận giá trị âm).

• Ước lượng xấp xỉ bậc 2 k₃ (Schulman 2020):

D_KL⁽³⁾ = 12 r_t² (Luôn không âm ≥ 0, phương sai cực thấp, rất ổn định khi huấn luyện LLM).

• Ước lượng không chệch không âm k₂:

D_KL⁽²⁾ = e^-r_t - 1 + r_t

Yêu cầu

Viết hàm:

def compute_token_kl_penalty(
    logprobs: torch.Tensor,
    ref_logprobs: torch.Tensor,
    estimator: str = "k3"
) -> torch.Tensor:
    pass

• logprobs: Tensor bất kỳ shape (B, S) chứa log xác suất của chính sách hiện tại.
• ref_logprobs: Tensor cùng shape (B, S) của mô hình tham chiếu.
• estimator: "k1", "k2", hoặc "k3".
• Trả về Tensor cùng shape (B, S) chứa ước lượng độ lệch KL tại từng token.

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