💬 NLP

PriFT: Prior-Support Guided Supervised Fine-Tuning

PriFT (Prior-Support Guided Fine-Tuning) improves supervised fine-tuning generalization and RL initialization by deriving stable token reweighting signals from a frozen pretrained model to avoid the self-reinforcing dynamics caused by using the online model's distribution.

Original authors: Ke Wang, Shuangqi Li, Mathieu Salzmann, Pascal Frossard

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ke Wang, Shuangqi Li, Mathieu Salzmann, Pascal Frossard

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.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 have a brilliant student (the AI model) who has already read millions of books and learned a vast amount of general knowledge. This is the pretrained model. Now, you want to teach this student a specific new skill, like solving complex math problems or writing code. This teaching phase is called Supervised Fine-Tuning (SFT).

Usually, teachers show the student a textbook with the correct answers and say, "Memorize this line by line." However, the paper argues that this "line-by-line" memorization has a flaw: the student might try to memorize answers that don't actually make sense to them based on what they already know. This leads to the student overfitting—memorizing the specific textbook so well that they fail when faced with slightly different questions.

The Problem: The "Moving Target" Teacher

Recent attempts to fix this involved a "smart teacher" who looks at what the student is currently learning and decides, "Okay, this answer makes sense to the student right now, so let's focus on it. That other answer is confusing, so let's ignore it."

The paper calls this online reweighting. The problem, according to the authors, is that this teacher is unstable. As the student learns, their brain changes. The "smart teacher" is actually the student's own brain looking in a mirror.

The Trap: If the student initially likes a certain type of answer, the teacher keeps reinforcing it. If the student dislikes another type, the teacher stops teaching it.
The Result: The student becomes a "rich-get-richer" machine. They get really good at the few things they already liked, but they lose the ability to explore new, diverse ways of solving problems. They become rigid and lose the broad knowledge they had before.

The Solution: PriFT (The "Frozen" Reference)

The authors propose a new method called PriFT (Prior-Support Guided Fine-Tuning).

Instead of asking the student's current, changing brain to decide what to learn, PriFT asks the student's original, pre-training brain (the "frozen reference") for advice.

Think of it like this:

The Old Way: You ask a student who is currently stressed and confused, "What should I study?" They might say, "Only the stuff I already know!" because they are scared of the new stuff.
The PriFT Way: You ask the student's "past self" (the version before they started this specific class), "What parts of this new lesson align with what I already know?" The past self gives a stable, calm answer: "Here are the concepts that fit well with your foundation. Focus on these, but don't ignore the rest."

This "past self" provides a Prior Support signal. It tells the model: "This token (word) is supported by your original knowledge, so it's safe to learn. That other token is a stretch, so be careful."

How PriFT Works (Two Flavors)

The paper introduces two ways to use this "past self" advice:

PriFT-prob (The Probability Check): It looks at how likely the "past self" was to say a specific word. If the past self was very confident, it gives that word a high weight. If the past self was unsure, it gives it a lower weight.
PriFT-mass (The Crowd Check): Sometimes, a word might be "easy" (the past self is 99% sure), but that doesn't mean it's the most important word to learn. PriFT-mass looks at the "cumulative mass." It asks, "Is this word supported by at least half of the possible options the past self considered?" This prevents the model from only learning the super-easy, obvious words and forces it to pay attention to harder, more important reasoning steps.

The Results: Why It Matters

The authors tested this on three difficult tasks: Math, Coding, and Medical Questions.

Better Generalization: Models trained with PriFT didn't just memorize the training data; they became better at solving new problems they hadn't seen before.
More Diversity: Unlike the "online" methods that made the model rigid, PriFT kept the model's thinking diverse. It was like keeping a toolbox with many different tools, rather than just one hammer.
Better Prep for Reinforcement Learning (RL): Often, after SFT, researchers use Reinforcement Learning (like a video game where the AI learns by trial and error) to get even smarter. The paper found that PriFT creates a much better "starting point" for this next stage. Because PriFT didn't narrow the model's focus too early, the model had more "room to grow" when it started playing the RL game.

The Bottom Line

The paper claims that by using a frozen, stable reference (the model's original knowledge) to guide the learning process, rather than letting the model's current, changing state dictate the learning, we get a smarter, more flexible AI. It's the difference between a student who blindly follows their current mood versus a student who consults their solid foundation of knowledge to decide what to learn next.

Key Takeaway: PriFT helps AI models learn new skills without forgetting who they are or becoming too narrow-minded, leading to better performance in math, coding, and medicine, and setting them up for even greater success in future training.

Technical Summary: PriFT (Prior-Support Guided Fine-Tuning)

Problem Statement

Supervised Fine-Tuning (SFT) is a standard and computationally efficient method for adapting pretrained Large Language Models (LLMs) to downstream tasks, often serving as the initialization for Reinforcement Learning (RL). However, SFT suffers from a generalization gap compared to RL. As an off-policy method, SFT fits fixed demonstration data token-by-token. When target tokens in the demonstrations are poorly aligned with the model's pretrained distribution, this process can lead to overfitting and the distortion of pretrained knowledge.

Recent approaches attempt to mitigate this by assigning non-uniform training weights to tokens based on their alignment with the model's predictive distribution (e.g., using target token probability or entropy). However, these methods typically compute these weights using the online model currently being fine-tuned. The paper argues that this creates a critical flaw: the predictive distribution of the online model rapidly drifts from the pretrained distribution during training. Consequently, the weighting signal becomes entangled with the optimization trajectory, inducing a self-reinforcing dynamic where the model concentrates learning signals on tokens it initially preferred, while suppressing others. This reduces sampling diversity and limits the model's capacity for further on-policy optimization (RL).

Methodology: PriFT

To address the instability of online reweighting, the authors propose PriFT (Prior-support guided Fine-Tuning). The core insight is that token weights should be derived from a frozen pretrained reference model rather than the evolving online model. This provides a stable estimate of "prior support"—the extent to which a target token is supported by the model's knowledge before task-specific adaptation.

Framework

PriFT modifies the standard SFT objective by introducing token weights $m_t$ derived from the frozen pretrained model $\pi_{\theta_{pt}}$ :
$L_{PriFT}(\theta_{on}) = -\sum_{t=1}^{T} u(\pi_{\theta_{pt}}(\cdot | x, y_{<t}), y_t) \log \pi_{\theta_{on}}(y_t | x, y_{<t})$
where $u(\cdot)$ is a token-utility function. By computing weights from the frozen reference, PriFT decouples weight estimation from the optimization trajectory, preventing the self-reinforcing bias observed in online methods.

Instantiations

The paper introduces two specific instantiations of the utility function $u$ :

PriFT-prob: Uses the pretrained target-token probability directly.
$m_t = \pi_{\theta_{pt}}(y_t | x, y_{<t})$
This scales the gradient magnitude based on how plausible the token was under the pretrained distribution. It preserves the standard SFT update direction but downweights tokens with low prior support.
PriFT-mass: Uses cumulative probability mass under the pretrained distribution to reduce bias toward "easy" tokens with sharp distributions.
$m_t = \mathbb{I}[u^{(MASS)}_t \geq 0.5]$
where $u^{(MASS)}_t = \sum_{v: \pi_{\theta_{pt}}(v) \leq \pi_{\theta_{pt}}(y_t)} \pi_{\theta_{pt}}(v)$ .
This metric measures the fraction of the probability mass assigned to tokens less likely than the target. A threshold of 0.5 selects tokens supported by at least half of the pretrained probability mass. Unlike PriFT-prob, this results in a binary selection mask where all selected tokens receive equal weight, preventing the over-emphasis of trivial high-confidence tokens (e.g., function words) while retaining informative tokens from harder contexts.

Computational Overhead

PriFT requires a single forward pass of the frozen pretrained model to precompute token weights or masks. This adds a relative overhead of approximately $1/(3E)$ (where $E$ is the number of epochs), which can be performed offline. Unlike KL-regularization methods that require access to the reference distribution during every optimization step, PriFT only requires storing scalar weights or binary masks, avoiding significant memory overhead.

Key Contributions

Identification of Signal Drift: The authors demonstrate that online token statistics rapidly drift from the pretrained distribution, creating a self-reinforcing loop that concentrates learning on initially preferred tokens and degrades generalization.
Stable Reweighting via Frozen Reference: They show that deriving token weights from a frozen pretrained model consistently outperforms using the online model across various reweighting rules (probability-based and entropy-based).
PriFT Framework: The proposal of a simple, general framework that precomputes weights from a frozen reference, decoupling the reweighting signal from the optimization trajectory.
Empirical Validation: Extensive experiments showing PriFT achieves state-of-the-art results among SFT baselines and provides superior initialization for subsequent RL training.

Experimental Results

The authors evaluated PriFT on mathematical reasoning, code generation, and medical question answering using models such as Qwen2.5-Math, Qwen3, and LLaMA-2.

Mathematical Reasoning: On Qwen2.5-Math-7B, PriFT-mass achieved an average accuracy (Avg@16) of 40.67 and Pass@16 of 69.60, outperforming the strongest prior baseline (TALR) by 3.97 and 8.75 points, respectively. Notably, while online reweighting methods (like DFT) often improved Avg@16 at the cost of Pass@16, PriFT improved both metrics simultaneously.
Code Generation: On Qwen3-4B-Instruct, PriFT variants outperformed standard SFT and other token-reweighted baselines, particularly on the challenging LiveCodeBench splits. PriFT-mass showed the smallest performance degradation on HumanEval+ compared to the original checkpoint, suggesting better preservation of existing capabilities.
Medical QA: PriFT-mass achieved the best overall performance on MMLU-Medical and MedMCQA benchmarks, outperforming other token-reweighted methods.
RL Initialization: Models initialized with PriFT showed stronger performance after subsequent RL training (DAPO) compared to those initialized with standard SFT or online reweighting methods (DFT). For instance, on Qwen3-8B-Base, PriFT-mass yielded a 9.57 point improvement in Avg@16 after RL compared to the baseline SFT initialization.

Significance and Claims

The paper claims that PriFT addresses a fundamental limitation in token-reweighted SFT: the entanglement of the weighting signal with the optimization trajectory. By utilizing the pretrained model as a stable source of "prior support," PriFT:

Preserves Sampling Diversity: It avoids the self-reinforcing concentration of probability mass on a narrow set of reasoning paths, maintaining broader entropy distributions and better Pass@k performance.
Reduces Distribution Shift: It minimizes the KL divergence from the pretrained model during fine-tuning compared to online reweighting methods.
Enhances RL Readiness: It provides a stronger initialization for RL by preserving the model's capacity for further on-policy optimization, suggesting that "good SFT" is not just about immediate performance but about maintaining the diversity required for future learning.

The authors note limitations, including the focus on moderate-sized open models and specific task domains, leaving the exploration of larger models and broader utility functions for future work. They do not claim PriFT solves all generalization issues but position it as a robust, generalizable improvement to the SFT paradigm.