From Dense to Dynamic: Token-Difficulty Driven MoEfication of Pre-Trained LLMs

TL;DR

DynaMoE is a post-training framework that transforms pre-trained dense LLMs into dynamic, token-difficulty-aware Mixture-of-Experts models, enabling customizable efficiency-accuracy trade-offs with minimal fine-tuning.

Contribution

It introduces a token-difficulty-driven routing mechanism for adapting pre-trained LLMs into dynamic MoE models with low fine-tuning cost.

Findings

Achieves flexible model variants with different accuracy-performance trade-offs.

Uses only 10B tokens for adaptation, significantly less than full training.

Maintains similar accuracy to baseline methods with reduced fine-tuning cost.

Abstract

Training large language models (LLMs) for different inference constraints is computationally expensive, limiting control over efficiency-accuracy trade-offs. Moreover, once trained, these models typically process tokens uniformly, regardless of their complexity, leading to static and inflexible behavior. In this paper, we introduce a post-training optimization framework, DynaMoE, that adapts a pre-trained dense LLM to a token-difficulty-driven Mixture-of-Experts model with minimal fine-tuning cost. This adaptation makes the model dynamic, with sensitivity control to customize the balance between efficiency and accuracy. DynaMoE features a token-difficulty-aware router that predicts the difficulty of tokens and directs them to the appropriate sub-networks or experts, enabling larger experts to handle more complex tokens and smaller experts to process simpler ones. Our experiments demonstrate that DynaMoE can generate a range of adaptive model variants of the existing trained LLM with a single fine-tuning step, utilizing only $10B$ tokens, a minimal cost compared to the base model's training. Each variant offers distinct trade-offs between accuracy and performance. Compared to the baseline post-training optimization framework, Flextron, our method achieves similar aggregated accuracy across downstream tasks, despite using only $\frac{1}{9}\text{th}$ of their fine-tuning cost.