DEMix Layers: Disentangling Domains for Modular Language Modeling

TL;DR

This paper presents DEMix layers, a modular approach for domain-specific conditioning in language models, improving adaptability, efficiency, and generalization across multiple domains without retraining.

Contribution

Introduction of DEMix layers, enabling modular, domain-specific expert networks in language models that can be added, removed, or mixed dynamically without retraining.

Findings

Reduces test perplexity across domains

Enhances training efficiency for large LMs

Allows rapid domain adaptation and control

Abstract

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.