Ưu tiên có thể diễn giải thông qua Mô hình Phần thưởng Đa mục tiêu và Hỗn hợp Chuyên gia

Community Article Published September 29, 2024

Học tăng cường từ phản hồi của con người (RLHF) đã nổi lên như phương pháp chính để điều chỉnh các mô hình ngôn ngữ lớn (LLM) với sở thích của con người. Quá trình RLHF thường bắt đầu bằng việc huấn luyện một mô hình phần thưởng (RM) sử dụng dữ liệu sở thích của con người. Các RM thông thường được huấn luyện trên các phản hồi theo cặp cho cùng một yêu cầu của người dùng, với xếp hạng tương đối cho biết phản hồi nào được con người ưa thích hơn. RM được huấn luyện đóng vai trò như một đại diện cho sở thích của con người. Tuy nhiên, do bản chất hộp đen của RM, kết quả đầu ra của chúng thiếu khả năng diễn giải, vì con người không thể hiểu trực quan tại sao RM lại cho rằng một phản hồi là tốt hay không.

Vì RM hoạt động như đại diện cho sở thích của con người, nên chúng cần có khả năng diễn giải được để đảm bảo rằng các quy trình quyết định nội bộ của chúng phù hợp với sở thích của con người và để ngăn chặn việc khai thác phần thưởng trong việc điều chỉnh LLM. Để xây dựng RM với các sở thích có thể diễn giải, chúng tôi đề xuất một phương pháp hai giai đoạn:

i) huấn luyện Mô hình Phần thưởng Đa mục tiêu Xếp hạng Tuyệt đối (ArmoRM) với dữ liệu xếp hạng tuyệt đối đa chiều, mỗi chiều tương ứng với một mục tiêu có thể diễn giải được (ví dụ: trung thực, độ dài dòng, an toàn);

ii) sử dụng chiến lược Hỗn hợp Chuyên gia (MoE) với một mạng cổng tự động chọn các mục tiêu phần thưởng phù hợp nhất dựa trên ngữ cảnh. Chúng tôi đã huấn luyện hiệu quả một ArmoRM với Llama-3 8B và một mạng cổng bao gồm một MLP nông trên ArmoRM. Mô hình được huấn luyện của chúng tôi, ArmoRM-Llama3-8B, đạt được hiệu suất tiên tiến trên RewardBench, một điểm chuẩn đánh giá RM cho mô hình ngôn ngữ. Đáng chú ý, hiệu suất của mô hình của chúng tôi vượt trội hơn phương pháp LLM-as-a-judge với giám khảo GPT-4, và gần bằng hiệu suất của mô hình phần thưởng Nemotron-4 340B lớn hơn nhiều.

Trong bài báo này, chúng tôi khám phá vai trò của mô hình phần thưởng (RM) trong khuôn khổ Học tăng cường từ Phản hồi của con người (RLHF). RM đóng một vai trò quan trọng trong việc điều chỉnh các mô hình ngôn ngữ lớn (LLM) vì chúng cung cấp một cách có thể mở rộng để tích hợp sở thích của con người vào quá trình huấn luyện của mô hình, hướng dẫn việc tối ưu hóa các chính sách của chúng. Khuôn khổ RLHF phổ biến nhất là khuôn khổ dựa trên học tăng cường sâu (DRL), như được phát triển trong các nghiên cứu quan trọng [Christiano et al., 2017; Ouyang et al., 2022; Bai et al., 2022]. Khuôn khổ này hoạt động theo ba giai đoạn:

Thu thập dữ liệu sở thích;
Mô hình hóa phần thưởng dựa trên mô hình Bradley-Terry [Bradley and Terry, 1952];
Tối ưu hóa chính sách bằng cách sử dụng Tối ưu hóa Chính sách Gần đúng (PPO) [Schulman et al., 2017] và mô hình phần thưởng được xây dựng trong giai đoạn 2. Khuôn khổ này đã đạt được thành công to lớn trong việc đào tạo sau ChatGPT [Ouyang et al., 2022] và Claude [Bai et al., 2022]. Những ý tưởng này cũng mở rộng sang các phương pháp khác, chẳng hạn như tinh chỉnh lấy mẫu từ chối [Dong et al., 2023; Gulcehre et al., 2023] và học sở thích trực tiếp lặp đi lặp lại [Xiong et al., 2023; Guo et al., 2024; Xie et al., 2024]. Trong tất cả các khuôn khổ RLHF này, khả năng của mô hình phần thưởng là rất quan trọng vì nó ảnh hưởng trực tiếp đến chất lượng của LLM được điều chỉnh.

Phương pháp mô hình hóa phần thưởng phổ biến nhất dựa trên ước lượng khả năng tối đa (MLE) của mô hình Bradley-Terry (BT) [Bradley and Terry, 1952]. Mặc dù được sử dụng rộng rãi, mô hình BT khá hạn chế trong khả năng nắm bắt sở thích phức tạp của con người [Munos et al., 2023; Swamy et al., 2024; Ye et al., 2024]. Ngoài vấn đề về khả năng, các RM phổ biến, như mô hình BT, thường là các mô hình hộp đen đưa ra điểm số hoặc sở thích mà không cung cấp lời giải thích có thể diễn giải được, khiến nó dễ bị hiện tượng khai thác phần thưởng được quan sát rộng rãi [Skalse et al., 2022; Singhal et al., 2023; Chen et al., 2024], trong đó LLM được điều chỉnh tạo ra các phản hồi có phần thưởng cao (được RM xếp hạng) nhưng không phù hợp với sở thích thực tế của con người [Gao et al., 2023; Lin et al., 2023; Coste et al., 2023]. Một ví dụ đáng chú ý về điều này là độ dài dòng, trong đó LLM được điều chỉnh tạo ra các phản hồi dài hơn mức cần thiết vì RM ưu tiên độ dài, bất kể chất lượng [Singhal et al., 2023; Wang et al., 2024a; Chen et al., 2024].

Trong nghiên cứu này, chúng tôi nhằm mục đích nâng cao mô hình phần thưởng bằng cách làm cho chúng có thể diễn giải được hơn [Molnar, 2020] và có thể điều khiển được [Wong et al., 2021]. Sử dụng độ dài dòng đã đề cập ở trên làm ví dụ, giả sử đầu ra của RM có thể phân tách được, nghĩa là nó gán điểm cao cho một phản hồi do hai yếu tố: 40% cho tính hữu ích và 60% cho độ dài. Trong trường hợp này, chúng ta có thể thấy rằng RM có thể bị ảnh hưởng bởi độ dài dòng. Hơn nữa, nếu RM có thể điều khiển được, chúng ta có thể điều chỉnh quy trình ra quyết định của nó để dựa trên điểm số 100% vào tính hữu ích. Điều này sẽ bất kể độ dài của phản hồi, do đó giảm thiểu độ dài dòng. Việc nâng cao khả năng diễn giải của RM cũng cho phép con người xác minh xem RM có quy trình quyết định nội bộ tương tự như con người khi hoạt động như đại diện cho sở thích của con người hay không. Chúng tôi tin rằng quá trình tương tác giữa người và AI này có thể đảm bảo rằng RM phù hợp với các giá trị và sở thích của con người, làm cho LLM được điều chỉnh RM đáng tin cậy và mạnh mẽ hơn.

Ở cấp độ cao, chúng tôi đề xuất một phương pháp hai giai đoạn, đầu tiên huấn luyện một RM đa mục tiêu và sau đó học một lớp cổng vô hướng hóa các mục tiêu phần thưởng theo cách hỗn hợp chuyên gia. Sau đó, chúng tôi xác nhận tính hiệu quả của nó bằng thực nghiệm bằng cách huấn luyện một RM như vậy với Llama-3 8B [Meta, 2024] và đạt được hiệu suất tiên tiến trên RewardBench, một điểm chuẩn để đánh giá RM.

Triển khai ArmoRM: Chúng tôi sử dụng kiến trúc Llama-3 8B [Meta, 2024] và khởi tạo xương sống mô hình với các tham số từ RM Bradley-Terry của Llama-3 8B được huấn luyện bởi Dong et al. [2024]. Chúng tôi nối thêm một lớp tuyến tính vào xương sống và huấn luyện nó với tổn thất hồi quy trong khi giữ cho xương sống bị đóng băng. Việc huấn luyện liên quan đến 19 mục tiêu (bao gồm tính hữu ích, tính chính xác, độ dài dòng, v.v.) từ 8 bộ dữ liệu.

Triển khai MoE: Lớp cổng là ReLU MLP gồm 3 lớp ẩn với 1024 đơn vị ẩn. Đối với chỉ số tương quan, chúng tôi áp dụng tương quan Spearman [Spearman, 1904] và sử dụng UltraFeedback [Cui et al., 2023] làm phân phối dữ liệu tham chiếu. Biến tỷ lệ được khởi tạo với giá trị 100 và lớp cổng được huấn luyện với xương sống LLM được giữ cố định. Việc huấn luyện được thực hiện trên 10 bộ dữ liệu sở thích theo cặp.

Việc huấn luyện ArmoRM (giai đoạn mô hình hóa phần thưởng đa mục tiêu) chỉ liên quan đến việc huấn luyện lớp tuyến tính cuối cùng (tức là thăm dò tuyến tính), vì vậy chúng tôi lưu các đặc trưng được trích xuất từ xương sống cục bộ và sau đó tiến hành thăm dò tuyến tính với bộ giải hồi quy tuyến tính của Scikit-learn trên CPU. Đối với giai đoạn MoE, chúng tôi cũng lưu các đặc trưng cục bộ và sau đó huấn luyện lớp cổng trên một GPU NVIDIA A6000 duy nhất.

Lớp cổng được huấn luyện bằng trình tối ưu hóa AdamW [Loshchilov and Hutter, 2019] với tốc độ học 0,001 trong 10.000 bước với kích thước lô là 1024. Chúng tôi cũng áp dụng trình lập lịch tốc độ học phân rã cosine.

RewardBench [Lambert et al., 2024] là điểm chuẩn đầu tiên được xây dựng để đánh giá mô hình phần thưởng cho mô hình ngôn ngữ. Nó bao gồm một tập hợp các nhiệm vụ đa dạng được thiết kế để đánh giá hiệu suất của mô hình phần thưởng cho việc điều chỉnh LLM, bao gồm bốn loại chính (Trò chuyện, Trò chuyện khó, An toàn, Lý luận) và một loại bộ trước. Mỗi loại bao gồm nhiều bộ dữ liệu với dữ liệu sở thích theo cặp, trong đó mỗi cặp bao gồm một phản hồi văn bản được chọn và bị từ chối. Điểm tổng thể được tính là trung bình có trọng số trên năm loại, trong đó bốn loại chính có trọng số 1.0 và loại bộ trước có trọng số 0.5.

Bảng 1 so sánh hiệu suất của phương pháp của chúng tôi (ArmoRM + MoE) với các mô hình phần thưởng khác. Một số quan sát chính có thể được rút ra từ những kết quả này:

Mô hình của chúng tôi vượt trội hơn đáng kể so với Llama-3 8B Bradley-Terry RM, cung cấp xương sống LLM cho mô hình của chúng tôi. Điều này chứng tỏ hiệu quả của thiết kế ArmoRM của chúng tôi và cơ chế cổng MoE trong việc cải thiện hiệu suất của mô hình phần thưởng.
Mô hình của chúng tôi cũng vượt trội hơn phương pháp LLM-as-a-judge [Zheng et al., 2023] với giám khảo GPT-4, cho thấy rằng mô hình của chúng tôi có thể được sử dụng như một sự thay thế rẻ hơn cho GPT-4 trong nhiều công việc chú thích.
Mô hình của chúng tôi với 8B tham số có hiệu suất gần ngang bằng với RM Nemotron-4 340B Wang et al. [2024b], một mô hình phần thưởng khổng lồ với 340B tham số. Điều này làm nổi bật sức mạnh và tiềm năng của phương pháp mô hình hóa phần thưởng của chúng tôi.

Trong nghiên cứu này, chúng tôi đã giải quyết vấn đề quan trọng về khả năng diễn giải trong mô hình phần thưởng cho RLHF trong bối cảnh điều chỉnh LLM với sở thích của con người. Chúng tôi đã đề xuất một phương pháp hai giai đoạn mới, bao gồm ArmoRM và chiến lược MoE với mạng cổng. ArmoRM của chúng tôi, được huấn luyện với Llama-3 8B, đã đạt được hiệu suất tiên tiến trên RewardBench, chứng tỏ hiệu quả của phương pháp mô hình hóa phần thưởng của chúng tôi.

Upvote