Pear Sft Preparation

PEAR: SFT Preparation for Reinforcement Learning

Problem

Standard supervised fine-tuning (SFT) optimizes for behavior policy distribution—the policy that generated training data. But downstream RL operates under the target policy, creating distributional mismatch.

Models trained with uniform SFT may not prepare well for RL. The mismatch between offline SFT optimization and online RL sampling creates sub-optimal initialization.

Core Concept

PEAR (Policy Evaluation-inspired) reweights SFT loss using importance sampling ratios between target and behavior policies. Rather than training uniformly, concentrate on tokens whose continuations remain plausible under the target policy.

This ensures offline updates focus on trajectories the RL stage will actually revisit, improving final post-RL performance.

Architecture Overview

• Importance Sampling Weights: Compute target/behavior policy likelihood ratios
• Token-Level Reweighting: Apply importance weights to individual token losses
• Suffix Ratio Computation: Use future continuations to estimate target likelihood
• Block-Level Stabilization: Partition sequences for gradient stability
• Sequence-Level Granularity: Uniform trajectory-wide weights for simplicity

Implementation

Step 1: Compute Importance Sampling Weights

Calculate policy likelihood ratios for reweighting.

def compute_importance_weights(batch, behavior_model, target_model):
    """Compute importance sampling weights from behavior to target policy."""
    weights = []

    for sequence in batch:
        tokens = tokenize(sequence)
        sequence_weight = 1.0

        for token_idx in range(len(tokens)):
            # Compute probability under behavior policy
            behavior_logits = behavior_model.get_logits(tokens[:token_idx])
            behavior_probs = F.softmax(behavior_logits, dim=-1)
            behavior_prob = behavior_probs[tokens[token_idx]].item()

            # Compute probability under target policy (being optimized)
            target_logits = target_model.get_logits(tokens[:token_idx])
            target_probs = F.softmax(target_logits, dim=-1)
            target_prob = target_probs[tokens[token_idx]].item()

            # Importance ratio
            ratio = target_prob / (behavior_prob + 1e-6)
            sequence_weight *= ratio

        weights.append(sequence_weight)

    return torch.tensor(weights)

Step 2: Token-Level Reweighting

Apply importance weights to individual token losses.

def token_level_reweighting(model, batch, behavior_model, target_model, learning_rate=1e-4):
    """Reweight SFT loss using token-level importance sampling."""
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    for sequence in batch:
        tokens = tokenize(sequence)

        # Compute suffix ratios for token-level weights
        suffix_weights = []
        cumulative_ratio = 1.0

        for token_idx in range(len(tokens) - 1, -1, -1):
            # Suffix ratio: probability of tokens[token_idx:] under target vs behavior
            suffix_target = compute_suffix_prob(target_model, tokens, token_idx)
            suffix_behavior = compute_suffix_prob(behavior_model, tokens, token_idx)

            suffix_ratio = suffix_target / (suffix_behavior + 1e-6)
            suffix_weights.insert(0, suffix_ratio)
            cumulative_ratio *= suffix_ratio

        # Apply weights to loss
        for token_idx in range(len(tokens)):
            target_token = tokens[token_idx]
            context_tokens = tokens[:token_idx]

            # Forward pass
            logits = model.get_logits(context_tokens)
            loss = F.cross_entropy(logits, target_token, reduction='none')

            # Weight loss by importance ratio
            weighted_loss = loss * suffix_weights[token_idx]
            weighted_loss.backward()

        optimizer.step()
        optimizer.zero_grad()

Step 3: Block-Level Stabilization

Partition sequences for more stable gradient updates.

def block_level_reweighting(model, batch, importance_weights, block_size=64, learning_rate=1e-4):
    """Apply block-level importance weighting for gradient stability."""
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    for sequence, weight in zip(batch, importance_weights):
        tokens = tokenize(sequence)

        # Partition into blocks
        for block_start in range(0, len(tokens), block_size):
            block_end = min(block_start + block_size, len(tokens))
            block_tokens = tokens[block_start:block_end]

            # Compute loss for block
            block_loss = 0
            for token_idx, token in enumerate(block_tokens):
                context = tokens[:block_start + token_idx]
                logits = model.get_logits(context)
                loss = F.cross_entropy(logits, token)
                block_loss += loss

            # Apply weight to entire block
            weighted_block_loss = (block_loss / len(block_tokens)) * weight

            weighted_block_loss.backward()

        optimizer.step()
        optimizer.zero_grad()

    return model

Step 4: Training Loop with PEAR

Integrated SFT training with importance reweighting.

def train_sft_with_pear(model, training_data, behavior_model, num_epochs=3, learning_rate=1e-4):
    """Train model with importance-sampled loss for RL preparation."""
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    for epoch in range(num_epochs):
        # Batch processing
        for batch in create_batches(training_data, batch_size=32):
            # Compute importance weights
            weights = compute_importance_weights(batch, behavior_model, model)

            # Token-level reweighting
            for sequence, weight in zip(batch, weights):
                tokens = tokenize(sequence)

                for token_idx in range(len(tokens)):
                    context = tokens[:token_idx]
                    target = tokens[token_idx]

                    logits = model.get_logits(context)
                    loss = F.cross_entropy(logits, target)

                    # Weight by importance ratio
                    weighted_loss = loss * weight.item()
                    weighted_loss.backward()

            optimizer.step()
            optimizer.zero_grad()

    return model

Practical Guidance

Hyperparameter Configuration

When to Use

• Preparing models for post-training with RL
• Improving AIME/competition math performance via subsequent GRPO
• Scenarios where behavior data distribution differs from target policy
• Maximizing final performance after expensive RL tuning
• Any SFT→RL pipeline where offline initialization matters

When Not to Use

• When RL is not planned (standard SFT sufficient)
• Behavior and target policies are very similar
• Extreme importance weights due to distribution shift
• Online learning scenarios (PEAR is offline-only)
• Data from uniform random behavior (weights meaningless)

Common Pitfalls

1. Extreme importance weights: Distribution mismatch causes unbounded weights. Clip or use variance reduction.
2. Weight instability: Early epochs with poor target model estimation. Start with behavior model as target.
3. Data quality sensitivity: Importance sampling amplifies noise in low-probability examples. Filter outliers.
4. Insufficient samples: Token-level ratios need stable estimates. Ensure large behavior dataset.

Reference

Good SFT Optimizes for SFT, Better SFT Prepares for Reinforcement Learning https://arxiv.org/abs/2602.01058