Scalable Multi-Domain Adaptation of Language Models using Modular Experts

## Abstract

Abstract Domain-specific adaptation is critical to maximizing the performance of
pre-trained language models (PLMs) on one or multiple targeted tasks,
especially under resource-constrained use cases, such as edge devices.
However, existing methods often struggle to balance domain-specific
performance, retention of general knowledge, and efficiency for training
and inference.
To address these challenges, we propose Modular Domain Experts (MoDE).
MoDE is a mixture-of-experts architecture that
augments a general PLMs with modular, domain-specialized experts.
These experts are trained independently and composed together via a
lightweight training process.
In contrast to standard low-rank adaptation methods, each MoDE expert consists of several transformer layers which scale better with
more training examples and larger parameter counts.
Our evaluation demonstrates that MoDE achieves comparable target
performances to full parameter fine-tuning while achieving 1.65%
better retention performance.
Moreover, MoDE’s architecture enables flexible sharding
configurations and improves training speeds by up to 38% over
state-of-the-art distributed training configurations.

## 1 Introduction

Recent advances in large-scale Pre-trained Language Models (PLMs) have
showcased impressive generalization capabilities Brown et al. (); Chowdhery et al. (); Anil et al. (); Team et al. ().
However, when applied to specialized domains such as medical, legal, or
financial sectors, these general-purpose models often require further
fine-tuning to maximize performance on target
domains Huang et al. (); Li et al. (); Singhal et al. ().

A straightforward approach to domain adaptation is full-parameter
fine-tuning, where the entire model is further trained on domain-specific
data Houlsby et al. (); Bapna et al. ().
While this method provides strong performance on target domains,
full-parameter fine-tuning may lead to catastrophic forgetting
where the model loses previously learned capabilities by overfitting to the
target domain Goodfellow et al. (); Kirkpatrick et al. ().
Additionally, this method is memory-intensive to serve in multi-domain
settings as each domain has a unique set of parameters, incurring a
significant parameter loading overhead when switching between
domains Hu et al. ().
In such cases, the cost of frequent “context switches”
significantly impacts performance, making full-parameter fine-tuning
impractical for scalable, efficient deployment Dhar et al. ().

To address the issues of forgetting and memory-efficiency,
parameter-efficient fine-tuning methods have been proposed, such as adapter
modules Houlsby et al. (), LoRA Hu et al. (), and
CoDA Lei et al. ().
These methods introduce a small number of trainable parameters and keep the
original model frozen during training.
Through targeted parameter updates,
these approaches are both computationally efficient and effective at
retaining prior knowledge Biderman et al. ().
Despite these benefits, parameter-efficient methods are limited in their expressive potential and
struggle to scale effectively across large domains and
datasets Niederfahrenhorst et al. ().

In this work, we propose Modular Domain Experts (MoDE), a scalable method for
multi-domain adaptation. MoDE is inspired by the Mixture of Experts
(MoE) approach (detailed in Section 2), and extends PLMs by
introducing modular, composable domain-specialized experts.
In MoDE, each expert consists of several transformer layers that
operate in parallel to the “backbone” PLM.
During training, the backbone remains frozen, and
each expert is independently trained for its respective domain.
For multi-domain deployment, MoDE combines multiple experts with
different specializations by placing them in parallel to the backbone
PLM, allowing for strong performance across diverse domains.
A lightweight fine-tuning steps teaches the experts to collaborate
effectively to provide strong performance across all target domains.
We further take advantage of parallelism in MoDE’s
architecture to enable new sharding strategies that improve training
efficiency, and provide privacy benefits during inference.

In summary, we make the following contributions:

- •
We introduce the novel Modular Domain Experts (MoDE) architecture
designed to address the scalability, catastrophic forgetting, and
memory constraints in multi-domain adaptation.
- •
We evaluate the performance of MoDE and demonstrate that it
outperforms LoRA by 1.4% and full parameter fine-tuning by 0.6% on a
multi-domain task.
- •
We analyze the training efficiency of MoDE, showing better
scalability compared to LoRA as the number of training examples and
parameters increases.
- •
We demonstrate how MoDE’s design enables flexible sharding
strategies that accelerate training, achieving up to 38% speedup over
standard distributed training configurations.

Figure: (a) Model architecture.
Refer to caption: x1.png

## 2 Related Work

Domain-Adaptive Pre-training.
MoDE is designed for scenarios where large domain-specific unlabeled
datasets are available, a process often referred to as “domain-adaptive
pre-training”, “continual learning”, “continued pre-training”, or
“further pre-training” Shi et al. (); Azerbayev et al. (); Colombo et al. (); Agarwal et al. ().
Beyond single-domain adaptation, recent research has expanded into
multi-domain adaptation, which presents additional
challenges Saunders (); Wu et al. ().
To the best of our knowledge, the most relevant works are
parameter-efficient fine-tuning methods, which are described in detail in
the next paragraph.

Parameter-Efficient Fine-Tuning.
Significant progress has been made in parameter-efficient fine-tuning
methods, which update only a small subset of model
parameters Houlsby et al. (); He et al. ().
Techniques such as LoRA Hu et al. () and
QLoRA Dettmers et al. () enhance language model performance by
introducing trainable low-rank decomposition matrices.
However, these low-rank approaches often fall short when adapting to
domains that require high expressiveness, such as mathematical
reasoning or coding Biderman et al. (); Niederfahrenhorst et al. ().

Mixtures of Experts (MoE) architectures are promising for
expanding the capacity of language models while keeping the computational
cost of training constant.
MoE models learn a routing function that selectively activates a subset of
experts, enabling sparse computation Du et al. (); Zhou et al. ().
Unlike MoDE, MoE models are primarily used to enhance pre-training
performance Dai et al. (); Team (); xAI ().
Recent approaches extend MoE concepts to improve adaptation to new
domains, such as life-long learning with distribution-specialized
experts Chen et al. (), conditional
adapters Lei et al. (), AdaMix Wang et al. (), and
MixPHM Jiang & Zheng ().
These approaches modify PLMs at a fine granularity (e.g., by modifying the
feed-forward network in the transformer layers).
In contrast, MoDE applies MoE at the transformer-level, offering a
scalable and expressive solution for multi-domain adaptation.

Furthermore, MoE-inspired approaches have shown potential in selecting
domain-specific adapters by routing input sequences, which enhances model
performance across diverse domains while alleviating catastrophic
forgetting Feng et al. ().
In this work, we focus on token-level routing, with sequence-level routing
as a natural extension for future research.

Distributed Training.
Scaling PLMs to billions of parameters requires partitioning parameters and
models and training data across multiple processors Dean et al. (); Li et al. (); Barham et al. ().
Common sharding strategies, including model parallelism and
data parallelism, evenly divide inputs and parameters across
accelerators Rajbhandari et al. (); Lepikhin et al. (); Alabed et al. ().
To address memory overheads and communication bottlenecks, researchers have
developed specialized tools to optimize the sharding configuration of a
model for the available
processors Lepikhin et al. (); Alabed et al. ().

Typically, distributed training and serving frameworks implement the
Single Program, Multiple Data (SPMD) model of computation which
enables building models and training programs as if developing for a single
processor with a large pool of memory.
However, SPMD has two key limitations:
(i) all operations (i.e., layers of a model) execute sequentially,
precluding the execution of different operations in parallel on different
processors, and
(ii) all operations must use all devices, which may be inefficient (e.g.,
by incurring significant communication overhead).
In contrast, the Multiple Program, Multiple Data (MPMD)
computational model has the ability to run multiple programs in parallel on
different devices, enabling fine-grained control over parallelism which has
the potential for more efficient execution compared to
SPMD Zheng et al. ().
We exploit parallelism in MoDE’s structure to implement an
MPMD-enabled sharding strategy that accelerates training by up to 38% over
SPMD model parallelism.

## 3 Method

Our proposed model architecture augments PLMs with modular, domain-specialized experts and aims to achieve the following design goals:

- •
Domain specialization. Experts must provide strong
performance and high expressiveness on their target domains.
Moreover, composing multiple experts with a PLM should result in strong performance on each of the target domains.
- •
Capability retention. The model must preserve the
generalist capabilities of the PLM and avoid catastrophic forgetting
which may cause performance regressions on other tasks.
- •
Efficient training and inference. The model architecture
must be efficient to train on large clusters, and provide advantages
for inference on edge devices.

### 3.1 Model Design

To meet these design goals, we decompose the language model into multiple
MoDE blocks (see Figure 1(a)). Each $i$th
block contains:
(i) a few backbone transformer layers,
$f_{\text{bb}}^{i}(\cdot)$, which are initialized from PLMs, and
(ii) a set of $N$ experts where each expert has the same (small) number
of transformer layers, and the $i$th block of the $j$th expert is given by
$f_{\text{e}}^{i,j}(\cdot)$.
All backbone and expert layers run in parallel, with their outputs combined
as a weighted sum.
The weights for each component are determined by a lightweight gating
function, $g^{i}(\cdot)$.

For each $i$-th block, the input $x\in\mathbb{R}^{d}$ and the output $y\in\mathbb{R}^{d}$ are defined as:

$$ $\displaystyle y$ $\displaystyle=\alpha_{\text{bb}}\,f_{\text{bb}}^{i}(x)+\sum_{j=1}^{N}\alpha_{ \text{e}}^{j}\,f_{\text{e}}^{i,j}(x)$ (1) $\displaystyle\quad\text{where}\quad[\alpha_{\text{bb}},\alpha_{\text{e}}^{1}, \alpha_{\text{e}}^{2},\cdots,\alpha_{\text{e}}^{N}]=g^{i}(x)$ $$

Backbone Model.
The backbone transformer layers, $f_{\text{bb}}^{i}(x)$ are derived from
PLMs. In each MoDE block, the number of backbone layers is determined
by dividing the total number of layers in the original PLM by the total
number of MoDE blocks.
The number of MoDE blocks is an important configuration
hyperparameter:
more
blocks enables more frequent
synchronization between the backbone and expert layers which provides
greater flexibility.
In Table 3, we conduct an ablation to assess the
impact of the number of blocks on accuracy.
In MoDE, the backbone layers are initialized from PLMs and remain
frozen during training to mitigate the issue of catastrophic forgetting.

Experts.
Each expert, $f_{\text{e}}^{i,j}(x)$ consists of several transformer
layers, and the number of transformer layers impacts the expert’s
performance.
Table 3 indicates that increasing the number of expert
layers improves performance on the target domain at the cost of higher
computational cost.
We use transformer layers for experts for two reasons:
(i) transformer blocks scale more efficiently than finer-grained designs,
such as modifications within attention matrices, and
(ii) they simplify deployment, as there are no modifications to the
transformer architecture itself.
In contrast, fine-grained methods often face deployment
challenges Yi et al. ().

Gating Function.
The outputs from the backbone and expert layers are combined through the gating layer, $g^{i}(x)$.
In this work, we use a simple token-level gating function.
This function consists of linear layer followed by a softmax for
normalization.
The gating function takes $x$ as input and outputs a vector of length $N+1$
as follows:

$$ $\displaystyle g^{i}(x)=\texttt{softmax}(\texttt{Linear}(x))$ (2) $$

The simplicity of MoDE’s gating function enables efficient
composition of multiple experts, as shown in Table 1.
However, future work could explore more advanced designs, such as sparse
gating and sequence routing, to further optimize memory usage and
computational efficiency.

Figure: (a) A model-parallel partitioning in which an entire MoDE model is split across all TPUs.
Refer to caption: x3.png

### 3.2 Training Procedure

MoDE has a two stage training procedure (Figure 1(b)): $(i)$ train a single expert independently on each domain, and $(ii)$ compose different experts into a single model to improve multi-domain performance.

Single Modular Expert Training.
For each domain, we independently train an MoDE model with a single expert. During training, only the expert layers and gating layers are updated, while the backbone layers remain frozen. Freezing the backbone not only reduces the computational cost of backpropagation but also helps mitigate catastrophic forgetting.

Composing Experts.
After training the modular experts on each domain, we can compose them into a single MoDE model for multi-domain adaptation. While the experts contain domain-specific knowledge, their outputs can interfere when combined, as they are pre-trained individually. To address this, we apply a lightweight fine-tuning step using data from all domains to learn the optimal weights for the gating function. Although freezing the experts requires fewer training examples during composition Appendix C, we have demonstrated that unfreezing the experts during this step further enhances multi-domain adaptation performance Table 1.

### 3.3 Parallelizing Block Execution

We exploit parallelism in MoDE’s structure to enable flexible
sharding configurations supported by the MPMD model of computation.
These sharding configurations improve training efficiency due to:

- 1.
The ability to execute the backbone and experts on distinct,
smaller sets of devices which reduces communication overheads
(Figure 2).
- 2.
The ability to configure the backbone and experts independently to
maximize overall efficiency, when their number of parameters are
different.

While MPMD has the potential to reduce communication overheads by running
blocks on fewer devices, the merges in the MoDE model architecture
are points of synchronization that require communication across all
devices.
To determine whether MPMD is beneficial for training, we evaluate whether
MPMD configurations can afford the cost of resharding between meshes in
Section 4.4.

While MPMD has the potential to reduce communication overheads by executing
blocks on smaller sets of devices, the linear combination of the outputs
from the backbone and expert blocks requires are points of synchronization
that require communication across all devices.
To assess the effectiveness of MPMD in training, we evaluate if the
benefits of MPMD configurations outweigh the costs associated with
resharding between device meshes, as discussed in
Section 4.4.

Moreover, MoDE’s architecture enables several benefits during
inference for privacy and performance.
By using MPMD to run MoDE with the backbone and expert parameters on
different devices, expert parameters for sensitive domains (e.g., personal
information, medical data, etc.) can remain private (e.g., by executing on a
user’s phone).
To switch to another MoDE model adapted to different task, only the
new expert weights need to be loaded into memory because the backbone PLM’s
weights are shared. This reduces the overhead of loading model weights and works well with
existing techniques developer to serve many adapters concurrently, such as
SLoRA Sheng et al. ().

## 4 Experiments

### 4.1 Experiment Setup

PLM Configuration.
Our PLM model consists of 1.58 billion parameters distributed across 18 transformer layers, comparable to smaller open-source models such as Gemma Team et al. (), Phi Gunasekar et al. (), and Llama 3.2 Meta (). The model is pre-trained on a high-quality dataset that spans a diverse range of natural language use cases, similar to GPT-3 Brown et al. (), GLaM Du et al. (), and LLaMA Touvron et al. ().

Multi-Domain Datasets.
We prepare two target domain datasets, Code and Math, to evaluate multi-domain adaptation, and one retention dataset, English, to measure catastrophic forgetting:

- •
Code consists of code samples retrieved from real-world applications.
- •
Math tests the model’s mathematical reasoning capabilities and is similar to the GSM Cobbe et al. () and MATH Hendrycks et al. () datasets.
- •
English contains literature texts with a distribution distinct from Code and Math.

Additionally, we create a mixed dataset, Math + Code, for model training, which contains an equal number of samples from the Math and Code datasets.

Evaluation Metric.
As our focus is on domain-adaptive pre-training with large amounts of unlabeled data, we adopt next-token prediction accuracy as the primary metric to assess the effectiveness of our methods. This metric has been shown to correlate strongly with downstream performance Kaplan et al. (); Hoffmann et al. (). In addition, since the evaluation spans multiple domains and includes a retention dataset, we use the average performance across all datasets to rank the models.

Baselines.
We benchmark MoDE against two widely-used domain adaptation methods: full-parameter fine-tuning Houlsby et al. () and LoRA Hu et al. ().
We exclude similar low-rank methods, such as CoDA Lei et al. () and QLoRA Dettmers et al. (), which prioritize computational efficiency and exhibit lower accuracy than LoRA.
Additionally, we conduct a comprehensive hyperparameter search for LoRA (Table 2) and report results for the best-performing LoRA configuration.

Chosen MoDE Configuration. We configure MoDE with three blocks, each containing two transformer layers per expert. For a detailed comparison of different configurations, refer to Table 3.

Training Details All methods are trained on a cluster of 128 TPUv3 accelerators across 8 servers. We use a batch size of 128, a learning rate of 0.001,
and train for 50k steps.

### 4.2 Multi-Domain Adaptation

Overall Performance.
We evaluate MoDE, LoRA, and full-parameter fine-tuning on the Math, Code, and English test sets.
We train the baseline models
directly on the mixed dataset, Math + Code.
For MoDE, we follow the training procedure described in Section 3.2, where individual MoDE models are first trained on Math and Code training sets independently, and then composed using the Math + Code dataset.
As shown in Table 1, MoDE with two experts (one for Math and one for Code) achieves the best overall performance and retention capabilities. On average, MoDE outperforms full-parameter fine-tuning by 0.59% and LoRA by 1.41%.
Notably, MoDE achieves target domain performance comparable to full-parameter fine-tuning, outperforming it by 0.28% on Math while trailing by 0.06% on Code, and provides a 1.55% improvement in retention on the English dataset.
Compared to LoRA, MoDE delivers a higher accuracy on the target domains, with a 1.68% improvement on Math and 2.12% on Code, and surpasses LoRA by 0.45% in English retention.

**Table 1: Multi-domain performance comparison of several methods which demonstrates that MoDE provides the best overall performance and capability retention. We adapt each method to the Math + Coding dataset, and compare the performance on the test sets for math, coding, and English retention.**
| Method | Accuracy |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Method | Parameters | Math | Code | English | Average |
| Full-parameter Fine-tuning | 1.583 B | 77.18 | 67.89 | 47.95 | 64.34 |
| LoRA | 1.585 B | 75.79 | 65.71 | 49.05 | 63.52 |
| MoDE 1$\times$Uninitialized Expert | 1.979 B | 76.71 | 67.30 | 49.17 | 64.39 |
| MoDE 2$\times$Uninitialized Experts | 2.376 B | 76.87 | 67.39 | 49.47 | 64.58 |
| MoDE 2$\times$Frozen Experts | 2.376 B | 76.99 | 66.94 | 49.07 | 64.33 |
| MoDE 2$\times$Experts | 2.376 B | 77.47 | 67.83 | 49.50 | 64.93 |
| vs. Full-parameter<br>Fine-tuning | $\uparrow$0.793 B | $\uparrow$0.29 | $\downarrow$0.06 | $\uparrow$1.65 | $\uparrow$0.59 |
| vs. LoRA | $\uparrow$0.789 B | $\uparrow$1.68 | $\uparrow$2.12 | $\uparrow$0.45 | $\uparrow$1.41 |

MoDE Configurations.
We address two questions in exploring different configurations:

- 1.
Which configuration of MoDE modular experts provides the best performance?
- 2.
Which composition strategy yields the best multi-domain performance?

For the first question, we evaluate MoDE with only one uninitialized expert and vary two hyperparameters: the number of MoDE blocks, and the number of transformer layers per expert within each block. As shown in Table 3, increasing the total number of expert layers leads to higher accuracy in the target domain.
To balance the number of added parameters, domain-specific accuracy, and capability retention, we configure the MoDE model with three MoDE blocks, each containing two transformer layers per expert (six layers per expert total).

For the second question, we compare against three alternative MoDE configurations.
We first measure the benefit of re-using experts by training MoDE configurations with one or two randomly initialized experts:
MoDE 1$\times$Uninitialized Expert measures the techniques performance when training a new expert on the multi-domain dataset from scratch, and MoDE 2$\times$Uninitialized Experts accounts for the difference in the number of parameters.
We also measure how much fine-tuning on the multi-domain dataset benefits accuracy.
For this configuration (MoDE 2$\times$Frozen Experts), we train only the gating function and re-use the weights single-domain training: one expert trained on Math and the other trained on Code.
According to Table 1, we find that more experts, initialization, and multi-domain fine-tuning all improve performance.

In addition, we examine how much mixed data from Math + Code is required for strong performance. As discussed in Appendix C, freezing the experts during composition yields better training sample efficiency compared to other methods.
This finding highlights the potential of our approach for privacy-preserving training, where different entities can independently train domain-specialized experts and later combine them into a coherent model, without sharing expert weights and only sharing sharing a small amount of training data.

Figure: Figure 3: Scalability of adaptation methods. We find that MoDE scales better than LoRA as the number of training examples and the number of parameters added increases. In the left figure, we increase the number of adapter parameters for LoRA by increasing the rank and MoDE by increasing the number of expert layers, and find that MoDE provides higher accuracy than LoRA with more trainable parameters. In the right figure, we generate versions of the Code with different numbers of training examples, and train a LoRA adapter and a MoDE expert for the same number of training steps on each. Although LoRA provides better accuracy on small datasets up to $\sim$1k training examples, we find that MoDE’s accuracy is better on large datasets, demonstrating that MoDE scales better with more training data than LoRA.
Refer to caption: x5.png

### 4.3 Scalability of MoDE

We evaluate the scalability of MoDE along two dimensions: scaling with additional parameters and scaling with increased training examples. We use LoRA as the baseline in both experiments and conduct experiments on a single domain using the Code dataset for training and evaluation.

Scalability with Additional Parameters.
We compare the accuracy improvements of MoDE and LoRA as additional parameters increases. For LoRA, we increase parameters by increasing the rank of the decomposition matrices, ranging from 8 to 2,048. For MoDE, we increase the number of transformer layers per expert while keeping the number of MoDE blocks constant. As shown in Figure 3, while LoRA is parameter-efficient, it struggles to convert additional parameters into higher accuracy. In contrast, MoDE shows a continuous improvement in accuracy as more parameters are added, indicating that MoDE scales more effectively with more parameters.

Scalability with Training Examples.
We evaluate how well MoDE and LoRA leverage additional training data to improve accuracy. Different versions of the Code dataset were created, each containing varying amounts of training examples. Both MoDE and LoRA were trained on these dataset versions, with the number of training steps keep the same across all experiments. As shown in Figure 3, we observe that LoRA performs better on smaller datasets (hundreds to thousands of training examples), but its accuracy plateaus as the number of training examples exceeds one thousand. In contrast, MoDE continues to improve as the number of training examples increases, demonstrating superior scalability with larger training sets.

Figure: (a) The cost of resharding between meshes is lower than the cost of model parallelism with experts.
Refer to caption: x7.png

### 4.4 Accelerating Training

Through our experiments, we find several MPMD-enabled
sharding configurations which increase training
speeds by up to 38% over SPMD model-parallelism.
The results confirm that the reduction in communication costs by running
backbones and experts on fewer accelerators outweighs the added costs of
resharding across meshes when merging the intermediate outputs of models.

All experiments use 8 TPUv5e accelerators which are managed by the Pathways
orchestration layer which supports MPMD Barham et al. ().
We implement MPMD training by extending the Pax distributed machine learning
framework which is
implemented in Jax Bradbury et al. (). While we use TPUs and a Jax-based framework as our training environment, we
underscore that the key idea of using MPMD to increase training speed in distributed settings extends to other accelerators (e.g., GPUs) and
training systems such as PyTorch Paszke et al. (). We measure the training performance, but note that many of the performance improvements enabled by MPMD sharding strategies may also benefit model serving. In details, we compare the following configurations:

- 1.
MPMD divides the TPUs
into different meshes to which backbone or experts are assigned.
- 2.
SPMD creates a single
mesh which consists of all 8 TPUs, backbone and experts.
- 3.
SPMD-pjit uses our training library’s to parallelize model training.

Cost of Resharding.
We compare the following configurations, which train MoDE using 4
experts where each consists of 6 transformer layers divided evenly into two
blocks:

- 1.
1 mesh uses model parallelism to shard the model across a
single mesh.
- 2.
3 mesh assigns the backbone to 4 TPUs and each expert
shares 2 TPUs with another expert. The backbone and the experts
are sharded on their meshes using model parallelism.
- 3.
5 mesh assigns the backbone to 4 TPUs, and one expert to
each of the remaining TPUs. The backbone is sharded with model
parallelism.

We find that the cost of resharding between meshes is lower than the cost
of model parallelism (Figure 4(a)).
The reduction in training step latency from $329$ ms for 1 mesh to $204$
ms, a 38% reduction, results due to the reduction in communication caused
by the MPMD sharding configuration.
Because the latency of the 3 mesh and 5 mesh
configurations is similar, we conclude that the cost to reshard does not
increase with the number of meshes.

Configuring the Expert Size.
We configure a mixture of 2 expert divided into two blocks and change the
total number of transformer layers added.
For example, setting the expert block size to 2 transformer layers adds
8 total transformer layers to the model.
The backbone is assigned to 6 TPUs with a model parallel sharding, and each
expert is assigned to one unoccupied TPU.

We find that SPMD performs better than MPMD when the expert size is small
or large (Figure 4(b)) due to under-utilization
of the TPUs.
When the expert size is small (e.g., 8 added layers total) and we train
using MPMD, the expert blocks finish executing before the backbone block,
so the expert TPUs remain idle until the backbone block completes and the
merge begins.
Similarly, when the expert size is large (e.g., 20 added layers total),
the backbone blocks finish before the expert blocks so the backbone’s TPUs
idle until the expert complete.
In contrast, SPMD executes the backbone and expert blocks sequentially on
all TPUs, ensuring that each TPU is always utilized.

When the runtime of a backbone block is similar to the runtime of the
expert blocks, MPMD training results in high utilization of the
accelerator.
Several settings impact the TPU utilization using MPMD which can be
configured to ensure high utilization, such as the number of accelerators
allocated to each mesh, the size of the expert blocks, and the sharding
strategy on each mesh.

Cost of Merges.
We train a mixture of 4 experts with 6 transformer layers each.
The backbone is assigned to a mesh of 4 TPUs, and each expert is assigned
to 1 unoccupied TPU.
We change the number of merges which determines the number of blocks (e.g.,
3 merges divide the backbone and each expert into three blocks).
We find that MPMD training speed decreases as the number of merges
increases due to the added communication costs incurred by each merge
(Figure 4(c)).

## 5 Conclusions

Adapting pre-trained language models (PLMs) to complex, multi-domain
settings is increasingly important as these models are deployed in
specialized environments requiring diverse capabilities.
To address this challenge, we propose Modular Domain Experts (MoDE), a
scalable technique for adapting PLMs to multi-domain tasks.
MoDE introduces modular, domain-specialized experts while preserving
the general knowledge of the PLM by keeping its weights frozen.
MoDE allows experts to scale with both the number of parameters and
the amount of training data, outperforming standard parameter-efficient
methods like LoRA by 1.4% on a challenging multi-domain dataset.
Additionally, MoDE delivers competitive performance compared to
full-parameter fine-tuning, achieving 1.5% better retention of general
capabilities and 0.6% higher accuracy across all subdomains.
To further optimize training efficiency, MoDE ’s architecture supports
flexible sharding strategies through MPMD, resulting in up to 38% faster
training compared to standard distributed training methods.
The ability to compose and reuse modular experts represents a significant
advancement in adapting PLMs to complex domains.
We hope that our work inspires further research into building modular,
expert-based language models.

Acknowledgments.
We would like to thank Ionel Gog and the Pathways Team for their support
and feedback on this work.

## Appendix A Ablation of LoRA Configurations

In Table 2, we provide an ablation study of
different LoRA configurations when adapting to the Code.
We vary the rank of the low-rank matrices introduced for adaptation from 8
to 2048.
While a rank of 32 provides the best performance on the Code
dataset, we find that a LoRA rank of 16 provides the best balance between
adaptation to Code and capability retention on the
English dataset.
We also test whether to apply LoRA to the feed-forward networks (FFN)
and the token embeddings in addition to the attention matrices.
We find that applying LoRA to the FFN slightly improves accuracy on
Code by 0.27% at the cost of a 0.86% decrease in accuracy on
English retention.

**Table 2: Ablation of LoRA configurations when adapting to the Code dataset.**
| Configuration | Accuracy |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Rank | FFN | Embeddings | Code | English | Average |
| 8 | N | N | 76.02 | 48.76 | 62.39 |
| 16 | N | N | 76.17 | 48.64 | 62.40 |
| 32 | N | N | 76.51 | 46.99 | 61.75 |
| 64 | N | N | 76.25 | 48.25 | 62.25 |
| 128 | N | N | 76.25 | 47.75 | 62 |
| 256 | N | N | 76.05 | 46.62 | 61.34 |
| 512 | N | N | 75.84 | 44.39 | 60.12 |
| 1024 | N | N | 76.04 | 43.91 | 59.97 |
| 2048 | N | N | 66.55 | 32.27 | 49.41 |
| 16 | Y | Y | 76.23 | 47.04 | 61.64 |
| 16 | Y | N | 76.44 | 47.78 | 62.11 |
| 16 | N | Y | 75.81 | 48.03 | 61.92 |
| 16 | N | N | 76.17 | 48.64 | 62.40 |

## Appendix B Ablation of MoDE Configurations

We investigate how MoDE hyperparameters impact accuracy.
In Table 3, we vary the number of blocks in a MoDE model and the number of expert transformer layers per block.
When adapting to the Code dataset, we find that the number of blocks impacts both domain-specific and retention performance, indicating that tuning the number of blocks impacts model performance.
We further find that performance on Code increases with the number of expert transformer layers which corresponds to the number of added parameters, but does not seem to significantly impact English retention.

**Table 3: Ablation of MoDE model configurations when fine-tuning on a coding dataset. We find that accuracy on the coding dataset increases with the number of expert layers. We select the underlined configuration with 3 blocks and a coding expert with 6 transformer layers to strike a balance between domain-specific accuracy, capability retention, and the number of parameters added.**
| Model Configuration | Accuracy |  |  |  |
| --- | --- | --- | --- | --- |
| Blocks | Expert Layers | Code | English Retention | Average |
| 2 | 1 | 76.98 | 47.87 | 62.43 |
| 2 | 2 | 77.17 | 47.86 | 62.51 |
| 3 | 1 | 77.26 | 47.49 | 62.38 |
| 3 | 2 | 77.51 | 47.95 | 62.73 |
| 6 | 1 | 77.79 | 47.61 | 62.70 |
| 6 | 2 | 77.96 | 47.65 | 62.80 |

## Appendix C Multi-Domain Data Efficiency

We explore how much multi-domain Math + Code data is required for strong performance on the target Math and Code datasets as well as English retention.
We follow a similar experimental setup to Section 4.3, but adapt the model to the multi-domain Math + Code dataset instead of the single-domain Code dataset.
A smaller amount of multi-domain data is desirable,
as it allows each domain owner to contribute minimal data while ensuring that their experts can collaborate effectively with those from other domains.

We compare against two MoDE variations alongside the baselines.
The first variation, Frozen MoDE, freezes experts during composition, while the second, Uninitialized MoDE, skips single-domain training entirely.
As shown in Figure 5, Frozen MoDE delivers the best performance when less mixture data is available. However, with more training data, the standard MoDE configuration achieves the highest accuracy across all domains.

Figure: Figure 5: Efficiency of mixture data. We examine how much mixture data is required for training. For left to right, we present the next token accuracy on Code, Math, and English. The x-axis for each subplot is the number of examples from Math + Code mixture data.
Refer to caption: x10.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. What is the time/memory cost regarding training the MoE model?

2. Have the authors conducted experiments on more tasks, and used more experts?

3. Can we analyze the selection probability for different experts given a certain task's input?

4. How to make a fair comparison with a baseline that has the size parameter size as the MoE model?
soundness: 2
strengths: 1. Using mixture-of-experts for multi-task learning is novel and interesting.

2. Experiemtal results on three datasets as well as the ablation studies show the effectiveness of the proposed approach.

3. The motivation is clear and Figure 1 is helpful to understand the proposed method. 

4. Related work are comprehensive to let me get the literature.
summary: This work aims to address domain-specific adaptation with pre-trained language models using a mixture-of-experts architecture that augments a general PLM with modular, domain-specialized experts. These experts are trained independently and composed together via a lightweight training process. For a multi-domain task, experts are modularly composed to enhance the model's performance. A lightweight fine-tuning process updates the experts and the gating function to improve the performance on the target task. Experiemntal results on three tasks shows the effectiveness of the proposed method.
weaknesses: 1. The selection probability of the gating function is unclear to me. How does the gating function select experts that represent different tasks when an input from certain task is fed into the model? It is crucial to understand how can the MoE model use the domain-specific knowledge from different experts.

2. Since the MoE will significantly increase the parameter size, e.g., in Table 1. It is unfair to compare the performance of MoE model which has more parameters with those have fewer parameters. 

3. The time cost in training a MoE model is very important. Althoght this work uses some model-parallel partitioning to accelerate the training. It is also unclear the time and memory cost within the proposed method.

4. Besides, the analysis of how the relationships between the number of experts and the top-k selection affects the performance is also miss in this work. Which has strong relationship to the efficiency and effectiveness of a MoE approach.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: - How are the experts initialized?
- What’s the relation between SPMD vs. tensor parallel, and MPMD vs. tensor parallel + data parallel?
- Typos: Line 016: “a general PLMs” → “general PLMs”
soundness: 2
strengths: - Simple idea of initializing specialized experts for specific domains and then merge them for overall performance improvement.
- Efficient implementation of modular expert training.
summary: Dealing with the multi-domain adaptation problem, this paper proposes a simple idea of combining domain-specialized experts as an MoE model (w/o sparse activation). The backbone model is frozen during training as a shared expert to avoid catastrophic forgetting on original domains. Experiments are conducted on three domains with token classification accuracy as the metric. Results show performance improvements.
weaknesses: My current ratings are based on the following weaknesses, and I am open to hearing from the authors if I'm wrong and re-evaulate the paper.

- Limited catastrophic forgetting verification. Only the *English* corpus is utilized for verifying the effectiveness of dealing with catastrophic forgetting (Code + Math → English). More setups would help understand the effectiveness. e.g., Code → Math + English, Math → Code + English, English → Code + Math, Code + English → Math, Math + English → Code.
- Potential unfair comparisons. Since MoDE is not sparsely activated, there are more available parameters participating in the computation **during inference**. It is unclear whether the performance improvement comes from the parameter scaling or the proposed method.
- Unclear settings of the pretrained 1.5B model and the multi-domain datasets.
    - The baseline performance of the pretrained 1.5B model is missing, and it is not clear how LoRA and full-param SFT would improve the performance. Do you train models from scratch on specific datasets or use a pretrained language model? If the 1.5B model is pretrained already, why not present its performance in Table 1? If the model is pretrained from scratch by your own, why not utilize a smaller pretrained model like LLaMA 3.0B or Phi? Are they too good to diminish the performance improvements?
    - It is not clear whether these datasets are seen during the pre-training stage.
    - Besides, if you utilize open-sourced datasets, please consider disclose details about the datasets. e.g., specific subsets of a open-sourced pre-training corpus.
- The paper proposes a similar idea of ModuleFormer and BTX, but they are not discussed or compared as baselines.
    - Shen, Y., Zhang, Z., Cao, T., Tan, S., Chen, Z., & Gan, C. (2023). Moduleformer: Learning modular large language models from uncurated data. *arXiv preprint arXiv:2306.04640*.
    - Sukhbaatar, S., Golovneva, O., Sharma, V., Xu, H., Lin, X. V., Rozière, B., ... & Li, X. (2024). Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM. *arXiv preprint arXiv:2403.07816*.
    - Maybe sparse upcycling should also be a baseline: Komatsuzaki, A., Puigcerver, J., Lee-Thorp, J., Ruiz, C. R., Mustafa, B., Ainslie, J., ... & Houlsby, N. (2022). Sparse upcycling: Training mixture-of-experts from dense checkpoints. arXiv preprint arXiv:2212.05055.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Are the precise details about the domain-specific datasets such as Code, Math, and English available?
2. Why the parameter count is same for Full-parameter Fine-tuning and LoRA in Table 1?
soundness: 2
strengths: 1. Considering the scale of current LLMs and the required cost to train such models, modular mixture-of-expert (MoE) architectures for multi-domain adaptation while maintaining the general-purpose capabilities of the backbone model, as proposed in this paper is quite important. 

2. Apart from performance, the paper also considers parallelization during training, which is another important aspects of current LLMs. Although the paper only shows for a particular infrastructure setting and only for training, the proposed parallelization results in an important speed-up.
summary: This paper introduces Modular Domain Experts (MoDE), a mixture-of-expert model that leverages separately trained domain experts. In this architecture, each MoDE expert consists of several transformer layers where experts process tokens in parallel to the backbone model and combined outputs per "blocks". In terms of training, MoDE includes (1) separately training domain experts in their domain data, and (2) fine-tuning all experts with unified data. The paper focuses on maximizing the performance of expert domains while maintaining the general performance of the backbone model in resource-constrained settings. In addition to the performance, the paper proposes parallelizing block execution, enabling speed-up during training.
weaknesses: 1. The paper disregards a large body of work such as [1], [2], [3], [4], [5] that are extremely relevant and similar to the proposed method in many aspects. Some of these works propose parameter-efficient MoE architecture, and others propose modular MoEs where experts are trained separately.

2. Connected to 1st point, the proposed MoDE architecture only compared with full model fine-tuning and simple LoRA methods while there are multiple baselines that should be considered such as BTX [1], BAM [2] or MoE-LoRA methods. 

3. For both full model fine-tuning and LoRA, only one variant where the backbone model is fine-tuned using the combined dataset (Code+Math) is presented. Another set of comparisons where the backbone model is fine-tuned per domain separately, and also fine-tuned using combined dataset but including English data as well required. 

4. It is unclear why the paper does not include any benchmark on Code, Math, and general language understanding in English. Only next-token prediction per domain is presented however, these results may not correlate with downstream benchmarks in the same degree. 

5. For the selected MoDE configuration, the additional number of parameters for only 2 domains is 0.789B which corresponds to approximately 50% of the pretrained parameter count (0.789/1.583). This is quite a high number of parameters when the number of new domains gets higher. 

[1] Li, Margaret, et al. "Branch-train-merge: Embarrassingly parallel training of expert language models." arXiv preprint arXiv:2208.03306 (2022).
[2] Zhang, Qizhen, et al. "BAM! Just Like That: Simple and Efficient Parameter Upcycling for Mixture of Experts." arXiv preprint arXiv:2408.08274 (2024).
[3] Zadouri, Ted, et al. "Pushing mixture of experts to the limit: Extremely parameter efficient moe for instruction tuning." arXiv preprint arXiv:2309.05444 (2023).
[4] Wu, Xun, Shaohan Huang, and Furu Wei. "Mixture of lora experts." arXiv preprint arXiv:2404.13628 (2024).
[5] Dou, Shihan, et al. "Loramoe: Revolutionizing mixture of experts for maintaining world knowledge in language model alignment." arXiv preprint arXiv:2312.09979 4.7 (2023).

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: The effectiveness of MoDE is sensitive to the configuration, yet the criteria for selecting configurations remain confusing. In Table 3,
 *  How do you reach the conclusion that configuring three MoDE blocks, each containing two transformer layers per expert strikes a balance? 
* The performance increases with the number of blocks increases, why do you select 3 instead of 6?
soundness: 3
strengths: 1. The proposed method demonstrates superior performance over existing methods like standard LoRA, excelling in knowledge retention capabilities and offering better scalability as expert parameters and training data increase.
2. The architecture allows for flexible integration with different sharding strategies, such as the MPMD-enabled sharding strategy, effectively accelerating training speeds.
summary: This paper introduces a mixture-of-experts (MoE) architecture, named Modular Domain Experts (MoDE), to efficiently adapt pre-trained language models to multi-domain tasks. The architecture is designed to optimize the trade-off between domain-specific performance, knowledge retention, and training efficiency. Additionally, the architecture supports integration with sharding strategies for improving computational efficiency.
weaknesses: 1. The baselines compared in this work are limited to standard LoRA and SFT. The concept of integrating modular domain expert is somewhat similar to prior papers on MoE and LoRA-MoE architecture (e.g., [1][2][3]), while the difference is MoDE applies transformer block as expert instead of feed-forward networks or LoRA block. The paper may need to add discussions of these differences to clarify its novelty. Furthermore, including other MoE or LoRA-MoE architectures as baselines would help demonstrate the superiority of the proposed approach.

2. Many low-level details are omitted, such as training and test datasets, hyperparameters of baseline methods and the two-stage training procedure. Including these details would help understand the method and facilitate future research endeavors. 

3. Although there is no consensus, models >10b parameters are generally considered LLMs. The backbone pretrain model in this paper is only 1.5b, it would be better to try multiple models with larger parameter sizes and apply other PLMs such as Phi and LLaMA to demonstrate scalability.

4. In Table 1, the baseline methods are trained on the mixed dataset, whereas MoDE is trained on Math and Code, respectively, and then on the mixed data. This seems unfair as it has more training steps.

[1] LoRAMoE: Alleviate World Knowledge Forgetting in Large Language Models via MoE-Style Plugin

[2] MOELoRA: An MOE-based Parameter Efficient Fine-Tuning Method for Multi-task Medical Applications

[3] MixLoRA: Enhancing Large Language Models Fine-Tuning with LoRA-based Mixture of Experts