RLHF hoạt động thế nào, và tại sao lại là PPO

Supervised fine-tuning dạy được mô hình bắt chước. Nó không dạy được mô hình giỏi hơn thứ nó bắt chước.

Đó là toàn bộ lý do RLHF tồn tại. Nếu bạn có một tập dữ liệu gồm những câu trả lời tốt, bạn huấn luyện mô hình tái tạo lại chúng từng token một — đó vẫn là dự đoán token kế tiếp, chỉ với dữ liệu đẹp hơn. Nhưng phần lớn cái làm nên một câu trả lời tốt lại không viết ra được thành một chuỗi mục tiêu. Nó mang tính so sánh: câu này hữu ích hơn câu kia, câu này từ chối trong khi lẽ ra không nên, câu này đúng về kỹ thuật nhưng đọc không nổi. Không ai viết ra được câu trả lời đúng, nhưng gần như ai cũng chọn được câu nào hay hơn trong hai câu.

RLHF là bộ máy biến phép so sánh đó thành gradient. PPO là phần giữ cho gradient không phá hỏng mô hình trên đường đi.

Ba giai đoạn

RLHF không phải một thuật toán. Nó là ba lượt huấn luyện xếp chồng lên nhau, và nhầm lẫn giữa chúng là nguồn gốc thông thường của mọi hiểu lầm về nó.

Giai đoạn 1 — supervised fine-tuning. Lấy một mô hình pretrained và fine-tune trên dữ liệu minh hoạ: prompt kèm câu trả lời do người viết hoặc phê duyệt. Đây là huấn luyện cross-entropy thông thường. Nhiệm vụ của nó không phải làm mô hình giỏi, chỉ là làm nó trả lời đúng định dạng để hai giai đoạn sau có thứ dùng được mà lấy mẫu. Gọi kết quả là mô hình SFT.

Giai đoạn 2 — reward model. Lấy nhiều câu trả lời cho cùng một prompt từ mô hình SFT, cho người xếp hạng, rồi huấn luyện một mô hình riêng dự đoán thứ hạng đó dưới dạng một số vô hướng.

Giai đoạn 3 — reinforcement learning. Để mô hình SFT sinh ra câu trả lời, chấm điểm bằng reward model, rồi cập nhật mô hình sao cho điểm cao hơn. Đây là chỗ PPO xuất hiện.

Những quyết định thiết kế đáng chú ý đều nằm ở giai đoạn 2 và 3.

Giai đoạn 2: học một reward từ phép so sánh

Con người chấm điểm tuyệt đối rất tệ. Bảo mười người chấm một câu trả lời từ 1 đến 10 thì bạn nhận về mười thang đo khác nhau. Hỏi họ trong hai câu thì câu nào hơn, mức đồng thuận cao hơn hẳn. Nên dữ liệu ưu tiên được thu thập theo cặp: với prompt x, một câu được chọn y_w và một câu bị loại y_l.

Reward model r thường chính là mô hình SFT, thay đầu ra ngôn ngữ bằng một đầu ra vô hướng. Nó được huấn luyện bằng hàm mục tiêu Bradley-Terry, phát biểu rằng xác suất một người thích y_w hơn y_l là sigmoid của hiệu hai điểm số:

# reward_chosen, reward_rejected: (batch,) số vô hướng từ reward model
loss = -torch.nn.functional.logsigmoid(reward_chosen - reward_rejected).mean()

Ba hệ quả rơi ra từ hàm mất mát này, và cả ba đều quan trọng về sau.

Reward chỉ có ý nghĩa tương đối so với các câu trả lời khác của cùng một prompt. Không có gì neo giá trị tuyệt đối lại — cộng một hằng số vào mọi điểm thì hàm mất mát không đổi. Nên những con số reward model xuất ra không so sánh được giữa các prompt khác nhau, và "reward tăng" chỉ có nghĩa trong phạm vi một prompt.

Nó là một hàm xấp xỉ đã học, không hoàn hảo, cho sở thích của con người, huấn luyện trên một mẫu hữu hạn. Ở bất cứ đâu policy đi lạc ra ngoài phân phối mà reward model từng thấy, đầu ra của nó là phỏng đoán.

Và nó bị đóng băng trong giai đoạn 3. Policy sẽ tối ưu ngược lại một hàm cố định, không hoàn hảo, suốt hàng nghìn bước. Nhớ kỹ điều này.

Tại sao lại cần reinforcement learning?

Đây là phiên bản thẳng thắn của câu hỏi: bạn đã có một số vô hướng cho mỗi chuỗi sinh ra rồi. Sao không fine-tune luôn trên những câu điểm cao và kết thúc ở đó?

Bạn làm được, và đó gần như chính là rejection sampling hay best-of-n fine-tuning. Nó hiệu quả, và đơn giản hơn nhiều. RL xứng đáng với độ phức tạp của nó vì một lý do: reward đến một lần duy nhất, ở cuối cả chuỗi, và nó không khả vi theo tham số của mô hình.

Nghĩ xem mô hình thực sự làm gì. Ở mỗi bước nó chọn một token từ một phân phối trên toàn bộ từ vựng. Việc lấy mẫu không khả vi — bạn không thể lan truyền ngược qua "và rồi nó chọn token thứ 4.712". Reward model chấm điểm văn bản đã hoàn chỉnh, nên không có mục tiêu ở cấp token nào để tính cross-entropy.

Reinforcement learning chính là bộ công cụ cho đúng dạng bài toán này: hành động rời rạc, phản hồi không khả vi, và công lao phải được gán ngược lại cho một chuỗi lựa chọn. Ánh xạ sang thuật ngữ RL:

• state — prompt cộng với mọi thứ đã sinh ra tới thời điểm đó
• action — token kế tiếp
• policy — chính mô hình ngôn ngữ
• reward — điểm của reward model, trả về ở token cuối cùng

Policy gradient, và chỗ nó hỏng

Ý tưởng policy gradient cơ bản khá trực giác. Tăng log-probability của những hành động dẫn tới reward cao, giảm của những hành động dẫn tới reward thấp, đánh trọng số theo mức độ tốt hơn kỳ vọng:

loss = -(log_prob * advantage).mean()

advantage là reward trừ đi một mức nền — hành động này tốt hơn bao nhiêu so với mức bạn thường mong đợi ở trạng thái đó. Trừ đi mức nền không làm đổi hướng gradient tính trung bình, nhưng giảm phương sai của nó rất mạnh, và đó là lý do mọi bản cài đặt thực tế đều có một value network dự đoán mức nền ấy.

Vấn đề nằm ở độ dài bước. Gradient này chỉ đúng với chính policy đã sinh ra dữ liệu. Đi một bước lớn thì policy mới đã khác đủ nhiều để các trajectory đã thu thập không còn mô tả nó nữa, và lần cập nhật kế tiếp được tính từ thông tin đã cũ. Ở mô hình ngôn ngữ, kiểu hỏng này rất kịch tính: mô hình tìm thấy một vùng suy biến có reward cao, sụp về đó, rồi bắt đầu nhả ra cùng một đoạn văn nịnh nọt bất kể prompt là gì. Một khi nó ngừng sinh ra mẫu đa dạng, nó không hồi phục được, vì nó không còn khám phá bất cứ thứ gì mà reward model chấm khác đi.

Vậy nên bạn cần bước nhỏ. Nhưng "dùng learning rate nhỏ" là một công cụ thô — cùng một learning rate sẽ quá lớn ở vùng này và nhỏ đến lãng phí ở vùng khác.

PPO: chặn bước cập nhật, không phải chặn learning rate

Cái hay của PPO là ràng buộc bước cập nhật theo mức thay đổi của phân phối đầu ra, chứ không theo khoảng cách tham số đã dịch chuyển.

Định nghĩa tỉ lệ xác suất giữa policy mới và policy đã sinh ra dữ liệu:

ratio = pi_new(a | s) / pi_old(a | s)

Tỉ lệ bằng 1 nghĩa là policy không đổi với token đó. Lớn hơn 1 nghĩa là policy mới có xu hướng sinh ra token đó nhiều hơn. Hàm mục tiêu thay thế có cắt ngưỡng là:

ratio = torch.exp(new_log_probs - old_log_probs)   # exp của hiệu hai log
unclipped = ratio * advantages
clipped = torch.clamp(ratio, 1 - eps, 1 + eps) * advantages
policy_loss = -torch.min(unclipped, clipped).mean()

với eps thường bằng 0.2. Đọc kỹ chữ min, vì đó là toàn bộ mẹo, và nó không đối xứng.

Khi advantage dương — token này tốt hơn kỳ vọng — hàm mục tiêu muốn tăng tỉ lệ. Cắt ngưỡng chặn nó lại ở 1 + eps. Qua điểm đó gradient bằng không: mô hình không được thưởng thêm vì đẩy xác suất token này cao hơn nữa trong lần cập nhật này. Nó đã dịch chuyển đủ rồi.

Khi advantage âm, hàm mục tiêu muốn giảm tỉ lệ, và cắt ngưỡng chặn sàn ở 1 - eps. Cùng một logic, theo chiều ngược lại.

Vì lấy giá trị nhỏ hơn giữa số hạng bị cắt và không bị cắt, PPO mang tính bi quan. Nó không bao giờ để việc cắt ngưỡng làm một bước cập nhật trở nên hấp dẫn hơn, chỉ có thể kém hơn. Và cách cắt này một chiều theo một nghĩa hữu ích: nếu tỉ lệ đã vọt xa khỏi khoảng theo hướng sai, số hạng không cắt được chọn và gradient vẫn kéo nó về. Cắt ngưỡng ngăn bạn chạy đi mất; nó không ngăn bạn quay về.

Cái lợi thực tế là bạn có thể lấy vài bước gradient trên cùng một lô dữ liệu đã sinh trước khi nó trở nên cũ — bình thường bạn chỉ được một bước — và điều đó cực kỳ đáng giá khi việc sinh ra lô dữ liệu ấy nghĩa là chạy inference qua một mô hình lớn.

Dây xích KL

Cắt ngưỡng giới hạn policy dịch chuyển bao xa trong mỗi lần cập nhật. Nó không nói gì về chỗ policy sẽ dừng lại sau mười nghìn lần. Cứ để mặc, mô hình sẽ vui vẻ trôi tới nơi nào đó mà reward model chấm cao còn con người thấy vô dụng, bởi reward model là một hàm xấp xỉ, và mọi hàm xấp xỉ đều có kẽ hở.

Đây không phải chuyện giả định. Tối ưu đủ mạnh ngược lại một reward model và bạn nhận về những đầu ra nhồi đầy các đặc điểm bề mặt mà nó tương quan với chất lượng: rào đón quá mức, bullet point khắp nơi, nhắc lại câu hỏi, và trên hết là độ dài, vì câu trả lời dài chấm điểm cao hơn trong hầu hết tập dữ liệu ưu tiên, và reward model đã học được điều đó.

Cách chữa là một khoản phạt KL so với mô hình SFT đã đóng băng, tính trên từng token:

# reference = mô hình SFT đã đóng băng; cả hai đều là log-prob của token đã lấy mẫu
kl = new_log_probs - ref_log_probs
rewards = reward_model_score - kl_coef * kl

Policy giờ tối ưu điểm của reward model trừ đi một khoản phí cho mỗi token mà nó lệch khỏi mô hình nó xuất phát. Hệ số quy định tỉ giá: một đơn vị trôi dạt đáng giá bao nhiêu reward. Đặt cao quá thì không có gì thay đổi; thấp quá thì mô hình đi lạc và luồn lách reward.

Hai chi tiết đáng biết. Điểm của reward model là một con số duy nhất ở cuối chuỗi, trong khi số hạng KL bị tính ở mọi token — tín hiệu reward tổng thể dày đặc về KL và thưa thớt về sở thích thật. Và trên thực tế KL thường được ước lượng bằng một estimator không chệch có phương sai thấp hơn, thay vì hiệu log-ratio thô, vì bản ngây thơ đủ nhiễu để gây hại.

Toàn bộ vòng lặp

for batch in prompts:
    # 1. sinh, và ghi lại log-prob của chính policy đã sinh ra chúng
    responses, old_log_probs = policy.generate(batch)

    # 2. chấm điểm
    scores = reward_model(batch, responses)             # một số vô hướng mỗi câu
    ref_log_probs = reference_model.log_probs(batch, responses)
    kl = old_log_probs - ref_log_probs
    rewards = scores - kl_coef * kl                     # theo từng token

    # 3. gán công lao
    values = value_head(batch, responses)
    advantages = compute_gae(rewards, values, gamma, lam)

    # 4. vài epoch trên cùng một lô -- đây là thứ mà cắt ngưỡng mua cho bạn
    for _ in range(ppo_epochs):
        new_log_probs = policy.log_probs(batch, responses)
        ratio = torch.exp(new_log_probs - old_log_probs)

        policy_loss = -torch.min(
            ratio * advantages,
            torch.clamp(ratio, 1 - eps, 1 + eps) * advantages,
        ).mean()

        value_loss = ((value_head(batch, responses) - returns) ** 2).mean()
        loss = policy_loss + vf_coef * value_loss

        loss.backward()
        optimizer.step()

Tức là bốn mô hình cùng nằm trong bộ nhớ một lúc: policy đang được huấn luyện, bản reference đóng băng, reward model, và value head. Đó là lý do RLHF được nhớ đến vì tốn kém và rắc rối chẳng kém gì vì thành quả nó mang lại.

Những gì thực sự hỏng

Reward hacking là kiểu hỏng nổi bật nhất, và nó không phải một lỗi bạn sửa được bằng cách cẩn thận hơn. Reward model là xấp xỉ từ mẫu hữu hạn cho sở thích con người; tối ưu một hàm xấp xỉ đủ mạnh thì luôn tìm ra khoảng cách giữa nó và thứ nó đại diện. Nhìn đường reward đi lên tự nó không nói lên điều gì — nó đi lên y hệt nhau dù mô hình đang tốt lên hay đang học cách khai thác kẽ hở.

Length bias xứng đáng có tên riêng vì nó xuất hiện quá đều đặn. Người gán nhãn hơi thích câu trả lời đầy đủ, reward model học "dài hơn thì tốt hơn" như một tương quan dễ nắm, và policy học cách nói dài dòng. Vài bản cài đặt trừ thẳng một khoản phạt độ dài; phần lớn chỉ theo dõi số token trung bình.

KL sụp theo cả hai chiều. Phạt quá nặng thì bạn tiêu tốn rất nhiều tính toán để tái tạo lại đúng mô hình SFT. Phạt quá nhẹ thì policy rời khỏi phân phối mà reward model hiểu được, nơi điểm số của nó chỉ là phỏng đoán, và nó tối ưu ngược lại nhiễu.

Value function bất ổn. Value head khởi đầu chẳng biết gì, nên advantage giai đoạn đầu chủ yếu là nhiễu, sinh ra những bước cập nhật tệ đúng lúc policy dễ bị uốn nắn nhất. Làm nóng value head trước khi cho policy dịch chuyển là cách chữa phổ biến.

Sau đó là gì

Điều dễ phàn nàn nhất về quy trình trên là nó quá đồ sộ so với việc nó làm. Huấn luyện một reward model, rồi chạy một vòng lặp RL với bốn mô hình thường trú, cân bằng cẩn thận với một số hạng KL, chỉ để rút tín hiệu ra từ dữ liệu vốn chỉ là những cặp "câu này hơn".

Direct Preference Optimization là câu trả lời nổi tiếng cho chuyện đó. Phần dẫn giải cho thấy policy tối ưu dưới một hàm mục tiêu reward có ràng buộc KL tồn tại ở dạng đóng, và dạng đó sắp xếp lại được để biểu diễn reward theo policy và mô hình reference. Thay ngược vào hàm mất mát Bradley-Terry thì reward model triệt tiêu. Còn lại là một hàm mất mát có giám sát tính thẳng trên các cặp ưu tiên:

chosen_logratio = policy_chosen_logps - reference_chosen_logps
rejected_logratio = policy_rejected_logps - reference_rejected_logps
loss = -torch.nn.functional.logsigmoid(beta * (chosen_logratio - rejected_logratio)).mean()

Không reward model, không vòng lặp lấy mẫu, không value head. Hai mô hình thay vì bốn, và một vòng huấn luyện bình thường.

Nó không hẳn tốt hơn. DPO chỉ huấn luyện trên tập ưu tiên cố định, nên nó không bao giờ nhìn thấy đầu ra hiện tại của chính policy — mà khả năng sinh ra, chấm điểm, rồi học từ thứ mô hình tạo ra ngay lúc này mới đúng là thứ online RL mang lại. Đó là lý do các biến thể online tồn tại, và là lý do PPO chưa biến mất.

Và đó là cách hữu ích để nhìn toàn bộ chuyện này: PPO không phải trọng tâm của RLHF. Trọng tâm là chuyển sở thích thành tín hiệu huấn luyện. PPO là một câu trả lời cho câu hỏi làm sao đi theo tín hiệu đó mà không ngã, và hiểu tại sao nó cắt ngưỡng cho bạn nhiều hơn là thuộc lòng rằng nó có cắt.

Tham khảo

• Training language models to follow instructions with human feedback — bài InstructGPT, mô tả kinh điển về quy trình ba giai đoạn
• Proximal Policy Optimization Algorithms — chính PPO, ra đời trước tất cả những thứ trên năm năm
• Direct Preference Optimization — phần dẫn giải loại bỏ reward model
• Deep reinforcement learning from human preferences — công trình trước đó, thiết lập việc học một reward model từ phép so sánh