Train LLMs to simultaneously act as reasoning agents and reward models through recycled on-policy rollouts, eliminating separate reward infrastructure while achieving 9.7% gains on reasoning tasks.
Train LLMs to simultaneously act as reasoning agents and reward models through recycled on-policy rollouts, eliminating separate reward infrastructure while achieving 9.7% gains on reasoning tasks.
Train Language Models with Co-Evolving Policy and Reward
Outcome
Build a unified LLM training system that eliminates separate reward models by recycling generated rollouts to train policy and reward judgment within a single model, reducing computational overhead while improving reasoning and reward accuracy.
Problem Context
Standard reinforcement learning for language models relies on two separate systems: a policy that generates responses and a reward model that evaluates them. This separation creates inefficiencies:
Expensive human preference data collection for reward modeling
Computational redundancy during training and inference (two models instead of one)
Reward model quality often lags behind policy improvements, creating optimization misalignment
Training data from policy improvements gets discarded rather than recycled
SPARK addresses these by creating a single model that simultaneously learns to reason better and judge answers more accurately, using the same rollouts for both objectives.
Core Concept
The central insight is co-evolution through data recycling: instead of separate systems, one model learns multiple complementary objectives from the same generated rollouts. Three interrelated training tasks create a feedback loop:
Pointwise scoring teaches the model to recognize when individual responses are correct
Pairwise comparison develops preference discrimination between response qualities
Reflection training enables the model to fix its own mistakes through self-correction
When the reward component improves, it produces better policy gradients. Better policies generate higher-quality rollouts. Better rollouts train more accurate reward judgment. This creates compounding gains without external infrastructure.
Architecture Overview
The training pipeline consists of four sequential phases per iteration:
Rollout generation: Sample multiple candidate responses from the policy for each input
On-policy evaluation: Assign advantage scores using verifiable ground truth (exact match, execution success, or reference comparison)
Data construction: Transform rollouts into three complementary datasets for pointwise, pairwise, and reflection objectives
Unified optimization: Backpropagate a combined loss that balances policy improvement with reward accuracy and drift regularization
The model architecture remains standard (transformer-based LLM or VLM). Training infrastructure differs only in objective composition, not in fundamental components.
Implementation
Phase 1: Rollout Generation and Evaluation
Generate multiple candidate responses per input and compute standardized advantage scores that will inform all downstream training objectives.
import torch
from torch.utils.data import DataLoader
defgenerate_rollouts(model, inputs, num_candidates=4, temperature=0.7):
"""Generate multiple candidate responses per input."""
rollouts = []
for input_text in inputs:
candidates = []
for _ inrange(num_candidates):
# Generate with sampling
output = model.generate(
input_text,
temperature=temperature,
max_length=256,
do_sample=True
)
candidates.append(output)
rollouts.append({
'input': input_text,
'candidates': candidates,
'generated_at': 'current_policy'
})
return rollouts
defcompute_advantages(rollouts, ground_truth_labels, baseline='mean'):
"""Assign advantage scores to candidates using verifiable rewards."""
processed = []
for rollout in rollouts:
input_text = rollout['input']
candidates = rollout['candidates']
labels = ground_truth_labels[input_text]
# Compute individual correctness scores
scores = []
for candidate in candidates:
# Use exact match, execution, or reference-based comparison
score = evaluate_correctness(candidate, labels)
scores.append(score)
# Standardize advantagesif baseline == 'mean':
baseline_val = sum(scores) / len(scores)
elif baseline == 'min':
baseline_val = min(scores)
else:
baseline_val = 0.0
advantages = [s - baseline_val for s in scores]
processed.append({
'input': input_text,
'candidates': candidates,
'scores': scores,
'advantages': advantages,
'normalized_scores': [(s - min(scores)) / (max(scores) - min(scores) + 1e-8)
for s in scores]
})
return processed
defevaluate_correctness(candidate, ground_truth):
"""Verify answer correctness using exact match or semantic similarity."""ifisinstance(ground_truth, list):
# Multiple acceptable answersreturnfloat(candidate.strip() in [g.strip() for g in ground_truth])
else:
# Single referencereturnfloat(candidate.strip() == ground_truth.strip())
Phase 2: Multi-Objective Data Construction
Transform single rollouts into three complementary datasets. Each serves a different learning objective, enabling the model to develop specialized subcomponents.
defconstruct_training_data(processed_rollouts):
"""Create pointwise, pairwise, and reflection datasets from rollouts."""
pointwise_data = []
pairwise_data = []
reflection_data = []
for rollout in processed_rollouts:
input_text = rollout['input']
candidates = rollout['candidates']
scores = rollout['normalized_scores']
advantages = rollout['advantages']
# POINTWISE: (input, response) -> correctness scorefor candidate, score inzip(candidates, scores):
pointwise_data.append({
'input': input_text,
'response': candidate,
'target_score': score,
'task_type': 'pointwise'
})
# PAIRWISE: Compare two responses and predict preferencefor i inrange(len(candidates)):
for j inrange(i + 1, len(candidates)):
better_idx = i if scores[i] > scores[j] else j
worse_idx = j if better_idx == i else i
pairwise_data.append({
'input': input_text,
'response_a': candidates[better_idx],
'response_b': candidates[worse_idx],
'label': 0, # response_a is better'task_type': 'pairwise'
})
# REFLECTION: Train model to self-correct wrong answersfor candidate, score inzip(candidates, scores):
if score < 0.5: # Incorrect response
reflection_data.append({
'input': input_text,
'incorrect_response': candidate,
'task': 'Generate corrected response',
'task_type': 'reflection'
})
return {
'pointwise': pointwise_data,
'pairwise': pairwise_data,
'reflection': reflection_data
}
defprepare_batch(examples, tokenizer, max_length=512):
"""Tokenize and format batch for unified training."""
batch = {'input_ids': [], 'attention_mask': [], 'labels': []}
for example in examples:
task_type = example.get('task_type', 'pointwise')
if task_type == 'pointwise':
# Format: "Question: {input}\nAnswer: {response}\nScore: {score}"
text = f"Question: {example['input']}\nAnswer: {example['response']}\nScore:"
target = str(int(example['target_score'] * 100))
elif task_type == 'pairwise':
# Format for preference prediction
text = f"Compare:\nA: {example['response_a']}\nB: {example['response_b']}\nBetter:"
target = "A"else: # reflection
text = f"Fix: {example['incorrect_response']}\nCorrected:"
target = ""# Will be generated
encoded = tokenizer(
text,
truncation=True,
max_length=max_length,
padding='max_length',
return_tensors='pt'
)
batch['input_ids'].append(encoded['input_ids'].squeeze())
batch['attention_mask'].append(encoded['attention_mask'].squeeze())
return {k: torch.stack(v) for k, v in batch.items()}
Phase 3: Unified Loss Computation
Combine multiple objectives into a single loss function that balances policy improvement, reward accuracy, and KL divergence regularization.
Self-correction capability; scale with error rate in rollouts
Beta KL (divergence)
0.05–0.2
Higher prevents policy drift; 0.1 standard baseline
Num candidates per input
2–8
More candidates = better advantage estimates; diminishing returns at 4+
Num rollout steps per epoch
2–8
Balance between on-policy freshness and efficiency
Learning rate
1e-5 to 5e-5
Standard LM fine-tuning range
Baseline method
'mean' or 'min'
'mean' reduces variance; 'min' emphasizes best rollouts
When to Use SPARK
Objective tasks with verifiable rewards: Math, coding, information retrieval where ground truth is checkable
Scaling reasoning without external reward models: When human preference annotation is bottlenecked
Reducing inference cost: Single unified model vs. separate policy + reward infrastructure
On-policy training preferred: When distribution drift is a concern and fresh rollouts are feasible
Limited preference data: When collecting human feedback is expensive or unavailable
Multi-scale reasoning: Combining pointwise judgment with pairwise preference and self-reflection
When NOT to Use SPARK
Subjective tasks without ground truth: Content that lacks verifiable correctness signals (creative writing, aesthetic judgments)
Preference-based RLHF only: If your primary goal is aligning with human subjective preferences and you have preference data, DPO or IPO may be more direct
Fully offline data: SPARK requires on-policy rollout generation; fully offline scenarios suit standard SFT or offline RL
Extreme efficiency constraints: Rollout generation per epoch adds computational cost compared to single forward pass SFT
Real-time deployment requirements: Multi-candidate generation and reflection add latency at inference
Noisy reward signals: Pointwise and pairwise objectives suffer if ground truth evaluation is unreliable or partially labeled
Low-resource settings: KL regularization with reference models requires memory for two model copies
Common Pitfalls
1. Unbalanced loss weights across objectives: If pointwise weight dominates, the model optimizes for score prediction at the expense of generation quality. Start with equal weights and adjust after 1–2 epochs of training.
2. Insufficient rollout diversity: Using low temperature (< 0.5) during generation reduces candidate diversity. Meaningless comparisons occur when all candidates are similar. Maintain temperature >= 0.7.
3. Ignoring baseline selection: Using a fixed baseline (e.g., 0.0) instead of per-batch standardization inflates advantages early in training. Always subtract batch mean or min from scores.
4. Reference model staleness: If the reference model for KL divergence isn't updated periodically, policy drift regularization becomes ineffective. Refresh every 2–3 epochs or when validation accuracy plateaus.
5. Over-weighting reflection loss on clean data: If your evaluation function is perfect (100% of rollouts labeled correctly), reflection training wastes capacity. Scale reflection weight down when error rate < 20%.
6. Batch size inconsistency across objectives: Combining pointwise (which scales to 4N candidates) with pairwise (scales to 6N comparisons) in fixed batch_size can cause memory spikes. Use dynamic batching or separate sub-batches.
7. Mixing on-policy and off-policy data: Rollouts generated from old policies have different advantage scales. Always regenerate rollouts before each training epoch rather than recycling old data.
Reference
Synergistic Policy And Reward Co-Evolving Framework for Large Language Models
