Research map / Reinforcement learning / RL from feedback

Learning Goal-Conditioned Representations for Language Reward Models

NeurIPS 2024 · Reinforcement learning › RL from feedback · approach: RL for LLM reasoning · 0 citations (OpenAlex)

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.

Trains LM reward models with an added goal-conditioned contrastive loss on hidden representations (Q-function-like), improving math and HH reward AUROC/accuracy and enabling Q-value filtering and guided-decoding steering.

Problem

Reward models for RLHF on language models are trained only with a sequence-level preference loss, so their intermediate representations do not encode how likely a partial generation is to reach a good goal state; it is unclear how representation learning can help RLHF for LMs.

Approach

Jointly train the reward model with a contrastive goal-conditioned loss that pulls together representations of a source state and a future state from the preferred completion while pushing apart a source state and a state from the dispreferred completion, so the cosine similarity approximates a goal-conditioned Q-function usable for filtering and steering.

Key results

• Q-Function reward model improves math reward AUROC by roughly 0.05 (ID) to 0.09 on GSM8k/MATH and ~0.04 average on OOD datasets.
• Q-value filtering on 50 samples raises majority-vote accuracy by 1.2% on GSM8k and 3.9% on MATH, filtering up to 65.6% of generations (55.6% tokens saved on MATH).
• On Helpful-Harmless, the contrastive RM reaches 70.5% accuracy vs 68.2% for the baseline; steering Llama 3 improves helpfulness win rate by 9.6% over SFT guided decoding.

Measured effects reported in the paper

Surprising findings

• Random source/goal state sampling for the contrastive loss degraded OOD performance below the standard preference-ranking baseline.
• DPO performed on par with the contrastive-RM PPO policy on in-distribution math tasks.

Limitations

• Inference-time goal state is a mean of training preferred completions, a mismatch with the training-time goal states.
• Partial Q-values are correlated with partial reward scores, so the signals are entangled and hard to disentangle for tasks like creative writing.
• PPO gains are modest (1.7% weighted average) and the RM was trained on off-policy Mixtral generations.

Experimental setup

Datasets: OpenMathInstruct-1 GSM8k MATH MathQA Helpful-Harmless HelpSteer

Benchmarks: GSM8k MATH algebra222 GSM-Hard Asdiv mawps svamp Helpful-Harmless HelpSteer

Metrics: AUROC accuracy majority-vote accuracy win rate proportion correct tokens saved

Baselines: Codellama 7b Reward Llama 8b Reward DPO PPO with preference-ranking reward model SFT guided decoding

Builds on: Bradley-Terry preference reward modeling proximal policy optimization direct preference optimization OpenMath-CodeLlama

Closest papers by meaning

• GRAM: A Generative Foundation Reward Model for Reward Generalization (ICML 2025 · similarity 0.91)
• RM-Bench: Benchmarking Reward Models of Language Models with Subtlety and Style (ICLR 2025 · similarity 0.91)
• Regularizing Hidden States Enables Learning Generalizable Reward Model for LLMs (NeurIPS 2024 · similarity 0.90)
• Active Reward Modeling: Adaptive Preference Labeling for Large Language Model Alignment (ICML 2025 · similarity 0.90)
• 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)