Dùng RL controller nhỏ để điều phối LLM lớn thông minh hơn

Nếu bạn đã từng làm việc với các bài toán reasoning dùng LLM, chắc bạn biết cái trick phổ biến nhất hiện nay là test-time scaling tức là thay vì train model to hơn, mình cho nó generate nhiều lần rồi chọn câu trả lời tốt nhất (kiểu majority voting hoặc best-of-N). Cách này work khá tốt, nhưng cái giá phải trả là compute tăng tuyến tính theo số lần sample. Câu hỏi đặt ra là: có cách nào để sample thông minh hơn không, thay vì cứ sample đủ N lần cho chắc?

Một paper vừa ra tuần này "Small RL Controller, Large Language Model" đề xuất một hướng khá thú vị để giải quyết đúng vấn đề này.

Vấn đề với các adaptive sampling hiện tại

Hiện tại đã có một số phương pháp adaptive sampling, như ASC (Adaptive Sampling with Confidence) hay ESC. Ý tưởng chung là: thay vì luôn sample đủ N lần, mình dừng sớm nếu model đã đủ confident. Nghe có vẻ hợp lý, nhưng vấn đề là các phương pháp này thường dựa trên heuristic rules hoặc giả định về distribution của output kiểu như "nếu 70% câu trả lời giống nhau thì dừng". Cách này fragile, vì threshold đó bạn lấy từ đâu? Tune tay? Hardcode? Không có gì đảm bảo nó optimal cả.

Theo kinh nghiệm của mình khi làm với các hệ thống inference production, cái khó nhất không phải là accuracy mà là cân bằng giữa accuracy, latency, và cost ba cái này thường trade-off lẫn nhau, và không có một con số threshold nào fit cho mọi bài toán.

Controller nhỏ, quyết định lớn

Paper này formulate bài toán theo hướng khác hẳn: họ model adaptive sampling như một Markov Decision Process (MDP). Cụ thể:

• State: thống kê từ các câu trả lời đã sample được (phân phối câu trả lời, mức độ đồng thuận, v.v.)
• Action: tiếp tục sample thêm, hay dừng lại
• Reward: được thiết kế để jointly optimize cả accuracy, latency, và computation cost

Rồi họ train một RL controller cực kỳ nhẹ để học policy này. Cái hay là controller này chỉ nhìn vào statistics của final answers không cần access vào hidden states hay logits của LLM. Điều này có nghĩa là nó model-agnostic và đặc biệt hơn, nó có thể train và deploy trên CPU không cần GPU.

Mình thấy cái này hay ở chỗ: bạn có một LLM to đang chạy trên GPU cluster tốn kém, nhưng cái "não" quyết định có nên gọi nó thêm không lại là một model bé xíu chạy trên CPU. Về mặt kiến trúc hệ thống, đây là một separation of concerns rất clean.

Lagrangian relaxation khi optimization theory gặp thực tế

Một phần thú vị khác trong paper là họ chứng minh được rằng framework này có thể được diễn giải như Lagrangian relaxation của một constrained optimization problem với explicit budget constraints. Nói nôm na: bạn có thể đặt ra ràng buộc kiểu "tổng số samples không được vượt quá X" hoặc "latency không được quá Y milliseconds", và framework sẽ tự học cách trade-off để đáp ứng constraint đó.

Cái này quan trọng trong production vì SLA thường đến dưới dạng hard constraints, không phải soft preferences. Bạn không nói "mình muốn nhanh hơn một chút" bạn nói "p99 latency phải dưới 2 giây, không thì alert page on-call".

So sánh với baseline

Kết quả thực nghiệm cho thấy method này đạt được trade-off tốt hơn so với ASC và ESC trên cả ba chiều: accuracy, số sampling rounds, và tổng số samples cần thiết.

Ứng dụng thực tế anh em lưu ý

Nếu bạn đang build hệ thống RAG hoặc reasoning pipeline với LLM ở scale, đây là loại research đáng để theo dõi vì nó tackle đúng pain point của production:

• Cost: Mỗi lần gọi LLM là tiền. Nếu controller học được rằng câu hỏi đơn giản chỉ cần 2-3 samples thay vì 10, bạn tiết kiệm được 70-80% inference cost cho những câu đó.
• Latency: Với user-facing application, dừng sớm = response nhanh hơn = UX tốt hơn.
• Flexibility: Vì framework support explicit budget constraints, bạn có thể tune behavior khác nhau cho các use case khác nhau ví dụ strict latency budget cho chatbot, nhưng accuracy-first cho batch processing.

Cái mình thấy cần explore thêm là: controller này được train trên distribution nào? Nếu distribution của câu hỏi production shift so với training data, policy có còn work không? Đây là câu hỏi open mà mình nghĩ paper chưa đề cập kỹ, và nó sẽ là thách thức thực tế khi deploy.

Nhưng về overall direction, việc dùng một lightweight RL agent để orchestrate một heavy LLM là một pattern rất promising và thực ra không chỉ giới hạn ở sampling. Bạn có thể imagine tương tự cho routing (câu hỏi nào cần model to, câu nào model nhỏ đủ), hay tool use (có cần search thêm không). Cái framework MDP + RL controller nhỏ này có thể generalize ra nhiều bài toán meta-decision trong LLM inference pipeline.

Paper đang available trên arXiv, mình recommend đọc phần formulation MDP và phần Lagrangian relaxation nếu bạn quan tâm đến lý thuyết cách họ connect RL reward với constrained optimization khá elegant.