ShiQ: Bringing back Bellman to LLMs

We extend our gratitude to the reviewer for their valuable time and insightful comments and questions.

Weaknesses

Developing a backward recursion scheme for KL regularised Q learning is not necessary a significant contribution. 

The proof to our knowledge has not been done with $T_{max}$ as a random variable, although the proofs are similar with the fixed horizon (or max sequence lenght).

The authors do not expand upon how a reliable token-based reward model can be assumed when in practice this token-based feedback is generally not available.

and 

 q1) I would appreciate the authors comment on the merits of the proposed approach given the lack of token-based rewards. Perhaps the proposed approach provides a principled approach for credit assignment. If that is the case the right benchmarking comparison should be other approaches credit assignment rather than sequence-based LLM training methods.

Even though we do not have access to token-level rewards, thanks to the multi-turn aspect of our algorithm, we manage to propagate information and obtain some signal. For example, in HH, the reward model only provides a reward for the entire trajectory, and it works well. Finally, even if there are no datasets with rewards per token, we have multi-step rewards in BFCL, which allows for granularity in ShiQ and enables us to benefit from it.

It appears the computational effort of the shifted- Q algorithm grows exponentially in T_max. This may restrict its practical application to short sequences.

and 

  q2) Can the authors elaborate on the computational burden imposed by the length limit T_max ? Could there be ways to mitigate this ?

We disagree that ShiQ grows exponentially with $T_{max}$. If you are referring to Equation 10 with the two sums, it is possible to avoid an explosion in the implementation.

Our loss is an external sum of cumulative sums. In the first pass, we calculate the internal cumulative sum starting from the index at $T_{max}$. Each term, when traversing the index in reverse, is simply the previous term plus one term. At the end, we sum all the terms. Therefore, we do not have an explosion of loss terms, and the complexity is linear in $T_{max}$, not exponential (for calculating the loss terms, assuming they are already computed). However, we do pay in memory because we need to store the terms, but again, there is no exponential complexity!

We would like to express our sincere gratitude to the reviewer for their valuable time, insightful comments, and thoughtful questions. Weakness: (4 points)

w1) While the method appears theoretically sound, the derivation involves multiple transformations and notations that make it overly complex. It's challenging to grasp the full intuition behind each step.

Thanks for your remark. We tried to keep the notations as simple as possible for the RL side. We also included graphs to explain our approach just before Section 2.1, in the main text, and all proofs are self-contained. Each paragraph first provides the intuition behind our approach and then attempts to be more precise about exactly what we are doing. We would appreciate advice on any sections that may be complicated for the reader, to improve the readability of our paper

w2)The empirical evaluations are a major weakness.
Most experiments are on toy datasets, which doesn't fully demonstrate the method's effectiveness in real-world, large-scale LLM scenarios.

We kindly disagree on this point. There is extensive evaluations with lots of ablations for our method. The toy experiments are for giving intuitions. ( ie bandit and MPD setting) Then, there is experiments on 3 differents environments with HH, UltraFeedback and BFCL-v3 which is multi-turn setting so we do not agree that is not large scale LLM scenario, especially for BFCL-V3.

Crucially, the experimental figures lack error bars, making it difficult to assess the statistical significance and robustness of the results.

We will definitely add it for bandits and MDPs. For LLMs, a seed is very costly, and there are only a few papers that can afford to include error bars with agentic benchmarks.

W3)The significance of the results is questionable. In Figures 3 and 4, ShiQ's performance is comparable to, or sometimes even slightly worse than, the baseline CoPG. The benefits claimed in the introduction (e.g., sample efficiency) are not clearly demonstrated by the empirical results.

CoPG is generally a strong baseline. Then, in Figure 4, under the multi-turn setting, our method, ShiQ, clearly outperforms. Moreover, in terms of the Pareto front, our method is superior in both Figure 3 (single-turn setting) and Figure 4 (multi-turn setting).

W3)The paper's clarity suffers from inconsistent notation. It mixes RL and LLM notations throughout, which is confusing. A unified and simplified notation would greatly improve readability.

First, we introduce a rigorous notation that accounts for dense rewards at each step, similar to reinforcement learning. These notations are necessary to derive all the proofs. Then, in a second step, we simplify the notations for practitioners so they can easily implement the loss. The notations are simplified insofar as the reward is sparse and only occurs at the end of the sequence. At no point notations are inconsistent. Moreover, we always clarify the connection between the two notations.

From the three sub-methods proposed (Easing sampling, Improved initialization, Multi-step extension), which one do you believe contributes most significantly to ShiQ's performance, and why?

All points in our method are important, as emphasized in the appendix with the ablation study on the initialization shift and the multi-step aspect when the reward is sparse (see Figures 5 and 6 in the appendix).

The easy sampling aspect is very important because otherwise, we would need to load two models during sampling (and having same temperature), which is extremely costly when scaling up, and this is not acceptable. The idea of the paper is to have the LLM policy as the softmax of the logits, so a simple formulation at the infrastructure level. As highlighted in Figure 5 of the appendix, Shiq no ms (for no multistep) has no signal and does not increase the reward, so it is very important in an environment with sparse rewards. Perhaps the least important is the shift, as we can still achieve performance, but we always have this initialization problem for learning.

Beyond empirical performance, what are the primary theoretical advantages of your Q-learning-based approach compared to traditional policy-gradient methods in the context of LLM fine-tuning?

The advantages are twofold. First, compared to classical policy gradients, our method has theoretical convergence guarantees to the best policy in the offline setting. Second, it allows for fine-grained allocation, which is not the case with classical policy gradients

How would you address the scaling challenges if ShiQ were to be applied to significantly larger LLMs and datasets? What potential issues might arise in a large-scale deployment, and how would you mitigate them?

The whole idea of the paper is to avoid excessive computational burden, such as a value network, etc. Compared to a policy gradient, it typically fits on the same GPU partition because we only load one LLM without a target network. Therefore, our method scales well with the size of datasets, at least as well as methods like GRPO, CoPG, etc.

From a computational perspective, the loss is similar to cross-entropy when you think about it. Instead of a single term over the trajectory, we have a sum of cumulative inverse sums, which is fundamentally like cross-entropy. If implemented correctly, we calculate the cumulative sum first, starting from the term with the least value, i.e., at index T_max, and sum at the end. This does not cause the loss calculation to explode exponentially with $T_{max}$ but is $O(T_{max})$ for computing the loss terms, assuming the terms are already calculated.

Furthermore, when comparing to policy gradients, we sample multiple times to obtain a baseline in policy gradients, which can be very costly. This is not the case in our method.

Is ShiQ more suitable for preference learning tasks (like aligning with human preferences) or for fine-tuning LLMs on structured reasoning tasks, similar to chain-of-thought (CoT) training in DeepSeek-R1?

The question for us is not about the number of turns but rather depends on the granularity of the signal. If, in the CoT, we have information for each segment with a reward, then ShiQ makes complete sense. If we are simply doing reference learning in a single turn, then a bandit algorithm like DPO is better suited for simple LLMs, as there is no reason for it not to work well. If we have intermediate reward information, such as in code with passing tests, then ShiQ has a better chance of functioning effectively.

We thank the reviewer for their time and for their pertinent remarks and questions.

ShiQ's final formulation closely resembles DRO [1], without significant improvements.
The author claims two main improvements over DRO:

(1) Token-level and multi-step rewards: ShiQ handles these, but DRO could also be trivially extended by replacing single-step rewards with multi-step returns.

We disagree that ShiQ is similar to DRO for several reasons. DRO operates in a bandit setting and is single-step. In our paper, we use MDPs, and adapting from a bandit to an MDP DRO algorithm is not straightforward. Furthermore, we do not agree that there is no significant improvement, especially in the multi-turn case (see Figures 3 and 4).

(2) Memory-friendly value modeling: ShiQ replaces DRO's separate value network with the proposed value difference method, but no experiments compare the performance impact of value difference versus an additional value network.

q3) For smaller models (e.g., 7B/4B), where memory constraints are less severe, would using an additional value network, a tiny shared-parameter network, or the proposed value difference yield different results? Is an extra value network still necessary for optimal performance?

Multiplying large networks like LLMs is undesirable, as it strains hardware resources and leads to wasteful memory consumption. Additionally, this approach is challenging to implement, as noted in DRO, where they mention the need for many tricks, such as policy learning rate rescaling, and that parameter sharing was detrimental to good empirical results, thus requiring a separate network. This separate network also necessitates a corresponding optimizer state, which by default takes twice the network memory, unless a more complex optimizer is used.

For all these reasons, we chose to design an algorithm without target networks. Furthermore, to ensure a fair comparison, we use DRO-V (the version without target networks) from the DRO paper as a baseline.

q1) For the implementation details, during optimization, are  $\pi_l$ and $v_l$ in Eq.10 updated simultaneously? Is stop-gradient applied? If the gradient of ShiQ can be provided in the paper, it may be more convenient to understand the optimization objective.

Thank you for your question. This is a very good point. Indeed, $v_l$ and $\pi_l$ are updated simultaneously, as the loss is only on the logits of the policy. Additionally, there is no stop-gradient in our loss, because it is a residual approach. Regarding the computation of the gradient, we completely agree with your point and will add the gradient computation to our paper.

 q2) In the multi-turn experiments, does DRO use single-step rewards or multi-step returns? If DRO is modified to use multi-step returns, would it be equivalent to ShiQ (ignoring value modeling differences)? What would the empirical results show?

In multi-turn experiments, DRO uses a multi-step reward, which is the sum of the rewards. This is not equivalent to ShiQ, as ShiQ uses a sum of square terms rather than a square of a sum in the multi-turn setting of DRO. Additionally, ShiQ includes a third shift at initialization that DRO does not have. Finally, we do not use target networks. For all these reasons, our approach differs from DRO.

q4) Minor Issue: The subscripts in Equation 9 are incorrect？
    

Thanks for pointing this. You are totally write we will correct typos.