$Q\sharp$: Provably Optimal Distributional RL for LLM Post-Training
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866912658415943680 |
|---|---|
| author | Zhou, Jin Peng Wang, Kaiwen Chang, Jonathan Gao, Zhaolin Kallus, Nathan Weinberger, Kilian Q. Brantley, Kianté Sun, Wen |
| author_facet | Zhou, Jin Peng Wang, Kaiwen Chang, Jonathan Gao, Zhaolin Kallus, Nathan Weinberger, Kilian Q. Brantley, Kianté Sun, Wen |
| contents | Reinforcement learning (RL) post-training is crucial for LLM alignment and reasoning, but existing policy-based methods, such as PPO and DPO, can fall short of fixing shortcuts inherited from pre-training. In this work, we introduce $Q\sharp$, a value-based algorithm for KL-regularized RL that guides the reference policy using the optimal regularized $Q$ function. We propose to learn the optimal $Q$ function using distributional RL on an aggregated online dataset. Unlike prior value-based baselines that guide the model using unregularized $Q$-values, our method is theoretically principled and provably learns the optimal policy for the KL-regularized RL problem. Empirically, $Q\sharp$ outperforms prior baselines in math reasoning benchmarks while maintaining a smaller KL divergence to the reference policy. Theoretically, we establish a reduction from KL-regularized RL to no-regret online learning, providing the first bounds for deterministic MDPs under only realizability. Thanks to distributional RL, our bounds are also variance-dependent and converge faster when the reference policy has small variance. In sum, our results highlight $Q\sharp$ as an effective approach for post-training LLMs, offering both improved performance and theoretical guarantees. The code can be found at https://github.com/jinpz/q_sharp. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_20548 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | $Q\sharp$: Provably Optimal Distributional RL for LLM Post-Training Zhou, Jin Peng Wang, Kaiwen Chang, Jonathan Gao, Zhaolin Kallus, Nathan Weinberger, Kilian Q. Brantley, Kianté Sun, Wen Machine Learning Artificial Intelligence Computation and Language Reinforcement learning (RL) post-training is crucial for LLM alignment and reasoning, but existing policy-based methods, such as PPO and DPO, can fall short of fixing shortcuts inherited from pre-training. In this work, we introduce $Q\sharp$, a value-based algorithm for KL-regularized RL that guides the reference policy using the optimal regularized $Q$ function. We propose to learn the optimal $Q$ function using distributional RL on an aggregated online dataset. Unlike prior value-based baselines that guide the model using unregularized $Q$-values, our method is theoretically principled and provably learns the optimal policy for the KL-regularized RL problem. Empirically, $Q\sharp$ outperforms prior baselines in math reasoning benchmarks while maintaining a smaller KL divergence to the reference policy. Theoretically, we establish a reduction from KL-regularized RL to no-regret online learning, providing the first bounds for deterministic MDPs under only realizability. Thanks to distributional RL, our bounds are also variance-dependent and converge faster when the reference policy has small variance. In sum, our results highlight $Q\sharp$ as an effective approach for post-training LLMs, offering both improved performance and theoretical guarantees. The code can be found at https://github.com/jinpz/q_sharp. |
| title | $Q\sharp$: Provably Optimal Distributional RL for LLM Post-Training |
| topic | Machine Learning Artificial Intelligence Computation and Language |
| url | https://arxiv.org/abs/2502.20548 |