The conventional wisdom in LLM development has shifted decisively: data composition now matters more than raw scale. Yet the tools used to categorise pretraining corpora remain surprisingly crude, either human-defined taxonomies that impose ontological structure the model never asked for, or Euclidean clustering methods applied to embedding spaces that are fundamentally non-Euclidean. GEM (Geometric Entropy Mixing) addresses this mismatch directly, and the results are worth paying attention to.
The Problem with Existing Approaches
Modern text embeddings from models like BGE or RoBERTa are l2-normalised by construction, which means they live on a high-dimensional unit hypersphere. Cosine similarity is the natural metric in this space, not Euclidean distance. When you apply K-Means to these representations, you are imposing flat-space geometry on a curved manifold, and the consequences are predictable. Anisotropy, sometimes called the cone effect, causes representations to cluster in narrow sub-regions of the sphere. Euclidean clustering then collapses into a handful of dominant clusters that absorb the semantic long-tail, destroying exactly the diversity that makes diverse pretraining data valuable.
Taxonomy-based alternatives have a different failure mode. Assigning documents to human-defined categories using LLM classifiers is expensive, brittle, and semantically misaligned: the categories reflect how humans organise knowledge, not how a self-supervised model learns to represent it. Inter-annotator consistency on nuanced web data is poor even for frontier models, which makes the whole enterprise questionable at scale.
Key Contributions
GEM reframes data categorisation as a variational inference problem over a mixture of von Mises-Fisher (vMF) distributions on the unit hypersphere. The vMF family is the natural directional analogue of the Gaussian: its sufficient statistic is the cosine similarity between a point and the component mean direction, which aligns perfectly with how modern embeddings are trained and compared. The framework makes four concrete contributions:
- Balance-regularised hyperspherical objective. A mixing-balance regulariser is added to the entropy-regularised variational objective, acting on empirical assignment mass rather than the generative prior. This decoupling prevents the rich-get-richer feedback that standard EM produces, where large clusters accumulate probability mass and smaller semantic niches disappear.
- MM-based inference with monotonic ascent guarantees. The optimisation uses a Minorize-Maximize algorithm with a provable guarantee of monotonic objective improvement at each step. This is a meaningful theoretical property: many clustering objectives in practice are optimised with heuristics that offer no convergence guarantees.
- Teacher-student distillation for web-scale deployment. The full vMF mixture model (the teacher) is trained on a seed corpus, then used to pseudo-label the full corpus. These labels train a lightweight FastText classifier (the student), reducing inference to linear time complexity. This is the practical engineering that makes the approach deployable on trillion-token corpora.
- Geometric Influence Score (GIS) for interpretable taxonomies. GIS-guided sampling selects representative documents for each discovered cluster, producing human-readable category descriptions as a byproduct of the geometric optimisation rather than as a separate, expensive annotation step.
Experimental Results
Experiments use 1.1B-parameter LLaMA-style models trained on CommonCrawl-derived data for 25 billion tokens, a controlled budget that allows fair comparison across categorisation schemes. GEM-derived taxonomies are plugged into three mixing algorithms: DoReMi, RegMix, and a sensitivity-driven approach called Perf. Baselines include K-Means, WebOrganizer (domain and format variants), and TOS.
The headline numbers are credible rather than spectacular, which is appropriate given the difficulty of the task. Under DoReMi, GEM outperforms the strongest baseline (WebOrganizer) by 1.23 percentage points on Commonsense Reasoning and 1.76 points on Logic and Linguistics. Gains persist across Perf and RegMix, with GEM consistently leading on Science QA and Commonsense Reasoning. The one exception is K-Means holding a marginal edge on Logic and Linguistics under RegMix, which the authors acknowledge rather than paper over.
A more interesting diagnostic is the predictability analysis. The authors measure Spearman rank correlation between RegMix-predicted validation loss and ground-truth loss across held-out mixture vectors, repeated over ten independent splits. GEM produces higher and tighter correlation distributions than all baselines. This matters because it means GEM's coordinate system makes mixture search more sample-efficient: small perturbations to mixing weights produce consistent, predictable loss changes rather than noisy, entangled responses. This is arguably a stronger argument for GEM than the raw accuracy numbers, since it speaks to the controllability of the entire data curation pipeline.
The sensitivity analysis over cluster count K shows performance peaking around K=36, with degradation at K=48 attributed to over-fragmentation. This is a sensible result, and the authors are right to include it, though the optimal K will almost certainly be dataset and architecture dependent.
Limitations and Open Questions
Several honest limitations deserve attention. The experiments are conducted at 1.1B parameters and 25 billion tokens. Whether the geometric advantages of vMF clustering persist at 7B, 70B, or beyond is genuinely unknown. The authors flag this as future work, but it is a significant gap: the anisotropy structure of embeddings may shift as encoder models scale, and the optimal balance regularisation strength is likely sensitive to corpus composition.
The teacher-student distillation step introduces a potential fidelity loss that is not deeply characterised. FastText is a strong lightweight classifier, but the degree to which the student's boundaries approximate the teacher's hyperspherical partitions on out-of-distribution web data is not systematically studied. A calibration analysis here would strengthen the paper considerably.
The GIS-based taxonomy generation produces interpretable labels, and the t-SNE visualisations look reasonable, but the evaluation of interpretability is qualitative. Given that one of the stated motivations for GIS is to make the system useful to human practitioners, some form of user study or quantitative interpretability metric would be more convincing than visual inspection.
Finally, the comparison set is solid but not exhaustive. HDBSCAN-based approaches (as used in NVIDIA Climb) are mentioned in the related work but absent from the main comparison table, which is a notable omission given that HDBSCAN handles non-convex cluster shapes better than K-Means and would be a more challenging unsupervised baseline.
Despite these gaps, GEM makes a coherent and technically grounded argument that the geometry of the embedding space should determine the geometry of the clustering method. The combination of principled directional statistics, balance regularisation, and scalable distillation is a practically deployable package, and the predictability results suggest real downstream value beyond marginal accuracy improvements. For anyone working on pretraining data pipelines, this is worth reading carefully.
Full paper: arXiv:2605.26121