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
Venue: ICLR 2025 Poster (2025)
Keywords: Large language models
supervised fine-tuning
diversity
test-time scaling
Primary area: foundation or frontier models, including LLMs
TL;DR: We design a game-theoretic framework for preserving diversity and mitigating forgetting during the fine-tuning of large language models

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.

Published as a conference paper at ICLR 2025
PRESERVING DIVERSITY IN SUPERVISED FINE-TUNING OF LARGE LANGUAGE MODELS
Ziniu Li! Congliang Chen Tian Xu², Zeyu Qin', Jiancong Xiao
Zhi-Quan Luo and Ruoyu Sun 1,2,6.t
'The Chinese University of Hong Kong, Shenzhen
Shenzhen Research Institute of Big Data
Naajjng University
*Hong Kong University of Science and Technology University Pennsylvania
'Shenzhen International Center for Industrial and Applied Mathematics
fziniuli, congliangchen @linkk cubk. edu, en, xut aut@lamee aatellmmd..oo Lamda.nju...................u. cdu cn,
zeyu, gincconnect...ccmmmmmmmmmmm ust. hk, jexiaolupenr. edu, (luezq, aunuuoyuJ@uuhe edu.
ABSTRACT
Large Language Models (LLMs) typically rely on Supervised Fine-Tuning (SFT) to specialize 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 develop
data and promotes output diversity. To implement this framework, 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 CE while significantly enhaneing 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, maintaining diverse outputs encourages models to retain pre-trained knowledge throughout the training process."
INTRODUCTION
Large Language Models (LLMs) (OpenAI, 2023; Dubey et al., 2024) powerful generative models that excel across specialized tasks. Despite extensive training, pre-trained LLMs often generate tokens without truly understanding users' queries, resulting irrelevant answers. To enhance performance on specific tasks, researchers employ instruction tuning (Raffel et al., 2020; Wei et al., 2021; Chung al., 2024), also known as Supervised Fine-Tuning (SFT) (Ouyang al., 2022; Bai al., 2022); see Figure 1. This process involves fine-tuning models on high high-quality prompt-response pairs, with the Cross Eross-Entropy (CE) loss being the de facto choice for optimization.
However, is CE the most suitable choice for SFT? To answer this question, we must examine SFT's role in the broader context of model development. As the initial post-training stage, SFT teaches LLMs to follow instructions. Once these models can produce responses that are clear, interpretable, and verifiable for annotators, they become well-suited leveraging human feedback in downstream tasks via techniques like reinforcement learning (Li et al., 2024c). From this perspective, SFT's primary objective is to align the output space of pre-trained LLMs with the preferred formats of downstream tasks. This alignment lays the groundwork for subsequent learning paradigms, including
t: Corresponding author. Code available https: / / teiithub.cc. com/ / liziniu/GMM.

---

Published as a conference paper at ICLR 2025
(Rafailov et al., 2023), self- self-immrrvement with synthetie data (Adler et al., 2024), and test-time scaling strategies (Snell et al., 2024; Brown et al., 2024; Wu et al., 2024).
Throughout these post-training stages, output diversity emerges as critical factor that enables exploration and identification of higher-quality responses. Beyond its benefits for developers, users also value diversity it provides multiple options to meet their customized needs." Unfortunately, using CE loss in SFT undermines this goal: while pre-trained LLMs naturally generate diverse outputs (Brown al., 2020; Wang et al., 2024a), fine-tuning with CE loss often significantly reduces diversity (O" Mo'Mahony et al., 2024; Kim et al., 2024) (see our own evaluation in Table 2 in the Appendix). This limitation motivates us reconsider CE's application in LLMs.
Question: Marissa hiking 12-mile trail. She took hour to walk the first miles, then another hour walk next two miles. she wants average speed miles per hour, what speed (in miles per hour) does she need to walk the remaining distance?
Answer: ##
Pre-trained LLM
SFT via CE
Our Goal of SFT
g
X
#*
X
AN
X
Mariiss valsed..
X
M16
which
Marisss has sheady wallond...
go.
Bldyy Fysur walding
X
Mariasa walledd.... tB
X
##
Readability
Diversity
Figure 1: Illustration of diversity preservation in SFT. While pre-trained LLMs produce diverse outputs, these often lack proper formatting. Standard SFT using CE improves readability but reduces diversity. We aim to maintain output diversity while enhancing the readability of LLMs' responses.
In theory, CE maximizes the likelihood of the given label, which proved successful for classification models (Krizhevsky et al., 2012; He et al., 2016). However, this process inherently overlooks alternative possibilities, fundamentally misaligning with the nature of open-ended generation tasks in LLMs, where responses to identical queries can legitimately vary in format, style, or reasoning paths. While this limitation is less problematic during pre-training with massive datasets, it particularly severe in SFT scenarios where training data is comparatively limited in both size and coverage.
Furthermore, we identify an another key feature of LLMs: they are extensively pre-trained, with rich knowledge encoded in the network. This feature again misaligns with CE during fine-tuning: by suppressing alternative plausible outputs during training, CE can lead "knowledge forgetting" or "alignment tax", eroding the pre-trained model's s encoded knowledge (Ouyang et al., 2022; Luo et al., 2023; Li et al., 2024b). This suggests that "diversity reduction" and "knowledge forgetting" are be deeply interconnected in LLMs. To address these issues, regularization techniques like weight decay (Touvron et al., 2023; Burns et al., 2023) or noisy perturbations to embeddings (Jain et al., 2023) are commonly applied alongside CE loss. However, they have their own limitations (see discussion in Appendix D), highlighting the need for more principled solutions.
distribution matching problem. Specifically, we conceptualize the learning process as the transfer of probability mass from source tokens to target tokens within the supervised dataset. To model this, we introduce game-theoren framework that incorporates an auxiliary variable (i.e., a cotaa-oontroller,. which governs the direction and magnitude of probability transfer. This framework enables the development of learning strategies that preserve diversity and mitigate overfitting, as detailed in the main text.
Contributions. In this paper, we frame the SFT of LLMs
2 Al interfaces like ChatGPT address this need by offering features such regeneration buttons. "How to preserve diversity and interestingness" is recognized as an open problem in the talk "ChatGPT and The Art of Post- training" by Barret Zoph and John Schulman (Zoph & Schulman, 2025). We conjecture this process is irreversible in post-training and focus on addressing this problem in the first stage (i.c., SFT).
2

---

Published as a conference paper at ICLR 2025
From theoretical perspective, we demonstrate that our approach connects to distribution matching problem that involves reverse Kullback-Leible (KL) divergence minimization with maximum entropy regularization, thereby ensuring the benefits of diversity. Alongside this theoretical insight, we develop a practical training algorithm, GEM, for the game. GEM leverages some unique properties of LLMs, offering computational efficiency and scalability comparable to optimizing the CE loss.
We empirically validate the effectiveness of GEM, by fine-tuning pre-trained models ranging from 3B to 70B in size. We have two main findings. First, GEM enhances the model's ability generate more diverse responses, which translates into improved performance in test-time scaling. For example, when fine-tuning Llama-3.1-4
(Dubey al., 2024). GEM outperforms CE by achieving 5-point improvement in chatting and an 8-point improvement in code generation tasks through repeated sampling. Importantly, to match the performance of baselines, GEM often requires only 0.5x the sampling budget. Second, by preserving output diversity, GEM also mitigates forgetting, demonstrating an 83% reduction in alignment tax.
To summarize, our contributions are threefold:
We introduce a game-theoretie framework for SFT, enabling a controlled design to preserve diversity. Specifically, we theoretically prove that it can achieve the distribution matching with reverse KI. minimization and maximum entropy regularization.
We develop new training method GEM that effectively solves the game in practice. This algorithm offers computational efficiency advantages over previous methods designed for similar games.
We empirically validate that preserving diversity in distribution matching enhances test-time scaling performance while reducing forgetting and alignment tax.
2
PRELIMINARY
Large Language Models. Let f be the generative distribution modeled by the language model, where fy]lz) denotes the conditional distribution of response y given the prompt Typically, is parameterized by a Transformer (Vaswani et al., 2017), with parameters 8. A key feature of f is that the generation space ) is finite, which enables favorable optimization properties that we will discuss later. Note that y and x may be sequential, in which case an auto-regressive formulation is employed.
Supervised Fine- Tane-TTnnnggaa and the Gross-Entropp Method. To specialize in downstream tasks, LLM relies on SFT after pre-training. This process involves using a supervised dataset with high-quality prompt-response pairs (ar, y). sampled from the prompt distribution p and conditional data distribution p(y(z). The Cross Entropy (CE) loss is the de facto training objective for SFT, designed to maximize the likelihood of the training data. Formally, this is expressed as:
min (0) = Ex-ppyypp (x)) [log fo fo(zz))]
Here, the prompt distribution p is typically not modeled during SFT and can be treated as a constant; for simplicity, we omit it when the context is clear. A key feature of this approach is that it exclusively maximizes the likelihood of the observed data, disregarding alternative plausible responses. Please refer to Figure 10 in the Appendix for the theoretical understanding of CE.
3 CHALLENGES AND PRINCIPLES FOR SFT
Before exploring technical solutions, we establish guiding principles for SFT. We note that SFT is rarely the final stage of LLM development; subsequent phases such as preference learning (Rafailov et al., 2023; Azar et al., 2024; Wang et al., 2024b), reinforcement learning (Li et al., 2024c; Shao et al., 2024), and advanced inference-tine time strategies (Snell al., 2024) heavily depend on sampling and output diversity to explore high-quality solutions. This reliance on diversity underscores key challenge: while pre-trained L.LMs produce diverse outputs due to their broad knowledge bases, standard SFT practices particularly the use of CE loss-ofter reduces this diversity (O" Mahony et al., 2024; Wang et al., 2024a). Such reduction can lead to knowledge forgetting, aligning with the "alignment tax" phenomenon observed in (Bai et al., 2022; Ouyang et al., 2022).
We argue that preserving output diversity during SET can address these issues. Intuitively, the ability indicator of the richness its retained knowledge.
of model to generate diverse responses serves By maintaining diversity, the model is compelled to consider alternative plausible responses, which in turn necessitates that its internal parameters encode and retain relevant knowledge. To operationalize
3

---

Published as a conference paper at ICLR 2025
CE
unbounded transfer
adaptive termination
3
all-to-one update
distribution callapse
sparse update
diversity-knpppng
Figure 2: Comparison of learning schemes: CE GEM (3 0). The arrows illustrate the probability movement directions during the learning process, with Token 3 as the target token.
this insight, we propose the following guiding principle for SFT: learning from the data while preserving diversity." In the following sections, present technical insights and solutions.
4 PROBABILITY TRANSFER THEORY
In this section, we draw insights into algorithm design by examining the dynamics of CE. We will introduce new theory probability transfer. To illustrate this concept, we study simplified setting, where the prompt with 0₂ E RK being the "logit" and K being the vocabulary size. Let y
X' is fixed and given. We model the distribution fo (z) softmax (fx)
[K] denote the token
class to be learned. We begin by calculating the gradient of the CE loss for the given example:
(0) [1] -fo(2))),..... 1 fo(il=), - [o(K[>))
(1)
This indicates that, except for the label class whose logit increases by 1 - (i) all other tokens experience logit decrease proportional to their probabilities.
What makes this behavior particularly interesting? We interpret it through logit flow dynamics perspective, where logits redistributed among token classes during training. Let i-th token class being the "target", while other tokens being the "sources". We have the following observation.
Proposition 1. The gradient of CE specifies a logit flow map: each source token j transfers fo(j)æ) logits the target token Formally,
-VeCce(0) E
Wi+s Cie-j
(2)
j=1
W (6((
Giej 10
th position
j-th position
where Wie; acts as weighting factor, and Biej is vector with i-th element being and the j-th element being -1 and 0 otherwise. Furthermore, the logit flow satisfies conservation property, ensuring that logits redistributed from source tokens equal those received by the target token:
K
Logits from source tokens
Σ folilæ)
j=1,j$i
=1
fo(ilz) = Logits to the target token.
Here, we provide an example for numerical illustration. Consider so 00.4, 0.3, 0.2, 0.1] with label i = 3. The CE gradient decomposes as:
-V [-0.4, 0.3, 0.8, 0.1] = 0.4 x 0.4
0.3 x [0. 0.11 0 10,0,1,
That is, there are 3 flows: tokens 1, act as sources, transferring logits to the target 3. The theorem provides insight into CE's optimization dynamics for a single gradient step. We formalize CE's iterative optimization procedure in Figure 3.
We highlight CE's limitations and present our proposed techniques below.
'Whle preserving diversity may retain suboptimal responses, defer this issue to the RI stage, where the model refine its outputs through tral-and-errorr, leveraging feedhack from reward mechanisms.
4

---

Published as a conference paper at ICLR 2025
While there exists source token j o 2 with for [z) 0, continue the following steps.
Find any j with So, (il=) > 0
Decrease the logit for source token by learning rate 1) and weight Wie-f
(CE)
0x+1 0x [j] 1) Wiej
Increase the logit for the target token i similar manner.
8 [t] Ax [i] 77 Wi+j
Figure 3: CE's optimization procedure in the probability transfer viewpoint.
Limitation of CE: All-to-One Probability Transfer. CE loss implements a logit flow mechanism where probability mass from all non-target tokens is transferred the target token (c.f. the first bullet point in Figure 3). This approach penalizes all non-target tokens regardless of their semantic relevance or contextual appropriateness. For example, in the sentence "I like coffee", while "coffee" is the target, reducing the logit and probability of "tea"-a semantically related and contextually plausible token-may harm generalization. This limitation is especially critical in LLMs, they are extensively pre-trained and encode rich knowledge across many tokens. The all-to-one flow dynamic disrupts these carefully learned token relationships, potentially reducing output diversity and contributing to knowledge forgetting.
Proposed Technique 1: Sparse Update. To address this limitation, we propose a sparse update strategy. Instead of considering all source tokens, we select only pivotal tokens for probability transfer (c.f. the first bullet in Figure 4). Specifically, we introduce a simple approach: we identify the pivotal token the one with the highest model confidence: E argmax / [((rr).. The underlying intuition is clear: the model only corrects its most confident prediction if it is incorrect (i.c., does not match the target label). This targeted correction mechanism protects the probability distribution of minority cases, ensuring they maintain a reasonable chance of being sampled during generation.
However, this technique alone is insufficient. As training progresses and the pivotal token's 's probability diminishes through repeated updates, the flow mechanism naturally shifts to other high- coob-rrbbbbility tokens. This cascading effect eventually approaches the dense update behavior of CE probabilities of all source tokens are sequentially reduced. This observation highlights the need for principled approach terminate probability transfer at appropriate stage.
The CE optimization process lacks a natural termination point for probability transfer (c.f. the first line in Figure 3). The logit flow continues indefinitely until all source token probabilities approach zero, causing the distribution to collapse and concentrate entirely on the target token. This represents an undesirable convergence point that eliminates distribution diversity. Fundamentally, this issue stems from CE's implicit assumption that observed data should be assigned maximum likelihood, without preserving reasonable probability for alternative tokens.
Limitation of CE: Unbounded Probability Transfer.
Proposed Technique Adaptive Termination. To prevent distribution collapse, we introduce an intuitive stopping criterion: halt the probability transfer once the target token i becomes the most probable token the distribution (c.f. the first line in Figure 4). Formally, if E argmax f( (z), we stop further update. This supports the assumption that while the observed data should be adjusted to increase its likelihood, other possibilities should still be considered, so the probability of the observed data should not be forced to 1. Our stopping rule is designed to ensure that, after learning, greedy decoding (i.e., selecting the token with the highest probability) can output the correct label, which achieves the primary training objective. An additional advantage of this rule is that, due to early termination, it keeps the resultant distribution close to its initial state, thereby mitigating forgetting."
The two techniques outlined above form the foundation for the initial prototype design of our proposed algorithm, GEM (Game-theore) theoretic Entropy Maximization); see Figure 2 and Figure 4. The meaning and rationale behind the name will be explained in detail Section 5.
'This serves as form of implicit regularization, encouraging the model to remain close to its pre- me-trained state. Unlike heuristic approaches that merely keep the model close to its initial parameters which inevitably face trade-off between performance and closeness our design identifies clear constraint: ensuring greedy decoding output the information to be leamedd. Within this constraint, we maximize closeness the pre-trained model.
5

---

2019) is well-suited for this Published as a conference paper at ICLR 2025
While the target token & argmax fe, (-(z), continue the following steps.
Calculate the model's best prediction j argmax f(-(æ)
Decrease the logit for source token by learning rate 71 and weight wiej?
(GEM)
1 0
Wiej
Increase the logit for the target token similar manner.
0x+1 [ 8 [ 77 Wiej
Figure 4: GEM's optimization procedure (3 0) in the probability transfer viewpoint.
While this prototype conceptually sound and aligns with the guiding principle of SFT, lacks the flexibility needed to extend its ideas to neural network training scenarios. To address this limitation, have developed general mathematical framework that refines the approach, making it both more elegant and adaptable. We will discuss this framework in detail in the next section.
HAME-THEORETIC FORMULATION
In this section, we present a mathematical framework for fine-tuning approaches that preserve diversity. We begin by highlighting a key observation from Section 4: introducing a meta-level mechanism helps govern the the logit flow, such as determining sparse updates and defining stopping rules. Mathematically, this corresponds the usage of an auxiliary variable. This insight motivates the development of framework that integrates such an auxiliary variable. We find that a game- theoretic framework (Goodfellow et al., 2014; purpose, it naturally introduces an additional player.
Specifically, our proposed mathematical framework is formally presented as follows:
min
[GEM(J, A E, E, Eyr, pp( E x0((2) [log / (y (x) log f [0))]]
(3)
f
max Q(f.9) E,Eygene (1)) log f gene [x] B H(g(-1x)).
In this framework, the term ye represents the supervised label in the dataset, while y corre- sponds to the monel-ggeerated output. For clarity, y=eal aligns with the target token and yseae aligns with the source token in the probability transfer framework. The symbol refers to a distribution similar f. Notably, q serves as the auxiliary variable introduced to regulate the process, achieving the functions of sparse updates and proper termination introduced in Section 4.
simple explanation of our framework that we aim to maximize the log pag--obbabiltty f for real data, as expressed in
data while minimizing the log-pne pog-orbbayylity probability for model-generatee Equation (3). Simultaneously, we optimize the auxiliary variable q to enhance performance by treating log f as objective function. However, this overview omits key insights into the underlying approach. In the following section, we will provide a more detailed explanation of our framework.
5.1 UNDERSTANDING THE GAME
In this section, we explain the game-theoretis framework and connect to the probability transfer theory developed in Section 4. For clarity, adopt the same setting as in Section and calculate the gradient with respect to 0. For a specific sample where yeesl = i, the gradient is given by:
K
(fo. =
Σ
Wi+ es-j
(5)
1=1
Wi+j g(jlz))
The form is similar to the CE flow in Equation (2), except that the weights now design variables governed by q, introducing greater flexibility. We will explore the design of and show that this
framework generalizes both the CE formulation and the GEM prototype discussed in Section 4.

---

Published as a conference paper at ICLR 2025
model distribution
mete--conrolllrrsellctinng source tokens for probability transfer
f
f=
p=
A
A=5
highest-confident token only
emphasizes high- probability tokens
matches CE's assignment
approaches uniform selection
Figure Illustration of the meta-controller q-
Our design of follows an optimization problem in Equation (4):
(x) with j = argmax fi(-1x) if 3 0
argmax Q f) =
(6)
softmax(L/$ log (y(z²)
i
That is, is a shifted distribution derived from f. This transformation is visualized in Figure 5.
Note that y serves a "controller" for prioritizing of source tokens and the strength of their associated flows. When / = 0, (x(x) (i.e., the Dirac distribution) selects the single token with the highest probability
tokens. It determines both the selection
f. which corresponds exactly to the algorithm prototype in Section 4. When 3 = 1, becomes identical to f. reducing to the CE formulation. For intermediate values 3 E (0, 1). q represents a soften distribution where high- probability modes become more prominent, while low-probability regions are further suppressed. This results reduced contribution from minority tokens the probability transfer, thereby protecting their probabilities. Thus, for small values, the framework encourages sparse updates and safeguards tokens that may not appear in the current dataset but were learned during pre-training. When 3 > 1, the framework encourages leveraging minority tokens for probability transfer, diminishing diversity more significantly- an approach not considered in our framework.
A key feature of q its adaptability. Rather than implementing a uniform controller for all prompts, q adapts to each distribution's 's unique characteristics. This adaptability is particularly well-suited for LLM training, where prompts come from diverse tasks and distribution shapes vary considerably across different prompts, similiar to the challenge highlighted in (Song et al., 2023; Li et al., 2024c).
5.2 THEORETICAL GUARANTEE
Building on our earlier illustrative analysis with a single data point, we now formally present a theory with the real data distribution p.
Proposition 2. For data distribution satisfying p(ulæ) > O, with B > O, the game in Equations (3) and (4) posses a unique Nush equilibrium point:
{ f ==oftmaa(( * logp)
q==p
Furthermore, f* corrrepoons to the optimal solution to the distribution matching problem (with 1/3 = (7 + 1)), which minimizes the reverse KL divergence with entropy regularization:
f argmin E. [DkL(f(->)) ((-))))(((l))))))))))))))))))) -{([(())))) f
(8)
This result provides intuitive understanding of the distribution matching problem that frame- work addresses: drives the distribution f close to p while encouraging the diversity through entropy regularization. This is exactly the goal set in Section 3. We note that analysis relies / > as the equilibrium point is unique in this scenario. For 3 = 0 and is the Dirac distribution, we can still apply the stopping criterion discussed Section 4. However, the choice of is neither unique nor analytically tractable, introducing additional complexities that we leave for future work.
5.3 TRAINING ALGORITHM: GEM
In this section, we present a algorithm for training neural networks within our framework. For clarity, in the main text, state the token-level version, where (w](=) denotes the one-step conditional
7

---

Published as a conference paper at ICLR 2025
distribution, while the extension the sequence-level is provided in Appendix B. Specifically, parameterize fo using a Transformer and optimize its parameter directly. The overall algorithm is outlined in Algorithm 1, which incorporates two key features: single-model optimization and variance-reduced gradient estimation. These features ensure that our algorithm is highly scalable, requiring nearly the same GPU memory and computational speed optimizing the standard CE loss.
Single-Model Optimization. Recall that we have a closed-form solution for (see Equation (4)), which significantly simplifies the training procedure. Specifically, the update rules are as follows:
{
f for VOLGEM (So, qu) 10=6,
9k+1 argmax, @(f0r4:: 9) softmax(1/8 log fei-1)
Specifically, be computed simply by shifting the logits of fe, eliminating the need to maintain separate network for q. This reduces the memory burden of training. In contrast, GANs require an additional neural network (i.e., the discriminator), which must be trained alongside the main model.
Varience-reduced Gradient Estimation. Building on the above observation, we slightly adapt the notation to define the loss function. Let & denote the parameters of the distribution f. Given training samples (Zi, yi ma)),, ), the loss is defined as (with slight notation abuse):
(GEM(0) = Σ y/amme
9k (aysenn(zz))
[log fo (y (x,) log fo -0 [x:]]
A notable feature is that it computes the true expectation over 9k- This reduces the variance of the gradient estimator and improves the training stability. This again differs from GANs, where stochastic gradient estimation introduces randomness from both the data distribution and the generated distribution q. We note that these optimization properties arise from the fact that distributions in LLMs have finite support, while GANs typically operate in continuous domains.
Algorithm GEM
Input:
Dataset D = (Tis y )
1: for iteration k = I..... ; do
log for) 3: Loss LGEM (0) E, Eyrom Th (y (xx) [log fo(ysene(xs) log fo(
2:
Set qk softmax(1
4: Update 0k+1 O2 77. VOLGEM l==0. Generative model fox+1
Output:
EXPERIMENTS
In this section, present experiment results validating the effectiveness of our proposed framework, demonstrating that GEM matches the downstream performance of CE while offering two distinct advantages: (1) its diversity preservation enables a wider range of outputs, enhancing test-time scaling performance, and (2) it mitigates forgetting, effectively reducing the alignment tax. Extended results available in Appendix F.
Set-up. We fine-tune the pre-trained Llama-3. 1-B model with the Ul t rittraFeecbbae dataset (Cui al., 2024). This dataset contains prompts from instruction datasets like Evol-Instruct and UltraChat, and responses generated by models such as GPT-4 and Llama-2-7B/11 BB-7B///B/7BB-Chr Following (Yu al., 2023: Liu et al., 2023; Cui et al., 2024), cosine learning rate decay schedule, and use macro batch size of 128. The maximum sequence length, encompassing both the prompt and response, is set 2,048 tokens. Models are trained for three epochs. Detailed experimental settings are described in Appendix F.
the learning rate
10-5 employing
We implement GEM with 3 0.7. Our primary baseline the standard CE loss. Additionally, we explore variant of CE incorporating weight decay of 0.1, which has been commonly used in previous studies (Ouyang et al., 2022; Bai et al., 2022). We refer to this approach as CE + WD. The NEFT method (Jain et al., 2023), which perturbs the input embedding with random noise in fine-tuning to mitigate overfitting, has also been implemented.
6.1 IMPROVING DIVERSITY AND TEST-TIME SCALING
8

---

Published as a conference paper at ICLR 2025
=2x efficiency
=2x efficiency
CL+WD NEIT
48
NEFT GEM
80
GEM
70-
Winn Rato
Pass Rake
50-
40
24
01
04
952
mino
Sampling Budqut
Samsting Budget
(a) Chat
(b) Code generation
Figure 7: Performance test-time scaling. The results demonstrate that GEM achieves better perfor- mance with the same sampling budget and is more efficient in reaching comparable performance.
In this section, we demonstrate that through implicit entropy maximization, GEM enhances output diversity. This is evi- dent from the entropy values of the generative distribution: 0.42 (CE), 0.41 (CE+ + WD), 0.43 (NEFT), and 0.76 (GEM). Consequently, GEM is more likely to sample diverse solu- tions and has a higher chance of identifying better solutions during inference through repeated sampling. This aligns with the recent trend of test-time scaling (Snell et al., 2024; Brown et al., 2024; Wu et al., 2024). We illustrate this benefit us- ing two tasks: chat and code generation, while mathematical reasoning tasks are explored in Appendix F.
5L
cutput Diversty Score (120- Sen-eur
Figure Enhancing output diversity boosts the win rate when using BoN.
Chat. We assess the model's chat ability using the best- of-N sampling strategy. We prompt the trained models answer 805 questions from the A AlpacaEval dataset (Li et al., 2023). For each question, the model generates 32 responses, from which a reward model selects the best response. We employ the reward model F -aaarxXLLLMM3 RM-v 15, which has top performance on Rewa Bench (Lambert et al., 2024). The selected response compared with GPT4's response with the win rate reported in Figure 7 (a).
The evaluation reveals that weight decay does not lead to performance improvements. On the other hand, NEFT shows strong performance, partly attributed to its longer responses, as highlighted in (Jain et al., 2023). However, GEM outperforms CE, with a 3. 1-pont (6.6% relative) increase win rate. Additionally, we note that GEM achieves comparable performance to CE while requiring about 2x fewer sampling costs. To further understand this improvement, we evaluate output diversity using the Self-BLUE metric? across generated responses, as depicted in Figure 6. GEM not only increases output diversity but also enables the generation of higher-quality responses.
Code Generation. We consider the HumanEval (Chen et al., 2021) benchmark, in which the model is asked to generate Python codes for 163 questions, and the executor judges their correctness. The evaluation metric is the pass rate among generated responses. The results are reported in Figure 7 (b).
In this task, we observe that for pass@ 1 100 weight decay slightly improves the performance of CE, increasing from 75.5 to 76.6, while NEFT shows no significant improvement, maintaining performance of 75.6. In contrast, GEM achieves a 7-point improvement, reaching a performance of 82.5, which represents a 9.3% relative increase. Furthermore, GEM proves be highly efficient in test-time scaling, requiring only about half the sampling budget to achieve comparable performance.
We note that in our experiments, the performance using BON or pass rate serves as an estimate of self- improvement. Specifically, high-quality self-generee generated samples produced by GEM can be distilled back into the model, enhancing its zero-shot performance (see (Sessa et al., 2024)). Therefore, we believe GEM can exhibit more effective scaling in post-training, a topic we plan to explore in future
'https: / /twgg/ugggggg ing acee co/ s sfai wfairr///eeeeeee saac//saaiiie rx- LLaMAA3 RM- "Self-BLUE measures similarity among responses; so, we -
Self-BL.UP as metrie for diversity.
9

---

Published as a conference paper at ICLR 2025
Pre-trained
CE
CE+WD
NEFT
GEM
ARC
GSM8K
HellaSwag
MMLU
TruthfulQA
Winogrande
Figure 8: Performance tasks from the OpenILM leaderboarde The results indicate that GEM outperforms CE, demonstrating a lower alignment tax.
work. We also that in the above evaluation, GEM enhances output diversity without sacrificing its direct generation performance: in chat, CE (24.5) vs. GEM (24.9), and in code generation, CE (32.7) vs. GEM (32.6). This contrasts with other techniques, which often increase diversity the cost of performance degradation; see Appendix D for further discussion.
6.2 MITIGATING FORGETTING AND REDUCING ALIGNMENT TAX
In this section, we show that GEM also helps mitigate forgetting and reduce alignment tax. We evaluate the mod- els across six tasks: ARC, GSM8K, HellaSwag, MMLU, TruthfulQA, and WinoGrande, listed the OpenLLM leade rboard. Results are presented in Figure 8.
62.4
We observe that after SFT, performance declines most tasks for the baseline models, with CE showing the most sig- nificant drops. However, GEM exhibits different behavior: for tasks such as ARC and HellaSwag performance does decrease. On average across six tasks, GEM shows an 80% reduction in alignment tax, with 0.3-point drop for GEM compared to a 1.5-point drop for CE. This behavior aligns with our method, which introduces y encourage sparse updates of tokens and prevent forgetting. Further validation can be seen by examining the distance from the pre-trained model. Specifically, measure the parameter difference from the initialization using the fz-norm, as shown in Figure 9. The results indicate that GEM is closer to the mr-ttrained model in parameter space, leading an interesting conclusion that diversity-preservi updates help reduce the model's deviations during learning.
61.6
80.8
Figure 9:
Staying close to the pre- trained model helps mitigate forgetting.
7 CONCLUSION
In this paper, we identify key limitation of SFT via optimizing the CE loss: its tendency to encourage over-memorization of training data and the resultant reduction in output diversity, which could hinder subsequent development that relies on sampling to explore better responses. To address these limitations, we develop a new thee-thhereetc formulation. Central to this framework is meta-controller that encourages sparse updates and ensures proper termination of the learning process. We validate the advantages of the proposed framework in test-time scaling and its effectiveness in mitigating catastrophic forgetting and reducing alignment tax.
We focus on SFT in this paper, but our techniques may also be of interest for developing algorithms to preserve diversity and mitigate forgetting. In particular, the enhanced diversity achieved by our approach can be beneficial in several contexts: it enables exploration to improve performance limits with reinforcement learning (Li, 2025), facilitates self-distillation with best-of-n sampling (Sessa al., 2024), and enhances the diversity of synthetic data generation (Adler et al., 2024). Please Appendix G for the discussion. We see the potential of our method in these areas and plan to explore these topics in future work.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
Ziniu Li would like to thank Yao Fu for the discussion about diversity during the ICML 2024 poster session, well as Yushun Zhang for reading the manuscript and providing valuable feedback. The work Ruoyu Sun was supported by NSFC (No. 12326608); Hetao Shenzhen-Howe Kong WWWWSSKCCYB-2222
Science and Technology Innovation Cooperation Zone Project (No.
University Development Fund EDFF1001199) the Chinese University of Hong Kong. Shenzhen:
Guangdong Provincial Key Laboratory of Mathematical Foundations for Artificial Intelligence (2033B121201000) The work of Z.-Q. Luo was supported by the Guangdong Major Project of Basic and Applied Basic Research NN.2022300333000001 the Guangdong Provincial Key Laboratory of Big Data Computing, and the National Key Research and Development Project under grant 2022YFA 0039000. The work of Tian Xu was supported by the Fundamental Research Program for Young Scholars (PhD Candidates) of the National Science Foundation of China (623B2049).
ETHICS STATEMENT
Our work focuses on designing better algorithms for fine-tuning large language models, aiming to enhance their effectiveness and broaden their applications. In particular, the entropy regularizer introduce for distribution matching encourages more diverse outputs from language models. We do not foresee any direct negative impacts from this approach.
REPRODUCIBILITY STATEMENT
The proof of Proposition 2 is provided in Appendix C. Experiment details for reproducing our numerical results can found in Appendix E. Our code accessible at https: /ttp//ggtttt/ttt/ com/ Liziniu/CEM.
REFERENCES
Bo Adler, Niket Agarwal. Ashwath Aithal, Dong H Anh. Pallab Bhattacharya, Annika Brundyn, Jared Casper, Bryan Catanzaro, Sharon Clay, Jonathan Cohen, al. Nemotron-4 340b technical report. arXiv preprint arXiv:2406. 1704, 2024.
Brenna D Argall, Sonia Chernova, Manuela Veloso, and Brett Browning, A survey of robot learning from demonstration. Robotics and autonomous systems, 4765):469-483, 2009.
Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bilal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello. A general theoretical paradigm to understand learning from human preferences. In International Conference Artificial Intelligence and Statistics, pp. 4447-4455. PMLR, 2024.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.08822 2022.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Advances in Neural Information Processing Systems, 34:27381-27394, 2021.
Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoc V Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint arXiv:2407.27788 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems 33, pp. 1877-1901, 2020.
Collin Burns, Pavel Izmailov, Jan Hendrik Kirchner, Bowen Baker, Leo Gao, Leopold Aschenbrenner, Yining Chen, Adrien Ecoffet. Manas Joglekar. Jan Leike, et al. Weak-to-strong generalization: Eliciting strong capabilities with weak supervision. arXiv preprint arx 2023.
11

---

Published as a conference paper at ICLR 2025
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv preprint arXiv:2107.03374. 2021.
Zixiang Chen, Yihe Deng, Huizhuo Yuan, Kaixuan Ji, and Quanquan Gu. Self-play fin-ttuning converts weak language models to strong language models. arXiv preprint arXiv:2401.01333 2024.
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetunee language models. Journal of Machine Learning Research, 25(70):1-53, 2024.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reichiro Nakano, et al. Training verifiers to solve 20. math word problems. arXiv preprint 10.148210..1468,, 2021.
Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Bingxiang He, Wei Zhu, Yuan Ni, Guotong Xie, Ruobing Xie, Yankai Lin, et al. Ultrafeedback: Boosting language models with scaled ai feedback. In Forty-first International Conference on Machine Learming, 2024.
Abhimanyu Dubey, Otkrist Gupta, Ramesh Raskar, and Nikhil Naik. Maximum-entropy fine grained classification. Advances in neural information processing systems, 31, 2018.
Abhimanyu Dubey, Abhinay Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21788. 2024.
Team Gemma, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.0829. 2024.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farles Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Yanzhu Guo, Guokan Shang, Michalis Vazirgiannis, and Chloé Clavel. The curious decline of linguistic diversity: Training language models synthetic text. arXiv preprint arX1v:2311.09807, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 770-778, 2016.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in Neural Information Processing Systems 29, pp. 4565 4573, 2016.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In Proceedings of the 10th International Conference on Learning Representations, 2022.
Edward J Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, and Nikolay Malkin. Amortizing intractable inference in large language models. arXiv preprint arXiv:2310.0433 04.664433,, 2023.
Neel Jain, Ping-ych Chiang, Yuxin Wen, John Kirchenbauer, Hong-Min Chu, Gowthami Somepalli, Brian R Bartoldson, Bhavya Kailkhura, Avi Schwarzschidd, Aniruddha Saha, et al. Neftune: Noisy embeddings improve instruction finetuning. arXiv preprint mrXiiv23100059// 2023.
Alexia
The relativistic discriminator: key element missing from standard GAN. In Proceedings of the 7th International Conference on Learning Representations, 2019.
Liyiming Ke, Matt Barnes, Wen Sun, Gilwoo Lee, Sanjiban Choudhury, and Siddhartha S. Srinivasa. Imitation learning as f-divergence minimization. arXiv, 1905. 18888 2019.
12

---

Published as a conference paper at ICLR 2025
Jiyeon Kim, Hyunji Lee, Hyowon Cho, Joel Jang, Hyeonbin Hwang, Seungpil Won, Youbin Ahn, Dohaeng Lee, and Minjoon Seo. Knowledge entropy decay during language model pretraining hinders new knowledge acquisition. arXiv preprint arX1v:2410.01380, 2024.
Robert Kirk, Ishita Mediratta, Christoforos Nalmpantis, Jelena Luketina, Eric Hambro, Edward Grefenstette, and Roberta Raileanu. Understanding the effects of rihf on llm generalisation and diversity. arXiv preprint arXiv:2310.0645: 2023.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. Advances in neural information processing systems, 25, 2012.
Nathan Lambert, Valentina Pyatkin, Jacob Morrison, LJ Miranda, Bill Yuchen Lin, Khyathi Chandu, Nouha Dziri, Sachin Kumar, Tom Zick, Yejin Choi, al. Rewardbench: Evaluating reward models for language modeling, arXiv preprint 2rXiv:2203.1333 2024.
Jiaxiang Li, Siliang Zeng, Hoi-To Wai, Chenliang Li, Alfredo Garcia, and Mingyi Hong. Getting more juice out of the sft data: Reward learning from human demonstration improves sft for Ilm alignment. arXiv preprint 241::2405.11888 2024a.
Xuechen Li, Tianyi Zhang. Yann Dubois, Rohan Taori, Ishaan Gulrajani, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Alpacaeval: An automatic evaluator of instruction-follovinn models. https: /tt//ggtttutet com comm/aasu--lle com./tat.uu-aabaappaaaaa eval, 2023.
Yi-Chen Li, Fuxiang Zhang, Wenjie Qiu, Lei Yuan, Chengxing Jia, Zongzhang Zhang, Yang Yu, and Bo An. Q-adapter: Customizing pre-trained llms to new preferences with forgetting mitigation. arXiv preprint arXivv2407.03886, 2024b.
Ziniu for Ilms?
Li.
Can
better
cold-start
strategies
improve
training
https:/ bttps://tangible http///anngnggggeeeeeeeeeeeeeeee poolee-pooo----ooo 203 notiononitee
Benn-Beter--aee Col--ooll---e Start- Startttirrtgiess Improve RL- Training
177a0742251
2025. Notion Blog.
Ziniu Li, Tian Xu, Yushun Zhang, Zhihang Lin, Yang Yu, Ruoyu Sun, and Zhi-Quan Luo. Remax: A simple, effective, and efficient reinforcement learning method for aligning large language models. In Forty-first International Conference Machine Learning. 2024c.
Wei Liu, Weihao Zeng, Keqing He, Yong Jiang, and Junxian He. What makes good data for alignment? comprehensive study of automatic data selection in instruction tuning. arXiv preprint arXiv:2312.15685 2023.
Yun Luo, Zhen Yang, Fandong Meng, Yafu Li, Zhou, and Yue Zhang. An empirical study of catastrophic forgetting in large language models during continual fine-tuning. arXiv preprint (rXiv:2308.08747 2023.
Laura O' o'Mahony, Leo Grinsztajn, Hailey Schoelkopf, and Stella Biderman. Autributing mode collapse in the fine-tuning of large language models. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024.
OpenAl. Gpt-4 technical report. arXiv preprint 2303::3333087722 2023.
Weronika Ormaniec, Felix Dangel, and Sidak Pal Singh. What does it mean to be a transformer? insights from theoretical hessian analysis. arXiv preprint arXiv:2410. 10986 2024.
Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J Andrew Bagnell, Pieter Abbeel, Jan Peters, et al. An algorithmic perspective on imitation learning. Foundations and Trends@ in Robotics, 7(1-2): 1-179, 2018
Long Ouyang, Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems 35, pp. 27730 2022.
13

---

Published as a conference paper at ICLR 2025
Gabriel Pereyra, George Tucker, Jan Chorowski, Lukasz Kaiser, and Geoffrey Hinton. Regularizing neural networks by penalizing confident output distributions. arXiv preprint arXiv:1701.0663 2017.
Dean Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural Computation, 3(1):88-97, 1991.
Team Qwen. Qwen2.5: party of foundation models, September 2024. URL htwps:ttp/////cct //vee// github. httpub..o/bllgggweeeeeeee .ttp///ee/2222//m
Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D Manning, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. arXiv preprint arXiv:2305. 18290, 2023.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1 1-7 2020.
Stéphane Ross, Geoffrey J. Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the 14th International Conference on Artificial Intelligence and Statistics, pp. 627-635, 2011.
Pier Giuseppe Sessa, Robert Dadashi, Léonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Ramé, Bobak Shariari. Sarah Perrin, Abe Friesen, Geoffrey Cideron, et al. Bond: Aligning Ilms with best-of-n distillation. arXiv preprint arXiv:2407. 14622, 2024.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.0332 2024.
Ilia Shumailov, Zakhar Shumaylov, Yiren Zhao, Yarin Gal, Nicolas Papernot, and Ross Ander- son. The curse of recursion: Training generated data makes models forget. arXiv preprint arXiv:2305.17493, 2023.
Charlic Snell, Jachoon Lee, Kelvin Xu, and Aviral Kumar. Scaling I1m test-time compute optimally can be more effective than scaling model parameters. arXiv preprint arXiv:2408.033 03314, 2024.
Ziang Song, Tianle Cai, Jason D Lee, and Weijie J Su. Reward collapse in aligning large language models. arXiv preprint 27X.,, 2023.
Hao Sun and Mihaela van der Schaar. inverse-rligmment Inverse reinforcement learning from demonstrations for Ilm alignment. arXiv preprini arXiv:2405.15624, 2024.
Ruoyu Sun, Tiantian Fang, and Alexander Schwing. Towards a better global loss landscape of gans. Advances in Neural Information Processing Systems, 33:10186-10199 2020.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307. 092888. 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30, pp. 5998 6008, 2017.
Nino Vieillard, Tadashi Kozuno, Bruno Scherrer, Olivier Pietquin, Rémi Munos, and Matthieu Geist. Leverage the average: an analysis of kl regularization in reinforcement learning. In Advances in Neural Information Processing Systems 33, pp. 12163-12174, 2020.
Evan Wang, Federico Cassano, Catherine Wu, Yunfeng Bai, Will Song, Vaskar Nath, Ziwen Han, Sean Hendryx, Summer Yue, and Hugh Zhang. Planning in natural language improves Ilm search for code generation. arXiv preprint arXiv:2409.03333 2024a.
Mingzhi Wang, Chengdong Ma, Qizhi Chen, Linjian Meng, Yang Han, Jiancong Xiao, Zhaowei Zhang, Jing Huo, Weijie Su, and Yaodong Yang. Magnetic preference optimization: Achieving last-iterate convergence for language model alignment. arXiv preprint arXiv:2410. 0224400667111 2024b.
14

---

Published as a conference paper at ICLR 2025
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoe V. Le, Ed H. Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. In Proceedings of the 1 11st International Conference on Learning Representations, 2023.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoe V Le. Finetuned language models are zero-shot learners. arXiv preprint arXiv:2109.01652 2021.
Yuxiang Wei, Zhe Wang, Jiawei Liu, Yifeng Ding, and Lingming Zhang. Magicoder: Empowering code generation with iss-instruct. In Forty-first International Conference on Machine Learning. 2024.
Yangzhen Wu, Zhiqing Sun, Shanda Li, Sean Welleck, and Yiming Yang. An empirical analy- sis of compute-optimal inference for problem-solving with language models. arXiv preprint arXiv:2408.00724 2024.
Jiancong Xiao, Ziniu Li, Xingyu Xie, Emily Getzen, Cong Fang. Qi Long, and Weijie Su. On the algorithmic bias of aligning large language models with rlhf: Preference collapse and matching regularization. arXiv preprint arXiv:2405. 16455, 2024.
Lechao Xiao. Rethinking conventional wisdom in machine learning: From generalization to scaling. arXiv preprint arXiv:2409.15133 2024.
Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments. In Advances in Neural Information Processing Systems 33, pp. 15737-15749, 2020.
Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. arXiv preprint arXiv:2309. 12284, 2023.
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi- Qhi-Quan Luo. Why trans- formers need adam: A hessian perspective. arXiv preprint arXiv:2402. 16788, 2024.
Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pp. 1433-1438. Chicago, IL, USA, 2008
Barret Zoph and John Schulman. Chatgpt and the art of post-training. Google URL https://docs. htw.gglle.co/ppeesetttiioooooo.. LIG.rU/nwwyearreeeeeeeooooo/oeen
Docs, 2025.
A RELATED WORK
Supervised Fine-Tuning, The phenomenon of alignment tax in fine-tuning LLMs was first docu- mented by Ouyang et al. (2022) and Bai et al. (2022), who empirically demonstrated that the in-context learning capabilities of LLMs degrade after Reinforvement Learning from Human Feedback (RLHF). Subsequently, Kirk al. (2023) revealed that RLHF fine-tuning also leads a reduction in output diversity. These findings highlight the issue of overfitting during the fine-tuning process. Importantly, we emphasize that these limitations are not exclusive RLHF but also manifest in SFT alone, as evidenced by (Ouyang et al., 2022, Figure 29) and O' MaMonyy et al. (2024). Given that SFT is the initial step in the post-training pipeline, addressing these limitations at their source critical to improving the overall robustness and diversity of fine-tuned models.
Our argument--tal CE is not the best fit for aligns with the findings of a recent study (Xiao, 2024). This research highlights that pre-training with CE operates within "scaling-centric ceetrric""eri" regime, where the generalization loss aligns closely with the training loss, rendering the generalization gap negligible and making CE viable choice. However, in the context of SFT, CE loss falls into the "generalization-centii regime, where careful consideration of regularization becomes crucial.
Game- Game-Theoretie Theoretic Formulation in LLMs. Closely related 2024; Li et al., 2024a) explored improving CE-trained models using techniques like self-play from game theory. However, our approach differs in two key ways. First, introduce the maximum
work. recent studies (Chen et al.,
15

---

Published as a conference paper at ICLR 2025
entropy principle into distribution matching, while their work examines the standard distribution matching framework. Second, we focus on addressing the limitations of CE loss by designing methods that directly applies pre-trained models, whereas their methods are applied post-SFT.
Imitation Learning and Inverse Reinforcement Learning. Our work is closely connected to Imitation Learning (IL) (Argall et al., 2009; Osa et al., 2018), framework in which a learner acquires maninng strategies by mimicking expert demonstrations. In fact, SFT can be interpreted as a form of IL with deterministic transitions (Sun & van der Schaar, 2024; Li et al., 2024a). Specifically, the cross-entropy loss used in SFT corresponds to behavior cloning (BC) in IL (Pomerleau, 1991), method that directly replicates expert actions. However, our framework is more closely aligned with Generative Adversarial Imitation Learning (GAIL) (Ho & Ermon, 2016), which has been shown to generally outperform behavior cloning (Ke et al., 2019; Xu et al., 2020) by leveraging additional transition information to better match the expert policy distribution.
Our work shares conceptual similarities with maximum entropy Inverse Reinforcement Learning (IRL) (Ziebart et al., 2008), which introduces an entropy term to enhance the robustness of reward function estimation. However, key distinction lies in objectives: while maximum entropy IRL aims to recover reward function (under which the optimal policy is typically deterministic), our goal is to learn a generative distribution that captures diversity, rather than focusing solely on reward inference.
Entropy Regularization. Dubey et al. (2018) proposed that achieving zero CE loss is not essential for high accuracy. Instead, they suggested that a conditional probability distribution where the argmax corresponds to the correct class is sufficient. This concept motivates our use of entropy regularization, which allows for assigning probabilities to alternative options beyond the observed data. Prior to our work, Pereyra al. (2017) also explored entropy regularization in the context of neural network training. It is important to note that Pereyra et al. (2017) focused on image classification tasks, while our focus is on text generation where data is sequential context of LLMs, Hu et al. (2023) explored the maximum entropy regularization by using GFlowNet (Bengio et al., 2021), but their methods require reward function rather than supervised data.
nature and more challenging, In the
A.1 DIFFERENCE WITH SPIN
Closely related to our work, Chen et al. (2024) also presented a game formulation for fine- tun-ttunin LLMs and developed an algorithm called SPIN. However, their work differs fundamentally from ours. We highlight the key differences to help readers better understand the distinctions between our work and theirs:
Motivation: SPIN aims to further leverage supervised data after optimizing the model with the CE loss, as Chen et al. (2024) believe SFT does not fully exploit the supervised data. In contrast, focus the limitations of SFT, particularly poor diversity and knowledge forgetting, and propose preserving output diversity to address these issues.
Chen et al. (2024) develop a theory proving that the equilibrium point con- verges to the distribution matching problem of the data distribution, which is reasonable at the population level. However, in practice, data distributions usually have limited size and coverage. To address this, we propose incorporating entropy regularization into the distribution matching framework.
Formulation:
Technical approach: SPIN is inspired by techniques from DPO (Rafailov et al., 2023). Specifically, it introduces KL regularization in the algorithm design (rather than in the distribution matching formulation). As a result, SPIN is conceptually similar the online version of DPO, where supervised data is treated as the chosen data and model-generated data as the rejected data. In contrast, entropy regularization is introduced directly in our game formulation, rather than at the algorithm design level. Furthermore, our approach is built on the proposed probability transfer theory, and our algorithm leverages an auxiliary variable that governs the distribution matching process to preserve output diversity. In essence, the two algorithms exhibit entirely different optimization properties and dynamics.
Computational Efficiency: Our approach runs at the same speed and memory requirements CE, while SPIN requires more computational resources due to its need for online sampling and the introduction of a reference model in its KL-regularized optimization. As a result, SPIN runs slower than our method.
16

---

Published as a conference paper at ICLR 2025
B IMPLEMENTATION OF GEM
Generalized GEM. We note that we can generalize the formulation Equations (3) and (4) by
introducing transformation function h on f's objective:
min (GEM(F,9) A By (1-1 By ol lo)) (log (ysene())) log [x)
(9)
max 2V.4) (p) Doog [yy [x]]] + B 1.B.B.(400 ((((-12))
(10)
In this way, we have
-VoCGEM = E Wiej Bimji
(11)
j=1,izi
1 Q(jli) (log f(ysene(2) - log f (y)=))(()))
The choice of the function directly affects the magnitude of wij- If h(z) = 2 (ie., a linear function presented in the main text), we simply have ('(z) 1, resulting in uniform weighting for probability adjustments. This corresponds to the algorithmic design in Section 4. Inspired by (olicoeur-Martince 2019; Sun et al., 2020: Rafailov al.. 2023). we also explore the design f(z)) = log s iggeid((2)), then h'(z) = si gigoid(2).. In this case, Wie/ becomes adaptive: takes higher value when log fy xx) > log fy (eallala (x) and lower value otherwise. This adaptability ensures that the model focuses more on cases where it struggles to distinguish between real and generated data.
In our experiments, we actually implement the log s gmo d function for h, as it provides more adaptivity than the linear function. We note that the linear function also works well in practice, but the log s gmaid function yields slightly better performance in some tasks.
Extension Sequential Data. In the main text, we have derived the algorithm for the case y is non-sequential. We note that optimization in the sequential case could be highly difficult. With a little abuse of notations, let y (y),..., VT) y1T. Now, we can extend the formulation in Equation (3) to the following:
min f
Eg Wir wpt B (m)) [h (log log f(y (x) f(yis*1
(12)
max
E, Bygr (-) [log (viirr [=)] + 3 H+((+++))
Here, we encounter a challenge: the joint distribution of V1.7, as cascaded categorical distribution, is quite complicated. This results in the expectation Eyg cannot be easily calculated before.
To deal with the above challenges, we propose decomposing the multi-step sequential optimization problem into multiple single-step optimization problems and solve each efficiently. This is inspired by the data distribution "reset" trick introduced by (Ross et al., 201 1) in imitation learning, where the teacher first demonstrates few actions, and the student completes the reset. For our problem, we restrict the distribution matching to the case that the prefix samples up to time step / are drawn from the data distribution p and solves the optimization problem at the f-th time step as before. Its mathematical formulation is given below:
{
max L Ce0 (J.9)
-E
Σ B p2) Eyg pp(-) Bygen ~ 12 1=1
[
(13)
where A [h (log f(y x vital) log (ggfeal vfae))))]]]]
The main advantage of this formulation is that for each sub-problent still have access to the conditional distribution, allowing the previously discussed computational advantages to remain applicable. The same idea applies the training of distribution q, still have the closed-form solution that q(-1z, vite)) softmax(1/3 log f-|2,, yre=)))).
We outline the proposed procedure for dealing with sequential data in Algorithm and provide its Py PyTorch implementation below. We examine the importance of the "reset" technique in Appendix F.
17

---

Published as a conference paper at ICLR 2025
Algorithm 2 GEM for Sequential Data
Input: Dataset D ((1s, V1..... ,
Initialize D 0
2: for sample index i do for timestep index
b "Reset" data distribution
3:
I,..., do
3 rj + (yi real .... v
yi yt
D Du ((E
4: Call Algorithm 1 on D Output: Generative model fo
def gem_loss (logits, labels, beta=0.7, 7, ignore_index--100, h="linear"):
shift_logits logits :-1, :1 contiguous
shift_labels labels 1:] cont.ogouuuuus ()
mask - shift labells != Lndex
shift logits shift_logits [mask]
shitt_labell shift_labet [mask]
10
with torch (orcchn)__gra((==
torch.gather.
12
shift_logits, dimm-1, index=shift_labbel mndxesshfftiaabllss.seseeeee unsqueeze (-1)
13
squueee
14
Dogits_diff shift_logite logits_on_last lamellablss..m.aaaaaaa unsqueeze (-1) if == "lincar";
17
weights torch.ones. like (logits_diff)
18
clif h = "log_sigmoie"
19
weights = F. s.sigmoid (0.00 01 logits_diff)
else:
raise ValueError (h)
22
25
gene_Log Log_ = T_log__oottenn (shift woggtttttttttt, dim=-1)
probs = torch.
#.1og_aoftmax (shi ii fogiitits / beta, dim=-1)
26
) ..detach()
27
28
real_ real_log_probs Norch.gather
gene_log_proke dim=-l, index-shift_labes insqueezee (-1)
squeeze
31
32
loss = -torch. sum
33
probs + weights (real Legl__log_rroe unsqueeze (-1)
gene gene__eg prg_pro dim=-i
) .eean
35
36
return
loss
Listing Pytorch Code of GEM
To understand the above implementation, we note that we leverage the gradient analysis in Section 5.3: first, we calculate the re-weighting term Lines 9-20. Then, we calculate the difference log- probabilities in Lines 22-33. Note that log-siggnnn sigmoid function. This ensures that the function behaves nearly linearly.
coefficient of 0.01 to scale the input in the
We observe that the computational complexity of GEM is nearly equivalent to that optimizing CE loss. To clarify, the computational cost of CE involves two primary steps: forward pass through the Transformer to compute the distribution fe. followed by calculating the likelihood log fo fo(2)2) The forward pass, which entails multiple layers of matrix multiplications, the primary computational bottleneck.
Computational Complexity Analysis:
18

---

Published as a conference paper at ICLR 2025
Similarly, GEM requires one forward pass through the Transformer to compute the distributions and f. and then calculate the relative difference log fo (z) log fo f)()))=)))) [z). As with CE. the forward pass is the main bottleneck in GEM. Backpropagation is performed in comparable manner for both methods, resulting in GEM achieving nearly the same training speed as CE.
In terms of memory consumption, GEM requires storing additional distribution 9, which occupies the same amount of memory as f. For example, in our setup with batch size of 4, a sequence length of 2048, and vocabulary size of 128k, g requires only about GB of memory. This is negligible compared with the memory consumed by other training components such as gradients, optimizer states, and activation caches, which can collectively exceed 100 GB.
C PROOF
Proposition 3. For the entropy-regulare KL minimization problem Equation (8), in the function space, we have the optimal solution:
(U(xx) = Zz p(yla)44(++))
where Zy is normalization constant Ey ply (z/))))/((+
The proof is based on the optimality condition of constrained optimization. Its proof can be found in the previous literature (see, e.g., (Vieillard et al., 2020, Appendix A)). We note that the above closed-form solution cannot be applied in practice because do not have access to the density function of the data distribution p.
Proof of Proposition 2. Proof of the first part. When h is linear function, we have that (f)
= By Eygro -p(-12))) Bygne (1(1))))))))))))))))))))))))))) [log (v real (2)) - log (y)ono [x]]
(1( E, Bygun pp(lp)) (y) [log (y real |at]] - Byy upl- (n) B [log I (y [x)]
By (o( [log [ (a))))))) ET Boge (((( [log I (1 (ysenelar)
For any x E X, we have that
OC p P-9
(14)
0f
f
When = p. we have Oc/0f = 0, i.c., a stationary point. By examining Equation (7), we have that f* softmax(3 l) = wof.taa((((/ * log )"), q is also stationary point. This proves that the point in Equation (7) is a stationary point for the Nash game.
To prove that it is the unique stationary point, exists amother equilibrium point addition to the one in Equation (7). This implies that either g* * or 1* * cooftaa x (1/3 log /). In the former case, the gradient in Equation (14) cannot be zero, which leads to a contradiction. the latter case, q² cannot be the optimal solution to Equation (6), again resulting
contraction-based argument. Assume there
contradiction. Thus, the original claim regarding uniqueness holds.
Proof of the second part. It is based on Proposition 3. As analyzed in Proposition 3, for 3 = 1/(7 + 1), the optimal f in Equation (7) corresponds the the optimal solution of minimizing reverse KL with entropy regularization.
D DISCUSSION
D.1 TEMPERATURE ADIUSTMENT FOR OUTPUT DIVERSITY
Temperature adjustment is widely used trick for encouraging output diversity. However, this trick applied during the inference stage, which differs from our focus on preserving output diversity during training. Moreover, the mechanisms for increasing diversity fundamentally different.
19

---

Published as a conference paper at ICLR 2025
When the temperature of the softmax distribution is increased, the probability mass shifts away from the mode (the most likely outcomes) toward the tail (less likely outcomes). While this can enhance diversity, it also carries significant drawback: in LLMs, the tail of the distribution often contains many nonsensical or low-quality tokens. As result, increasing the tail probability can lead undesirable incoherent outputs.
In contrast, GEM employs a different mechanism to protect diversity. By promoting sparse updates and preserving tail probabilities, GEM ensures that diversity is maintained without disproportionately amplifying nonsensical or low-quality tokens. This approach provides a more controlled and effective way balance diversity and output quality during training.
D.2
THEORY OF CE
Theory of CE
CE Loss (Empirical)
Loss (Population)
min
y 1og
max E yop log
input-labet
P: prompt distributton
July [u)): conditional prediction O: parameters of neural network
pl [x]] the conditional data distribution to learn
Eq Equivalence
CE can be used to learn a distribution
If the data samples are "abundant" "
Forward KL Divergence
X
min Egepp KL(p( [x)S(( (a)) constant
Classification
Pre-training [huge data)
SFT
(one label sample is enough)
(data limited)
Distribution Matching
Figure 10: Theory CE. The figure illustrates the connection between the empirical (top left) and population (top right) versions of CE, well as its relationship to distribution matching (bottom right). also highlights the applicable regime of CE (bottom left).
CE is fundamental loss function in machine learning that serves method for learning probability distributions. As illustrated in Figure 10, CE has both empirical and population formulations, where the empirical version minimizes the negative log-likelihood of predictions given inputs across training data, while the population version maximizes the expected log-like likkliihoood across the true data distribution. In particular, its population formulation is equivalent to minimizing the KL divergence between the true conditional distribution and the model's predicted distribution (plus constant), making CE essentially distribution matching technique. CE's effectiveness depends critically on the quantity of data available to approximate the underlying distribution.
As discussed our introduction, CE performs well for classification tasks where even a single labeled example per class can provide sufficient signal, and for pre-training scenarios with abundant data. However, it typically underperforms for SFT LLMs where data is limited and generation tasks are inherently open-ended.
D.3 CE WITH WEIGHT DECAY
Weight decay is widely used technique for mitigating overfitting, particularly effective in training convolutional neural networks (CNNs). However, we observed that it is not always effective when applied to training generative models with Transformers. We hypothesize two primary reasons for this discrepancy:
Architecture.
The structural characteristics of CNNs make them more homogeneous in design, resulting in relatively uniform loss landscape across different parameter blocks within the network.
20

---

Published as a conference paper at ICLR 2025
In contrast, Transformers exhibit heterogeneous properties, as noted in recent works (Zhang et al., 2024; Ormaniec et al., 2024), leading to significantly varied loss landscapes across parameter blocks. As result, a uniform weight decay applied to all parameters in Transformers may be suboptimal, since different parameter blocks might require specialized weight decay strategies.
Task. CNNs in classification tasks are designed to learn predictor that outputs unique prediction, whereas generative models aim to learn a distribution. For generative models, directly regularizing the distribution is preferable. Weight decay, which regularizes parameters indirectly, may not effectively serve this purpose.
Based on these observations, we argue that entropy regularization is better suited for training generative models with Transformers. Unlike weight decay, entropy regularization directly focuses on the target distribution, and its influence can effectively backpropagate to specific parameters through adaptive optimizers like Adam, accommodating the heterogeneity of Transformer architectures. A more deeper exploration of this topic is left for future work.
D.4 CE WITH ENTROPY REGULARIZER
Inspired by our framework, one may propose an alternative approach for preserving output diversity is to introduce an entropy regularization with the CE loss. We discuss this approach in this section.
Formally, we have that
max Ee Ey-pt a)) [log f(y|=)] +y Byy log f fu/2=)]
(15)
(P.(P.))) Foonsstat
=H(()
This corresponds to the forward KL. with maximum entropy regularization. This corresponds to the forward KI. divergence with maximum entropy regularization. We note that this optimization problem does not have a closed-form solution. To better understand its behavior, we observe that this formulation leads an improper increase in tail probabilities when maximizing the entropy, as illustrated in Figure 11. In the context of LLMs, this increase often translates into nonsensical tokens in the vocabulary. leading undesirable generation outputs. A concrete example is provided in Table I. in which we directly sample sentences.
Table I: Examples of generated sentences for models trained with the altraaeeedbar dataset. The symbol "[...]" indicates good sentences that have been omitted to save space. This example shows that CE with entropy regularization can generate nonsensical tokens (shown in red) due to the increase in tail probabilities.
Prompt
Whats the largest star in our galaxy?
CE
Hello! Atlantis is legendary city that was said to have existed in ancient Greece. According to the story, it was a highly advanced and prosperous city that was located on an island in the ocean. [...]
CE+Entropy
Hello! Atlantis Documentary is a 2019 American documentarychoFLIce film directed by Já oblík and produced by Werner Herzog. The film explores the history and legacy of Atlantis, an ancient Greek city-state that was said to have calendar knowledge and advanced technology, through interviews with scholars and histori- ans.ython
GEM-LS
Hello there! As a helpful, respectful, and honest assistant, I'd be happy to help you explore the fascinating topic of Atlantis! Atlantis is an ancient Greek myth that tells the story a legendary realm said to have existed in the Atlantic Ocean, west of the Pillars of Hercules. [...]
The core issue of CE with an entropy regularizer arises because the gradient the entropy regularizer can dominate for tokens with low probabilities. Specifically, we have that
-op --
(x. of
0H af
-(1 log J).
where the division is element- wise. Consequently, for tokens with low probabilities in both ff and p. i.e., (x: 0 and p(x:) Rs 0, the gradient from the forward KI. divergence, -p(x)/f(x)) could be of
21

---

Published as a conference paper at ICLR 2025
constant order. In contrast, the gradient from the entropy regularizer, (-1 log f(x)) -c as f()) + O. This imbalance disproportionately increases the tail probabilities, leading to undesirable effects.
In contrast, the proposed method, GEM, does not have this issuc. This is because both distribution matching and entropy regularization in GEM are defined jointly over the generative distribution f. rather than being applied separately to the data distribution and generative distribution f. As shown in the proof of Proposition the gradient of GEM is given by:
--C -
-
aʃ J+
where the first term, f. is identical to that of CE with entropy regularization, but the second term q/f. is unique GEM. Since q = softmax(1/$ log f). q is more squeezed distribution than f. Consequently, for f(x) = 0, we have q(z)/ fx) < 1. This ensures that the gradient is not dominated in the cww-robbbilityy (tail) region, preventing improper increases in tail probabilities. Thus, GEM
achieves a more balanced optimization.
D.5 RELATION WITH GAN
We clarify that while our framework shares similar structure with GANs (Goodfellow et al., 2014; Lolicoeur-Martnneuu 2019), it has totally different meaning. Specifically, GANs originally designed for image generation tasks, where discriminator network is additionally introduced to measure the distance between distributions of real and generated images. We do not follow this storyline. Unlike in GANs, where measuring distances between image distributions is challenging, computing distances between token distributions language models is simple due to the discrete nature of token distributions. As we have explained, the introduction of the variable in our framework to control the direction and magnitude of probability transfer during distribution matching objective that is distinct from the goal of GANs.
We note that Proposition 3 also has an important computational implication. While the distribution matching problem in Equation (8) is thcoretically well-defined, it is computationally intractable in practice. In contrast, our game-theoretic formulation provides a feasible alternative. To understand why, consider the decomposition of the reverse KL divergence:
DKL. (((x))))) = EE [wlx log f(y(x))
Σ (w/xx) log f()x - E f() p(ylæ)
p(y(
In practice, we only have access finite samples drawn from p rather than direct knowledge of log ((gl().).. Consequently, the term Ey f(y]))) log p(y(<)) cannot be easily estimated from finite samples, rendering the formulation in Equation (8) impractical. This difficulty contrasts with the CE formulation, which corresponds finite samples. We note that a similar situation arises in GANs (Goodfellow et al., 2014), where the game-the theorete approach effectively solves the Jensen- Shannon divergence minimization problem, even though the Jensen- -Sannnon divergence itself cannot be directly estimated from data samples.
forward KL divergence that can be easily estimated from
22

---

Published as a conference paper at ICLR 2025
CE+Entropy
x
X
X
X
GEM-LS
Figure 11: Distributions of next-token probabilities for trained models with the UltraFeedback dataset, presented from top bottom: CE, CE+Entropy, GEM-LS. The prompt "Give me single-digit number". The top 300 probabilities shown with a subsampling rate of 20 for clear visualization. dotted line indicates the probability threshold of 10-4, The figure demonstrates that the CE+Entropy model has longer tail with higher probabilities assigned to some nonsensical tokens, marked with crosses.
23

---

Published as a conference paper at ICLR 2025
E EXPERIMENT DETAILS
All experiments are conducted using A800-80GB GPUs with the DeepSpeed distributed training framework, utilizzig ZeRO-2 and gradient checkpointing without offloading. We use flash-attention-2 with deterministic backward for reproducibility. The experiments are based on the pretrained Llama- 3-8B model, using Adam as the optimizer with a global batch size of 128. Following (Yu al., 2023; Liu et al., 2023; Cui et al., 2024), the learning rate is set to 2 x 10-5, with warm-up ratio of 0.03 and cosine learning rate decay. Training is performed over epochs. All supervised datasets tokenizer. When generation of responses is required for evaluation, we use the vLLM to accelerate inference.
formatted into the chat format using the Llama-3.
E.1 ULTRAFEEDBACK
We use the dataset filtered by HuggingfaceH4 team, which is available at co datasetss /ausets//uggiii Higgguiiiii wig.ginngaaaa 4ceH4/ /ul tettrrrffebbck www.biiiiiiieeee
'he dataset contains 61,135 training samples and 1,000 test samples. For training, we set the maximum sequence length 2,048, dropping longer sequences and padding shorter ones. To achieve a global batch size of 128, we use a per-device batch size of 4, a gradient accumulation step of 4, and GPUs. The training times takes about 24 GPU hours for all methods. For the CE method, we have tuned hyperparameters for weight decay and entropy regularization, selecting values from { {0.1. 0.01, 0.001 }. In both cases, value of 0.1 provided the best overall results. For NEFT, we use a noise scale hyperparameter of as recommended by (Jain et al., 2023).
Evaluation metrics, including perplexity, and entropy, are based on these 1,000 test samples. For entropy calculation, we compute the conditional entropy, whose expectation can be calculated exactly, and average over the sequence.
For the chatting task, we use the 805 test questions from the AlpacaEval 1 dataset and employ the reward model Fs IXLLMAA. -wt.aMMM-----..-11- The maximum generation sequence length set to 2048. For each question, 32 samples generated with the configuration t tempera ture=0 top 50, top p O. 9. To calculate the win rate, we use the Saadley-Teryy model:
exp(r(x, v))
Ply > y x)) exp(r(x,, y)) + exp(r(x, y'))
We use GPT-4 generated responses as baseline for calculating the win rate, specifically the 1pt4..106..revii version.
For the code generation task, there are 164 test questions for HumanEval, We use the prompt from (Wei et al., 2024):
You exceptionally intelligent coding assistant that consistently delivers accurate and reliable responses to user instructions. @ @ Instruction {instruction]
For each question, 200 responses are generated with the configuration temperature-0.f top_k=50, <oo.p==.... 9 to estimate the pass rate. The evaluation scripts are from https:// github, com/ wwwsseeuuueeeeeee
www.//s-euuuemmaaaaa.eee------ee wttp:///mma..iddrr//eeoeoeooeeee wwt.o///mai//////iie//ee/eeeeeea mmain/eeperiii eeeeeepprrrmmeee ttettcooeeee.eeee py.
F ADDITIONAL RESULTS
F. DIVERSITY REDUCTION OF SFT VIA CE
We provide empirical evidence that while SFT using CE loss improves the instruction-followinng ability of models for downstream tasks, it significantly reduces the output diversity of pre-trained models. To demonstrate this, we conduct experiments using the pre-trained L.IM, Llama-3.1-80
httpps //gips:://iiii tht.///ttttttttto/e thub. www/tatsu-ttaaacecocooocco.a tatstaattt/aaaaaaa paca_evaa /b1 ob/mm/ in/eesul hsspppppg11100 preview/model http:////mmmdeeeoottuuueetee _outtutss json
24

---

Published as a conference paper at ICLR 2025
employing few-shot learning teach the model to follow instructions and answer questions. We compare the output diversity of the pre-trained LLM with that of the same model fine-tuned using prompt both models answer 100 questions from the AlpacaEval dataset
CE loss. Specifically, (Li et al., 2023), generating 32 responses for each question. Following (Kirk et al., 2023), we then measure output diversity using the following metries: 1) N-gram diversity: the proportion of distinct n-grams in single response (intra-diversity); 2) Self-BLEU diversity: calculated as 100 minus the Self-BLEU score (inter-diversity), where one response is treated reference among multiple generated responses; 3) Sentence-Beil diversity: the cosine dissimilarity between pairs of responses the embedding space. All criteria range from 0 100 (with Sentence-BEi diversity scaled by multiplying by 100), and higher values indicate greater diversity.
The results, presented Table 2, show that fine-tuning with CE loss reduces output diversity compared to the pre-trained L.M. For completeness, we also report the output diversity of GEM in this setting. GEM demonstrates better output diversity than the CE-tuned model and performs comparably to the pre-trained LLM-in fact, it shows slight improvement. We believe this marginal improvement is partially attributable to the few-shot prompting strategy, which slightly restricts the output space of the pre-trained LLM.
Table 2: Output diversity of pre-trained, CE fine-tuned, and GEM fine-tuned LLMs. Higher values indicate greater diversity.
N-gram Diversity
Self-BLEU Diversity
Sentence-BERI Diversity
Pre-trained LLM (5-shot)
23.3
73.6
14.2
CE
22.5
71.1
13.9
GEM
25.8
75.4
15.4
F.2 SENSITIVITY HYPER-PARAMETERS
In this main text, our theory suggests that the hyper-paramette parameter ,3 controls the strength of entropy regularization, with entropy regularization. We believe such guidance is sufficient for practitioners. Additionally, we provide sensitivity studies on the chat task in Table 3. We observe that as 3 + 1, the performance of GEM converges to that of CE, and as 3 + O, GEM benefits from diversity preservation.
corresponding to no regulation and / - corresponding to strong
Table 3: Performance on Chat task of Llama-3.1-4 1-BB when trained with the UltraFeedback dataset.
Random Sampling
BON@32
CE
24.5
47.3
GEM (3 = 1.0)
24.7
47.5
GEM (3 0.9)
24.7
48.3
GEM (3 0.8)
24.8
49.4
GEM (3 0.7)
24.9
50.4
GEM (3 0.6)
24.3
50.1
F.3 MATH REASONING TASKS
In this section, we demonstrate that GEM is also effective in mathematical reasoning tasks. We evaluate the model trained with the U1 raFeedback dataset on the GSMSK (Cobbe et al., 2021) task. We use the following prompt to generate responses:
Your task to answer the question below. Give step-by-step reasoning before you answer, and when you't re ready to answer, please use the format "The answer is: Question: {question}
We consider three metrics: greedy decoding, majority voting (Wang et al., 2023), and pass rate. For the latter two metrics, report the performance across 16 samples with temperature of 0.6 and
25

---

Published as a conference paper at ICLR 2025
top-p of 0.9. The results are shown in Table 4. We find that GEM mitigates overfitting, evidenced by improved performance greedy decoding. Additionally, by preserving output diversity, GEM enhances performance in majority voting and pass rate.
Table 4: Performance on GSM8K (Cobbe et al., 2021) of Llama-3. 1-88B when trained with the UltraFeedback dataset. For majority voting and pass rate, we use 16 samples.
Greedy Decoding
Majority Voting
Pass Rate
CE
49.7
67.3
88.8
CE +WD
50.0
65.9
89.5
NEFT
48.4
66.3
89.1
GEM
51.3
69.5
91.1
F.4 SCALABILITY
In this section, demonstrate the scalability of our method across models of varying sizes, ranging from 3B to 70B parameters, including Qwen-2.5-3B and Qwen-2.5-7B (Qwen. 2024), Gemma-2-9B (Gemma et al., 2024), and Alama-3.1-700 (Dubey al., 2024). For all models except Llama-3.1 1- 70B, we apply LoRA (Hu et al., 2022) with a rank of 16 due to computational constraints, while maintaining the same settings and hyperparameters used in our experiments with Llama-3,1 1--BB.. The results presented in Figure 12 show that the benefits observed with Llama-3.1 1--BB extend to other models as well, confirming the broader effectiveness of GEM.
CE4W0
NEPT
GEN
Chatting (BON))32)
OpenLUM Leaderboard (Avenged over tases)
Chatting (BCNP)32)
HumanEral
Openi Lesecher Leacherbourd (Averaged over taske)
(PasseP100)
(Pass8P100)
(a) Qwen-2.5-3B
(b) Qwen-2.5-7B
CC-WD
KETT
GIN
Chatting (BNN822)
HumonBval (Pas0:100)
OpenLUM Leaderboard (Averaged over tasii)
Chatting (30N@16)
HumanSwil (Pase@6S))
ODenLLM Leacer jiveraged over tuska)
(c) Gemma-2 2--9B
(d) Llama-3.1 (LoRA)
Figure 12: Performance across different architectures and model sizes. The results show that GEM consistently outperforms CE.
F.5 IMPORTANCE OF DATA RESET
In this section, we extend the technical challenge in Section 5 regarding how to solve the problem in Equation (12). To recap, the problem is that
E,, Eyg ~pt (n) (n) Bygee (a)) ( (log (yyss (x:) log log/ (()
s.t. argmax E, E ( [log ((1rr + 1/3 H((T(-1x))
key challenge is that the expectation Egene[ cannot be calculated as easily as before. Worse still, Monte Carlo estimation as used in (Chen et al., 2024; Li et al., 2024a), by drawing samples from the distribution, does not provide an accurate gradient estimate. fundamental difficulty in this stochastic approximation arises from the distribution shift between the SFT data and the pre-trained
26

---

Published as a conference paper at ICLR 2025
distribution. To better understand this, refer to examples provided in Table 5. We observe that SFT data typically has finite-lnnghh sequences, while the pre-trained distribution produces samples that repetitive and can even be infinite in length.
This causes issues: assuming the probability of any token is lower-bounded by a small number E (0, 1), this means that log ff ( (x) approaches -oc when T goes infinite for pre-training distribution data. While this might seem acceptable as the gradient would reduce the probability of such samples, the challenge that the sample size is vast: for the Llama-3-8B model, the vocabulary size is 128k, and with a typical sequence length of 2048, the sample space size is 128000201s This makes difficult for the model to find effective directions for improvement. To validate this claim, we directly implemented the idea of stochastic approximation and found that training failed after 80 optimization steps (with 10k samples),, and the model could not generate good responses; see Table 5. In fact, techniques in (Chen et al., 2024; Li et al., 2024a) are usually applied to models after SFT, where the distribution shift between the model and data smaller. This technical remark is also discussed in the online forum ht / githhb..thh.ggggggggggggg cww/coo///////o huluuclaaa aml CTLL/SSIN//// sueues/////e / issuecomment 2062926711
G FUTURE WORK
Below, we explore GEM's potential applications.
Better Cold-Start for Online RI. Training. SFT typically serves as a cold start for online RL training. To fully harness the potential of RL, exploration across diverse generated responses is critical. However, CE often reduces diversity, making it less effective in this context. Generalized Experience Replay (GEM) helps maintain diversity when learning from cold-start data, ultimately enhancing RL training performance. This idea is explored in recent blog post https:/wwww.notoonso/Caan -Better- Cold-Start-Strateee Siaaaa-ttrrattees--ppoeeeee
Training-for-LLMM 17aa0742a5 68082286 6c867c55bbbbb..
Preference Collapse in RLHE SFT-trained models can be refined through Reinforcement Learning from Human Feedback (RLHF) to better align with human values (Ouyang et al., 2022; Bai et al., 2022). In this context, Xiao al. (2024) studied the impact of SFT models on preference learning in RLHF, demonstrating that if an SFT model collapses (i.c., becomes biased toward certain outputs with near-certain probability), it can further lead preference collapse alignment. Their findings underscore the importance of addressing collapse during the SFT stage.
Synthetic Data Generation. SFT-trained models are often used as synthetic data generators for self- improvement (see. e.g., (Adler et al., 2024; Dubey et al., 2024)). In this context, maintaining output diversity is essential. By generating wide range of diverse outputs, models can explore various potential solutions, reducing the risk of overfitting and uncovering better-performing strategies. Our experiments about best-of-n against& a reward model Section is inline of this topic.
Mode Collapse. When models are repeatedly fine-tuned on text generated by their predecessors, linguistic diversity gradually erodes. This recursive process amplifies existing errors and biases, with each successive generation inheriting the limitations of its predecessors. This phenomenon, known as mode collapse, has been extensively documented in prior studies (Guo et al., 2023; Shumailov et al., 2023). By introducing entropy regularizer, expect to help mitigate mode collapse in the self-improvement.
Thiis method is computationally slow due to sampling responses. In fact, these computational resources of 80 steps actually more than those required for GEM for epoch tuning of 60k samples.
27

---

Published as a conference paper at ICLR 2025
Table 5: Examples of generated sentences from pre-trained models. The symbol "[...]" indicates sentences have been omitted to save space for the ground truth label, and the repetitive sentences for the pretrained model. Examples illustrate the distribution shift between SFT data and the pre-train pre-trained distribution. In SFT data, ground truth labels have a finite length (marked by an explicit <EOS> in the example), whereas the pre-trained distribution often generates repetitive sentences without clear endpoint.
Prompt
I would like to know more about Clyde Hill Washington...whatt can you tell me?
Ground Truth Label
Clyde Hill a small, affluent city located King County, Washington, United States. is situated on the castern shore of Lake Washington, north of Bellevue, and east of Kirkland. It is known for its stunning views of the lake and the Cascade Mountains.
Clyde Hill was incorporated in the year 1953 and is home to ap- proximately 2,800 residents. The city has mix of residential and commercial areas, with large estates and luxurious mansions, making it one of the wealthiest ZIP codes in the United States. The median household income is around $200,000, and the median home value is roughly $1.6 million.
[..]]
Overall, Clyde Hill, Washington, offers unique and luxurious lifestyle to its residents. The combination of picturesque scenery, luxury living, and excellent amenities makes it truly special place. <EOS>
Pretrained Model
would like to know more about Clyde Hill Washington...what can you tell me?
I would like to know more about Clyde Hill Washington...watt can you tell me?
would like to know more about Clyde Hill Washington...whtt can you tell me?
I would like to know more about Clyde Hill Washington...what can you tell me?
[..]]
GEM without reset trick
Thank you for reaching out with your queries on Clyde Hill Washing- ton! As of 2021, the estimated median home value in Clyde Hill, Wash- ington, is $3,750,000 [...] Thank you for trusting me with your queries and I hope was able to help. Stay in touch and talk to you soon! HexatrigsimalleusfmmcCC ())dom((etssdde každ r
GEM
Clyde Hill is small, affluent neighborhood located in the eastern part of Washington, D.C., bordering the Potomac River. It known for its stunning views of the Washington Monument, the National Mall, and the Potomac River. The area is home to mix of embassies, diplomatic missions, and luxury residential buildings.
Clyde Hill also home to the historic Clyde Hill House, which was built in 1929 and served the official residence of the U.S. Am- bassador to France from 1933 to 1946. The house now a private residence and is not open the public.
[..]]
Overall, Clyde Hill is unique and beautiful neighborhood that offers glimpse into the history and elegance of Washington, D.C. <EOS>
28


Reviews:
[{"review_id": "hZjafVPLX6", "reviewer": "Reviewer_xjk9", "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.", "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?\n- 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?\n- For the CE + entropy baseline [1], what coefficient γ did you use to weight the regularization? Did you conduct a grid search?\n- Line 33: I suggest replacing \"Despite extensive pre-trained, \" with \"Despite extensive pre-training, \" or \"Despite being extensively pre-trained, \".\n- Line 423: Entropy regularization supports Principle 1.\n\n[1] Abhimanyu Dubey, Otkrist Gupta, Ramesh Raskar, and Nikhil Naik. Maximum-entropy fine grained classification. Advances in neural information processing systems, 31, 2018.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- SFT is a timely and impactful area of research as it is an essential phase of the LLM pipeline.\n- The proposed GEM method is grounded in theory, effectively tackles overfitting, and improves the output diversity of Llama3-8B.\n- The experimental setup is modern.\n- The paper is clear and well-structured.", "weaknesses": "- GEM is only evaluated on Llama3-8B, which might particularly benefit from GEM. \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cgbuUiSCh6", "reviewer": "Reviewer_yyC9", "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.", "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?\n2. Is there any performance degradation when increasing the temperature at inference time?\n3. 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.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The work is well-motivated, and the writing is clear.\n2. They demonstrated increased diversity in the model's output through high performance on the test-time compute set.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6Jjhc3dhBh", "reviewer": "Reviewer_4XM8", "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.", "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?\n\n2. 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?\n\n3. What is the functional difference between reverse KL and normal KL?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This article is well-written, with clearly organized content and significant research value.\n2. The proposed GEM loss is simple and effective; this straightforward loss improvement simultaneously enhances diversity and accuracy.\n3. 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.\n2. There is no comparison of training costs across different methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oHQyEwci5x", "reviewer": "Reviewer_wHN6", "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.", "questions": "* In scenarios with limited training data, does GEM perform consistently, or would CE still be preferable due to its simplicity?\n* Would the regularization in GEM, especially the entropy component, introduce a non-trivial increase in compute or memory requirements, especially when scaling?\n* 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.\n* In cases where GEM aims to enhance diversity, could it inadvertently introduce biases in generation, particularly when generating responses with multiple valid answers?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n* 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.\n* 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\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)