Skip to content

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

🕒 Published (v1): 2025-01-01 · Source: ICLR · Venue: ICLR 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Q-SFT casts offline Q-learning as weighted supervised fine-tuning by encoding Q-values directly into token probabilities, eliminating the need to add or reinitialize value heads on pretrained LLMs/VLMs. The weighted cross-entropy loss uses Bellman backup targets as per-token weights, enabling stable value learning that smoothly inherits pretrained knowledge. The method is evaluated across language games, web tasks, VLM navigation (ALFWorld with LLaVA-1.6), and robotic manipulation.

Problem

Standard offline Q-learning methods require regressing value functions via temporal-difference (TD) error, which is unstable at scale and necessitates separately initialized value heads—discarding the likelihood-based parameterization of pretrained LLMs/VLMs. This mismatch between maximum-likelihood pretraining and Q-value regression causes existing value-based RL methods to underperform SFT on multi-turn language tasks, while SFT itself cannot stitch suboptimal trajectories into a better policy.

Method

Q-SFT reframes Q-learning as a weighted cross-entropy (WCE) loss over the model's existing token distribution. The key insight: if token probabilities are trained with weights equal to Bellman backup values \(\hat{B}^* p_{\bar\theta}(a|s) = r + \gamma \max_{a'} \frac{p_{\bar\theta}(a'|s')}{\pi_\beta(a'|s')}\), the resulting likelihoods \(\hat{p}_\theta(a|s)\) provably satisfy \(Q^*(s,a)\pi_\beta(a|s) \leq \hat{p}_\theta(a|s) \leq Q^*(s,a)\)—i.e., they are conservative Q-value estimates. No new heads or weight resets are needed; the value model is initialized directly from pretrained logits. A separately trained behavior clone \(\pi_\phi\) (standard SFT) is used for policy extraction at inference time: \(\hat\pi(a|s) \propto \pi_\phi(a|s)\exp(\beta \hat{p}_\theta(a|s))\). The label-smoothing structure of the WCE loss (residual weight distributed equally over non-selected actions) provides implicit regularization against out-of-distribution actions.

Key Contributions

  • First value-based offline RL algorithm that fine-tunes LLMs/VLMs without any architectural modification (no value heads, no weight reinitialization).
  • Theoretical proof that the WCE objective produces conservative Q-value estimates bounded by \([\pi_\beta Q^*, Q^*]\), with performance bounds competitive with state-of-the-art offline RL.
  • Empirical demonstration across six task types: language games (Chess, Wordle, 20 Questions), WebShop, VLM navigation (ALFWorld with LLaVA-1.6), and robotic manipulation.
  • Scaling analysis showing Q-SFT retains pretrained knowledge better than ILQL as model size grows (GPT2-medium → large → xl).

Results

  • ALFWorld (VLM, LLaVA-1.6): Q-SFT achieves 39/21/19/14/18/21% success across 6 task types vs. ILQL 28/7/0/5/2/15% and SFT 38/15/0/11/0/18%; competitive with GPT4-V ReAct (45/19/17/7/12/24%) despite ~30× fewer parameters.
  • WebShop: Q-SFT scores 0.63 vs. SFT 0.55, offline ArCHer 0.57, ReAct (GPT3.5) 0.60.
  • 20 Questions: Q-SFT returns −13.1 vs. ILQL −14.2, SFT −17.3 (dataset average), ReAct −13.2.
  • Wordle: Q-SFT −2.11 vs. ILQL −2.08 (near-tied), SFT −3.81.
  • Robotic manipulation (Pick Object): Q-SFT 94% vs. Q-Transformer 92%, CQL 78%, BC 44%.
  • Scaling (20Q, 10% data): Q-SFT increasingly outperforms ILQL as model size grows from 345M → 774M → 1558M parameters.
  • Sample efficiency (robotics): Q-SFT reaches Q-Transformer final performance faster in the low-data regime.

Limitations

  • Inference-time policy extraction requires both \(\pi_\phi\) and \(\hat{p}_\theta\) to be available simultaneously, doubling memory/compute relative to standard SFT.
  • Bellman backup weights require access to \(\pi_\beta(a'|s')\) estimates, which must themselves be approximated; errors in the behavior policy clone propagate into value estimates.
  • Evaluated on GPT2-medium (345M) and LLaVA-1.6; scaling to frontier-scale models (billions+) is not empirically validated.
  • Online RL extension is identified as future work; the method is offline-only.
  • Chess and Wordle experiments use randomly initialized weights, removing the pretraining advantage the method theoretically provides.

Relevance to Vision-Language Models

Q-SFT is directly applicable to VLMs: the paper validates it on LLaVA-1.6 for embodied navigation (ALFWorld), showing that multi-turn RL fine-tuning of a VLM via weighted SFT outperforms both ILQL and standard SFT without any architectural change. This is significant for VLM practitioners because standard RLHF/DPO methods target single-turn responses, while Q-SFT enables offline multi-turn optimization (sequential tool use, navigation, robotic control) while fully leveraging visual pretraining. The architecture-agnostic nature means it can be dropped into any autoregressive VLM, and the authors explicitly cite VLA models (e.g., RT-2) as the next target—making this a stepping stone toward RL-trained vision-language-action systems.