$Q\sharp$: Provably Optimal Distributional RL for LLM Post-Training
Authors: Jin Peng Zhou
Kaiwen Wang
Jonathan Daniel Chang
Zhaolin Gao
Nathan Kallus
Kilian Q Weinberger
Kianté Brantley
Wen Sun
Venue: NeurIPS 2025 poster ()
Keywords: large language models
distributional RL
reasoning
Primary area: reinforcement_learning
TL;DR: We introduce a theoretically-grounded distributional RL algorithm for LLM post-training that demonstrates improvement upon prior work on both synthetic and mathematical reasoning tasks.

Abstract: 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 \url{https://github.com/jinpz/q_sharp}.

Q#: Provably Optimal Distributional RL for LLM Post-Training
Jin Peng Zhou" Kaiwen Wang" Jonathan Chang' Zhaolin Gaol, Nathan Kallus Kilian Q. Weinberger', Kianté Brantley', and Wen Sun
Commell University Harvard University Netffix *Databricks
Abstract
Reinforcement learning (RL) post-training is crucial for LLM alignment and reasoning, but existing policy-based methods, such PPO and DPO, fall short fixing shortcuts inherited from pre-training. In this work, we introduce Qr, 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 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, QE outperforms prior baselines in math reasoning benchmarks while maintaining smaller KL divergence the reference policy. Theoretically, we establish 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 QF as an effective approach for post-training LLMs, offering both improved performance and theoretical guarantees. The code can be found https: //g thub. com/ jtnp//ijn////// q sharp.
aggregated online dataset. Unlike prior
Introduction
Reinforcement learning (RL) post-training is crucial step in training large language models (LLMs), aligning their generations with human preferences [1] and enhancing their reasoning capabilities [2, 3]. This stage typically follows supervised learning (next-token prediction), where the model further trained to maximize expected cumulative reward while minimizing KL divergence from the reference policy ne obtained via supervised learning. The KL penalty plays critical role by keeping the model close to π rof mitigating issues such as reward hacking and catastrophic forgetting.
Most state-of-ftt oh-the-rrt LLMs [4, 5, 6] are post- trainer using policy-based RL methods, which update model weights via stochastic gradient descent using algorithms like RLOO [7], PPO [8], and DPO [9]. However, these methods are computationally expensive, requiring full backpropagation through the LLM during training. In this paper, we propose a more efficient alternative: a value-based RL approach that guides the generations of the reference policy nf using learned value function, without modifying N- model weights. This approach is particularly attractive because, for many use much smaller models to learn value functions for guidance. For instance, in our experiments (Section 3.2), show that 1B parameter value model can effectively steer and improve 70B parameter LLM.
tasks, evaluating generations is casier than producing them [4. 10], suggesting
Existing vasue- methods for LLM post-training, such as CD [11] and VAS [12], fall short of faithfully optimizing the Kl--constrained RI. objective. These approaches guide ge using Q=" -th
'qqual contribution. Correspondence kr24))skkk37) edu.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Reward-to-go (xc -)
Algorithm Sketch
Reference
Representative Results
Rowaighted Policies
distribution
Reward distribution
poticy
ce icursi
function policy
8:000)
Daxeline
16
70
Collect data sward-te--oo distribution
Diarrgnnce
8.0
6.6.
Sours)
Bussline
(our))
Roward-o-gg
(Gr,v) R Stra [exp(=/1]] (-xmssllex(((///)
1) oe (1)(() (F))
Raceline
5
Figure I: (Left) A sketch of our post-training algorithm (Qz) based on distributional RL. Qi alternates between learning Z the reward-to-ge distribution of ग and using the induced policy to collect new data and further improve the distributional estimate. (Right) Evaluation result on the GSMSK dataset [19]. We see that QF achieves both higher accuracy and lower KL compared to prior value- algorithms [11, 12].
based
expected reward-to-ge under -d without KL regularization does not guarantee convergence to the optimal policy π='", In contrast, under the classical KI KL-regularized RL framework, we show that is provably optimal to guide ne using Q*.", the expected reward-to-go under the optimal policy ,,, which accounts for KL. regularization. This theoretical insight ensures convergence to n -" and addresses the shortcomings of previous methods. As we demonstrate empirically and theoretically, prior approaches can lead to suboptimal rewards or large KL divergence- -issues that our algorithm, Qi, provably avoids.
Our method exploits special properties of Q*.*" in deterministic MDPs and iteratively trains a model estimate through supervised distributional learning such MLE. The iterative training procedure is motivated by the classic imitation learning algorithm DAgger [13], which addresses covariate shift and ensures that the learned Q*1 estimator remains accurate when used to guide go at inference time. This distributional learning approach not only enhances empirical performance but also enables second-order style regret bounds- instance-dependent bounds that adapt to the variance of the model's generation.
Qi differs from traditional RL methods in two key aspects. First, we avoid complex temporal dif- ference (TD) learning [14] or Q-learning techniques [15, 16), instead relying on direct supervised learning of fixed critic. Second, while we adopt a distributional perspective, Q# conceptually sim- pler than classical distributional RL algorithms like C51 [17]: we directly learn outcome distributions via supervised maximum likelihood, without invoking distributional Bellman updates. We elaborate on this and related works in Appendix A. In summary, our contributions are as follows:
1. We propose Qi. principled algorithm for KL-regularized RL in deterministic MDPs, which includes LLMs, based on guiding re with the soft Q" learned with distributional RL (Section 2.2).
2. We prove siraaceedeeendeen PAC bounds for convergence to the optimal policy, which only requires realizability in the function class (Section 4).
3. We show that value-based post-training, which includes Qt, can fix biases and shortcuts in a star-graph environment [18], while popular policy-based methods cannot (Section 3.1).
4. We provide extensive experiments on math reasoning tasks that validate the effectiveness of our method at maximizing reward while maintaining small KL deviations from the reference policy (Section 3.2).
2 Method
2.1
Preliminaries
We study KL-regularized reinforcement learning (RL) in deterministic Markov Decision Processes (MDPs), where large language model (LLM) post-training a motivating special case. An MDP is defined by state space X, action space ) horizon H, transition kernels (Pi,... PT)) with PA. X x y > A(X)) and known reward functions .......,m where rh X' x y - TR. policy (TT. (T...... TH) consists of decision rules na X (()). For given 7 > 0, the KL-regularized

---

value of a policy is defined as
V = B [E (2)), - (TTL((T(
mh (xa))
(I)
classical result shows that Kl.-regularized RI. can be solved via soft Bellman equations [20]. Starting from h = H and proceeding backward, we define:
H+1 (1) = 0, "A T (y 2) x o (y exp((i "Q"" Vy =n Bywattt) exp(1) (x, 1)).
@x v) - (x, + Be'n (x h+1 (x')],
(2,
(2)
of Moreover, 0% (x: y) This expresses the optimal policy as softmax over 0 weighted by " is the maximal expected KL-regularized return starting from (z,, y) at time h. We now focus on deterministic MDPs, which covers LLM post-training and other structured generation tasks such as diffusion models [21].
Assumption 2.1. The transitions Pa are deterministic.
Under this assumption, the value function simplifies signifieantly: "o-
BByway (w)) (exp(1) "ra(s,)) + +1 (x')) = Eg- (ex((1 Σ1zA re( M+)) | J. = x),
where Equation (3) is due to the determinism of Ph. while Equation (4) follows by recursively unrolling until the final step. Note that although VAt (2%) corresponds to the soft value of the optimal policy, its recursion expressed via expectations over मम We summarize this in the following known result [22, 23, 21]:
Theorem 2.2. Under Assumption 2.1 we have (xn) = 17 In 'Eex Fe 9c)) | an.]] and (Ir,, yn. In E_.<(ex((() Σizn Fe( v:)) TA, Uh).
This shows V and Q are simple functionals of Z the cumulative reward distribution of ne where the functional is fPP) = 17 In Ep exp(X/n). In other words, if we learn the cumulative reward distribution of -ee then we can directly compute V** and Q*"7, without any dynamic programming.
This offers several benefits. First, we do not require temporal difference (TD) learning (i.e.. boot- strapping) which notoriously unstable with deep networks [24] and requires completeness-type assumptions to guarantee convergence in theory [25]. Second, fitting the reward-to-ge distribution Z" or regressing Eg (ep(() Eizh Tt]] a standard supervised learning task with a fixed target, which is much more stable in practice and well-understood in theory. Notably, there is no bootstrap- ping or changing targets which is what renders deep RL fragile. Third, we can apply distributional RL methods, where we directly fit the distribution Z via supervised learning (e.g., maximum likelihood). Importantly, our approach does not involve distributional Bellman equation nor distributional TD update, which are known to be contrratract under certain metrics [17]. Prior work has shown that fitting Z* in this manner yields benefits in representation learning [17, 26], lower variance updates [27], and second-order bounds [28, 29].
Applicability to LLMs. Our deterministic MDP framework directly models LLM post-training as a special case [4]. The initial state I1 corresponds the input prompt, each intermediate state 27, is the current generation prefix, and the action yn. is the next token. The policy thus reflects the LLM's autoregressive decoding process. The transition function deterministic: Ps. (2 yn) Thyii. which simply appends the new token to the prefix. In many post-training settings, the reward is sparse, meaning only FH nonzero. Under this assumption, Theorem 2.2 simplifies "","((n (ar,, yh.) = n In (ex((()(1) (exp((x=x((xx rr yn)) was yr). For example, the reward may indicate solution correctness in math tasks or reflect user preference in dialogue, as determined by learned reward model.
Inference with cumulative reward distribution. Let 2* denote the conditional distribution over cumulative rewards under rollouts from ,e that is, Z v) D Σizh re( Bt) .y I, 2hh
y, where the trajectory (Fh, Wh..... ... TH. yı1) is sampled under x ne and 2 denotes equality in distribution. Combining Theorem 2.2 and Equation (2), the optimal policy can be rewritten in terms
3

---

of Z" as , x)) & mh (y 2)B==~z; (2.w) (xxp(//n). This motivates defining a general family of policies induced by any distribution Z : x y 7 A (R) via
hh
7. (1 x) ( (x)))=2 (exp
(5)
2T²", we can approximate the optimal policy by estimating Z' with 2 = Z using distributional learning techniques such as maximum likelihood estimation (MLE), and then instantiating n 2..1 This forms the core of our proposed Q: algorithm.
Since π²10
2.2
Algorithm QF
We propose Q-Sharp (Qt), distributional value- algorithm for KL-regularized RL in deter- ministic MDPs. QF iteratively collects data from progressively improved policies to approximate the target distribution Z* (Algorithm 1). In this section, we describe Qt in practical terms for deep neural networks and LLMs; in Section 4, we formalize it using online learning oracles and prove convergence under a mild realizability assumption.
Let ZÍ : X'x'' A(R) denote a paramet-
Algorithm Q1
ric conditional distribution with parameters A. Given sample R E TR (e.g,, drawn from Z*) and a model prediction Z, let L(R, 2) be distributional loss for training the model. We denote by 0 the parameter that minimizes the distance between Z and Z For example, if Z (zz, y) is Ber P (z. y)), we can parameter- ize Z (z, y) by neural network that outputs a scalar estimate $ of p (a, y). The natural loss in this case is binary cross-entropy (BCE): r In (1 r) In(1 P)- This binary setup is appropriate for tasks such math or multiple-choice questions where the reward is binary. If the reward distribution has no known parametric form, one can use a paramaraettric model (e-g., a histogram that discretizes the reward space) trained via max- imum likelihood (MLE) [17]: L (r, é) = In 2 fidx(r)], where idx(r) returns the in- dex of the bin containing r, and 2 denotes the probability estimate for bin i. In general, Qi can incorporate any distributional RL loss function [30]. Once 2o closely approximates Z", instantiate near-optimal policy n. via Equation (5). In Section 4, prove that this procedure converges to the optimal policy under a mild realizability assumption.
1: Input: reference policy x=
2: Initialize parameters 0 of conditional distribu- tion 7 : x - A RR) and dataset Dn, = for all h.
3: for = 1, -1,2,... 2,... until convergence do Let π² f Z2 be policy induced by
4:
Zar (using Equation (5)). for 1,2..... N do
5:
6:
Sample a switching time h ~ [H]. Roll-in with x² for h - steps. Resume trajectory with re from Ti-
7:
Lice (r. p)
8:
9:
Let R2 denote cumulative rewards after
time
Add (It, Mt, R2) to D., V h. 11: end for
10:
12: Update H by minimizing the distribu- tional loss on the aggregated data:
gh+1 +- arg ming Σn ED, [C(Rn, 7 (9 B-))
13:
end for Output: Final 0*
14:
Then, the key idea of Qt is iterative data-collection and update process. At iteration k, with current parameters 2² we deploy the induced policy * := = to gather new data. Specifically, we roll in with x for h steps to reach a state 2h. then switch re complete the trajectory. The cumulative reward from step h to the end, denoted RA,K, is a sample from Zj (r,). We add these samples to the dataset and update 0 via gradient descent on the distributional loss. This process repeats until convergence.
Our iterative approach similar in spirit DAgger [13], Aggre [31, 32], and RLGF [33], which likewise mitigate distribution shift to ensure the learned estimator remains accurate at test time. In contrast, prior value-based methods such as CD [11] and autropy-regularize PRM [34] train their estimators only on data from re While such an estimator may perform well on n refss 's distribution, it offers no guarantee of accuracy when used steer fffs generation inference time.
Comparison with CD and VAS. The most closely related value-based baselines are CD [11] and VAS [12], yet they exhibit three critical limitations. (i) Incorrect value target. Both methods re-weight ref using @xt unregularized Q Q-function of oo thereby ignoring the KI. term, shown in Section 4, this choice can yield policies that are either sub-optimal in reward or far from re Qt

---

Pro-Train
REINFORCE RMO DR (ours]
100
Path
Goal
Easy Takens
Accurscy [%)
prompt:
"edge set 10
60
correct response: "31, 4, 9, '1"
40
shortcut response: "31, 1, 20, 17,
X
Difficult Token
20
Graphs
(a) The star-graph G(5,5).
(b) Generalization accuracies.
Figure 2: (a) The star-graph with degree d and path length & = 5. Next-token prediction, the standard pre- pre-training loss, learns sub-optimal shortcut that selects random first node and follows it until the end. (b) Accuracies on held-out test graphs for various models. Policy- baseed post-training methods such REINFORCE and RPO (a variant of DPO) still exhibit the shorteut and has test accuracy 1/d, while our value-based Qt fixes the shortcut and achieves near-perfect accuracy.
instead employs the principled target Q* and is guaranteed converge to n* under realizability. (ii) Offline training. CD and VAS fit their value functions on a fixed dataset, whereas Q= alternates data collection and updates, improving robustness distribution shift [13, 31]. (iii) Squared-loss regression. Both baselines learn Q= with an l loss, implicitly assuming homoskedastic Gaussian rewards. Q7 leverages distributional RL losses, which are theoretically more sample-efficient [35, 29] and empirically superior [17, 26).
Relation to actor-critic methods. Although Q5 learns a value function, its target V*. (or Q*:?) is fixed throughout training. Standard actor-critic algorithms PPO) continuously update V or QT as т evolves, and rely on bootstrap-basse TD updates. In contrast, Q: trains the value network via distributional supervised learning (e.g., MLE), thereby avoiding the instability of changing targets.
Relation to DPO [9]. While the form of Equation 5 resembles DPO's policy expression, their derivations and scopes fundamentally different. DPO begins from the same KL-regularized RL objective but, without exploiting the deterministic transition structure, operates at the sequence level, corresponding to the one-step
(H 1). Its policy is given by (r x)) Z(#) n n
a:) cxp (22, where y denotes full completion and (ar) is the partition function over all possible sequences. When H 1, our Q* reduces to the reward r, and the DPO expression naturally follows as a special case. However, the DPO partition function (xx) is intractable to normalize, and practical implementations must rely on pairwise preference data (e.g., Bradley-Terry modeling) to bypass it.
Inference with multiple Because the learned distribution 20 is independent of 1). single trained network can support any choice 17 at inference time simply by plugging it into Equation (5).
Experiments
3.1 Star-Graph
We begin with the star-graph task from [18], illustrated in Figure 2(a). A star-graph G(d, E) has paths of length é from central node. Given a start/goal node and the graph edges, the LM must generate a valid path. Though seemingly simple, [18] showed that next-token pre-training often learns faulty shortcut: the model picks the first node at random (correct with probability 1/d) and follows the path, yielding a test accuracy of only 1/d. This highlights the limitations of next-token prediction on planning tasks. [36] also showed that the task embeds the "sparse parity" problem determining whether the sum gradient-basse based optimizers and is widely studied in learning theory and optimization [37, 38, 39, 40].
binary string is even or odd which is known to be difficult for
Can this shortcut be fixed during post-training? We evaluate REINFORCE [41], DPO [9], RPO [42], and Qt. reporting test accuracies in Figure 2 (b). Qi consistently corrects the shorteut, achieving near-perfect accuracy, even for long paths (G(2, 20)) or large degrees (G(5, 5)). CD [11] achieves similar performance Qt. In contrast, policy-baser methods like REINFORCE and RPO fail to
5

---

Table I: Comparison of QF with गh and CD baseline on GSM8K (Left) and MATH (Right). For both Llama 3 and Llama 3.1 8B, Qi consistently improves both pass@l and majority voting accuracy upon baselines while incurring minimal KI. deviation.
x
Llama 3 8B
Llama 3.1 8B
Llama 3 8B
Llama 3.1 8B
Methods
xtd
CD
QF
gnd
CD
Q1
Methods
x-
CD
Q11
CD
QF
paaa01
69.1
77.8
78.4
82.9
84.5
85.1
paas01
25.4
24.9
27.1
43.9
45.3
46.7
maj108 1081
85.8
87.2
88.1
90.5
90.9
91.4
maj1e8
34.3
34.3
37.9
57.0
59.0
60.1
KL-Divergence
6.39
2.65
7.43
3.67
KL-Divergence
15.27
7.14
26.8
8.69
fix the shortcut, plateauing at 1/d accuracy. DPO performs worst, often collapsing the policy to zero accuracy by suppressing both chosen and rejected paths-a failure mode also noted by RPO. These results suggest that once shorteuts are learned, policy-hased methods struggle to unlearn them, reinforcing the effectiveness of value-based approaches like Q1 and CD for LLM post-training. Please see Appendix C for a more detailed discussion on why REINFORCE and RPO cannot fix shortcuts and implementation details.
3.2 Math Reasoning
Datasets. We evaluate on two mathematical reasoning benchmarks: GSM8K [19], dataset of grade school arithmetic word problems, and MATH [43], which features more challenging high school competition problems. We split each training set 90%-10%% for training and validation. Test performance is reported on the full GSM8K test set and a 500-sample subset of MATH (MATH-500), following prior work [44, 45]. In Appendix G. we also evaluate Q$ on AIME-24 24 dataset.
Models. We experiment with Llama 3 [5] and Qwen 2.5 [46] model families, both of which are competitive math reasoning tasks and span wide range of parameter scales. Due to space constraints, we report results for Llama 3 in the main text and defer Qwen 2.5 results to Appendix G. Unless otherwise noted, the Q function in Q1 is parameterized and initialized with Llama 3.2 IB model, and we use 77 = 0.1, which yields consistent and strong performance. We run Q# for two iterations, after which performance converges. Additional details on model configurations and QF training are provided in Appendices D and E.
Evaluation metrics. We report single sample accuracy (pass01) and majority voting accuracy (maj10k) pass@1 evaluates one sampled generation per problem against the ground truth, while maj jajk@k checks if the most frequent answer among k samples is correct. We use k = 8, temperature T 0.8, and nucleus sampling p = 0.9. The evaluation prompt template is provided in Appendix
Main results. Table presents Qi performance on GSM8K (Left) and MATH (Right) with ,e as Llama 3 or 3.1 8B. Although both have 8B parameters, Llama 3.1 performs significantly better. Across all settings, Qi consistently improves over ne boosting pass@1 by up to 9% on GSM8K with just IB additional parameters. We also compare against the CD baseline [11], which incorrectly uses QT .0 to guide π| Qi outperforms CD on both accuracy metrics while maintaining lower KL divergence. Overall, Q: Pareto-dominates CD in the KL-regularized RL setting by achieving higher reward and lower KL. We note that CD [11] and VAS [12] are concurrent work and differ only in minor aspects such sampling strategy. Therefore, we use CD canonical baseline for empirical comparison. Since Qi is complementary to policy-based methods, we further evaluate its effectiveness when guiding a PPO-trained model, as shown in Appendix I.
Larger nd and QE sizes. We evaluate how performance scales with larger rf and Qi models on MATH (Table 2). Using 70B Ilama 3 and 3.1 as ae significantly boosts baseline pass@1 (45.6% and 60.6%, respectively). Remarkably, a IB Q1 still improves these large models c.g. by 2.5% pass@1 and 3.5% maj 108 for Llama 3.1. Increasing Qt to 3B yiclds further gains, demonstrating scalability, Compared to Table (right), we note that with 9B total parameters (8B पम + IB Qt), the majmmm 108 accuracy already matches the pass01 of the 70B ne in Table 2, suggesting a promising low-resource alternative. For Llama 3, pass@1 1 improves while maj108 slightly drops, likely due to increased generation diversity benefiting harder problems but redueing consistency easier ones.
Qi as reward model. Beyond guiding मd generation, Qt's token-level Q function can also assess how good a complete generation is among many. We compute @(generation, EOS) by applying Q: as a reward model (Qi-RM) on GSM8K and MATH, using both ne and Qi generations. Table reports two settings: Qt-RM Best of 8 (selects top-scoring sample) and Q maj 168 (aggregates

---

Table 2: Performance of π re and Q2 on MATH with larger N/ and Q model sizes. Q of size IB is capable of guiding 70B x model. Increasing Qi model sizes to 3B also leads to noticeably better performance for Llama 3.1 70B.
x
Llama 3 70B
Llama 3.1 70B
Q: Model
None
Llama 3.2 IB
Llama 3.2 3B
None
Llama 3.2 IB
Llama 3.2 3B
pase@1 f
45.6
46.4
46.7
60.6
63.1
64.1
ma.j 128 f
55.6
55.5
55.3
69.0
72.5
72.7
XI-Divergence
3.12
5.15
4.98
4.99
@ass@1 vs. KL-Divergence
vs. KL-Divergence
CD
0.80
0.75
pass@1
0.70
10a
0.65
0.60
0.55
0.50
10
0#
10a KL-Divergence
10
10a KL-Divergence
10
Figure 3: Performance tradeoff between CD and QF on the GSM8K validation set. (Left) pass@1 vs. KL divergence: QF outperforms CD with higher accuracy and lower KL. (Right) KL divergence across vf. CD's KL increases rapidly and deviates from ,t while Qi remains stable.
majority voting with scores). Qi-RM maj108 consistently improves over vanilla maj1@8 and Best of 8 yields more than 10% gains over pass@1 for nd The reward model can be used on both zel and Qł own generations to further improve performance, which suggests the (same) reward model has generalizability for evaluating diverse generations.
Effect of 17. Figure shows the performance- KL tradeoff between CD and Q1 on the GSMSK validation set. (Left) Increasing KL can improve pass01 for both methods, but Qt consistently achieves a better Pareto frontier. (Right) CD is highly sensitive to ry: as v increases, its KL grows rapidly and performance degrades below that of nd In contrast, Q$ remains stable and requires less tuning of 17.
We ablate several design choices in Table on the GSM8K and MATH validation sets using pass@1 accuracy. The "Prefix" column tests training on all t 2 h prefixes after switching to re (Algorithm 1, Line 10), opposed to only f h. Though this breaks IIDD assumptions, the increased training data improves QF performance by up to 4%. We compare two parameterizations of Q*."; Q-type, which computes Q*."(z, (z. y) for all y, and V-type, which predicts Q=' 9) for specific y. V-type outperforms Q-type, likely due to its lower parameter count and per-token computation. Details in Appendix D. We also compare distributional Qi with MSE-based regression, which underperforms as expected under Bernoulli rewards. Finally, more iterations of Algorithm yield marginal gains, with performance saturating after two iterations, which we adopt by default.
Ablations.
Qualitative comparison. Figure shows side-by-side generations from -h and Qi on math reasoning tasks. While both models often begin with similar KL deviation Q1 typically corrects re's mistakes and produces more coherent reasoning. This behavior reflects QF 's ability to assign higher value correct tokens, thereby steering generation more effectively at critical decision points. Additional examples are provided in Appendix K.
with Qf's low
Beyond math reasoning. To further validate the generality of Qi beyond mathematical reason- ing tasks, we evaluate its performance on QuALITY [47], a challenging multiple-choice reading comprehension benchmark with long-form passages drawn from Project Gutenberg. As shown in
7

---

Table 3: Performance of x- and Q1 with Qi reward model. The reward model de- termine the best generation among all genera- tions for problem and consistently improves maj 128 for π get and Qi own generations.
Table 4: Ablations of Q (last row) on pass@1 with various configurations the validation set GSM8K and MATH. The improvement suggests that our design choices all contribute positively to the final performance.
Setting
GSM8K
Llama
Type
Oye.
GSMMK
LJama MATH
Methods
n
Q:
x"
Q:
Single
Dist
80.5
64.5
pacsel
All
Q
Dist
81.4
66.4
naj1ea
85.8
88. 1
57.0
60.1
All
V
MSE
81.4
65.4
QI-RM Best 8
85.9
86.0
54.0
54.0
V
Dist
82.3
67.4
(2p 108
Appendix H, Table 6, we compare Qi with nd and CD baseline using Qwen 2.5 and Llama 3.1.
Specifically, Qwen 2.5 1B guides Qwen 2.5 7B and Llama 3.2 1B guides Llama 3.1 8B. Across both architectures, Qt consistently improves upon प in all evaluation metrics, demonstrating its robustness beyond the mathematical domain.
Theory
4.1
CD VAS are sub-optimal for KI.rregularized RI.
First, CD and VAS both propose to reweight (. [x) with the unregularized Q-function of nt
CD x)) x Tf x) exp(Q" (z,v)/n)
(6)
where recall that (m,, Unh) " = E Tt In ynl. Comparing with Equation (2), we can already see that CD, does not match the optimal policy "1,,, as QT" can be arbitrarily far from Q In particular, CD may fail to optimize the KL-regularized RL objective and exhibit two failure cases, which we demonstrate with simple MDP in Figure 4. First, we show that CD fails to maximize expected reward this MDP, even as the KL-regularizer yj decays to zero.
Theorem 4.1. Under Figure 4, CD learns to always select the left sub-tree as n - 0, which gives sub-optimal reward of 0.1, while x learns to always select the right sub-tree and chooses the path that gives reward
Proof. First, for CD, we have QT" (31, 0.1 and Q (51, RR) 0.05. Hence, CD's probabil- PL cxp((01/v))
ity of selecting the left sub-tree is pz (tt.01//m)ppexpppoooo
which converges to as - 0. Next, nj In exp(1/n) + 0.95). Hence, Qf's probability of selecting the left sub-tree is (0 0000 which converges to as + 0. Thus, CD learns the sub-optimal path.
for Q4, have Q*''(z1, al) = 0.1 and Q0*"
Pa. exp(0.117)
Next, show that CD also incurs higher KL than Q1.
Zi 5(0.))
Z Ber(0.05)
Theorem 4.2. Under Figure 4, CD's KL con- verges to In /P)) while Qi's KL converges to Im(1/pR) as 17 - 0. Thus if PL < Pn. CD converges higher KL than Q1.
Pt
PR
L
JR
0.05
0.95
Proof. As shown in Theorem 4.1, CD learns to select the left sub-tree while Q# learns to select the right sub-tree as 17 + 0. Then, the KLs simply follow by definition.
== 0.1
T=1
r
Figure 4: A tree MDP where edges are labeled with refts action probability. ge goes to the left sub-tree w.p. PL and the right sub-tree w.p. Pa. where Pt, PR > O. The left sub-tree gives r = 0.1 w.p. 1. In the right sub-tree, gee chooses reward I w.p. 0.05 and chooses reward 0 w.p. 0.95.
In sum, we proved that Figure 4, CD both incurs higher KL and achieves a lower sub-optimal reward compared to Qi- Thus, Qt generally Pareto-dominates CD the reward-KI. trade- off, which matches our empirical findings.

---

4.2 Performance Guarantee for Q
We prove that the learned policy by Qi guaranteed converge to the optimal policy with enough samples. This result holds in rich-observation MDPs where the size of the state space can be exponentially large or infinite, long as the mild realizability assumption holds.
To setup, let F. be a distributional function class for modeling Z', -the reward-to-go distribution under =d Each element of F has type f (J1: -..-.... fH) and fn X' y A( Vmax)). For the purpose of analysis, we assume access to a no-regret online learning oracle for the maximum likelihood (MLE) loss, which proceeds follows: for each iteration
1,2, ..,,, K, given any
(Kh,kr kb. RA. h==1 the oracle outputs Z₁ E F s.t.
E (log Zz (Ri,k Thik: Vh,,k) ) log 2h,,(RRs, Thikk vh,,k) S Regmte
Here, Regmie (K) denotes the cumulative regret the MLE oracle after iterations. No-regret online learning is well-studied in the literature [48, 49] and is standard tool when reducing decision making to supervised learning [13, 50, 35]. For example, if J is finite and satisfies realizability, then Vovk's aggregating algorithm ensures that Reg (K) < In((Fl) [51].2
Assumption 4.3 (Realizability). Z' F.
The following algorithm a slightly modified version of Algorithm I amenable for theoretical analysis. The only differences with Algorithm are: (1) we use the MLE oracle to learn Zz. and (2) for purpose of local exploration, we play a random action at the switching time h before following re to the end of the trajectory [31].
We now state main PAC bound for Q#.
Theorem 4.4. Fix any E (0, Vimax] and E (0, 1). Under Assumptions 2.1 and 4.3, Algorithm ensures w.p. at least 1 - 6, setting 3 ln(1/6) Regmte (K)). we have
K E (V*M - w** S AV (v h= ER=
B max Es B),
where CV 11 (x, = Exr Pan -nif(A) Var(exp( (2A )/1)) (F
is the coefficient of variation of exp((Z 9h)) and En == (exp((Vmmn Q7 (In, M1 (m)ll is the envelope of exp(() (Vmax - (In, /n)) both under ne
(
We highlight this applies to rich- observation MDPs where our only re- quirement for F is realizability. Our bound only scales with the function class's complexity, i.e., In((FF), and does not contain structural complexity measures. In contrast, prior bounds in RL theory require stronger assumptions such as Bellman completeness [52, 53, 54, 55, 56, 57, 29], even in deterministic MDPs [58], and/or scale with structural complexity measures such as coverabil- ity [59, 60]. eluder dimension [61, 55), and certain rank related complexity mea- sures [62, 63, 64].
Algorithm Qi (Theory Version)
1: Input: reference rf iteration count K, regularizer 1).
2: Initialize 2, randomly.
3: for k = 1, ........ K do
4: Let Th + π 7n.9
5: for step
1.2..... H do Roll-in with x² for h steps and see Th.c Play random action yo, and transit to a 1.k.
6:
7:
8:
Resume trajectory with rf from # 1,k-
9:
Let RA, be cumulative rewards after time h.
end for
11: Input {TA.K, Mr.k, RA. heju to MLE oracle.
12: Receive Z% from MLE oracle.
13: end for Output: 2,...,2. 2x.
Computational efficiency. Algorithm is model-free and computationally effi- cient. In contrast, prior model-fres algorithms for rich-observation MDPs perform exploration with version spaces and are computationally hard [62, 65, 55, 59, 29]. Thus, Theorem 4.4 shows that Algorithm 2 achieves both statistical and computational efficiency under mild assumptions by simply operating within the Kl-regularieed RL framework, which is of great relevance for post-training. We remark that [66] observed similar benefits in offline RL while we study the harder online setting.
14:
Suppose rewards-tor to-go under गh in [0, Vmmax] w.p. 1.
aa S 6 is short for S Cb for universal constant C.

---

Second-order guarantee. Thanks to the distributional perspective, Theorem 4.4 is second-order bound [28, 29). Its leading term O(y Σm= E 1 v)) aggregates coefficients of variation.
In the worst case this is O(V Σx= K but when 2 has small or zero variance the term vanishes, leaving the lower-order O (maxne]11 E7, In(K)), logarithmical in K. Thus the bound adaptively improves from (VVK) instance-dependent: when Q* = V it equals 1, climinating the exponential dependence on In general, we can tolerate an 17 that is as small as the worst V Q under rollouts from n- which is reminiscent of the condition required for first-order or small-loss bounds [50, 35, 57J.
((In(K) in benigh instances. Interestingly, the envelope term EA, is also
Bernoulli rewards simplification. For closed-ended tasks (e.g. math or multiple-choice) the reward- to-go (z, y) is Bernouli,, (z. y) ~ (en(pk(z; y)). Then the CV term can be bounded by CV S Sugi
Pw)/Ph and the envelope term becomes [1/palle___)ee)) (wr1), which notably does not have exponential dependence on 1 1/n. Thus, as long as the reference model g has sufficient probability of solving the problem, our bound can be made independent of n7. Finally, we note that the ceaabuttoona-realizabbii only parameter of a Bernoulli distribution is its mean; also the MLE loss reduces to the binary cross-entropy loss [50, 57]. We present the corollary below and the proof in Appendix B.I.
condition can also be weakened mean- recan-realizabiiity. since the
Corollary 4.5. Suppose reward-to- gos are Bernoulli: (z, v) Ber(pa Then, under the setup of Theorem 4.4, the bound can be simplified to:
Ek (V- V** S A(V E Ee un Unil((A)
1 1=pa(#A.20) Ph (mh
+ [H
(1 Il (statt - 8).
Modification for Regret Bound. It is possible to turn Theorem 4.4 into regret bound by replacing random action in Line 7 of Algorithm 2 with a no-regret contextual bandit oracle, where "context" is 24, action is yn and "reward" is Rg This is alike the steps needed to convert Aggre VaTe's PAC bound into a regret bound [31]. Our theory can be interpreted as a regret/PAC reduction from KL-regularized RL in deterministic MDPs no-regret online learning, which mirrors the type of imitation learning guarantees obtained for Aggre VaTe [31].
Remark:
5 Limitations & Conclusion
Our results focus on deterministic MDPs including LLM post-training, where the optimal action- -vauue Q** is a simple functional of the reference return distribution -=d and Theorem 2.2 applies directly. For domains with stochastic transitions such as multi-agent game playing where the next state depends on the (potentially unpredictable) behavior of the other player, Q* need to be learned via lemporal-diffrrent methods, which typically rely on the stronger Bellman--oomllttnnes assumption and may introduce additional training instability. In summary. QF offers a principled and practical avenue for post-training LLMs. It combines distributional-RI objective with supervised regression, enjoys provable convergence under mild assumptions, and consistently surpasses prior value- based baselines synthetic planning and math-reasoning benchmarks achieving higher accuracy while maintaining lower KL divergence from the reference policy.
Acknowledgment
JP7 is supported by a grant from the Natural Sciences and Engineering Research Council of Canada (NSERC) (567916). KW is supported by Google PhD Fellowship. 7G is supported by LinkedIn- Cornell Grant, Wen Sun is supported by NSF 11S-2154711, NSF CAREER 2339395 and DARPA LANCER: LeArning Network CybERagents. This research is also supported by grants from the National Science Foundation NSF (IIS- 1846210, IIS-2107161. and IIS- 17242222 HDR-21183100 the Cornell Center for Materials Research with funding from the NSF MRSEC program (DMR-1719875), DARPA, arXiy, Linkedln, Google, and the New York Presbyterian Hospital.
10

---

References
[1] Paul Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30, 2017.
[2] Amrith Sethur, Chirag Nagpal, Adam Fisch, Xinyang Geng, Jacob Eisenstein, Rishabh Agarwal, Alekh Agarwal, Jonathan Berant, and Aviral Kumar. Rewarding progress: Scaling automated process verifiers for Ilm reasoning. arXiv preprint arXiv:2410.08146, 2024.
[3] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang. Xiao Bi, et Deepseek-rl: Incentivizing reasoning capability in Iims via reinforcement learning. arXiv preprint arXiv:2501. 2025.
[4] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744, 2022.
[5] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Ilama 3 herd of models. arXiv preprint arXiv:240..21783, 2024.
[6] Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. Gemma 2: Improving open language models at a practical size. arXiv preprint vrXiv:2408.00115 2024.
[7] Wouter Kool, Herke van Hoof, and Max Welling. Buy 4 reinforce samples, get baseline for free! 2019.
[8] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
[9] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. Advances in Neural Information Processing Systems, 36, 2024.
[10] Jing-Cheng Pang, Pengyuan Wang, Kaiyuan Li, Xiong-Hul Chen, Jiacheng Xu, Zongzhang Zhang, and Yang Yu. Language model iafroimproveeent by reinforcement learning contempla- tion. arXiv preprint arXiv:2305. 14483, 2023.
[11] Sidharth Mudgal, Jong Lee, Harish Ganapathy, YaGuang Li, Tao Wang, Yanping Huang, Zhifeng Chen, Heng-Tze Tze Cheng, Michael Collins, Trevor Strohman, et al. Controlled decoding from language models. arXiv preprint arXiv:2310.1700 2023.
[12] Seungwook Han, Idan Shenfeld, Akash Srivastava, Yoon Kim, and Pulkit Agrawal. Value augmented sampling for language model alignment and personalization. arXiv preprint arXiv:2405.06633 2024.
[13] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth interna- tional conference artificial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011.
[14] Gerald Tesauro. Practical issues in temporal difference learning. Advances in neural information processing systems, 1991.
[15] Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-leaming, In Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
[16] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in neural information processing systems, 33:11799 1191, 2020.
11

---

[17] Mare G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforce- ment learning. In International conference on machine learning, pages 449-458 PMLR, 2017.
[18] Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. arXiv preprint arXiv:2403.0668 06966,, 2024.
[19] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110. 14168, 2021.
[20] Brian D Ziebart, Andrew I. Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pages 1433-1438. Chicago, II., USA, 2008.
[21] Carles Domingo--Eniicn Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprini arXiv:2409.08861, 2024.
[22] Alexandre Piché, Valentin Thomas, Cyril Ibrahim, Yoshua Bengio, and Chris Pal. Probabilistic planning with sequential monte carlo methods. In International Conference Learning Representations, 2018
[23] Xiner Li, Yulai Zhao, Chenyu Wang, Gabriele Scalia, Gokcen Eraslan, Surag Nair, Tom- maso Biancalani, Shuiwang Ji, Aviv Regev, Sergey Levine, et al. Derivative-free guidance in continuous and discrete diffusion models with soft value-based decoding. arXiv preprint arXiv:2408.088 08255, 2024.
[24] Hado Van Hasselt Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Modayil. Deep reinforcement learning and the deadly triad. arXiv preprint 2018.
[25] Rémi Munos and Csaba Szepesvári. Finite-time bounds for fitted value iteration. Journal of Machine Learning Research, 9(5), 2008
[26] Clare Lyle, Marc G Bellemare, and Pablo Samuel Castro. comparative analysis of expected and distributional reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 4504-45 2019.
[27] Mark Rowland, Yunhao Tang, Clare Lyle. Rémi Munos, Mare Bellemare, and Will Dab- ney. The statistical benefits of quantile temporal-din differecee learning for value estimation. In International Conference on Machine Learning, pages 19210-29231 PMLR, 2023.
[28] Kaiwen Wang, Nathan Kallus, and Wen Sun. The central role of the loss function in reinforce- ment learning. arXiv preprint arXiv:2409.12799 2024.
[29] Kaiwen Wang, Owen Oertell, Alekh Agarwal, Nathan Kallus, and Wen Sun. More benefits of being distributional: Second-order bounds for reinforcement learning. International Conference of Machine Learning, 2024.
[30] Marc G Bellemare, Will Dabney, and Mark Rowland, Distributional reinforcement learning. MIT Press, 2023.
[31] Stephane Ross and J Andrew Bagnell. Reinforcement and imitation learning via interactive no-regret learning. arXiv preprint arXiv: 1406.5997 2014.
[32] Wen Sun, Arun Venkatraman, Geoffrey Gordon, Byron Boots, and J Andrew Bagnell. Deeply aggrevated: Differentiable imitation learning for sequential prediction. In International confer- ence on machine learning, pages 3309-3318 PMLR, 2017.
[33] Jonathan D Chang, Kiante Brantley, Rajkumar Ramamurthy, Dipendra Misra, and Wen Sun. Learning to generate better than your 11m. arXiv preprint arXiv:2306. 20t.:2006.11166, 2023.
[34] Hanning Zhang, Pengcheng Wang, Shizhe Diao, Yong Lin, Rui Pan, Hanze Dong, Dylan Zhang, Pavlo Molchanov, and Tong Zhang. Entropy-regularizzed process reward model. arXiv preprint arXiv:2412.11000 2024.
12

---

[35] Kaiwen Wang, Kevin Zhou, Runzhe Wu, Nathan Kallus, and Wen Sun. The benefits of being distributional: Small-loss bounds for reinforcement learning. Advances in Neural Information Processing Systems, 36, 2023.
[36] Edward S Hu, Kwangjun Ahn, Qinghua Liu, Haoran Xu, Manan Tomar, Ada Langford, Dinesh Jayaraman, Alex Lamb, and John Langford. Learning to achieve goals with belief state transformers. arXiv preprint arXiv:2410.23500 2024.
[37] Shai Shalev-Shwartz, Ohad Shamir, and Shaked Shammah. Failures of gradient-based deep learning. In International Conference Machine Learning, pages 3067-3075. PMLR, 2017.
[38] Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750-21766 2022.
[39] Emmanuel Abbe and Colin Sandon. Polynomial-timg timoee universality and limitations of deep learning. Communications on Pure and Applied Mathematics, 76(1 1):3493 -3499, 2023.
[40] Yiwen Kou, Zixiang Chen, Quanquan Gu, and Sham Kakade. Matching the statistical query lower bound for k-sparse parity problems with sign stochastic gradient descent. Advances in Neural Information Prucessing Systems, 37: 13001-113037 2024.
[41] Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzich Fadace, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker. Back basics: Revisiting reinforce style optimization for learning from human feedback in llms. arXiv preprint 2402,44000 2024.
[42] Richard Yuanzhe Pang. Weizhe Yuan, Kyunghyun Cho, He He, Sainbayar Sukhbaatar, and Jason Weston. Iterative reasoning preference optimization. arXiv preprint 10Xvv22404.19733, 2024.
[43] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint urXiv:2103.0332 2021.
[44] Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. arXiv preprint 2tt:::330520050 2023.
[45] Peiyi Wang. Lei Li, Zhihong Shao, Runxin Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce Ilms step-by-step without human annotations. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1426-9439 2024.
[46] An Yang, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, al. Qwen2. 5 technical report. arXiv preprint arXiv:2412.15115, 2024.
[47] Richard Yuanzhe Pang, Alicia Parrish, Nitish Joshi, Nikita Nangia, Jason Phang, Angelica Chen, Vishakh Padmakumar, Johnny Ma, Jana Thompson, He He, et al. Quality: Question answering with long input texts, yes! arXiv preprint arXiv:2112.08008, 2021.
[48] Nicolo press, 2006.
and Gâbor Lugosi. Prediction, learning, and games. Cambridge university
[49] Francesco Orabona. modern introduction to online learning. arXiv preprint 1rX1v:1912.13213. 2019.
[50] Dylan J Foster and Akshay Krishnamurthy. Efficient first-order contextual bandits: Prediction, allocation, and triangular discrimination. Advances in Neural Information Processing Systems, 34:18907-18919, 2021.
[51] Vladimir G Vovk. A game of prediction with expert advice. In Proceedings of the eighth annual conference on Computational learning theory, pages 51-60, 1995.
13

---

[52] Jinglin Chen and Nan Jiang. Information thhooretic considerations in batch reinforcement learning. In International Conference on Machine Learning, pages 1042-1051. PMLR, 2019.
[53] Yuanhao Wang. Ruosong Wang, and Sham Kakade. An exponential lower bound for linearly realizable mdp with constant suboptimality gap. Advances in Neural Information Processing Systems, 34:9521-9533 2021.
[54] Dylan J Foster, Akshay Krishnamurthy, David Simehi-Levi, and Yunzong Xu. Offline rein- forcement learning: Fundamental barriers for value function approximation. arXiv preprint arXiv:21 10919, 2021.
[55] Chi Jin, Qinghua Liu, and Sobhan Miryoosefi. Bellman eluder dimension: New rich classes of rl problems, and sample-efficien algorithms. Advances in neural information processing systems, 34:1 14413406-13418, 2021.
[56] Jonathan Chang, Kaiwen Wang, Nathan Kallus, and Wen Sun. Learning bellman complete representations for offline policy evaluation. In International Conference on Machine Learning, pages 1938-29711 PMLR, 2022.
[57] Alex Ayoub, Kaiwen Wang, Vincent Liu, Samuel Robertson, James McInerney, Dawen Liang, Nathan Kallus, and Csaba Szepesvari. Switching the loss reduces the cost in batch reinforcement learning. In Forty-first International Conference on Machine Learning, 2024.
[58] Runzhe Wu, Ayush Sekhari, Akshay Krishnamurthy, and Wen Sun. Computationally efficient rl under linear bellman completeness for deterministic dynamics, arXiv preprint arXiv:2406. 11810, 2024.
[59] Tengyang Xie, Dylan Foster, Yu Bai, Nan Jiang, and Sham M Kakade. The role of coverage in online reinforcement learning. arXiv preprint wrXiv:22100041577 2022.
[60] Zakaria Mhammedi, Dylan J Foster, and Alexander Rakhlin. The power of resets in online rein- forcement learning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[61] Daniel Russo and Benjamin Van Roy. Eluder dimension and the sample complexity of optimistic exploration. Advances in Neural Information Processing Systems, 26, 2013.
[62] Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Contextual decision processes with low bellman rank are pac-Icarnable In International Conference Machine Learning, pages 1704-1713. PMLR, 2017.
[63] Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Model-based rl in contextual decision processes: Pac bounds and exponential improvements over model-free approaches. In Conference on learning theory, pages 2898-2933. PMLR, 2019.
[64] Simon Du, Sham Kakade, Jason Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, and Ruosong Wang. Bilinear classes: A structural framework for provable generalization in rl. In International Conference on Machine Learning, pages 2826-2836. PMLR, 2021.
[65] Christoph Dann, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert Schapire. On oracle-effcient pac rl with rich observations. Advances in neural information processing systems, 31, 2018.
[66] Masatoshi Uehara, Nathan Kallus, Jason D Lee, and Wen Sun. Offline minimax soft-q-learning under realizability and partial coverage. Advances in Neural Information Processing Systems, 36, 2023.
[67] Kevin Yang and Dan Klein. FUDGE: Controlled text generation with future discriminators. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani- Tur, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou, editors, Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human I Language Technologies, pages 3511 1-3535, Online, June 2021. Association for Computational Linguistics.
14

---

[68] Stephen Zhao, Rob Brekelmans, Alireza Makhzani, and Roger Grosse. Probabilistic inference in language models via twisted sequential monte carlo. arXiv preprint arXiv:240.117466 2024.
[69] Will Dabney, Mark Rowland, Mare Bellemare, and Rémi Munos. Distributional reinforce- ment learning with quantile regression. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
[70] Jesse Farebrother, Jordi Orbay, Quan Vuong, Adrien Ali Taīga, Yevgen Chebotar, Ted Xiao, Alex Irpan, Sergey Levine, Pablo Samuel Castro, Aleksandra Faust, et al. Stop regressing: Training value functions via classification for scalable deep rl. arXiv preprint wrXi:22403.09500 2024.
[71] Kaiwen Wang, Dawen Liang, Nathan Kallus, and Wen Sun. Risk-sensitive rl with optimized certainty equivalents via reduction to standard rl. arXiv preprint arXiv:2403.0633 2024.
[72] Ke Sun, Bei Jiang, and Linglong Kong. How does return distribution in distributional reinforce- ment learning help optimization? arXiv preprini arXiv:2209. 220::.453.. 2022.
[73] Ke Sun, Yingnan Zhao, Wulong Liu, Bei Jiang, and Linglong Kong. Distributional reinforcement learning with regularized wasserstein loss. Advances in Neural Information Processing Systems, 37:63184 632211, 2024.
[74] Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
[75] Alisa Liu, Xiaochuang Han, Yizhong Wang, Yulia Tsvetkov, Yejin Choi, and Noah A. Smith. Tuning language models by proxy. First Conference on Language Modeling, 2024.
[76] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pages 19274-19286. PMLR, 2023.
[77] Wei Xiong, Hanze Dong, Chenlu Ye, Han Zhong, Nan Jiang, and Tong Zhang. Gibbs sampling from human feedback: A provable kl-constrained framework for rlhf. CoRR, 2023.
[78] Richard Yuanzhe Pang, Weizhe Yuan. He He. Kyunghyun Cho, Sainbayar Sukhbaatar, and Jason Weston. Iterative reasoning preference optimization. Advances in Neural Information Processing Systems, 37:1 37:116617-116677,, 2024.
[79] Zhaolin Gao, Jonathan Chang, Wenhao Zhan, Owen Oertell, Gokul Swamy, Kianté Brantley, Thorsten Joachims, Drew Bagnell, Jason Lee, and Wen Sun. Rebel: Reinforcement learning via regressing relative rewards. Advances in Neural Information Processing Systems, 37:52354- 52400. 2025.
[80] Tengyang Xie, Dylan Foster, Akshay Krishnamurthy, Corby Rosset, Ahmed Awadallah, and Alexander Rakhlin. Exploratory preference optimization: Harnessing implicit qt-approximation for sample-efficient rlhf. arXiv preprint arXiv:2403. 2004, 2024.
[81] Jae Hyeon Cho, Minkyung Park, and Byung-Jun Lee. Vpo: Leveraging the number of votes in preference optimization. arXiv preprint arXiv:2410.22891, 2024.
[82] Shenao Zhang, Donghan Yu, Hiteshi Sharma, Han Zhong, Zhihan Liu. Ziyi Yang, Shuohang Wang, Hany Hassan, and Zhaoran Wang. Self-exploring language models: Active preference elicitation for online alignment. arXiv preprint 2405 /9332, 2024.
[83] Roy Fox, Ari Pakman, and Naftali Tishby. Taming the noise in reinforcement learning via soft updates. arXiv preprint arXiv: 12.08662, 2015.
[84] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[85] Ofir Nachum, Yinlam Chow, Bo Dai, and Lihong Li. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. Advances in neural information processing systems, 32, 2019.
15

---

[86] Ofir Nachum, Bo Dai, Ilya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans. Al- gaedice: Policy gradient from arbitrary experience. arXiv preprint arXiv:1912.02274, 2019.
[87] Dylan Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv: 134877 2021.
[88] Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time. iv. Communications on pure and applied mathematics, ((2)):83-2122 1983.
[89] Edward $ Hu, Kwangjun Ahn, Qinghua Liu, Haoran Xu, Manan Tomar, Ada Langford, Jayden Teoh, Bryon Xu, David Yan, Dinesh Jayaraman, et al. The belief state transformer. arXiv preprint arXiv:2410.23506, 2024.
[90] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 18))9, 2019.
[91] Ilya Loshchilov, Frank Hutter, et al. Fixing weight decay regularization in adam. arXiv preprint arXiv:1711.05101,5 5, 2017.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Please see the last paragraph of our introduction which includes both our claims and the sections where the claim is supported.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations in Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
17

---

Justification:
Please see Appendix B for the complete assumptions and proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Please see Appendix C, D, E for the information needed to reproduce our results.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
18

---

Answer: [Yes]
Justification:
We provide our code with instructions in the supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/nips. wwtppcccccc/cceeccccccccccccclct httaa//ggiieeee/eeuuiiieeeee /Coo//eoomissssseeeee crriissssoppiiyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: All training and test details are provided in Section 3.1, 3.2 and Appendix C, D, E, F.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[No]
Justification:
The experiments are reported on datasets that are significantly large.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
19

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The compute resources used are reported in Appendix C.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Ethics.
The research conducted in the paper conforms with the NeurIPS Code of
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper presents work whose goal is advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
20

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The benchmarks and data splits are publicly available. All licenses are respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
21

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Assets will be released, and instructions and details are included for reproduction.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
22

---

16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https:/ //ttrr//eeu../eeeee Wwweuripsscc/ooofeeeee /LLLL for what should or should be described.
23

---

A Related Works
From the empirical side, the most relevant works are controlled decoding (CD; [1I]) and value augmented sampling (VAS; [12D. These two works both propose to guide the reference policy nt with QT .0, the expected reward-to-go under n. withou! KL regularization. As discussed in Section 4.1, guiding with Q , is not principled for the KL-regularized RL problem and can lead to both sub-optimal reward and large KL from nt In contrast, we propose to guide Ned with Q*"", the expected reward-to-go under the optimal policy with KL regularization, which is the correct closed-form of the optimal policy, A recent work [34] proposed a process reward model (PRM) of a similar form as our Q* but their PRM is applied to steps instead of tokens, and they do not use distributional RL or iterative training (i.c., data aggregation).
In terms of reweighting ,d with classifier scores, FUDGE [67] is another closely related work but their derivation is based on Bayes rule and FUDGE does not solve KL-rregularieed RI.. Sequential Monte Carlo (SMC) methods [22, 68] also reweight r' distribution with a twist function, where the optimal twist function is analogous to our Q* One key difference that SMC performs resampling while we directly combine logits of -d and exp(Q".n) to avoid importance sampling, which has higher variance. Finally, none of these prior works apply distributional RL losses [17, 69, 70, 57] or online data aggregation [13] to learn Q* which we showed to be beneficial in our ablations. Indeed, CD and VAS both use square loss regression over a fixed offline dataset. We also remark that risk-sensitive RL has been an important application of distributional RL [69, 71] and extending QF along those lines is a promising future direction.
We also discuss some of the recent advances in stable distributional RL. [72] shows that the categorical distributional RL loss, which we employ for our theory and experiments, enjoys smoothness and optimization stability under a bounded logit condition. [73] introduces Sinkhorn distributional RL loss which computationally efficient alternative for Wasserstein distance, and was shown to be more stable for multi-dimensions rewards. [69] proposed 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 [74] identified as 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-Icarning or tricks such 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.
We also cite some tangentially related works. Proxy tuning [75] and speculative decoding [76] both use a small model to guide the logit distribution of a large ne model. Speculative decoding is focused on maximizing the large model's likelihood, which does not relate to any extrinsic rewards. In our framework, the classifier model can be any size relative to n. although deeper investigation into the computational benefits of using a small classifier is a promising direction for future work. We note that the star-graph problem can also be solved during pre- pre-training by also predicting backwards via the belief state transformer [36].
Finally we discuss previous post-training methods for LLMs. First, online iterative DPO [77. 78], REBEL [79], PPO [8], etc. are based on policy gradient and require a good reset distribution which only guarantees local optimality. XPO [80], VPO [81], SELM [82], etc. treat this as an exploration setting but requires solving non-convex optimization oracles and relies on strong structure conditions such as coverability / eluder / linearity, similar to the theoretical works like [55, 59]. Instead, we approach post-training in a fundamentally different angle and solve it via simple computationally tractable regression and mle oracles, without any strong structural conditions or reset distribution assumptions.
From the theoretical side, KL-regularized RL is closely related to soft RL or maximum entropy RL which are well-studied [20, 83, 84, 22]. The optimal policy decomposition in deterministic MDPs is also known in prior works [23, 21]. Our contribution is an algorithm that provably learns Q* using distributional RL [17] and data aggregation [13]. This enables us to prove reduction of KL-regularized RL (in deterministic MDPs) to no-regret online learning, which ensures convergence to the optimal policy with realizability being the only assumption for function approximation. Notably we are able to avoid more stringent conditions such as completeness or structural MDP conditions which are ubiquitous in the current literature [53, 55, 56, 35, 29, 57, 59]. [66] observed similar benefits in offline RL, while we provide guarantees for the harder online RL setting.
24

---

Complementary to our online, KL-regularized setting, DualDICE [85] and AlgaeDICE [86] tackle the aariinarancce "curse of horizon" that arises when one performs importance weighting for long tra- jectories in offline RI.. Both methods replace per-step importance weights with caationaryddistribenn density ratios, learned through dual (Lagrangian) formulation, and have shown empirical success on low-dimensional continuous-contro) benchmarks, although learning is also shown to be difficult in high-ding high-dimensional control tasks [56]. Because we continually collect on-policy data and constrain updates via an explicit KL penalty-which already limits distribution shift-we do not need such ratio estimation; nonetheless, density-ratio approaches remain a promising orthogonal direction for variance reduction in purely offline LLM post-training.
We remark that our theoretical guarantees are quite similar structure to that of Aggre VaTe [31, 32], which is reduction of imitation learning to no no-regret online learning. Besides the obvious difference in problem setting, another improvement from our work using distributional RI. theory to prove second-order bounds. Notably, we are able to prove second-order bounds without any completeness assumptions that were required in [35, 28, 29].
B Proofs
In this section, we provide the full proof for Theorem 4.4. Theorem 4.4. Fix any (0, V and 8 E (0, 1). Under Assumptions 2.1 and 4.3, Algorithm 2 ensures w.p. at least 1 - 6, setting 3 lm(1/0) + Regmle we have
Σk (V* - 0)) S AV (v Σ/ h- Σ CV (x,, CV v) B maXke(M) E, -3);
Var(exp(2)
where CVA (I, y) == En wark -ni((A((AA)
is the coefficient of variation
E Exexpl (2....B
of exp (In, w/)//)), and En = |eep(( (Vmax - ( yn () 1)/m)))///lll (l (r is the envelope of exp(( (ep((Mmmx (5n,, Ux)),))))/)) both under refe
Proof. Fix any 7) E (0, Vm.x). Let Qn,(Iz, v) = 17 2- exp(z/n) denote the induced soft Q function from the distributional estimate Ze. Let ( x) oc कt (y) |x) exp(Qn.e(z))) denote the induced policy from Qh,k- Then,
V v**. .0
h= + 1Q (oh,
(xp, x Q " "Khn T + (KL( oH. (In)) (LL (xn))
H Exd 12 31 (Ih, )1 - r)KL(TT (mT) (mx)) - (Qna(En, (K (zx)) || (
+ Qu T//)
(xn, x1)1
(10
Eth Byt 10 (5h, R/ x -T (x (f)())) (Qn,e(en, Tx7) 7)KLI(TT (x%)) (xn)))
+ (In Th) 2110" (xn, +
EE Exl 12 t7 (In, ) - Qi (#h, T ) + Qu, (AAA( T/:) - Q2 (Zh,
Σ h=1 Bow wot [ma .x 12 17 (I/,, T) (1, T]
< 2A Σ/ By 1 - UUni(() 1 (xr, vn. - Qn. (In, wn)
where (i) is by the performance difference lemma in the soft MDP (Lemma B.2): (ii) is by Donsker- Varadhan (Lemma B. which proves that (M*) arg max, E,, QQ, (Z), X) (L(x(xx)
afd (xx)). Now, bound the difference between the optimal and learned Q functions: 1 (z, v) Qa v))
" In (e exp(z/n) In (Z exp(z/n) (1)
S (11 + Vinx//)) CV:
exp(Vmarr exp(Vmar/1)-1 H (I, 8)
Eawap(a.s) expls/n)
(1) + VInai ( (exp==
exp(Vmaz//)) H2 (P. "..(Q,""<."))))
25

---

where (i) is by Lemma B.4 and the fact that 2z E [ and (I, v) (Z y), Zn y)) is the Hellinger distance between the learned Zn, and optimal zz.
Thus, if we let Th, Uh ~ xh on (A) denote the distribution of rolling in with n& until In and taking a random yh ~ (A), then we have:
Σ V.n - V**
2A EAE 11=1 E Ek=1 Un v.) - Qn, U.]
AV max Σx=1 ER=
(esp/) Mh) exp((2 " (2b.20%)/1) 11, (2,, w)
exp /n)
AVmax
1= E Eor Uni(((A)
CV h. (Fh, E ER= B Batto (2 vhh)
+ AVmax
exp(Vmase/)) exp(Qn (20 1/m) Unit,))))A))
E E Unit, [H (xn,, Wh)]-
The final step is to bound the summed Hellinger square terms. This can be done via Multiplica- tive Azuma's inequality and [87, Lemma A.14], which shows that for any d (0, 1), have
En, Batt, Unir(A) (Zh, Uh)] I Et, H2, (Z yh.k) + ln(1/6) & Regmte (K) + ln which recall is exactly the definition of B. This finishes the proof of Theorem 4.4.
O
Lemma B.1 (Donsker- Varadhan's 's Variational Formula; [88]). For any prior p A()), consider the KL-regularized optimization:
n* = arg max (E V(T) == (,,QQ()) - 7)L(((()) II p(0)))
The optimal policy x² is given by (0) x p(0) exp(Q(0)) and it has value V(**) = n In En-p exp (x)(Q(0)/n)
Lemma B.2 (Soft Performance Difference Lemma (PDL)). For any I and n,
H [T fat1 fa un)l - 7KL of (z:)).
Va - (1
Σ
h=1
For any n. T'.
V - V* Σ Ex 10% (w). - Q% (In, nKL((H ll (xx)) (KL(TT(IN) ll (h
h=1
Proof. Let (L(TT(ZN)) = (L(TA(xN) (xx)) denote KL-divergence w.r.t xe Then,
V -V
Ei= Ex [Th (KL(TA(IN)) - (TT(
DM Ex [TM 7/LL(THH (I KLL (x++
Q7 (In, -]L((T(TA)) + (
+ Q (#o+ x²) - Q (In, n²) + (KL((T((NN) (TL()(( Σ" Ex [T Q (I1, 9n) - Q (zh, + (L (zx)) - (TH((T
Σm [r(rn (TKL(Th=11
yn) (D)) E Ex 1Q (IH,
Q7 (I,, + W)L (xx)) (TA
Lemma B.3. For any two numbers x, v E (xxp((),exp(()) we have
[In(x) - In(y)) < (1 + b a)
b 2 2. then max(1,
b-a T-) S 3(b a).
26

---

Proof. x 2 y, then In(z) =
If < y. then In In(2x) = -Im(1+(x-))))) /v)) By premise, we have 0 2 -ー 2 exp(a-b)-14 Note that - (n(++2) is convex
and is thus upper bounded by the line connecting (0, 0) and (exp(a b) - a-a), Le.. In +2) -
b-o [- for 0 2 exp(a b) 1. Thus, In(1 (æ w)/w) S T-exp(a-b) æーツ Thus,
h-a
we've ve shown that [In(x) - In(y)] S max (1, T-exp(a-
b-a
Finally, since T-oxp(-m] S 1
when x 2 have max(1, b- S max
-
Lemma B.4. For any distributions p. 9 on (a, b), we have
In (
a)
Var,(e") Ege
H(p, + expl))-exp(a) H (1 Ege²
where (1 - [.(((()) v(()))² is the squared Hellinger distance.
Proof. By Lemma B.3, we have [In P(=)e* - In E. a())e* S (1+b- a) Σ. (x()))--(e))) E. a(aje=
By Lemma B.5, we have that the numerator is bounded by / (e) + (exp(b) (xp(a))H?(p,.d))
Lemma B.5 (Second-Order Lemma). Suppose p. are distributions on the interval (a, b). Then, we have
l- - x VVrr(p) x) + (b ((6-a) (f
Proof. Define p/. q as the normalized distributions on [0, 1). is the law of x'= (X (X-a)/(b- -b) where ~ p. Then, we have
1-
क = (b a)lF' - 7
S (b - a) Vara\\/aa (1
VVrlp) 4) + (b- Hf (p. 4),
where the S step is due to the second-order lemma of [28].
B.1 Case of Bernoulli reward-to-go
In this section, we focus on problems where (Ir y) Ber(pa (z, y)) is a Bernoulli distribution, which is common for closed-ended problems such as math or multiple choice. Here, the envelope term can be bounded follows:
Lemma B.6. If (1, = (er(pa(z, v)), then we have Vmma = 1 for all nj > O, we have
cxp((1 Qx," v))/n) s 1/px(z, v)-
Proof. Fix x, y and let p = Pr. (x, v). Then, it suffices to show that
1/n In(pexp(1/n) + p) s ln(1/p).
This is indeed true because
exp(1/11)
1/n ex(pe -1-p) = In
pexp(1/n) p
1
In
1mm(1p))
p+ (1 - p) (exp(11/7)
We can also bound the coefficient of variance in terms of the Bernoulli parameter. Lemma B.7. If 2,, (2,y) 9) (er(p,(x, v)), then for all r > O, we have
/ Var(expp (Z, v)/n))))))))))))))))))))))))vv E Elexp( (2, v)/n)
S v p1-p)/
27

---

Proof. Fix 2, y and let (x, v). Then, the variance term is:
Var(exp 2 u) w)/n)
- (B
( p) - (pexp(1/n) (1
(1 p) - p² exp(2/n) - 21p
- p) - (1 2p)
p(1 exp((()) + (1 - (1-p)p- 2pexp(l/n) p)
= p(1 ()xpp(22/)) + exxp(1/n))
p(1 ()xp((1//)- - 1)
Thus, the CV is:
pp(1- (((xpp)))(11/) 1)² pexp(1/1) + +1-1
Vl(L-p)(exp(() 1)
p(1 p)
p(exp(1/n)-11) 1I <
((1 - p)/p.
P
28

---

Pre-Train
REINFORCE
DPO
RPO
CD
OF (ours)
100
80
Accuracy [%6]
60
40
20
G(2,5)
G(5,5) Star Graphs
G(3.8)
Figure 5: Full results for star-graph experiments. The empty bar for G(5, 5) and G(3, 8) are for DPO, which pushed down both the chosen and reject paths resulting in 0 accuracy.
C Additional Discussion and Implementation Details for Star-Graph
The shortcut behavior, also known as the Clever Hans Trick [18], in the star-graph task arises directly from the regr-segee next-token prediction objective. Specifically, the model minimizes loss by memorizing the first token seen during training and following the corresponding edge, achieving low training егтоr but generalizing poorly at test time when the initial token is not provided. This leads to brittle, shortcut- bhsse policy.
Policy-based methods such as REINFORCE and RPO attempt to correct this by upweighting high- reward trajectories. However, because their loss is still based the product of next-token prediction probabilities, the same as in pretraining, they are vulnerable to the same shortcut and require exponentially many samples via gradient descent on the policy correct once it is learned (Theorem of [89]).
In contrast, Q1 does not depend myopic next-token supervision. Instead, it learns predict the cumulative reward-to -o-goo from each (prefix, token) pair under the reference policy, and uses this to guide generation toward optimal completions. This token-level value modeling allows Q2 predict future outcome and assign higher value to early tokens that lead to long-term reward. In other words, Qi 's loss function is directly trained perform planning, making it robust the Clever Hans Trick [18] that undermines next-token-base methods. As shown in Figure 5, both Qi and CD are able to solve the star-graph task near-perfectly, while policy-based methods perform at random-gues: level.
follow the setup of [18] and reused their official code for producing the star-graph results. We used the GPT-2 small model for graphs G(2, 5), G(5,5 5) and the GPT-2 medium model for G(3, 8) [90].3 We first pretrain these models with next-token prediction on a pretraining set of 200k random graphs and correct paths. We call this the resultant model the "pre-trained" model, and observed by [18], these models have the Clever Hans shortcut so they do not generalize well on unseen test graphs. We highlight that this is failure in generalization, since the pre-trained model achieves near-perfect accuracy on the training set but only 1 1/d accuracy on the test set.
In order to fix the Clever Hans shortcut, we perform post-training with two common baselines REINFORCE [41] and DPO [9], RPO [42] - as well as our algorithm Qr. The post-training is done on another set of 200k random graphs. For REINFORCE, the reward function use is 1 if the response is correct, and -0.1 if incorrect. We noticed that the incorrect reward is too negative, this causes model collapsing to accuracy of 0. For DPO and RPO, sampled pairwise responses (Vebosen- Freject) where Yehosen is the correct path and Project is an incorrect shortcut path sampled from the pretrained model. For Qt. also trained the classifier on the same dataset of pairwise responses, where correct paths are marked with reward and incorrect responses are marked with reward O. Throughout, used the AdamW optimizer with weight decay 0.1 and batch size of 256, and trained for 10 epochs. The learning rates were 2.5e 4 for pre-training: 1e for REINFORCE; le 4 for DPO and RPO; le for classifier- CD and Qt. All models are trained on single A 100 or H100 GPU. All models were evaluated on a separate test set of 20k graphs, using top-k 10 and
MModels from https://huggingffcc.oo//// compummtiooocmmmim//m/mmoco and https//huggingfiie.coooopat community/gpt2-mediu
29

---

temperature 1.0. For Qi and CD, we use = 0.1. We found that DPO often pushed down the probabilities of both the chosen and reject paths, leading to poor performance even on the training set; RPO fixed this issue and so we report the RPO numbers.
D Additional Model Details
re models. All models we use in the experiments the "Instruct" versions. That is, Llama -aaaa-----Inntttee and we use the default chat template and
8B refers to 1aa--11amaMMeta--tta system message from Meta interact with them.
Qi models. Two variants for QF are implemented and experimented: Q-type and V-type. Specifically, the Q-type takes input of a partial generation æ: and computes Q* (1c, y) for all y in the vocabulary of the -h model whereas the V-type takes input of concatenated x and a specific token ý and outputs a single value that represents Q*,"(x, 0)). Because of the key difference, Q-type therefore can efficiently calculate Q* with just one forward pass and its model architecture can also be identical to the original LLM. V-type, however, has prohibitive inference cost with naive implementation since requires making [V] forward passes at every decoding step to calculate the full Q function. In the paragraph below, diseuss efficient implementation to address this issue. For Q-type, initialize the model directly from Llama 3.2 1B and for V-type, we replace the last layer of Llama 3.2 1B with randomly initialized fully connected layer with output size of Therefore, V-type Qf also has slightly fewer number of parameters than Q-type. We by default use V-type QF in our experiments.
Efficient inference with V-type. To speed up inference for V-type, we note that not all tokens in the vocabulary are worth computing its value since for any partial generation most tokens have extremely low probability from ाम as the next token candidate. In our preliminary experiments, have found that only computing the values for the top 20 tokens ranked by х/ give similar performance compared to computing for all tokens. Additionally, we also note that the values for these tokens can be computed in one forward pass. To accomplish this, input partial generation x and the top 20 candidate next tokens together, modify the attention mask so that the candidate tokens do not attend to each other but still æ. This allows to compute the values for these top tokens in just one additional forward pass without any approximation.
E
Q# Training Settings
We collect 16 samples for each question in the training set and label every sample either as correct (1) or incorrect (0) based on the final answer. The first round of training data is collected with just ne For training Qi model, we filter out samples from questions where all samples are either correct or incorrect. we use a learning rate of 2e and weight decay of 0.01 with Adam optimizer [91]. The model is trained for 5 epochs. We train Q# for two iterations converges. In the second iteration, we repeat the above data collection procedure and concatenate the training data from the first round. The model is always trained from scratch between iterations.
observe performance
F Additional Evaluation Details
We evaluate all methods and models with zero-shot prompting. The prompt template is PProblem \n In\n(0} Write your answer inside Wboxed[(O. An AnAnSolution: where {0} is replaced by the actual question from the dataset. The MATH-500 dataset can also be found at Huggingface 4
G Math Reasoning Results on Qwen 2.5
We conduct experiments using Qwen 2.5 [46], where 1.5B model guides the 7B version on GSM8K, MATH and AIME-24 (Table 5). All other configurations mirror those used with Llama 3. We find that Qi consistently outperforms both re and CD across all datasets, achieving higher accuracy with
30

---

lower KL divergence. Compared to Table 1, Qwen 2.5 yields stronger overall performance, likely due its stronger base model, demonstrating that Q# generalizes well across model families.
Table 5: Comparison of QE with xe and CD baseline on GSM8K (left), MATH (middle) and AIME-24 (right) with Qwen 2.5.
Dataset
GSM8K
MATH
AIME-24
Methods
re
CD
Q1
CD
QF
CD
Q1
pass@1 1
76.1
79.0
83.5
58.6
60.7
61.9
9.3
13.5
14.1
maj108 108
92.9
93.1
93.8
72.8
74.2
74.8
16.7
16.7
20.0
KL-Divergence
5.37
4.10
7.07
6.46
9.95
9.23
H
Results on QuALITY
In Table 6, we show the results of Qf on QuALITY [47], a challenging multiple-choice reading comprehension benchmark with long-form passages drawn from Project Gutenberg. Q# consistently performs better than baselines.
Table 6: Comparison of Qf with re and CD baseline on QuALITY with Qwen 2.5 and Llama 3.1.
gee
Qwen 2.5 7B
Llama 3.1 8B
Methods
CD
QF
CD
Q1
pass01 f
64.5
64.2
68.1
73.5
75.1
75.9
maj108+
72.0
66.3
73.3
79.3
79.3
81.1
KL-Divergence J
12.32
7.90
9.23
8.88
Comparison with Policy-based Methods
Qf can serve as lightweight complement to policy-based approaches. Specifically, Q2 can guide both the base reference policy and policies trained via reinforcement learning such PPO. To empirically assess this, we present results on the MATH dataset where Qi is instantiated with Qwen 2.5 1.5B model and used guide: (1) the base Qwen 2.5 7B reference model and (2) a PPO-trained version of the same model. As shown in Table 7, Qf consistently improves both pass$1 and maj 188 for each policy. particular, when applied to the PPO-trained policy, QE reduces the KL divergence from nd while further boosting accuracy. We also note qualitative distinction: PPO improves pass@1 but slightly reduces maj1 168, indicating that its generations tend to lower entropy and less diverse. QH, in contrast, improves both metrics while maintaining closer aligmment with re
In terms of efficiency, Qł is significantly lighter to train. PPO requires approximately 20 hours on H100 GPUs, whereas Qi training completes in roughly 5 hours on a single H100 GPU, thanks to its supervised learning objective and the use of much smaller model. These findings suggest that Q= can effectively enhance performance while maintaining closer alignment with the reference policy, demonstrating its practical advantage as complementary lightweight module.
Table 7: Comparison of Qf with PPO-trained models and their guided variants on the MATH dataset.
Methods
n
n ref
Qi
PPO
PPO + QF
pass01 1
58.6
61.9
68.4
71.1
maji08 f
72.8
74.8
72.4
73.4
KI.-Divergence +
6.46
69.52
60.53
31

---

Computational Complexity and Runtime Comparison of QF
Qi and other value-based baselines such as CD [11] have the same computational complexity, Compared to generating responses solely with ne value-based approaches additionally use the guidance model to compute Q function at every decoding step. That it increases complexity by the ratio of the guidance model's 's size to that of -t Since the guidance model can be much smaller size compared to ,h the overhead mild. For instance, guiding a Llama with Llama increases complexity by 12.5%.
Additionally, efficiently implemented value-based guidance for Q: in Hugging Face using LogitProcessor and key-value caches. On an Nvidia A6000, generating one response on test set of MATH takes 4. 40s for πh and 181 for Qt. slightly exceeding 12.5% possibly due to sequential Q function computation in LogitProcessor. The code for our implementation can be found in the supplementary materials.
K Qualitative Examples
In Figure 6 and the ones below it, we qualitative visualize side by side generation results from x and QF on Llama 3 8B GSM8K and Llama 3.1 8B MATH settings. In the first example of Figure 6, we observe nd and Q$ start off similarly by calculating the total number of cookies Shannon eat but ाh starts to roughly guess the answer (highlighted in red) without calculating the answer precisely, where Qt calculates the answer step by step (in blue). In the second MATH example, gh first states some confusing statement (highlighted in orange) and then makes a mistake addition (5 + 5) instead of multiplication when calculating all the possibilities. Qi. however, correctly solves the problem by multiplying the options (5 * 5) for Paul and Jesse and then adds the single case when using the number of 2, arriving the correct answer.
More examples can be found subsequent pages with various levels of difficulty, We highlight that Qi can still make confusing statements similar to xh even in solution that arrives the correct final answer. For example, in the Cecilia puppy example (the first after Figure 6), similar to ne Q: also makes division of 1 when it should be multiplication.
32

---

Problem from GSM8K: Shannon makes her madcleine cookies and cats 2 night treat. She make enough cookies to last her for 30 days by storing them in the freezer. Her recipe makes dozen madeleine cookies. How many dozens cookies will she need make so she has enough for 30 days?
Incorrect Solution from n² Since Shannon eats cookies per night, she will eat 2-30 60 cookies over 30 days. She will need dozen cookies, since 7 12 84 >> 60. Final Answer: The final answer is hope is correct.
Correct Solution from Qf: Shannon eats 2 cookies per night, she will eat 2 30 = 60 cookies in 30 days. Since dozen cookies is 12, we can divide 60 by 12 to find out how many dozens she needs. 60 + = answer 5 dozens.
Counting & Probability Problem from MATH: Paul and Jesse each choose a number random from the first six primes. What the probability that the sum of the numbers they choose even?
Incorrect Solution from _ To find the probability that the sum of the numbers Paul and Jesse choose is
need consider the possible combinations of even and odd numbers. The first six primes are 2, 11, 13. Since is only even prime number, we always and one or more odd numbers. For the sum to be even, both numbers must be either even or odd. Paul chonses the even numher (2), Jesse can choose any of the add numbers (3, 5, 7, 11, 13), resulting in 5 possibilities. If Paul chooses any of the odd numbers, Jesse can choose the even number (2), resulting in more possibilities. Therefore, there The total number possible combinations The probability that the sum of the numbers Paul and Jesse choose is even is 10/36, which simplifies to
one even number
5 = 10 possibilities where the is even. 36 choices Paul and choices Jesse).
5
5/18. answer is:
18
Correct Solution from Qi: The first six primes are 2, 3, 5, 7, 11, and 13. For the sum of the numbers Paul and Jesse choose to be even, they must either both choose numbers. There is only one even prime number, which is 2. Paul has 6 options, and Jesse also has options. The total number of possibilities is 6 6= 36. For the where both choose the even number, there is only one possibility (2, 2). For the case where both choose add numbers, there are options for Paul and options for Jesse. The total number of possibilities 5 * 25. So, the total number of possibilities where the sum is even is 25 = 26. The probability that the sum of the numbers they
numbers both choose odd
13
choose is even is 26/36 13/18. The answer
Figure 6: Example generations from re and Qi on GSM8K and MATH questions. oe consistently fails to solve the problems compared to Qt. The generations from re and Q: usually start off with similar reasoning chain-of-thought but ,e gradually makes unjustified reasoning leap or simply guess of the manually added for ease of visualization. Blue represents correct reasoning steps or answers, red represents clearly incorrect ones and orange represents ambiguous or minor mistake that could potentially lead to clearly incorrect steps or answers.
whereas Q$ tends to be more logical for solving the problems. Colors
33

---

Problem from GSM8K: Cecilia just bought new puppy. According to her veterinarian, she has to feed the puppy cup of dog food every day for the first 180 days. Then she has feed the puppy cups of dog food every day for the rest of its life. If bag dog food contains 110 cups, how many bags of dog food will Cecilia use the first year?
Incorrect Solution from n For the first 180 days, Cecilia needs 180 180 cups of dog food. For the
rest of the year, she needs 365 180 = 92.5 cups of dog food. So in total she needs 180 + 92.5 = 272.5
2
cups of dog food. If one bag of dog food contains I 110 cups, then she will need 272.5 dog food. Final Answer: The final answer is 2.5. hope is correct.
2.5 bags of
110
Correct Solution from Q:: The first 180 days, she will need to feed her puppy 180 180 cups of dog food. After that, will need to feed her puppy 365-180=185 days cups each, for total of 2 185 = 370 cups of dog food. So, the total amount of dog food needed in the first year is 180 + 370 = 550 cups. Since each bag dog food contains 110 cups, the number bags needed
550
110
Problem from GSM8K: According its nutritional info, bag of chips has 250 calories per serving. 300g bag has 5 servings, how many grams can you if your daily caloric target is 2000 and you have already consumed 1800 calories?
Incorrect Solution from x= Let's break this down step by step!
The bag of chips has 5 servings, and each serving has 250 calories. To find the total calories in the bag, multiply servings x 250 calorics/serving 1250 calories
number servings by calories per serving:
2. You have alrcady consumed 1800 calories, and your daily calorie target is 2000. To find out how many calories you consume, subtract the calories you't 've already had from your daily target: 2000 1800 200 calories
Since each serving of chips has 250 calories, you can consume:
200 calories + 250 calories/serving 0.8 servings
4. Since there are 5 servings in the bag, you can eat:
0.8 servings x 300g/bag 240g
So. you can eat up to 240g the bag chips.
Here's the answer:
240
Correct Solution from Q#: A delicious problem!
Let's break it down step by step:
1. The bag of chips has 250 calories per serving, and there are 5 servings in the bag. servings 1250 calories.
The total calories the bag are: 250 calorics/serving
3. You have already consumed 1800 calories, so you have 2000 - 1800 = 200 calories left for the day.
Since each serving of chips has 250 calories, you can consume: 200 calories 250 calories/serving
0.8 servings.
Since are 5 servings
bag, you can eat: 0.8 servings 300g/5 servings 0.8 60g 48g.
So, you can eat up 48g of the bag chips.
Enjoy your snacking, but don't 't overdo it!
34

---

Problem from GSM8K: Nate's dog can dig six holes day. He digs for 14 days while Nate on vacation. When Nate gets home, he starts filling in 9 holes day, but the dog keeps digging new holes every night. How many weeks does take him fill all the holes?
Incorrect Solution from r=. Nate's dog digs 6 holes a day for 14 days, so the total number of holes dug is 6 x 14 84. After Nate gets home, he starts filling in the holes, and the dog keeps digging new holes every night. So the number of holes fill is always 84 6. Nate fills 9 holes day, the number of days takes to fill all the holes
84 + 6
=
9
Correct Solution from Q:: The dog digs holes day, and Nate is away for 14 days, so the dog digs 6 14 84 holes. Then, Nate starts filling in 9 holes day, but the dog digs 6 new holes every night, so 3 holes are filled each day. So, each day more holes are filled in. The number of holes want to find holes.
filled in after days is 84 3x. We want find the smallest such that 84 3x is O, the smallest x such that 84. 3x 84 true when x 28, so takes 28 days to fill Since there are days week, it takes
weeks fill
holes.
Prealgebra Problem from MATH: bookcase has shelves total of 24 books. The top shelf mystery books. The middle shelf has 10 math books. The bottom shelf has science books. Two hooks are now taken off each shelf. What fraction of the books remaining on the three shelves are math books? Express your answer common fraction.
Incorrect Solution from x = find the fraction the books remaining on the three shelves that are math books, we first need to find the total number of books remaining on the three shelves. Initially, the bookcase had 24 books. Two books are taken off each shelf. Since there are shelves, total of 6 books are taken off.
So, the total number of books remaining on the three shelves is: 24 (initial books)-
(books removed) books.
The middle shelf had 10 math books. Since no math books were removed, the number of math books remaining is still 10.
Now, we find the fraction of the books remaining the three shelves that are math books. To do this, we divide the number of math books remaining by the total number books remaining.
Fraction of math books = Number of math books Total number of books = 10/ 18
We can simplify this fraction by dividing both the numerator and denominator by their greatest common divisor, which 2.
Simplified fraction 5
So, the fraction of the books remaining on the three shelves that are math books is: 5/9
remaining books are math books, we first need to find the total number of books remaining on the shelves after two books are taken off each shelf. Top shelf initially had books. After removing books, has 8 2 books.
Correct Solution from Q#: To find the fraction
Middle shelf initially had 10 books. After removing 2 books, it has 10 2 8 books.
Bottom shelf initially had books. After removing 2 books, has
books.
So the total number of books remaining is 6 + 8 4= 18 books.
Initially, there were math books. Since only the middle shelf has math books, the number math books remains the same at books.
The fraction of the remaining books that are math books is the number math books (8) divided by the total number of remaining books (18).
Fraction 8 18 =
So the answer is: 9
35

---

Intermediate Algebra Problem from MATH: The asymptotes of a hyperbola are y 2z and y 17 2x. Also, the hyperhola passes through the point (4, 7). Find the distance hetween the foci of the hyperbola.
Incorrect Solution from x² To find the distance between the foci of the hyperbola. we first need to determine the center and the standard form of the hyperbola. The asymptotes given by 2x 2x. Their point intersection will the center the hyperbola.
find the intersection point, we set two equations equal to cach other: 2x - = 17 2x 4x 20 x = Then substitute the value of x into one of the asymptote equations to find the corresponding value: y 2(5) - 3 7 So, the center of the hyperbola is (5, 7). Next, we need to find the standard form of the hyperbola. The slopes of the asymptotes 12, the standard form of the hyperbola given by:
and y
10
y 7)
(x 5)2
6a
Since the slopes of the asymptotes are +2, we know that b/a = 2. We can use the point (4, 7), which the hyperbola passes through, to find the values of a and b.
7)
5)2
a²
4
(-1)²
b²
ba
=1
b =
Since b/a 2, have 6/2 1/2. Now find the values of a and 6².
a² (1/2)²=11/
Now that we have the values of a² b², we can find distance hetween foci the hyperbola using the equation:
0² a² +b²
0² = 1/4+1
2 5/1
== 55/1 v5/2
The distance between the faci of the hyperhola 2c = 2/5/2 /5. Therefore, the distance between
the foci of the hyperbola is v5
Correct Solution from Q#: find the equation of the hyperbola, we first need to find the center and the standard form of the hyperbola. The center is the point of intersection of the two asymptotes. To find we can set the two equations equal to each other and solve for x:
2x Combine like terms: 1r 20 Divide by 4: 5 Now substitute this value of x into one of the asymptote equations find y: 2x
3 y = 2(5) 3
So, the center the hyperbola is at (5,7). Since slopes 82
of the asymptotes are and -2, the standard form of the hyperbola is: (g-A)² (.-h) = where of the hyperbola. To find the values of and b,
a²
(h,k) is the of the asymptotes are 2 and -2. The slopes of the asymptotes are given by 4. Therefore: & = 2 We also know that the point (4,7) is on the hyperbola. We can substitute this point into the equation ger: (T-7)²
the fact that the slopes
(4-5)2 This simplifies to: = 1² Now we can substitute this value be into 8
the equation
2 to get: & = 2
Now can substitute the values of and b into the standard
form of the hyperbola to get: (w-7)² (a-5)² The distance between the foci hyperbola given by 2c, where is the distance from the center to each focus. The relationship between b, and given by: c² a² + = 2 +
/5 The distance between the foci is: 2c = 2
the distance between
hyperbola is 2v5.
36

---

Algebra Problem from MATH: Simplify the following expression simplified fraction:
V845
/5
\80
v5
5
V845
+
4
/80
Incorrect Solution from n²'; To simplify the expression
we need follow
v5
5
the order operations. First, we simplify the terms inside the square root:
/80
v16
x
5
v5
5v3
v5
4v5
4
55
4
v845
v13
x 65
/13x x3 = 13v5
V45 v V
v5= 3
v5 +
13v5
+3v5
Now, we can rewrite the expression as:
We can simplify the numerator by
v5
finding common denominator: V5
13v5
9V5
52v5
108v5
169v5
+ +35
Now,
9
36
36
36
36
169
we can rewrite the expression as:
36
To simplify the expression, multiply the numerator
v5
169.5
and denominator by V5:
36
169v5
V5
169 x
845
Now, we can rewrite
V5
36
845
expression
We simplify the square root by finding the square root of the numerator and
36
845
5 x
v5
13
v5x
13
13v5
denominator:
So, the simplified fraction
62
6
6
Correct Solution from Q#: To simplify the given expression, we first simplify terms within square root in the numerator.
V5
V80 V16
1V5
1
V815
/I3
65 113 13. 13v5
15
v9.5
3
Now, we substitute these simplified expressions back into the original expression:
/845
v5
13v5
45
+
+v5
/80
9
4
9
V5
V5
We can simplify the expression further by combining like terms in the numerator:
55
13v5
9/5 + 52/5 108
169/5
355
4
36
36
169
v5
v5
v5
36
Finally, take the square root of the fraction:
169
169
13
36
V36
13
The answer
37


Reviews:
[{"review_id": "qbeMzRFP0C", "reviewer": "Reviewer_2Dby", "summary": "This paper proposes Q#, a novel distributional reinforcement learning algorithm designed for post-training LLMs. It addresses shortcomings of existing methods like PPO and DPO, which fail to fix shortcuts inherited from pre-training. By introducing a value-based approach, Q# uses a KL-regularized Q-function to guide the model without modifying its weights. The method is shown to outperform previous baselines in math reasoning tasks while maintaining a lower KL divergence from the reference policy.", "questions": "1. What would the performance look like in real-world LLM tasks beyond math reasoning (e.g., commonsense reasoning)?\n2. How sensitive is Q# to hyperparameter choices and is there a trade-off between performance and stability in different configurations?", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. The paper provides provable guarantees for the convergence of the Q# algorithm, a significant theoretical advancement in KL-regularized RL, which ensures that the method avoids pitfalls of previous approaches.\n2. Q# consistently outperforms existing methods like CD and VAS in benchmarks, showing its ability to fix biases and improve performance in math reasoning tasks.\n3. The algorithm shows that smaller models can effectively guide large models, which could lead to more resource-efficient post-training approaches.\n\nWeakness:\n1. The experiments focus heavily on math reasoning tasks and a star-graph environment, leaving other potential applications of Q# unexplored.\n2. While Q# works with models as small as 1B parameters, the authors do not explore the computational overhead involved when scaling up to very large LLMs (e.g., 70B parameters).", "was_revised": true, "final_justification": "After reading the author's rebuttal, I have decided to keep my score unchanged.", "rebuttal": "Dear Reviewer 2Dby,\n\nThank you for your constructive feedback and we respond to individual points below.\n\n**Other potential applications of Q♯ beyond math reasoning**\n\n**Authors' reply:** Thank you for the valuable comment! Following your suggestion, we conducted experiments on QuALITY [1], a multiple-choice QA benchmark 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.\n\n|$\\pi^{\\text{ref}}$||Qwen 2.5 7B|||Llama 3.1 8B||\n|-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Methods|$\\pi^{\\text{ref}}$|CD|$Q\\sharp$|$\\pi^{\\text{ref}}$|CD|$Q\\sharp$|\n|pass@1 ↑|64.5|64.2|**68.1**|73.5|75.1|**75.9**|\n|maj1@8 ↑|72.0|66.3|**73.3**|79.3|79.3|**81.1**|\n|KL-Divergence ↓|-|12.32|**7.90**|-|9.23|**8.88**|\n\nThe results suggest that Q♯ can also consistently improve performance on QA tasks for models such as Qwen 2.5 and Llama 3.1 with smaller KL divergence than CD baseline.\n\n> While Q♯ works with models as small as 1B parameters, the authors do not explore the computational overhead involved when scaling up to very large LLMs (e.g., 70B parameters).\n\n**Authors' reply:** Thank you for the note. We would like to clarify that our approach is designed to efficiently guide large LLMs (such as Llama 70B) using much smaller Q♯ models (1B or 3B parameters), rather than scaling Q♯ itself to the same size as reference policy. As shown in Table 2, even a 1B Q♯ model improves the performance of a 70B reference model on MATH, both in terms of pass@1 and majority voting accuracy.\n\nWe also analyze runtime implications in Appendix H. The overhead incurred by Q♯ is approximately proportional to the size ratio between the Q♯ model and the reference model. This overhead becomes relatively smaller as the reference model scales up. For instance, when guiding Llama 8B with a 1B Q♯ model using our Huggingface Transformers implementation, the additional runtime is around 20%. This cost remains modest even in larger-scale settings and is justified by the observed performance gains.\n\n**Sensitivity of Q♯ to hyperparameter choices**\n\n**Authors' reply:** Thank you for raising this important point. Q♯ demonstrates strong robustness to hyperparameter choices, especially when compared to the CD baseline. As shown in Figure 3, the performance of CD is highly sensitive to the choice of $\\eta$: as $\\eta$ decreases, CD’s pass@1 accuracy drops sharply while its KL divergence increases significantly. In contrast, Q♯ maintains stable performance across a range of $\\eta$ values and consistently achieves a more favorable trade-off between accuracy and KL divergence. This indicates that Q♯ not only provides improved performance but also greater stability under different configurations, reducing the need for extensive hyperparameter tuning.\n\n**Citations**\n\n[1] Pang, Richard Yuanzhe, Alicia Parrish, Nitish Joshi, Nikita Nangia, Jason Phang, Angelica Chen, Vishakh Padmakumar et al. \"QuALITY: Question Answering with Long Input Texts, Yes!.\" In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 5336-5358. 2022."}, {"review_id": "k5qFSng0Qg", "reviewer": "Reviewer_bUKk", "summary": "This paper proposes Q♯, a new value-based algorithm for KL-regularized RL, tailored for post-training LLMs. Unlike popular policy-based methods (like PPO, DPO) that often fail to correct shortcuts learned during pre-training, Q♯ learns the optimal KL-regularized Q-function using distributional RL, then uses it to guide the reference policy. The approach treats LLM post-training as a deterministic MDP, fitting well with the autoregressive nature of language generation. Theoretically, it reduces KL-regularized RL to no-regret online learning, providing variance-dependent PAC bounds that guarantee convergence. Empirically, on tasks like star-graph planning, GSM8K, and MATH, Q♯ outperforms existing baselines, achieving higher reasoning accuracy with lower KL divergence from the reference policy.", "questions": "I think this paper has a strong theoretical foundation, but it would be even stronger with more empirical support. For example:\n- How does the empirical performance of Q♯ compare to policy-based methods (like PPO and DPO) on GSM8K and MATH?\n- How does Q♯ compare empirically to VAS?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### **Strengths**\n\n- While previous value-based approaches (like CD, VAS) guided policies using unregularized Q-functions (which does not solve the KL-regularized objective), Q♯ explicitly uses the optimal regularized Q⋆,η and provides rigorous convergence guarantees.\n- Avoids unstable TD bootstrapping by directly learning the cumulative reward distribution, framing it as supervised learning. This leads to better stability and faster convergence, especially when the reference policy has low variance.\n- On the star-graph task, Q♯ corrects the \"Clever Hans\" shortcut that policy-based methods like REINFORCE or DPO fail to fix.\n- On GSM8K and MATH datasets, Q♯ consistently improves pass@1 and majority voting accuracy over CD and the original reference model, while also incurring lower KL divergence.\n- Unlike CD, whose performance is highly sensitive to η and can suffer large KL spikes, Q♯ is empirically much more stable, reducing the burden of hyperparameter tuning.\n- The learned Q♯ can also evaluate complete generations as a reward model, improving Best-of-N sampling and providing generalizable scoring across different generations.\n\n&nbsp;\n\n### **Weakness**\nWhat I understood as the claimed advantages of the proposed method over policy-based methods:\n(1) It does not require updating all the large model parameters, only needs to learn the Q function.\n(2) It can correct biases and shortcuts learned during pre-training.\n\nWhat I’m most curious about is:\nWhy can’t policy-based methods fix such shortcuts, whereas Q♯ can?\nSection 3.1’s star-graph experiment is interesting, but it’s a simple toy environment, and the paper essentially says “it worked when we tried it.” It does not clearly explain why only Q♯ is able to correct these shortcuts. I’d like to understand the reason behind this more deeply.\n\nAlso, what would be an example of such a shortcut (or hacky solution) in math? Is there an example from GSM8K or MATH where Q♯ uniquely fixes a mathematical shortcut that the policy-based methods fail to fix?\n\nLooking at Appendix G’s qualitative examples, it seems that the benefits of Q♯ come more from general advantages of LLM post-training — such as fixing calculation errors, misunderstanding the problem-solving process, or skipping essential logical steps — rather than specifically correcting “shortcuts.” So it’s still unclear to me in what empirical ways Q♯ is actually better than policy-based methods.\n\nAlong similar lines, the star-graph planning task, GSM8K, and MATH actually seem to be quite different in nature. The star-graph task appears to require more explicit planning, while GSM8K and MATH seem to focus on reasoning. Because of this, showing a toy experiment on star-graph planning and then moving on to GSM8K and MATH feels a bit disconnected.", "was_revised": true, "final_justification": "After reading the rebuttal and carefully considering the authors' responses, I have decided to raise my score to a 5.", "rebuttal": "Dear Reviewer bUKk,\n\nThank you so much for your detailed review and suggestions for us to improve the paper. We respond to individual points below.\n\n> What I understood as the claimed advantages of the proposed method over policy-based methods: (1) It does not require updating all the large model parameters, only needs to learn the Q function. (2) It can correct biases and shortcuts learned during pre-training.\n\n> How does the empirical performance of Q♯ compare to policy-based methods (like PPO and DPO) on GSM8K and MATH?\n\n**Authors' reply:** Thank you for the helpful question and the understanding is correct. We also add that value-based methods like Q♯ can serve as a lightweight complement to policy-based approaches. Specifically, Q♯ can guide both the base reference policy and policies trained via reinforcement learning such as PPO.\n\nTo empirically assess this, we present results on the MATH dataset where Q♯ is instantiated with a Qwen 2.5 1.5B model and used to guide: (1) the base Qwen 2.5 7B reference model and (2) a PPO-trained version of the same model. As shown in the table below, Q♯ consistently improves both pass@1 and maj1@8 for each policy. In particular, when applied to the PPO-trained policy, Q♯ reduces the KL divergence from the reference policy while further boosting accuracy. We also note a qualitative distinction: PPO improves pass@1 but slightly reduces maj1@8, indicating that its generations tend to be lower entropy and less diverse. Q♯, in contrast, improves both metrics while maintaining closer alignment with the reference.\n\n|Dataset|MATH||||\n|-|:-:|:-:|:-:|:-:|\n|Methods|$\\pi^{\\text{ref}}$|$\\pi^{\\text{ref}}$ + $Q\\sharp$|PPO|PPO + $Q\\sharp$|\n|pass@1 ↑|58.6|61.9|68.4|**71.1**|\n|maj1@8 ↑|72.8|**74.8**|72.4|73.4|\n|KL-Divergence ↓|-|**6.46**|69.52|60.53|\n\nIn terms of efficiency, Q♯ is significantly lighter to train. PPO requires approximately 20 hours on 4×H100 GPUs, whereas Q♯ training completes in roughly 5 hours on a single H100 GPU, thanks to its supervised learning objective and the use of a much smaller model. These findings suggest that Q♯ can effectively enhance performance while maintaining closer alignment with the reference policy, demonstrating its practical advantage as a complementary lightweight module.\n\n**Why Q♯ can fix shortcuts while REINFORCE and RPO cannot**\n\n**Authors' reply:** Thank you for the insightful question. The shortcut behavior, also known as the Clever Hans Trick [1], in the star-graph task arises directly from the auto-regressive **next-token prediction** objective. Specifically, the model minimizes loss by memorizing the first token seen during training and following the corresponding edge, achieving low training error but generalizing poorly at test time when the **initial** token is not provided. This leads to a brittle, shortcut-based policy.\n\nPolicy-based methods such as REINFORCE and RPO attempt to correct this by upweighting high-reward trajectories. However, because their loss is still based on the product of next-token prediction probabilities, the same as in pretraining, they are vulnerable to the same shortcut and require **exponentially many samples** via gradient descent on the policy to correct it once it is learned (Theorem 1 of [2]).\n\nIn contrast, Q♯ does not depend on myopic next-token supervision. Instead, it learns to predict the **cumulative reward-to-go** from each (prefix, token) pair under the reference policy, and uses this to guide generation toward optimal completions. This token-level value modeling allows Q♯ to predict future outcome and assign higher value to early tokens that lead to long-term reward. In other words, Q♯’s loss function is directly trained to perform *planning*, making it robust to the Clever Hans Trick [1] that undermines next-token–based methods. As shown in Appendix Figure 5, both Q♯ and CD are able to solve the star-graph task near-perfectly, while policy-based methods perform at random-guess level. We will incorporate this discussion into the revised version.\n\n> Also, what would be an example of such a shortcut (or hacky solution) in math? Is there an example from GSM8K or MATH where Q♯ uniquely fixes a mathematical shortcut that the policy-based methods fail to fix?\n\n**Authors' reply:** In our experiments on GSM8K and MATH, we did not observe specific, easily isolatable shortcuts of the same form as those seen in the star-graph setting. This may be due to the more flexible and open-ended nature of mathematical reasoning.\n\nTo provide intuition, we describe a contrived example inspired by integration by parts. Consider the problem $\\int x e^x , dx$, which requires selecting the appropriate decomposition into $u$ and $dv$. One path (e.g., choosing $u = x$, $dv = e^x dx$) leads to a clean and solvable integral. The other path (e.g., $u = e^x$, $dv = x dx$) results in a recursive or more complex expression. A model trained purely via next-token prediction loss may mimic observed solutions without internalizing the decision-making criteria, potentially choosing either path arbitrarily. In contrast, a value-based approach like Q♯ can learn to distinguish between such paths based on expected reward, consistently selecting the option that leads to successful solutions.\n\n**Benefits of Q♯ from qualitative examples**\n\n**Authors' reply:** Thank you for the detailed observation. Our goal with the qualitative examples are mainly to show that Q♯ can correct the reference policy’s mistakes while maintaining low KL divergence. As shown, Q♯ often starts with a similar prefix as $\\pi^{\\text{ref}}$, but when a divergence occurs, it typically avoids errors that $\\pi^{\\text{ref}}$ commits—such as arithmetic mistakes or flawed logical steps. This behavior reflects Q♯’s ability to assign higher value to correct tokens, thereby steering generation more effectively at critical decision points. We will clarify the description of qualitative examples more clearly in the final version.\n\n> Along similar lines, the star-graph planning task, GSM8K, and MATH actually seem to be quite different in nature. The star-graph task appears to require more explicit planning, while GSM8K and MATH seem to focus on reasoning. Because of this, showing a toy experiment on star-graph planning and then moving on to GSM8K and MATH feels a bit disconnected.\n\n**Authors' reply:** We agree with the reviewer that the star-graph task differs in nature from math reasoning benchmarks like GSM8K and MATH. Our intention in including the star-graph experiment is to use it as a controlled and interpretable environment to highlight a core limitation of next-token prediction loss. The star-graph setting clearly illustrates how value-based methods, such as Q♯, can overcome this limitation by considering long-term outcomes.\n\nWe then turn to GSM8K and MATH to evaluate Q♯ in more realistic and complex reasoning tasks. While these domains do not exhibit the same explicit planning structure as the star-graph, they still involve critical decision points where early mistakes can propagate—analogous in spirit to the planning failures observed in star-graph. Our empirical results across multiple model sizes demonstrate that Q♯ consistently outperforms CD in these benchmarks, reinforcing the broader applicability and effectiveness of our value-based approach.\n\n> How does Q♯ compare empirically to VAS?\n\n**Authors' reply:** Thank you for the question. CD and VAS are concurrent works that adopt the same core approach: both use the unregularized Q function $Q^{\\pi^{\\text{ref}}}$—i.e., the expected reward-to-go under $\\pi^{\\text{ref}}$ without KL regularization—to reweight the reference policy. As we show in Section 4.1, this choice leads to fundamental theoretical limitations: neither CD nor VAS guarantees convergence to the optimal KL-regularized policy, and both can suffer from suboptimal reward and excessive KL divergence.\n\nIn terms of implementation, CD and VAS differ only in minor aspects such as sampling strategy (e.g., top-k filtering), but these do not affect the underlying objective. As a result, we treat them as representative of the same class of methods, and use CD as a canonical baseline for empirical comparison. We will clarify this connection more explicitly in the final version of the paper.\n\n**Citations**\n\n[1] Bachmann, Gregor, and Vaishnavh Nagarajan. \"The pitfalls of next-token prediction.\" arXiv preprint arXiv:2403.06963 (2024).\n\n[2] Hu, Edward S., Kwangjun Ahn, Qinghua Liu, Haoran Xu, Manan Tomar, Ada Langford, Dinesh Jayaraman, Alex Lamb, and John Langford. \"The belief state transformer.\" arXiv preprint arXiv:2410.23506 (2024)."}, {"review_id": "pifYNa6Q7Q", "reviewer": "Reviewer_3hUn", "summary": "The authors introduce $Q\\sharp$ a novel RL algorithm for LLM post-training. $Q\\sharp$  works in deterministic MDPs and makes use of value learning. Unlike other value learning approaches, $Q\\sharp$  doesn't use a TD objective. The authors show bounds of the algorithm as well as empirical results. Indeed $Q\\sharp$ beats the Controlled Decoding algorithm (CD) with relative consistency (although the improvements are very marginal.)", "questions": "* Is the performance difference between CD and $Q\\sharp$ significant?\n* How does $Q\\sharp$ compare to other algorithms like GRPO on math reasoning tasks?\n* The authors note that their algorithm works for scenarios with deterministic transition dynamics. Can the authors highlight an environment with stochastic dynamics in the limitations to make it more intuitive for the reader what the limitations of the algorithm are?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n* The algorithm is novel and improves upon the baseline presented (CD).\n* Value estimation in an LLM-RL setting is an important and challenging topic.\n* PAC bounds are provided.\n* Algorithm tested with both Llama and Qwen family of models.\n\nWeaknesses:\n* There are some concerns around the robustness of the evaluations on GSM8k and MATH500. There are no mentions of the number of seeds used to train the models (I'm assuming just a single seed), and it's therefore impossible to know whether the difference in score is significant, or how robust the learning algorithms are.\n* It's also unclear how this method compares to other household algorithm, both in terms of training speed and final converged performance. Here, adding another baseline line like GRPO would add a lot of scientific value to the paper.\n\nGiven the above strengths and weaknesses I can only recommend borderline accept, but I'm happy to increase my score if the authors show results with more seeds and another baseline (ideally something like GRPO).", "was_revised": true, "final_justification": "Authors included extra evaluations and implemented an additional baseline, so I've increased my score to 5.", "rebuttal": "Dear Reviewer 3hUn,\n\nThank you for your thoughtful feedback and we respond to individual points below.\n\n> There are some concerns around the robustness of the evaluations on GSM8k and MATH500. There are no mentions of the number of seeds used to train the models (I'm assuming just a single seed), and it's therefore impossible to know whether the difference in score is significant, or how robust the learning algorithms are.\n\n> Is the performance difference between CD and Q♯ significant?\n\n**Authors' reply:** Thank you for raising this important point. For a direct comparison between CD and Q♯, we used the same trained model and varied only the value calculation used at inference. As a result, the primary source of randomness in evaluation comes from sampling variability during generation. To quantify this, we computed standard deviations via empirical bootstrapping over the GSM8K test set. For pass@1 and maj1@8, the standard deviations are approximately 0.7% and 0.3% respectively, suggesting that the performance gains of Q♯ over CD are still likely to be meaningful.\n\nTo further assess robustness and generality, we assess CD and Q♯ performance on two additional benchmarks: AIME-24, a more challenging math reasoning dataset, and QuALITY [1], a multiple-choice QA benchmark with context passages from Project Gutenberg. On AIME-24, we used Q♯ (Qwen 2.5 1.5B) trained on MATH to guide Qwen 2.5 7B. For QuALITY, we trained with two settings: Qwen 2.5 1B guiding Qwen 2.5 7B and Llama 3.2 1B guiding Llama 3.1 8B.\n\n|Dataset||AIME-24||\n|-|:-:|:-:|:-:|\n|Methods|$\\pi^{\\text{ref}}$|CD|$Q\\sharp$|\n|pass@1 ↑|9.3|13.5|**14.1**|\n|maj1@8 ↑|16.7|16.7|**20.0**|\n|KL-Divergence ↓|-|9.95|**9.23**|\n\n|$\\pi^{\\text{ref}}$||Qwen 2.5 7B|||Llama 3.1 8B||\n|-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Methods|$\\pi^{\\text{ref}}$|CD|$Q\\sharp$|$\\pi^{\\text{ref}}$|CD|$Q\\sharp$|\n|pass@1 ↑|64.5|64.2|**68.1**|73.5|75.1|**75.9**|\n|maj1@8 ↑|72.0|66.3|**73.3**|79.3|79.3|**81.1**|\n|KL-Divergence ↓|-|12.32|**7.90**|-|9.23|**8.88**|\n\nThese results across multiple datasets and model families provide evidence that Q♯ is generally robust and can consistently improve over both the reference policy and CD, while maintaining lower KL divergence.\n\n> It's also unclear how this method compares to other household algorithm, both in terms of training speed and final converged performance. Here, adding another baseline line like GRPO would add a lot of scientific value to the paper.\n\n> How does Q♯ compare to other algorithms like GRPO on math reasoning tasks?\n\n**Authors' reply:** Thank you for the helpful suggestion and we conduct additional experiments on the MATH dataset using PPO as a household algorithm to compare performance and training speed.\n\nIn these experiments, Q♯ is instantiated with a Qwen 2.5 1.5B model and used to guide both: (1) the original Qwen 2.5 7B reference model, and (2) the same model after PPO post-training. As shown in the table below, Q♯ consistently improves performance in both settings. Notably, when applied to the PPO-trained policy, Q♯ further increases accuracies and reduces the KL divergence — indicating improved alignment with $\\pi^{\\text{ref}}$ and better performance. The results suggest that Q♯ can also be complementary to policy-based approaches.\n\nWe additionally note a qualitative distinction: PPO improves pass@1 but slightly reduces maj1@8, indicating that its generations tend to be lower entropy and less diverse. Q♯, in contrast, improves both metrics.\n\n|Dataset|MATH||||\n|-|:-:|:-:|:-:|:-:|\n|Methods|$\\pi^{\\text{ref}}$|$\\pi^{\\text{ref}}$ + $Q\\sharp$|PPO|PPO + $Q\\sharp$|\n|pass@1 ↑|58.6|61.9|68.4|**71.1**|\n|maj1@8 ↑|72.8|**74.8**|72.4|73.4|\n|KL-Divergence ↓|-|**6.46**|69.52|60.53|\n\nIn terms of training speed, Q♯ is significantly lighter. PPO requires approximately 20 hours on 4×H100 GPUs, whereas Q♯ training completes in roughly 5 hours on a single H100 GPU, thanks to its supervised learning objective and the use of a much smaller model. These findings suggest that Q♯ can effectively enhance performance, demonstrating its practical advantage as a complementary lightweight module.\n\n> The authors note that their algorithm works for scenarios with deterministic transition dynamics. Can the authors highlight an environment with stochastic dynamics in the limitations to make it more intuitive for the reader what the limitations of the algorithm are?\n\n**Authors' reply:** One example environment with stochastic dynamics in the context of LLMs is multi-agent interaction, such as game playing with another external player. While the LLM can generate its own actions, the next state depends on the (potentially unpredictable) behavior of the other player, making the overall transition dynamics stochastic. We will include this example in the Limitations section of the final version to make the scope of our assumptions more concrete and intuitive.\n\nWe will update the final version of the paper with the above discussion and evaluation results. Please let us know if you have any further questions!\n\n**Citations**\n\n[1] Pang, Richard Yuanzhe, Alicia Parrish, Nitish Joshi, Nikita Nangia, Jason Phang, Angelica Chen, Vishakh Padmakumar et al. \"QuALITY: Question Answering with Long Input Texts, Yes!.\" In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 5336-5358. 2022."}, {"review_id": "RKtkT9WinA", "reviewer": "Reviewer_H6aA", "summary": "This manuscript proposes a novel value-based RL algorithm of Q# for LLM post-training. Concretely, Q# employes an offline way to collect data and update the value model, which yields more accurate value estimation and thus better performance. Theoretically imporvement are provided to show the rationale of Q#. Emprical study is conducted on Llama and Qwen verifies that Q# is superior to ref model and CD.", "questions": "This paper focuses on KL regularization in LLM training. However, many recent RL-based post-training algorithms for LLMs have removed the KL term, as it has been observed to hinder the learning of CoT reasoning [1]. The manuscript would be more convincing if the authors could empirically demonstrate the benefit of KL regularization in enhancing reasoning capabilities, particularly in long-form generation.\n\n&nbsp;\n\n\n[1] *Yu, Qiying, et al. \"Dapo: An open-source llm reinforcement learning system at scale.\" arXiv preprint arXiv:2503.14476 (2025).*", "limitations": "N/A", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n1. The paper introduces a novel value-based RL algorithm, Q\\#, designed to enhance LLM post-training.\n\n2. Theoretical analysis is provided to support the effectiveness of Q\\# and justify its design.\n\n3. The proposed method is evaluated on multiple LLM backbones, including LLaMA and Qwen, demonstrating its general applicability and superiority.\n\n\n## Weaknesses\n\n1. The paper is difficult to follow in its current form. It is unclear which baseline algorithm Q\\# builds upon, or how Q\\# is derived step by step. For example, the formulation in Eq.~(5) closely resembles that of DPO, yet the connection is not well explained.\n\n2. The experimental comparison is limited and lacks strong baselines. Q\\# is only compared against the reference model (pre-trained but not post-trained) and CD. It would significantly strengthen the paper to include comparisons with more widely-used RL methods such as PPO or GRPO.\n\n3. The empirical evaluation is restricted to two relatively simple math reasoning datasets (MATH and GSM8K). Additional evaluations on more challenging benchmarks such as AIME would provide stronger evidence for reasoning.", "was_revised": true, "final_justification": "The authors provided more empirical results and clarify the experimental settings.", "rebuttal": "Dear Reviewer H6aA,\n\nThank you so much for your review and constructive suggestions for us to improve the paper. We respond to individual points below.\n\n> The paper is difficult to follow in its current form. It is unclear which baseline algorithm Q# builds upon, or how Q# is derived step by step. For example, the formulation in Eq.~(5) closely resembles that of DPO, yet the connection is not well explained.\n\n**Authors' reply:** Thank you for this valuable feedback. We agree that better clarifying the derivation of Q♯ and its relationship to DPO will strengthen the paper. Below, we provide a clearer derivation and explain the connection and key differences.\n\n**Derivation of Q♯:** Q♯ is derived from the standard KL-regularized RL objective in deterministic MDPs, as given in Eq. (1). Following classical results [1], the optimal policy at step $h$ takes the form (Eq. (2)): $\\pi^{\\star,\\eta}\\_h(y|x) \\propto \\pi^{\\text{ref}}\\_h(y|x) \\exp(\\frac{1}{\\eta} Q^{\\star,\\eta}\\_h(x, y))$ where $Q^{\\star, \\eta}\\_h(x, y)$ is the soft action-value function. For deterministic MDPs, such as LLM decoding, we can recursively expand $Q^{\\star, \\eta}\\_h(x, y)$ using the deterministic transition function (concatenating token $y$ to the prefix $x$). This yields Eq. (4) and substituting back the optimal policy expression yields Eq. (5), which defines a token-level policy that reweights $\\pi^{\\text{ref}}$ by the expected exponentiated cumulative reward-to-go.\n\n**Connection to DPO [2] (and the H = 1 case):** While the form of Eq. (5) resembles DPO’s policy expression, the derivation and scope are fundamentally different. DPO starts with the same KL-regularized RL objective but, without exploiting the deterministic transition structure, it operates at the **sequence level**, corresponding to the **one-step case (H = 1)**. Its policy is \n$\\pi_r(y|x) = \\frac{1}{Z(x)} \\pi^{\\text{ref}}(y|x) \\exp(\\frac{1}{\\beta} r(x,y))$ where $y$ is a full completion and $Z(x)$ is the partition function over all possible sequences [2]. When $H = 1$, our $Q^{\\star}$ simply equals the reward $r$, and the DPO formula follows as a special case. Additionally, DPO partition function is intractable to normalize, typically requiring pairwise preference data to bypass it (Bradley-Terry modeling).\n\nIn contrast, Q♯ explicitly solves the **multi-step KL-regularized RL problem**. By leveraging deterministic MDP structure, we derive a tractable, closed-form expression for the optimal *token-level* policy in Eq. (5). Empirically, as shown in star-graph experiments (Appendix Figure 5), the difference between these paradigms matters in practice. DPO, being policy-based and sequence-level, cannot fix the shortcut behavior, while value-based methods like Q♯ achieve near-perfect accuracy by explicitly modeling the reward-to-go and guiding token-level decisions.\n\nIn summary, Q♯ is not built upon DPO or any other baseline. It is an independent derivation from first principles of KL-regularized RL in deterministic MDPs [1,3-5]. Compared to DPO, Q♯ addresses the more general and challenging $H > 1$ case, and to the best of our knowledge, it is the first algorithm for LLM post-training that provably guides the reference policy using the optimal KL-regularized Q-function in a tractable token-level form. We will revise the paper to make these distinctions and derivation clearer.\n\n> The experimental comparison is limited and lacks strong baselines. Q# is only compared against the reference model (pre-trained but not post-trained) and CD. It would significantly strengthen the paper to include comparisons with more widely-used RL methods such as PPO or GRPO.\n\n**Authors' reply:** Thank you for the helpful comment! Following your suggestion, we conduct additional experiments on the MATH dataset using PPO as a baseline.\n\nIn these experiments, Q♯ is instantiated with a Qwen 2.5 1.5B model and used to guide both: (1) the original Qwen 2.5 7B reference model, and (2) the same model after PPO post-training. As shown in the table below, Q♯ consistently improves performance in both settings. Notably, when applied to the PPO-trained policy, Q♯ further increases accuracies and reduces the KL divergence — indicating improved alignment with $\\pi^{\\text{ref}}$ and better performance. The results suggest that Q♯ can also be complementary to policy-based approaches.\n\nWe additionally note a qualitative distinction: PPO improves pass@1 but slightly reduces maj1@8, indicating that its generations tend to be lower entropy and less diverse. Q♯, in contrast, improves both metrics.\n\n|Dataset|MATH||||\n|-|:-:|:-:|:-:|:-:|\n|Methods|$\\pi^{\\text{ref}}$|$\\pi^{\\text{ref}}$ + $Q\\sharp$|PPO|PPO + $Q\\sharp$|\n|pass@1 ↑|58.6|61.9|68.4|**71.1**|\n|maj1@8 ↑|72.8|**74.8**|72.4|73.4|\n|KL-Divergence ↓|-|**6.46**|69.52|60.53|\n\nIn terms of efficiency, Q♯ is significantly lighter to train. PPO requires approximately 20 hours on 4×H100 GPUs, whereas Q♯ training completes in roughly 5 hours on a single H100 GPU, thanks to its supervised learning objective and the use of a much smaller model. These findings suggest that Q♯ can effectively enhance performance, demonstrating its practical advantage as a complementary lightweight module.\n\n> The empirical evaluation is restricted to two relatively simple math reasoning datasets (MATH and GSM8K). Additional evaluations on more challenging benchmarks such as AIME would provide stronger evidence for reasoning.\n\n**Authors' reply:** Thank you for the suggestion. We conduct additional evaluation on AIME-24 using the same Q♯ model checkpoint (Qwen 2.5 1.5B) trained on MATH to guide the Qwen 2.5 7B reference model, without any further finetuning. The results are shown below. We observe that Q♯ significantly improves performance over the reference policy, despite being trained on a different dataset, indicating generalization capability. Furthermore, Q♯ consistently outperforms the CD baseline in both accuracy metrics while maintaining a lower KL divergence, supporting its effectiveness.\n\n|Dataset||AIME-24||\n|-|:-:|:-:|:-:|\n|Methods|$\\pi^{\\text{ref}}$|CD|$Q\\sharp$|\n|pass@1 ↑|9.3|13.5|**14.1**|\n|maj1@8 ↑|16.7|16.7|**20.0**|\n|KL-Divergence ↓|-|9.95|**9.23**|\n\nIn addition, following Reviewer 2Dby’s suggestion on studying other potential applications of Q♯ beyond math reasoning, we conducted experiments on QuALITY [6], a multiple-choice QA benchmark with context passages 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 QA tasks for both Qwen and Llama.\n\n|$\\pi^{\\text{ref}}$||Qwen 2.5 7B|||Llama 3.1 8B||\n|-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Methods|$\\pi^{\\text{ref}}$|CD|$Q\\sharp$|$\\pi^{\\text{ref}}$|CD|$Q\\sharp$|\n|pass@1 ↑|64.5|64.2|**68.1**|73.5|75.1|**75.9**|\n|maj1@8 ↑|72.0|66.3|**73.3**|79.3|79.3|**81.1**|\n|KL-Divergence ↓|-|12.32|**7.90**|-|9.23|**8.88**|\n\n**KL regularization may hinder CoT reasoning**\n\n**Authors' reply:** Thank you for raising this and highlighting concurrent work such as DAPO [7], which we will discuss and cite in the final version. We agree that the role of KL-regularized RL in post-training—particularly for long-form CoT reasoning—remains nuanced, and that in some settings, strong constraints on the reference policy could limit the exploratory behavior needed for emergent reasoning capabilities.\n\nNevertheless, KL-regularized RL problem remains a valuable and widely used formulation that underpins many practical algorithms, including DPO [2], CD [8], and VAS [9]. Our goal in this paper is to revisit this fundamental setting and address it in a principled and provably optimal way. Prior algorithms such as CD and VAS can be viewed as approximate or sub-optimal methods for KL-regularized RL, while Q♯ directly solves for the optimal solution in a tractable, token-level form.\n\nFrom a theoretical perspective, KL regularization plays a crucial role in tractability. In particular, it avoids the need for Bellman completeness assumption, which is typically required for convergence in standard (unregularized) RL [10]. Our analysis shows that under KL regularization, one can obtain provable convergence under the much milder realizability condition. Thus, even when the function class is far from Bellman complete, KL regularization offers favorable theoretical properties. We will clarify our scope and motivation in the final version, and thank the reviewer again for highlighting this distinction.\n\n**Citations**\n\n[1] Ziebart, Brian D., et al. \"Maximum entropy inverse reinforcement learning.\" Aaai. Vol. 8. 2008.\n\n[2] Rafailov, Rafael, et al. \"Direct preference optimization: Your language model is secretly a reward model.\" Advances in neural information processing systems 36 (2023): 53728-53741.\n\n[3] Domingo-Enrich, Carles, et al. \"Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control.\" arXiv preprint arXiv:2409.08861 (2024).\n\n[4] Piché, Alexandre, et al. \"Probabilistic planning with sequential monte carlo methods.\" International Conference on Learning Representations. 2018.\n\n[5] Li, Xiner, et al. \"Derivative-free guidance in continuous and discrete diffusion models with soft value-based decoding.\" arXiv preprint arXiv:2408.08252 (2024).\n\n[6] Pang, Richard Yuanzhe, et al. \"QuALITY: Question answering with long input texts, yes!.\" arXiv preprint arXiv:2112.08608 (2021).\n\n[7] Yu, Qiying, et al. \"Dapo: An open-source llm reinforcement learning system at scale.\" arXiv preprint arXiv:2503.14476 (2025).\n\n[8] Mudgal, Sidharth, et al. \"Controlled decoding from language models.\" arXiv preprint arXiv:2310.17022 (2023).\n\n[9] Han, Seungwook, et al. \"Value augmented sampling for language model alignment and personalization.\" arXiv preprint arXiv:2405.06639 (2024).\n\n[10] Munos, Rémi, and Csaba Szepesvári. \"Finite-Time Bounds for Fitted Value Iteration.\" Journal of Machine Learning Research 9.5 (2008)."}]

Author rebuttal:

Meta-review: This paper introduces Q#, a novel and theoretically principled value-based algorithm for the KL-regularized RL post-training of LLMs. The core contribution is leveraging distributional RL to directly learn the optimal KL-regularized Q-function, which provides a more sound objective than prior value-based methods like CD and VAS. The reviewers all agreed that the paper has a strong theoretical grounding , where the authors provided PAC bounds under a mild realizability assumption. Some concerned were raised regarding the empirical evaluation, requesting comparisons to stronger policy-based baselines like PPO and broader testing on more diverse and challenging datasets. The authors addressed these points in their rebuttal, with new experiments on the AIME and QuALITY benchmarks and a direct comparison against PPO. Therefore, I recommend accepting this paper, contingent on the authors incorporating the valuable new experiments, analysis, and clarifications from their rebuttal into the final camera-ready version.
Decision: Accept (poster)