GEM: How to Preserve Diversity During SFT

Paper (ICLR 2025): Preserving Diversity in Supervised Fine-Tuning of Large Language Models.

The Problem with Cross-Entropy

The standard SFT loss is the cross-entropy (CE) loss. From the perspective of the token distribution space, its optimization objective pulls the token distribution at every position toward the one-hot distribution (the ground truth). This is an all (source tokens) -to- one (target token) update, and its optimization objective has a tendency to reduce entropy.

CE loss pulls every position's token distribution toward the one-hot target (Figure 2 in the paper)

GEM: From All-to-One to Sparse Updates

GEM's idea is to change the all-to-one update into a sparse update: only the source tokens in the high-probability region are required to move closer to the target token, while the low-probability region is largely left alone.

The implementation of the GEM loss is quite simple: introduce a softmax distribution q with temperature β (which does not participate in gradient updates), and multiply it by logprob_pred - logprob_truth.

Algorithm 1: GEM (from the paper)

Listing 1: PyTorch code of GEM (from the paper)

Concretely, the auxiliary distribution is defined as:

q_{k}(y \mid x) = \operatorname{softmax}\!\Big( \frac{1}{\beta} \log f_{\theta_k}(y \mid x) \Big)

And the GEM loss is:

\mathcal{L}_{\mathrm{GEM}}(\theta) = \sum_i \sum\nolimits_{y^{\mathrm{gene}}} q_k(y^{\mathrm{gene}} \mid x_i) \cdot \Big[ \log f_\theta(y^{\mathrm{gene}} \mid x_i) - \log f_\theta(y^{\mathrm{real}}_i \mid x_i) \Big]

That is, q acts as a multiplicative weighting over the difference between the predicted log-probability and the ground-truth log-probability at each position.

The Only Hyperparameter: β

The only parameter in GEM is β:

As β approaches 0, the update becomes sparser, promoting diversity more strongly;
When β = 1, GEM degenerates into the standard SFT (CE) loss;
The value of β needs to be selected through experiments; the authors use β = 0.7 in the paper.

Experimental Results

Table 2: Output diversity of pre-trained, CE fine-tuned, and GEM fine-tuned LLMs — higher values indicate greater diversity

From the experimental results, whether it is output diversity, the Best-of-N trend on task metrics, or the effect of subsequent RL training, GEM outperforms SFT. In the Best-of-N experiments, GEM's efficiency even reaches 2× that of SFT. In addition, GEM comes with a bonus: it mitigates catastrophic forgetting.

The Math: Nash Equilibrium

Mathematically, the Nash equilibrium of GEM is reverse KL + entropy regularization:

f^{\star} = \operatorname{argmin}_{f} \; \mathbb{E}_{x} \Big[ D_{\mathrm{KL}}\big( f(\cdot \mid x) \,\|\, p(\cdot \mid x) \big) - \gamma \, \mathcal{H}\big( f(\cdot \mid x) \big) \Big]

However, its effect differs from directly optimizing SFT + entropy regularization — the latter causes the probabilities of long-tail, low-quality tokens to inflate, whereas GEM does not exhibit this behavior.

Figure 11: Next-token probability distributions of CE, CE+Entropy, and GEM-LS — CE+Entropy assigns high probabilities to nonsensical long-tail tokens (marked with crosses), while GEM-LS does not

Who Uses It

In Alibaba's NeurIPS 2025 paper Teaching Language Models to Reason with Tools, the GEM algorithm is used in the SFT stage.

Closing Thoughts

This paper views SFT through the lens of probability transition, which I find quite interesting. I previously wrote Understanding SFT from an RL Perspective — readers who are interested may also want to take a look.

References

[1] Preserving Diversity in Supervised Fine-Tuning of Large Language Models (ICLR 2025) — arxiv.org/abs/2408.16673

[2] GEM official PyTorch implementation — github.com/liziniu/GEM

[3] Teaching Language Models to Reason with Tools (NeurIPS 2025) — arxiv.org/abs/2510.20342