Q-SFT: Q-Learning for Language Models via Supervised Fine-Tuning (2411.05193v2)

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 LLMs 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 LLMs. 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 LLMs, 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.

• The paper introduces Q-SFT, an offline RL algorithm that combines supervised learning with Q-learning by treating token logits as Q-values, avoiding non-stationary targets.
• Empirical evidence shows Q-SFT outperforms SFT and value-based RL in tasks like language games and robotic control, demonstrating better integration of pre-training knowledge.
• Q-SFT offers practical benefits by fine-tuning LLMs without architecture changes, facilitating efficient deployment in complex applications like dialogues and navigation.

Overview of Q-SFT: Q-Learning for LLMs via Supervised Fine-Tuning

The paper introduces Q-SFT, an innovative algorithm that adapts offline reinforcement learning (RL) principles to fine-tune LLMs by drawing parallels with supervised fine-tuning (SFT). The proposed method incorporates Q-learning into the supervised learning framework by conceptualizing Q-values as probabilities derived from token logits in trained LLMs. This technique is designed to address the scaling and instability issues observed in traditional value-based RL applications to large models while preserving the pretraining benefits of LLMs. The results indicate that Q-SFT delivers comparable, if not superior, performance to state-of-the-art RL and SFT strategies while maintaining a theoretically solid foundation.

Key Contributions

1. Algorithm Innovation: The authors propose an offline RL algorithm that ingeniously combines supervised and Q-learning methods. By modifying the traditional supervised learning objective with weighted cross-entropy, the Q-SFT approach estimates Q-values without regressing from non-stationary targets.
2. Empirical Evidence: The paper provides empirical results across several domains, indicating that Q-SFT can outperform both contemporary SFT and value-based RL algorithms in tasks involving both LLMs and VLMs. This performance superiority is noticeable in diverse tasks, ranging from natural language games to robotic manipulation.
3. Theoretical Foundations: The theoretical backing for the Q-SFT method claims performance bounds competitive with state-of-the-art algorithms, based on the properties of learned likelihoods as conservative estimates of authentic Q-values.
4. Practical Implementation: Q-SFT can be implemented without altering the model architecture of pre-trained LLMs or VLMs, an advantage over existing methods requiring additional value heads or weight reinitialization.

Strong Numerical Results and Claims

• The introduction of Q-SFT purportedly results in a demonstrable improvement across several task metrics. For example, the paper claims enhanced outcomes in language games, with Q-SFT outperforming other methods by significant margins in some tasks, such as chess and wordle. Specifically, the success rate of robotic control tasks has been reported to be competitive, showcasing about a 16% improvement over previous best results.
• The main assertion is that Q-SFT integrates prior pre-training knowledge more effectively than analogous methods, which is evidenced by smaller data requirements and higher success rates during initial training phases.

Implications and Speculations

By fusing RL and SFT approaches, Q-SFT offers a pathway to leverage the strengths of both settings while mitigating individual limitations. The ability to retain and capitalize on the pretrained model's knowledge, without necessitating model architecture alterations, bears significant operational benefits. This could facilitate more efficient deployment in real-world applications, particularly those involving complex multi-turn interactions such as dialogues and autonomous navigation.

Future potential researchers may explore expansions of Q-SFT to include vision-language-action (VLA) models, potentially providing even greater benefits in contexts that require an integrated sensory and action framework. Additionally, the translation of Q-SFT to online reinforcement learning environments represents a fertile ground for investigation, potentially broadening its applicability to dynamic scenarios involving real-time decision-making.

In conclusion, Q-SFT addresses critical scalability and efficiency challenges in applying RL to LLMs by recast Q-learning in terms amenable to robust supervised learning paradigms, setting a precedent for future explorations into hybrid RL methodologies.