SED-SFT: Selectively Encouraging Diversity in Supervised Fine-Tuning

TL;DR

SED-SFT tackles the problem that standard SFT with Cross-Entropy collapses generation diversity, hindering subsequent RL exploration. It introduces a lightweight, selective entropy regularization term combined with a Top-$k$ based masking strategy, operationalized through $\mathcal{L}_{SED-SFT}(\theta) = \sum_t [ -\log \pi_{\theta}(y^*_t|x,y^*_{<t}) + \lambda M_t \mathcal{L}_{DE}(\pi_{\theta}(y^*_t|x,y^*_{<t})) ]$ with $\mathcal{L}_{DE}(p) = (p-\tfrac{1}{2})^2$, and $M_t = \mathbf{1}\{P_{Top-k}(t) < \tau\}$ where $P_{Top-k}$ is computed from the top-$k$ token probabilities. Empirically, across two backbones and eight math benchmarks, SED-SFT yields average RL improvements of $+2.06$ points on Llama-3.2-3B-Instruct and $+1.20$ points on Qwen2.5-Math-7B-Instruct, with negligible SFT overhead and higher sentence-level diversity (lower Self-BLEU) than CE and several baselines. This demonstrates that selective diversification during SFT can meaningfully enhance downstream RL performance in mathematical reasoning tasks while preserving predictive accuracy where it matters. The work contributes a practical, tunable mechanism to balance diversity and precision in SFT for improved RL outcomes.

Abstract

Supervised Fine-Tuning (SFT) followed by Reinforcement Learning (RL) has emerged as the standard post-training paradigm for large language models (LLMs). However, the conventional SFT process, driven by Cross-Entropy (CE) loss, often induces mode collapse, where models over-concentrate on specific response patterns. This lack of distributional diversity severely restricts the exploration efficiency required for subsequent RL. While recent studies have attempted to improve SFT by replacing the CE loss, aiming to preserve diversity or refine the update policy, they fail to adequately balance diversity and accuracy, thereby yielding suboptimal performance after RL. To address the mode collapse problem, we propose SED-SFT, which adaptively encourages diversity based on the token exploration space. This framework introduces a selective entropy regularization term with a selective masking mechanism into the optimization objective. Extensive experiments across eight mathematical benchmarks demonstrate that SED-SFT significantly enhances generation diversity with a negligible computational overhead increase compared with CE loss, yielding average improvements of 2.06 and 1.20 points in subsequent RL performance over standard CE-based baselines on Llama-3.2-3B-Instruct and Qwen2.5-Math-7B-Instruct, respectively. The code is publicly available at https://github.com/pppa2019/SED-SFT

Figures (3)

• Figure 1: The comparison of Cross-Entropy, pure diversity-encouraging, and SED-SFT. The dashed line token boxes indicate the tokens will be masked in SED-SFT, i.e., the tokens with low exploration space. SED-SFT achieves a balance between accuracy and diversity by avoiding encouraging the masked tokens.
• Figure 2: The heatmap for the probability of the labels on each position. The case is a simple math problem and evaluated on Qwen-0.5B
• Figure 3: The visualization of the cumulative probability with different k values in Top-k.