# Mixture-of-Modules:
Reinventing Transformers as Dynamic Assemblies of Modules

Zhuocheng Gong1111Equal Contributions. Ang Lv2111Equal Contributions. Jian
Guan3,4222Corresponding authors. Junxi Yan3 Wei Wu4 Huishuai Zhang1
Minlie Huang3 Dongyan Zhao1222Corresponding authors. Rui Yan2222Corresponding
authors.
1Peking University 2Renmin University 3Tsinghua University 4Ant Group
<EMAIL_ADDRESS><EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS>
<EMAIL_ADDRESS>

###### Abstract

Is it always necessary to compute tokens from shallow to deep layers in
Transformers? The continued success of vanilla Transformers and their variants
suggests an undoubted “yes”. In this work, however, we attempt to break the
depth-ordered convention by proposing a novel architecture dubbed mixture-of-
modules (MoM), which is motivated by an intuition that any layer, regardless
of its position, can be used to compute a token as long as it possesses the
needed processing capabilities. The construction of MoM starts from a finite
set of modules defined by multi-head attention and feed-forward networks, each
distinguished by its unique parameterization. Two routers then iteratively
select attention modules and feed-forward modules from the set to process a
token. The selection dynamically expands the computation graph in the forward
pass of the token, culminating in an assembly of modules. We show that MoM
provides not only a unified framework for Transformers and their numerous
variants but also a flexible and learnable approach for reducing redundancy in
Transformer parameterization. We pre-train various MoMs using OpenWebText.
Empirical results demonstrate that MoMs, of different parameter counts,
consistently outperform vanilla transformers on both GLUE and XSUM benchmarks.
More interestingly, with a fixed parameter budget, MoM-large enables an over
38% increase in depth for computation graphs compared to GPT-2-large,
resulting in absolute gains of 1.4 on GLUE and 1 on XSUM. On the other hand,
MoM-large also enables an over 60% reduction in depth while involving more
modules per layer, yielding a 16% reduction in TFLOPs and a 43% decrease in
memory usage compared to GPT-2-large, while maintaining comparable
performance. 111Code is available at https://github.com/gzhch/MoM

##  1 Introduction

Transformer-based language models (Vaswani et al., 2017) have demonstrated
remarkable abilities across a wide range of challenging natural language tasks
(Bubeck et al., 2023). In addition, the success of Transformer in natural
language processing (NLP) is also inspiring innovations in other fields such
as computer vision (Peebles & Xie, 2023; Agostinelli et al., 2023) and
biomedicine (Singhal et al., 2023; Madani et al., 2023). A Transformer
architecture typically consists of stacked layers that are identical in
structure, whereby layers are organized in the order of depth, using the
output of the previous layer as the input for the next. While this design
convention has been widely accepted as a matter of course in the Transformer
era, we challenge it by reconsidering whether the static and depth-ordered
organization can fully unleash the potential of Transformers, given the well-
known issues of over-parameterization (Zeng et al., 2023) and efficiency
(Raposo et al., 2024).

Before us, some rudimentary studies have touched on the question–they dissect
Transformer into modules such as attention heads and feed-forward networks
(FFNs) and allow relatively flexible module call order. For example, Mixture-
of-Experts (MoE) (Shazeer et al., 2017)) sets up multiple FFNs within the same
layer and activates a specific subset during inference. Early-exiting (Zhou et
al., 2020; Xin et al., 2020; Schuster et al., 2022) and Mixture-of-Depths
(MoD) (Raposo et al., 2024)) bypass certain layers when computing each token.
On the one hand, these efforts indeed lead to improvements in terms of either
efficacy or efficiency through the introduction of dynamic mechanisms into the
vanilla structure of Transformers, and thus corroborate our questioning
regarding the established convention; on the other hand, they still follow the
depth-ordered paradigm (i.e., tokens are passed from shallow layers to deep
layers), leaving significant room for better architectures.

Figure 1: Mixture-of-Modules reinvents Transformers as dynamic assemblies of
modules. In (b), we illustrate the ongoing construction of an MoM model during
the forward computation. The assembly lasts $H$ rounds, with the current
illustration showcasing progress in the third round. For each token, routers
select the best $K$ attention modules, denoted as $m^{\text{A}}_{k}$, and the
best $K$ feed-forward network modules, denoted as $m^{\text{F}}_{k}$, from a
module set $\mathcal{M}$ (including “SKIP” modules). These selected modules
collectively constitute assembled modules $\mathcal{F}^{\text{A}}$ and
$\mathcal{F}^{\text{F}}$, which are then appended to the existing computation
graph. Detailed notations are presented in §3.

In this work, we completely disrupt the traditional practice in the design of
Transformers by breaking down the depth-ordered organization. Numerous studies
have indicated that knowledge in Transformers is often dispersed across
multiple FFNs in different layers (Geva et al., 2021; McGrath et al., 2023; Lv
et al., 2024), and many attention heads serve similar or identical functions,
such as copying specific token information towards the end position of the
input (Olsson et al., 2022; Wang et al., 2023). Encouraged by this evidence,
we pose the question of whether the computation of a token can “move” freely
across layers, that is the token can be computed by flowing to modules in
deeper layers, sticking to modules of the same layer, or even going back to
modules in previous layers. To answer the question, we propose a novel
architecture dubbed Mixture-of-Modules (MoM) in which the core idea is to
define a neural network as dynamic assemblies of modules derived from vanilla
Transformer, as depicted in Figure 1.

The basis of MoM is a finite set of modules. Each module is defined by multi-
head attention (MHA), a feed-forward network (FFN) (including Add & Norm), or
a specialized module labeled “SKIP”. Each MHA or FFN module is identical in
structure and different in parameterization. SKIP enables skip operations for
arbitrary tokens at arbitrary time steps. Given a token, each time two routers
select modules from the set and integrate the modules into the computation
graph during the forward pass. Hence, the whole computation graph of the token
is formed as an assembly of modules, and the routers learn to optimize the
organization of the modules in the assembly. We introduce a two-phase approach
for training MoM models. In the first phase, we pre-train a vanilla
Transformer on a large-scale corpus. Then, in the second phase, we decompose
the pre-trained Transformer into modules as a warm-up of MoM, randomly
initialize the routers, and further update both the modules and the routers
under the mechanism of dynamic assembly. By this means, we can both enhance
parameter utilization and accelerate the convergence of the model.

MoM has three major advantages over existing Transformer-based architectures:
(1) it provides a unified framework for various Transformer variants,
incorporating popular methods such as mixture-of-experts, early-exiting, and
mixture-of-depths as special cases. The framework sheds light on architecture
design for future works; (2) it brings unprecedented flexibility in forward
computation. With the dynamic assembly mechanism, “depth” and “parameter
count” are no longer entangled as they are in the conventional sense. One can
build powerful architectures by either enlarging the module pool (i.e.,
increasing parameter count) or increasing the depth (parameter count can be
fixed). Hence, MoM offers a dynamic and learnable approach to reducing
redundant parameters in Transformers; and (3) it offers efficient structures
that achieve performance comparable to vanilla Transformers but require
significantly fewer FLOPs and less memory in forward computation.

We pre-train MoM in three sizes–122M (small), 346M (medium), and 774M (large)–
using OpenWebText (Gokaslan & Cohen, 2019), and assess their performance with
GLUE (Wang et al., 2018a) and XSUM (Narayan et al., 2018a). Empirical results
indicate that (1) MoMs, across all the three parameter counts, consistently
outperform vanilla GPT-2 models on both text understanding and generation
tasks; (2) parameters are quite redundant in vanilla Transformers. One can
develop an MoM that is at least 30% deeper than a vanilla GPT-2, resulting in
at least 1.4 absolute gain in GLUE and at least 1 absolute gain in XSUM.
Moreover, one can further remove $50$% of the MHA modules and $25$% of the FFN
modules from the module pool of the model, while maintain comparable
performance222After removing $50$% of the MHA modules and $25$% of the FFN
modules, the validation loss increased by $0.065$, the GLUE average score
dropped from $81.98$ to $81.34$, and the ROUGE average in XSUM dropped from
$19.34$ to $18.55$.; and (3) for those concerned with efficiency, MoM-large
can reduce TFLOPs by 16% and memory usage by 42% in forward computation, while
maintaining comparable performance to GPT-2-large, via properly increasing the
number of modules and compressing the model depth.

Our contributions are three-fold: (1) proposal of Mixture-of-Modules to
disrupt the depth-ordered convention in Transformer construction, and reinvent
Transformers as dynamic assemblies of modules; (2) empirical verification of
the efficacy of Mixture-of-Modules on GLUE and XSUM; and (3) a series of new
insights into the over-parameterization issue of vanilla Transformers, and
their implications for future architecture design. The code for implementing
MoM is open-sourced at LABEL:??.

##  2 Related works

MoM owns a dynamic mechanism of module selection and combination, and thus is
related to conditional computation techniques (Bengio et al., 2013; Davis &
Arel, 2014; Cho & Bengio, 2014). Existing work on conditional computation can
be categorized into two groups: _dynamic depth_ and _dynamic width_. In these
fields, terms such as gating and routing are used interchangeably, hereafter
referred to as “routers” for clarity in presentation.

As a typical approach in dynamic depth, Early-exiting (Graves, 2016; Figurnov
et al., 2017; Schuster et al., 2022) accelerates model inference through
terminating forward computation at intermediate layers. The decision to exit
often relies on confidence-based metrics (Elbayad et al., 2020; Varshney et
al., 2023; Xin et al., 2020) or pre-determined strategies (Liu et al., 2020;
Corro et al., 2023). With some degree of generalization, Layer-skip
(Srivastava et al., 2015; Wang et al., 2018b; Bapna et al., 2020) represents a
more adaptive variant of early-exiting, enabling certain layers to be skipped
without terminating the entire forward computation. Existing works mainly
facilitate it by training a router (Zeng et al., 2023; Raposo et al., 2024) or
layer pruning (Yang et al., 2024; Kim et al., 2024). Finally, if we view
parameter copying as a particular way to increase network depth with
controlled model size, then some parameter sharing methods (Dehghani et al.,
2019; Lan et al., 2020), wherein certain modules or layers share parameters,
also fall in the dynamic depth group.

In terms of dynamic width, Mixture-of-Experts (MoE, (Shazeer et al., 2017;
Lepikhin et al., 2021; Fedus et al., 2022)) is a representative method. An MoE
model conceptualizes an FFN module as an “expert” for storing knowledge.
Comprising multiple such experts, an MoE layer replaces the traditional FFN
layer within Transformers, aiming for superior performance in handling
knowledge-related tasks. During forward computation, a router network
dynamically assigns each token to the top $K$ experts out of a total of $N$
experts, thereby increasing the maximum network width by $K$ times. Other
dynamic width methods, such as CODA (Lei et al., 2024) and CoLT5 (Ainslie et
al., 2023), use similar routing mechanism to select whether a token passes
through a heavy or light pathway for not only each FFN layer but also each
attention layer.

MoM breaks the depth-ordered paradigm followed by existing approaches when
performing forward computation. It not only unifies a number of approaches
described above but also offers a more flexible and learnable way to achieve
conditional computation.

##  3 Methodology

The idea of Mixture-of-Modules (MoM) is inspired by the theory presented in
“the society of mind” by Marvin Minsky (Minsky, 1986), which explains the true
intelligence as certain and very special ways of combinations of simple and
modular units (in the book, they are termed “agents”). In §3.1, we first
provide an overview of MoM. Then, we detail the assembly of modules and the
routers in §3.2 and §3.3. After that, we present the training procedure of MoM
in §3.4. Finally in §3.5, we show that MoM unifies various techniques of
dynamic computation allocation within Transformers as special cases.

###  3.1 Mixture-of-Modules (MoM)

Before delving into the details, we first give a brief description of the
workflow of MoM. MoM views the construction of an $H$-depth transformer as an
$H$-step iterative assembly process. In each assembly step, router
$\mathcal{R}$ dynamically selects $K$ modules from a module set $\mathcal{M}$
for each token. Then these selected modules are assembled guided by the
assembling function $\phi$. Formally, MoM can be defined by a 5-tuple
$<\mathcal{M},\mathcal{R},\phi,K,H>$.

$\mathcal{M}$ is the set that contains all possible modules, where modules are
defined as atomic units that could be assembled. There are two types of
modules in a Transformer model, _i.e._ , the multi-head self-attention module
(MHA) and the feed-forward network module (FFN), denoted as $m^{\text{A}}$ and
$m^{\text{F}}$, respectively. In addition, denote the input hidden state of a
Transformer layer as $\mathbf{x}\in\mathbb{R}^{d}$, we include a special
module $m^{\text{S}}:\mathbf{x}\mapsto\mathbf{x}$ in $\mathcal{M}$, which
means the absence of an operation applied to the token, allowing for skipping
one round of computation. Therefore, in MoM, we have:

$\displaystyle\mathcal{M}=\\{m^{\text{A}}_{i}\\}_{i=1}^{N_{\text{A}}}\cup\\{m^{\text{F}}_{i}\\}_{i=1}^{N_{\text{F}}}\cup\\{m^{\text{S}}\\},$
(1)

where $N_{\text{A}}$ are $N_{\text{F}}$ refer to numbers of MHAs and FFNs,
respectively.

$\mathcal{R}$ is a router responsible for dynamically selecting appropriate
modules from $\mathcal{M}$ and assembling them into the computation graph. We
use distinct routers for MHAs and FFNs, denoted as $\mathcal{R}^{\text{A}}$
and $\mathcal{R}^{\text{F}}$ respectively. The output of the router
$\mathcal{R}^{\mathcal{X}}$ is an $(N_{\mathcal{X}}+1)$-dimensional
distribution wherein each item represents the weight assigned to each module
$m^{\mathcal{X}}_{i}$ as well as $m^{\text{S}}$. Formally:

$\displaystyle\mathcal{R}^{\mathcal{X}}:\quad\mathbf{x}\mapsto\mathbf{r}^{\mathcal{X}},$
(2)
$\displaystyle\mathbf{x}\in\mathbb{R}^{d},\quad\mathbf{r}^{\mathcal{X}}\in\mathbb{R}^{N_{\mathcal{X}}+1},\quad\mathcal{X}\in\\{\text{A},\text{F}\\}.$

An MoM model is dynamically assembled step by step. In each step, based on the
output of $\mathcal{R}_{\mathcal{X}}$, $K_{\mathcal{X}}$ modules are selected
from $\mathcal{M}$ and assembled together. The assembly process lasts $H$
steps, as detailed in §3.2. Further elaboration on these routers, including
their architecture and the working pipeline, is deferred to §3.3. Notably, in
MoM, dynamic assembly occurs at the token level, wherein each token is
independently and dynamically assigned by routers to appropriate modules for
processing.

###  3.2 Dynamic assembly of modules

We delve into how an MoM model is dynamically assembled. The construction is
an iterative process where in the $h$-th step (i.e., the $h$-th layer of the
model being constructed), we have the input $\mathbf{x}_{h}$. The subscript
$h$ is omitted when there’s no ambiguity. The router $\mathcal{R}$ selects $K$
modules with the largest routing weight. We denote the indices of the selected
modules as
$\mathcal{K}_{\mathcal{X}}=\\{i|r_{i}\in\text{TopK}(\mathbf{r}^{\mathcal{X}})\\}$.
Then the selected modules are assembled together through the assembling
function $\phi$. Formally,

$\displaystyle\phi:\quad<\mathcal{M},\mathcal{R}^{\mathcal{X}},\mathbf{x}_{h}>\
\mapsto\
\mathcal{F}^{\mathcal{X}}_{h},\quad\mathcal{X}\in\\{\text{A},\text{F}\\},$ (3)

where $\mathcal{F}^{\mathcal{X}}_{h}$ represents the assembled modules. These
assembled modules transform the input $\mathbf{x}_{h}$ into the output
$\mathbf{x}_{h+1}$. We hope the role of the $h$-th step of MoM assembly is
somewhat akin to the $h$-th Transformer block in the conventional sense.
Therefore, we establish two rounds of routing and assembling in each assembly
step: one for MHA and the other for FFN. The forward computation of MoM models
at the $h$-th step assembly can be represented as:

$\displaystyle\mathbf{u}_{h}=\mathcal{F}^{\text{A}}_{h}(\mathbf{x}_{h})+\mathbf{x}_{h},$
(4)
$\displaystyle\mathbf{x}_{h+1}=\mathcal{F}^{\text{F}}_{h}(\mathbf{u}_{h})+\mathbf{u}_{h}.$

We employ Pre-norm in MoM, which normalizes the input before feeding to
assembled modules $\mathcal{F}^{\mathcal{X}}$. The dynamic assembly process is
depicted in Figure 1(b). We now introduce the detailed formalization for
$\mathcal{F}^{\text{A}}$ and $\mathcal{F}^{\text{F}}$, respectively.

Assembly of attention modules ($\mathcal{F}^{\text{A}}$). We begin by
considering the scenario where the $m^{\text{S}}$ module is not selected by
routers. Suppose that an MHA module contains $Z$ individual heads, then the
assembly of $K_{\text{A}}$ MHA modules (i.e., the computation process of
$\mathbf{o}=\mathcal{F}^{\text{A}}(\mathbf{x})$) is defined as:

$\displaystyle\mathbf{o}=\mathbf{a}\sum_{k\in\mathcal{K}_{\text{A}}}r^{\text{A}}_{k}\cdot\mathbf{W}_{k}^{O},$
(5)
$\displaystyle\mathbf{r}^{\text{A}}=\mathcal{R}^{\text{A}}(\mathbf{x})=\left(r^{\text{A}}_{1},\ldots,r^{\text{A}}_{k},\ldots,r^{\text{A}}_{K_{A}}\right),$
$\displaystyle\mathbf{a}=\left(\mathbf{x}\sum_{k\in\mathcal{K}_{\text{A}}}\mathbf{W}_{k,z}^{V}\right)\cdot$
$\displaystyle\texttt{softmax}\left(\frac{(\mathbf{X}\sum_{k\in\mathcal{K}_{\text{A}}}\mathbf{W}_{k,z}^{Q})(\mathbf{X}\sum_{k\in\mathcal{K}_{\text{A}}}\mathbf{W}_{k,z}^{K})^{\top}}{\sqrt{d_{\text{head}}}}\right),$

where $\mathbf{X}\in\mathbb{R}^{L\times d}$ is the input representation of the
sequence,
$\mathbf{W}_{k,z}^{Q},\mathbf{W}_{k,z}^{K},\mathbf{W}_{k,z}^{V}\in\mathbb{R}^{d\times
d_{\text{head}}}$, and $\mathbf{W}_{k}^{O}\in\mathbb{R}^{d_{\text{head}}\times
d}$ are weight matrices with $d_{\text{head}}=d/Z$. When the $m^{\text{S}}$
module is selected, the operation of $\mathcal{F}^{\text{A}}$ only involves
the remaining $K_{\text{A}}-1$ attention modules.

Assembly of feed-forward networks ($\mathcal{F}^{\text{F}}$). The assembly of
$\mathcal{F}^{\text{F}}$ is more modular, where the outputs of $K_{\text{F}}$
modules are simply weighted and aggregated. When the $m^{\text{S}}$ module is
not selected, $\mathcal{F}^{\text{F}}$ can be formalized as follows:

$\displaystyle\mathcal{F}^{\text{F}}(\mathbf{u}):=\sum_{k\in\mathcal{K}_{\text{F}}}r^{\text{F}}_{k}\cdot
m^{\text{F}}_{k}(\mathbf{u}),$ (6)
$\displaystyle\mathbf{r}^{\text{F}}=\mathcal{R}^{\text{F}}(\mathbf{x})=\left(r^{\text{F}}_{1},\ldots,r^{\text{F}}_{k},\ldots,r^{\text{F}}_{K_{F}}\right).$

When $m^{\text{S}}$ is chosen, likewise, only $K_{\text{F}}-1$ FFNs form the
$\mathcal{F}^{\text{F}}$.

###  3.3 MoM router ($\mathcal{R}$)

In prior approaches, routing occurs as a one-step decision-making process
within a layer. However, in MoM which possesses a dynamically constructed
computation graph, each decision is interdependent with the preceding ones,
influencing the entire forward computation. Consequently, the router in MoM
necessitates an awareness of past decisions. To model such dependency, we
employ a gated recurrent unit (GRU, (Cho et al., 2014)) as the backbone of
routers. Two routers in MoM are identical in structure. At each assembly step,
the GRU in the $\mathcal{R}^{\mathcal{X}}$ maintains an
$\mathbf{s}^{\mathcal{X}}_{h}$ as the hidden state of the GRU network. This
state is recurrently updated as follows:

$\displaystyle\mathbf{s}^{\text{A}}_{h}=\texttt{GRU}^{\text{A}}(\mathbf{x}_{h},\mathbf{s}^{\text{A}}_{h-1}),$
(7)
$\displaystyle\mathbf{s}^{\text{F}}_{h}=\texttt{GRU}^{\text{F}}(\mathbf{u}_{h},\mathbf{s}^{\text{F}}_{h-1}).$

The weights assigned to each module by $\mathcal{R}^{\mathcal{X}}$ are
computed as:

$\displaystyle\mathbf{r}^{\mathcal{X}}=\mathbf{W}_{\mathcal{X}}\mathbf{s}^{\mathcal{X}}_{h},$
(8)
$\displaystyle\mathbf{W}_{\mathcal{X}}\in\mathbb{R}^{(N_{\mathcal{X}}+1)\times
d},\quad\mathcal{X}\in\\{\text{A, F}\\}.$

###  3.4 Training approach

A straightforward approach is to pre-train an MoM model initialized from
scratch. This approach, however, suffers from a degeneration issue, as the
learned functions of modules become homogeneous, making router training
challenging. To address the issue, we propose a two-phase training approach.
In the first phase, we pre-train a vanilla Transformer where modules acquire
distinct functionalities. Then, in the second phase, we initialize the module
set $\mathcal{M}$ with the pre-trained modules and initialize the routers from
scratch. Subsequently, we continue training both modules and routers using the
same data and objective as in the first phase. Through empirical studies, we
find that the two-phase training method improves the specialization of module
functionalities and accelerates router convergence.

###  3.5 MoM as a unified framework

A compelling property of MoM is that it unifies a wide range of Transformer-
based dynamic computation allocation architectures. With specific
configurations, layer-skip (e.g., early-exiting, mixture-of-depths, etc.),
parameter sharing, and mixture-of-experts can be viewed as special cases.

#### Layer-skip.

The key idea is to skip layers according to certain criteria which can either
be defined heuristically Liu et al. (2024) or learned from data Zeng et al.
(2023); Raposo et al. (2024). Within the MoM framework, layer-skip can be
formulated as a special cluster of assembly functions $\phi$, namely:

$\displaystyle\phi_{\text{layer-
skip}}(\mathcal{M},\mathcal{R}^{\mathcal{X}},\mathbf{x}_{h})=\begin{cases}m_{h}^{\mathcal{X}}&\text{if}\
c_{\text{skip}}(h)=1\\\ m^{S}&\text{if}\ c_{\text{skip}}(h)=0\end{cases},$ (9)

where $c_{\text{skip}}(\cdot)$ is the criterion that decides whether to skip
the $h$-th layer or not. Note that the technique of early-exiting Graves
(2016); Figurnov et al. (2017) can be viewed as a special case of layer-skip,
where once a layer is skipped, all subsequent layers will be skipped too.

#### Parameter sharing.

We consider parameter sharing that shares weights across modules and does not
involve reparameterization techniques. Under this restriction, the sharing
paradigm can be defined as a criterion function $c_{\text{share}}:i\mapsto j\
(j\leq i)$, representing using the same weights for module
$m_{i}^{\mathcal{X}}$ and module $m_{j}^{\mathcal{X}}$. Within MoM, parameter
sharing can be formulated as:

$\displaystyle\phi_{\text{parameter-
sharing}}(\mathcal{M},\mathcal{R}^{\mathcal{X}},\mathbf{x}_{h})=m_{c_{\text{share}}(h)}^{\mathcal{X}}.$
(10)

#### Mixture-of-Experts.

MoE splits FFN into experts, and the experts are not shared across different
layers. In MoE, routing is only performed on FFN modules, thus the computation
of MHA is the same as that of a vanilla Transformer. The assembly function for
MoE can be written as:

$\displaystyle\phi_{\text{MoE}}(\mathcal{M},\mathcal{R}^{\mathcal{X}},\mathbf{x}_{h})=\begin{cases}m_{h}^{\text{A}}&\text{if}\
\mathcal{X}=\text{A}\\\
\phi_{\text{MoM}}(\mathcal{M}_{h},\mathcal{R}^{\text{F}},\mathbf{x}_{h})&\text{if}\
\mathcal{X}=\text{F}\\\ \end{cases},$ (11)

where $\mathcal{M}_{h}=\\{m_{h,i}^{\text{F}}\\}_{i=1}^{N_{\text{F}}}$ is the
collection of experts for layer $h$.

Figure 2 illustrates the forward computation process across different methods,
offering an intuitive presentation of the versatility and universality in MoM.

Figure 2:  Visualization of forward computation of five models, where each
consists of only two layers just for demonstration purposes. The switch icon
symbolizes the selective execution of one (in Layer-skip) or more (in MoE and
MoM) subsequent computation pathways.

##  4 Experiments

###  4.1 Experimental setup

We implement language models with MoM since language modeling is a challenging
task requiring both language understanding and generation ability, thereby
effectively evaluating MoM and baselines. Below, we elaborate on the
implementation details of MoM, baselines, and evaluation setups.

#### Implementation details.

We conduct experiments on three model scales, which we denote as MoM-small,
MoM-medium, and MoM-large, respectively. These models contain 122M, 346M, and
774M parameters, respectively. Detailed configurations can be found in
Appendix A. The vanilla Transformers used for initializing MoM are official
GPT-2 checkpoints downloaded from HuggingFace333https://huggingface.co/openai-
community/gpt2. $K$ and $H$ represent two hyper-parameters of MoM. We denote a
configuration where $K=a,H=b$ as K$a$H$b$. If the skip module is included in
$\mathcal{M}$, we append the suffix S to K$a$H$b$.

In the two-phase training, we exploit OpenWebText (Gokaslan & Cohen, 2019) as
the pre-training dataset, and pre-process the data with the same pipeline as
nanoGPT (karpathy, 2023). OpenWebText contains $9$ billion tokens after
tokenization, from which $4$ million tokens are randomly sampled as the
validation set. The training sequence length for every input is $1,024$. We
set the learning rate to $1$e-$3$ with a warm-up ratio of 0.1 throughout the
two phases, and do not use dropout. All models are trained on 8 $\times$ A100
GPUs with a total batch size of 8 $\times$ 64\. Two training phases require
$20$k and $10$k optimization steps, respectively.

In practice, considering the large search space characterized by $H$ and $K$,
we confine the architecture search space to a practical scale with a
“chunking” strategy. An MoM is divided into several chunks. Each chunk is
independent and parameterized with identical $H$ and $K$. We present a
detailed description and specific configuration in Appendix A, and an
empirical analysis of the efficacy of chunking in Appendix B.

Methods | MoM Config | Parameter | Computation | Memory | Validation | Validation | GLUE | XSUM
---|---|---|---|---|---|---|---|---
$\phi$ | (K$a$H$b$) | Count | Cost (TFLOPs) | Cost (Gb) | Loss | Perplexity | (Average) | (Average)
small
GPT2 | K1H4 | 122M | 2.92 | 2.98 | 3.10 | 22.22 | 75.32 | 14.26
MoD | K1H4S | 122M | - | - | 3.22 | 25.11 | 72.24 | 9.71
MoE | K2H4 | 283M | 3.81 (+30.5%) | 2.98 (+0.0%) | 3.07 | 21.18 | 77.25 | 14.18
MoE (share) | K2H4 | 122M | 3.81 (+30.5%) | 2.98 (+0.0%) | 3.14 | 23.41 | 75.82 | 14.15
$\text{MoM}_{\text{E}}$ | K3H1S | 122M | 2.45 (-16.1%) | 2.45 (-17.8%) | 3.16 | 23.59 | 75.92 | 14.17
$\text{MoM}_{\text{I}}$ | K3H2S | 122M | 3.49 (+19.5%) | 2.63 (-11.7%) | 3.03 | 20.79 | 77.81 | 14.24
$\text{MoM}_{\text{P}}$ | K2H6S | 122M | 5.04 (+72.6%) | 3.34 (+12.1%) | 2.98 | 19.59 | 78.22 | 15.19
medium
GPT2 | K1H4 | 346M | 8.28 | 4.74 | 2.81 | 16.69 | 80.49 | 18.14
MoD | K1H4S | 346M | - | - | 2.99 | 19.82 | 76.17 | 14.81
MoE | K2H4 | 921M | 11.37 (+37.3%) | 4.74 (+0.0%) | 2.80 | 16.53 | 80.47 | 17.75
MoE (share) | K2H4 | 346M | 11.37 (+37.3%) | 4.74 (+0.0%) | 2.82 | 16.81 | 80.35 | 17.59
$\text{MoM}_{\text{E}}$ | K3H1S | 346M | 6.80 (-17.9%) | 3.33 (-29.7%) | 2.83 | 16.91 | 80.41 | 17.11
$\text{MoM}_{\text{I}}$ | K3H2S | 346M | 10.20 (+23.2%) | 3.80 (-19.8%) | 2.77 | 15.89 | 81.03 | 18.66
$\text{MoM}_{\text{P}}$ | K2H6S | 346M | 16.23 (+96.0%) | 5.69 (+20.0%) | 2.72 | 15.18 | 81.93 | 19.30
large
GPT2 | K1H4 | 774M | 17.76 | 7.20 | 2.66 | 14.33 | 84.47 | 20.35
MoD | K1H4S | 774M | - | - | 2.81 | 16.62 | 81.49 | 18.62
MoE | K2H4 | 2100M | 25.43 (+43.2%) | 7.20 (+0.0%) | 2.64 | 14.17 | 84.43 | 20.63
MoE (share) | K2H4 | 774M | 25.43 (+43.2%) | 7.20 (+0.0%) | 2.65 | 14.22 | 83.83 | 20.39
$\text{MoM}_{\text{E}}$ | K3H1S | 774M | 14.84 (-16.4%) | 4.13 (-42.6%) | 2.66 | 14.50 | 83.39 | 20.46
$\text{MoM}_{\text{I}}$ | K3H2S | 774M | 20.31 (+14.5%) | 5.15 (-28.5%) | 2.64 | 13.92 | 84.49 | 21.73
$\text{MoM}_{\text{P}}$ | K2H6S | 774M | 36.07 (+103.1%) | 9.24 (+28.3%) | 2.60 | 13.21 | 85.90 | 22.36

Table 1: Comprehensive comparison between MoMs and baselines. We highlight the
best results in bold and underline the second-best results. Appendix C
includes the detailed performance on GLUE and XSUM.

#### Baselines.

In addition to the vanilla Transformer model (Radford et al., 2019), the
following models (or methods) are also employed as baselines: (1) MoD (Raposo
et al., 2024): a layer-skip method proposed recently that dynamically routes
around Transformer blocks.444As official code is unavailable until the
submission, we follow the paper to implement MoD ourselves. (2) MoE: the
mixture-of-experts architecture utilized by Mixtral. We implement the model
with the open-sourced
code.555https://github.com/huggingface/transformers/blob/v4.36.1/src/transformers/models/mixtral/modeling_mixtral.py
(3) MoE (share): a variant of MoE in which all layers share the same set of
experts. We involve this model as a baseline because unlike the standard MoE
that has more parameters, MoE-share has the same number of parameters with the
vanilla Transformer model, making the comparison more fair. Moreover, it also
sheds light on how well the MoE architecture can utilize a fixed budget of
modules.

Note that all the above-mentioned methods are special cases within the MoM
framework with various configurations. Furthermore, we explore a wide range of
MoM instances defined by K$a$H$b$S, where $a\leq 4$ and $b\leq 6$. We examine
all MoM instances within the search space (as detailed in §4.3), and spotlight
three distinct models for comparison against other baselines:

$\bullet$ $\text{MoM}_{\text{P}}$ (K2H6S) represents a Performant MoM model
after tuning $K$ and $H$.

$\bullet$ $\text{MoM}_{\text{E}}$ (K3H1S) significantly enhances Efficiency
compared to vanilla Transformers, while maintaining acceptable performance.

$\bullet$ $\text{MoM}_{\text{I}}$ (K3H2S) serves as a midpoint in
configurations between the two preceding models. This model is positioned
between $\text{MoM}_{\text{E}}$ and $\text{MoM}_{\text{P}}$. We aim to
highlight performance and efficiency Interpolation as feature of MoM through
configuration interpolation.

#### Evaluation settings.

We employ GLUE benchmark (Wang et al., 2018a) to evaluate the language
understanding ability and XSUM (Narayan et al., 2018b) to evaluate the text
generation ability. All models are fine-tuned with a learning rate of
$2$e-$5$. The sequence is $128$ for GLUE and $1024$ for XSUM. For smaller GLUE
sub-datasets (CoLA, STS-B, MRPC, and RTE), we set the batch size to 32 and
train for 3 epochs. For larger datasets (MNLI, QNLI, QQP, and SST-2), we
utilize a batch size of 64 and perform training for a total of $8,000$
gradient steps. For XSUM, we set the batch size to 64 and train for 3 epochs.
For efficiency evaluation, we report inference TFLOPs and memory usage. TFLOPs
are calculated using DeepSpeed FLOPs profiler (DeepSpeed, 2023) and memory
consumption is calculated with PyTorch toolkits (pytorch, 2023).

###  4.2 Main results

Table 1 reports the evaluation results. Our analysis yields the following
conclusions:

#### MoM unleashes the potential of Transformers and our initial motivation is
confirmed.

When maintaining the number of parameters, $\text{MoM}_{\text{P}}$ is
characterized by the deepest computation graph ($H$). Across all model scales,
$\text{MoM}_{\text{P}}$ consistently outperforms all baselines on both GLUE
and XSUM by significant margins. The enhanced performance of
$\text{MoM}_{\text{P}}$ validates our initial motivations: (1) the traditional
depth-ordered layer organization is sub-optimal; (2) improvements can be
realized through two key modifications to the computation graph, including
dynamic module organization and improved parameter utilization.

$\text{MoM}_{\text{E}}$ is characterized by its minimum depth ($H$). By
strategically selecting appropriate modules at each assembly step,
$\text{MoM}_{\text{E}}$ strives to reduce memory and computation costs while
maintaining performance. Although $\text{MoM}_{\text{E}}$ is slightly
surpassed by a vanilla Transformer, it outperforms MoD, another efficiency-
driven method by large margins.

Besides, we observe that $\text{MoM}_{\text{I}}$ archives a decent performance
by slightly outperforming vanilla GPT-2. Comparing to vanilla GPT-2,
$\text{MoM}_{\text{I}}$ consumes no more than 25% extra computation but save
at least 11.7% memory across all scales, indicating that
$\text{MoM}_{\text{I}}$ achieves a good balance between performance and
efficiency.

Figure 3: How validation loss varies with respect to $N_{\text{A}}$ and
$N_{\text{F}}$, comparing to MoM (medium) with $N_{\text{A}}=N_{\text{F}}=4$.

#### MoM models provide insights into the over-parameterization issue.

With GPT-2-medium as initialization, we develop a series of MoM models, each
defined by different pairs of ($N_{\text{A}},N_{\text{F}}$), with both values
not exceeding $4$. We assess the validation loss increase for each model
relative to the benchmark model where $N_{\text{A}}=4$ and $N_{\text{F}}=4$,
as illustrated in Figure 3. In this experiment, we set $K$ equal to the number
of modules to make full use of the module parameters. Interestingly, the FFN
and MHA modules exhibit different degrees of redundancy. Specifically, when
$N_{\text{F}}$ is fixed and the number of MHAs is gradually reduced, a
significant increase in loss is not observed until $N_{\text{A}}$ is reduced
from $2$ to $1$, suggesting considerable redundancy in the MHAs of
Transformers. In contrast, when fixing $N_{\text{A}}$ and reducing the number
of FFNs gradually, each time of removing an FFN leads to evident loss
increase, indicating FFNs are less over-parameterized. These quantitative
findings align with previous research suggesting that the parameterization of
attention can be simplified to enhance efficiency while maintaining
performance (DeepSeek-AI, 2024; Shazeer, 2019).

#### As the parameter size scales up, MoM models enjoy consistent gains in
both performance and efficiency.

When we look into the difference across different scales, we observe that (1)
the performance gain of MoM is stable; (2) $\text{MoM}_{\text{E}}$-medium and
$\text{MoM}_{\text{E}}$-large exhibit more significant reductions in resource
costs comparing to $\text{MoM}_{\text{E}}$-small. These observations across
different scales reinforce our previous motivation: Transformers are over-
parameterized, which becomes more evident as the model size increases.

###  4.3 Insights from hyper-parameter search

Figure 4 shows how the validation loss for MoM-small and MoM-medium varies
with respect to different settings of $K$ and $H$ ($K\in\\{1,2,3,4\\}$,
$H\in\\{1,2,3,4,5,6\\}$). From this experiment, we have the following
observations and insights: (1) allowing more modules to be assembled at each
step (i.e., larger $K$) and more rounds of assembling actions (i.e., larger
$H$) generally leads to better performance, indicating that Transformer-based
models benefit from a larger computation graph even if the parameter size
remains the same. However, (2) the benefits of increasing $K$ and $H$ become
marginal when $K>2$ and $H>1$. Comparing K3H1 to K2H6, we can see that the
validation loss is comparable, while K3H1 performs slightly worse on
downstream tasks as discussed in §4.2. However, K3H2 improves efficiency by
flattening the depth, making it a good choice that balances performance and
efficiency. Flattening modules from different depths to the same depth cancels
computation dependencies of each other. This characteristic brings an extra
benefit because the computation of modules from the same depth can be
parallelized. This technique has been validated and adopted in MoE
applications (Fedus et al., 2022; Lepikhin et al., 2021) (called expert
parallelism) and can be easily extended to further accelerate MoM (K3H2).

Figure 4: Validation loss for MoM-small and MoM-medium under different settings of $K$ and $H$. | Phase 1 | Phase 2 | Val. Loss | PPL.
---|---|---|---|---
(1) | MoM | - | 2.85 | 17.26
(2) | Vanilla | - | 2.81 | 16.69
(3) | MoM | MoM | 2.78 | 16.20
(4) | Vanilla | MoM | 2.72 | 15.18
(5) | Vanilla | MoM-same | 3.13 | 22.87
Table 2: Performance comparison under different training setups. When training
MoM from scratch, we set the total gradient steps of phase 1 to 20k. The total
steps of phase 2 for all ablations are 10k.

###  4.4 Impact of two-phase training

We investigate how the two-phase training strategy influences model
performance using $\text{MoM}_{\text{P}}$-medium as case studies.
Specifically, we consider the following training strategies: (1) training MoM
from scratch with 20k steps; (2) training vanilla Transformer with 20k steps;
(3) training MoM from scratch with 30k steps; (4)
$\text{MoM}_{\text{P}}$-medium that is trained with the two-phase strategy;
and (5) training MoM with the two-phase strategy but in the second phase, MHA
modules and FFN modules are forced to be identical, respectively. Table 2
shows the results. First, comparing (1) with (2), we find that the absence of
weight initialization compromises the training quality of MoM, making it worse
than the vanilla Transformer, underscoring the importance of initializing
module weights with a well-trained vanilla Transformer model for MoM. A
plausible explanation for the results is that training without warm-up leads
to homogeneous modules that hurt the convergence of the routers (as
illustrated in Appendix D). The explanation is further justified by comparison
between (4) and (5), as when we force all modules to be identical, the
performance of MoM also drops dramatically even with the two-phase training
strategy. As expected, training with more steps can enhance performance (cf.,
comparison between (1) and (3)), but the two-phase strategy is still the
better choice when we compare (3) with (4).

##  5 Conclusions

In this work, we propose Mixture-of-Modules (MoM), a novel architecture that
reinvents transformers as a collection of individual modules and the dynamic
assembly process conducted with these modules. This novel view offers us an
opportunity to explore a wide range of different configurations of model
architecture and unify a series of transformer variants. With exhaustive
experiments, we not only validate the effectiveness of MoM by both significant
efficiency and performance gains but also reach new insights about
Transformers.

## Limitations

Our current design of the router still has room for improvement. Unlike MoE
wherein the router makes one-time decisions about which experts to select, the
router of MoM is responsible for conducting multi-step decision-making. In
this scenario, instructing the router to make correct decisions continuously
is a hard problem since the decision space grows exponentially with the
increase of assembly steps. The current implementation has not considered this
question and has not explicitly encouraged or discouraged the router to make
some choices, thus, we are not sure whether the learned routing decisions are
optimal or not. In the future, we will explore using techniques like
reinforcement learning or neural architecture search to design more
sophisticated routers.

## References

  * Agostinelli et al. (2023) Andrea Agostinelli, Timo I. Denk, Zalán Borsos, Jesse Engel, Mauro Verzetti, Antoine Caillon, Qingqing Huang, Aren Jansen, Adam Roberts, Marco Tagliasacchi, Matt Sharifi, Neil Zeghidour, and Christian Frank.  Musiclm: Generating music from text, 2023.
  * Ainslie et al. (2023) Joshua Ainslie, Tao Lei, Michiel de Jong, Santiago Ontañón, Siddhartha Brahma, Yury Zemlyanskiy, David Uthus, Mandy Guo, James Lee-Thorp, Yi Tay, et al.  Colt5: Faster long-range transformers with conditional computation.  _arXiv preprint arXiv:2303.09752_ , 2023.
  * Bapna et al. (2020) Ankur Bapna, Naveen Arivazhagan, and Orhan Firat.  Controlling computation versus quality for neural sequence models, 2020.
  * Bengio et al. (2013) Yoshua Bengio, Nicholas Léonard, and Aaron Courville.  Estimating or propagating gradients through stochastic neurons for conditional computation, 2013.
  * Bubeck et al. (2023) Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, Harsha Nori, Hamid Palangi, Marco Tulio Ribeiro, and Yi Zhang.  Sparks of artificial general intelligence: Early experiments with gpt-4, 2023.
  * Cho & Bengio (2014) Kyunghyun Cho and Yoshua Bengio.  Exponentially increasing the capacity-to-computation ratio for conditional computation in deep learning, 2014.
  * Cho et al. (2014) Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio.  Learning phrase representations using RNN encoder–decoder for statistical machine translation.  In Alessandro Moschitti, Bo Pang, and Walter Daelemans (eds.), _Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pp. 1724–1734, Doha, Qatar, October 2014. Association for Computational Linguistics.  doi: 10.3115/v1/D14-1179.  URL https://aclanthology.org/D14-1179.
  * Corro et al. (2023) Luciano Del Corro, Allie Del Giorno, Sahaj Agarwal, Bin Yu, Ahmed Awadallah, and Subhabrata Mukherjee.  Skipdecode: Autoregressive skip decoding with batching and caching for efficient llm inference, 2023.
  * Davis & Arel (2014) Andrew Davis and Itamar Arel.  Low-rank approximations for conditional feedforward computation in deep neural networks, 2014.
  * DeepSeek-AI (2024) DeepSeek-AI.  Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024.
  * DeepSpeed (2023) DeepSpeed.  Deepspeed flops profiler, 2023.  URL https://www.deepspeed.ai/tutorials/flops-profiler/.
  * Dehghani et al. (2019) Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser.  Universal transformers.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=HyzdRiR9Y7.
  * Elbayad et al. (2020) Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli.  Depth-adaptive transformer, 2020.
  * Fedus et al. (2022) William Fedus, Barret Zoph, and Noam Shazeer.  Switch transformers: scaling to trillion parameter models with simple and efficient sparsity.  _J. Mach. Learn. Res._ , 23(1), jan 2022.  ISSN 1532-4435.
  * Figurnov et al. (2017) Michael Figurnov, Maxwell D. Collins, Yukun Zhu, Li Zhang, Jonathan Huang, Dmitry Vetrov, and Ruslan Salakhutdinov.  Spatially adaptive computation time for residual networks.  In _2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 1790–1799, 2017.  doi: 10.1109/CVPR.2017.194.
  * Geva et al. (2021) Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy.  Transformer feed-forward layers are key-value memories.  In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih (eds.), _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pp. 5484–5495, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.emnlp-main.446.  URL https://aclanthology.org/2021.emnlp-main.446.
  * Gokaslan & Cohen (2019) Aaron Gokaslan and Vanya Cohen.  Openwebtext corpus.  http://Skylion007.github.io/OpenWebTextCorpus, 2019.
  * Graves (2016) Alex Graves.  Adaptive computation time for recurrent neural networks, 2016.
  * karpathy (2023) karpathy.  nanogpt, the simplest, fastest repository for training/finetuning medium-sized gpts., 2023.  URL https://github.com/karpathy/nanoGPT.
  * Kim et al. (2024) Bo-Kyeong Kim, Geonmin Kim, Tae-Ho Kim, Thibault Castells, Shinkook Choi, Junho Shin, and Hyoung-Kyu Song.  Shortened llama: A simple depth pruning for large language models, 2024.
  * Lan et al. (2020) Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut.  Albert: A lite bert for self-supervised learning of language representations.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=H1eA7AEtvS.
  * Lei et al. (2024) Tao Lei, Junwen Bai, Siddhartha Brahma, Joshua Ainslie, Kenton Lee, Yanqi Zhou, Nan Du, Vincent Zhao, Yuexin Wu, Bo Li, et al.  Conditional adapters: Parameter-efficient transfer learning with fast inference.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Lepikhin et al. (2021) Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen.  {GS}hard: Scaling giant models with conditional computation and automatic sharding.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=qrwe7XHTmYb.
  * Liu et al. (2020) Yijin Liu, Fandong Meng, Jie Zhou, Yufeng Chen, and Jinan Xu.  Faster depth-adaptive transformers.  In _AAAI Conference on Artificial Intelligence_ , 2020.  URL https://api.semanticscholar.org/CorpusID:229283620.
  * Liu et al. (2024) Yijin Liu, Fandong Meng, and Jie Zhou.  Accelerating inference in large language models with a unified layer skipping strategy, 2024.
  * Lv et al. (2024) Ang Lv, Yuhan Chen, Kaiyi Zhang, Yulong Wang, Lifeng Liu, Ji-Rong Wen, Jian Xie, and Rui Yan.  Interpreting key mechanisms of factual recall in transformer-based language models, 2024.
  * Madani et al. (2023) Ali Madani, Ben Krause, Eric R. Greene, Subu Subramanian, Benjamin P. Mohr, James M. Holton, Jose Luis Olmos, Caiming Xiong, Zachary Z. Sun, Richard Socher, James S. Fraser, and Nikhil Naik.  Large language models generate functional protein sequences across diverse families.  _Nature Biotechnology_ , 41(8):1099–1106, 2023.  doi: 10.1038/s41587-022-01618-2.  URL https://doi.org/10.1038/s41587-022-01618-2.
  * McGrath et al. (2023) Thomas McGrath, Matthew Rahtz, Janos Kramar, Vladimir Mikulik, and Shane Legg.  The hydra effect: Emergent self-repair in language model computations, 2023.
  * Minsky (1986) M. Minsky.  _The Society of Mind_.  Touchstone book. Simon and Schuster, 1986.  ISBN 9780671607401.  URL https://books.google.com.hk/books?id=veVOAAAAMAAJ.
  * Narayan et al. (2018a) Shashi Narayan, Shay B Cohen, and Mirella Lapata.  Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pp. 1797–1807, 2018a.
  * Narayan et al. (2018b) Shashi Narayan, Shay B. Cohen, and Mirella Lapata.  Don’t give me the details, just the summary! Topic-aware convolutional neural networks for extreme summarization.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , Brussels, Belgium, 2018b.
  * Olsson et al. (2022) Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah.  In-context learning and induction heads.  _Transformer Circuits Thread_ , 2022.  https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html.
  * Peebles & Xie (2023) William Peebles and Saining Xie.  Scalable diffusion models with transformers.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)_ , pp. 4195–4205, October 2023.
  * pytorch (2023) pytorch.  Pytorch profiler, 2023.  URL https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html.
  * Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al.  Improving language understanding by generative pre-training.  2018\.
  * Radford et al. (2019) Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.  Language models are unsupervised multitask learners.  2019\.  URL https://api.semanticscholar.org/CorpusID:160025533.
  * Raposo et al. (2024) David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro.  Mixture-of-depths: Dynamically allocating compute in transformer-based language models, 2024.
  * Schuster et al. (2022) Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler.  Confident adaptive language modeling.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), _Advances in Neural Information Processing Systems_ , volume 35, pp. 17456–17472. Curran Associates, Inc., 2022.  URL https://proceedings.neurips.cc/paper_files/paper/2022/file/6fac9e316a4ae75ea244ddcef1982c71-Paper-Conference.pdf.
  * Shazeer (2019) Noam Shazeer.  Fast transformer decoding: One write-head is all you need, 2019.
  * Shazeer et al. (2017) Noam Shazeer, *Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=B1ckMDqlg.
  * Singhal et al. (2023) Karan Singhal, Shekoofeh Azizi, Tao Tu, S. Sara Mahdavi, Jason Wei, Hyung Won Chung, Nathan Scales, Ajay Tanwani, Heather Cole-Lewis, Stephen Pfohl, Perry Payne, Martin Seneviratne, Paul Gamble, Chris Kelly, Abubakr Babiker, Nathanael Schärli, Aakanksha Chowdhery, Philip Mansfield, Dina Demner-Fushman, Blaise Agüera y Arcas, Dale Webster, Greg S. Corrado, Yossi Matias, Katherine Chou, Juraj Gottweis, Nenad Tomasev, Yun Liu, Alvin Rajkomar, Joelle Barral, Christopher Semturs, Alan Karthikesalingam, and Vivek Natarajan.  Large language models encode clinical knowledge.  _Nature_ , 620(7972):172–180, 2023.  doi: 10.1038/s41586-023-06291-2.  URL https://doi.org/10.1038/s41586-023-06291-2.
  * Srivastava et al. (2015) Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber.  Highway networks, 2015.
  * Varshney et al. (2023) Neeraj Varshney, Agneet Chatterjee, Mihir Parmar, and Chitta Baral.  Accelerating llama inference by enabling intermediate layer decoding via instruction tuning with lite, 2023.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
  * Wang et al. (2018a) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman.  Glue: A multi-task benchmark and analysis platform for natural language understanding.  In _International Conference on Learning Representations_ , 2018a.
  * Wang et al. (2023) Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt.  Interpretability in the wild: a circuit for indirect object identification in GPT-2 small.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=NpsVSN6o4ul.
  * Wang et al. (2018b) Xin Wang, Fisher Yu, Zi-Yi Dou, Trevor Darrell, and Joseph E. Gonzalez.  Skipnet: Learning dynamic routing in convolutional networks.  In _Computer Vision – ECCV 2018: 15th European Conference, Munich, Germany, September 8-14, 2018, Proceedings, Part XIII_ , pp. 420–436, Berlin, Heidelberg, 2018b. Springer-Verlag.  ISBN 978-3-030-01260-1.  doi: 10.1007/978-3-030-01261-8˙25.  URL https://doi.org/10.1007/978-3-030-01261-8_25.
  * Xin et al. (2020) Ji Xin, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin.  DeeBERT: Dynamic early exiting for accelerating BERT inference.  In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault (eds.), _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pp. 2246–2251, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.204.  URL https://aclanthology.org/2020.acl-main.204.
  * Yang et al. (2024) Yifei Yang, Zouying Cao, and Hai Zhao.  Laco: Large language model pruning via layer collapse.  _ArXiv_ , abs/2402.11187, 2024.  URL https://api.semanticscholar.org/CorpusID:267751181.
  * Zeng et al. (2023) Dewen Zeng, Nan Du, Tao Wang, Yuanzhong Xu, Tao Lei, Zhifeng Chen, and Claire Cui.  Learning to skip for language modeling, 2023.
  * Zhou et al. (2020) Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei.  Bert loses patience: Fast and robust inference with early exit.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 18330–18341. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/d4dd111a4fd973394238aca5c05bebe3-Paper.pdf.

##  Appendix A More implementation details

Table 3 lists the configurations of MoM-small/medium/large.

| MoM-small | MoM-medium | MoM-large
---|---|---|---
Initialization model | GPT2-small | GPT2-medium | GPT2-large
Hidden size | 768 | 1024 | 1280
Total number of FFN/MHA modules | 12 | 24 | 36
Number of attention heads | 12 | 16 | 20
Max sequence length | 1024 | 1024 | 1024
Vocabulary size | 50257 | 50257 | 50257
Table 3: Model configurations for MoM-small/medium/large.

In practice, we segment MoM into equally-sized chunks, each containing $4$ MHA
modules and $4$ FFN modules, namely $N=4$. Within each chunk, we execute the
MoM assembly process as presented in §3. We restrict the search space of each
chunk by setting $K\leq 4$ and $H\leq 6$, which results in $4\times 6=24$
combinations in total.

Then we elaborate on the initialization of chunked MoM. Taking a 8-layer
vanilla transformer as an example, the architecture is sliced as: the
bottom/top 2 layers remain the same, and modules in the middle 4 layers form a
MoM block, wherein we conduct the iterative assembly process. We denote this
chunking strategy as [1-1-4-1-1] where “1” represents the standard Transformer
block and “4” represents a chunk whose $N$ equals to 4. Empirically, we find
this setting to be stable across various choices of K$a$H$b$. A detailed
experimental analysis of different chunking strategies can be found in
Appendix B. Similarly, for MoM-small, MoM-medium and MoM-large, we use the
chunking strategies of [1-1-4-1-4-1], [1-1-1-4-1-4-1-4-1-4-1-1], and
[1-4-1-4-1-4-1-4-1-4-1-4-1-4-1], respectively.

##  Appendix B Chunks

Chunking Strategies | MoM Config | Val. Loss
---|---|---
[1-1-4-1-1] | K1H4S | 3.27
[1-1-4-1-1] | K2H4S | 3.22
[1-6-1] | K1H6S | 3.45
[1-6-1] | K2H6S | 3.21
[8] | K1H8S | 4.63
[8] | K2H8S | 3.23
[4-4] | K1H4S | 5.59
[4-4] | K2H4S | 3.22
Table 4: Applying different chunking strategies on an 8-layer MoM. These
models follow the same two-phase training procedure and the total training
steps of the second phase is 5k.

In this section, we study the impact of different chunking strategies on MoM
performance. This experiment is conducted on an 8-layer MoM. Except for
[1-1-4-1-1], we include several alternatives: [4-4] (two successive MoM blocks
with $N=4$), [1-6-1] (the top and bottom one layer are kept unchanged, and
modules in the middle 6 layers form an MoM with $N=6$), and [8] (all modules
form a big MoM with $N=8$). Table 4 shows the results of different chunking
strategies. When $K=1$, strategies other than [1-1-4-1-1] exhibit unstable
training curves and bad performance. This is because the routers need to make
multi-step decisions in the search space. A larger search space (the increase
of $N$) and more assembly steps (the increase of $H$) all lead to a harder
task for the routers to find the correct path in the search space. Things are
much better when $K=2$, where all strategies converge quite well, indicating
that we can develop a full MoM architecture (i.e.,[8]) without the chunking
strategy. In practice, we still apply the chunking strategy (i.e.,
[1-1-4-1-1]), because the strategy allows us to vary $K$ in a larger range,
and thus we can study MoM with more configurations. The experiment
demonstrates the necessity of manually restricting the search space of MoM so
that the decision-making burden for the routers would be relieved.

##  Appendix C Detailed downstream evaluation results

Table 5 presents the evaluation results for each sub-task of GLUE across
different models and Table 6 presents the results on XSUM with respect to
other ROUGE metrics.

Method | SST-2 | COLA | MRPC | QQP | QNLI | RTE | MNLI-(m/mm) | STS-B | average
---|---|---|---|---|---|---|---|---|---
small
GPT2 | 92.09 | 26.27 | 85.11 | 83.09 | 87.55 | 62.45 | 79.90/78.74 | 82.67 | 75.32
MoD | 87.73 | 21.66 | 81.43 | 82.01 | 81.79 | 64.62 | 75.96/75.22 | 79.76 | 72.24
MoE | 89.68 | 45.87 | 82.05 | 84.09 | 85.94 | 65.52 | 79.46/78.74 | 83.87 | 77.25
MoE (share) | 89.53 | 37.29 | 82.01 | 83.49 | 85.21 | 64.94 | 78.49/77.99 | 83.43 | 75.82
$\text{MoM}_{\text{E}}$ | 90.83 | 31.83 | 82.92 | 83.49 | 84.79 | 69.68 | 78.52/77.40 | 83.84 | 75.92
$\text{MoM}_{\text{I}}$ | 90.02 | 47.21 | 82.90 | 84.07 | 85.34 | 66.43 | 79.82/79.72 | 84.79 | 77.81
$\text{MoM}_{\text{P}}$ | 91.40 | 46.16 | 82.68 | 84.62 | 86.49 | 67.51 | 80.85/80.18 | 84.06 | 78.22
medium
GPT2 | 94.15 | 48.18 | 86.00 | 85.94 | 90.41 | 64.98 | 84.02/83.92 | 86.77 | 80.49
MoD | 89.45 | 37.43 | 84.97 | 84.21 | 84.88 | 64.26 | 78.94/77.94 | 83.47 | 76.17
MoE | 91.86 | 49.20 | 85.51 | 86.39 | 89.09 | 68.12 | 84.44/83.49 | 86.17 | 80.47
MoE (share) | 92.78 | 49.44 | 84.56 | 86.40 | 89.35 | 66.79 | 84.00/83.39 | 86.42 | 80.35
$\text{MoM}_{\text{E}}$ | 91.40 | 48.48 | 87.48 | 86.92 | 89.11 | 67.51 | 83.73/83.01 | 86.06 | 80.41
$\text{MoM}_{\text{I}}$ | 92.46 | 50.19 | 87.31 | 86.43 | 89.35 | 68.94 | 84.44/83.46 | 86.66 | 81.03
$\text{MoM}_{\text{P}}$ | 92.88 | 53.61 | 87.64 | 86.64 | 89.75 | 71.06 | 84.69/83.98 | 87.14 | 81.93
large
GPT2 | 94.15 | 60.04 | 88.74 | 87.88 | 91.89 | 75.45 | 86.80/85.98 | 89.30 | 84.47
MoD | 91.51 | 52.86 | 87.29 | 87.20 | 89.57 | 67.87 | 85.07/84.38 | 87.66 | 81.49
MoE | 94.32 | 61.19 | 88.54 | 88.36 | 92.04 | 71.68 | 87.38/86.94 | 89.45 | 84.43
MoE (share) | 93.88 | 61.43 | 88.12 | 88.17 | 91.56 | 68.31 | 87.01/86.73 | 89.28 | 83.83
$\text{MoM}_{\text{E}}$ | 93.64 | 59.26 | 88.25 | 87.59 | 91.29 | 72.23 | 85.20/84.78 | 88.26 | 83.39
$\text{MoM}_{\text{I}}$ | 93.69 | 62.25 | 89.19 | 88.12 | 92.36 | 74.98 | 87.22/86.78 | 89.90 | 84.94
$\text{MoM}_{\text{P}}$ | 94.42 | 64.49 | 89.56 | 88.68 | 92.94 | 77.69 | 87.68/86.98 | 90.70 | 85.90

Table 5: Detailed evaluation on the GLUE benchmark. We follow the previous evaluation setting Radford et al. (2018), for SST-2, QNLI, RTE, and MNLI, we report accuracy as the metric. For MRPC and QQP, we report the F1 score. For STS-b, we report the combined score of Pearson correlation and Spearman correlation. Method | ROUGE-1 | ROUGE-2 | ROUGE-L | ROUGE-AVG
---|---|---|---|---
small
GPT2 | 20.8 | 5.05 | 16.92 | 14.26
MoD | 14.45 | 2.89 | 11.80 | 9.71
MoE | 20.56 | 5.26 | 16.71 | 14.18
MoE (share) | 20.64 | 5.17 | 16.64 | 14.15
$\text{MoM}_{\text{E}}$ | 20.62 | 4.98 | 16.91 | 14.17
$\text{MoM}_{\text{I}}$ | 20.50 | 5.31 | 16.90 | 14.24
$\text{MoM}_{\text{P}}$ | 21.73 | 6.14 | 17.72 | 15.19
medium
GPT2 | 25.07 | 8.35 | 21.00 | 18.14
MoD | 20.89 | 6.04 | 17.52 | 14.81
MoE | 24.58 | 8.20 | 20.49 | 17.75
MoE (share) | 24.36 | 8.19 | 20.22 | 17.59
$\text{MoM}_{\text{E}}$ | 24.56 | 7.04 | 19.72 | 17.11
$\text{MoM}_{\text{I}}$ | 26.29 | 8.21 | 21.47 | 18.66
$\text{MoM}_{\text{P}}$ | 26.61 | 9.05 | 22.24 | 19.30
large
GPT2 | 28.16 | 9.92 | 22.98 | 20.35
MoD | 26.08 | 8.38 | 21.41 | 18.62
MoE | 28.54 | 9.89 | 23.47 | 20.63
MoE (share) | 28.47 | 9.47 | 23.23 | 20.39
$\text{MoM}_{\text{E}}$ | 28.48 | 9.65 | 23.24 | 20.46
$\text{MoM}_{\text{I}}$ | 30.91 | 10.21 | 24.08 | 21.73
$\text{MoM}_{\text{P}}$ | 31.38 | 10.77 | 24.93 | 22.36
Table 6: Detailed evaluation on the XSUM dataset. ROUGE is employed as the
evaluation metrics.

##  Appendix D Router analysis

###  D.1 Architecture design

Figure 5: Training curves of MoM-small (K1H4) with {GRU, MLP} routers.

We study the implementation of a key component in MoM: the routers. We
substitute the GRU within the router with a simple two-layer MLP, eliminating
the interaction among router decision states. Our exploration of the router’s
impact involves two setups: (a) initializing MoM from scratch, and (b)
employing the two-phase training approach. Here are some intriguing results.
As depicted in Figure 5, when initializing from scratch, both router
structures exhibit nearly identical loss curves. However, under setting (b),
training MoM with the MLP router becomes unstable marked by spikes in gradient
magnitude throughout the training. This instability suggests the router’s
inability to establish a consistent assembly plan for tokens. When
initializing from scratch, the required capability may not be learned
efficiently by these modules, as they often develop homogeneous
functionalities that waste the parameters. Conversely, in setting (b), where
modules are initialized with specialized functions, the optimization
progresses smoothly and converges quickly.

###  D.2 Learned router patterns

Figure 6: Routing patterns of MoM-medium.

We are curious whether the router follows specific patterns when choosing and
assembling modules. We visualize the transition probabilities between modules
(Figure 6) to answer this. The first observation is that the router does not
degrade into simply memorizing the original shallow-to-deep order but jumps
across modules as expected. For example, a common routing path in Figure 6 is
$(2,2,1,3,4,2)$ for FFN modules and $(3,1,3,2,S,4)$ for MHA modules (number
represents the module index). Another observation is that the loads of
different modules are imbalanced. In some cases, specific modules are hardly
used. Unlike MoE, which uses an auxiliary loss to balance the loads across
different experts (Fedus et al., 2022), we do not see a positive effect by
adding a balance loss to MoM. Adding regularization alleviates the imbalance
issue at the cost of performance degradation (by increasing validation
perplexity by 1.8 points). We posit an intuitive explanation: within the
language model framework, tasks that can be decomposed into numerous sub-tasks
may exhibit various levels of difficulty. Consequently, some sub-tasks
necessitate more engagement of modules with specific processing capabilities,
thus contributing to the observed imbalance.