Learning What to Learn: Stage-Specific Data Sets for SFT-then-RL in Small Language Model Reasoning

A new study proposes a difficulty-aware framework for post-training Small Language Models that assigns reasoning data to distinct supervised fine-tuning and reinforcement learning stages, arguing that the two phases serve fundamentally different purposes [1]. The work, posted to arXiv on June 3, addresses a gap in how researchers prepare data for the common SFT-then-RL pipeline used to improve reasoning in Small Language Models [1]. The authors contend that supervised fine-tuning is better suited for acquiring reasoning skills the model has not yet mastered, while reinforcement learning is more effective at consolidating skills the model can already partially access [1]. Based on that principle, the framework organizes training data into three stage-specific sets: an Acquisition Set, a Consolidation Set, and a Recycled Set [3]. For hard samples in the SFT stage, the researchers introduce a Bridge mechanism that transforms raw teacher-generated reasoning traces into more learnable supervision [1]. The mechanism evaluates each reasoning step by importance, jumpiness, and difficulty, then applies operations such as keep, expand, compress, drop, or localize to produce capacity-aligned supervision for the smaller models [3]. After SFT, samples that yield all-zero rewards during RL — meaning the current policy fails to produce any correct solution under the rollout budget — are excluded from direct policy optimization [3]. Instead of handling these failures within RL through additional search or hinting, the framework routes them into a Recycled Set. A stronger teacher converts the failures into diagnostic, repair, and new reasoning trajectory supervision for the next SFT stage, creating an iterative loop in which SFT expands the model’s reasoning boundary and RL consolidates partially accessible skills [3]. The approach reflects a broader re-examination of the SFT-RL relationship. Other recent work has shown that models initialized from stronger SFT checkpoints can underperform those from weaker ones after identical RL training, a finding attributed to distribution mismatch between offline SFT data and the policy optimized during online RL [5]. Separate research has formulated the integration as a meta-learning problem, treating SFT as an upper-level teacher and RL as a lower-level student to provide targeted guidance that directly supports reward optimization [4]. Experiments on two SLMs across five reasoning benchmarks showed that the proposed method consistently improved over representative SFT, distillation, and RL baselines [1]. The results underscore the importance of coordinating data difficulty across the two stages rather than optimizing each in isolation [1].