Research map / Reinforcement learning / RL from feedback
Q-Probe: A Lightweight Approach to Reward Maximization for Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Q-probing, which trains a small linear probe on frozen LLM embeddings to reweight k sampled completions via softmax, optimizing reward with reward-model or importance-weighted policy-gradient losses.
Problem
Adapting pre-trained language models to maximize task-specific rewards without the cost of finetuning, and without white-box access (e.g., API models).
Approach
Keep the LM frozen, learn a linear value probe on prompt-completion embeddings, and use it to softmax-reweight k sampled completions; train the probe with an importance-weighted policy gradient loss that contrasts samples from the same prompt, matching the inference procedure.
Key results
- On Code-LLaMA-7B, Q-probe with L_PG reaches 0.46 MBPP-Test vs 0.38 greedy baseline; 5-shot plus Q-probe reaches 0.52.
- On preference data, Q-probe (L_QP) reaches 50.10% win rate vs DPO 44.97% and offline PPO 44.07%; KTO+Q-probe reaches 55.01% vs KTO 51.46%.
- On GSM8K with Code-LLaMA-7B, Q-probe L_PG reaches 0.45 vs 0.29 greedy baseline.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| train the Q-probe with the direct policy-gradient loss L_PG instead of reward regression L_Q (ablation) | Q-probe on Code-LLaMA-7B · MBPP-Test | expected return: 0.38 -> 0.46 | Helped |
| train the Q-probe with the direct policy-gradient loss L_PG instead of classification loss L_CE (ablation) | Q-probe on Code-LLaMA-7B · MBPP-Test | expected return: 0.40 -> 0.46 | Helped |
| combine 5-shot prompting with Q-probe L_PG | Code-LLaMA-7B · MBPP-Test | expected return: 0.46 -> 0.52 | Helped |
| combine 5-shot prompting with Q-probe L_PG | Code-LLaMA-7B · HumanEval | expected return: 0.34 -> 0.39 | Helped |
| use Q-probe L_CE instead of L_Q (ablation) | Q-probe on Code-LLaMA-7B · GSM-8K | expected return: 0.36 -> 0.43 | Helped |
| use Q-probe L_PG instead of L_CE (ablation) | Q-probe on Code-LLaMA-7B · GSM-8K | expected return: 0.43 -> 0.45 | Helped |
| add Q-probe L_QP to the SFT model | LLaMA-7B after SFT · HH + OpenAssistant + SHP | win rate (%): 37.86 -> 50.10 | Helped |
| add Q-probe L_QP to the KTO-finetuned model | KTO-finetuned LLaMA-7B · HH + OpenAssistant + SHP | win rate (%): 51.46 -> 55.01 | Helped |
| use 2- or 3-layer MLP probes instead of the 1-layer linear probe (ablation) | Q-probe on preference data · human preference learning benchmarks | win rate | Hurt |
| replace Code-LLaMA embeddings with text-embedding-3-small embeddings (ablation) | Q-probe on gpt-3.5-turbo-1106 · MBPP-Test | expected return: 0.68 -> 0.66 | Mixed |
| replace Code-LLaMA embeddings with text-embedding-3-small embeddings (ablation) | Q-probe on gpt-3.5-turbo-1106 · HumanEval | expected return: 0.64 -> 0.54 | Hurt |
Surprising findings
- Q-probe on gpt-3.5-turbo embeddings gave only modest gains (about 3% over base), and text-embedding-3-small embeddings yielded no gains over baseline.
- Larger 2- or 3-layer MLP probes overfit and harmed preference-learning performance compared with a linear probe.
Limitations
- Inference cost scales with k samples drawn from the base model.
- Gains are smaller on stronger API base models (about 3%) and depend on embedding quality.
Experimental setup
Datasets: MBPP HumanEval GSM8K Anthropic HH OpenAssistant Stanford Human Preferences (SHP)
Benchmarks: MBPP HumanEval GSM8K Anthropic HH + OpenAssistant + SHP win rate
Metrics: expected return (pass rate) win rate (%)
Baselines: Baseline (Pass@1) Greedy decoding 5-shot on successes SFT on successes Prompt RM Finetune RM PPO (offline) DPO KTO
Builds on: DPO KTO Bradley-Terry reward modeling LoRA
Closest papers by meaning
- Discriminative Policy Optimization for Token-Level Reward Models (ICML 2025 · similarity 0.90)
- Vanishing Gradients in Reinforcement Finetuning of Language Models (ICLR 2024 · similarity 0.89)
- ToolRL: Reward is All Tool Learning Needs (NeurIPS 2025 · similarity 0.89)
- Implicit Reward as the Bridge: A Unified View of SFT and DPO Connections (NeurIPS 2025 · similarity 0.88)
- Q-SFT: Q-Learning for Language Models via Supervised Fine-Tuning (ICLR 2025 · similarity 0.88)
- Quantile Reward Policy Optimization: Alignment with Pointwise Regression and Exact Partition Functions (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.