UniGeM: Unifying Data Mixing and Selection via Geometric Exploration and Mining

TL;DR

UniGeM proposes a geometry-inspired approach to data curation that unifies macro-distribution balancing and micro-quality selection as a manifold approximation problem. It operates in two stages: Stage-I Macro-Exploration uses stability-driven clustering to determine an intrinsic manifold resolution $K^*$ and global sampling budgets, while Stage-II Micro-Mining refines within each cluster via probe-based semantics, a structural penalty, and a cohesion gate to select high-quality, on-manifold samples. The framework is backed by a Wasserstein-based theoretical bound linking global quantization and local pruning, and is validated by training 8B and 16B MoE models on 100B tokens, yielding $2.0\times$ data efficiency and improved reasoning and multilingual generalization over baselines. Empirically, UniGeM outperforms adapted SOTA baselines on code benchmarks and exhibits robust cross-language generalization, with careful ablations showing the indispensable roles of both stages and the value of the intrinsic resolution, cohesion, and structure penalties. The work demonstrates that preserving manifold topology and local dependencies during data curation can meaningfully enhance large-model performance while reducing data requirements, suggesting practical pathways for scalable, structure-preserving data pipelines.

Abstract

The scaling of Large Language Models (LLMs) is increasingly limited by data quality. Most methods handle data mixing and sample selection separately, which can break the structure in code corpora. We introduce \textbf{UniGeM}, a framework that unifies mixing and selection by treating data curation as a \textit{manifold approximation} problem without training proxy models or relying on external reference datasets. UniGeM operates hierarchically: \textbf{Macro-Exploration} learns mixing weights with stability-based clustering; \textbf{Micro-Mining} filters high-quality instances by their geometric distribution to ensure logical consistency. Validated by training 8B and 16B MoE models on 100B tokens, UniGeM achieves \textbf{2.0$\times$ data efficiency} over a random baseline and further improves overall performance compared to SOTA methods in reasoning-heavy evaluations and multilingual generalization.

Figures (11)

• Figure 1: Average downstream performance: Random sampling vs. UniGeM for pre-training an 8B MoE model. The inset radar chart shows that UniGeM achieves stronger multilingual performance across programming languages.
• Figure 2: Overview of the UniGeM framework. The curation pipeline operates in two hierarchical stages: (Left) Stage-I (Macro-Exploration) identifies the intrinsic manifold resolution $K^*$ via topological stability (Algorithm \ref{['alg:intrinsic_resolution']}) and derives global mixing weights $\mathbf{r}$ through a softmax over geometric scores (Eq. \ref{['eq:budget_softmax']}), ensuring comprehensive global coverage across diverse semantic regions. (Right) Stage-II (Micro-Mining) computes sub-cluster sampling weights $W(S_j)$ by combining the inherited budget $r_k$, a semantic score modulated by a structural penalty $\exp(-\lambda \mathcal{L}_{struct}(S_j))$ (Eq. \ref{['eq:structural']}), and a cohesion gate $[\beta_{S_j}+\epsilon]$ (Eq. \ref{['eq:geo_gate']}). This stage refines selection within each region to downweight off-manifold outliers while preserving representative local structural dependencies.
• Figure 3: Impact of Macro-Exploration (Stage-I). UniGeM (red) outperforms Meta-rater (green) and the w/o Stage-I variant (blue) across 7 languages, indicating that global clustering improves multilingual generalization.
• Figure 4: Intrinsic Resolution Selection. The stability index $J_{\text{stab}}(K)$ (blue line) rises and then plateaus, attaining its maximum at $K^*=72$. The marginal gain (orange bars) diminishes significantly beyond this point.
• Figure 5: Empirical distribution of a representative subset of sub-clusters in $C_{37}$. Sub-clusters are mapped by Semantic Score ($P_{S_j}$, X-axis) and Structural Penalty (Y-axis), with bubble size representing Cluster Cohesion. The color gradient indicates the Sampling Probability$W(S_j)$, illustrating the budget allocation toward sub-clusters.