# DEMix Layers: Disentangling Domains for Modular Language Modeling

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

## Key facts

- Cited by 132 papers in the map (data through 14 Oct 2026).
- Builds on 53 of its 84 references that are in the map.
- First posted on arXiv 11 Aug 2021.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [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)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Gradient Episodic Memory for Continual Learning](https://mlsanity.com/paper/1706.08840/) (2017)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Progressive Neural Networks](https://mlsanity.com/paper/1606.04671/) (2016)
- [RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models](https://mlsanity.com/paper/2009.11462/) (2020)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [BitFit: Simple Parameter-efficient Fine-tuning for Transformer-based Masked Language-models](https://mlsanity.com/paper/2106.10199/) (2021)
- [Plug and Play Language Models: A Simple Approach to Controlled Text Generation](https://mlsanity.com/paper/1912.02164/) (2019)
- [CTRL: A Conditional Transformer Language Model for Controllable Generation](https://mlsanity.com/paper/1909.05858/) (2019)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Beyond English-Centric Multilingual Machine Translation](https://mlsanity.com/paper/2010.11125/) (2020)
- [Defending Against Neural Fake News](https://mlsanity.com/paper/1905.12616/) (2019)
- [Neural Text Generation with Unlikelihood Training](https://mlsanity.com/paper/1908.04319/) (2019)
- [A Dataset of Information-Seeking Questions and Answers Anchored in Research Papers](https://mlsanity.com/paper/2105.03011/) (2021)
- [MAD-X: An Adapter-Based Framework for Multi-Task Cross-Lingual Transfer](https://mlsanity.com/paper/2005.00052/) (2020)
- [DExperts: Decoding-Time Controlled Text Generation with Experts and Anti-Experts](https://mlsanity.com/paper/2105.03023/) (2021)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [One Billion Word Benchmark for Measuring Progress in Statistical Language Modeling](https://mlsanity.com/paper/1312.3005/) (2013)
- [The Pushshift Reddit Dataset](https://mlsanity.com/paper/2001.08435/) (2020)

## Built on by (most cited first, 30 of 132)

- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [Continual Learning of Large Language Models: A Comprehensive Survey](https://mlsanity.com/paper/2404.16789/) (2024)
- [Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints](https://mlsanity.com/paper/2212.05055/) (2022)
- [Continual Pre-training of Language Models](https://mlsanity.com/paper/2302.03241/) (2023)
- [Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models](https://mlsanity.com/paper/2208.03306/) (2022)
- [A Review of Sparse Expert Models in Deep Learning](https://mlsanity.com/paper/2209.01667/) (2022)
- [Lifting the Curse of Multilinguality by Pre-training Modular Transformers](https://mlsanity.com/paper/2205.06266/) (2022)
- [Continual Pre-Training of Large Language Models: How to (re)warm your model?](https://mlsanity.com/paper/2308.04014/) (2023)
- [Lifelong Pretraining: Continually Adapting Language Models to Emerging Corpora](https://mlsanity.com/paper/2110.08534/) (2021)
- [ERNIE-ViLG 2.0: Improving Text-to-Image Diffusion Model with Knowledge-Enhanced Mixture-of-Denoising-Experts](https://mlsanity.com/paper/2210.15257/) (2022)
- [Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM](https://mlsanity.com/paper/2403.07816/) (2024)
- [Time Waits for No One! Analysis and Challenges of Temporal Misalignment](https://mlsanity.com/paper/2111.07408/) (2021)
- [Soft Merging of Experts with Adaptive Routing](https://mlsanity.com/paper/2306.03745/) (2023)
- [Test-Time Training on Nearest Neighbors for Large Language Models](https://mlsanity.com/paper/2305.18466/) (2023)
- [AdapterSoup: Weight Averaging to Improve Generalization of Pretrained Language Models](https://mlsanity.com/paper/2302.07027/) (2023)
- [Continual Learning of Natural Language Processing Tasks: A Survey](https://mlsanity.com/paper/2211.12701/) (2022)
- [Modular Deep Learning](https://mlsanity.com/paper/2302.11529/) (2023)
- [SILO Language Models: Isolating Legal Risk In a Nonparametric Datastore](https://mlsanity.com/paper/2308.04430/) (2023)
- [ELLE: Efficient Lifelong Pre-training for Emerging Data](https://mlsanity.com/paper/2203.06311/) (2022)
- [Z-ICL: Zero-Shot In-Context Learning with Pseudo-Demonstrations](https://mlsanity.com/paper/2212.09865/) (2022)
- [In-context Learning with Retrieved Demonstrations for Language Models: A Survey](https://mlsanity.com/paper/2401.11624/) (2024)
- [Lifelong Language Pretraining with Distribution-Specialized Experts](https://mlsanity.com/paper/2305.12281/) (2023)
- [Towards Lifelong Learning of Large Language Models: A Survey](https://mlsanity.com/paper/2406.06391/) (2024)
- [MoRAL: MoE Augmented LoRA for LLMs' Lifelong Learning](https://mlsanity.com/paper/2402.11260/) (2024)
- [NLPeer: A Unified Resource for the Computational Study of Peer Review](https://mlsanity.com/paper/2211.06651/) (2022)
- [Scaling Expert Language Models with Unsupervised Domain Discovery](https://mlsanity.com/paper/2303.14177/) (2023)
- [Investigating Continual Pretraining in Large Language Models: Insights and Implications](https://mlsanity.com/paper/2402.17400/) (2024)
- [Knowledge Card: Filling LLMs' Knowledge Gaps with Plug-in Specialized Language Models](https://mlsanity.com/paper/2305.09955/) (2023)
- [Recent Advances of Foundation Language Models-based Continual Learning: A Survey](https://mlsanity.com/paper/2405.18653/) (2024)
- [D-CPT Law: Domain-specific Continual Pre-Training Scaling Law for Large Language Models](https://mlsanity.com/paper/2406.01375/) (2024)

## Abstract (arXiv)

We introduce a new domain expert mixture (DEMix) layer that enables conditioning a language model (LM) on the domain of the input text. A DEMix layer is a collection of expert feedforward networks, each specialized to a domain, that makes the LM modular: experts can be mixed, added or removed after initial training. Extensive experiments with autoregressive transformer LMs (up to 1.3B parameters) show that DEMix layers reduce test-time perplexity, increase training efficiency, and enable rapid adaptation with little overhead. We show that mixing experts during inference, using a parameter-free weighted ensemble, allows the model to better generalize to heterogeneous or unseen domains. We also show that experts can be added to iteratively incorporate new domains without forgetting older ones, and that experts can be removed to restrict access to unwanted domains, without additional training. Overall, these results demonstrate benefits of explicitly conditioning on textual domains during language modeling.