Reward Model Serving

The reward model scores policy completions — it’s the signal that drives RLHF learning. Its serving infrastructure affects both training speed and alignment quality.

---

1. Types of Reward Signals

Type	How it works	Quality	Speed	Cost
Trained reward model	Separate LLM trained on human preferences	High	Moderate (model inference)	High (needs training data)
Rule-based verifier	Code execution, math checker	Perfect (for verifiable tasks)	Fast	Low
LLM-as-judge	Another LLM scores the completion	Moderate	Slow (full generation)	Moderate
Human evaluation	Actual humans rate completions	Highest	Very slow	Very high

Trained Reward Model

# Reward model: same architecture as policy, but with a scalar head
class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.base = base_model
        self.reward_head = nn.Linear(hidden_size, 1)  # scalar output
    
    def forward(self, input_ids, attention_mask):
        hidden = self.base(input_ids, attention_mask).last_hidden_state
        # Score from the last token's representation
        reward = self.reward_head(hidden[:, -1, :])
        return reward.squeeze(-1)

Rule-Based Verifier (DeepSeek-R1 Approach)

def verify_math(completion, ground_truth):
    """Extract answer from completion, compare to ground truth."""
    extracted = extract_answer(completion)  # regex or structured parsing
    return 1.0 if extracted == ground_truth else 0.0

def verify_code(completion, test_cases):
    """Execute code in sandbox, run test cases."""
    result = sandbox.execute(completion, timeout=10)
    passed = sum(1 for tc in test_cases if result.check(tc))
    return passed / len(test_cases)

Rule-based rewards are preferred when possible — they’re accurate, fast, and don’t need a separate model.

---

2. Serving Architecture

Co-Located (Same GPUs)

GPU 0-3: Policy model (TP=4)
GPU 4-5: Reward model (TP=2, smaller model)
GPU 6-7: Reference model (TP=2)

Pros: No network transfer for scoring. Cons: Fewer GPUs for policy model.

Separate Serving

Policy cluster: GPU 0-7 (generation + training)
Reward API: Separate vLLM instance on GPU 8-9
  → Score completions via HTTP/gRPC

Pros: Independent scaling. Cons: Network latency for scoring.

No Separate Model (GRPO with Verifiers)

Policy cluster: GPU 0-7 (generation + training)
Verifier: CPU process (math checker, code executor)
  → No GPU needed for reward!

This is why DeepSeek-R1’s approach is infrastructure-efficient — rule-based rewards don’t need GPU resources.

---

3. Reward Model Quality

Issue	Symptom	Mitigation
Reward hacking	Model exploits reward model weaknesses	Ensemble multiple reward models
Length bias	Longer responses get higher rewards	Normalize by length
Distribution shift	Reward model degrades as policy improves	Periodically retrain reward model
Overoptimization	Quality improves then degrades past optimum	KL penalty, early stopping

For sandbox management in agent RL, see Agent RL Infrastructure.

---

Self-Check

• Compare trained reward models vs rule-based verifiers
• Implement a reward model with a scalar head
• Design a serving architecture for co-located vs separate reward model
• Explain reward hacking and how to mitigate it

---

References

• Training language models to follow instructions (Ouyang et al., 2022)
• DeepSeek-R1 (2025)