mlsanityData to 7 Oct 2026

Mixture-of-Modules: Reinventing Transformers as Dynamic Assemblies of Modules

arXiv 2407.06677 ↗First posted 9 Jul 2024cs.CLcited by 11

Not on the map. Why?

Abstract (arXiv)

Is it always necessary to compute tokens from shallow to deep layers in Transformers? The continued success of vanilla Transformers and their variants suggests an undoubted "yes". In this work, however, we attempt to break the depth-ordered convention by proposing a novel architecture dubbed mixture-of-modules (MoM), which is motivated by an intuition that any layer, regardless of its position, can be used to compute a token as long as it possesses the needed processing capabilities. The construction of MoM starts from a finite set of modules defined by multi-head attention and feed-forward networks, each distinguished by its unique parameterization. Two routers then iteratively select attention modules and feed-forward modules from the set to process a token. The selection dynamically expands the computation graph in the forward pass of the token, culminating in an assembly of modules. We show that MoM provides not only a unified framework for Transformers and their numerous variants but also a flexible and learnable approach for reducing redundancy in Transformer parameterization. We pre-train various MoMs using OpenWebText. Empirical results demonstrate that MoMs, of different parameter counts, consistently outperform vanilla transformers on both GLUE and XSUM benchmarks. More interestingly, with a fixed parameter budget, MoM-large enables an over 38% increase in depth for computation graphs compared to GPT-2-large, resulting in absolute gains of 1.4 on GLUE and 1 on XSUM. On the other hand, MoM-large also enables an over 60% reduction in depth while involving more modules per layer, yielding a 16% reduction in TFLOPs and a 43% decrease in memory usage compared to GPT-2-large, while maintaining comparable performance.

Key facts. Cited by 11 papers in the map (data through 7 Oct 2026). Builds on 39 of its 52 references that are in the map. First posted on arXiv 9 Jul 2024.

Builds on

39 of 52 references.

References · 39

  1. 2017Attention Is All You Needcited by 59,932
  2. 2022Scalable Diffusion Models with Transformerscited by 5,931
  3. 2018GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understandingcited by 4,854
  4. 2014Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translationcited by 4,016
  5. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
  6. 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,014
  7. 2019ALBERT: A Lite BERT for Self-supervised Learning of Language Representationscited by 2,772
  8. 2023Sparks of Artificial General Intelligence: Early experiments with GPT-4cited by 2,579
  9. 2013Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computationcited by 2,031
  10. 2022Large Language Models Encode Clinical Knowledgecited by 1,807
  11. 2019SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systemscited by 1,718
  12. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,601
27 more
  1. 2018Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarizationcited by 1,308
  2. 2020Transformer Feed-Forward Layers Are Key-Value Memoriescited by 1,307
  3. 2022Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 smallcited by 1,130
  4. 2024DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelcited by 1,093
  5. 2022In-context Learning and Induction Headscited by 930
  6. 2019Fast Transformer Decoding: One Write-Head is All You Needcited by 700
  7. 2018Universal Transformerscited by 613
  8. 2016Adaptive Computation Time for Recurrent Neural Networkscited by 479
  9. 2017SkipNet: Learning Dynamic Routing in Convolutional Networkscited by 328
  10. 2020DeeBERT: Dynamic Early Exiting for Accelerating BERT Inferencecited by 322
  11. 2022Confident Adaptive Language Modelingcited by 297
  12. 2020BERT Loses Patience: Fast and Robust Inference with Early Exitcited by 267
  13. 2015Highway Networkscited by 257
  14. 2024Mixture-of-Depths: Dynamically allocating compute in transformer-based language modelscited by 242
  15. 2019Depth-Adaptive Transformercited by 227
  16. 2023The Hydra Effect: Emergent Self-repair in Language Model Computationscited by 145
  17. 2024LaCo: Large Language Model Pruning via Layer Collapsecited by 124
  18. 2016Spatially Adaptive Computation Time for Residual Networkscited by 118
  19. 2023SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inferencecited by 89
  20. 2023CoLT5: Faster Long-Range Transformers with Conditional Computationcited by 84
  21. 2023Conditional Adapters: Parameter-efficient Transfer Learning with Fast Inferencecited by 69
  22. 2013Low-Rank Approximations for Conditional Feedforward Computation in Deep Neural Networkscited by 44
  23. 2024Interpreting Key Mechanisms of Factual Recall in Transformer-Based Language Modelscited by 30
  24. 2020Controlling Computation versus Quality for Neural Sequence Modelscited by 27
  25. 2014Exponentially Increasing the Capacity-to-Computation Ratio for Conditional Computation in Deep Learningcited by 24
  26. 2023Learning to Skip for Language Modelingcited by 19
  27. 2024Accelerating Inference in Large Language Models with a Unified Layer Skipping Strategycited by 17

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 5; 2025: 3; 2026: 3.

11 papers.

  1. 2024Layerwise Recurrent Router for Mixture-of-Expertscited by 21
  2. 2025D2MoE: Dual Routing and Dynamic Scheduling for Efficient On-Device MoE-based LLM Servingcited by 17
  3. 2025Coupling Experts and Routers in Mixture-of-Experts via an Auxiliary Losscited by 11
  4. 2024More Expressive Attention with Negative Weightscited by 9
  5. 2025Autonomy-of-Experts Modelscited by 7
  6. 2024Duo-LLM: A Framework for Studying Adaptive Computation in Large Language Modelscited by 6
  7. 2024PEAR: Position-Embedding-Agnostic Attention Re-weighting Enhances Retrieval-Augmented Generation with Zero Inference Overheadcited by 1
  8. 2024Extracting and Combining Abilities For Building Multi-lingual Ability-enhanced Large Language Modelscited by 1
  9. 2026UniPool: A Globally Shared Expert Pool for Mixture-of-Expertscited by 1
  10. 2026Mixture of Layers with Hybrid Attentioncited by 0
  11. 2026CORTEX: Learning to Share and Specialize in Dense Language Modelscited by 0