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

arXiv: https://arxiv.org/abs/2407.06677 · HTML: https://mlsanity.com/paper/2407.06677/

## Key facts

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

## Builds on (most cited first, 30 of 39)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Large Language Models Encode Clinical Knowledge](https://mlsanity.com/paper/2212.13138/) (2022)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization](https://mlsanity.com/paper/1808.08745/) (2018)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [SkipNet: Learning Dynamic Routing in Convolutional Networks](https://mlsanity.com/paper/1711.09485/) (2017)
- [DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference](https://mlsanity.com/paper/2004.12993/) (2020)
- [Confident Adaptive Language Modeling](https://mlsanity.com/paper/2207.07061/) (2022)
- [BERT Loses Patience: Fast and Robust Inference with Early Exit](https://mlsanity.com/paper/2006.04152/) (2020)
- [Highway Networks](https://mlsanity.com/paper/1505.00387/) (2015)
- [Mixture-of-Depths: Dynamically allocating compute in transformer-based language models](https://mlsanity.com/paper/2404.02258/) (2024)
- [Depth-Adaptive Transformer](https://mlsanity.com/paper/1910.10073/) (2019)
- [The Hydra Effect: Emergent Self-repair in Language Model Computations](https://mlsanity.com/paper/2307.15771/) (2023)
- [LaCo: Large Language Model Pruning via Layer Collapse](https://mlsanity.com/paper/2402.11187/) (2024)
- [Spatially Adaptive Computation Time for Residual Networks](https://mlsanity.com/paper/1612.02297/) (2016)

## Built on by (most cited first, 11 of 11)

- [Layerwise Recurrent Router for Mixture-of-Experts](https://mlsanity.com/paper/2408.06793/) (2024)
- [D2MoE: Dual Routing and Dynamic Scheduling for Efficient On-Device MoE-based LLM Serving](https://mlsanity.com/paper/2504.15299/) (2025)
- [Coupling Experts and Routers in Mixture-of-Experts via an Auxiliary Loss](https://mlsanity.com/paper/2512.23447/) (2025)
- [More Expressive Attention with Negative Weights](https://mlsanity.com/paper/2411.07176/) (2024)
- [Duo-LLM: A Framework for Studying Adaptive Computation in Large Language Models](https://mlsanity.com/paper/2410.10846/) (2024)
- [Autonomy-of-Experts Models](https://mlsanity.com/paper/2501.13074/) (2025)
- [PEAR: Position-Embedding-Agnostic Attention Re-weighting Enhances Retrieval-Augmented Generation with Zero Inference Overhead](https://mlsanity.com/paper/2409.19745/) (2024)
- [Extracting and Combining Abilities For Building Multi-lingual Ability-enhanced Large Language Models](https://mlsanity.com/paper/2410.07825/) (2024)
- [UniPool: A Globally Shared Expert Pool for Mixture-of-Experts](https://mlsanity.com/paper/2605.06665/) (2026)
- [Mixture of Layers with Hybrid Attention](https://mlsanity.com/paper/2605.09516/) (2026)
- [CORTEX: Learning to Share and Specialize in Dense Language Models](https://mlsanity.com/paper/2609.34449/) (2026)

## 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.