simpo-loss-function

name	simpo-loss-function
description	Implement SimPO loss with length-normalized rewards and target margin.

SimPO Loss Function Implementation

Overview

SimPO (Simple Preference Optimization) implements a preference optimization objective that uses length-normalized average log probability as an implicit reward, with a target reward margin component.

Key Formula

L_SimPO(πθ) = -E_(x,yw,yl)~D log σ(β/|yw| log πθ(yw|x) - β/|yl| log πθ(yl|x) - γ)

Components

1. Length-Normalized Reward

• Formula: r_SimPO(x, y) = β/|y| * log πθ(y|x)
• Purpose: Average log probability per token, prevents length bias
• Why: Aligns training with generation metric (which uses average log likelihood for beam search)

2. Bradley-Terry Objective

• Formula: p(yw ≻ yl | x) = σ(r(x, yw) - r(x, yl) - γ)
• Purpose: Probabilistic ranking between winning and losing responses
• Function: σ is sigmoid function

3. Target Reward Margin (γ)

• Purpose: Ensure reward difference exceeds a target threshold
• Effect: Improves generalization by enforcing margin between classes
• Typical range: 0.5 to 1.5

Implementation Details

Computing Log Probabilities

# log_probs shape: (batch_size, seq_len)
# Sum across sequence dimension to get total log probability
log_prob_sum = log_probs.sum(dim=1)  # (batch_size,)

# Divide by sequence length for normalization
seq_lengths = (input_ids != pad_token_id).sum(dim=1)  # (batch_size,)
avg_log_prob = log_prob_sum / seq_lengths.float()  # (batch_size,)

Computing Reward Differences

# Batch structure: pairs of (winning, losing) responses
batch_size = avg_log_probs.shape[0]
winning_rewards = avg_log_probs[:batch_size//2]
losing_rewards = avg_log_probs[batch_size//2:]

# Reward difference with margin
reward_diff = beta * winning_rewards - beta * losing_rewards - gamma

Computing Loss

# Bradley-Terry with sigmoid
import torch.nn.functional as F

sigmoid_term = torch.sigmoid(reward_diff)
loss = -torch.log(sigmoid_term).mean()

Common Pitfalls

1. Not using length normalization: Creates bias toward longer sequences
2. Wrong batch structure: Ensure paired winning/losing responses
3. Missing average in log probabilities: Use sum/length, not just sum
4. Gradient flow: Ensure no detach() breaks gradients to model

Hyperparameters

• β: Temperature parameter, typically 2.0-2.5
• γ: Target margin, typically 0.3-1.6, depends on setting
• learning_rate: Usually small, 1e-6 to 5e-7

References

• SimPO Paper: Section 2.3 "The SimPO Objective"
• Length normalization: Equation (3) in paper
• Gradient analysis: Appendix F