PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zihao Huang, Jundong Zhou, Xingwei Qu, Qiyang Min, Ge Zhang

arXiv:2601.21420v1cs.LG

TL;DR

LLMs process every token uniformly even though sequences differ in predictability and reasoning demands. ConceptMoE dynamically merges semantically similar tokens into concepts and evaluates the resulting architecture under matched compute and parameter budgets. It consistently improves performance across language, vision-language, long-context, and continual-training settings while reducing attention and cache costs.

  • Problem

    LLMs allocate equal computation to all tokens despite differences in semantic weight, leaving a need for adaptive concept-level computation.

  • Method

    ConceptMoE uses a learnable chunk module to merge semantically similar tokens into concepts before compute-intensive modeling, with MoE enabling matched-FLOP and matched-parameter evaluation.

  • Results

    +0.9 points on language pretraining, +2.3 points on long context understanding, and +0.6 points on multimodal benchmarks are reported, with +5.5 points for continual-training conversion with layer looping.

  • Takeaways & Limitations

    ConceptMoE reduces attention computation by up to R^2× and KV cache by up to R×, while reaching 175% prefill and 117% decoding speedups at R = 2.

  • Takeaways & Limitations

    At R = 4, forced merging of semantically different tokens loses critical information; R should be calibrated to dataset characteristics, with R = 1.5 to R = 2 effective for typical pretraining corpora.

Abstract

from arXiv · show

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.

1 Introduction

ConceptMoE replaces uniform token-level processing with adaptive concept-level computation, while enabling controlled comparisons and practical efficiency gains across training settings.

  • Motivation: LLMs allocate equal computation to every token despite large differences in semantic importance and predictability.This can waste computation on routine predictions while under-resourcing semantically dense content.
  • Approach: ConceptMoE merges consecutive semantically similar tokens into concept representations before compute-intensive processing.A learnable chunk module measures inter-token similarity to identify boundaries and preserve finer granularity for distinct tokens.
  • Evaluation: ConceptMoE enables fair comparison with standard MoE by matching total parameters and average per-token FLOPs after reallocating saved computation.The design uses MoE to adjust activated parameters independently of total parameters, isolating architectural benefits.
  • Results: +5.5 points improvement is achieved in continual-training conversion with layer looping, while from-scratch training reaches +6.4 points overall.The framework is evaluated across language pretraining, vision-language training, and pretrained-MoE conversion.
  • Efficiency: 175% prefill speedups and 117% decoding speedups are observed on long sequences at R = 2.ConceptMoE also reduces attention map computation by up to R^2× and KV cache by up to R×.
  • Practical adoption: ConceptMoE requires a lightweight chunk module and minor decoder modifications, supporting integration into existing MoE architectures.The modifications include additional QKV projectors in the last 4 layers and apply to pretraining and continual training.

2 Related Work

Prior approaches compress token sequences through vocabulary expansion, fixed or heuristic merging, and adaptive byte-level chunking, but each leaves important limitations.

  • Vocabulary expansion: A 100× vocabulary expansion produces only 1.3× compression, while massive vocabularies increase training and inference costs.Vocabulary expansion can improve compression and downstream performance, but further scaling is described as impractical.
  • Token-level merging: Fixed-length and rule-based token merging cannot adapt compression to varying information density.Information-sparse tokens should merge more aggressively, whereas information-rich tokens should retain finer granularity.
  • Byte-level models: Byte-level methods require aggressive or sophisticated merging, with approaches differing in adaptivity, end-to-end training, and compression control.BLT uses a pretrained auxiliary model for entropy-based merging, while H-Net reports 9× compression with an end-to-end dynamic chunking module.

3 Approach

ConceptMoE compresses semantically similar token sequences into concept representations before expensive MoE processing, while preserving finer granularity for information-rich tokens. Its chunking, dechunking, and joint-decoding design supports adaptive compute allocation, target compression control, and reduced attention-related costs.

  • Overview: ConceptMoE comprises encoder, chunk, concept-model, dechunk, and decoder modules, with the concept model carrying most computation.The encoder, concept model, and decoder use stacked MoE modules; chunking reduces the sequence entering the concept model.
  • Joint decoding: Joint decoding injects concept information into decoder attention through additional query, key, and value projections for every token.Each token is guaranteed an associated concept, and the added projections have negligible parameters because the decoder is shallow.
  • Chunk module: The chunk module merges consecutive tokens with high semantic similarity and preserves smaller chunks around semantic shifts requiring more careful processing.Learnable query and key projections measure adjacent-token cosine similarity, with boundary probabilities identifying chunk starts.
  • Chunk module: The auxiliary loss constrains the token compression ratio to a target R = N/M ≥1 while allowing sample-level allocation between difficult and easy examples.Aggregated batch statistics let difficult samples use lower compression and easy samples higher compression; λ controls the loss weight.
  • Chunk module: Random boundary perturbations during training mitigate excessive compression caused by evaluation-time distribution shifts.The perturbations simulate evaluation conditions that otherwise can speed inference while substantially degrading performance.
  • Dechunk module: Dechunking maps each concept back to its associated token positions, applies exponential moving averages, and prevents information leakage.The index mappings assign every token a concept, while EMA can eliminate boundaries when neighboring concepts jointly predict subsequent tokens effectively.
  • Compute reallocation: Concept-level compression reduces compute-intensive processing and lowers attention-map computation and KV-cache requirements.At compression ratio R, the concept model processes fewer tokens; the architecture therefore reduces quadratic attention-map computation and cache length.

4 Experiments

Experiments evaluate ConceptMoE across language, vision-language, and controlled compute settings, showing consistent gains while preserving matched total parameters and per-token FLOPs.

  • Experimental setup: All comparisons maintain identical total parameters and per-token FLOPs, excluding attention-map computation, to isolate compute-allocation effects.The evaluation spans language pretraining, vision-language training, and downstream benchmarks.
  • Experimental setup: The training protocol uses TPP = 400, with 243B tokens for the 12B-parameter model and 559B tokens for the 24B-parameter model.Both configurations activate eight experts per token.
  • Main results: ConceptMoE consistently outperforms standard MoE across most evaluated metrics under matched compute and parameter budgets.Both model pairs differ in whether computation is allocated at the concept or token level.
  • Vision-language training: Vision-language experiments compress both visual and textual tokens, using R = 2 to preserve total activated compute.The setup uses a 60B-parameter MoE configuration and applies the third compute-reallocation strategy.
  • Vision-language training: At R = 2, multimodal training yields lower loss for ConceptMoE by 0.01 during PT, 0.017 on text data during CT, and 0.012 on image-text data.The model compresses text less and images more while maintaining overall R = 2.
  • Downstream evaluation: ConceptMoE outperforms MoE by 0.9 points on text benchmarks, 2.3 on long-context tasks, and 0.6 on multimodal tasks.The long-context gains include improved Needle-task performance, while some fine-grained visual tasks decline slightly.

4.3 Train from CT

Continual-training conversion from MoE to ConceptMoE preserves baseline performance without increasing matched compute, while layer looping produces substantial downstream gains and inference speedups.

  • Post-SFT results: CT conversion preserves baseline performance for ConceptMoE-top15, while layer looping increases Open Benchmark performance by 5.5 points.ConceptMoE-top15 improves by 0.4 points, whereas ConceptMoE-top11-loop8 provides the larger gain.
  • Inference efficiency: At R = 1.5, converted ConceptMoE enables up to 43.6% prefill acceleration and 53.3% decoding acceleration.The stabilized continual-training compression ratio supports direct inference speedup.
  • Conversion procedure: ConceptMoE conversion adds a chunk module, a dechunk module, and zero-initialized QKV projectors in the last four self-attention layers.The conversion starts from a pretrained MoE and then applies continual training with long-context stages.
  • Compute allocation: The converted configurations keep average FLOPs and total parameters nearly identical through expert activation or intermediate-layer looping.The compared schemes increase activated experts from 8 to 15 or from 8 to 11 while looping 8 layers.
  • Compression and distribution shift: During PT, evaluation uses R = 1.81 and reduces FLOPs by 12.5%, whereas CT stabilizes at R = 1.5 because evaluation and CT distributions align more closely.The differing ratios indicate distribution shift between PT data and later evaluation or CT data.
  • Post-SFT results: ConceptMoE-top11-loop8 gains 5.5 points overall after SFT, including +8.3 in reasoning, +12.2 in math, and +6.4 in code.The gains occur without increasing FLOPs; training from scratch adds a further 0.9 points.

4.4 Significant inference speedup

ConceptMoE delivers substantial inference speedups over MoE on long sequences, including quality-oriented models with doubled depth, by exploiting compressed attention and KV-cache processing.

  • Latency results: 117.1% decoding speedups are achieved by efficiency-oriented ConceptMoE configurations, while prefill speedups reach 24.7%.The evaluation uses MoE-A10B-300B as baseline and compares five configurations on Hopper GPUs.
  • Controlled configurations: ConceptMoE-2L-top8-R2 matches MoE in FLOPs and total parameters while doubling the number of layers.The configuration notation specifies layer multiplication, activated experts, and compression ratio.
  • Latency results: 175% prefill speedups are measured across sequence lengths from 4K to 1024K, with decoding evaluated for KV-cache lengths from 4K to 64K.Decoding measurements use batch size 256.
  • Mechanism: The speedup advantage grows on long sequences because R = 2 reduces attention-map computation quadratically and KV-cache processing linearly.This benefit persists even when ConceptMoE doubles the number of layers.

4.5 Structure ablation

Ablations show that adaptive chunking, cosine-based routing, auxiliary-loss tuning, and joint decoding each affect downstream quality differently from training loss alone.

  • Auxiliary loss: λ = 0.03 is selected because larger auxiliary-loss weights degrade training loss while all tested settings remain close to the target compression ratio of 2.The tested values are 0.03, 0.1, 0.5, and 1.0.
  • Router type: 36.4 average downstream performance for the cosine router substantially exceeds the linear router’s 34.4 despite the linear router’s 0.003 lower training loss.The comparison uses ConceptMoE-A0.5B-12B trained for 243B tokens at R = 2.
  • Joint decoding: Joint decoding raises downstream performance from 35.1 to 36.4 average while slightly increasing converged training loss by 0.002.The component adds negligible computation and incorporates concept information through additional QKV projections.

4.6 Boundary Noise for Robustness

Boundary noise improves agreement between training and evaluation compression and can improve downstream scores, while experiments show that excessive target compression harms reasoning and math performance.

  • Boundary robustness: Noise regularization moves the boundary-probability mean toward 1/R, including 0.667 for R = 1.5, stabilizing evaluation compression ratios.Without noise, probabilities cluster near 0.5 and evaluation perturbations can increase compression beyond the intended ratio.
  • Boundary noise: 30.3 average score is achieved with τ = 4, a +1.4 improvement over baseline despite higher training loss.Bernoulli noise is more effective than the alternatives, and τ = 6 is selected for later experiments to limit training-loss impact.
  • Target compression ratio: R = 2 reaches 50.8 average downstream performance, whereas R = 4 reaches 47.7 and shows a 0.013 training-loss gap at convergence.The experiments use MoE-A1B-24B as baseline with compute reallocation strategy 3.
  • Target compression ratio: At R = 4, reasoning scores 51.3 versus 56.8 and math scores 46.3 versus 50.0, showing particularly large losses on complex tasks.The results indicate that aggressive compression can disrupt complex reasoning patterns.
  • Target compression ratio: For typical pretraining corpora, R = 1.5 to R = 2 is reported as an effective balance between efficiency and information preservation.The paper attributes the boundary to dataset semantic redundancy rather than unconditional maximization of compression.

5 Conclusion

ConceptMoE combines learnable semantic chunking with compute-matched MoE evaluation to improve performance across language, vision-language, long-context, and continual-training settings while reducing inference costs.

  • Conclusion: ConceptMoE dynamically merges semantically similar tokens into concepts through learnable chunking.The approach preserves finer granularity for semantically distinct content.
  • Conclusion: +0.9 points on language pretraining, +0.6 points on vision-language training, and +2.3 points on long-context understanding are reported under compute-matched comparisons.The MoE architecture reallocates saved computation while matching baseline FLOPs and total parameters.
  • Conclusion: +5.5 points with layer looping are achieved in continual-training conversion.This result is reported alongside improvements across the other evaluated settings.
  • Conclusion: Attention computation falls by up to R^2× and KV cache by up to R×, with R = 2 prefill speedups up to 175% and decoding speedups up to 117%.The paper presents these reductions as efficiency gains beyond benchmark performance.
  • Conclusion: Minimal architectural changes support straightforward integration into existing MoE systems.The conclusion frames adaptive concept-level processing as improving both effectiveness and efficiency.

A Evaluation benchmark

The evaluation materials organize OpenBench datasets by difficulty and model scale across different domains. Tables 8 and 9 provide the corresponding benchmark listings.

  • Models with 12B and 24B total parameters use relatively easy evaluation sets.
  • Models with 60B and 90B total parameters are evaluated on both easy and hard sets.
  • OpenBench evaluation sets are organized across different domains.

B Code

The code implements ConceptMoE with an encoder, a compressed concept model, and a decoder. It predicts chunk boundaries from adjacent-token similarity, merges selected concepts, and restores them to the original sequence for joint decoding.

  • A 1024-token input is compressed to 701 tokens before entering the concept model.The hidden size in the example is 2048.
  • The chunk module computes boundary probabilities from cosine similarity between adjacent projected hidden states.Queries and keys are normalized after separate linear projections, and probabilities are clamped to [0,1].
  • Boundaries are selected when boundary probability is at least 0.5, with the first element forced to be a boundary.
  • The dechunk module recursively merges concept representations using boundary probabilities and restores them from 701 to 1024 positions.The merge is written as a loop for clarity and can be accelerated through parallel scan.
  • The decoder adds the restored concept representation to the hidden state and performs joint decoding.The code applies a straight-through estimator to selected probabilities before the decoder.
  • ConceptMoE uses a 2-layer encoder, 23-layer concept model, and 2-layer decoder.The implementation defines transformer layers with indices 0–1, 2–24, and 25–26.
Loading 2601.21420v1…