pear-sft-preparation

name	pear-sft-preparation
title	Good SFT Optimizes for SFT, Better SFT Prepares for Reinforcement Learning
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2602.01058
keywords	["SFT","Reinforcement Learning","Importance Sampling","Loss Reweighting","AIME"]
description	Improve post-RL performance by reweighting SFT loss using importance sampling. Prioritize training examples that match the target policy distribution, not the behavior policy. Achieves 14.6% Pass@8 gains on AIME.

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=-)
            target_prob = target_probs[tokens[token_idx]].item()

            
            ratio = target_prob / (behavior_prob + )
            sequence_weight *= ratio

        weights.append(sequence_weight)

     torch.tensor(weights)