Minibatch Selection via Partition Matroid Constrained Gradient Matching

2026-06-05

AI Summary

This work addresses the challenge of balancing domain coverage and convergence efficiency when training large language models on heterogeneous data. The authors propose PartitionSel, a novel method that integrates partition matroid constraints with validation-guided gradient matching to enable efficient cross-domain minibatch selection. By simultaneously respecting per-domain budget limits and minimizing sample redundancy, PartitionSel leverages weak submodular optimization and orthogonal matching pursuit to provide theoretical approximation guarantees. Experimental results demonstrate that the approach significantly outperforms baseline strategies in fine-tuning Qwen2.5 and Llama-3, effectively reducing the number of conflicting gradient pairs within batches and enhancing the compatibility of training updates.

Abstract

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.

Problem

• minibatch selection
• heterogeneous data
• domain coverage
• convergence speed
• large language models

Innovation

• partition matroid
• gradient matching
• minibatch selection
• weak submodularity
• cross-domain optimization