Elamur Research⌕
☰
home / papers / 2026 / UniGeM: Unifying Data Selection and Mixing via Geometric Exploration and Mining

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

Paper · 2026 · Changhao Wang, Yunfeiyu, Xinhao Yao et al.

Summary

Data quality now limits LLM pre-training more than data volume. Domain mixing and sample selection are usually handled as separate steps, which can damage the logical and hierarchical structure of code data. UniGeM combines both steps in one geometric pipeline. It needs no proxy model and no external reference dataset; curation is treated as approximating the manifold of the corpus's embeddings.

In Stage-I (Macro-Exploration), the corpus is clustered. The number of clusters is set by checking how stable cluster scores remain across resolutions, and a softmax over geometric cluster scores assigns each cluster its share of the sampling budget. In Stage-II (Micro-Mining), each cluster is split into sub-clusters, which are scored by an LLM probe, a structural penalty, and a cohesion gate.

The authors trained 8B and 16B MoE models from scratch on a 100B-token code-and-text mix. UniGeM reached about twice the data efficiency of random sampling and outperformed adapted SOTA baselines, especially on reasoning-heavy and multilingual code benchmarks. This points to a cheaper, structure-aware way to curate pre-training data.

Key points

  • Stage-I cluster features. Each cluster gets four geometric proxies: cohesion, size, average sequence length, and entropy of language identifiers.
    • These are combined into one score, with weights set by spectral consensus.
    • The cluster count K is chosen by maximizing Kendall's τ rank stability between neighbouring resolutions. The result was K* = 72.
  • Stage-II sub-cluster weighting. Each sub-cluster's weight multiplies four terms:
    • the budget inherited from its parent cluster;
    • an LLM-derived semantic score;
    • a rectified Mahalanobis structural penalty on length and entropy;
    • a sigmoid cohesion gate that compares the sub-cluster with its parent.
  • Models. Both the 8B and 16B MoE models have 1.4B active parameters and were trained on 100B tokens.
    • UniGeM achieved 2.0× data efficiency over random sampling.
    • UniGeM-16B scored 39.5 after one epoch, 6.6 points above Random.
  • Limitations stated by the authors.
    • Evaluation was mostly on code; general web-text mixtures are untested.
    • Embedding and clustering costs become substantial at trillion-token scale.
    • The pipeline runs only as a static pre-processing step.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Renmin University of ChinaBeijing · China114 works · 16,118 citations🎓University of Chinese Academy of SciencesBeijing · China147 works · 8,480 citations🏢Ant GroupHangzhou · China42 works · 1,131 citations🎓Politecnico di TorinoTurin · Italy11 works · 18 citations

Also on the byline: Institute of Microelectronics.

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksInformation Systems · 1,437 works

Cite

Cite as Changhao Wang, Yunfeiyu, Xinhao Yao, Jiaolong Yang, Lu Yu, Junpeng Fang, Chaobo Li, Riccardo Cantoro, Qing Cui and Jun Zhou. 2026. UniGeM: Unifying Data Selection and Mixing via Geometric Exploration and Mining. Findings of the Association for Computational Linguistics: ACL 2026.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.findings-acl.1037

@inproceedings{wang2026unigem,
  title = {{UniGeM: Unifying Data Selection and Mixing via Geometric Exploration and Mining}},
  author = {Changhao Wang, Yunfeiyu, Xinhao Yao, Jiaolong Yang, Lu Yu, Junpeng Fang, Chaobo Li, Riccardo Cantoro, Qing Cui and Jun Zhou},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-acl.1037},
  url = {https://doi.org/10.18653/v1/2026.findings-acl.1037}
}

References

40 references, 12 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov. 2024. Meta flow matching: Integrating vector fields on the wasserstein manifold. arXiv preprint arXiv:2408.14608. — [arXiv:2408.14608](https://arxiv.org/abs/2408.14608) find
  2. Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, and 1 others. 2021. Program synthesis with large language models. arXiv preprint arXiv:2108.07732.
  3. Federico Cassano, John Gouwar, Daniel Nguyen, Sydney Nguyen, Luna Phipps-Costin, Donald Pinckney, Mark Chen. 2021. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374.
  4. Mayee Chen, Nicholas Roberts, Kush Bhatia, Jue Wang, Ce Zhang, Frederic Sala, and Christopher Ré. 2023. Skill-it! a data-driven skills framework for understanding and training language models. Advances in Neural Information Processing Systems, 36:36000– 36040. find
  5. Shizhe Diao, Yu Yang, Yonggan Fu, Xin Dong, Dan Su, Markus Kliegl, Zijia Chen, Peter Belcak, Yoshi Suhara, Hongxu Yin, and 1 others. 2025. Climb: Clustering-based iterative data mixture bootstrapping for language model pre-training. arXiv preprint arXiv:2504.13161. — [arXiv:2504.13161](https://arxiv.org/abs/2504.13161) find
  6. Yilun Du, Katie Collins, Josh Tenenbaum, and Vincent Sitzmann. 2021. Learning signal-agnostic manifolds of neural fields. Advances in Neural Information Processing Systems, 34:8320–8331. find
  7. Z Feng. 2020. Codebert: A pre-trained model for program-ming and natural languages. arXiv preprint arXiv:2002.08155.
  8. Samir Yitzhak Gadre, Georgios Smyrnis, Vaishaal Shankar, Suchin Gururangan, Mitchell Wortsman, Rulin Shao, Jean Mercat, Alex Fang, Jeffrey Li, Sedrick Keh, and 1 others. 2024. Language models scale reliably with over-training and on downstream tasks. arXiv preprint arXiv:2403.08540.
  9. Robert M. Gray and David L. Neuhoff. 2002. Quantization. IEEE transactions on information theory, 44(6):2325–2383. find
  10. Alex Gu, Baptiste Rozière, Hugh Leather, Armando Solar-Lezama, Gabriel Synnaeve, and Sida I Wang. 2024. Cruxeval: A benchmark for code reasoning, understanding and execution. arXiv preprint arXiv:2401.03065.
  11. Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, and 1 others. 2023. Textbooks are all you need. arXiv preprint arXiv:2306.11644. find
  12. Daya Guo, Shuo Ren, Shuai Lu, Zhangyin Feng, Duyu Tang, Shujie Liu, Long Zhou, Nan Duan, Alexey Svyatkovskiy, Shengyu Fu, and 1 others. 2020. Graphcodebert: Pre-training code representations with data flow. arXiv preprint arXiv:2009.08366.
  13. Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, and 1 others. 2022. Training compute-optimal large language models. Advances in Neural Information Processing Systems, 35:30016–…
  14. Jing Huang, Diyi Yang, and Christopher Potts. 2024. Demystifying verbatim memorization in large language models. arXiv preprint arXiv:2407.17817. find
  15. Naman Jain, King Han, Alex Gu, Wen-Ding Li, Fanjia Yan, Tianjun Zhang, Sida Wang, Armando Solar-Lezama, Koushik Sen, and Ion Stoica. 2024. Livecodebench: Holistic and contamination free evaluation of large language models for code. arXiv preprint arXiv:2403.07974.
  16. Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361.
  17. Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferrandis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, and 1 others. 2022. The stack: 3 tb of permissively licensed source code. arXiv preprint arXiv:2211.15533. find
  18. Michel Ledoux. 2001. The concentration of measure phenomenon. 89. American Mathematical Soc. find
  19. Jiazheng Li, Lu Yu, Qing Cui, Zhiqiang Zhang, Jun Zhou, Yanfang Ye, and Chuxu Zhang. 2025. MASS: Mathematical data selection via skill graphs for pretraining large language models. arXiv preprint arXiv:2503.14917. — [arXiv:2503.14917](https://arxiv.org/abs/2503.14917) find
  20. Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Marc Marone, Christopher Akiki, Jia Li, Jenny Chim, and 1 others. 2023. Starcoder: may the source be with you! arXiv preprint arXiv:2305.06161. find
  21. Ruihang Li, Yixuan Wei, Miaosen Zhang, Nenghai Yu, Han Hu, and Houwen Peng. 2024. Scalingfilter: Assessing data quality through inverse utilization of scaling laws. arXiv preprint arXiv:2408.08310. find
  22. Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, and 1 others. 2024. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434. find
  23. Fengze Liu, Weidong Zhou, Binbin Liu, Zhimiao Yu, Yifan Zhang, Haobin Lin, Yifeng Yu, Bingni Zhang, Xiaohuan Zhou, Taifeng Wang, and 1 others. 2025. Quadmix: Quality-diversity balanced data selection for efficient llm pretraining. arXiv preprint arXiv:2504.16511. — [arXiv:2504.16511](https://arxiv.org/abs/2504.16511) find
  24. Sören Mindermann, Jan M Brauner, Muhammed Razzak, Mrinank Sharma, Andreas Kirsch, Winnie Xu, Benedikt Höltgen, Aidan N Gomez, Adrien Morisot, Sebastian Farquhar, and 1 others. 2022. Prioritized training on points that are learnable, worth learning, and not yet learnt (rho-loss). In International Con… find
  25. Vardan Papyan, XY Han, and David L Donoho. 2020. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40):24652– 24663. find
  26. Gabriel Peyré, Marco Cuturi, and 1 others. 2019. Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning, 11(5-6):355–607. find
  27. Phil Pope, Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. 2021. The intrinsic dimension of images and its impact on learning. In International Conference on Learning Representations (ICLR). find
  28. Roland Schäfer. 2017. Accurate and efficient general-purpose boilerplate detection for crawled web corpora. Lang. Resour. Eval., 51(3):873–889. find
  29. ByteDance Seed, Yuyu Zhang, Jing Su, Yifan Sun, Chenguang Xi, Xia Xiao, Shen Zheng, Anxiang Zhang, Kaibo Liu, Daoguang Zan, and 1 others. 2025. Seed-coder: Let the code model curate data for itself. arXiv preprint arXiv:2506.03524.
  30. Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Morcos Ariyo. 2022. Beyond neural scaling laws: beating power law scaling via data pruning. Advances in Neural Information Processing Systems, 35:19523–19536.
  31. Hong Ye Tan and Emma Slade. 2025. Dataset distillation as pushforward optimal quantization. arXiv preprint arXiv:2501.07681. — [arXiv:2501.07681](https://arxiv.org/abs/2501.07681) find
  32. Pablo Villalobos and 1 others. 2022. Will we run out of data? an analysis of the limits of scaling datasets in machine learning. arXiv preprint arXiv:2211.04325. find
  33. Cédric Villani. 2009. Optimal transport: old and new, volume 338. Springer. find
  34. Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V Le, Tengyu Ma, and Adams Wei Yu. 2023a. Data selection for language models via importance resampling. arXiv preprint arXiv:2302.03169. find
  35. Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang, Quoc V Le, Tengyu Ma, and Adams Wei Yu. 2023b. Doremi: Optimizing data mixtures speeds up language model pretraining. In NeurIPS. find
  36. Zichun Yu, Spandan Das, and Chenyan Xiong. 2024. Mates: Model-aware data selection for efficient pre-training with data influence models. Advances in Neural Information Processing Systems, 37:108735– 108759. find
  37. Paul L Zador. 1982. Asymptotic quantization error of continuous signals and the quantization dimension. IEEE Transactions on Information Theory, 28(2):139–149. find
  38. Chi Zhang, Huaping Zhong, Kuan Zhang, Chengliang Chai, Rui Wang, Xinlin Zhuang, Tianyi Bai, Qiu Jiantao, Lei Cao, Ju Fan, Ye Yuan, Guoren Wang, and Conghui He. 2025a. Harnessing diversity for important data selection in pretraining large language models. In The Thirteenth International Conference on… find
  39. Yanzhao Zhang, Mingxin Li, Dingkun Long, Xin Zhang, Huan Lin, Baosong Yang, Pengjun Xie, An Yang, Dayiheng Liu, Junyang Lin, and 1 others. 2025b. Qwen3 embedding: Advancing text embedding and reranking through foundation models. arXiv preprint arXiv:2506.05176.
  40. Xinlin Zhuang, Jiahui Peng, Ren Ma, Yinfan Wang, Tianyi Bai, Xingjian Wei, Qiu Jiantao, Chi Zhang, Ying Qian, and Conghui He. 2025. Meta-rater: A multi-dimensional data selection method for pre-training language models. In Proceedings of the 63rd Annual Meeting of the Association for Computational L… find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/wang-2026-unigem-unifying-data-selection-mixing/