KPop: Phương pháp huấn luyện RL ổn định cho mô hình MoE quy mô lớn

Nội dung dịch chi tiết

Trong quá trình triển khai mô hình ngôn ngữ lớn (LLM), một hiện tượng phổ biến là mô hình hoạt động tốt ở giai đoạn tiền huấn luyện nhưng lại trở nên bất ổn, giảm điểm thưởng hoặc thậm chí sụp đổ sau khi tinh chỉnh bằng học tăng cường (RL). Nguyên nhân cốt lõi nằm ở sự khác biệt giữa engine huấn luyện (thường dùng PyTorch, độ chính xác cao) và engine suy luận (thường dùng TensorRT, ONNX hoặc lượng tử hóa), dẫn đến sự không nhất quán giữa huấn luyện và suy luận (Training-Inference Mismatch).

Đối với mô hình hỗn hợp chuyên gia (MoE), vấn đề này càng nghiêm trọng do cơ chế định tuyến (routing) khác nhau giữa hai giai đoạn. Phương pháp cũ là IcePop cố gắng khắc phục bằng cách đặt ngưỡng tỷ lệ xác suất cố định để lọc token, nhưng thực tế cho thấy cách này không thích ứng được với quá trình huấn luyện động và vô tình loại bỏ các token hiếm quan trọng.

KPop ra đời để giải quyết hạn chế này bằng cách sử dụng độ lệch KL (KL divergence) nhị phân đối xứng để đo lường sự khác biệt giữa phân phối huấn luyện và suy luận. Cơ chế này cho phép hệ thống tự động điều chỉnh: với các token hiếm, KPop cho phép sai số lớn hơn để khuyến khích khám phá; với các token tần suất cao, nó yêu cầu sự kiểm soát nghiêm ngặt hơn.

Kết quả thực nghiệm trên các mô hình MoE quy mô lớn cho thấy KPop duy trì sự ổn định vượt trội so với IcePop. Trong các tác vụ lập trình dài hạn (SWE-bench Verified), KPop giúp điểm thưởng tăng từ 0,54 lên 0,68 và duy trì huấn luyện ổn định hơn 800 bước, ngay cả khi độ phân kỳ của token lớn.

Một phát hiện thú vị từ nghiên cứu là việc chỉ cập nhật 70% - 80% số lượng token vẫn đủ để mô hình hội tụ tốt. Điều này cho thấy nhiều tín hiệu gradient ở cấp độ token thực chất là nhiễu hoặc dư thừa. Việc chọn lọc token thông minh không chỉ giúp ổn định quá trình huấn luyện mà còn mở ra hướng đi mới để tối ưu hóa hiệu suất tính toán trong tương lai.