🤖 machine learning

Capacity-Dependent Effects of Data Selection for Reasoning

This paper demonstrates that the effectiveness of likelihood-based data selection for reasoning fine-tuning is capacity-dependent, where high-likelihood data benefits smaller models with faster initial gains while low-likelihood data yields superior long-term performance for larger models.

Original authors: Cuong Dang, Hoang Anh Just, Ruoxi Jia

Published 2026-08-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Cuong Dang, Hoang Anh Just, Ruoxi Jia

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 are trying to teach a robot how to solve a tricky math problem. You have a super-smart teacher who can solve it perfectly, but you also have a student who is just starting out. In the world of artificial intelligence, this is called "distillation." The goal is to take the knowledge from a giant, powerful model (the teacher) and shrink it down into a smaller, faster model (the student) so it can run on everyday devices. But here's the catch: the teacher can generate many different ways to solve the same problem. Some answers are simple and look a lot like what the student would have guessed on their own. Others are complex, wild, and very different from the student's current thinking.

For a long time, researchers thought the best way to teach the student was to show them the answers that were easiest for them to understand—the ones that felt familiar. It seemed logical: if you try to teach a beginner calculus using advanced physics, they might get confused and give up. But a new study suggests this "easy is best" rule might be a trap. The real secret isn't about making things easy; it's about matching the difficulty of the lesson to the size of the student's brain. If the student is tiny, easy lessons work great. But if the student is huge and powerful, they actually need the hard, confusing lessons to reach their full potential. This paper explores exactly how to pick the right lessons for the right robot, ensuring we don't waste time or computing power.

The researchers behind this study, Cuong Dang, Hoang Anh Just, and Ruoxi Jia, set out to test a popular idea called "likelihood-based selection." In simple terms, "likelihood" is just a score that tells us how likely a model is to say a specific sentence. If a model thinks, "I would definitely say this," that's a high-likelihood answer. If it thinks, "I would never say this," that's a low-likelihood answer. The paper investigates a specific question: Should we always train our AI students using the high-likelihood answers (the ones that feel familiar), or should we sometimes force them to learn from the low-likelihood answers (the ones that feel strange and difficult)?

To find the answer, the team ran a series of experiments using math problems as their playground. They took a massive pool of math questions and had powerful "teacher" models generate many different solutions for each one. Then, they trained "student" models of various sizes—ranging from very small (1.5 billion parameters) to quite large (8 billion parameters)—using two different strategies. One group of students only saw the "easy" answers (high-likelihood), while the other group only saw the "hard" answers (low-likelihood). They watched how these students performed over time, checking their scores on various math benchmarks like AIME, AMC, and MATH.

What they discovered was a fascinating pattern that depends entirely on the size of the student's brain. For the smallest models, the "easy" strategy was a clear winner. These small students learned quickly and steadily when given answers that were close to what they already knew. When they were forced to look at the "hard" answers, they got lost. They would start repeating the question over and over, or just copy the answer without understanding the steps, failing to improve. It was as if they were trying to climb a mountain but didn't have the legs to make the first step; they just slipped back down.

However, the story changed completely for the larger, more powerful models. These big students started out doing better with the "easy" answers, just like the small ones. But as training continued, something magical happened. The students trained on the "hard" answers began to catch up and eventually overtake the others. While the "easy" group hit a ceiling and stopped improving, the "hard" group kept climbing. They were able to absorb the complex, unfamiliar reasoning and use it to solve problems they couldn't touch before. The researchers call this the "Fast-Fit / Slow-Gain" pattern. The small models "fast-fit" to the easy data and stop there, while the large models experience a "slow-gain" where the difficult data eventually leads to much stronger reasoning skills.

The paper suggests that there is no single "best" way to pick training data. It all depends on the computing budget and the size of the model. If you have a small model or limited time, stick to the high-likelihood, easy-to-learn data. It's efficient and gets you quick results. But if you have a large model and the patience to train it for a long time, you should embrace the low-likelihood, difficult data. It's like a marathon runner: a beginner needs a flat, easy path to build confidence, but an elite athlete needs the steep, rocky trails to break their personal bests.

The authors back this up with both real-world experiments and a theoretical framework. They showed that small models simply lack the "capacity" to understand the gap between their current knowledge and the teacher's complex answers. The difficult data is too far away, and the small model can't bridge the gap. But large models have enough "room" in their brains to stretch and learn from that distance. The paper concludes that data selection for reasoning shouldn't be a one-size-fits-all rule. Instead, it should be a careful balance: match the difficulty of the lesson to the capacity of the learner, and you'll get the best results.

Technical Summary: Capacity-Dependent Effects of Data Selection for Reasoning

Problem Statement

In the supervised fine-tuning (SFT) of Large Language Models (LLMs) for reasoning tasks, a critical challenge arises from the heterogeneity of candidate responses for a single instruction. Even when generated by strong teacher models, responses vary significantly in how well they align with the student model's current distribution. Recent work, such as GRAPE (Zhang et al., 2025), has proposed selecting high-likelihood responses (those closest to the student's current policy) to ensure stable optimization. Conversely, traditional active learning and data selection literature often advocates for low-likelihood (hard) examples to push models beyond their current capabilities.

This paper investigates a fundamental conflict: Does high-likelihood data universally yield the best reasoning performance, or does its efficacy depend on the student model's capacity and training duration? The authors challenge the assumption that "easier" (high-likelihood) data is always preferable, positing that the value of likelihood-based selection is contingent on whether the student possesses sufficient representational capacity to absorb challenging supervision.

Methodology

The authors conduct controlled experiments across a spectrum of model scales and training durations to isolate the effects of data likelihood and model capacity.

Datasets and Models: Experiments utilize the MATH12K dataset. Teacher models (Qwen2.5-72B and Gemma-3-27B-it) generate multiple candidate responses for each instruction. Student models range from 1.5B to 8B parameters, including various Qwen and Llama families.
Selection Strategies: Two distinct datasets are constructed for each instruction:
- $D_{high}$ : Responses with the highest log-likelihood under the base student model.
- $D_{low}$ : Responses with the lowest log-likelihood under the base student model.
Training Protocol: Models undergo full-parameter SFT for 5 epochs. Performance is evaluated using pass@1 on 10 diverse mathematical reasoning benchmarks (e.g., AIME, AMC, GPQA, MATH).
Analysis: The study analyzes learning dynamics through two lenses:
1. Teacher View: Measuring the log-likelihood of student outputs under the teacher model to assess alignment with the teacher distribution.
2. Original Model View: Measuring the log-likelihood of student outputs under the student's own base model to assess distribution shift.
Theoretical Framework: The authors provide a theoretical analysis of capacity-constrained distillation using linear classifiers. They model the student's feasible parameter space as a subspace and analyze how gradient flow transfers knowledge from teacher to student based on the projection of the data span onto this subspace.

Key Results

1. The "Fast-Fit / Slow-Gain" Pattern

The experiments reveal a clear capacity-dependent dichotomy in learning dynamics:

Small Models (e.g., 1.5B, 3B): High-likelihood data provides a distinct early advantage ("Fast-Fit"). These models achieve rapid, stable improvements and often reach their peak performance early. Low-likelihood data fails to provide meaningful gains; the models struggle to absorb the supervision, often resulting in repetitive or shallow outputs that do not improve reasoning capabilities.
Large Models (e.g., 7B, 8B): While high-likelihood data still offers strong early gains, low-likelihood data demonstrates a "Slow-Gain" trajectory. Initially lagging, these models eventually surpass the performance of those trained on high-likelihood data given sufficient training duration (e.g., after 2–5 epochs). The low-likelihood data pushes large models toward the teacher distribution, unlocking stronger reasoning abilities that high-likelihood data cannot access.

2. Learning Dynamics and Distribution Shift

Teacher Alignment: Small models trained on low-likelihood data fail to move closer to the teacher distribution, sometimes drifting further away. In contrast, larger models successfully reduce the teacher loss when trained on low-likelihood data, indicating effective knowledge transfer.
Stability: Small models exhibit unstable training dynamics and distribution shifts when exposed to low-likelihood data, often reverting to rote repetition. Larger models adapt in a more stable and focused manner, successfully traversing the optimization landscape to reach higher global peaks.

3. Theoretical Insights

The theoretical analysis formalizes these observations. It posits that knowledge transfer is constrained by the intersection of the data span (what the data reveals about the teacher) and the student's feasible parameter space (what the student can represent).

For high-capacity students, the feasible space is large enough to capture the knowledge gap exposed by difficult (low-likelihood) data.
For low-capacity students, the feasible space is a strict subspace. If the knowledge exposed by low-likelihood data lies outside this subspace, the student cannot learn it, rendering the supervision ineffective. High-likelihood data, being closer to the student's current representation, lies within this feasible space, allowing for incremental but limited improvement.

Key Contributions

Capacity-Dependent Data Selection: The paper demonstrates that high-likelihood data is not universally optimal. The effectiveness of likelihood-based selection is governed by an interaction between data distribution, model capacity, and computing budget.
Learning Dynamics Analysis: The authors characterize the "Fast-Fit / Slow-Gain" phenomenon, showing that small models benefit from easy data for rapid convergence, while large models require hard data for long-term capability expansion.
Theoretical Framework: A capacity-constrained view of distillation is developed, showing that successful transfer requires the exposed teacher knowledge to lie within the student's representable subspace. This explains why low-likelihood data fails for small models but succeeds for large ones.

Significance and Claims

The paper concludes that data selection strategies for reasoning should not be treated as a one-size-fits-all rule. Instead, practitioners must consider the specific constraints of their setup:

For small models or limited compute budgets: High-likelihood data is preferable for achieving fast, stable, and efficient early improvements.
For large models with sufficient training budgets: Low-likelihood data becomes increasingly valuable, offering the potential for significantly higher final performance by pushing the model toward the teacher's distribution.

The authors emphasize that the "best" data selection strategy is context-dependent, determined jointly by the student's capacity and the available training duration. This work refines the understanding of distillation and data engineering, moving beyond simple heuristics to a more nuanced, capacity-aware approach.