DEMix Layers: Disentangling Domains for Modular Language Modeling
This work addresses domain-specific challenges in language modeling, offering modularity for improved generalization and flexibility, though it is incremental in building on existing mixture-of-experts approaches.
The paper tackles the problem of domain adaptation in language models by introducing DEMix layers, which condition the model on input text domains, resulting in reduced test-time perplexity, increased training efficiency, and rapid adaptation with minimal overhead.
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.