🤖 AI

Towards On-Policy SFT: Distribution Discriminant Theory and its Applications in LLM Training

This paper introduces Distribution Discriminant Theory (DDT) to enable On-Policy Supervised Fine-Tuning, proposing In-Distribution Finetuning and Hinted Decoding techniques that achieve generalization performance surpassing offline RL algorithms like DPO and SimPO while retaining the computational efficiency of SFT.

Original authors: Miaosen Zhang, Yishan Liu, Shuxia Lin, Xu Yang, Qi Dai, Chong Luo, Weihao Jiang, Peng Hou, Anxiang Zeng, Xin Geng, Baining Guo

Published 2026-03-17

📖 5 min read🧠 Deep dive

CC0 1.0

Original authors: Miaosen Zhang, Yishan Liu, Shuxia Lin, Xu Yang, Qi Dai, Chong Luo, Weihao Jiang, Peng Hou, Anxiang Zeng, Xin Geng, Baining Guo

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a brilliant but very particular student (the AI) how to solve math problems. You have two main ways to teach them:

The "Textbook" Method (Supervised Fine-Tuning / SFT): You give the student a stack of answer keys from a textbook. You say, "Memorize these answers exactly."
- The Problem: The textbook answers might be written in a style the student doesn't naturally use, or they might contain tiny errors. If you force the student to memorize them blindly, they might forget their own natural way of thinking or start making up nonsense just to match the book. They become a "parrot" rather than a thinker.
The "Practice Test" Method (Reinforcement Learning / RL): You let the student try to solve problems on their own. When they get it right, you give them a gold star. When they get it wrong, you take it away.
- The Problem: This is incredibly slow and expensive. You have to run thousands of practice tests to get a few gold stars. Also, if the student gets stuck, you might not know why they failed, making it hard to fix.

The Big Idea:
The authors of this paper asked: "Can we get the speed of the Textbook Method but the smart, natural thinking of the Practice Test Method?"

They realized the answer is yes, but only if we stop forcing the student to memorize everything in the textbook and start teaching them only the parts that match their natural brain.

Here is how they did it, broken down into three simple concepts:

1. The "Vibe Check" (Distribution Discriminant Theory)

Imagine the student has a specific "vibe" or natural rhythm when they speak. Some answers in the textbook feel like they were written by a robot; others feel like they were written by a human.

The authors created a mathematical "Vibe Check" (called DDT). It's a tool that scans every single word in a textbook answer and asks: "Does this word feel like something our student would naturally say?"

If the word feels natural (In-Distribution), the tool says, "Good! Keep this."
If the word feels forced or weird (Out-of-Distribution), the tool says, "Wait, this doesn't sound like you. Let's ignore this part."

The Analogy: Think of it like a DJ mixing music. If a song has a beat that clashes with the current vibe, the DJ doesn't just play it louder; they fade it out or skip it so the dance floor stays smooth.

2. The "Smart Highlighter" (In-Distribution Finetuning / IDFT)

Once they have the "Vibe Check," they built a new training method called IDFT.

In normal training, if the student gets a word wrong compared to the textbook, the teacher screams, "NO! Change your brain!" This causes the student to panic and forget everything they already knew (this is called Catastrophic Forgetting).

With IDFT, the teacher is much smarter:

Easy/Normal Words: "Great job! Keep doing what you're doing." (Reinforces their natural style).
Weird/Out-of-Sync Words: "That's not quite your style. Let's not stress about this specific word right now." (Protects their brain from getting confused).

The Analogy: Imagine you are learning to paint. A normal teacher forces you to copy a masterpiece exactly, even the messy brushstrokes you don't understand. The IDFT teacher says, "Copy the colors and shapes you understand well, but if the artist used a weird technique you don't get, just leave a blank space there for now. Don't ruin your own painting style trying to force it."

3. The "Ghostwriter" (Hinted Decoding)

Sometimes, the textbook answer is correct, but the student just can't generate it because it sounds too foreign. The authors created a trick called Hinted Decoding.

Imagine the student is writing an essay. They know the facts (the answer is correct), but they are struggling with the style (how to say it).

Old Way: The student tries to copy the textbook style and fails.
New Way (Hinted Decoding): The teacher whispers the facts to the student ("The answer is 42"), but tells the student, "Now, write the rest of the story in your own voice."

The AI generates the answer using its own natural style, but guided by the correct facts. It's like having a ghostwriter who knows the facts but writes in the author's unique voice.

Why is this a Big Deal?

It's Fast: It's as fast as reading a textbook (SFT), not slow like running thousands of practice tests (RL).
It's Safe: It prevents the AI from "forgetting" how to be a helpful assistant while learning new math tricks.
It's Better than RL: Surprisingly, their method actually performed better than the expensive Reinforcement Learning methods on math and coding tasks, using way less computer power.

In a Nutshell:
The paper says, "Don't force the AI to be a robot that mimics textbooks. Instead, teach it using a filter that only lets in the knowledge that fits its natural personality. This way, it learns faster, remembers more, and stays smart."

1. Problem Statement

The paper addresses the performance gap between Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) in Large Language Models (LLMs).

The Gap: While SFT is computationally efficient, it often yields inferior generalization compared to RL (e.g., RLHF, RLVR). This is primarily attributed to RL's use of on-policy data (data generated by the model itself), which preserves the model's native distribution and mitigates catastrophic forgetting. In contrast, standard SFT forces the model to fit external data equally, regardless of whether that data aligns with the model's internal distribution, potentially damaging pre-trained knowledge structures.
The Challenge: RL is often infeasible due to the difficulty of obtaining reliable, verifiable feedback (e.g., in open-ended tasks) or its sparse reward signals leading to low learning efficiency.
The Goal: To bridge the SFT-to-RL chasm by enabling On-Policy SFT. The objective is to retain SFT's data efficiency while enhancing generalization by aligning the training process with the model's own distribution at both the data and objective levels.

2. Methodology

The authors propose a framework built on a new theoretical foundation and two complementary techniques.

A. Distribution Discriminant Theory (DDT)

DDT provides a rigorous mathematical framework to quantify whether a token or sequence is "in-distribution" (sampled from the model's own distribution) or "out-of-distribution" (OOD).

Hypothesis Testing: The problem is framed as distinguishing between $H_0$ (token sampled from model distribution $p_t$ ) and $H_1$ (token from external distribution $q_t$ ).
Optimal Statistic: Using Signal Detection Theory (SDT) and maximizing the Signal-to-Noise Ratio (SNR), the authors prove that the Centered Log-Likelihood (CLL) is the optimal criterion for discrimination.
- Formula: $\phi_t = \log p_t(x_t) + H[p_t]$ , where $H[p_t]$ is the entropy of the model's distribution at step $t$ .
- Insight: Unlike raw log-likelihood ( $\log p_t$ ), CLL subtracts the entropy, effectively normalizing for context difficulty. This eliminates variance caused by the inherent complexity of the prompt, making the statistic robust for detecting true distributional shifts.
Sequential Properties: The cumulative score of CLL over a sequence behaves as a zero-mean martingale under $H_0$ (in-distribution) and exhibits a negative linear drift under $H_1$ (OOD), allowing for error bounds and reliable detection.

B. Application 1: In-Distribution Finetuning (IDFT)

IDFT is a loss-level modification to standard SFT that uses the DDT statistic to adaptively reweight the training objective.

Mechanism: It introduces a modulation coefficient $\gamma_t = \exp(-\phi_t)$ $γ_{t} = exp (- ϕ_{t})$ .
- In-Distribution Tokens ( $\phi_t \approx 0$ ): $\gamma_t \approx 1$ , maintaining standard learning.
- OOD Tokens ( $\phi_t \ll 0$ ): $\gamma_t > 1$ , triggering gradient suppression to prevent the model from overfitting to noise or data that disrupts its native structure.
- Strong-Domain Tokens ( $\phi_t > 0$ ): $\gamma_t < 1$ , enhancing knowledge consolidation.
Goal: To prevent catastrophic forgetting by suppressing harmful gradients from OOD samples while reinforcing high-confidence, in-distribution reasoning steps.

C. Application 2: Hinted Decoding

Hinted Decoding is a data-level technique that transforms fixed training datasets into on-policy samples aligned with the model's distribution.

Process: Given a question and a ground-truth answer, the algorithm generates a response that maintains the correctness of the answer but adopts the stylistic distribution of the model.
Mechanism: It employs a two-stream decoding approach:
1. Imitator Stream: Conditioned on the question and answer (providing correctness).
2. Model Stream: Conditioned only on the question (providing native style).
Adaptive Mixing: The final token distribution is a geometric mixture of the two streams, weighted by the entropy of the Immitator's output.
- High entropy (uncertainty) $\rightarrow$ Rely more on the Model Stream (preserving style).
- Low entropy (certainty) $\rightarrow$ Rely more on the Immitator (ensuring correctness).
Safety Mechanism: Includes a "False Positive" detection mechanism to switch back to standard decoding once the Chain-of-Thought (CoT) delimiter is reached, ensuring the final answer remains consistent with the reasoning.

3. Key Contributions

Distribution Discriminant Theory (DDT): A novel theoretical framework establishing the optimality of Centered Log-Likelihood (CLL) for distribution discrimination, grounded in signal detection theory and martingale properties.
IDFT Algorithm: A loss-level method that dynamically suppresses gradients for OOD tokens, effectively mitigating catastrophic forgetting without requiring RL infrastructure.
Hinted Decoding: A training-free inference algorithm that re-aligns static datasets to the model's native distribution, enabling "on-policy" training from offline data.
Empirical Validation: Demonstrating that this SFT-based framework can surpass prominent offline RL algorithms (DPO, SimPO) in generalization performance while maintaining SFT's computational efficiency.

4. Experimental Results

The authors evaluated their framework on advanced LLMs (Qwen2.5, DeepSeek-R1, Mistral, LLaMA) across mathematical reasoning and code benchmarks.

Performance vs. Baselines:
- IDFT significantly outperformed standard SFT, Dynamic Fine-Tuning (DFT), and Entropy-Adaptive Fine-Tuning (EAFT) on rigorous math benchmarks (AIME, AMC, Olympiad).
- HD + IDFT (Hinted Decoding + IDFT) consistently surpassed offline RL methods (Rejection Sampling, DPO, SimPO) across Math-C, Math-G, and General Reasoning benchmarks.
Efficiency: The proposed method achieved superior results with lower computational budgets (fewer GPU hours) compared to RL baselines, which often require extensive rollouts.
Catastrophic Forgetting: The framework successfully preserved general capabilities (MMLU, ARC) while learning new domain-specific skills, a common failure point for standard SFT on strong instruct models.
Ablation Studies: Confirmed that the adaptive weighting in Hinted Decoding and the continuous modulation in IDFT are critical; static thresholds or simple re-weighting failed to match performance.

5. Significance and Impact

Bridging the Gap: The work proves that the performance gap between SFT and RL is largely due to distributional misalignment, not just the learning algorithm itself. By aligning data and objectives with the model's distribution, SFT can achieve RL-level generalization.
Practical Alternative: It offers a viable, efficient alternative to RL for domains where reward modeling is difficult, expensive, or impossible (e.g., complex reasoning, agent tasks without verifiable ground truth).
Theoretical Foundation: Provides a rigorous statistical basis (DDT) for understanding and controlling distributional shifts in LLM training, moving beyond heuristic loss modifications.
Open Source: The authors have open-sourced the code, facilitating further research into on-policy SFT and distribution-aware training.

In summary, this paper redefines the SFT paradigm by introducing a theoretically grounded method to distinguish and prioritize in-distribution data, enabling LLMs to learn effectively from static datasets without the overhead and instability of traditional RL.