Minibatch selection for Language Models via Partition Matroid Constrained Gradient Matching

Training large language models (LLMs) on heterogeneous data requires selecting minibatches that balance convergence speed with coverage across domains. Existing methods either select samples independently within each domain or rely on computationally expensive proxy models to learn continuous domain weights. We propose PartitionSel, a cross-domain minibatch selection approach that maximizes a validation-guided gradient-matching utility under per-domain budgets encoded as a partition-matroid constraint. By coupling the per-domain budgets through a single utility, PartitionSel is designed to reduce redundancy in selections across domains. The proposed objective is weakly submodular and admits an orthogonal matching pursuit algorithm with provable approximation guarantees. Empirically, we evaluate PartitionSel for minibatch selection during the fine-tuning of Qwen2.5 and Llama-3 on MetaMathQA and Mol-Instructions. PartitionSel achieves robust gains over per-domain and domain-agnostic baselines on both benchmarks. It also reduces the number of conflicting gradient pairs within each batch, indicating that the cross-domain coupling translates into more compatible training updates.

Training large language models requires mixing data from many different domains, such as mathematics, science, and chemistry. However, deciding which training examples to include in each minibatch is difficult: selecting data independently within each domain can lead to conflicting learning signals, while existing approaches that jointly optimize across domains are often computationally expensive. We introduce PartitionSel, a new method for jointly selecting training examples across multiple domains. Instead of treating each domain separately, our method chooses examples that work well together by matching the model’s learning updates to a small validation set. At the same time, it ensures that every domain is fairly represented through domain-specific selection budgets. We also show theoretically that this optimization problem has useful mathematical structure, allowing us to design an efficient algorithm with provable guarantees. We evaluate our approach on mathematical reasoning and molecular generation tasks using Qwen2.5 and Llama-3 models. PartitionSel consistently improves performance across many benchmarks while reducing conflicts between training gradients from different domains. These results suggest that jointly selecting data across domains can make large language model training both more efficient and more effective.