On-policy SFT and Rejection Sampling Fine-Tuning

931 단어·2 분·원문(.md)

The problem mentioned in Off-Policy Distillation is the method of forcibly memorizing others' answers. In other words, it's a limitation of rote learning where one can only memorize and cannot fix problems that arise outside of the predefined path. Let's first look at On-Policy SFT (Rejection Sampling Fine-Tuning, RFT), a transitional yet powerful methodology in practice, which emerged to overcome this limitation.

Does performance improve if we only select good samples generated by the current model and re-SFT, even without RL? #

Yes, it does. Tremendously. Especially in domains where clear scoring criteria (Verifiers) exist, such as solving math problems or coding, it shows miraculous improvements.

In the previous chapter, we mentioned that memorizing the Teacher model's answers creates an Imitation Gap.

This is because the Teacher's brain structure (parameter space) and the Student's brain structure are different.

However, RFT (Rejection Sampling Fine-Tuning) discards the Teacher and makes the Student think multiple times on its own.

For a single problem, the Student's temperature is increased, and it's instructed to solve it in, say, 100 different ways.

95 of them will be nonsense, and by chance, 5 solutions will emerge that reach the correct answer through its own on-policy logic.

These 5 are what we call "perfectly fitting" logical developments that align perfectly with the student's brain structure.

If these are selected and used to SFT itself again, it's not about forcibly memorizing, but rather a process of transforming into a realization that it vaguely knew itself.


Generate N -> Score -> Filter -> SFT Pipeline #

Let's break down how to build this pipeline by stack.

Step 1: Generate N (Mass Generation) #

  • stack: vLLM

Instead of giving the model the same question and asking it only once, we provide n=56 and temperature=0.7 to make it generate 50 diverse reasoning paths for a single question in parallel.

Step 2: Score & Filtering (Scoring and Filtering) #

  • stack: Python Verfiler, Regex Parser, Unit Test, LLM-as-a-Judge

This is the most crucial step. It filters out the garbage from the 50 answers and keeps only the truly correct ones.

  • Math: Run a Python regex to extract only the \boxed{answer} part and check if it matches the actual correct answer (Rule-Based Grader).
  • Coding: Execute the generated code in a sandbox and run Unit Tests.
  • General Conversation: If there's no single correct answer, use a smarter model like GPT-5-5 or an LLM as a Judge to score whether the answer followed the instructions well.

Step 3: SFT (Retraining) #

  • stack: TRL SFTTrainer

Gather only the high-quality On-Policy data that passed filtering and run another round of training with the SFTTrainer code we saw in the previous SFT section.

import re
from vllm import LLM, SamplingParams

# 1. 현재 Student 모델 로드 (On-policy)
model = LLM(model="my-student-8B-model")

prompt = "문제: 12와 18의 최소공배수를 구하고 최종 답은 \boxed{} 안에 넣어줘."

# 2. Generate N: 동일 문제에 대해 20개의 다른 풀이 생성
sampling_params = SamplingParams(n=20, temperature=0.8, max_tokens=512)
outputs = model.generate([prompt], sampling_params)

# 3. Score & Filter: 파이썬 정규식 Verifier
good_samples = []
ground_truth = "36"

for output in outputs[0].outputs:
    text = output.text
    # \boxed{숫자} 형태를 찾는 정규식 파서
    match = re.search(r'\\boxed\{(.+?)\}', text)
    
    if match:
        extracted_answer = match.group(1).strip()
        if extracted_answer == ground_truth:
            # 정답을 맞춘 '풀이 과정(Reasoning Path)'만 수집
            good_samples.append({
                "instruction": prompt,
                "output": text # 자기 자신이 만들어낸 성공적인 풀이 과정
            })

print(f"20개 중 정답을 도출한 훌륭한 샘플 {len(good_samples)}개 확보 완료. SFT로 넘어갑니다.")

Why do we use complex RL when this is so good? #

RFT is easy to implement and performs well, but it has clear limitations.

  1. No Partial Credit (Binary Reward): In SFT, if it's not the correct answer, it's wrong. What if the reasoning process was brilliant up to 90%, but there was one small addition error? In RFT, it's considered incorrect and discarded entirely. The model cannot learn from that valuable brilliant reasoning ability.
  2. Waste of Resources (Loss of Negative Signal): The model also needs to learn information about how it thinks incorrectly, like a notebook of wrong answers. Since SFT is done only with correct data, it cannot impose constraints like "don't do it this way!"

We know it got it right, but can't we subtly distinguish between better and less good solutions to precisely refine the model's weights? When it generates an incorrect answer, can't we penalize it to prevent it from going down such a foolish path again?

What we'll explore next time to solve this problem are RLHF (PPO, DPO), which provides sophisticated rewards to the model through mathematical formulas, and reasoning-centric reinforcement learning, GRPO.

It's a type of post-training where the model deepens its own thinking.

AI/onsft.md