PaperlayerView exact revision on arXiv

Source-linked AI summary

DEMix Layers: Disentangling Domains for Modular Language Modeling

Suchin Gururangan, Mike Lewis, Ari Holtzman, Noah A. Smith, Luke Zettlemoyer

arXiv:2108.05036v2cs.CLcs.AI

TL;DR

Dense language-model training provides limited domain modularity and can be costly to adapt without forgetting prior data. The paper introduces DEMix layers with domain-specialized feedforward experts that can be mixed, added, or removed after training. Experiments report improved performance across training and novel domains, with modular adaptation and restriction capabilities, subject to limitations from shared parameters and inference-time mixing costs.

  • Problem

    Dense training can limit generalization, increase adaptation costs, and make adapting to new domains without forgetting or restricting exposed domains difficult.

  • Method

    DEMix layers replace transformer feedforward layers with domain-specialized experts and condition computation on known or estimated input domains.

  • Results

    DEMix layers improve performance in training and novel test-time domains while enabling experts to be mixed, added, or removed after training.

  • Takeaways & Limitations

    Explicit domain conditioning provides a modular language model that can adapt after training and restrict unwanted domains at inference time.

  • Takeaways & Limitations

    Removing experts cannot guarantee full forgetting because self-attention and input-embedding parameters remain shared across domains.

Abstract

from arXiv · show

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.

1 Introduction

Dense language-model training treats data uniformly, limiting generalization, efficiency, and flexibility across domains. DEMix layers introduce domain-specialized experts that can be mixed, added, or removed at inference time while improving performance in training and novel domains.

  • Motivation: Dense training updates all parameters across all data, leaving domain variation to be implicitly discovered.This approach assumes models can fit all domains equally well.
  • Motivation: Dense training can emphasize accessible or prevalent subsets, limiting generalization to less prevalent domains.Updating every parameter also becomes more expensive as model size grows.
  • Proposal: The paper proposes modular language models with components specialized to distinct domains and customizable by mixing, adding, or removing components after training.The design targets rapid adaptation without updating the entire model.
  • Proposal: A DEMix layer replaces a transformer feedforward layer with one domain-specific expert per domain, conditioning computation on the input domain.Known domains use their labels during training, while inference can estimate the domain when it is unknown.
  • Results: Mixing experts improves performance on heterogeneous, unseen, and training-domain test data, while adding or removing experts supports adaptation and restricting unwanted domains.The reported modular functions preserve generalization benefits and allow predictable inference-time changes.

2 Multi-Domain Corpus

The study uses a provenance-tagged, largely English multi-domain corpus designed for language modeling and reproducible evaluation. It separates eight training domains from eight novel domains to test distribution shift and unclear domain boundaries.

  • Corpus design: The corpus was constructed with explicit provenance metadata and contains substantial metadata-tagged text for language modeling.Its licensing supports reproducibility, and an API is provided for downloading and preprocessing the data.
  • Domain definition: The study operationalizes domain coarsely through document provenance, such as the dataset used to access the text.This choice is described as easy and intuitive, while the broader definition of domain remains debated.
  • Training domains: The training corpus contains eight largely English domains varying in complexity and coverage.These domains are listed at the top of Table 1 along with training and evaluation sizes.
  • Novel domains: The novel-domain collection contains eight largely English domains that may or may not align with the training domains.These domains evaluate generalization under a more challenging distribution shift with less-clear boundaries.

3 DEMIX Layer

DEMix replaces transformer feedforward layers with domain-specific experts routed by document-level metadata. This modular design specializes parameters by domain while supporting inference-time mixing, addition, and removal of experts.

  • Domain routing: Each DEMix layer contains one feedforward expert per training domain, with every token in a sequence routed to the expert selected by its domain label.The model uses domain metadata rather than learning token-level routing; eight training domains yield eight experts.
  • Domain routing: Unlike token-level mixture-of-experts methods, DEMix assigns all tokens in a sequence to the same domain expert.This document-level routing follows an observable segmentation of the training data and avoids learned routing and associated load-balancing techniques.
  • Architecture and efficiency: Replacing every transformer feedforward layer with DEMix layers increases total parameters and memory costs while reducing GPU latency and training runtime costs.Experts retain the original feedforward dimensions, so the effective parameter count grows even though computation is partitioned across GPUs.
  • Architecture and efficiency: DEMIX achieves the same or slightly higher TFLOPs/GPU throughput as DENSE models at the same total FLOPs per update, despite adding significantly more parameters.The efficiency gain comes from synchronizing expert parameters only among GPUs assigned to the same domain.

4 In-Domain Performance

Experiments compare DEMix against dense and domain-informed baselines under a fixed runtime budget. DEMix generally improves in-domain performance, with the clearest gains in smaller models, while domain heterogeneity explains important exceptions.

  • Overall in-domain results: Domain balancing consistently helps DENSE training and is especially important when domain prevalence is imbalanced.Balancing improves most domains but hurts REDDIT relative to the DENSE baseline, where REDDIT provides far more text than the other domains.
  • Overall in-domain results: Benefits from domain tokens or DEMix layers are clearest for the smallest model and remain smaller but consistent as model size increases.As models grow, DENSE increasingly catches up to models with additional domain information on average.
  • Domain-level variation: At 1.3B parameters per GPU, most training domains benefit from naive DEMix, but WEBTEXT, REALNEWS, and REDDIT perform worse than the DENSE baseline.The authors attribute variation to domain heterogeneity and differing similarity between domains, which can make parameter sharing advantageous.
  • Domain heterogeneity: DEMIX experts achieve their highest affinity on their assigned domains, while WEBTEXT, REALNEWS, and REDDIT experts show relatively high affinity across many domains.Expert affinity also correlates positively with bigram overlap between expert and target domains (r=0.40, t=3.45, p=0.001).
  • Domain heterogeneity: Because discrete domain labels can be too rigid, the paper introduces mixtures over domain experts to improve performance on heterogeneous domains.The mixture variant relaxes single-expert selection for mixed or unknown domains.

5 Mixing Experts at Inference Time

When domain labels are unavailable, ambiguous, or novel, DEMIX estimates domain posteriors and uses them to mix expert outputs at inference time. Across known and unseen domains, weighted mixtures improve over individual experts, dense baselines, and simple averaging, with cached priors performing best on novel domains.

  • Dynamically Estimating Domain Membership: DEMIX replaces hard expert selection with parameter-free, test-time mixture coefficients derived from posterior beliefs over input domains.The weights can be based on uniform, updated, or cached priors, and introduce no new parameters.
  • Dynamically Estimating Domain Membership: The model estimates domain posteriors from the input history and combines domain-conditioned expert predictions to produce the hidden representation.For CORD-19, the example assigns the highest probabilities to medical and news domains before mixing their expert outputs.
  • Novel Domain Performance: Cached-prior DEMIX mixtures achieve the best average perplexity across eight unseen domains and outperform dense baselines and individually selected experts.The comparison uses existing models only; no new models are trained for this experiment.
  • Visualizing Domain Membership: Posterior estimates identify known domains but also reveal substantial overlap, while novel domains distribute probability across multiple training domains.Reddit is estimated as approximately 80% REDDIT, 11% WEBTEXT, and 8% REALNEWS; CORD-19 spans MED, REALNEWS, and 1B, while Yelp Reviews spans REVIEWS, WEBTEXT, and REDDIT.
  • Novel Domain Performance: Weighted expert ensembling outperforms simple averaging, indicating that sparse posterior-weighted mixtures are important for heterogeneous domains.The cached-prior approach lets smaller models match or outperform much larger dense models.
  • In-Domain Performance: Cached-prior mixtures improve performance across all known domains and model scales, with average gains of 10.7% on REDDIT, 2.4% on WEBTEXT, and 1.9% on REALNEWS.The largest improvements occur on heterogeneous domains, where provenance categories may not align with the most effective domain boundaries.

6 Adaptive Pretraining with New Experts

DEMix-DAPT adapts language models to new domains by training only newly added experts, avoiding the forgetting associated with dense continued pretraining. Across novel-domain experiments, it generally matches or outperforms dense adaptation while adding about 10% of the model’s parameters.

  • Motivation: DENSE-DAPT updates all model parameters and may be expensive, while also risking forgetting from earlier training phases.DAPT uses unannotated in-domain text to improve task performance; dense continued pretraining can be difficult for large models and smaller computational budgets.
  • Method: DEMix-DAPT adds and trains a new expert for target data while leaving existing experts and shared parameters unchanged.The new expert can be initialized from the pretrained expert nearest to the target domain, and only its parameters are updated.
  • Results: DEMix-DAPT reduces perplexity on novel domains without forgetting, whereas DENSE-DAPT improves the target domain while progressively worsening training-domain performance.This contrast is illustrated for adaptation to a single additional domain and is also reported for larger models in the appendix.
  • Results: DEMix-DAPT outperforms DENSE-DAPT on some domains, including CORD-19 and ACL PAPERS, and closely approaches it on GUTENBERG.The authors suggest the method is especially helpful when the target domain strongly aligns with an existing expert.
  • Efficiency: 10% of total DEMix-model parameters comprise the additional expert, compared with updating all parameters under DENSE-DAPT.The reported parameter comparison concerns adding one expert for a new domain.
  • Results: Adding eight adapted experts reduces perplexity across all domains, including previously seen domains, without forgetting.Average reductions on previously seen domains include 2.4% on MED, 1.8% on REALNEWS, and 2% on REDDIT.

7 Language Models with Removable Parts

DEMix layers support inference-time removal of domain experts to restrict model behavior. Removing an expert degrades performance on its associated domain and often approaches the performance of a model retrained without that domain, but shared parameters prevent guarantees of complete forgetting.

  • Motivation: Dense training makes it difficult to restrict access to undesirable domains after training.The paper motivates domain removal partly by the presence of undesirable content such as hate speech and extremism in pretraining datasets.
  • Mechanism: DEMix experts assigned to unwanted domains can be disabled and unused at inference time.This provides lightweight control over which training domains contribute to predictions.
  • Limitation: Shared self-attention and input-embedding parameters mean expert removal does not guarantee that content from the removed domain is fully forgotten.The paper identifies establishing bounds on this residual access as future work.
  • Evaluation: The removal evaluation compares all experts active, one expert deactivated, and a model retrained without the corresponding domain.The retrained comparison replaces the removed domain with GUTENBERG.
  • Results: Removing a domain expert harms performance on its associated domain, in most cases approaching a model that was not exposed to that domain.For WEBTEXT and REALNEWS, the deactivated-expert model even underperforms the retrained comparison.

8 Related Work

The related work connects DEMix layers to metadata-conditioned language modeling, inference-time control, multilingual parameterization, continual learning, adapters, and multi-domain models. DEMix-DAPT particularly parallels parameter-efficient adaptation and model expansion approaches.

  • Metadata: Prior work uses document metadata for topic modeling, language-model adaptation, document representations, and inference-time control of generation.Some transformer language models prepend metadata to the input text, similarly to the paper’s domain-token setting.
  • Inference-Time Control: Inference-time control methods have used expensive procedures, densely trained models, or multiple experts, and the latter can also be applied to DEMix layers.Experts trained on different domains may steer text generation when used with DEMix layers.
  • Multilinguality: Multilingual research motivates language-specific parameters, while DEMix investigates analogous domain-specific parameters for language models.The paper notes that language boundaries are often clearer than domain boundaries.
  • Continual Learning: DEMIX-DAPT is a continual-learning approach that can be combined with regularization, meta-learning, episodic memory, or data replay.These methods are presented as related techniques for learning incrementally on new data.
  • LM Adapters: DEMIX-DAPT resembles adapter methods because both freeze the pretrained model and update a small number of additional parameters.The paper also suggests combining expert addition with adapter modules to improve efficiency.
  • Multi-Domain Models: Multi-domain modeling has been studied in statistical and neural machine translation, as well as with smaller models and explicit domain labels.Prior work includes both supervised domain conditioning and dense training in multi-domain settings.

9 Conclusion

The paper presents DEMix layers as a modular language-model design for rapid inference-time adaptation. Experts can be mixed, added, or removed to address heterogeneous, new, or unwanted domains.

  • Conclusion: DEMix layers provide inference-time modularity and a rapidly adaptable language-model system.This design addresses limitations attributed to dense training.
  • Conclusion: Experts can be mixed for heterogeneous or unseen domains, added to incorporate new domains, and removed to restrict unwanted domains.These operations modify domain participation at inference time.
  • Future Work: The design assumes coarse provenance labels or other metadata for identifying domains in pretraining data.The paper proposes unsupervised domain discovery as an alternative direction.
  • Future Work: Datasets with many diverse subdomains and long-tailed prevalence may require automatic clustering or hierarchical experts.The experiments study datasets with a few large domains, leaving narrower and less prevalent distributions as a future setting.

A.1 Collecting Domains

The authors collect most domains from associated sources, with specialized acquisition pipelines for Twitter, Gutenberg, and breaking-news data.

  • Most domains use the associated sources listed in Table 1 without modification.
  • Twitter data comes from the Twitter Academic API, while Gutenberg data uses the gutenberg-dammit scraping tool.
  • Breaking-news data is gathered from sources rated highly factual, using Newspaper3K to scrape the latest 1,000 articles from each site.After deduplication, the collection contains about 20K articles from 400 news sources.

A.2 Dataset Anonymization

The appendix documents dataset anonymization, domain-posterior estimation, and additional evaluation materials. It also reports that DENSE-DAPT improves novel-domain perplexity while substantially degrading performance on training domains.

  • Dataset Anonymization: Anonymization uses regexes targeting common user-identifiable patterns and replaces matches with dummy tokens.The regexes and corresponding dummy tokens are listed in Table 8.
  • Domain Posterior Estimation: Domain posteriors are tracked over development-data blocks, becoming usually stable after about 50 blocks.The experiments conservatively use 100 blocks to compute each domain posterior, including for training and novel domains.
  • DENSE-DAPT Evaluation: DENSE-DAPT improves performance on novel domains across model sizes but causes a large performance hit on training domains.Table 13 measures average perplexity changes in training and novel domains relative to the original DENSE language model; negative values indicate improvement.
  • Experimental Documentation: The appendix provides pretraining hyperparameters for 125M, 350M, 760M, and 1.3B-parameter-per-GPU language models.The hyperparameters are reported as identical for DEMIX and DENSE training.
Loading 2108.05036v2…