Diversity as a Reward: Fine-Tuning LLMs on a Mixture of Domain-Undetermined Data

TL;DR

This work investigates how semantic data diversity can boost cross-domain abilities in large language models when domain labels are unavailable. It develops DaaR, a self-supervised framework that (i) synthesizes model-aware domain centroids, (ii) trains a lightweight domain-prediction probe and a diversity predictor on frozen embeddings, and (iii) uses a closed-loop, entropy-based diversity reward to select data for fine-tuning. Theoretical analysis via importance sampling explains why diversity-aware selection improves multi-domain performance and why naive, label-dependent approaches can be suboptimal. Empirical results across multiple models and seven benchmarks show that DaaR achieves state-of-the-art averages, with notable gains in challenging math and coding tasks, while maintaining efficiency and robustness across scales and tasks. This work highlights the potential of data-diversity guided, feedback-driven data-model co-design for robust, domain-agnostic LLM fine-tuning.

Abstract

Fine-tuning large language models (LLMs) using diverse datasets is crucial for enhancing their overall performance across various domains. In practical scenarios, existing methods based on modeling the mixture proportions of data composition often struggle with data whose domain labels are missing, imprecise or non-normalized, while methods based on data selection usually encounter difficulties in balancing multi-domain performance. To address these challenges, in this work, we investigate the role of data diversity in enhancing the overall abilities of LLMs by empirically constructing contrastive data pools and theoretically deriving explanations. Building upon the insights gained, we propose a new method that gives the LLM a dual identity: an output model to cognitively probe and select data based on diversity reward, as well as an input model to be tuned with the selected data. Extensive experiments show that the proposed method notably boosts performance across domain-undetermined data and a series of foundational downstream tasks when applied to various advanced LLMs. We release our code and hope this study can shed light on the understanding of data diversity and advance feedback-driven data-model co-design for LLMs.

Figures (21)

• Figure 1: The illustration highlights our observations and the proposed DaaR method. Observation shows the t-SNE visualization of embeddings for data samples with different distributions, leading to varying LLM's evaluation performance. In the domain-undetermined data selection scenario, DaaR introduces a dual-identity framework for LLMs: ① An output model to probe and select data based on diversity reward, ② An input model to be tuned with the selected data. The diversity probe is trained on model-aware synthetic data, enabling domain discrimination and diversity reward prediction.
• Figure 2: Training loss and validation process of the two stages of DaaR on Qwen2-7B, more detailed results in Appendix \ref{['sec:appendix-dynamics']}.
• Figure 3: Semantic cosine similarity across different domains for generated samples of size 10.
• Figure 4: The distribution of the score by Alpagasus.
• Figure 5: The distribution of the tags by Instag.

Theorems & Definitions (2)

• Proposition 3.2
• Proposition C.1