Efficiently pre-training language models with mixtures of cluster-oriented, trainability-aware experts

Fang Dong, Jixian Zhou, Hengjie Cao, Mengyi Chen, Yifeng Yang, Ruijun Huang, Mingzhi Dong, Yujiang Wang, Dongsheng Li, David A. Clifton, Robert P. Dick, Qin Lv, Fan Yang, Tun Lu, Ning Gu, Li Shang

Keywords: Mixture of Expert, Language models, Clustering, Trainability

TL;DR: We utilize the feature cluster structure and the singular spectra of feature and Jacobian spaces, and obtained a MoE with high expert-level trainability for language model.

Abstract:

Language models (LMs) are pre-trained on large-scale corpora from diverse data sources, encapsulating knowledge across various domains, with their feature spaces often displaying clustering structures. The mixture of experts (MoEs) approach is commonly used to scale up model learning capabilities to handle such complexities; however, the fine-grained learning dynamics at the expert level remain largely unexplored. This work analyzes the spatial and temporal characteristics of these clustering structures and examines their impact on the fine-grained trainability of individual experts. Our analysis builds on the singular spectrum of the feature and Jacobian spaces leading to two key observations. First, a few top singular vectors from the feature matrix are sufficient to capture the layer-wise feature cluster patterns. More interestingly, the maximum singular value of the Jacobian matrix reveals conflicts between different feature clusters, and experts exhibit varying levels of trainability, completing their learning asynchronously during training.
Inspired by these insights, we proposed mixtures of cluster-guided, trainability-aware experts (MO-CTE), with
an efficient routing method to mitigate inter-cluster conflicts to improve expert trainability and
a simple yet effective criterion for early stopping low-trainability experts, thus reducing total training costs.
We evaluate the proposed MO-CTE across extensive datasets and tasks. Experimental results indicate that MO-CTE accelerates convergence by approximately 37\% in test perplexity and 30\% in downstream tasks, and improves performance by 3.6\% over baselines when consuming similar computation resources.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 5701