ModuleFormer: Modularity Emerges from Mixture-of-Experts

Motivation

• Earlier modular SMoE approaches relied on curated data to learn expert functionality, whereas ModuleFormer targets uncurated pretraining data.Curated data is described as difficult to obtain and scale.Source: §1 · p. 2

Contributions

• ModuleFormer induces modularity from uncurated data using feedforward experts, stick-breaking attention, and specialized losses.The mutual-information loss balances module loads during pretraining, while load concentration supports later module selection and pruning.Source: §1 · p. 2

Efficiency

• Sparse activation lets ModuleFormer match dense-model performance with 50% latency and more than 2× token throughput.The architecture also uses a smaller memory footprint than dense models.Source: §1 · pp. 2–3

Extendability

• Adding modules enables learning new domain knowledge while preserving previous modules, reducing exposure to catastrophic forgetting during adaptation.The paper motivates this design because full-model fine-tuning is costly and can cause forgetting.Source: §1 · pp. 2–3

Specialisation

• Fine-tuning can specialize a subset of modules for a downstream task, allowing unused modules to be pruned without sacrificing performance.This supports lightweight deployment with a configurable performance–model-size trade-off.Source: Figure 1 · p. 2, §1 · p. 3

Pretrained Language Models and finetuning

• LoRA and related parameter-efficient methods reduce adaptation cost, but the paper reports weaker performance when learning a substantially different language.Full-model fine-tuning becomes increasingly difficult as language models grow.Source: §2 · p. 3

Sparse Mixture of Experts for Modularity

• Neural Module Networks require predefined module functionality and expert labels, limiting their practicality beyond task-specific settings.Their modules are dynamically assembled according to the input reasoning structure.Source: §2 · p. 3
• DEMIX uses domain-specific experts but requires curated domain labels, while unknown-domain inference can scale linearly with the number of modules.These constraints motivate modularity induced from domain-agnostic data.Source: §2 · p. 3

Length-extrapolation in Attention

• Stick-breaking attention builds on prior length-extrapolation work while encoding positional information without positional embeddings.Unlike ALiBi, it is described as avoiding attention suppression of distant terms by positional bias.Source: §2 · p. 3

Continual Learning

• Continual-learning methods address forgetting, but duplicated task-specific models complicate inference on domain-agnostic data.The paper presents these approaches as orthogonal to ModuleFormer and potentially combinable with it.Source: §2 · pp. 3–4

Neural Network Pruning

• Prior SMoE load balancing uses every expert, making expert pruning more likely to sacrifice performance.ModuleFormer introduces load concentration to bias adaptation toward fewer modules.Source: §2 · p. 4

3.1 Preliminary: Mixture of Experts

• An SMoE router predicts module probabilities and evaluates only the selected top-k experts, reducing computation during training and inference.Unselected experts have zero gate values and need not be evaluated.Source: §3.1 · p. 4
• A ModuleFormer block replaces the feed-forward layer with SMoE and self-attention with top-k Mixture-of-Attention modules.Each block specifies feedforward module count, top-k routing, and hidden dimension; attention modules are likewise sparsely activated.Source: §3.1 · p. 4

3.2 Stick-breaking Self-Attention head

• Stick-breaking attention computes each token’s attention to previous and current tokens through key–query matching and a stick-breaking distribution.The distribution is efficiently computed with cumulative sums in log-space.Source: §3.2 · pp. 4–5
• By selecting the latest matching token, stick-breaking attention implicitly encodes position and removes the need for explicit positional embeddings or relative biases.The formulation is intended to simplify length extrapolation beyond the pretraining sequence length.Source: §3.2 · pp. 4–5
• Each stick-breaking head uses query, key, value, and output projections, with key and value matrices shared across heads.When attention and embedding dimensions match, identity key and value matrices can remove shared parameters.Source: §3.2 · p. 5

4.1 Load Balancing during Pretraining

• Mutual-information load balancing maximizes marginal module entropy while minimizing conditional entropy, distributing expert usage across batches but concentrating each token on fewer modules.The method contrasts token-module mutual information with task-module alignment and uses the two entropy terms to balance overall load and token-level routing concentration.Source: §4.1 · p. 5

4.2 Load Concentration during Finetuning

• Load concentration during finetuning encourages fewer experts, allowing modules with low task frequency to be pruned after counting their usage on training or validation data.The module frequency f_m is treated as task importance, and experts below a chosen threshold can be removed.Source: §4.2 · p. 6

4.3 Inserting new Modules for Continual Learning

• Continual learning inserts randomly initialized modules while freezing previous modules, largely avoiding forgetting but leaving router misrouting as a residual failure mode.The router freezes A and B, trains only A′, and applies norm regularization because larger A′ magnitudes increase new-expert usage without restricting expert capacity.Source: §4.3 · p. 6

5.1 Language Modeling

• MoLM matches dense models of similar computation while using about 25% active parameters per token, reducing latency to 50% and doubling throughput.MoLM-4B-K2 is comparable to dense models around 1.3B parameters, while MoLM-4B-K4 and MoLM-8B-K2 are comparable to dense models around 2.7B parameters.Source: §5.1 · pp. 7–8, Table 2 · p. 7
• Mutual Information loss outperforms removing load balancing and matches ST-MoE loss, while stick-breaking attention outperforms RoPE-based attention in the ablation study.The study also reports that learnable absolute positional embeddings were unstable, whereas RoPE and stick-breaking attention were more stable.Source: §5.1 · p. 8, Table 4 · p. 8
• ModuleFormer’s expert distributions vary by domain, with similar domains showing smaller KL divergence and MI loss encouraging stronger input-category–expert correlation than ST-MoE balancing.The reported comparison associates MI loss with stronger modularity than the ST-MoE load balancing loss.Source: §5.1 · pp. 8–9, Figure 2 · p. 8

5.2 Learning New Knowledge with New Modules

• In continual joint pre-training, sparse models experience less interference and ModuleFormer performs consistently better than the baseline under efficient tuning.The authors attribute the lower interference mainly to the sparse architecture rather than to having more trainable parameters.Source: §5.2 · p. 9, Table 5 · p. 9
• In continual lifelong pre-training, routing regularization trades plasticity against stability as new-expert usage is restricted, while LoRA falls behind on both dimensions.Increasing the regularization weight improves stability but reduces plasticity; lowering LoRA rank or adding weight decay does not improve either dimension.Source: §5.2 · pp. 9–10, Table 6 · p. 9

5.3 Finetuning and Pruning Modules

• Pruning 40%–50% of parameters preserves HumanEval performance, whereas layer-wise pruning causes a visible performance drop.The fine-tune-then-prune strategy achieves similar performance to prune-then-fine-tune while requiring only one fine-tuning process and allowing dynamic pruning ratios.Source: §5.3 · pp. 10–11, Figure 3 · p. 10
• After fine-tuning with load concentration loss, roughly half of the modules activate at less than 0.3% of the most frequently used experts, enabling domain-specialized pruning.The pruning procedure ranks experts by normalized activation frequency and removes modules below a threshold.Source: §5.3 · pp. 10–11
• Load concentration loss improves pruning performance when 30%–75% of parameters remain, while omitting it performs better above 75% remaining parameters.This trade-off supports flexible deployment across devices with different capacity constraints.Source: Figure 4 · p. 10, §5.3 · p. 11

6 Conclusion and Limitations

• 50% lower latency and more than 2× throughput let MoLM match dense LLM performance with a smaller memory footprint.MoLM also shows reduced susceptibility to catastrophic forgetting, supports new modules for new-language learning, and enables pruning of unused task-specialized modules without sacrificing performance.Source: §6 · p. 11
• ModuleFormer introduces stick-breaking attention heads, mutual-information load balancing, and load-concentration losses for modular pretraining and finetuning.The architecture is paired with feedforward modules and methods for inserting and pruning modules.Source: §6 · p. 11
• MoLM uses more parameters than dense models to achieve comparable performance, reflecting optimization difficulty from discrete SMoE gating.How to learn optimal gating remains an open question.Source: §6 · p. 11

A Stick-breaking Self-attention

• As the context window grows, the stick-breaking attention cumulative mass approaches 1.This follows because each β_j,t lies between 0 and 1, so the residual product diminishes with longer context.Source: A Stick-breaking Self-attention · p. 15

C.1 Dataset Details

• CC-100 [Wenzek et al., 2020, Conneau et al., 2020] collects CommonCrawl texts in 100 languages; joint training combines newly learned languages with English at 4:1 and uses 6 billion tokens.Source: §C.1 · p. 15

C.2 Benchmark Details

• mLAMA evaluates comprehensive factual knowledge through fill-in-the-blank questions and two incorrect candidate answers.The correct sentence is assessed by whether it has the lowest perplexity.Source: §C.2 · p. 15

C.3 Baselines Details

• GPT-Neo is a GPT-2-like dense model pretrained on The Pile, the same corpus used in the experiments.Source: §C.3 · p. 16
• LoRA provides small- and medium-rank parameter-efficient adaptation baselines by controlling the number of trainable parameters.Source: §C.3 · p. 16

C.4 Continual Training Hyperparameters

• Table 8 reports the continual-training hyperparameters.Source: Table 8 · p. 16

C.5 Continual Training Ablation

• With proper routing regularization, a partially trainable router strictly outperforms a fully trainable router.Source: §C.5 · p. 16
• Joint training generally achieves better overall performance than lifelong training, except for some new-language abilities under low regularization.Source: §C.5 · p. 16

D.1 Finetuning Hyperparameters

• Pruning-method comparisons normalize expert frequency either by the maximum frequency within each layer or by total layer activation frequency.Source: Figure 5 · p. 17

• Max and Sum achieve comparable pass@10 and pass@100 results.Source: §D.2 · p. 17
• Max achieves better pass@1 results than Sum.Source: §D.2 · p. 17
• The paper uses Max as its major pruning method because of its pass@1 advantage.Source: §D.2 · p. 17