adu2021 Technique: Instance-Level Margin Adaptation in DPO
Direct Preference Optimization (DPO) is a widely-used preference learning method, but it treats all training examples uniformly with a single fixed temperature parameter. In practice, some preference pairs are obvious (one response is clearly much better), while others are subtle (slight quality differences). MADPO addresses this by using reward models to detect preference difficulty and adapt the temperature per sample.
The key insight is that margin (the strength of preference) varies across examples. For hard discrimination tasks with low margins, amplifying the learning signal helps the model learn the distinction. For easy cases with high margins, dampening prevents overconfident updates that lead to overfitting.
Core Concept
MADPO extends DPO with two stages:
Reward Model Training: Learn the preference margin for each training pair (how much stronger is one response over another?)
Adaptive Weighting: Use these margins to modulate the DPO loss, creating instance-level temperature adaptation that improves both convergence and generalization.
Architecture Overview
- Reward Model Stage: Train a standard reward model on preference pairs
- Margin Computation: Extract preference margins from learned model
- Loss Reweighting: Apply margin-dependent weights to DPO loss
- Joint Optimization: Train policy and reward model together
Implementation Steps
Train a reward model to capture preference margins.
def train_reward_model(model, preference_pairs, optimizer, num_epochs=5):
"""
Train a reward model that predicts how much stronger one response
is compared to another.
Args:
model: Reward model
preference_pairs: List of (query, chosen, rejected) tuples
optimizer: PyTorch optimizer
num_epochs: Training epochs
Returns:
losses: Training curve
"""
import torch
import torch.nn.functional as F
losses = []
for epoch in range(num_epochs):
epoch_loss = 0.0
for query, chosen, rejected in preference_pairs:
# Get reward scores for both responses
chosen_reward = model(query, chosen)
rejected_reward = model(query, rejected)
# Margin: how much we prefer chosen over rejected
margin = chosen_reward - rejected_reward
# Loss: encourage margin > 0 (chosen better) with margin ~1.0
loss = F.relu(1.0 - margin) # Contrastive loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss += loss.item()
losses.append(epoch_loss / len(preference_pairs))
return losses
Compute preference margins and create weight function.
def compute_preference_margins(reward_model, preference_pairs):
"""
Compute the margin (strength of preference) for each pair.
Args:
reward_model: Trained reward model
preference_pairs: List of (query, chosen, rejected) tuples
Returns:
margins: Array of margin values
"""
import torch
margins = []
with torch.no_grad():
for query, chosen, rejected in preference_pairs:
chosen_score = reward_model(query, chosen)
rejected_score = reward_model(query, rejected)
margin = (chosen_score - rejected_score).item()
margins.append(margin)
return margins
def create_margin_weights(margins, margin_threshold=0.5, scaling_factor=2.0):
"""
Create adaptive weights based on preference margins.
Hard pairs (low margin) get amplified; easy pairs (high margin) get dampened.
Args:
margins: Array of margin values
margin_threshold: Threshold separating hard from easy pairs
scaling_factor: How much to amplify/dampen
Returns:
weights: Per-sample weight array
"""
import numpy as np
margins = np.array(margins)
# Normalize margins to [0, 1]
normalized_margins = (margins - margins.min()) / (margins.max() - margins.min() + 1e-8)
# Weight function: amplify low-margin (hard) cases
# Dampen high-margin (easy) cases
weights = np.where(
normalized_margins < margin_threshold,
# Hard pairs: amplify signal
1.0 + (margin_threshold - normalized_margins) * scaling_factor,
# Easy pairs: dampen signal
1.0 - (normalized_margins - margin_threshold) * (scaling_factor / 2)
)
return weights
Implement Margin Adaptive DPO loss.
def madpo_loss(model, query, chosen, rejected, margin_weight=1.0, beta=0.5):
"""
Compute margin-adaptive DPO loss.
Args:
model: Policy model
query: Input query
chosen: Preferred response
rejected: Non-preferred response
margin_weight: Instance-level weight from margin adaptation
beta: Temperature parameter (inverse)
Returns:
loss: Weighted DPO loss
"""
import torch
import torch.nn.functional as F
# Get log probabilities from policy
chosen_logp = model.get_log_prob(query, chosen)
rejected_logp = model.get_log_prob(query, rejected)
# Standard DPO formulation
logits = beta * (chosen_logp - rejected_logp)
# DPO loss with adaptive weighting
loss = -F.logsigmoid(logits)
# Apply margin-based weight
weighted_loss = margin_weight * loss
return weighted_loss
Full training loop with margin adaptation.
def train_madpo(policy_model, reward_model, preference_pairs,
policy_optimizer, reward_optimizer, num_epochs=5):
"""
Train policy and reward model jointly using margin-adaptive DPO.
Args:
policy_model: Language model policy
reward_model: Reward model
preference_pairs: List of (query, chosen, rejected) tuples
policy_optimizer: Optimizer for policy
reward_optimizer: Optimizer for reward model
num_epochs: Training epochs
Returns:
metrics: Training metrics
"""
import torch
metrics = {'policy_loss': [], 'reward_loss': []}
for epoch in range(num_epochs):
policy_epoch_loss = 0.0
reward_epoch_loss = 0.0
# Stage 1: Train reward model to estimate margins
for query, chosen, rejected in preference_pairs:
chosen_reward = reward_model(query, chosen)
rejected_reward = reward_model(query, rejected)
margin = chosen_reward - rejected_reward
# Reward model loss
reward_loss = torch.nn.functional.relu(1.0 - margin)
reward_optimizer.zero_grad()
reward_loss.backward()
reward_optimizer.step()
reward_epoch_loss += reward_loss.item()
# Stage 2: Compute margins and train policy with adaptation
margins = compute_preference_margins(reward_model, preference_pairs)
weights = create_margin_weights(margins)
for (query, chosen, rejected), weight in zip(preference_pairs, weights):
# MADPO loss with margin weighting
policy_loss = madpo_loss(
policy_model, query, chosen, rejected,
margin_weight=weight, beta=0.5
)
policy_optimizer.zero_grad()
policy_loss.backward()
policy_optimizer.step()
policy_epoch_loss += policy_loss.item()
metrics['policy_loss'].append(policy_epoch_loss / len(preference_pairs))
metrics['reward_loss'].append(reward_epoch_loss / len(preference_pairs))
return metrics
Practical Guidance
| Aspect | Recommendation | Notes |
|---|---|---|
| Margin threshold | 0.3-0.7 | Separates hard/easy pairs; adjust based on margin distribution |
| Scaling factor | 1.5-3.0 | Controls amplification/dampening strength |
| Reward model architecture | Simple linear or small MLP | Lightweight; avoids overfitting |
| Training order | Reward model then policy | First estimate margins, then adapt |
| When to use | Diverse preference distributions | Especially useful with naturally hard/easy pairs |
| When NOT to use | Homogeneous easy preferences | Adaptation provides limited benefit if all pairs have similar difficulty |
| Common pitfall | Unstable margin estimation | Use validation set to monitor reward model quality |
When to Use MADPO
- Preference datasets with variable annotation quality or confidence
- Fine-tuning where margin signals indicate task difficulty
- Scenarios where cost is sensitive to confidence (harder to judge = more important)
When NOT to Use MADPO
- Small preference datasets where stable estimation is difficult
- Tasks where all pairs are similarly difficult
- Real-time training with tight budget constraints
Common Pitfalls
- Reward model drift: Monitor reward model validation loss separately
- Margin scale mismatch: Normalize margins before creating weights
- Extreme weight values: Clip weights to avoid unstable gradients
- Overfitting to reward model: Use separate validation set for margin estimation