$Q\sharp$: Provably Optimal Distributional RL for LLM Post-Training

Dear Reviewer gBqj,

Thank you so much for your detailed review and suggestions for us to improve the paper. We respond to individual points below.

Could the authors conduct experiments with more models to enhance the generality of the approach?

Authors' reply: Thank you for the suggestion! Table 2 of the paper shows that both Llama 1B and Llama 3B can effectively guide Llama 70B, further improving its performance despite its strength over 8B models, supporting the generality of our approach.

Additionally, as suggested by Reviewer GoFf, we tested Q# on a long-context benchmark, where it also performs well (see our response to Reviewer GoFf for details). We further extended experiments to the Qwen model series—please see the next point for results.

Evaluation with different prompting methods like COT or TOT.

Authors' reply: We conducted experiments using Qwen 2.5 1.5B to guide Qwen 2.5 7B model on GSM8K and MATH with COT prompting ("Let's think step by step"):

This shows that Q# consistently outperforms baselines across both datasets while maintaining lower KL penalty, and highlights the generality of Q# across different model families and prompting methods. We will add these results in the final version.

Deterministic MDP assumption may limit applicability in settings with stochasticity. Can Q# be extended to MDPs with stochastic transitions?

Authors' reply: In this work, we focus on deterministic MDPs since the $Q^{\star,\eta}$ is a simple functional of the cumulative rewards distribution $Z^{\pi^{ref}}$, and that they capture the LLM post-training setting very well. If the MDP has stochastic transitions, then $Q^{\star,\eta}$ won't be a functional of $Z^{\pi^{ref}}$ (i.e. Thm 2.2 won't hold) and we will need to learn $Q^{\star,\eta}$ via temporal difference learning. Theoretically, this will require the bellman completeness assumption, which is much stronger than our realizability. Practically, TD-learning with function approximation is much less stable than online regression. Thus, while Q# can be extended to stochastic MDPs via TD-learning, this direction presents additional challenges and is less directly relevant to the LLM post-training application we focus on.

Could the authors provide more detailed comparisons of computational complexity and runtime performance between Q♯ and baseline methods?

Authors' reply: Q# and the CD baseline have the same computational complexity and runtime since both use the guidance model to compute a Q function alongside $\pi^{\text{ref}}$. Compared to generating responses solely with $\pi^{\text{ref}}$, Q# increases complexity by the ratio of the guidance model’s size to that of $\pi^{\text{ref}}$. For instance, guiding Llama 8B with Llama 1B increases complexity by ~12.5%.

We efficiently implemented Q# in Hugging Face using LogitProcessor and key-value caches. On an Nvidia A6000, generating one response on test set of MATH takes 4.10s for $\pi^{\text{ref}}$ and 5.18s for Q#, slightly exceeding 12.5% possibly due to sequential Q function computation in LogitProcessor. We will open-source our implementation in the final version.

it could benefit from referencing some of the most recent advances in stable distributional RL.

Authors' reply: Thanks for the suggestion, we’ll include more references about DistRL in the final version. For example, [1] shows that the categorical DistRL loss, which we employ for our theory and experiments, enjoys smoothness and optimization stability under a bounded logit condition. [2] introduces a Sinkhorn DistRL loss which is a computationally efficient alternative for Wasserstein distance, and was shown to be more stable for multi-dimensional rewards. [3] proposed a KL-regularized categorical loss which they showed is empirically more stable in Atari games. However, these references all apply TD-learning with function approximation and replay buffers, which (Sutton and Barto, 2018) identified as a deadly triad that is notoriously difficult to scale, requiring many tricks such as double Q-learning and target networks. In contrast, our work obviates the need for TD-learning or tricks such as the target network by leveraging the special form of Q* in deterministic KL-regularized MDPs, which perfectly captures the LLM post-training application we focus on.

[1] Sun et al. "How Does Return Distribution in Distributional Reinforcement Learning Help Optimization?."

[2] Sun et al. "Distributional reinforcement learning with regularized wasserstein loss."

[3] Li et al. "Efficient distributional reinforcement learning with Kullback-Leibler divergence regularization."

Dear Reviewer B1zk,

Thank you for your encouraging review and strong support! We address individual questions below.

On the right hand side of the first equation should be $Q^{\star,\eta}_h$.

Authors' reply: Thanks for catching this typo! We will make sure to fix it in the final version of the paper.

The term $Reg_{mle}$ has not been formally defined. In context it is clear that it represents the regret but it should be formally defined.

Authors' reply: The reviewer is fully correct that $Reg_{mle}$ represents the regret of the MLE oracle after K iterations, which is currently defined in Line 377. We will make sure to highlight this definition in the revision.

Please let us know if you have any other questions!

Dear Reviewer GoFf,

Thank you for the detailed review and constructive feedback! We truly appreciate the time and effort that have been invested in providing constructive comments. Please find our responses below.

Star-graph results for claiming policy-based method shortcomings in planning tasks

Authors' reply: We appreciate the reviewer’s feedback and agree that our claim about policy-based methods picking up shortcuts should be framed more narrowly. Our findings are specific to the star-graph task, where we observe that policy-based methods tend to exploit pretraining shortcuts rather than learning effective planning strategies. While this provides concrete evidence of the issue in this setting, we acknowledge that further investigation is needed to generalize this claim beyond the studied task. We will clarify this limitation in the final version of the paper.

Additional related work such as [Nachum'19a][Nachum'19b]

Authors' reply: Thank you for pointing out DICE, which we’ll certainly add a discussion and citation in our final version. DualDICE (for off-policy evaluation [Nachum'19a]) and AlgaeDICE (for off-policy learning [Nachum'19b]) are motivated by the curse of horizon in offline RL where the variance of importance weighting grows exponentially in $H$. DICE obviates this explosion in variance by Lagrangianizing the linear program of $Q^\pi$ and fitting an additional nuisance function $w$, which is the density ratio between the visitation measure under $\pi$ and the data distribution $\nu$. DICE was shown to be effective in low-dimensional control tasks [Nachum’19a, Nachum’19b], although learning the $w$ function was also shown to be difficult in high-dimensional control tasks [Chang’22]. Exploring DICE for LLM post-training is a promising direction to pursue for future work and is beyond the scope of our paper.

the work does not delve into more unique properties of LLM post-training (like long contexts)

Authors' reply: Thank you for the valuable comment! Following your suggestion, we conducted experiments on a long context benchmark, QuALITY [Pang’22], a multiple-choice QA dataset with context passages in English sourced from Project Gutenberg. We experimented with two settings: Qwen 2.5 1B guiding Qwen 2.5 7B and Llama 3.2 1B guiding Llama 3.1 8B.

The results suggest that Q# can also consistently improve performance on long context tasks for models such as Qwen 2.5 and Llama 3.1 with smaller KL divergence than CD baseline.

Additionally, we conducted new experiments on the star-graph $G(2, 20)$ to stress-test long context planning capabilities, where the LLM has to perform 20 backup reasoning steps to solve the task. We observe that Q# solves $G(2, 20)$ with nearly 100% accuracy while all policy-based post-training approaches get stuck at 50% accuracy and fail to fix the short-cut, which matches our existing results on $G(2,5),G(5,5)$ and $G(3,8)$ in the paper. We will add these new results to the final version of the paper.

Can the proposed approach be extended to stochastic MDPs?

Authors' reply: In this work, we focus on deterministic MDPs since the $Q^{\star,\eta}$ is a simple functional of the cumulative rewards distribution $Z^{\pi^{ref}}$, and that they capture the LLM post-training setting very well. If the MDP has stochastic transitions, then $Q^{\star,\eta}$ won't be a functional of $Z^{\pi^{ref}}$ (i.e. Thm 2.2 won't hold) and we will need to learn $Q^{\star,\eta}$ via temporal difference learning. Theoretically, this will require the bellman completeness assumption, which is much stronger than our realizability. Practically, TD-learning with function approximation is much less stable than online regression. Thus, while Q# can be extended to stochastic MDPs via TD-learning, this direction presents additional challenges and is less directly relevant to the LLM post-training application we focus on.

Additional Citations:

[Nachum'19a] AlgaeDICE: Policy Gradient from Arbitrary Experience.

[Nachum'19b] DualDICE: Behavior-agnostic Estimation of Discounted Stationary Distribution Corrections.

[Chang’22] Learning bellman complete representations for offline policy evaluation.

[Pang’22] QuALITY: Question Answering with Long Input Texts, Yes!