[Paper Review] Preserving Diversity in Supervised Fine-Tuning of Large Language Models

Ziniu Li, Congliang Chen+5 more|arXiv (Cornell University)|Aug 29, 2024

Magnetic confinement fusion research4 citations

TL;DR

This paper proposes GEM (Generative Entropy-regularized Matching), a novel supervised fine-tuning method for large language models that replaces cross-entropy loss with reverse Kullback-Leibler divergence minimization combined with entropy regularization. GEM reduces overfitting and significantly improves generation diversity, leading to up to 10-point performance gains on code generation and math reasoning benchmarks using best-of-n sampling, outperforming standard cross-entropy and weight decay baselines on Llama-3-8B.

ABSTRACT

Large Language Models (LLMs) typically rely on Supervised Fine-Tuning (SFT) to specialize in downstream tasks, with the Cross Entropy (CE) loss being the de facto choice. However, CE maximizes the likelihood of observed data without accounting for alternative possibilities. As such, CE usually leads to reduced diversity in the model's outputs, which hinders further development that requires sampling to explore better responses. To address this limitation, this paper introduces a new game-theoretic formulation for SFT. In this framework, an auxiliary variable is introduced to regulate the learning process. We prove that the proposed game-theoretic approach connects to the problem of reverse KL minimization with entropy regularization. This regularization prevents over-memorization of training data and promotes output diversity. To implement this framework, we develop GEM, a new training algorithm that is computationally efficient as CE by leveraging some unique properties of LLMs. Empirical studies of pre-trained models from 3B to 70B parameters show that GEM achieves comparable downstream performance to CE while significantly enhancing output diversity. This increased diversity translates to performance gains in test-time compute scaling for chat and code generation tasks. Moreover, we observe that preserving output diversity has the added benefit of mitigating forgetting, as maintaining diverse outputs encourages models to retain pre-trained knowledge throughout the training process.

Motivation & Objective

• To address overfitting and limited generation diversity in supervised fine-tuning (SFT) of large language models (LLMs), which commonly result from aggressive updates of cross-entropy (CE) loss.
• To develop a principled alternative to CE loss that encourages flatter, more diverse output distributions while preserving data fidelity.
• To improve generalization and performance in downstream tasks—especially those relying on advanced decoding strategies like best-of-n sampling—by enhancing model diversity.
• To demonstrate the effectiveness of entropy-regularized reverse KL divergence minimization as a tractable and effective alternative to standard SFT training.
• To provide a scalable, efficient method compatible with standard SFT pipelines that can be applied to both general and domain-specific instruction tuning.

Proposed method

• Proposes GEM (Generative Entropy-regularized Matching), a new training objective based on minimizing reverse Kullback-Leibler (KL) divergence between the data distribution and the model’s generative distribution, with an entropy regularizer to prevent over-memorization.
• Uses a tractable optimization scheme that avoids adversarial training, making it computationally efficient and compatible with standard SFT frameworks.
• Introduces a loss function that balances data likelihood and output diversity by maximizing entropy while still fitting the supervised data.
• Employs a reparameterization trick to enable differentiable optimization of the entropy-regularized reverse KL objective during fine-tuning.
• Applies the method to Llama-3-8B using both general (UltraFeedback) and domain-specific (MetaMathQA, Magicoder-OSS-Instruct) datasets to evaluate generalization and diversity.
• Uses best-of-n sampling and majority voting to evaluate diversity gains, with pass@100 and accuracy metrics on code and math benchmarks.

Figure 1 : Illustration of the difference between the standard CE and the proposed method GEM for SFT of LLMs.

Experimental results

Research questions

• RQ1 Can replacing cross-entropy loss with entropy-regularized reverse KL divergence minimization reduce overfitting in supervised fine-tuning of LLMs?
• RQ2 Does the proposed GEM method improve generation diversity in LLMs without requiring domain-specific data?
• RQ3 To what extent does GEM enhance performance on downstream tasks that benefit from diverse outputs, such as math reasoning and code generation?
• RQ4 How does GEM compare to standard CE loss and weight decay regularization in terms of generalization and robustness across different datasets?
• RQ5 Can entropy regularization mitigate mode collapse and improve performance under advanced decoding strategies like best-of-n sampling?

Key findings

• On the UltraFeedback dataset for instruction tuning, GEM reduced perplexity and improved IFEval performance, indicating less overfitting compared to CE loss.
• For math reasoning on GSM8K and MATH, GEM-LS achieved 2.6-point and 1.6-point improvements in BON@32 over CE, respectively, demonstrating enhanced diversity and generalization.
• On code generation benchmarks, GEM-LS achieved a 9.7-point (14.7% relative) improvement over CE in Pass@100 on HumanEval and an 8.0-point (11.1% relative) gain on MBPP.
• On the more challenging MATH benchmark, GEM-LS improved performance by 1.9 points (8.0% relative) in greedy decoding and 1.7 points (5.8% relative) in majority voting, showing robustness to difficult tasks.
• Weight decay regularization showed inconsistent improvements, while GEM consistently outperformed CE across all benchmarks and settings, confirming its effectiveness in reducing overfitting and increasing diversity.
• The results indicate that GEM’s entropy regularization effectively mitigates overfitting and enables better performance with best-of-n sampling, validating its utility in test-time computation and self-improvement pipelines.

Figure 2 : Performance of using advanced generation strategies such as best-of-n and majority voting in chatting (left), math reasoning (middle) and code generation (right) tasks.