# Reward Model for sunoGPT - Complete Guide

**Comprehensive guide for training, evaluating, and using reward models for music generation quality assessment.**

---

## Table of Contents

1. [Overview](#overview)
2. [Architecture](#architecture)
3. [Training](#training)
4. [Evaluation](#evaluation)
5. [Usage](#usage)
6. [Design Decisions](#design-decisions)
7. [Troubleshooting](#troubleshooting)
8. [Technical Details](#technical-details)
9. [Testing](#testing)
10. [References](#references)

---

## Overview

### What Is This?

A reward model that scores music generation quality, trained on human preference data (chosen vs rejected pairs). Used for:
- **Generation ranking**: Select best among multiple candidates
- **Quality filtering**: Remove low-quality generations
- **RLHF/PPO**: Future reinforcement learning from human feedback
- **Evaluation**: Measure model improvements

### Key Features

✅ **Token-level outputs** - Rewards for each position `(batch, seq_len)`  
✅ **Sequence-level training** - Trained on sequence preferences (A > B)  
✅ **Label smoothing** - Handles ~70% noisy preference data  
✅ **Vectorized masking** - Excludes prompt + padding efficiently  
✅ **Progressive visualization** - See how rewards evolve every ~750 tokens  
✅ **RAD-ready** - Token outputs enable reward-guided sampling

### When to Use

- After pre-training GPT model on music generation
- When you have preference data (human labels or AI feedback)
- For ranking/filtering generations
- As preparation for RLHF

---

## Architecture

### RewardHead Module

Located in `modules/gpt.py`, outputs scalar rewards:

```python
class RewardHead(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.ln = LayerNorm(config.n_embd)
        self.reward_proj = nn.Linear(config.n_embd, 1, bias=False)
    
    def forward(self, x: torch.Tensor):
        # x: (batch, seq_len, n_embd)
        # Output: (batch, seq_len, 1) - one reward per token
        return self.reward_proj(self.ln(x))
```

**Key**: Outputs rewards for **each token position**, not just one scalar per sequence.

### Model Output

When `use_reward_head=True`, the GPT model returns:

```python
reward_logits = model(X, return_logits=True)  # (batch, seq_len)
```

One reward value for each token in the sequence.

### Why Token-Level Outputs?

Even though we train with sequence-level labels, token-level outputs provide:

1. **Flexibility**: Can aggregate rewards however needed (mean, max, last-10, etc.)
2. **Interpretability**: See which parts of sequences get high/low rewards
3. **RAD Inference**: Enable reward-guided sampling (from RAD paper)
4. **Debugging**: Identify where quality drops

---

## Training

### Quick Start

```bash
python train_reward_model.py \
  --data_dir /path/to/dpo/data \
  --out_dir /path/to/checkpoints \
  --preload_checkpoint /path/to/pretrained_gpt.pt \
  --use_reward_head True \
  --learning_rate 1e-6 \
  --label_smoothing 0.1 \
  --batch_size 8
```

Or use a config file with optimized settings for noisy preference data:

```bash
python train_reward_model.py configs/reward_model_test.py
```

### Training Configuration

**Dataset**: Paired preferences (chosen vs rejected sequences)
- Dataset 0: Rejected/negative samples
- Dataset 1: Chosen/positive samples
- Batch format: `[rejected, chosen, rejected, chosen, ...]`

**Loss Function**: Bradley-Terry with label smoothing

```python
loss = -log(sigmoid(reward_chosen - reward_rejected)) * (1 - label_smoothing)
       -log(sigmoid(reward_rejected - reward_chosen)) * label_smoothing
```

**Hyperparameters** (Optimized for ~70% noisy labels):

```python
learning_rate = 1e-6      # Very conservative for stability
weight_decay = 0.1        # Strong regularization
label_smoothing = 0.1     # Handles noisy preferences
warmup_iters = 5_000      # Longer warmup
grad_clip = 0.5           # Strong gradient clipping
batch_size = 8            # Must be even (for pairing)
```

### Training Script

```bash
python train_reward_model.py \
  --data_dir /path/to/dpo/data \
  --out_dir /path/to/checkpoints \
  --preload_checkpoint /path/to/pretrained_gpt.pt \
  --use_reward_head True \
  --learning_rate 5e-6 \
  --label_smoothing 0.1 \
  --max_iters 1000 \
  --batch_size 8
```

### Expected Training Metrics

With ~70% noisy labels:

```
iter 0:   loss 0.69, acc 50%, margin 0.2
iter 100: loss 0.55, acc 65%, margin 0.5
iter 500: loss 0.45, acc 70%, margin 0.8
iter 1000: loss 0.40, acc 72%, margin 1.0
```

**Accuracy**: Caps at ~70-75% (limited by label noise)  
**Margin**: Positive and increasing (chosen > rejected on average)

### Checkpoint Saving

Training saves only:
- `last_ckpt_infer.pt` - Latest checkpoint, ready for inference

Use this for evaluation!

---

## Evaluation

### Quick Start

```bash
cd /home/tony/Work/neon_2/sunoGPT

python scripts/eval_reward_model.py \
  --checkpoint /path/to/last_ckpt_infer.pt \
  --data_dir /path/to/dpo/data
```

### What Gets Generated

Creates timestamped directory with:

```
reward_eval_results/2025-11-02_14-30-45/
├── config.json                # Evaluation config
├── test_cases.txt            # Sample pair results (10 pairs)
├── reward_progression.png    # Progressive rewards every ~1.2s
├── validation_stats.json     # Full validation metrics
├── margin_distribution.png   # Margin histogram
├── summary.txt               # Human-readable summary
└── errors.txt                # Error log (if any)
```

### Progressive Reward Visualization

The key output: `reward_progression.png`

Shows **cumulative average reward** as sequence grows:
- Step 1: avg(rewards[0:750])
- Step 2: avg(rewards[0:1500])
- Step 3: avg(rewards[0:2250])
- etc.

**Two curves per subplot**:
- Green: Chosen sequence progressive reward
- Red: Rejected sequence progressive reward

**What to look for**:
- ✅ Green consistently above red (model works!)
- ✅ Relatively stable curves (consistent assessment)
- ✅ Clear separation (confident discrimination)

### Evaluation Options

```bash
# More test cases
--n_test_cases 20

# More visualization pairs
--n_viz_pairs 10

# Higher temporal resolution (every 0.5s)
--token_interval 12

# Limit validation size
--max_val_samples 100
```

### Interpreting Results

**test_cases.txt**:
```
Pair 0:
  Chosen:   0.8453 - upbeat pop song...
  Rejected: 0.2341 - slow ballad...
  Margin:   +0.6112
  Correct:  ✓

Accuracy: 80.0% (8/10)
```

**Accuracy**:
- >70%: Excellent
- 60-70%: Good (expected with noisy labels)
- <60%: Needs improvement

**Margin**:
- >1.0: Strong, confident
- 0.5-1.0: Good, clear signal
- 0.2-0.5: Moderate
- <0.2: Weak

---

## Usage

### Load Trained Model

```python
from scripts.reward_eval_utils import load_reward_model

model, model_args = load_reward_model(
    "path/to/last_ckpt_infer.pt",
    device="cuda"
)
```

### Get Rewards for Samples

```python
from scripts.reward_eval_utils import extract_scalar_rewards

# Token-level rewards
reward_logits = model(X, return_logits=True)  # (batch, seq_len)

# Scalar reward (averaged over valid tokens)
scalar_reward = extract_scalar_rewards(
    reward_logits, Y, [loss_start_index]
)
```

### Rank Multiple Generations

```python
candidates = [gen1, gen2, gen3, ...]
rewards = []

for candidate_X, candidate_Y, start_idx in candidates:
    output = model(candidate_X, return_logits=True)
    reward_logits = output["reward_logits"]
    end_indices = compute_loss_end_indices(candidate_Y, reward_logits.shape[1])
    scalar = extract_scalar_rewards(reward_logits, [start_idx], end_indices)
    rewards.append(scalar.item())

best_idx = np.argmax(rewards)
best_generation = candidates[best_idx]
```

### Filter Low-Quality Samples

```python
threshold = 0.5

for sample_X, sample_Y, start_idx in dataset:
    output = model(sample_X, return_logits=True)
    reward_logits = output["reward_logits"]
    end_indices = compute_loss_end_indices(sample_Y, reward_logits.shape[1])
    scalar = extract_scalar_rewards(reward_logits, [start_idx], end_indices)
    
    if scalar > threshold:
        keep_sample(sample)
    else:
        discard_sample(sample)
```

---

## Design Decisions

### Sequence-Level Training, Token-Level Outputs

**Why different granularities?**

**Labels are sequence-level**:
- You know: Sequence A > Sequence B
- You don't know: A[token_50] > B[token_50]

**Outputs are token-level**:
- Enables flexible aggregation
- Supports RAD inference
- Provides interpretability

**Training approach**:
```python
# Get token-level outputs
output = model(X, return_logits=True)
reward_logits = output["reward_logits"]  # (batch, seq_len)

# Average to match label granularity
loss_end_list = compute_loss_end_indices(Y, reward_logits.shape[1])
scalar_rewards = extract_scalar_rewards(reward_logits, loss_start_list, loss_end_list)

# Sequence-level Bradley-Terry loss
loss = -log(sigmoid(chosen - rejected))
```

This is correct! Don't force token-level training when you only have sequence-level labels.

### Label Smoothing for Noisy Data

With ~70% label accuracy (noisy preferences):

```python
label_smoothing = 0.1  # Interpolate 10% toward opposite preference
```

Prevents overfitting to mislabeled pairs while still learning the signal.

### Conservative Hyperparameters

```python
learning_rate = 1e-6   # 10x lower than standard
weight_decay = 0.1     # 10x higher for regularization
grad_clip = 0.5        # Stronger clipping
```

Stability > speed when dealing with noisy labels.

### Conditional Last Token Removal

```python
# In modules/gpt.py forward_old():
if not self.config.use_reward_head:
    x = x[:, :-1]  # Only for standard training
```

Reward model keeps full sequence to score all generated tokens.

---

## Troubleshooting

### Issue: TypeError - tuple instead of tensor

**Cause**: Model returning (text_logits, semantic_logits, coarse_logits) instead of rewards

**Fix**: Ensure `use_reward_head=True` in model config
```python
model_args["use_reward_head"] = True
```

### Issue: Shape mismatch (14079 vs 14080)

**Cause**: Conditional last token removal not working

**Fix**: Already handled in `extract_scalar_rewards` - robust to both cases

### Issue: Dtype mismatch (BFloat16 vs Float)

**Cause**: Model parameters not all converted to bfloat16

**Fix**: Use `.to()` with both device and dtype:
```python
model = model.to(device=device, dtype=torch.bfloat16)
```

### Issue: Accuracy stuck at 50%

**Possible causes**:
- Data pairing incorrect (check dataset structure)
- Model not training (check loss decreasing)
- Labels too noisy (check data quality)

**Solutions**:
- Verify `load_dpo_pair=True` in data loading
- Check training logs for loss decrease
- Try longer training

### Issue: Margin near zero

**Possible causes**:
- Model needs more training
- Learning rate too high
- Label noise too high

**Solutions**:
- Train longer (more iterations)
- Lower learning rate (try 5e-7)
- Increase label smoothing (try 0.15)

---

## Technical Details

### Masking Logic

Rewards computed only on **valid** tokens:

1. **After loss_start_index**: Excludes prompt/conditioning
2. **Before padding**: Excludes padding tokens (Y == -1)

```python
def create_valid_mask(Y, loss_start_index_list, seq_len):
    # Position mask
    mask = positions >= start_indices
    
    # Padding mask (Y[j]==-1 means X[j+1] is padding)
    y_valid = (Y[:, 0, :] != -1)
    y_valid_shifted = F.pad(y_valid, (1, 0), value=True)
    
    # Combined
    mask = mask & y_valid_shifted
    return mask
```

**Fully vectorized** - no Python loops!

### Padding Handling

Y is shifted by 1 relative to X:
- `Y[j]` predicts from `X[j]` to `X[j+1]`
- If `Y[j] == -1`, then `X[j+1]` is padding

The code handles this shift:
```python
y_valid_shifted = F.pad(y_valid, (1, 0), value=True)  # Prepend True
```

### Scalar Reward Extraction

```python
# Token-level rewards
reward_logits = model(X)  # (batch, seq_len)

# Create mask for valid tokens
mask = create_valid_mask(Y, loss_start_list, seq_len)

# Average over valid tokens only
masked_rewards = reward_logits * mask
scalar = masked_rewards.sum(dim=1) / mask.sum(dim=1)
```

### Critical Bugs Fixed

**Bug 1: Off-by-One (Last Token Removal)**
- Problem: Last token removed unconditionally
- Fix: Conditional removal (not for reward model)
- Impact: Now includes complete generated sequence

**Bug 2: Padding Contamination**
- Problem: Padding tokens included in average
- Fix: Detect padding via Y, exclude from mask
- Impact: Clean averaging, no noise

**Bug 3: Dtype Mismatch**
- Problem: Mixed float32/bfloat16
- Fix: Convert all parameters to bfloat16
- Impact: Consistent computation

---

## Testing

### Unit Tests (12/12 Passing)

Run with:
```bash
cd /home/tony/Work/neon_2/sunoGPT
PYTHONPATH=. python tests/test_reward_model.py
```

**Tests include**:
1. ✓ Reward head initialization
2. ✓ Reward head forward pass
3. ✓ Bradley-Terry loss
4. ✓ Gradient flow
5. ✓ Freeze base model
6. ✓ Paired preference data
7. ✓ Reward extraction
8. ✓ extract_scalar_rewards correctness (exact: 12.00==12.00)
9. ✓ extract_scalar_rewards edge cases
10. ✓ Label smoothing
11. ✓ Token-level reward loss
12. ✓ Token-level with padding

### Numerical Verification

Tests use position-based rewards to verify exact correctness:
```
Sample 0: mean(positions[5:20]) = 12.00 == 12.00 ✓
Sample 1: mean(positions[8:15]) = 11.00 == 11.00 ✓
Sample 2: mean(positions[10:12]) = 10.50 == 10.50 ✓
```

Masking logic is mathematically verified!

---

## Design Decisions

### Why Sequence-Level Loss (Not Token-Level)?

**Your data**:
- Labels: "Sequence A is better than B" (sequence-level)
- No information about individual token quality

**Token-level training would assume**:
- A[token_0] > B[token_0]
- A[token_1] > B[token_1]
- ... for all tokens

**This is wrong!** Example:
```
Sequence A (chosen):  [good, good, BAD,  good] → overall good
Sequence B (rejected): [ok,   ok,   GOOD, ok]   → overall worse

Sequence-level: A > B ✓ (correct)
Token-level at pos 2: A[2] < B[2] ✗ (false assumption!)
```

**Solution**: Train with sequence-level loss, keep token-level outputs.

### Token-Level Infrastructure (Implemented but Disabled)

The code includes token-level training functions:
- `token_level_reward_loss()` - Per-token Bradley-Terry
- Configuration flags: `use_token_level_loss`, `lambda_token`
- Full implementation and tests

**Why implemented?**
- Future-proofing: Ready if you get token-level labels
- RAD research: Can experiment with token-level objectives
- Flexibility: Can enable with one flag

**Why disabled?**
```python
use_token_level_loss = False  # Labels are sequence-level
```

Matches your data granularity.

### RAD (Reward-Augmented Decoding)

Inspired by [arXiv:2310.09520](https://arxiv.org/abs/2310.09520), token-level outputs enable:

```python
# During generation:
reward_at_position = reward_model(context)[0, -1]
p_new ∝ p_base * exp(alpha * reward_at_position)
next_token = sample(p_new)
```

**Future work**: Implement RAD sampling for controlled generation.

---

## Usage Examples

### Example 1: Evaluate Single Pair

```python
import torch
from scripts.reward_eval_utils import load_reward_model, extract_scalar_rewards

# Load model
model, _ = load_reward_model("path/to/last_ckpt_infer.pt")

# Load samples (use your data loading)
X_chosen, Y_chosen, start_idx = load_chosen_sample()
X_rejected, Y_rejected, _ = load_rejected_sample()

# Get rewards
with torch.no_grad():
    with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
        output_c = model(X_chosen, return_logits=True)
        output_r = model(X_rejected, return_logits=True)

# Extract reward logits from dict
rewards_c = output_c["reward_logits"]
rewards_r = output_r["reward_logits"]

# Compute end indices and extract scalars
end_c = compute_loss_end_indices(Y_chosen, rewards_c.shape[1])
end_r = compute_loss_end_indices(Y_rejected, rewards_r.shape[1])
scalar_c = extract_scalar_rewards(rewards_c, [start_idx], end_c)
scalar_r = extract_scalar_rewards(rewards_r, [start_idx], end_r)

print(f"Chosen: {scalar_c.item():.3f}")
print(f"Rejected: {scalar_r.item():.3f}")
print(f"Margin: {(scalar_c - scalar_r).item():.3f}")
print(f"Correct: {scalar_c > scalar_r}")
```

### Example 2: Progressive Rewards

```python
# Get token-level rewards
reward_logits = model(X, return_logits=True)[0]  # (seq_len,)

# Get valid positions
mask = create_valid_mask(Y, [start_idx], len(reward_logits))[0]
valid_pos = torch.where(mask)[0]

# Compute progressive rewards (every 750 tokens)
for end_offset in range(750, len(valid_pos), 750):
    tokens = valid_pos[:end_offset]
    progressive_reward = reward_logits[tokens].mean()
    print(f"At {end_offset/25:.1f}s: {progressive_reward:.3f}")
```

### Example 3: Best-of-N Sampling

```python
# Generate N candidates
candidates = [generate() for _ in range(N)]

# Score each
rewards = []
for X, Y, start_idx in candidates:
    output = model(X, return_logits=True)
    reward_logits = output["reward_logits"]
    end_indices = compute_loss_end_indices(Y, reward_logits.shape[1])
    scalar = extract_scalar_rewards(reward_logits, [start_idx], end_indices)
    rewards.append(scalar.item())

# Select best
best = candidates[np.argmax(rewards)]
```

---

## File Reference

### Core Implementation

| File | Purpose |
|------|---------|
| `modules/gpt.py` | RewardHead class, GPT with reward mode |
| `train_reward_model.py` | Training script with Bradley-Terry loss |
| `utils/dpo_data_utils.py` | DPO paired data loading |
| `tests/test_reward_model.py` | 12 comprehensive unit tests |

### Evaluation

| File | Purpose |
|------|---------|
| `scripts/eval_reward_model.py` | Main evaluation script |
| `scripts/reward_eval_utils.py` | Helper functions for evaluation |

### Configuration

| File | Purpose |
|------|---------|
| `configs/reward_model_test.py` | Example configuration |

---

## References

### Papers

- **Reward-Augmented Decoding** ([arXiv:2310.09520](https://arxiv.org/abs/2310.09520))
  - Deng & Raffel, 2023
  - Token-level rewards for controlled generation
  - Inspiration for architecture design

- **Bradley-Terry Model**
  - Classical preference modeling
  - P(A > B) = sigmoid(reward_A - reward_B)

### Related Work

- **DPO** (Direct Preference Optimization) - Used for comparison
- **InstructGPT** - RLHF with reward models
- **RLHF** - Reinforcement Learning from Human Feedback

---

## Quick Command Reference

### Training
```bash
python train_reward_model.py \
  --data_dir /path/to/dpo/data \
  --out_dir /path/to/checkpoints \
  --preload_checkpoint /path/to/pretrained_gpt.pt \
  --use_reward_head True \
  --batch_size 8
```

### Evaluation
```bash
cd /home/tony/Work/neon_2/sunoGPT
python scripts/eval_reward_model.py \
  --checkpoint /path/to/last_ckpt_infer.pt \
  --data_dir /path/to/dpo/data
```

### Testing
```bash
PYTHONPATH=/home/tony/Work/neon_2/sunoGPT python tests/test_reward_model.py
```

---

## Performance Expectations

### Training
- **Time**: ~1-2 hours for 286 iterations on 16 GPUs
- **Accuracy**: 65-72% on validation (with ~70% noisy labels)
- **Margin**: 0.5-1.2 (positive, increasing over training)

### Evaluation
- **Time**: ~30-60 seconds for test cases + visualization
- **Output**: Timestamped directory with plots and stats
- **Files**: 4-6 files (test cases, plots, summary, optional errors)

### Inference
- **Speed**: ~1-2ms per sample on GPU
- **Memory**: Depends on model size (~4GB for 1B params)

---

## Advanced Topics

### Token-Level Training (Optional)

If you ever get token-level labels, enable with:
```python
use_token_level_loss = True
lambda_token = 1.0
```

This provides 10-100x more training signal but requires token-level preference data.

### Freezing Base Model

To train only the reward head (faster):
```python
freeze_base_model = True
learning_rate = 1e-4  # Can use higher LR
```

Trains only ~1M parameters instead of full model.

### Multi-GPU Training

Supports DDP and FSDP:
```python
fsdp = True
sharding_strategy = "full_shard"
grad_checkpointing = True
```

### Custom Reward Aggregation

Instead of mean, try:
```python
# Last 10 tokens only
scalar = reward_logits[:, -10:].mean()

# Max reward
scalar = reward_logits.max()

# Weighted by position
weights = torch.linspace(0.5, 1.0, len(reward_logits))
scalar = (reward_logits * weights).sum() / weights.sum()
```

---

## FAQ

**Q: Why not use DPO instead of reward model?**  
A: Reward model enables ranking, filtering, and future RLHF. DPO is for direct policy optimization.

**Q: Can I use this for text generation?**  
A: Architecture is general - designed for music but adaptable to any sequential generation.

**Q: How much data do I need?**  
A: At least 1000 preference pairs. More is better. Quality > quantity.

**Q: What if my labels are >90% accurate?**  
A: Reduce label_smoothing to 0.05 or 0.0, can use higher learning rate.

**Q: Can I fine-tune on new data?**  
A: Yes! Set `preload_checkpoint` to existing reward model, train on new pairs.

**Q: Does this work for real-time generation?**  
A: Yes for ranking. For RAD inference (reward-guided sampling), need to implement sampling logic.

---

## Acknowledgments

Implementation based on:
- DPO training infrastructure (train_dpo.py)
- RAD paper insights (token-level outputs)
- Bradley-Terry preference modeling
- Extensive debugging and iteration

---

**Complete reward model system for sunoGPT - train, evaluate, and deploy!** 🎉