Q-SFT: Q-Learning for Language Models via Supervised Fine-Tuning

Authors: Joey Hong, Anca Dragan, Sergey Levine
Primary area: reinforcement learning
Keywords: offline reinforcement learning, language models, dialogue, robotics

Abstract:
Value-based reinforcement learning (RL) can in principle learn effective policies for a wide range of multi-turn problems, from games to dialogue to robotic control, including via offline RL from static previously collected datasets. However, despite the widespread use of policy gradient methods to train large language models for single turn tasks (e.g., question answering), value-based methods for multi-turn RL in an off-policy or offline setting have proven particularly challenging to scale to the setting of large language models. This setting requires effectively leveraging pretraining, scaling to large architectures with billions of parameters, and training on large datasets, all of which represent major challenges for current value-based RL methods. In this work, we propose a novel offline RL algorithm that addresses these drawbacks, casting Q-learning as a modified supervised fine-tuning (SFT) problem where the probabilities of tokens directly translate to Q-values. In this way we obtain an algorithm that smoothly transitions from maximizing the likelihood of the data during pretraining to learning a near-optimal Q-function during finetuning. Our algorithm has strong theoretical foundations, enjoying performance bounds similar to state-of-the-art Q-learning methods, while in practice utilizing an objective that closely resembles SFT. Because of this, our approach can enjoy the full benefits of the pretraining of language models, without the need to reinitialize any weights before RL finetuning, and without the need to initialize new heads for predicting values or advantages. Empirically, we evaluate our method on both pretrained LLMs and VLMs, on a variety of tasks including both natural language dialogue and robotic manipulation and navigation from images.

Decision: accept (poster)
Reviewer scores: 8, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730551700925

thread_type:
official_review

<thread_object>
creator: reviewer_zr9o

<content>
summary:
This paper proposes Q-SFT, a new approach that enables fine-tuning of LLMs and VLMs within an offline RL framework but using a supervised learning-style loss function. The method matches or outperforms previous approaches across a range of LM fine-tuning benchmarks.

strengths:
- The motivation is clear: the instability of the Q-learning objective is a well-known challenge in the RL literature.
- The proposed method is straightforward and demonstrates strong empirical performance, with improvements in sample efficiency and scalability over previous methods.
- The evaluation covers a wide range of benchmarks.

weaknesses:
It seems there are some limitations to the proposed approach that are not addressed in the paper.
- During inference, the approach requires double the number of forward passes—one for the original pretrained LM and one for the fine-tuned LM. This can be challenging in practical scenarios where real-time decision-making is essential, such as robotic manipulation.
- In Theorem 4.1., the lower bound for the learned p_theta is Q* multiplied by pi_beta. If pi_beta is small, the resulting pi_theta can be overly conservative compared to the true Q*.

questions:
.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730577520296

thread_type:
official_review

<thread_object>
creator: reviewer_chgq

<content>
summary:
This paper proposes a new RL training method for LLMs that take advantage of the pretraining. It is based on Q-learning, but instead of attaching a new value head, the method uses the existing token probabilities as value function. There is some theoretical guarantees and diverse experimental results involving multi-steps interactions.

strengths:
- Novelty: the method seems to be quite novel. The way it works is quite different from a typical Q-learning where L2 loss is used. Instead it relies on a cross-entropy loss that is modified. Given the importance of RLHF in LLMs and the need of multi-step interaction, a study like this can be impactful and important.
- The experiments are quite diverse, ranging from LLM-based text tasks to VLM-based robotic manipulation. It is also compared to a diverse set of baselines, which helps to solidify the claims.
- There is a proof about a theoretical guarantee, which is always nice to have. The method also leads to an improvement in performance in practice too, where the gap increasing as the model scales. Given the lack of value-based approaches in LLMs, it would be interesting to see if this method will adopted widely.
- Paper is well written and easy to follow.

weaknesses:
- I felt like the writing of the experimental results were a bit too short. It would have been good to have more ablation studies to understand why the method works, and in which situations it is more optimal. For example, since the main motivation is that the value head is not initialized from scratch, what will happen if we do that but keep the loss the same. Also there is not much on the training details, such as the number of training steps etc. I think the other parts can be condensed to make space. 
- Another important part that was too short was the method itself. The method is introduced between L230-261, which is about only half page. It doesn’t give enough explanation about why this method should work, so it was quite hard to understand it. For example, can you give more insight into why the proposed method should be more stable? 
- Reuse of the pretrained weights is emphasized as the main motivation, but some of the tasks actually train the model from scratch and others are not in natural language, which is a bit conflicting. There is still an improvement when trained from scratch, which is good, but there is a lack of explanation why it works better.

questions:
- Some typos: L39 citation, L84 sentence, L194 “performing”, Fig2 top left, 
- Eq2 should have a negative sign? Also tab1 -2.08 should be bold, not -2.11. 
- The figure 1 was bit too small and hard to read.
- L153 defines RL as MDP, but what about POMDP?
- About the theoretical guarantee, how good is the lower bound? A probability from a policy can be quite low for large vocabulary, in which case then the bound is not that tight?
- How many models needed to be loaded during training? There is the main model, also a target model and the original model. So is that 3? How does that affect memory usage?
- Given that some baselines use different base models, it would be nice to put them in the table. This will help if they perform better because their base model is stronger or not.
- In fig3, the error bars correspond to different seeds?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730697354160

thread_type:
official_review

<thread_object>
creator: reviewer_wnh5

<content>
summary:
This paper proposes a method that leverages the LLM logits learned via supervised fine-tuning to approximate the Q-learning objective. Authors argue that traditional Q-learning setup suffers from discarding the logit head, and learning a new head for Q values, so the proposed method directly translates the learned logits to Q values, with theoretical analysis on its bound and approximation. Experiments include text-based games, image-based navigation and robotic manipulation, and show some benefits of the proposed method over the prior value-based RL method.

strengths:
- Using probability prediction in the Q-learning objective is cool, and it sounds particularly beneficial for large vocabulary/action space. 
- The implementation of the method is straightforward: learning two models using different objectives, and use both at inference time

weaknesses:
- I found the pitch very confusing: it emphasizes multi-turn problems like dialogue, but the actual math formulation treats each turn as a single token in an utterance (i.e., the action space is the vocabulary space). The generation process is to generate an utterance token by token, which is conventional, and it’s confusing where the multi-turn setup plays a role other than using this generative model repeatedly to generate different utterances at different parts of the dialogue. In other words, the value function here is not the reward per utterance as stated in the intro (line 074), and question answering is also not a single-step problem as the answer is also generated token by token.
- Experimental setup seems simplified in terms of 1) for language-based tasks, the game nature makes the effective vocabulary very small and unambiguous, 2) the well-defined reward function setup in text games do not translate into practical problems, 3) why no supervised learning results from prior work as a direct comparison?; 4) most of the datasets already have RL method outperform supervised learning which is contradictory to the intro “offline RL fine-tuning of LLMs does not reliably outperform supervised fine-tuning”. Therefore, it’s unclear whether the improvement brought by the proposed method should be contrasted to what exactly.
- The writing can be further improved in terms of clarity, accuracy, and fixing typos. A few examples below: 
   - Table 1: number from the proposed method is bold in the Wordle column, but it’s not the best performance method
   - Define h used in line 162
   - Prime should be on a, instead of argmax in line 218
   - Consistently use “RL” instead of having “RI” occurred 
   - Citation should be within a pair of parentheses: line 037, 039

questions:
Please see the previous section.

Response to the rebuttal: (posted here since the rebuttal ended)

Thanks for all the clarification and discussion -- I raised my scores accordingly. It would be very valuable to revise the paper to address the points we discussed, especially the clear definition on step vs turn, the potential to scale better than existing RL approaches, the setup and comparable work of SFT baselines, motivation to alleviate the downsides of using RL, etc.. The rebuttal was helpful to better understand this work, and it would definitely be better if readers only need to read the paper itself to sufficiently appreciate this work.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700390107

thread_type:
official_review

<thread_object>
creator: reviewer_oykw

<content>
summary:
This paper proposed Q-SFT, a method that integrates Q-learning with SFT for LLMs in the multi-turn RL setting. To exploit both the representation and the logits in LLMs learned in the pretraining stage and bypass the poor scaling effect of TD-style objectives, the authors proposed a novel weighted cross-entropy loss that embeds the learning of the Q-value into the weights. In this way, the logits prior in the pretrained LLMs can be leveraged, and no new head is needed to learn the Q-values. Theoretical analysis has shown the guarantee of the new learning objective leading to a conservative estimation of the value function. Experiments on several scenarios have demonstrated the effectiveness and efficiency of the proposed method.

strengths:
- The work targets at the multi-turn RL scenario for LLMs, which is vital for the development of complex reasoning and agentic use in LLMs.
- The proposed algorithm is backuped with a solid theoretical motivation.
- The experimenting scenarios have a broad coverage, and the performance is well-demonstrated.

weaknesses:
Please see the Questions below.

questions:
- While the theoretical analysis shown in Section 4.3 demonstrates that the Q value learned with Equation (3) has the bounds independent from the hyperparameter \gamma, I wonder what the ablating effect \gamma brings in practice: Will certain configuration of \gamma contribute the learning of a better Q-value estimation? This question has another motivation: In Equation (3), some of the probability mass is allocated to the dummy actions a', and the weight ratio between the actions in the demonstration trajectory and the dummy ones has one degree of freedom determined by \gamma. It would therefore be great to sweep \gamma to study the ablation effect of the weight ratio.
   - And by the way, how is the \gamma set in the demonstrated experiments?

- According to Equation (3), learning the Q-value function requires a learned behavior policy \pi_\beta in prior. How will the quality of the behavior policy impact the learning of the Q-values, and the performance of the ultimate policy? While I appreciate the experiments in Figure 4 where only 10% of the training data are used to get the \pi_\beta, I wonder how the performance grows with the quality of the initial behavior policy. For example, if the initial behavior policy is near perfect, will it be that the ultimate policy, although weighted with the learned Q-values, performs similarly with the initial one?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732693867360

thread_type:
public_comment

<thread_object>
creator: ~tadashi_kozuno1

<content>
title:
Questions about the proof of Theorem 4.1

comment:
Hi. Thank you very much for the interesting paper. While reading the proof of Theorem 4.1, I was not able to follow some reasonings. It is very much appreciated if the authors could clarify.

1. How did you derive the inequality just below Eq (6) from Eq (6)? I guess $\mathfrak{B}^\star \widehat{p}^k$ is assumed to be less than 1, but it is not clear for me why it is less than 1.
2. What is exactly meant by "taking the fixed point" in the proof? For example, in Line 861, simply expanding $\widehat{p}^k / \pi_\beta$ according to the inequality in Line 861 does not seem to lead to the conclusion in Line 863.

Regarding the first question, I think $\mathfrak{B}^\star \widehat{p}^k \leq 1$ does not hold in general for any $k$. I can construct a counterexample with a two-state two-action MDP and uniform $\widehat{p}^1$. (I might be wrong though.)

Regarding the second question, suppose a simple case where both rewards and behavior policy are constant (say $r$ and $\pi$ everywhere). Then simply expanding $\widehat{p}^k / \pi_\beta$ in RHS of Line 861 according to the inequality in Line 861, I end up with something like $\sum_{n=1}^{k} \gamma^k \pi^{-k} r$, which diverges if $\gamma > \pi$. Am I missing something?
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a new way to do RL in multi-turn dialogue setting where the use of RL is less common compared to the vast success of RL in single-turn case. The proposed approach frames Q-learning as a cross-entropy problem with suitably selected weights such that the trained policy can be used to recover conservative bounds on the Q-value. 

While I think the mathematical reasoning here is not important, the reviewers appreciated the mathematical analysis. However, upon going over the mathematical derivations in the paper I noticed that certain steps were not justified. While equation 5 is correct, the next equation isn't justified. It along with the WCE loss assumes that $\mathcal{B}^\star p_\theta$ is in $[0, 1]$ but it isn't clear why this is the case even with assumption 4.1. In particular, if for some $a'$ we have $\pi_\beta(a' \mid s') = 0$, then for a uniform $p_\theta$ we have $\mathcal{B}^\star p_\theta = \infty$. Note that the public comment also points this out. Now, if we assume $\mathcal{B}^\star p_\theta \in [0, 1]$, then we get can complete the proof. These steps are missing in the proof but are as follows:

$\alpha^{(k+1)}(s, a) \ge r(s, a) + \gamma E_{s' \sim P(. | s, a)}[ \max_{a'} \alpha^{(k)}(s', a')]$,

where $\alpha^{(t)}(s, a) = \frac{p^{(t)}(a \mid s)}{\pi_\beta(a \mid s)}$ for any $t \in \mathbb{N}$. Now if you assume convergence that $\lim_{t\rightarrow \infty} \alpha^{(t)}(s, a) = \alpha(s, a)$, we get:

$\alpha(s, a) \ge r(s, a) + \gamma E_{s' \sim P(. | s, a)}[ \max_{a'} \alpha(s', a')]$,

using Bellman optimality, 
$Q^\star(s, a) = r(s, a) + \gamma E_{s' \sim P(. | s, a)}[ \max_{a'} Q^\star(s', a')]$

and defining $\delta(s, a) = \alpha(s, a) - Q^\star(s, a)$ and $\delta = \inf_{s, a}\delta(s, a)$, one gets:

$\delta(s, a) \ge \gamma E_{s' \sim P(. | s, a)}[ \max_{a'} \alpha(s', a') - \max_{a'} Q^\star(s', a')] \ge \gamma E_{s' \sim P(. | s, a)}[ \min_{a'} \left( \alpha(s', a') - Q^\star(s', a')\right)] \ge \gamma \delta$,

where the last line uses $\max_x f(x) - \max_x g(x) \ge \min_x (f(x) - g(x))$. Finally, taking inf on LHS means that $\delta \ge \gamma \delta$, which implies $\delta \ge 0$ if $\gamma \in (0, 1)$. Or, $\alpha(s, a) \ge Q^\star(s, a)$. I'd request authors to add justification for why $B^\star p \in (0, 1)$ and these missing steps.

Strengths:
1. Method is simple
2. Experiments are presented for several domains
3. Multi-turn LLM reasoning is an important topic and the use of RL for this has been somewhat less successful. This approach proposes a solution for this.
4. Reviewers appreciated the theoretical result but see below.

Weakness:
1. The theoretical results are at best unjustified and at worst potentially incorrect
2. Requires 3 copies of the model. This makes it expensive in practice as even 2x increase in memory is costly and makes a given approach undesirable. This is important especially given the motivation of the authors to make the Q-learning approach user-friendly.
3. The reviewer ChGQ requested results where the value head is randomly initialized but the authors did not include this study in response. I understand that experiments can be hard to complete during the rebuttal period but this one would have been nice to see.

Overall, my biggest concern is the mathematical foundation of the approach. However, given that the approach works well empirically and is simple and should easy to try, I am leaning towards a weak accept.