Preserving Diversity in Supervised Fine-Tuning of Large Language Models

Authors: Ziniu Li, Congliang Chen, Tian Xu, Zeyu Qin, Jiancong Xiao, Zhi-Quan Luo, Ruoyu Sun
Primary area: foundation or frontier models, including LLMs
Keywords: large language models, supervised fine-tuning, diversity, test-time scaling

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.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1730344684935

thread_type:
official_review

<thread_object>
creator: reviewer_whn6

<content>
summary:
The paper introduces a novel method for supervised fine-tuning of large language models (LLMs) called Generative and Entropy-regularized Matching (GEM). GEM replaces the traditional Cross Entropy (CE) loss with a reverse Kullback-Leibler (KL) divergence minimization enhanced with entropy regularization. The approach seeks to mitigate overfitting and increase diversity in model outputs. The authors demonstrate that GEM outperforms CE in instruction-following, math reasoning, code generation, and creative writing tasks, using Llama-3-8B as a model. The method also claims computational efficiency by optimizing a single model and reducing sampling steps while preserving output quality.

strengths:
* The paper addresses known limitations of CE in supervised fine-tuning, such as overfitting and limited diversity, which are crucial challenges for deploying LLMs in diverse tasks.
* GEM's performance is evaluated across multiple tasks and datasets, showing consistent improvements over CE in both general-purpose and domain-specific fine-tuning.

weaknesses:
* While the authors attempt to adapt GEM for sequential data, the proposed solution (using a data distribution “reset” trick) might introduce limitations for real-time applications. This part could benefit from further empirical validation.
* The performance of GEM depends on parameters like the $\beta$ term in entropy regularization. The paper lacks a sensitivity analysis to show how robust GEM is to these hyperparameters
* Although GEM offers an alternative to Cross Entropy, the idea of using entropy regularization to promote diversity is not new, and reverse KL divergence has also been previously explored in generative modeling. The novelty in combining these may be seen as incremental, especially without substantial empirical differentiation from prior methods.

questions:
* In scenarios with limited training data, does GEM perform consistently, or would CE still be preferable due to its simplicity?
* Would the regularization in GEM, especially the entropy component, introduce a non-trivial increase in compute or memory requirements, especially when scaling?
* The paper focuses on the Llama-3-8B model, but it is unclear if GEM’s performance gains generalize across various architectures or if they are specific to this model’s configuration.
* In cases where GEM aims to enhance diversity, could it inadvertently introduce biases in generation, particularly when generating responses with multiple valid answers?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730387336063

thread_type:
official_review

<thread_object>
creator: reviewer_4xm8

<content>
summary:
This paper introduces a new SFT loss, GEM, based on the ideas of reverse KL loss and entropy. This loss can mitigate the overfitting issues of cross-entropy loss in the SFT process. Across a variety of task datasets, GEM demonstrates improved diversity and higher accuracy.

strengths:
1. This article is well-written, with clearly organized content and significant research value.
2. The proposed GEM loss is simple and effective; this straightforward loss improvement simultaneously enhances diversity and accuracy.
3. The experiments are thorough, with validation across various mainstream tasks for large models.

weaknesses:
1. Compared to CE Loss, GEM may introduce additional hyperparameters, which could make training more challenging.
2. There is no comparison of training costs across different methods.

questions:
1. I understand that the introduced GEM can make the model's output distribution more diverse; however, increased diversity usually results in lower accuracy on domain-specific datasets or a higher probability of hallucination errors. Why does GEM improve both generalization and accuracy simultaneously?

2. What is the function of h in Equation 2? Is h necessary? If not, what would happen if h were omitted? And if h is optional, why are only the two variants mentioned in the paper permissible?

3. What is the functional difference between reverse KL and normal KL?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730509728560

thread_type:
official_review

<thread_object>
creator: reviewer_yyc9

<content>
summary:
The use of cross-entropy (CE) loss during the supervised fine-tuning stage often results in overfitting and a reduction in output diversity. To address this limitation, this work introduces a novel training method called GEM. GEM tackles the challenge of distribution matching through reverse KL divergence minimization and incorporates maximum entropy regularization to promote output diversity. The authors demonstrated the effectiveness of GEM across multiple scenarios, showcasing its ability to maintain high performance while reducing overfitting. Additionally, they highlighted that GEM offers benefits for test-time computation, where maintaining diversity is crucial.

strengths:
1. The work is well-motivated, and the writing is clear.
2. They demonstrated increased diversity in the model's output through high performance on the test-time compute set.
3. They designed a method that is tractable for sequential data.

weaknesses:
1. The goal seems to be heavily focused on enhancing the diversity of generations. I believe it would be beneficial to further evaluate how the proposed method performs on tasks where diversity is not as critical. For instance, it would be helpful to show how it performs on tasks where factual accuracy is important.

questions:
1. This is not a critical question, but is there any method to demonstrate that the performance improvement in test-time computation truly results from the increased diversity in the model's output?
2. Is there any performance degradation when increasing the temperature at inference time?
3. Could you also show performance results using LLM-as-a-judge in the Creative Writing section? While diversity is an important dimension, I believe it is also necessary to evaluate other aspects such as the coherency of the writing.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730577529828

thread_type:
official_review

<thread_object>
creator: reviewer_xjk9

<content>
summary:
The authors address the overfitting and decrease in output diversity that arise when using the cross-entropy loss during supervised fine-tuning by introducing a novel distribution matching method called *Generative and Entropy-regularized Matching of distributions* (GEM). GEM is based on two core principles: (1) the model should assign higher probabilities to observed data without overtraining, and (2) as supervised datasets cannot encompass all possible data and thus only cover a limited distribution, the model should learn from both the supervised data and its own generated samples. The first principle is addressed by maximizing the entropy, while the second is achieved by minimizing the reverse KL divergence between the model and target distributions. The authors showed that GEM effectively improves the output diversity, mitigates overfitting, and outperforms the cross-entropy loss (with and without weight decay and entropy regularization) as well as NEFT on multiple benchmarks, including IFEval, HumanEval, GSM8K, MATH, and MBPP.

strengths:
- SFT is a timely and impactful area of research as it is an essential phase of the LLM pipeline.
- The proposed GEM method is grounded in theory, effectively tackles overfitting, and improves the output diversity of Llama3-8B.
- The experimental setup is modern.
- The paper is clear and well-structured.

weaknesses:
- GEM is only evaluated on Llama3-8B, which might particularly benefit from GEM. 
- While Best-Of-N (BON) is useful to show the diversity of the output, it is not a practical solution and therefore does not represent the performance of the model.

questions:
- In addition to Majority Voting (MV) and Best-Of-N (BON), can you report the performance using an LLM as a judge on the same 32 samples?
- Out of the 32 generated samples, how prevalent is the selected answer by MV? In other words, can you report the ratio of the selected answers when using majority voting?
- For the CE + entropy baseline [1], what coefficient γ did you use to weight the regularization? Did you conduct a grid search?
- Line 33: I suggest replacing "Despite extensive pre-trained, " with "Despite extensive pre-training, " or "Despite being extensively pre-trained, ".
- Line 423: Entropy regularization supports Principle 1.

[1] Abhimanyu Dubey, Otkrist Gupta, Ramesh Raskar, and Nikhil Naik. Maximum-entropy fine grained classification. Advances in neural information processing systems, 31, 2018.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732208051798

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response and Summary of Changes

comment:
We sincerely thank all reviewers and area chairs for their efforts in reviewing our paper and providing valuable feedback. We have carefully addressed each concern raised by the reviewers and made corresponding revisions to our paper. We highlight the revision in red. Below, we summarize the key changes in our revision:

- **Section 4.2 (Page 5):** We emphasize the computational advantages of GEM, including *single-model optimization* and *accurate gradient estimation*. This approach allows GEM to effectively solve the challenging reverse-KL-based distribution matching problem and scale efficiently to models with 70B parameters, for which prior GAN-style algorithms may struggle to solve effectively.
- **Section 5.1 (Page 9):** We visualize GEM's test-time scaling performance across different generation budgets, showcasing the benefits of its enhanced diversity. Notably, to match the performance of baselines, GEM often requires only 0.5x the sampling budget.
- **Section 5.2 (Page 10):** We discuss how GEM enhances both diversity and generalization. Additionally, Appendices D.1 and D.2 have been revised to explain why the classical weight decay technique and the naive CE approach with entropy regularization often fall short in achieving these improvements.
- **Appendix B (Page 16 - Page 17):** We detail the implementation of GEM and analyze its computational efficiency, demonstrating that it matches the computation costs of CE.
- **Appendix F.1 (Page 24 - Page 25):**
    - *Quality of Creative Writing:* We use LLM-as-a-judge to evaluate the quality of responses in writing poems and stories. GEM improves both diversity and overall response quality.
    - *Quality of Chatting Responses:* For the chatting task, we employ LLM-as-a-judge to evaluate response quality. GEM outperforms the baselines, largely due to its ability to mitigate overfitting.
    - *Self-Consistency Analysis:* We analyze self-consistency in majority voting by calculating the ratio of majority votes. The results show that GEM achieves high self-consistency across diverse responses, consistent with its generative learning mechanism.
- **Appendix F.3 (Page 27):** We perform a sensitivity analysis of GEM's hyper-parameter to assess their impact. The results demonstrate that GEM is robust to the hyper-parameter.
- **Appendix F.4 (Page 27 - Page 28):** We show the limitations of temperature adjustment as an ad-hoc inference-stage trick for increasing output diversity. Furthermore, we demonstrate that by incorporating diversity into the training process, GEM achieves better quality and diversity compared than this trick.
- **Appendix F.5 (Page 28)**: We show that GEM outperforms CE consistently in low-data regimes. 
- **Appendix F.6 (Page 28):** We provide 24 new evaluation results comparing baseline methods and our approach across models ranging from 3B to 70B parameters, including Qwen-2.5-3B, Qwen-2.5-7B, Gemma-2-9B, and Llama-3.1-70B. The findings highlight that GEM consistently outperforms CE, while baseline methods often fall short of delivering improvements over CE.

We believe that the revisions outlined above have improved the quality of our paper, thanks to the insightful comments from the reviewers. We hope these new results can address the reviewers' concerns and further strengthen the contributions of our work. Thank you once again for your valuable feedback.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733310927378

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Rebuttal Discussion

comment:
Dear Reviewers, ACs, SACs, and PCs,

We sincerely thank you for your efforts in reviewing our paper! We would like to take this opportunity to summarize the rebuttal discussion to facilitate further discussion among reviewers and ACs.

We received four review comments, most of which recognize the soundness, presentation, and contribution of our paper, with evaluations ranging from "good" to "excellent". We are pleased that our responses successfully addressed the concerns raised by Reviewer 4XM8 and Reviewer xjk9. Although we did not receive follow-up feedback from Reviewer yyC9 and Reviewer wHN6, we greatly appreciate their initial comments and have provided detailed responses to address potential misunderstandings. Below, we summarize their main concerns and our responses:

- Reviewer yyC9 recommended evaluating factual accuracy alongside diversity. In response, we provided evidence of our method's performance on instruction-following and math reasoning tasks under greedy decoding, where there is no randomness and the focus is on factual accuracy. Reviewer yyC9 also proposed the use of LLM-as-a-Judge. To address this, we employed the Llama-3.1-70B-Instruct model to evaluate 96,000 responses, demonstrating that GEM not only improves diversity but also enhances response quality, as claimed. We believe these results address Reviewer yyC9's concerns effectively.
- Reviewer wHN6 raised questions about the scalability of GEM to other architectures and its performance on larger models. This is an important point, and we have tested GEM on models with up to 70B parameters across various modern architectures, including Qwen, Llama, and Gemma. Our experiments confirm that GEM performs well across architectures, further validating its robustness. Importantly, GEM operates at the same computational speed as CE, making it a reliable and practical method for practitioners. Additionally, as clarified in our response to Reviewer wHN6, our proposed techniques are new and not covered in prior literature, underscoring the technical contribution of our work.


All review comments have helped us improve the quality of our paper. We have provided over 3 technical discussions and clarifications and presented more than 10 experiment results in the revision (detailed in the "General Response and Summary of Changes" section). These revisions further emphasize our paper's three key contributions:

1. We introduce the framework of entropic distribution matching for SFT of LLMs, addressing the limitations of overfitting and distribution collapse associated with the default CE method.
2. We design a new training method, GEM, that addresses the training challenges of reverse KL divergence minimization and maximum entropy regularization within our framework. Furthermore, GEM is supported by a theoretical guarantee of convergence.
3. Across models ranging from 3B to 70B, we empirically demonstrate that GEM mitigates overfitting of training data and improves the diversity of responses, leading to increased efficiency in test-time scaling.

In summary, we believe our work presents a new framework for fine-tuning LLMs, with techniques that may also be valuable for other directions, such as self-improvement through best-of-n distillation, addressing preference collapse in RLHF, and more (as detailed in Appendix G of our paper).

Thank you again for your time and consideration! We hope the above clarifications are helpful.

Sincerely,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a new distribution matching method called Generative and Entropy-regularized Matching of distributions (GEM), which is based on reverse KL divergence and entropy. It aims to address the issues of overfitting and reduced output diversity that often occur when using cross-entropy (CE) loss in the SFT process. The motivation is clear, the expression is well-articulated, and the experiments are thorough and comprehensive. Although the reviewer raised concerns regarding training costs and the generalizability of the experimental models, the authors conducted thorough experiments during the rebuttal and effectively addressed the issues. Overall, this is a solid work, so the AC recommends accept.