📊 statistics

InfoSFT: Learn More and Forget Less with Information-Aware Token Weighting

InfoSFT is a principled token-weighting scheme for supervised fine-tuning that concentrates learning signals on maximally informative, medium-confidence tokens to improve generalization across diverse tasks while better preserving the model's pre-existing capabilities compared to standard SFT and existing mitigation methods.

Original authors: Mahdi Sabbaghi, George Pappas, Adel Javanmard, Hamed Hassani

Published 2026-05-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Mahdi Sabbaghi, George Pappas, Adel Javanmard, Hamed Hassani

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

The Big Picture: Teaching a Robot New Tricks

Imagine you have a very smart robot (a Large Language Model) that has already learned a lot about the world. Now, you want to teach it a specific new skill, like solving complex math problems or writing computer code. You do this by showing it examples of experts doing the task. This process is called Supervised Fine-Tuning (SFT).

However, the paper points out a problem with the standard way we do this. It's like a teacher trying to teach a student by forcing them to memorize every single example in a textbook, no matter how easy or hard it is.

The Problem: The "One-Size-Fits-All" Trap

The paper identifies two main issues with the standard approach:

Overfitting (The Parrot Effect): The robot tries to memorize every example perfectly, even the ones that are very strange or unlikely for it to understand. It's like a student who memorizes the exact wording of a math problem but doesn't understand the logic, so they fail when the numbers change.
Catastrophic Forgetting (The Amnesia Effect): In trying to learn these new, difficult tricks, the robot accidentally "forgets" the things it already knew well. It's like a chef who, in trying to learn a new, fancy recipe, forgets how to boil water or chop onions.

The Old Solutions: Too Much or Too Little

Previous methods tried to fix this by filtering out the "hard" examples (the ones the robot finds unlikely).

The Analogy: Imagine a teacher who only lets the student practice on problems they can already solve easily.
The Result: The student gets very confident and stable, but they never learn anything new because they never struggled with the difficult concepts they needed to master.

The New Solution: InfoSFT (The "Goldilocks" Teacher)

The authors propose a new method called InfoSFT. Instead of treating all examples the same, or ignoring the hard ones, InfoSFT acts like a wise teacher who knows exactly how much to push the student.

The Core Idea: The "Sweet Spot" of Confidence
InfoSFT assigns different "weights" (importance scores) to every word the robot generates, based on how confident the robot is about that word:

Too Easy (High Confidence): If the robot is already 99% sure of a word (like "the" or "is"), InfoSFT says, "You already know this; stop wasting time on it." It gives these words zero weight.
Too Hard (Zero Confidence): If the robot is completely lost and the word is extremely unlikely, InfoSFT says, "This is too confusing right now; we'll skip it for now to avoid breaking your brain." It gives these words very low weight.
Just Right (Medium Confidence): If the robot is unsure but has a good guess (maybe 50-80% confidence), InfoSFT says, "This is the perfect learning moment! Focus here!" It gives these words the highest weight.

The Analogy:
Think of learning to ride a bike.

Standard SFT is like forcing you to ride on a flat sidewalk (too easy) and then suddenly throwing you into a hurricane (too hard). You either get bored or crash.
InfoSFT is like a coach who puts you on a gentle hill. You are wobbly and unsure (medium confidence), so you learn the most. If the hill is too steep, the coach stops you. If the ground is flat, the coach lets you coast.

Why This Works Better

The paper shows that by focusing on these "medium-confidence" moments, the robot learns the new behavior (like math or coding) much faster and better than before.

Better Generalization: The robot learns the logic of the task, not just the specific examples. It can solve new problems it hasn't seen before.
Less Amnesia: Because the robot isn't forced to drastically change its entire personality to memorize weird examples, it keeps its original skills (like safety and general conversation) intact.

The "One-Line" Magic

The authors emphasize that this isn't a massive overhaul. It's a tiny tweak to the math used to train the robot. They simply changed the formula that decides how much attention to pay to each word. It's like changing the volume knob on a stereo: instead of playing all songs at the same volume, you turn up the volume on the songs you need to hear and turn down the ones you already know.

Summary

InfoSFT is a smarter way to teach AI. It stops the AI from memorizing the easy stuff and getting confused by the impossible stuff. Instead, it focuses all its energy on the "just right" moments where learning actually happens. This helps the AI master new skills without forgetting the old ones.

Technical Summary: InfoSFT

Problem Statement

Supervised Fine-Tuning (SFT) is the standard approach for teaching Large Language Models (LLMs) new behaviors using offline expert demonstrations. However, standard SFT suffers from two critical limitations:

Weak Test-Time Generalization: Standard SFT uniformly fits all samples, including those with low likelihood under the base model. These low-likelihood tokens incur large negative-log-likelihood losses, forcing the model to overfit specific samples rather than learning the underlying target behavior. Consequently, SFT often underperforms compared to on-policy methods (like RL) in generalization benchmarks.
Catastrophic Forgetting: Adapting to new behaviors using low-likelihood offline data induces substantial shifts in the model's output distribution. This large distributional shift is correlated with the degradation of prior capabilities (e.g., safety, instruction following, or general reasoning), a phenomenon known as catastrophic forgetting.

Existing mitigation strategies, such as filtering or down-weighting low-likelihood data, often suppress the very novel behaviors the base model needs to learn, creating a tradeoff between learning new tasks and retaining old ones.

Methodology: InfoSFT

The authors propose InfoSFT, a principled token-weighting scheme designed to concentrate learning signals on "maximally informative, medium-confidence tokens." The method is derived from a proximal update framework, optimizing the tradeoff between fitting the expert distribution and staying close to the base model under a fixed KL-divergence budget.

Theoretical Derivation

The authors frame SFT as a policy gradient problem where the weighting of a sample depends on its likelihood $q = \pi_\theta(y^*|x)$ under the current model.

Standard SFT assigns a uniform weight ( $w=1$ ), effectively treating all tokens equally.
DFT (Likelihood-Weighted SFT) assigns weights proportional to likelihood ( $w \propto q$ ), which stabilizes training but underweights informative low-likelihood samples.
InfoSFT derives an optimal weighting rule that prioritizes fitting the population distribution of expert data rather than memorizing individual samples. The optimal weight for a token with likelihood $q$ $q$ is derived as:
$w_{info}(q) \propto q \left( \text{logit}(p) - \text{logit}(q) \right)_+$
Where:
- $p$ is an estimate of the average expert confidence (a calibration constant).
- $\text{logit}(q) = \log(\frac{q}{1-q})$ .
- $(\cdot)_+$ denotes clipping to ensure non-negative weights.

Key Characteristics

Middle-Confidence Focus: The weighting function peaks for tokens with intermediate confidence. It assigns high weights to tokens that are unlikely under the base model (and thus informative) but not so unlikely that they cause instability.
Vanishing Weight for Extremes:
- For highly confident tokens ( $q \to 1$ ), the correction term vanishes, avoiding unnecessary reinforcement of tokens the model already predicts well.
- For extremely unlikely tokens ( $q \to 0$ ), the weight vanishes, preventing instability.
Implementation: InfoSFT requires only a one-line modification to the standard token-wise loss function. It replaces the uniform weight with the derived weighting factor.

Key Contributions

Optimal Update under KL Budget: The paper theoretically demonstrates that, when constrained to remain close to the base model (fixed KL budget), the optimal update strategy assigns the highest weights to samples with intermediate confidence, rather than uniform or likelihood-proportional weights.
InfoSFT Algorithm: A simple, novel SFT variant that incorporates this weighting rule. It emphasizes informative low-likelihood tokens while giving little weight to already-solved tokens.
Complementarity with SFT: The authors show that InfoSFT and standard SFT are complementary. SFT is effective for introducing very unlikely formats (e.g., new reasoning tags like <thought>), while InfoSFT is superior for refining the model once those formats become more likely.

Experimental Results

The authors evaluated InfoSFT across math, code, and Chain-of-Thought (CoT) reasoning tasks using diverse model families (Qwen-2.5, Llama-3.1) and sizes.

Generalization Performance: InfoSFT consistently outperformed both vanilla SFT and DFT on test-time benchmarks (MATH500, AMC, HumanEval, MultiPL-E). Notably, it achieved higher pass@8 scores, indicating better diversity and robustness in generating correct solutions among multiple samples.
Catastrophic Forgetting: In experiments measuring the tradeoff between new-task accuracy and prior capabilities (Science Q&A and Tool Use), InfoSFT achieved a superior Pareto frontier. It reached higher new-task accuracy than SFT and DFT while preserving prior capabilities (e.g., HumanEval, IFEval) more effectively.
Entropy Control: Unlike DFT, which often leads to entropy collapse (deterministic outputs), InfoSFT maintains higher output entropy, which is crucial for downstream alignment and sampling-based inference.
CoT Reasoning: In CoT fine-tuning, a two-stage approach (1 epoch SFT followed by 1 epoch InfoSFT) yielded the best results on AIME, leveraging SFT to learn the format and InfoSFT to refine the reasoning.

Significance and Claims

The paper claims that token weighting is a key ingredient for reliable supervised post-training. By moving beyond uniform fitting or simple likelihood weighting, InfoSFT addresses the fundamental tension between learning new behaviors from offline data and preserving pre-existing capabilities. The method offers a principled, low-overhead solution that improves generalization and mitigates catastrophic forgetting without requiring complex data filtering or on-policy sampling. The authors position InfoSFT as a robust alternative that bridges the gap between the stability of SFT and the generalization benefits of on-policy methods.