ConceptMoE: Adaptive Token-to-Concept Compression for Implicit Compute Allocation

TL;DR

ConceptMoE tackles the inefficiency of uniform token-level computation in large language models by introducing adaptive token-to-concept compression. A learnable chunk module merges semantically similar tokens into concepts, processed by a compute-heavy concept module, with a dechunking path and joint decoding to preserve information flow. The framework enables fair comparisons by reallocating saved compute to match baseline FLOPs and total parameters, demonstrating consistent improvements across language, long-context, and vision-language tasks, plus substantial inference speedups (e.g., up to $175\%$ prefill and $117\%$ decoding at $R=2$). These results show that adaptive concept-level processing can improve both effectiveness and efficiency of large models, with straightforward integration and practical applicability in continual training settings.

Abstract

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.

Figures (11)

• Figure 1: Overview of ConceptMoE, with details of chunk and dechunk modules.
• Figure 2: Training dynamics of loss and compression ratio. (a) Loss difference between ConceptMoE and MoE during language model pretraining (PT). (b) Loss difference during multimodal continue training(CT), separated into image-text data (mmloss) and text-only data(textloss). (c) Compression ratio evolution for image tokens and text tokens during multimodal training.
• Figure 3: From MoE to ConceptMoE illustration. The blue blocks denote the original MoE components. We add a chunk module and a dechunk module. In addition, in the last four self-attention layers, we insert an extra QKV projector initialized to zeros to enable joint decoding of concepts and tokens.
• Figure 4: Evolution of evaluation metrics across three benchmarks throughout the training process.
• Figure 5: Inference latency speedup over MoE for prefill and decoding. The prefill plot uses sequence length on the x axis, and the decoding plot uses KV cache length on the x axis with batch size 256. The y axis reports speedup relative to MoE in percent. ConceptMoE-$x$L-top$y$-R$z$ matches MoE in FLOPs and total parameters, where $x$ is the layer multiplier, $y$ is the number of activated experts per MoE block with baseline 8 and increased to match FLOPs, and $z$ is the compression ratio.