# Conditional computation in neural networks: principles and research trends

Simone Scardapane Corresponding author. Remaining authors are listed in
alphabetical ordering. S. Scardapane is partly funded by Sapienza grants
RM1221816BD028D6 (DESMOS) and RG123188B3EF6A80 (CENTS). P. Minervini is partly
funded by ELIAI (The Edinburgh Laboratory for Integrated Artificial
Intelligence), EPSRC (grant no. EP/W002876/1), an industry grant from Cisco,
and a donation from Accenture LLP. DIET Department, Sapienza University of
Rome Alessandro Baiocchi  DIAG Department, Sapienza University of Rome
Alessio Devoto  DIAG Department, Sapienza University of Rome  Valerio Marsocci
CNAM Conservatoire national des arts et métiers  Pasquale Minervini  School of
Informatics, University of Edinburgh  Jary Pomponi  DIET Department, Sapienza
University of Rome

###### Abstract

This article summarizes principles and ideas from the emerging area of
applying conditional computation methods to the design of neural networks. In
particular, we focus on neural networks that can dynamically activate or de-
activate parts of their computational graph conditionally on their input.
Examples include the dynamic selection of, e.g., input tokens, layers (or sets
of layers), and sub-modules inside each layer (e.g., channels in a
convolutional filter). We first provide a general formalism to describe these
techniques in an uniform way. Then, we introduce three notable implementations
of these principles: mixture-of-experts (MoEs) networks, token selection
mechanisms, and early-exit neural networks. The paper aims to provide a
tutorial-like introduction to this growing field. To this end, we analyze the
benefits of these modular designs in terms of efficiency, explainability, and
transfer learning, with a focus on emerging applicative areas ranging from
automated scientific discovery to semantic communication.

Note: This tutorial paper is currently under review at Intelligenza
Artificiale (IOS Press). Feedback can be submitted to
<EMAIL_ADDRESS>
##  1 Introduction

In the last twenty years, neural networks (NNs) have undergone two opposing
trends. On one hand, the number of practical applications has continued to
grow, fueled by successes in, among others, language modeling [37], drug
design [92], rendering, and language-vision reasoning [63]. On the other hand,
their design has crystallized around a very small set of layers (e.g., multi-
head attention) and principles (e.g., permutation equivariance), while the
focus has shifted on scaling up their training, both in terms of data and
parameters [37]. Apart from scale, maybe less than a dozen components and
variations thereof are enough to categorize the vast majority of neural
networks deployed nowadays.

Among these design principles, sequentiality has remained a key component.
NNs, be they convolutional, recurrent, or transformers, are composed of a
stack of differentiable operations, which are activated in sequence for each
input to be processed. Stated in another way, their computational graph, i.e.,
the sequence of primitive operations executed on the underlying hardware, is
fixed beforehand when instantiating them. Because NNs have continued to scale
in depth and width, this has led to several issues in terms of computational
performance and efficiency [22]. Standard techniques to make NNs more
efficient only address this partially, by replacing the original network with
other static models having less layers (distillation [94], pruning [42]), less
precision per parameter, or by approximating each weight matrix (e.g., via
low-rank factorization).

By viewing neural networks as computing systems, this behavior is counter-
intuitive. Hardware components are designed based on their expected peak
usage, while only executing a fraction of their resources at any given time
(e.g., memory). Similarly, software libraries and operating systems are
composed of millions of lines of code, only a handful of which are selected
and run in a given moment, thanks to the use of branches, loops, and
conditional execution. Recently, a large body of literature has flourished on
embedding similar sparse modularity principles in the design of neural
networks [66]. Based on the original idea of conditional computation [5], this
has blossomed into a large number of practical implementations, ranging from
mixture-of-experts (MoEs) [84] to dynamic mechanisms to select tokens and
layers in transformers [51]. This tutorial is intended as an organized,
uniform overview and entry point into this growing body of literature.

The benefits of developing networks that can dynamically adjust their
computational graph goes beyond memory or time efficiency. More and more,
neural networks are treated in the same way as software [40], and deploying
them requires the possibility of quickly debugging their predictions [11],
continually fine-tuning them on new data, and transferring parts of their
knowledge from one network to the other [3], similarly to standalone software
libraries. As we will see, dynamically-activated neural networks provide a
principled way to improve both zero-shot transfer and generalization [69]
(Sec. 4.2) and explainability of the models (Section 4.3). This aligns with
the requirements of many novel applications of NNs, from scientific discovery
[92] to AI-native telecommunications [98, 85]. In particular, smart semantic
communication networks [85] envisioned in the so-called beyond 6G model
necessitate networks that can flexibly adapt to bandwidth and energy
constraints, while ensuring transferability and communication through separate
neural modules.

For the purpose of this tutorial, we consider three flavors of dynamism: (a)
neural components that can restrict their computation to a smaller subset of
input tokens (dynamic input sparsity); (b) layers that can selectively
activate sub-components for processing a token (dynamic width sparsity); and
(c) layers that can be completely skipped during their execution (dynamic
depth sparsity). As we show in Section 2, a simple mathematical formalism
encompasses all three cases. Section 2 also highlights how modularity can be
achieved with the addition of a small set of primitives to our networks’
toolkit, namely, the possibility of sampling in a differentiable way elements
from a set. We discuss in Section 2.2 the simplest technique to this end, the
Gumbel-Softmax trick [35, 49], and some common extensions.

We then proceed to discuss three notable implementations of these concepts in
Section 3: early-exit (EE) models, mixture-of-expert (MoE) layers, and token
selection mechanisms. While these models are generally discussed in separate
fashions (e.g., see [81, 50] for EEs, and [104, 20] for MoEs), viewing them as
specific instances of a general framework highlights many similar trends and
characteristics. In fact, we argue that designing networks with dynamically
activated components is not simply a matter of enhancing performance: in all
these cases, the resulting networks are more apt at adapting to variable
system’s constraints (e.g., decreased energy usage), specialization,
catastrophic forgetting, and multimodality. We build on these insights in
Sections 4 and 5, where we list potential research directions for these models
including adapting their computational cost and energy at inference and
training time in an elastic way, zero-shot transfer, robustness, and
explainability.

Relation to prior works: we do not claim to be the first to discuss these
ideas, nor to present them in a general way. We simply hope to provide a
simple, cohesive entry point into an extremely fascinating and growing
research direction in the literature. Recently, modular deep learning [66] has
been proposed as a general term for neural networks where computation is
functionally decomposed into units that are sparsely activated. Compared to
[66], our formalism is simpler and we focus on a smaller set of ideas: for
example, we do not consider input composition methods, nor soft routing
strategies, and we focus on viewing modularity as discrete sampling inside
neural networks. Hence, we provide an orthogonal view to [66], to which we
refer for a broader outlook on modularity and composition. We also focus on a
tutorial exposition, preferring simplicity and clarity to completeness. We
only describe the simplest techniques for differentiable sampling, and we
refer to [58] for a larger and more in-depth introduction to this field.
Additional entry points in the literature include [24] (for an overview up to
2021 of dynamism in NNs) and [22] for designing faster and more efficient
transformer models.

##  2 Three types of conditional computation

###  2.1 A general formalism for sparse modularity

Neural networks can be described as the composition of several trainable,
differentiable operations of the form:

$y=f(x,w)$ (1)

where $x$ denotes the input, $w$ the module’s parameters, and $f$ the specific
operation. Because of functional composition, $f$ can represent an operation
at any level of complexity, e.g., a basic linear algebra operation, a layer
(convolutional, recurrent, …), or a composite block, such as the classical
combination of token mixing, channel mixing, and layer normalization
operations found in transformers.

Most of our discussion focuses on transformer-like architectures, where the
input is composed of a set of tokens, representing parts of the complete input
[77] (e.g., subwords in a sentence or patches in an image). The input $x$ in
(1) is purposefully left vague: depending on the scenario and the layer, $x$
can be a single token (e.g., for the fully-connected blocks inside a
transformer), the entire set of tokens, a mini-batch of inputs, or other forms
of data aggregation.

We are interested in augmenting (1) to allow the dependence of the different
quantities to vary based on some conditioning input $x_{c}$, where $x_{c}$ can
be equal to $x$, some linear or non-linear projection of $x$ (in order to
decouple the conditioning input from the input used for the layer and increase
the flexibility of the dynamic mechanism), or some additional input
information (e.g., a latent token representing the language of the speaker in
an audio model [46]).

Figure 1: Input sparsity: A differentiable mechanism subsamples input tokens
to be processed by the later parts of the network (we show original image
patches in the figure, but the tokens can be equivalently be replaced by their
latent representations if we consider an intermediate layer of the
architecture).

Suppose we have available some trainable subsampling operation, defined over
the conditioning input $x_{c}$ and over some set $\mathcal{S}$, such that:

$\Gamma(\mathcal{S},x_{c})\subseteq\mathcal{S}\,.$ (2)

Practically, sets are always represented in some matrix format inside neural
networks (e.g., by stacking all elements row-wise), in which case we assume
the relevant quantities in (2) to follow the same conventions. As an example,
if $\mathcal{S}$ is represented by a matrix $\mathbf{S}$, (2) can be
implemented by row-wise masking of the corresponding elements:

$\Gamma(\mathbf{S},x_{c})=\mathbf{M}\odot\mathbf{S}$ (3)

where $\mathbf{M}$ is a binary mask of suitable shape with rows set to either
$\mathbf{0}$ or $\mathbf{1}$. Different implementations can be obtained
depending on whether the size of the output mask is known in advance (e.g.,
how many tokens to keep in a given layer), or must be estimated by $\Gamma$
itself. We will see later on in Section 2.2 how such an operation can be
implemented in practice.

This small addition allows us to implement several levels of modularity inside
our neural network. First, assume the input $x$ can be decomposed into smaller
components $x=\left\\{x_{1},\ldots,x_{n}\right\\}$. In many cases, this is a
natural decomposition, e.g., tokens inside a transformer, time instants for a
sequence, or frames for a video input. However, they can also correspond to
additional register tokens [18], to elements extracted from some external
memory (e.g., AdaTape [100]), or additional views or modalities. Then, dynamic
input sparsity (Fig. 1) is achieved by combining the layer with a subsampling
operation on the input:

$\text{Input sparsity}:f(\Gamma(x,x_{c}),w)$ (4)

This allows the layer to focus exclusively on input components that are
relevant to the current operation. As an example, consider an image with a
very wide, uniformly blue background: for the majority of tasks, we can
imagine the layer to be able to operate even when removing the vast majority
of tokens corresponding to such background [95]. This is, indeed, observed in
practice, as we’ll describe in Section 3.5.

Figure 2: Width sparsity: Different parts of a layer (e.g., experts) can be
activated based on the conditioning value.

Second, assume the weights themselves can be decomposed into blocks
$b=\left\\{w_{1},\ldots,w_{m}\right\\}$, and the entire function decomposed
as:

$f(x,w)=\oplus(\left\\{f_{1}(x,w_{1}),\ldots,f_{m}(x,w_{m})\right\\}),$ (5)

where $\oplus$ is a permutation equivariant operation that aggregates the
results of the different blocks. This is also a very natural decomposition in
practice: for example, $f_{i}$ could correspond to a single output channel in
a convolutive layer, with $\oplus$ being a concatenation. Or, $f_{i}$ can be a
head in a multi-head attention layer, with $\oplus$ being a concatenation
followed by a projection. In this case, what we call dynamic width sparsity
(Fig. 2) can be achieved by subsampling the blocks:

$\displaystyle\text{Width sparsity}:f(x,\Gamma(w,x_{c}))=$
$\displaystyle\oplus\left(\Gamma(\left\\{f_{1}(x,w_{1}),\ldots,f_{m}(x,w_{m}),x_{c}\right\\}\right)$
(6)

Imposing some form of structure over $w$ is fundamental since subsampling each
weight separately would result in a highly unstructured form of dynamic
pruning, which is generally very difficult to optimize and control [36]. Once
again, this formulation is very generic: for example, if each $f_{i}$
corresponds to an activation function, subsampling with cardinality $1$
corresponds to choosing the best activation function for each token or each
input, in a form of dynamic model selection mechanism [47]. If each $f_{i}$ is
an entire model, this can be used to route information across blocks with
different types or complexity (see, e.g., layer stitching [60]). We will see
in Section 3.2 a common implementation of this principle in the context of
MoEs, and defer more general discussions to Section 4.

Third, even if a weight decomposition is unavailable, dynamic computation can
still be achieved by considering the entire layer as a single block and
conditionally skipping its execution. This dynamic depth sparsity (Fig. 3) can
be achieved easily by rewriting the layer with an additional, untrainable
scalar weight $\sigma=1$, and writing:

$\displaystyle\text{Depth
sparsity}:f^{\prime}(x,w,\Gamma(\left\\{\sigma\right\\},x_{c}))=$ (7)
$\displaystyle\sigma(x_{c})f(x,w)+(1-\sigma(x_{c}))x$ (8)

where we assume the empty set $\emptyset$ to be implemented as $0$ in
practice. As a variant of this idea, we can skip entire blocks of layers by
exiting the network instead of adding a skip connection in (8). We will see
concrete implementations of this concept in Section 3.1.

Figure 3: Depth sparsity: A subset of the layers can be deactivated by the
sampling mechanism, like in early-exit networks.

###  2.2 Discrete sampling: the Gumbel-Softmax trick

Before proceeding, we discuss briefly the implementation of the subsampling
layer $\Gamma$. As a prototypical example, we restrict our analysis to
subsampling a single element from the set in input. In this case, a very
common implementation is the so-called Gumbel-Softmax (GS) trick [35], also
known as the concrete distribution [49]. In this section we only provide a
high-level overview, and we refer to [34] for a fuller exposition.

Figure 4: Overview of the Gumbel-Softmax trick. We show in orange
differentiable operations (the argmax’s gradient being zero almost
everywhere), and with dashed arrow the relaxed backward path.

First of all, we process the inputs with some trainable layer
$p(\mathcal{S},x)\in\mathbb{R}^{\lvert\mathcal{S}\rvert}$ to provide a real-
valued score for each element in $\mathcal{S}$, which is proportional to its
probability of being sampled. In this section we drop the subscript from the
conditioning input $x_{c}$ for readability. The implementation of $p$ depends
on the use case. For example, suppose that the conditioning input is a tensor
$x\in\mathbb{R}^{h\times w\times d}$ coming from intermediate output of a CNN
having width $w$, height $h$, and $d$ channels, while $\mathcal{S}$ is the set
of output channels in a convolutive layer, of which we want to select a single
one [31]. We associate to each channel a trainable vector $\mathbf{c}_{i}$,
and compute an affinity score via the dot product with a pooled representation
of the image [31]:

$\displaystyle\mathbf{v}=\text{MLP}\left(\sum_{i,j}x_{ij}\right)$ (9)
$\displaystyle
p(\mathcal{S},x)=\left[\mathbf{v}^{\top}\mathbf{c}_{1},\ldots,\mathbf{v}^{\top}\mathbf{c}_{\lvert\mathcal{S}\rvert}\right]$
(10)

Many other designs for $p$ are possible based on $x$ and $\mathcal{S}$. Let us
denote by $\mathbf{p}$ the output of $p$, and by $g_{i}$ samples from the so-
called Gumbel distribution [35]. We also assume that the elements of
$\mathcal{S}$ correspond to integers,
$\mathcal{S}=\left\\{1,\ldots,\lvert\mathcal{S}\rvert\right\\}$, in which case
we simply write $\Gamma_{\lvert\mathcal{S}\rvert}(x)$ for the subsampling over
the set of integers. Then, the following operation provides an unbiased sample
from $\mathcal{S}$, with the probability of sampling the $i$-th element
$p(s\in\mathcal{S})\propto\exp(p_{s})$ proportional to $p_{s}$:

$s=\underset{i}{\arg\max}\;\;\left\\{p_{i}+g_{i}\right\\}\,.$ (11)

Removing the Gumbel noise $g_{i}$, this corresponds to taking the element with
highest score $p_{i}$, while sampling provides a degree of freedom that is
helpful in exploring possible alternatives. Since the probabilities $p_{i}$
are implicitly trained via $p(\mathcal{S},x)$, the network can learn to select
the element from $\mathcal{S}$ which is most useful for the specific input
$x$. In the case of channel selection, for example, this can lead to
specializing single channels to specific types of inputs.

The operation in (11) is not easily trainable via gradient descent (since its
gradient will be zero almost everywhere), but it can be relaxed with a softmax
approximation:

$s_{i}=\frac{\exp(\left(p_{i}+g_{i}\right)/\tau)}{\sum_{j}\exp(\left(p_{j}+g_{j}\right)/\tau)}\,.$
(12)

The quality of the approximation can be controlled by the user-defined
parameter $\tau$, sometimes called the temperature [34]. The values in
$\mathbf{s}$ are not binary anymore, being generic convex combinations of the
(one-hot) representations from the elements in $\mathcal{S}$. However, a
common relaxation is to use the binary values in (11) during the forward pass
of the network, and relax the gradients to use the soft approximation in (12)
during the backward pass. This is called straight-through estimation (STE). A
high-level overview of the complete schema is given in Fig. 4.

The GS trick can be easily extended to sampling more than a single value by
replacing the $\arg\max$ with a top-k operation [43]. Suitable
generalizations, such as the entmax family [16], can also sample binary
vectors with a variable number of elements. The simplicity of the GS trick
makes it widespread in many applications, but several other types of sampling
layers can be found in the literature, especially for more complex
combinatorial spaces, for which we refer to [58]. A larger overview of
possible sampling operations for MoE layers is also given in Sec. 3.3.

##  3 Concrete implementations

The previous section has introduced a generic framework for designing networks
with dynamic computational graphs. In the literature, these ideas have
coalesced around a few, notable set of implementations. In this section we
overview three of the most popular ones: early exits (Sec. 3.1), MoE layers
(Sec. 3.2), and token sampling (Sec. 3.5) as examples of networks having
dynamic width, depth, and input, respectively. In all cases we focus on key
strengths and drawbacks of each approach, highlighting how they connect to the
framework of Sec. 2.

###  3.1 Early exits

Early exit neural networks (EENNs) were introduced based on the idea that not
all the inputs to a network require to go through all the layers of the model
to be correctly classified [88]. In fact, the accuracy of a network can
decrease with respect to the depth on particularly easy samples, a phenomenon
known as overthinking [79].

Consider a neural network (sometimes called the backbone network in this
context), which can be either trained from scratch or pre-trained. In an
early-exit framework, the backbone is augmented with auxiliary classifiers
(early exits) which are connected at intermediate outputs of the backbone.
These auxiliary blocks can be trained all at the same time or in a layer-wise
fashion [81], while at inference time they are used to halt the computation
when the model is confident enough about the prediction. This is an example of
dynamic depth sparsity (Section 2), except that an entire subset of layers is
skipped at the same time.

We introduce first the most common formulation of EENNs, in which early exits
are jointly trained while the early exit logic is defined manually with
additional hyper-parameters. As already stated, we can see a neural network as
a composition of $b$ sequential blocks:

$f(x)=\mathrm{f}_{b}\circ\mathrm{f}_{b-1}\circ\cdots\circ\mathrm{f}_{1}(x)$
(13)

where $\circ$ denotes function composition, $b$ is the last layer of the
backbone, and we remove the dependency on the parameters for simplicity. Where
to place the auxiliary classifiers is in general an hyper-parameter [81], but
here, for clarity, we add a classifier after each block, producing multiple
prediction functions:

$y_{i}(x)=\mathrm{c}_{i}\circ\left(\mathrm{f}_{i}\circ\mathrm{f}_{i-1}\circ\cdots\circ\mathrm{f}_{1}(x)\right)$
(14)

where $c_{i}$ is a small classifier (e.g., for a CNN, a typical design is to
have a global average pooling operation followed by a fully-connected layer,
while for a transformer we can apply a fully-connected layer on a specialized
class token). This formulation leads to an EENN which is capable of
classifying input samples at any stage of the computation through the sequence
$y_{1}(x),y_{2}(x),\ldots,y_{b-1}(x),f(x)$. Ideally, each classifier can
specialize on a subset of the training samples based on the complexity of the
input. In the simplest case, given a desired output $y$, all early exits can
be trained simultaneously by minimizing the sum of the losses for each
classifier:

$L=\alpha\text{CE}(y,f(x))+\sum_{i=1}^{b-1}\beta_{i}\text{CE}(y,y_{i}(x)),$
(15)

where CE is the cross-entropy loss, and $\alpha$, $\beta_{1}$, $\ldots$,
$\beta_{b-1}$ are hyper-parameters that balance the loss contribution from
each early exit. At inference time, an exit can be chosen by comparing, e.g.,
the predicted class probability or the entropy of the prediction to a user-
defined parameter acting as a threshold. Other strategies include combining
predictions based on trainable [80] or geometric [96] ensembling.

These networks have found applications in many fields recently. Concerning
NLP, many strategies with BERT have been proposed [110, 97, 113], varying from
patience-based EE [110] to extra classification layers [97]. For computer
vision, EE applications vary based on the strategy, including multi-exiting
[4, 44, 94, 8], and loss-weight adjustment [93]. Also, the tasks vary,
spanning semantic segmentation [44], video recognition [23], adversarial
training [41], and image compression [98]. Recently, EE strategies have been
proposed also for vision-language models [87].

Although this formulation of EENNs (or minor variations thereof) is very
common in the literature (see e.g., [88, 8, 110, 50, 82, 87]), it only fits
partially our view of dynamic sparsity, since the computational graph is
unchanged during training and the dynamism is only obtained via an external
heuristic. To this end, [80] proposed a mechanism to optimize the early exit
selection during training, called differentiable branching. Suppose we augment
each early exit with an additional output that we denote as $\gamma_{i}(x)$.
The key idea of differentiable branching is to define a new recursive output,
which is given by a soft composition of the original outputs:

$\tilde{y}_{i}(x)=\gamma_{i}(x)y_{i}(x)+(1-\gamma_{i}(x))\tilde{y}_{i+1}(x)$
(16)

where the recursion stops at $i=b$, where $\gamma_{i}(x)=1$ by definition. At
inference time, this can be turned into an early exit module by viewing
$\gamma_{i}(x)$ as a binary classifier and thresholding it at $0.5$. By
replacing the gates $\gamma_{i}(x)$ with samples from a GS distribution, we
return to the general setting described in Section 2. In fact, if the outputs
of $\gamma_{i}(x)$ are binary we can identify:

$\Gamma_{b-1}(x)=\begin{bmatrix}\gamma_{1}(x)\\\
(1-\gamma_{1}(x))\gamma_{2}(x)\\\ \vdots\\\
\left[\prod_{i=1}^{b-2}(1-\gamma_{i}(x))\right]\gamma_{b-1}(x)\end{bmatrix}$
(17)

with the subsampling block over the indexes of the early exits. This can also
be given a Bayesian interpretation under the stick-breaking process [68].
While the sampler in (17) is conditioned on the full input, other choices are
possible. For example, in multitask learning a separate task embedding can be
used to select different exits for different tasks [106], while in the case of
transformers only the class token can be used to decide an optimal exit
strategy [30]. Being able to train the choice of early exit is fundamental
because it allows us to regularize it further for, e.g., better inference-
accuracy trade-offs: we revisit this idea in Section 4.1. See [80, 50] for in-
depth reviews on EENNs and [80] for a fuller description of differentiable
branching.

###  3.2 Mixture-of-experts

Mixture-of-Experts (MoEs) models were introduced more than thirty years ago as
adaptive algorithms to perform dynamic ensembles of individual machine
learning algorithms, denoted as experts in that context [104]. Recently, MoE
layers in neural networks have gained popularity, especially for LLMs and
large vision modules [77]. They allow scaling the model’s parameters while
keeping the compute budget constant, and they offer the possibility of
distributing the training by decomposing the computation of a layer over
multiple GPUs [77]. As a side product, they offer dynamic width sparsity by
selectively activating only parts of the layer for each input. At the time of
writing, the largest open source language and vision models are based on a MoE
formulation [45, 21, 77, 17, 37].

In a MoE layer, the trainable parameters are grouped into topologically
identical blocks called experts. In the forward pass, a decision block (called
router in this context) selects a subset of expert blocks (typically, only a
fixed number of $k$ experts, with $k$ selected beforehand, is chosen) that
will be activated and routes the input data to the selected experts. More
formally, the MoE layer can expressed as:

$f(x,w)=\sum_{i=1}^{n}\gamma_{i}(x)f_{i}(x,w_{i})$ (18)

where $\gamma$ is the routing function and $f_{1},f_{2}\dots f_{n}$ are the
experts. In its simplest form, this is a sparse linear combination of the
outputs of the experts. This is immediately seen to be an instance of 6 with
$\Gamma_{n}(x)=\left[\gamma_{1}(x),\gamma_{2}(x),\ldots,\gamma_{n}(x)\right]$,
provided that the outputs of $\gamma_{i}(x)$ are sparse.

The unique advantage of MoEs is that only a fraction of the network
parameters, depending on the chosen value for $k$, are used for computation.
Such fraction can be tuned by changing the value of $k$. Additionally, one can
devise routing functions such that experts focus on different areas of the
input space, leading to specialized sub-networks (see Sec. 4.2 for a
discussion on specialization and Sec. 3.3 for an overview of routing
algorithms for MoE layers).

As stated earlier, the seminal idea of MoEs dates back to [39], where this
paradigm is regarded as an ensembling technique with soft gates. The first
work that adopts MoEs in a modern deep learning scenario is [83]. Here, the
gate becomes sparse, i.e., a fixed top-$k$ operation and the main goal is
scaling the number of parameters. Based on these initial results, in the past
few years the deep learning community witnessed a steep growth of MoE’s
success in conjuction with transformer architectures.

Unlike early exiting, where the routing function typically works at the level
of input samples (e.g., images for a CNN), a MoE layer as in (18) can work at
the level of individual tokens in a transformer. More specifically, in the
fully-connected layers of a transformer block, each token can be assigned to
one or more experts. Transformer MoEs have been proposed for NLP [21, 71], for
vision [77] and even for multimodal training [57]. Additional variants can be
applied to individual heads of an attention block [107] or adapters in a fine-
tuning context [105].

###  3.3 Routing in Mixture of Experts

Different routing strategies can be adopted to orchestrate the experts. These
are defined in the gating function $\gamma$ and usually leverage some
differentiable sampling method similar to the Gumbel-Softmax trick (Sec. 2.2).
Routing is a crucial aspect of any MoE architecture, as it must perform a
matching between input data and expert network, conditioned on the input data
itself. Decisions taken by the routing function govern the training. They can
determine faster convergence and expert specialization or cause expert
starvation, a peculiar situation where only a few experts are always activated
due to their ever-increasing expertise.

The routing algorithm usually adopts a dot product for computing similarity
between the embedding of each routed sample (e.g. a token in a transformer)
and $n$ trainable expert embeddings (one for each expert). Once this operation
has occurred, the assignment can happen according to various methods.

In _top- $k$ token choice_, each token is assigned to the $k$ most similar
experts, where $k$ is an arbitrary value. This is the original approach from
[83], and it usually requires additional regularization to avoid expert
starvation and training instability. The same approach has been adopted with
minor changes (e.g. balancing loss and the $k$ factor) in [73, 21]. _Top- $k$
expert choice_, proposed in [111], tries to solve the expert balancing problem
by assigning each expert to the $k$ most similar tokens, where $k$ is an
arbitrary value. Other works [15, 78] cast the expert assignment problem into
a _reinforcement learning_ framework and train over a routing policy to assign
each token to a single expert. Finally, [114] achieves surprisingly good
results by randomly assigning experts to routed samples. An empirical
comparison of different routing functions can be found in [48], while a
broader overview of MoE layers is given in [20].

###  3.4 Soft routing

We discuss here a variant of the MoE layer in (18), known as soft MoE, which
has become popular recently to avoid training instabilities in the layer. The
key idea is to use the gating function to combine the layer’s inputs [70] or
weights [56] instead of the model’s outputs. These variants also provide a
fixed compute budget which is decoupled from the total number of parameters,
but they are generally easier to implement and to the train.

First, denote by $x_{1},\ldots,x_{n}$ the tokens in input to the layer. The
first soft MoE variant we consider can be defined as [70]:

$f_{i}(x,w_{i})=f_{i}\left({\color[rgb]{0,1,1}\sum_{i=1}^{n}\gamma_{i}(x_{i})x_{i}},w_{i}\right)$
(19)

i.e., each expert is activated with a soft combination of input tokens.
Compared with standard top-$k$ routing, each expert is only activated once in
this formulation, irrespective of the number of tokens. The output of the
layer for a single token is then given by another soft combination, this time
over the experts’ outputs:

$f(x_{i})=\sum_{i=1}^{n}\gamma_{i}(x_{i})f_{i}(x,w_{i})$ (20)

By comparison, another soft MoE variant can be achieved by performing a soft
combination of the experts’ weights [56]. Denote by $f(x,w)$ the generic
architecture of a single expert, we have:

$f(x)=f\left(x,{\color[rgb]{0,1,1}\sum_{i=1}^{n}\gamma_{i}(x)w_{i}}\right)$
(21)

A thorough comparison of the relative benefits of hard vs. soft routing is
still lacking in the literature, but we highlight a few benefits in Sec. 4.

###  3.5 Token dropping and token merging

We conclude with a brief analysis of token selection mechanisms. In
particular, we consider two cases of token selection, which we refer to as
token dropping (removal of tokens from the model) and token merging
(combination of tokens). Both are conditional computation techniques that can
be employed in any transformer model. The fundamental idea behind token
dropping and token merging is that the input data often contains information
that is nearly useless for the final task. Since transformers operate on set-
based inputs with no fixed cardinality, one can dynamically reduce the number
of tokens being passed to each layer depending on the relevance of said tokens
for the final task.

More formally, let $\mathbf{X}$ be the set of $n$ tokens inside a transformer
model stacked horizontally. For the purpose of this section. the subsampling
operation $\Gamma$ can be implemented as a matrix multiplication of the input
matrix $\mathbf{X}$ and a mask $\mathbf{M}$:

$\mathbf{X}^{\prime}=\Gamma_{n}(\mathbf{X})=\mathbf{M}\mathbf{X}$ (22)

Depending on the properties of the matrix $\mathbf{M}$, we can obtain either
token dropping or token merging:

  * •

Token dropping: $\mathbf{M}$ is a matrix of shape $n^{\prime}\times n$ that
selects a subset of $n^{\prime}$ elements from the original tokens. This can
be achieved by setting each row to a one-hot vector.

  * •

Token merging: $\mathbf{M}$ is a binary matrix of shape $n^{\prime}\times n$,
with the requirement that each column sum to one. In this way, the output
tokens are combinations of the input ones, so each input token participates in
a single output token. This can be seen as a dynamic variant of the standard
pooling used in CNNs.

Both methods aim to reduce the computational cost of the forward pass of a
transformer by reducing the number of tokens. The mask $\mathbf{M}$ is usually
computed depending on the input data, $\mathbf{M}=\text{MLP}(\mathbf{X})$, in
such a way as to eliminate only redundant information, such as background
patches in an image. This is usually done either by small modules that are
added to the backbone model [75] or by evaluating some channels added to the
tokens for this purpose [51]. For token selection, when dealing with mini-
batches of inputs, the sub-sampling operation in (22) cannot be performed
easily in a vectorized way. In these cases, (22) can be replaced by a masking
operation as in (2), with the additional constraint that subsequent selection
modules can only select values that were not masked previously [75]. In these
cases, performance gain only materialize at inference time, while at training
time most operations are still executed in a masked way.

Token dropping and token merging have been largely explored in recent years;
some of the most impactful works on these topics in the field of vision
transformers are [51, 75] for token dropping, while [9, 61] explored token
merging. These techniques have also been employed in other fields where
transformers are predominant: in 3D computer vision token dropping has been
applied to point cloud transformers by [102], while in NLP [32] applied token
dropping in the pre-training of BERT, obtaining significant speedups.

##  4 Research directions

Designing neural networks with sparse modularity principles has a number of
benefits. The first and most studied is increased efficiency, both in training
and in inference (Sec. 4.1). However, modular networks show emergent
properties also in terms of specialization and transferability (Sec. 4.2), as
long as providing a blueprint for a new type of explainability techniques
(Sec. 4.3). We briefly overview these aspects in the following sections.

###  4.1 Efficiency in training and inference

Conditional computation has gained significant attention for accelerating
training and inference in deep learning models [22, 66]. The approaches can
vary based on several factors. For example, some research investigated the
sparsification of CNNs [90] and of transformers [36, 12], some others focused
on some specific aspects of the net, such as the gradients [19], the
backpropagation, the activations [13] and the attention layers [89, 14]. For
clarity, we focus here on efficiency aspects related to the three families of
models seen in Sec. 3, focusing on emerging trends and open challenges.

When training from scratch an EENN, the joint training described in Sec. 3.1
cannot provide faster training since all exits must be trained simultaneously.
Soft combinations of the early exits [80] combined with sampling tricks could
be useful, but this aspect has been scarcely explored in the literature,
possibly due to training instabilities and collapse. Still, EEs can provide
indirect benefits to training by accelerating the convergence of deep neural
networks [8, 88]. They can also improve the inference efficiency of large-
scale pre-trained models [25, 110], making them more discriminative [88], act
as a regularization technique [79], and possibly reduce training problems
(e.g. vanishing gradient phenomena) [62]. In addition, they can be trained in
a layerwise fashion if one starts from a pre-trained network [24].

Compared to EEs, MoEs were investigated mostly to speed-up training by
distributing the experts across different GPUs [20], or to accelerate
inference by activating a subset of experts for each forward pass. For the
former, tiling the model across separate machines requires customized
implementations [73]. In addition, to achieve an effective stable throughput,
training must avoid collapses of the routing function, and care must be taken
to balance the number of tokens which are sent to the different experts. This
last point can be achieved with the addition of so-called balancing losses,
e.g., batch prioritized routing [77, 15, 20]. Due to these problems, open-
source implementations of large MoE models have generally lagged behind
proprietary models, with some preliminary steps being taken in this direction
[101].

The benefits of applying MoEs can also vary depending on the specific
component which is being replaced. In [83], MoEs replace MLP layers to scale-
up transformers. A similar methodology, for vision, was presented in [77].
Other approaches proposed to replace attention layers [107], entire blocks
[86] of the net, and adapters [105]. By comparison, soft MoEs (Sec. 3.4) can
reach similar gains as standard MoE layers while being simpler to train.
Finally, the majority of MoEs are trained from scratch, but recently
moefication has emerged as an interesting research direction, in which a pre-
trained model is converted to a MoE variant by clustering the activations [72,
108]. These variants provide different accuracy-time trade-offs based on the
specific routing function. Additional emerging research trends are exploring
dynamic variants of routing to provide a flexible inference budget instead
[67, 95].

Finally, reducing the number of tokens with token selection techniques is a
straightforward strategy to improve the inference time of the network. Most
methods in this sense have focused on token dropping [27, 103, 9]: AdaViT [51]
proposes a light-weight decision network, attached to the backbone to produce
decisions on-the-fly; DynamicViT [75] proposes an attention-masking strategy
to block interaction of redundant tokens; MSViT [28] proposes a conditional
gating mechanism that selects the token scale for every image region; GTP-ViT
[99] introduces a Graph-based Token Propagation (GTP) to propagate less
significant tokens’ information to spatially and semantically connected
important tokens. During training, token selection is generally achieved by
masking the corresponding input elements (Sec. 3.5). Designing token selection
mechanisms that can improve training time (for a fixed compute budget) is
still an open challenge.

Some works have focused more on specific tasks, such as diffusion [10], long-
input sequences [1], or to better understand the general pattern of token
dropping [27]. For token merging, PatchMerger [76] proposes a module that
reduces the number of tokens by merging them between two consecutive
intermediate layers [76]. Recently, a hybrid solution, Token Fusion (ToFU)
[42], proposed to put together the benefits of both token pruning and token
merging. Token selection can also be combined with the other methods discussed
in this paper. As an example, Adaptive Computation Module (ACM) [95] is a
technique that combines token dropping with a dynamic width principle, in
which each token is allocated a variable width for each layer in the network.

###  4.2 Specialization

Modularity can also lead to specialization benefits in specific tasks [109,
53], such as language modeling [46, 52], cross-lingual transfer [64, 3, 63,
91], speech processing [52, 65], computer vision and multi-modality [74, 57,
2], and task generalization [55].

Specialization can happen in two ways: implicitly if no external information
is provided, or explicitly if additional information (e.g., the speaker
identity) is known. The key insight is that knowledge of, e.g., the task, the
domain, or the speaker provides latent information that can be used to
condition the routing blocks ($\Gamma$ in Sec. 2), thus specializing specific
parts of the network or specific components to different scenarios. This can
be achieved in different ways: manual routing in which different components
are pre-selected for the different latent vectors [66], direct conditioning of
the routing functions [112], entropy regularizers to force the distributions
w.r.t. to a specific latent vector to have low entropy [57], or weight sharing
across tasks [87]. Soft merging [56] also seems to offer specialization
benefits.

Studies on quantifying the quality of the resulting specializations were
carried out recently in the literature [101]. MoEs showed promising
performance due to their nature, and, specifically, hard routing facilitates
module specialization [66]. On the other hand, learned routing could lead to
sub-optimal results [54, 53] w.r.t to fixed routing, except for specific cases
[69]. As another example, [101] showed that many routing decisions tend to
ignore context, and they can be fixed during the early stages of training.
Generally speaking, understanding the interplay between routing and
specialization and the extent to which this specialization correlates with
human-understandable semantics and biological plausibility remain open
challenges that will require novel benchmarks and metrics [72, 109]. Recent
moefications work show that some form of emergent modularity may be found also
in pre-trained networks with no specific modularity bias [72, 108].

A benefit of having specialized sub-structures and components is that the
networks can perform better in multi-task learning [55] and multi-domain
scenarios [57], and they can potentially enable zero-shot transfer and
generalization of the resulting sub-structures. We list a few interesting
examples from the recent literature. PHATGOOSE [55] is a MoE specialized in
zero-shot generalization, thanks to a new post-hoc routing strategy. EMoE
focuses on the implicit modular structures (Emergent Modularity) of large pre-
trained transformers [72]. DSelect-k [29] presents a continuously
differentiable and sparse gate for multitask learning. Uni-Perceiver-MoE [112]
is a generalist conditional MoE that shows SOTA performance when compared to
specialized MoEs. LIMoE [57] is focused on language-image pre-training.
DeepSeekMoE [17] manages to ensure expert specialization for language models,
still reducing computational costs. Other approaches include modular submodels
to scale language models [6] and a class-aware channel pruning for queriable
NNs [38].

###  4.3 Explainability

Finally, modularity and sparsity can provide significant gains in
explainability, which is a significant issue when deploying systems [26]. In
particular, analyzing and plotting the routing decisions made by the
subsampling blocks $\Gamma$ almost always provide valuable insights into the
predictions. These include, but are not limited to, visualizing representative
tokens sent to each expert [57], visualizing the early exits distribution for
a given sequence in an autoregressive model [82], visualizing the tokens that
were never discarded in a network having token selection [59, 51, 75], or the
number of experts that were activated token-wise [95]. These techniques can
provide benefit also for specific tasks, such as object detection [13].

Also in this case, we lack principled benchmarks and frameworks to analyze the
resulting plots, which is an open problem in explainability in general [26].
In addition, for networks having thousands of blocks or experts, manually
analyzing each of them can be a time-consuming process. LLMs can potentially
help in automating this process [7]. From a mechanistic interpretability point
of view, fully understanding and being able to track the evolution of modules
in these networks could be a large step forward in better understanding
overfitting, generalization, and fine-tuning [33]. Finally, we note that being
able to plot and visualize internal decisions of the network provides a direct
way for providing feedback on, e.g., whether specific modules should be kept
or ignored, and to perform interventional analyses [11].

##  5 Conclusions and future trends

In this tutorial paper we have provided an introduction to the emerging field
of designing neural networks which are sparsely activated in a modular
fashion, via the use of conditional computation techniques. To this end, we
have provided both a general mathematical formalism to categorize these
approaches, and then an overview of several concrete implementations including
mixture-of-expert models and early exit neural networks. Although these models
have been investigated mostly for improving training and/or inference time, we
have discussed a number of additional emerging benefits from this approach,
including specialization, generalization, and explainability. Many of these
benefits are only starting to be investigated, opening up interesting avenues
of research.

Some common challenges have also emerged from our discussion: (a) being able
to control and adapt the inference and/or training time is still an open
challenge, with most techniques providing a fixed accuracy-time trade-off
(e.g., MoEs with top-$k$ routing); (b) sparse routing introduces balancing and
collapse challenges that must be taken care of; (c) most routing decisions are
taken only locally (e.g., layer-wise), while taking them globally requires
sampling from a combinatorially large search space requiring new sampling
techniques [58]; (d) outside of accuracy, benchmarks and metrics for
evaluating specialization and generalization of these models are still being
developed [53].

There are also several research directions that we were not able to touch due
to space constraints: these include scaling laws for modular models [15],
biological plausibility [81, 109], and exploiting these models in specific
applicative fields such as split computing [50] and semantic communication
[106].

## References

  * [1] J. Ainslie, T. Lei, M. de Jong, S. Ontanon, S. Brahma, Y. Zemlyanskiy, D. Uthus, M. Guo, J. Lee-Thorp, Y. Tay, Y.-H. Sung, and S. Sanghai.  CoLT5: Faster long-range transformers with conditional computation.  In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 5085–5100, Singapore, Dec. 2023. Association for Computational Linguistics.
  * [2] J. Andreas, M. Rohrbach, T. Darrell, and D. Klein.  Neural module networks.  In Proceedings of the 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 39–48, 2016.
  * [3] A. Ansell, E. Ponti, A. Korhonen, and I. Vulić.  Composable sparse fine-tuning for cross-lingual transfer.  In S. Muresan, P. Nakov, and A. Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1778–1796, Dublin, Ireland, May 2022. Association for Computational Linguistics.
  * [4] A. Bakhtiarnia, Q. Zhang, and A. Iosifidis.  Improving the accuracy of early exits in multi-exit architectures via curriculum learning.  In Proceedings of the 2021 International Joint Conference on Neural Networks (IJCNN), pages 1–8. IEEE, 2021.
  * [5] E. Bengio, P.-L. Bacon, J. Pineau, and D. Precup.  Conditional computation in neural networks for faster models.  arXiv preprint arXiv:1511.06297, 2015.
  * [6] F. Biadsy, Y. Chen, X. Zhang, O. Rybakov, A. Rosenberg, and P. Moreno.  A scalable model specialization framework for training and inference using submodels and its application to speech model personalization.  In Interspeech 2022. ISCA, 2022.
  * [7] S. Bills, N. Cammarata, D. Mossing, H. Tillman, L. Gao, G. Goh, I. Sutskever, J. Leike, J. Wu, and W. Saunders.  Language models can explain neurons in language models. 2023.  OpenAI Blog, 14, 2023.
  * [8] T. Bolukbasi, J. Wang, O. Dekel, and V. Saligrama.  Adaptive neural networks for efficient inference.  In Proceedings of the 34th International Conference on Machine Learning - Volume 70, ICML’17, page 527–536. JMLR.org, 2017.
  * [9] D. Bolya, C.-Y. Fu, X. Dai, P. Zhang, C. Feichtenhofer, and J. Hoffman.  Token merging: Your ViT but faster.  In Proceedings of the 2023 International Conference on Learning Representations (ICLR), 2023.
  * [10] D. Bolya and J. Hoffman.  Token merging for fast stable diffusion.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 4598–4602, 2023.
  * [11] A. Bontempelli, S. Teso, K. Tentori, F. Giunchiglia, and A. Passerini.  Concept-level debugging of part-prototype networks.  arXiv preprint arXiv:2205.15769, 2022.
  * [12] T. Chen, Y. Cheng, Z. Gan, L. Yuan, L. Zhang, and Z. Wang.  Chasing sparsity in vision transformers: An end-to-end exploration.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 19974–19988, 2021.
  * [13] X. Chen, Z. Liu, H. Tang, L. Yi, H. Zhao, and S. Han.  Sparsevit: Revisiting activation sparsity for efficient high-resolution vision transformer.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 2061–2070, June 2023.
  * [14] R. Child, S. Gray, A. Radford, and I. Sutskever.  Generating long sequences with sparse transformers.  arXiv preprint arXiv:1904.10509, 2019.
  * [15] A. Clark, D. De Las Casas, A. Guy, A. Mensch, M. Paganini, J. Hoffmann, B. Damoc, B. Hechtman, T. Cai, S. Borgeaud, G. B. Van Den Driessche, E. Rutherford, T. Hennigan, M. J. Johnson, A. Cassirer, C. Jones, E. Buchatskaya, D. Budden, L. Sifre, S. Osindero, O. Vinyals, M. Ranzato, J. Rae, E. Elsen, K. Kavukcuoglu, and K. Simonyan.  Unified scaling laws for routed language models.  In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 4057–4086. PMLR, 17–23 Jul 2022.
  * [16] G. M. Correia, V. Niculae, and A. F. T. Martins.  Adaptively sparse transformers.  In K. Inui, J. Jiang, V. Ng, and X. Wan, editors, Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 2174–2184. Association for Computational Linguistics, Nov. 2019.
  * [17] D. Dai, C. Deng, C. Zhao, R. Xu, H. Gao, D. Chen, J. Li, W. Zeng, X. Yu, Y. Wu, et al.  Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models.  arXiv preprint arXiv:2401.06066, 2024.
  * [18] T. Darcet, M. Oquab, J. Mairal, and P. Bojanowski.  Vision transformers need registers.  In Proceedings of the 2024 International Conference on Learning Representations (ICLR), 2024.
  * [19] T. Dettmers and L. Zettlemoyer.  Sparse networks from scratch: Faster training without losing performance.  arXiv preprint arXiv:1907.04840, 2019.
  * [20] W. Fedus, J. Dean, and B. Zoph.  A review of sparse expert models in deep learning.  arXiv preprint arXiv:2209.01667, 2022.
  * [21] W. Fedus, B. Zoph, and N. Shazeer.  Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.  The Journal of Machine Learning Research, 23(1):5232–5270, 2022.
  * [22] Q. Fournier, G. M. Caron, and D. Aloise.  A practical survey on faster and lighter transformers.  ACM Computing Surveys, 55(14s):1–40, 2023.
  * [23] A. Ghodrati, B. E. Bejnordi, and A. Habibian.  Frameexit: Conditional early exiting for efficient video recognition.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 15608–15618, 2021.
  * [24] Y. Han, G. Huang, S. Song, L. Yang, H. Wang, and Y. Wang.  Dynamic neural networks: A survey.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(11):7436–7456, 2021.
  * [25] Y. Han, Y. Pu, Z. Lai, C. Wang, S. Song, J. Cao, W. Huang, C. Deng, and G. Huang.  Learning to weight samples for dynamic early-exiting networks.  In European Conference on Computer Vision, pages 362–378. Springer, 2022.
  * [26] V. Hassija, V. Chamola, A. Mahapatra, A. Singal, D. Goel, K. Huang, S. Scardapane, I. Spinelli, M. Mahmud, and A. Hussain.  Interpreting black-box models: a review on explainable artificial intelligence.  Cognitive Computation, 16(1):45–74, 2024.
  * [27] J. B. Haurum, S. Escalera, G. W. Taylor, and T. B. Moeslund.  Which tokens to use? investigating token reduction in vision transformers.  In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 773–783, 2023.
  * [28] J. D. Havtorn, A. Royer, T. Blankevoort, and B. E. Bejnordi.  Msvit: Dynamic mixed-scale tokenization for vision transformers.  In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) Workshops, pages 838–848, October 2023.
  * [29] H. Hazimeh, Z. Zhao, A. Chowdhery, M. Sathiamoorthy, Y. Chen, R. Mazumder, L. Hong, and E. Chi.  Dselect-k: Differentiable selection in the mixture of experts with applications to multi-task learning.  Advances in Neural Information Processing Systems, 34:29335–29347, 2021.
  * [30] X. He, I. Keivanloo, Y. Xu, X. He, B. Zeng, S. Rajagopalan, and T. Chilimbi.  Magic pyramid: Accelerating inference with early exiting and token pruning.  In NeurIPS 2021 Workshop on Efficient Natural Language and Speech Processing, 2021.
  * [31] C. Herrmann, R. S. Bowen, and R. Zabih.  Channel selection using Gumbel Softmax.  In European Conference on Computer Vision, pages 241–257. Springer, 2020.
  * [32] L. Hou, R. Y. Pang, T. Zhou, Y. Wu, X. Song, X. Song, and D. Zhou.  Token dropping for efficient BERT pretraining.  In S. Muresan, P. Nakov, and A. Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 3774–3784, Dublin, Ireland, May 2022. Association for Computational Linguistics.
  * [33] Y. Huang, S. Hu, X. Han, Z. Liu, and M. Sun.  Unified view of grokking, double descent and emergent abilities: A perspective from circuits competition.  arXiv preprint arXiv:2402.15175, 2024.
  * [34] I. A. Huijben, W. Kool, M. B. Paulus, and R. J. Van Sloun.  A review of the Gumbel-max trick and its extensions for discrete stochasticity in machine learning.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(2):1353–1371, 2022.
  * [35] E. Jang, S. Gu, and B. Poole.  Categorical reparameterization with Gumbel-Softmax.  2017\.
  * [36] S. Jaszczur, A. Chowdhery, A. Mohiuddin, L. Kaiser, W. Gajewski, H. Michalewski, and J. Kanerva.  Sparse is enough in scaling transformers.  Advances in Neural Information Processing Systems, 34:9895–9907, 2021.
  * [37] A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford, D. S. Chaplot, D. d. l. Casas, E. B. Hanna, F. Bressand, et al.  Mixtral of experts.  arXiv preprint arXiv:2401.04088, 2024.
  * [38] Y.-H. Jin, K.-H. Lee, and D.-W. Choi.  Querynet: Querying neural networks for lightweight specialized models.  Information Sciences, 589:186–198, 2022.
  * [39] M. Jordan and R. Jacobs.  Hierarchical mixtures of experts and the em algorithm.  1993\.
  * [40] N. Kandpal, B. Lester, M. Muqeeth, A. Mascarenhas, M. Evans, V. Baskaran, T. Huang, H. Liu, and C. Raffel.  Git-theta: a git extension for collaborative development of machine learning models.  In Proceedings of the 40th International Conference on Machine Learning, ICML’23. JMLR.org, 2023.
  * [41] P. Karpikova, E. Radionova, A. Yaschenko, A. Spiridonov, L. Kostyushko, R. Fabbricatore, and A. Ivakhnenko.  Fiancee: Faster inference of adversarial networks via conditional early exits.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 12032–12043, 2023.
  * [42] M. Kim, S. Gao, Y.-C. Hsu, Y. Shen, and H. Jin.  Token fusion: Bridging the gap between token pruning and token merging.  In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), pages 1383–1392, January 2024.
  * [43] W. Kool, H. Van Hoof, and M. Welling.  Ancestral gumbel-top-k sampling for sampling without replacement.  The Journal of Machine Learning Research, 21(1):1726–1761, 2020.
  * [44] A. Kouris, S. I. Venieris, S. Laskaridis, and N. Lane.  Multi-exit semantic segmentation networks.  In European Conference on Computer Vision, pages 330–349. Springer, 2022.
  * [45] D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen.  Gshard: Scaling giant models with conditional computation and automatic sharding.  In Proceedings of the 2021 International Conference on Learning Representations (ICLR), 2021.
  * [46] Z. Lin, L. Wu, M. Wang, and L. Li.  Learning language specific sub-network for multilingual machine translation.  In C. Zong, F. Xia, W. Li, and R. Navigli, editors, Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 293–305, Online, Aug. 2021. Association for Computational Linguistics.
  * [47] H. Liu, K. Simonyan, and Y. Yang.  Darts: Differentiable architecture search.  2019\.
  * [48] T. Liu, M. Blondel, C. Riquelme, and J. Puigcerver.  Routers in vision mixture of experts: An empirical study.  arXiv preprint arXiv:2401.15969, 2024.
  * [49] C. J. Maddison, A. Mnih, and Y. W. Teh.  The concrete distribution: A continuous relaxation of discrete random variables.  2017\.
  * [50] Y. Matsubara, M. Levorato, and F. Restuccia.  Split computing and early exiting for deep learning applications: Survey and research challenges.  ACM Computing Surveys, 55(5):1–30, 2022.
  * [51] L. Meng, H. Li, B.-C. Chen, S. Lan, Z. Wu, Y.-G. Jiang, and S.-N. Lim.  Adavit: Adaptive vision transformers for efficient image recognition.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 12309–12318, 2022.
  * [52] G. Mialon, R. Dessì, M. Lomeli, C. Nalmpantis, R. Pasunuru, R. Raileanu, B. Rozière, T. Schick, J. Dwivedi-Yu, A. Celikyilmaz, E. Grave, Y. LeCun, and T. Scialom.  Augmented language models: a survey.  Transactions on Machine Learning Research, 6:1–35, 2023.
  * [53] S. Mittal, Y. Bengio, and G. Lajoie.  Is a modular architecture enough?  Advances in Neural Information Processing Systems, 35:28747–28760, 2022.
  * [54] M. Mohammed, H. Liu, and C. Raffel.  Models with conditional computation learn suboptimal solutions.  In I Can’t Believe It’s Not Better Workshop: Understanding Deep Learning Through Empirical Falsification, 2022.
  * [55] M. Muqeeth, H. Liu, Y. Liu, and C. Raffel.  Learning to route among specialized experts for zero-shot generalization.  arXiv preprint arXiv:2402.05859, 2024.
  * [56] M. Muqeeth, H. Liu, and C. Raffel.  Soft merging of experts with adaptive routing.  arXiv preprint arXiv:2306.03745, 2023.
  * [57] B. Mustafa, C. Riquelme, J. Puigcerver, R. Jenatton, and N. Houlsby.  Multimodal contrastive learning with limoe: the language-image mixture of experts.  Advances in Neural Information Processing Systems, 35:9564–9576, 2022.
  * [58] V. Niculae, C. F. Corro, N. Nangia, T. Mihaylova, and A. F. Martins.  Discrete latent structure in neural networks.  arXiv preprint arXiv:2301.07473, 2023.
  * [59] B. Pan, R. Panda, Y. Jiang, Z. Wang, R. Feris, and A. Oliva.  Ia-red2: Interpretability-aware redundancy reduction for vision transformers.  Advances in Neural Information Processing Systems, 34:24898–24911, 2021.
  * [60] Z. Pan, J. Cai, and B. Zhuang.  Stitchable neural networks.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 16102–16112, 2023.
  * [61] Z. Pan, B. Zhuang, H. He, J. Liu, and J. Cai.  Less is more: Pay less attention in vision transformers.  Proceedings of the AAAI Conference on Artificial Intelligence, 36(2):2035–2043, Jun. 2022.
  * [62] N. Passalis, J. Raitoharju, A. Tefas, and M. Gabbouj.  Efficient adaptive inference for deep convolutional neural networks using hierarchical early exits.  Pattern Recognition, 105:107346, 2020.
  * [63] J. Pfeiffer, G. Geigle, A. Kamath, J.-M. Steitz, S. Roth, I. Vulić, and I. Gurevych.  xGQA: Cross-lingual visual question answering.  In S. Muresan, P. Nakov, and A. Villavicencio, editors, Findings of the Association for Computational Linguistics: ACL 2022, pages 2497–2511. Association for Computational Linguistics, May 2022.
  * [64] J. Pfeiffer, N. Goyal, X. Lin, X. Li, J. Cross, S. Riedel, and M. Artetxe.  Lifting the curse of multilinguality by pre-training modular transformers.  In M. Carpuat, M.-C. de Marneffe, and I. V. Meza Ruiz, editors, Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 3479–3495, Seattle, United States, July 2022.
  * [65] J. Pfeiffer, A. Kamath, A. Rücklé, K. Cho, and I. Gurevych.  AdapterFusion: Non-destructive task composition for transfer learning.  In P. Merlo, J. Tiedemann, and R. Tsarfaty, editors, Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, pages 487–503. Association for Computational Linguistics, Apr. 2021.
  * [66] J. Pfeiffer, S. Ruder, I. Vulić, and E. M. Ponti.  Modular deep learning.  Transactions on Machine Learning Research, 11:1–76, 2023.
  * [67] M. Piórczyński, F. Szatkowski, K. Bałazy, and B. Wójcik.  Exploiting transformer activation sparsity with dynamic inference.  arXiv preprint arXiv:2310.04361, 2023.
  * [68] J. Pomponi, S. Scardapane, and A. Uncini.  A probabilistic re-intepretation of confidence scores in multi-exit models.  Entropy, 24:1, 2021.
  * [69] E. M. Ponti, A. Sordoni, Y. Bengio, and S. Reddy.  Combining modular skills in multitask learning.  arXiv preprint arXiv:2202.13914, 2022.
  * [70] J. Puigcerver, C. Riquelme, B. Mustafa, and N. Houlsby.  From sparse to soft mixtures of experts.  arXiv preprint arXiv:2308.00951, 2023.
  * [71] J. Puigcerver, C. Riquelme, B. Mustafa, C. Renggli, A. S. Pinto, S. Gelly, D. Keysers, and N. Houlsby.  Scalable transfer learning with expert models.  In Proceedings of the 2021 International Conference on Learning Representations (ICLR), 2021.
  * [72] Z. Qiu, Z. Huang, and J. Fu.  Emergent mixture-of-experts: Can dense pre-trained transformers benefit from emergent modular structures?  arXiv preprint arXiv:2310.10908, 2023.
  * [73] S. Rajbhandari, C. Li, Z. Yao, M. Zhang, R. Y. Aminabadi, A. A. Awan, J. Rasley, and Y. He.  DeepSpeed-MoE: Advancing mixture-of-experts inference and training to power next-generation AI scale.  In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 18332–18346. PMLR, 2022.
  * [74] J. Rajendran, A. Lakshminarayanan, M. M. Khapra, P. P, and B. Ravindran.  Attend, adapt and transfer: Attentive deep architecture for adaptive transfer from multiple sources in the same domain.  In Proceedings of the 2017 International Conference on Learning Representations (ICLR), 2017.
  * [75] Y. Rao, W. Zhao, B. Liu, J. Lu, J. Zhou, and C.-J. Hsieh.  Dynamicvit: Efficient vision transformers with dynamic token sparsification.  Advances in Neural Information Processing Systems, 34:13937–13949, 2021.
  * [76] C. Renggli, A. S. Pinto, N. Houlsby, B. Mustafa, J. Puigcerver, and C. Riquelme.  Learning to merge tokens in vision transformers.  arXiv preprint arXiv:2202.12015, 2022.
  * [77] C. Riquelme, J. Puigcerver, B. Mustafa, M. Neumann, R. Jenatton, A. Susano Pinto, D. Keysers, and N. Houlsby.  Scaling vision with sparse mixture of experts.  Advances in Neural Information Processing Systems, 34:8583–8595, 2021.
  * [78] C. Rosenbaum, T. Klinger, and M. Riemer.  Routing networks: Adaptive selection of non-linear functions for multi-task learning.  In Proceedings of the 2018 International Conference on Learning Representations (ICLR), 2018.
  * [79] S. Sarti, E. Lomurno, and M. Matteucci.  Anticipate, ensemble and prune: Improving convolutional neural networks via aggregated early exits.  Procedia Computer Science, 222:519–528, 2023.
  * [80] S. Scardapane, D. Comminiello, M. Scarpiniti, E. Baccarelli, and A. Uncini.  Differentiable branching in deep networks for fast inference.  In Proceedings of the 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 4167–4171. IEEE, 2020.
  * [81] S. Scardapane, M. Scarpiniti, E. Baccarelli, and A. Uncini.  Why should we add early exits to neural networks?  Cognitive Computation, 12(5):954–966, 2020.
  * [82] T. Schuster, A. Fisch, J. Gupta, M. Dehghani, D. Bahri, V. Tran, Y. Tay, and D. Metzler.  Confident adaptive language modeling.  Advances in Neural Information Processing Systems, 35:17456–17472, 2022.
  * [83] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  In Proceedings of the 2017 International Conference on Learning Representations (ICLR), 2017.
  * [84] T. Shen, M. Ott, M. Auli, and M. Ranzato.  Mixture models for diverse machine translation: Tricks of the trade.  In International Conference on Machine Learning, pages 5719–5728. PMLR, 2019.
  * [85] E. C. Strinati, P. Di Lorenzo, V. Sciancalepore, A. Aijaz, M. Kountouris, D. Gündüz, P. Popovski, M. Sana, P. A. Stavrou, B. Soret, et al.  Goal-oriented and semantic communication in 6g ai-native networks: The 6g-goals approach.  arXiv preprint arXiv:2402.07573, 2024.
  * [86] S. Tan, Y. Shen, Z. Chen, A. Courville, and C. Gan.  Sparse universal transformer.  In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 169–179. Association for Computational Linguistics, Dec. 2023.
  * [87] S. Tang, Y. Wang, Z. Kong, T. Zhang, Y. Li, C. Ding, Y. Wang, Y. Liang, and D. Xu.  You need multiple exiting: Dynamic early exiting for accelerating unified vision language model.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10781–10791, 2023.
  * [88] S. Teerapittayanon, B. McDanel, and H.-T. Kung.  Branchynet: Fast inference via early exiting from deep neural networks.  In Proceedings of the 2016 23rd International Conference on Pattern Recognition (ICPR), pages 2464–2469. IEEE, 2016.
  * [89] M. Treviso, A. Góis, P. Fernandes, E. Fonseca, and A. Martins.  Predicting attention sparsity in transformers.  In A. Vlachos, P. Agrawal, A. Martins, G. Lampouras, and C. Lyu, editors, Proceedings of the Sixth Workshop on Structured Prediction for NLP, pages 67–81. Association for Computational Linguistics, May 2022.
  * [90] T. Verelst and T. Tuytelaars.  Dynamic convolutions: Exploiting spatial sparsity for faster inference.  In Proceedings of the 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, 2020.
  * [91] T. Vu, A. Barua, B. Lester, D. Cer, M. Iyyer, and N. Constant.  Overcoming catastrophic forgetting in zero-shot cross-lingual generation.  In Y. Goldberg, Z. Kozareva, and Y. Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 9279–9300, Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational Linguistics.
  * [92] H. Wang, T. Fu, Y. Du, W. Gao, K. Huang, Z. Liu, P. Chandak, S. Liu, P. Van Katwyk, A. Deac, et al.  Scientific discovery in the age of artificial intelligence.  Nature, 620(7972):47–60, 2023.
  * [93] M. Wang, J. Mo, J. Lin, Z. Wang, and L. Du.  Dynexit: A dynamic early-exit strategy for deep residual networks.  In Proceedings of the 2019 IEEE International Workshop on Signal Processing Systems (SiPS), pages 178–183. IEEE, 2019.
  * [94] X. Wang and Y. Li.  Harmonized dense knowledge distillation training for multi-exit architectures.  Proceedings of the AAAI Conference on Artificial Intelligence, 35:10218–10226, 2021.
  * [95] B. Wójcik, A. Devoto, K. Pustelnik, P. Minervini, and S. Scardapane.  Adaptive computation modules: Granular conditional computation for efficient inference.  arXiv preprint arXiv:2312.10193, 2023.
  * [96] M. Wołczyk, B. Wójcik, K. Bałazy, I. T. Podolak, J. Tabor, M. Śmieja, and T. Trzcinski.  Zero time waste: Recycling predictions in early exit neural networks.  Advances in Neural Information Processing Systems, 34:2516–2528, 2021.
  * [97] J. Xin, R. Tang, J. Lee, Y. Yu, and J. Lin.  DeeBERT: Dynamic early exiting for accelerating BERT inference.  In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2246–2251. Association for Computational Linguistics, July 2020.
  * [98] Q. Xing, M. Xu, T. Li, and Z. Guan.  Early exit or not: Resource-efficient blind quality enhancement for compressed images.  In European Conference on Computer Vision, pages 275–292. Springer, 2020.
  * [99] X. Xu, S. Wang, Y. Chen, Y. Zheng, Z. Wei, and J. Liu.  Gtp-vit: Efficient vision transformers via graph-based token propagation.  In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), pages 86–95, January 2024.
  * [100] F. Xue, V. Likhosherstov, A. Arnab, N. Houlsby, M. Dehghani, and Y. You.  Adaptive computation with elastic input sequence.  In Proceedings of the 40th International Conference on Machine Learning, ICML’23. JMLR.org, 2023.
  * [101] F. Xue, Z. Zheng, Y. Fu, J. Ni, Z. Zheng, W. Zhou, and Y. You.  Openmoe: An early effort on open mixture-of-experts language models.  arXiv preprint arXiv:2402.01739, 2024.
  * [102] J. Yang, Q. Zhang, B. Ni, L. Li, J. Liu, M. Zhou, and Q. Tian.  Modeling point clouds with self-attention and gumbel subset sampling.  In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 3318–3327, 2019.
  * [103] H. Yin, A. Vahdat, J. M. Alvarez, A. Mallya, J. Kautz, and P. Molchanov.  A-vit: Adaptive tokens for efficient vision transformer.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10809–10818, 2022.
  * [104] S. E. Yuksel, J. N. Wilson, and P. D. Gader.  Twenty years of mixture of experts.  IEEE Transactions on Neural Networks and Learning Systems, 23(8):1177–1193, 2012.
  * [105] T. Zadouri, A. Üstün, A. Ahmadian, B. Ermiş, A. Locatelli, and S. Hooker.  Pushing mixture of experts to the limit: Extremely parameter efficient moe for instruction tuning.  arXiv preprint arXiv:2309.05444, 2023.
  * [106] G. Zhang, Q. Hu, Z. Qin, Y. Cai, and G. Yu.  A unified multi-task semantic communication system with domain adaptation.  In GLOBECOM 2022-2022 IEEE Global Communications Conference, pages 3971–3976. IEEE, 2022.
  * [107] X. Zhang, Y. Shen, Z. Huang, J. Zhou, W. Rong, and Z. Xiong.  Mixture of attention heads: Selecting attention heads per token.  In Y. Goldberg, Z. Kozareva, and Y. Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 4150–4162. Association for Computational Linguistics, Dec. 2022.
  * [108] Z. Zhang, Y. Lin, Z. Liu, P. Li, M. Sun, and J. Zhou.  MoEfication: Transformer feed-forward layers are mixtures of experts.  In S. Muresan, P. Nakov, and A. Villavicencio, editors, Findings of the Association for Computational Linguistics: ACL 2022, pages 877–890, Dublin, Ireland, May 2022. Association for Computational Linguistics.
  * [109] Z. Zhang, Z. Zeng, Y. Lin, C. Xiao, X. Wang, X. Han, Z. Liu, R. Xie, M. Sun, and J. Zhou.  Emergent modularity in pre-trained transformers.  In A. Rogers, J. Boyd-Graber, and N. Okazaki, editors, Findings of the Association for Computational Linguistics: ACL 2023, pages 4066–4083. Association for Computational Linguistics, July 2023.
  * [110] W. Zhou, C. Xu, T. Ge, J. McAuley, K. Xu, and F. Wei.  Bert loses patience: Fast and robust inference with early exit.  Advances in Neural Information Processing Systems, 33:18330–18341, 2020.
  * [111] Y. Zhou, T. Lei, H. Liu, N. Du, Y. Huang, V. Zhao, A. M. Dai, Q. V. Le, J. Laudon, et al.  Mixture-of-experts with expert choice routing.  Advances in Neural Information Processing Systems, 35:7103–7114, 2022.
  * [112] J. Zhu, X. Zhu, W. Wang, X. Wang, H. Li, X. Wang, and J. Dai.  Uni-perceiver-moe: Learning sparse generalist models with conditional moes.  Advances in Neural Information Processing Systems, 35:2664–2678, 2022.
  * [113] W. Zhu.  LeeBERT: Learned early exit for BERT with cross-level optimization.  In C. Zong, F. Xia, W. Li, and R. Navigli, editors, Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 2968–2980, Online, Aug. 2021. Association for Computational Linguistics.
  * [114] S. Zuo, X. Liu, J. Jiao, Y. J. Kim, H. Hassan, R. Zhang, T. Zhao, and J. Gao.  Taming sparsely activated transformer with stochastic experts.  In Proceedings of the 2022 International Conference on Learning Representations (ICLR), 2022.