GEM: Geometric Entropy Mixing for Optimal LLM Data Curation

LLM pre-training efficacy increasingly depends on data composition rather than sheer volume. Yet, optimal mixing is hindered by categorization flaws: human taxonomies suffer from ontological misalignment, and Euclidean clustering fails to address embedding anisotropy. We introduce GEM (Geometric Entropy Mixing), a framework reformulating data curation as a variational problem on the hypersphere augmented with a mixing-balance regularizer. By decoupling the generative prior and optimizing the objective via a provable MM (Minorize-Maximize) algorithm, GEM effectively counteracts the cluster collapse to discover balanced semantic structures invisible to Euclidean heuristics. We employ teacher-student distillation to scale this geometric fidelity to web-scale corpora and introduce the Geometric Influence Score (GIS) for interpretable taxonomy generation. Experiments with 1.1B-parameter models demonstrate that GEM establishes a new state-of-the-art when integrated into mixing strategies like DoReMi and RegMix, improving average downstream accuracy by up to 1.2% and offering a robust coordinate system for predictable data mixing.

Large language models learn from huge collections of text, but simply adding more text is not always enough. The mix of training data matters: it should cover many useful kinds of content without letting a few large topics dominate. Existing methods often sort web text using hand-designed categories or simple automatic grouping, but these can miss subtle topics or bury smaller ones. We developed GEM, a method that groups documents according to how their meanings are arranged in a model’s internal map while keeping the groups balanced. GEM first discovers these groups from a smaller sample, then trains a fast lightweight classifier to label the much larger corpus efficiently. It also selects representative examples from each group so people can understand and name the discovered topics. In our experiments, GEM helped existing data-mixing methods produce better training mixtures. Models trained with GEM-based mixtures generally performed better on science, commonsense, and logic/language tasks.