ConceptMoE: Adaptive Compute in MoE Models

• ConceptMoE is a framework for adaptive token compression that merges semantically similar tokens into concept embeddings for efficient computation.
• It leverages a dynamic chunking module to reduce attention FLOPs by up to R² and minimizes the KV cache footprint with minimal architectural changes.
• Empirical results demonstrate improved efficiency and performance across language, vision-language, and continual training tasks through adaptive compute reallocation.

ConceptMoE is a framework for adaptive and efficient allocation of computation in large-scale neural networks, notably within the Mixture-of-Experts (MoE) paradigm. Its central innovation is the dynamic merging (chunking) of semantically similar tokens into concept representations, enabling implicit token-level compute allocation and concept-driven model capacity scaling. This approach distinguishes meaningful, information-rich spans from trivially predictable tokens, allowing for selective and context-dependent compression before entering compute-intensive components. ConceptMoE substantially improves both efficiency (reduced FLOPs and memory footprint) and task performance in language, vision-language, and continual training regimes, while requiring minimal architectural modifications to existing MoE models (Huang et al., 29 Jan 2026).

1. Adaptive Token-to-Concept Compression: Motivation and Foundation

Modern LLMs process every token in a sequence with uniform computational cost across all self-attention and feed-forward layers, regardless of each token's informativeness. This uniformity leads to inefficient resource use where semantically trivial tokens (e.g., function words, punctuation) consume as much compute as those requiring deeper reasoning. ConceptMoE counters this by leveraging a learnable chunking module to merge contiguous, semantically similar tokens into single concept embeddings, compressing the token stream prior to heavy MoE layers. Concept tokens are positioned as high-level representations for regions demanding complex computations, while uninformative spans are compressed to minimize computational overhead (Huang et al., 29 Jan 2026).

2. Compression Ratio, Boundary Detection, and Theoretical Savings

ConceptMoE introduces the compression ratio $R$, defined as: $R = \frac{N}{M} \geq 1$ where $N$ is the input sequence length and $M$ the output concept token count post-chunking. The attention map computational cost and the Key-Value (KV) cache requirements scale as follows:

Mode	Baseline (per-token)	ConceptMoE, compression $R$
Attention FLOPs	$L_{\mathcal C} d N^2$	$L_{\mathcal C} d (N/R)^2$ ($\times \frac{1}{R^2}$)
KV cache	$2 L_{\mathcal C} d N$	$2 L_{\mathcal C} d (N/R)$ ($R = \frac{N}{M} \geq 1$0)

This results in up to $R = \frac{N}{M} \geq 1$1-times reduction in attention computation and $R = \frac{N}{M} \geq 1$2-times reduction in memory for the KV cache. Empirical measurements document prefill speedups up to 175% and decoding speedups up to 117% on long sequences at $R = \frac{N}{M} \geq 1$3 (Huang et al., 29 Jan 2026).

Boundary detection is implemented via learnable cosine similarity metrics between projected token representations. A token boundary for starting a new concept is chosen where the normalized distance $R = \frac{N}{M} \geq 1$4 between consecutive tokens exceeds 0.5: $R = \frac{N}{M} \geq 1$5 An auxiliary loss function $R = \frac{N}{M} \geq 1$6 regularizes the chunk module to achieve the target ratio $R = \frac{N}{M} \geq 1$7, weighted by a hyperparameter $R = \frac{N}{M} \geq 1$8 (typically 0.03).

3. Chunk Module, Joint Decoding, and Token-to-Concept Mapping

The ChunkModule comprises two learnable projections $R = \frac{N}{M} \geq 1$9 for token similarity assessment, a Bernoulli sampler for stochastic boundary creation, and an auxiliary loss for ratio targeting. Consecutive tokens within chunked boundaries are aggregated to concept embeddings $N$0. Post-processing applies exponential moving average (EMA) smoothing: $N$1 During dechunking, each original token $N$2 is mapped to its parent concept index $N$3, and the final token embedding is reconstructed as: $N$4 This joint decoding ensures all tokens benefit from deep concept-level computations derived from information-rich sub-sequences (Huang et al., 29 Jan 2026).

4. MoE Model Integration and Compute Reallocation Strategies

The ConceptMoE architecture is realized as: $N$5 where $N$6 and $N$7 are shallow encoder and decoder stacks, and $N$8 denotes a deep MoE-driven concept model. To preserve direct baseline comparison, the compute saved via chunking ($N$9) is redistributed in three ways that maintain total parameter count:

1. Increase expert count in MoE layers ($M$0).
2. Loop additional layers through $M$1 without expanding parameters.
3. Scale widths of self-attention and MoE hidden layers.

These compute reallocation strategies isolate architectural benefits of concept-driven MoE versus naïve token-level MoE for comprehensive benchmarking.

5. Empirical Results and Ablation Studies

ConceptMoE demonstrates robust improvements on multiple scales:

• Language Pretraining: +0.9 points over standard MoE on open-bench tasks.
• Long Context Handling: +2.3 points via adaptive chunking and joint decoding.
• Multimodal Benchmarks: +0.6 points using vision-language token chunking.
• Continual Training Conversion: +5.5 point gains via layer looping in the concept model for large (90B parameter) baselines.

Ablation studies identify critical components: dynamic chunking is +1.2 points over fixed merging, cosine similarity routers generalize better than linear, and joint decoding adds +1.3 points. Boundary sharpening parameter $M$2 yields optimal performance, and compression ratio $M$3 strikes the best tradeoff between speed and reasoning accuracy. At higher ratios ($M$4) performance degrades on mathematical and reasoning tasks (Huang et al., 29 Jan 2026).

6. Implementation and Integration Considerations

ConceptMoE requires minimal architectural changes. The chunk module itself consists of two linear layers and a Bernoulli sampler. Dechunking is a straightforward index remapping plus EMA smoothing. Four additional QKV projectors are inserted in decoder layers for continual training; these are initialized to zero. The core MoE layers remain unchanged, leveraging the concept-level representations as input, and compute can be dynamically reallocated as needed. Appendix A.2 in the paper provides a detailed PyTorch-style code sketch for integration (Huang et al., 29 Jan 2026).

7. ConceptMoE in Structured Graph Bottleneck Models

MoE-SGT extends the ConceptMoE paradigm to multimodal concept bottleneck modeling, integrating graph-based reasoning modules and explicit concept alignment. Here, human-interpretable concepts $M$5, fused from vision-language priors, propagate through answer–concept and answer–question heterogeneous graphs. A structure-injecting Graph Transformer (SGT) equipped with MoE layers models multi-level dependencies, dynamically allocating reasoning capacity per concept via expert routing (Lin et al., 1 Oct 2025). This framework achieves improved accuracy across single- and multi-label object recognition benchmarks, chest X-ray pathology detection, and vision-language QA tasks.

Model	Best Single-label Top-1 (%)	Multi-label ROC-AUC/F1
MoE-SGT	79.76 (CUB-200), 73.41 (ImageNet)	0.76/0.80 (MIMIC-CXR), 0.83/0.86 (CheXpert)
Sparse-CBM	80.27 (CUB-200), 71.62 (ImageNet)	0.72/0.73 (MIMIC-CXR), 0.80/0.81 (CheXpert)

This suggests that ConceptMoE, in its graph-integrated variant, excels when modeling structured inter-concept dependencies and dynamic cross-modal reasoning demands. A plausible implication is improved interpretability via explicit concept bottleneck alignment, while maintaining or exceeding standard model accuracy (Lin et al., 1 Oct 2025).

8. Limitations, Extensions, and Future Directions

ConceptMoE's compression mechanism relies on a static concept pool at inference, which may constrain flexibility in rapidly evolving or open-domain settings. Streaming LLMs for on-the-fly concept expansion, adaptive boundary detection algorithms, and more sophisticated MoE routing (e.g., top-$M$6 gating or load-balancing losses) are identified as promising future extensions. Generalization beyond annotated medical and vision-language benchmarks will likely require systematic retuning of graph architectures, expert allocation schemes, and compression strategies.

ConceptMoE represents an overview of adaptive concept-level processing and implicit compute allocation, advancing efficient, scalable, and interpretable modeling in both unimodal and multimodal domains (Huang et al., 29 Jan 2026, Lin et al., 1 Oct 2025).