ConceptMoE: Adaptive Token-to-Concept Compression for Implicit Compute Allocation
This addresses the efficiency and effectiveness limitations of large language models for AI practitioners by enabling adaptive computation allocation with minimal architectural modifications.
The paper tackles the problem of uniform computation allocation in large language models by introducing ConceptMoE, which dynamically compresses tokens into concept representations to perform implicit token-level compute allocation. The method achieves performance improvements of +0.9 to +5.5 points across various tasks while reducing attention computation by up to R²× and KV cache by R×, with measured speedups of up to 175% in prefill and 117% in decoding at R=2.
Large language models allocate uniform computation across all tokens, ignoring that some sequences are trivially predictable while others require deep reasoning. We introduce ConceptMoE, which dynamically merges semantically similar tokens into concept representations, performing implicit token-level compute allocation. A learnable chunk module identifies optimal boundaries by measuring inter-token similarity, compressing sequences by a target ratio $R$ before they enter the compute-intensive concept model. Crucially, the MoE architecture enables controlled evaluation: we reallocate saved computation to match baseline activated FLOPs (excluding attention map computation) and total parameters, isolating genuine architectural benefits. Under these conditions, ConceptMoE consistently outperforms standard MoE across language and vision-language tasks, achieving +0.9 points on language pretraining, +2.3 points on long context understanding, and +0.6 points on multimodal benchmarks. When converting pretrained MoE during continual training with layer looping, gains reach +5.5 points, demonstrating practical applicability. Beyond performance, ConceptMoE reduces attention computation by up to $R^2\times$ and KV cache by $R\times$. At $R=2$, empirical measurements show prefill speedups reaching 175\% and decoding speedups up to 117\% on long sequences. The minimal architectural modifications enable straightforward integration into existing MoE, demonstrating that adaptive concept-level processing fundamentally improves both effectiveness and efficiency of large language models.