margin-adaptive-dpo

name	margin-adaptive-dpo
title	Margin Adaptive DPO: Leveraging Reward Model for Granular Control in Preference Optimization
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2510.05342
keywords	["Direct Preference Optimization","DPO","Reward Models","Preference Learning","Temperature Adaptation"]
description	Adapt DPO temperature parameters per sample using reward model margins, amplifying learning signals for hard preference pairs while dampening easy ones.

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:

1. Reward Model Training: Learn the preference margin for each training pair (how much stronger is one response over another?)

2. 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 / (preference_pairs))
        metrics[].append(reward_epoch_loss / (preference_pairs))

     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

Reference

Paper: https://arxiv.org/abs/2510.05342