Improving Diversity in Language Models: When Temperature Fails, Change the Loss

Thank you for the detailed review and your feedback.

Decoding-based methods

In our paper, we deliberately focused on the temperature parameter (commonly referred to as the "diversity" parameter) because its theoretical analysis already presented some complexity. For this reason, we chose to limit the scope of this study to ensure a focused exploration of the temperature parameter. However, we agree that comparing empirically with other decoding-based methods could bring other insights.

As suggested, we investigated the effect of the following methods on the WritingPrompts dataset:

• top-p (or nucleus) sampling
• KL-Divergence Guided Temperature Sampling

We re-implemented the KL-Guided sampling from scratch, since the original implementation was not compatible with our PyTorch models. We used the same range for $\sigma$ as in the original paper.

Top-p

KL-Guided

(As $\sigma$ increases, KL-Guided approaches standard decoding.)

We can conclude that Top-p increases Precision at the cost of Recall. Since Top-p removes the least probable tokens, we could intuitevely expect a higher Precision.

KL-Guided seems to globally decrease both P&R. We believe that an explanation would be that this method has been mainly designed for conditional text generation, like summarization and question answering, which might exhibit specificites compared to our generation task.

Overall, we believe that these results provide more grounding on the necessity of using Recall-oriented losses as opposed to simple decoding-based methods.

General models

To answer your question, we trained a general-purpose instruction model on the Alpaca dataset using our proposed losses, $\lambda$-PR and $c$-Div, to demonstrate their effectiveness on broader instruction tuning tasks. We then evaluated this model on both the WritingPrompts and MathQA-Python datasets.

For MathQA-Python, which requires generating Python code from natural language questions, we used 3 in-context examples to prompt the model and evaluated it as in CodeContest. We used the same evaluation as for CodeContest. We chose MathQA over CodeContest because it offers more training data and is better suited to the capabilities of our models.

Alpaca MathQA-Python

On these results, we see that overall, our losses do not significantly impact Precision, meaning that generalization/in-context ability is not affected. However, we observe a significant increase in Recall, which is consistent with our previous findings.

Alpaca WritingPrompts

On the WritingPrompts dataset, we observe that general models trained with our losses on the Alpaca dataset exhibit a similar pattern to the specialized models. We see a significant increase in Recall, sometimes at the expense of Precision.

This suggests that our losses can consistently improve Recall even on more general models and tasks. We hope these additional experiments provide a more comprehensive view of the effectiveness of our proposed losses.

Suggestions

• We will add more details about P&R in the appendix, this should help clarify the evaluation process.
• We computed MAUVE metrics for the WritingPrompts dataset. We will add it to Table 1.

• For 243–244, we mean that image generation can be optimized for specific PR tradeoffs using f-divergences, as shown by Verine et al. (2023). However, these methods rely on assumptions that don’t hold for text due to its causal nature. We’ll clarify this in the final version.
• In code generation, pass@100–pass@1 reflects the gain from sampling multiple candidates over one. In your example, the metric is relevant: if samples differ only slightly (e.g., variable names), structure diversity is low and pass@1 is likely high. We’ll clarify this with more details on the evaluation process.

We would like to thank Reviewer 2owJ for the detailed review and suggestions. We are glad you appreciate the paper.

Following your suggestion, as well as those from other reviewers, we conducted additional experiments:

New dataset: MathQA-Python.
We trained a model on MathQA-Python and used the same evaluation as for CodeContest: Precision as pass@1, and Recall as the gain obtained from sampling (pass@100 - pass@1). We chose MathQA-Python over CodeContest because it offers more training data and is better suited to the capabilities of our models.

New model: Llama3.2-3b WritingPrompts.
We trained an additional model on the WritingPrompts dataset.

General-purpose instruction model.
We trained a general-purpose instruction model on the Alpaca dataset using our proposed losses, $\lambda$-PR and $c$-Div, to demonstrate their effectiveness on broader instruction tuning tasks and more expressive models. We then evaluated this model on both the WritingPrompts and MathQA-Python datasets.

We report the results below and will update Table 1 in the final version. While time constraints limited the scope, we believe these additional experiments provide a strong indication the effectiveness of our losses on more tasks, models and metrics.

Olmo1b MathQA-Python

An interesting observation is that $\lambda$-PR impacts Precision more severely than the other losses, but ultimately achieves a high Recall. This suggests that the resulting model offers better coverage of the target distribution.

Llama3.2-3b WritingPrompts

Note that for this larger model (compared to Olmo1b), we did not benchmark the TruncR loss due to its incompatibility with distributed training in the current implementation. We also observed that the model tuned with $\lambda$-PR was much more sensitive to the sampling temperature. We used $t=0.5$, as higher values led to some degeneracies. However, we could not investigate this further.

Alpaca WritingPrompts

Alpaca MathQA-Python

All Alpaca experiments used a temperature of $t=0.5$ to avoid degeneracies.

Analysis

We observe the same pattern as in the initial experiments. This confirms that our losses can achieve higher Recall than NLL, which we believe strengthens the claims and findings presented in the paper.

Section 6.2

• For Section 6.2, we will reformulate the text to clarify that the experiments are designed to verify the assumptions behind the theoretical analysis, not its results.

Experiments

• We used hyperparameters very similar to those described in the original Olmo paper (we only used a smaller learning rate of 1e-6 instead of 2e-6). To ensure comparability, we used the same optimization parameters for all losses.
• For the PR metrics, we used the exact same setup as in the original paper (Le Bronnec et al., 2024), i.e., GPT2-large as the embedding model.
• We trained all models under the same conditions (same number of epochs, same batch size, same optimizer, etc.) on the training set and reported the metrics on the validation set.

We will incorporate these details in the appendix of the paper.

Suggestions

• We will gladly incorporate the suggested reformulations, it will indeed improve the overall flow of the paper. We could indeed move some parts to the appendix and spend more time discussing the idea behind the model.

For a side note, in Theorem 4.2, the bounds rely on almost no assumption on the target distribution or the model, and the bound is not necessarily useful for large $t$, especially if $Z$ is low. For instance, when the model is uncertain (i.e., $Z$ is low), the bound may still be useful even for $t < 1$. But we agree that this is more a general theoretical result than a practical one.

We thank Reviewer DUE2 for the review and constructive feedback. We have added further experiments with more datasets and models; please refer to our response to Reviewer 2owJ for details. While the first question has only one relevant scenario, we believe the second and third questions require multiple scenarios to answer meaningfully.

The authors provide three scenarios to answer the three proposed questions, yet in the subsequent analysis, not all of these questions seem to be fully addressed by the three given scenarios. Therefore, additional experiments may be needed to ensure that the issues claimed in the three scenarios are adequately answered (e.g., in Section 6.4, only the WritingPrompts dataset task is discussed).

Following your suggestion, as well as those from other reviewers, we conducted additional experiments within the limited available time to further investigate the effect of our Recall-oriented losses. (Similar response as to Reviewer 2owJ) The new experiments include:

• New dataset: MathQA-Python.
  We trained a model on the MathQA-Python dataset. Using the same proxies for evaluating Precision and Recall as in CodeContest.

• New model: Llama3.2-3b WritingPrompts.
  We trained an additional model on the WritingPrompts dataset.

• General-purpose instruction model.
  We trained a general-purpose instruction model on the Alpaca dataset using our proposed losses, $\lambda$-PR and $c$-Div, to demonstrate their effectiveness on broader instruction tuning tasks. We then evaluated this model on both the WritingPrompts and MathQA-Python datasets.

We report the results in the tables below and will extend Table 1 accordingly in the final version of the paper.

Olmo1b MathQA-Python

An interesting observation is that $\lambda$-PR impacts Precision more severely than the other losses, but ultimately achieves a high Recall. This suggests that the resulting model offers better coverage of the target distribution.

Llama3.2-3b WritingPrompts

Note that for this model, which is larger than Olmo1b, we did not benchmark the TruncR loss, as the current implementation is not compatible with distributed training.
We also observed that the model tuned with $\lambda$-PR was much more sensitive to the sampling temperature. We used $t=0.5$, as higher values led to some degeneracies. However, due to limited time, we could not investigate this further.

Alpaca WritingPrompts

Alpaca MathQA-Python

We started from a pre-trained Llama3.1-8B model and used the same training setup as described in the paper. This yields a basic instruction-tuned model, capable of generalization (but still with limited capacity compared to SOTA models).
Note that for all experiments conducted on Alpaca, we used a temperature of $t=0.5$ to avoid degeneracies. As with previous experiments, we could not benchmark the TruncR loss due to its incompatibility with distributed training.

Analysis

We observe the same pattern as in the initial experiments. This confirms that our losses can achieve higher Recall than NLL, which we believe strengthens the claims and findings presented in the paper.

We thank Reviewer pQ4a for the thoughtful feedback and careful review. We are grateful for the time and effort. First, regarding the algorithmic questions raised by several reviewers: we will include detailed algorithm blocks for all computations in the appendix.

We address explicit questions below and will incorporate all clarity improvements in the next revision.

Additional experiments

Following your suggestion, we added experiments on the MathQA-Python dataset (details in response to reviewer DUE2). We also trained a general-purpose instruction model on the Alpaca dataset, and observed a similar trend as in our initial experiments. We hope these results further support our method's generalizability.

Computing the support

• In practice, the support size of $P$ is indeed intractable. However, we can still obtain a rough characterization of the sparsity of $P$ (introduced in Thm 4.2) by leveraging the sparsity of the token-level conditional distributions $P(\cdot \mid x_{<l})$. To this end, we approximate the reference $P$ using a strong model $Q_{\theta}$ (Llama3.1-8B in our experiments). For each sample and each token position $i$ in the solution, we compute the conditional distribution $Q_{\theta}(\cdot \mid x_{<i}) \in \mathbb{R}^V$, where $V$ is the vocabulary size. At each position, we then determine the smallest number of tokens $n_{\mathrm{top_p}}$ that account for $p \in {0.9, 0.95, 0.99}$ of the total probability mass. For each sample, we take the maximum of these $n_{\mathrm{top_p}}$ values across all positions, and finally, we compute the geometric mean of these maxima across the dataset to estimate the sparsity of $P$.

• We refer to this estimate as an "upper bound" since the true sparsity of $P$ is likely lower than what we obtain from the conditional distributions alone.

• We will make sure to include this explanation in the final version of the paper.

Confusing explanations

• Thanks for the feedback, we will better distinguish between PR over the full distribution and conditional PR.

Conditions for an increasing Recall

• Yes, all conditions mentioned in Prop. 4.3 should hold simultaneously. In the case $\rho \approx 1$, the temperature does not help, as the model is already close to the target distribution.

Eq.18 and TruncR loss.

• We indeed differentiate only the terms outside the indicator function. The notation $\bar{Q}$ indicates that we use the value of $Q$ in the implementation, but no gradient flows through it (achieved via detach() in PyTorch). In Eq. (18), this implies that the gradient of the sum is zero when the condition is not satisfied. We will add an algorithm block to clarify this in the final version.
• Regarding $\delta$, it characterizes the proportion of samples to keep. So, yes your interpretation is correct. During training, samples below the $\delta$ threshold will have their likelihood increased. To maintain a fixed proportion of selected samples, $\delta$ will gradually increase.

Prop 5.2 + teacher model

• Proposition 5.2 is proved in Appendix B.1, but is not clearly referenced in the main part. We will fix this. By "difficult," we refer to the fact that both of both sampling from $Q$ and computing $P(\cdot | x_{<l})$, since $P$ is unknown.
• In a different setup, you're right that $P$ could be approximated by a teacher model. But in that case, the objective would be to match the Recall of the target model to that of the teacher, not to the true distribution. This, along with the idea of using RLHF to improve Recall, are very interesting directions for future work, but we believe they deviate somewhat from the main focus of our paper and theoretical analysis. We will nonetheless add a discussion on this in the final version of the paper.

Fig 4 PR

• Unlike Figures 5, 6, and 7, which display empirical values of Precision and Recall (computed using the metrics from Le Bronnec et al., 2024), Figure 4 is a theoretical illustration of the effect of each loss on the PR-curve, as in Verine et al., 2023. It is parametrized by $\lambda$, and showing both $\alpha_\lambda$ and $\beta_\lambda$ is standard in that context. Our goal was to illustrate that optimizing the $\lambda$-PR objective increases the corresponding tradeoff value for a given $\lambda_0$. In particular, for $\lambda_0 < 1$ (i.e., below the line $y = x$), Recall increases more than Precision.

Fig 6 $\lambda$

• The three points corresponds to models trained with $\lambda$-PR for $\lambda=1$ (matching TailR loss), and some $\lambda<1$. We observe that $\lambda < 1$, Precision decreases while Recall increases, as expected.

Tasks sparsity

• In addition to CodeContest, we evaluated the sparsity of WritingPrompts (plotted in the table below), a creative task. The sparsity is lower than for CodeContest, which could be expected as this is a less constrained task.

Top-p | 0.9 | 0.95 | 0.99

• | - | - | - | |Supp(P)|/V | 3.86% | 7.80% | 24.9%