Training Reward Models

import torch
import torch.nn.functional as F

def compute_ranking_loss(
    chosen_rewards: torch.Tensor,
    rejected_rewards: torch.Tensor,
    margin: float = 0.0
) -> torch.Tensor:
    """
    Compute ranking loss for reward model training.
    
    This is the Bradley-Terry ranking loss: -log(sigmoid(r_chosen - r_rejected))
    
    The goal: make chosen_rewards > rejected_rewards by a comfortable margin.
    
    Args:
        chosen_rewards: Rewards for chosen responses, shape (batch_size,)
        rejected_rewards: Rewards for rejected responses, shape (batch_size,)
        margin: Optional margin to enforce minimum difference (usually 0)
    
    Returns:
        Ranking loss (scalar)
    """
    # Compute the difference: chosen should be higher than rejected
    # If margin > 0, we require chosen to be higher by at least margin
    logits = chosen_rewards - rejected_rewards - margin
    
    # Apply log-sigmoid for numerical stability
    # -log(sigmoid(x)) = log(1 + exp(-x)) = softplus(-x)
    # (PyTorch's softplus is more numerically stable than manually computing log(sigmoid))
    loss = F.softplus(-logits)
    
    return loss.mean()

# Let's see it in action!
print("Example: Computing ranking loss")
print("=" * 50)

# Create a batch of 4 examples
batch_size = 4
chosen_rewards = torch.tensor([2.0, 1.5, 3.0, 0.5])
rejected_rewards = torch.tensor([1.0, 1.0, 2.0, 1.0])

print("\nChosen rewards:  ", chosen_rewards.tolist())
print("Rejected rewards:", rejected_rewards.tolist())
print("\nDifferences:     ", (chosen_rewards - rejected_rewards).tolist())

loss = compute_ranking_loss(chosen_rewards, rejected_rewards)
print(f"\nRanking loss: {loss.item():.4f}")

# Accuracy: how often is chosen > rejected?
accuracy = (chosen_rewards > rejected_rewards).float().mean()
print(f"Accuracy: {accuracy.item():.2%}")

print("\nNotice:")
print("  - Examples 1-3: chosen > rejected → contributes low loss")
print("  - Example 4: chosen < rejected → contributes high loss (model wrong!)")
print("  - Overall accuracy is 75% (3 out of 4 correct)")

# Let's look at the individual losses
individual_losses = F.softplus(-(chosen_rewards - rejected_rewards))
print("\nIndividual losses per example:")
for i, (loss_val, diff) in enumerate(zip(individual_losses, chosen_rewards - rejected_rewards)):
    print(f"  Example {i+1}: diff={diff:+.1f} → loss={loss_val:.4f}")

def compute_ranking_loss_with_metrics(
    chosen_rewards: torch.Tensor,
    rejected_rewards: torch.Tensor,
    margin: float = 0.0
) -> dict:
    """
    Compute ranking loss AND all the useful metrics you want to track.
    
    This is what you'd call during training to get a full picture of how
    your model is performing.
    """
    loss = compute_ranking_loss(chosen_rewards, rejected_rewards, margin)
    
    # Accuracy: how often does model rank chosen higher?
    accuracy = (chosen_rewards > rejected_rewards).float().mean()
    
    # Mean rewards (useful for debugging)
    mean_chosen = chosen_rewards.mean()
    mean_rejected = rejected_rewards.mean()
    
    # Mean margin: how much higher is chosen vs rejected on average?
    mean_margin = (chosen_rewards - rejected_rewards).mean()
    
    return {
        "loss": loss,
        "accuracy": accuracy,
        "mean_chosen_reward": mean_chosen,
        "mean_rejected_reward": mean_rejected,
        "mean_margin": mean_margin,
    }

# Let's use the same example from before
metrics = compute_ranking_loss_with_metrics(chosen_rewards, rejected_rewards)

print("Training metrics for our example batch:")
print("=" * 50)
for k, v in metrics.items():
    if k == "loss":
        print(f"  {k}: {v.item():.4f} (lower is better)")
    elif k == "accuracy":
        print(f"  {k}: {v.item():.2%} (higher is better)")
    elif k == "mean_margin":
        print(f"  {k}: {v.item():.4f} (positive and growing is good)")
    else:
        print(f"  {k}: {v.item():.4f}")

print("\nInterpretation:")
print("  - Accuracy of 75% is decent (3 out of 4 correct)")
print("  - Mean margin of +0.5 means chosen responses score 0.5 higher on average")
print("  - During training, you'd watch these metrics improve over time")

import torch.nn as nn
from torch.utils.data import DataLoader
from transformers import AutoModel, AutoTokenizer, get_linear_schedule_with_warmup
from tqdm.auto import tqdm

class RewardModel(nn.Module):
    """
    A reward model for predicting human preferences.
    
    Architecture is simple:
    - Take a base language model (GPT-2, Llama, whatever)
    - Add a "value head" on top that projects to a single scalar
    - That scalar is the reward
    
    The base model processes the text and extracts meaning.
    The value head says "based on this meaning, how good is this response?"
    """
    
    def __init__(self, base_model, hidden_size):
        super().__init__()
        self.base_model = base_model
        
        # The value head: dropout for regularization, then linear layer to scalar
        self.value_head = nn.Sequential(
            nn.Dropout(0.1),  # 10% dropout to prevent overfitting
            nn.Linear(hidden_size, 1)  # hidden_size → 1 scalar reward
        )
    
    def get_rewards(self, input_ids, attention_mask):
        """
        Compute reward scores for input sequences.
        
        Args:
            input_ids: Token IDs, shape (batch_size, seq_len)
            attention_mask: Attention mask, shape (batch_size, seq_len)
        
        Returns:
            Reward scores, shape (batch_size,)
        """
        # Run the base model to get hidden states
        outputs = self.base_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True
        )
        hidden_states = outputs.last_hidden_state  # (batch_size, seq_len, hidden_size)
        
        # Get the LAST non-padding token's hidden state for each sequence
        # (This is where the model has seen the entire response)
        seq_lengths = attention_mask.sum(dim=1) - 1  # -1 because of 0-indexing
        batch_size = hidden_states.shape[0]
        last_hidden = hidden_states[
            torch.arange(batch_size, device=hidden_states.device),
            seq_lengths.long()
        ]
        
        # Project to scalar and squeeze out the last dimension
        return self.value_head(last_hidden).squeeze(-1)

# Let's create a reward model and see it in action
print("Creating a reward model from GPT-2...")
print("=" * 50)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load base model and tokenizer
base_model = AutoModel.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token  # GPT-2 doesn't have a pad token by default

# Create reward model
reward_model = RewardModel(base_model, hidden_size=768)  # GPT-2 has 768 hidden dims
reward_model.to(device)

print(f"\nReward model created!")
print(f"  Device: {device}")
print(f"  Base model: GPT-2 (124M parameters)")
print(f"  Hidden size: 768")
print(f"  Value head: 768 → 1 (just 769 parameters!)")

# Test it with a forward pass
test_texts = [
    "This is a helpful and informative response that answers the question clearly.",
    "I don't know, just Google it yourself."
]

print(f"\nTesting with example responses...")
inputs = tokenizer(test_texts, return_tensors="pt", padding=True, truncation=True)
inputs = {k: v.to(device) for k, v in inputs.items()}

with torch.no_grad():
    rewards = reward_model.get_rewards(inputs['input_ids'], inputs['attention_mask'])

print(f"\nTest rewards (before training, so basically random):")
for text, reward in zip(test_texts, rewards):
    print(f"  [{reward.item():+.4f}] \"{text[:60]}...\"")
    
print("\nNote: Rewards are random right now because the model is untrained.")
print("After training, the first response should get a higher reward!")

def train_reward_model(model, train_loader, eval_loader, config, device):
    """
    Complete reward model training loop.
    
    For each (prompt, chosen, rejected) triple in the dataset:
    - Compute rewards for both chosen and rejected
    - Calculate ranking loss: we want chosen > rejected
    - Backprop and update weights
    
    Args:
        model: RewardModel instance
        train_loader: DataLoader for training data
        eval_loader: DataLoader for evaluation (can be None)
        config: Training configuration dict
        device: torch device
    
    Returns:
        Trained model
    """
    
    # AdamW optimizer with weight decay (helps prevent overfitting)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config['learning_rate'],
        weight_decay=0.01
    )
    
    # Learning rate scheduler: linear warmup then decay
    total_steps = len(train_loader) * config['num_epochs']
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=config['warmup_steps'],
        num_training_steps=total_steps
    )
    
    model.train()
    best_eval_accuracy = 0.0
    
    for epoch in range(config['num_epochs']):
        epoch_metrics = {'loss': 0, 'accuracy': 0}
        
        progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}")
        
        for batch in progress_bar:
            # Move batch to device
            batch = {k: v.to(device) for k, v in batch.items()}
            
            # Forward pass for chosen responses
            chosen_rewards = model.get_rewards(
                batch['chosen_input_ids'],
                batch['chosen_attention_mask']
            )
            
            # Forward pass for rejected responses
            rejected_rewards = model.get_rewards(
                batch['rejected_input_ids'],
                batch['rejected_attention_mask']
            )
            
            # Compute loss and metrics
            metrics = compute_ranking_loss_with_metrics(
                chosen_rewards, rejected_rewards
            )
            loss = metrics['loss']
            
            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            
            # Gradient clipping (prevents exploding gradients)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            
            optimizer.step()
            scheduler.step()
            
            # Update metrics
            epoch_metrics['loss'] += loss.item()
            epoch_metrics['accuracy'] += metrics['accuracy'].item()
            
            progress_bar.set_postfix({
                'loss': f"{loss.item():.4f}",
                'acc': f"{metrics['accuracy'].item():.2%}"
            })
        
        # End of epoch
        avg_loss = epoch_metrics['loss'] / len(train_loader)
        avg_acc = epoch_metrics['accuracy'] / len(train_loader)
        print(f"Epoch {epoch+1} - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.2%}")
    
    return model

# Let's demonstrate what happens during a training step
print("Demonstrating a single training step")
print("=" * 50)

# Put the model in training mode
reward_model.train()
optimizer = torch.optim.AdamW(reward_model.parameters(), lr=1e-5)

# Create a synthetic training example
# Chosen: a helpful, informative response
# Rejected: a dismissive, unhelpful response
batch_chosen = tokenizer(
    ["The answer is 42. This is the result from Douglas Adams' novel 'The Hitchhiker's Guide to the Galaxy', where a supercomputer calculated it as the answer to the ultimate question of life, the universe, and everything."],
    return_tensors="pt", padding=True, truncation=True, max_length=64
)
batch_rejected = tokenizer(
    ["I don't know. Just look it up yourself."],
    return_tensors="pt", padding=True, truncation=True, max_length=64
)

batch_chosen = {k: v.to(device) for k, v in batch_chosen.items()}
batch_rejected = {k: v.to(device) for k, v in batch_rejected.items()}

# BEFORE training step
with torch.no_grad():
    chosen_r_before = reward_model.get_rewards(batch_chosen['input_ids'], batch_chosen['attention_mask'])
    rejected_r_before = reward_model.get_rewards(batch_rejected['input_ids'], batch_rejected['attention_mask'])

print("\nBefore training step:")
print(f"  Chosen reward:   {chosen_r_before.item():+.4f}")
print(f"  Rejected reward: {rejected_r_before.item():+.4f}")
print(f"  Margin:          {(chosen_r_before - rejected_r_before).item():+.4f}")

# Compute loss and backprop
chosen_r = reward_model.get_rewards(batch_chosen['input_ids'], batch_chosen['attention_mask'])
rejected_r = reward_model.get_rewards(batch_rejected['input_ids'], batch_rejected['attention_mask'])
loss = compute_ranking_loss(chosen_r, rejected_r)

optimizer.zero_grad()
loss.backward()
optimizer.step()

# AFTER training step
with torch.no_grad():
    chosen_r_after = reward_model.get_rewards(batch_chosen['input_ids'], batch_chosen['attention_mask'])
    rejected_r_after = reward_model.get_rewards(batch_rejected['input_ids'], batch_rejected['attention_mask'])

print(f"\nAfter training step:")
print(f"  Chosen reward:   {chosen_r_after.item():+.4f}")
print(f"  Rejected reward: {rejected_r_after.item():+.4f}")
print(f"  Margin:          {(chosen_r_after - rejected_r_after).item():+.4f}")
print(f"  Loss:            {loss.item():.4f}")

print("\nSee what happened?")
print("  - Chosen reward increased (model likes the good response more)")
print("  - Rejected reward decreased (model likes the bad response less)")
print("  - Margin increased (model is more confident in its ranking)")
print("\nThat's learning in action!")

# Here's a typical training configuration
config = {
    'learning_rate': 1e-5,
    'batch_size': 4,
    'num_epochs': 1,
    'warmup_steps': 100,
    'gradient_accumulation_steps': 4,
    'max_grad_norm': 1.0,
}

print("Reward Model Training Configuration")
print("=" * 50)
for k, v in config.items():
    print(f"  {k:30s} = {v}")

print("\nEffective batch size:", config['batch_size'] * config['gradient_accumulation_steps'])
print("\nThis configuration is conservative but reliable.")
print("It works for most reward modeling tasks without much tuning.")