# Prescriptive Machine Learning for Automated Decision Making: Challenges and
Opportunities

Eyke Hüllermeier
Institute of Informatics
University of Munich (LMU)
<EMAIL_ADDRESS>
###### Abstract

Recent applications of machine learning (ML) reveal a noticeable shift from
its use for _predictive_ modeling in the sense of a data-driven construction
of models mainly used for the purpose of prediction (of ground-truth facts) to
its use for _prescriptive_ modeling. What is meant by this is the task of
learning a model that stipulates appropriate decisions about the right course
of action in real-world scenarios: Which medical therapy should be applied?
Should this person be hired for the job? As argued in this article,
prescriptive modeling comes with new technical conditions for learning and new
demands regarding reliability, responsibility, and the ethics of decision
making. Therefore, to support the data-driven design of decision-making agents
that act in a rational but at the same time responsible manner, a rigorous
methodological foundation of prescriptive ML is needed. The purpose of this
short paper is to elaborate on specific characteristics of prescriptive ML and
to highlight some key challenges it implies. Besides, drawing connections to
other branches of contemporary AI research, the grounding of prescriptive ML
in a (generalized) decision-theoretic framework is advocated.

_Keywords_ Predictive modeling $\cdot$ Prescriptive modeling $\cdot$ Decision
theory $\cdot$ Bounded rationality $\cdot$ Uncertainty $\cdot$ Preference
learning $\cdot$ Counterfactual reasoning $\cdot$ AI ethics $\cdot$ Fairness

##  1 Introduction

Machine learning (ML) methodology, fueled with access to ever-increasing
masses of data and unprecedented computing power, has been the main driving
factor of recent progress in artificial intelligence (AI) and its applications
in various branches of science and technology, industry and business,
economics and finance, amongst others. In this regard, ML is most commonly
perceived as a means for _predictive modeling_ , that is, for the data-driven
construction of a model that is mainly used for the purpose of predicting
unknown facts in a specific context — albeit models may, of course, serve
other purposes, too, such as understanding and explanation, or may have a more
descriptive flavor. A predictive model, or “predictor” in ML jargon, is
trained in a supervised manner on cases encountered by the “learner” over the
course of time, such as emails categorized as spam or non-spam, and the model
is then used to make predictions in future situations, e.g., to automatically
mark new emails.

Looking at emerging applications of ML methodology, there is a visible shift
from predictive modeling to _prescriptive modeling_ , by which we mean the
task of learning a model that stipulates appropriate decisions or actions to
be taken in real-world scenarios. In fact, decisions are nowadays increasingly
automated and made by algorithms instead of humans, and most of these
_automated decision making_ (ADM) algorithms are trained on data using ML
methods. For example, think of decisions in the context of employees
recruitment, such as hiring or placement decisions [41]. Another important
case is the data-driven construction of individualized treatment rules in
personalized medicine [62], and further examples can be found in other
domains, including jurisdiction [33], finance [26], and disaster management
[60]. The distinction between predictive and prescriptive modeling can be
compared with the distinction between predictive (“What will happen?”) and
prescriptive (“How to make something happen?”) tasks as commonly made in
business analytics [8].

As shown by the above examples of ADM, machine learning is gaining in societal
significance and increasingly impacting on our daily life, a development that
comes with increasing demands for ML methodology: “prescriptive ML” is
supposed to (i) stipulate decisions in a _transparent_ , _principled_ (rather
than ad hoc), and _ethically justifiable_ manner, and (ii) to be sufficiently
aware of its own (in)competence, notably its (un)certainty about the right
course of action, so as to assure its _reliability_ and _responsibility_ [35,
13].

In addition to raising new demands, prescriptive modeling also exhibits
notable technical differences compared to standard ML settings. In particular,
the key assumption of an objective “ground truth” underlying every prediction,
revealed to the learner as training information and serving as a reference to
evaluate the correctness of predictions, is no longer justified. For example,
there is nothing like a ground-truth therapy plan for a patient, a ground-
truth dietary plan for an athlete preparing for a competition, or a “true”
abstract of an article in automatic text summarization [23]. Even when taking
the stance that _optimality_ may serve as a surrogate for truth, an optimal
therapy or plan could never be determined or verified, as it would require
counterfactual reasoning about (perhaps infinitely many) “possible worlds”
(e.g., how the patient would have developed under each possible therapy, or
the athlete would have performed under each possible dietry) [44].

As detailed further below, prescriptive ML also exhibits other characteristics
distinguishing it from predictive ML, and showing that standard ML methodology
does not fully comply with the demands and technical subtleties of
prescriptive modeling. What is needed, therefore, is an ML methodology
specifically tailored to the task of prescriptive modeling.

##  2 Predictive versus Prescriptive ML

Predictive modeling aims to capture the (stochastic) dependence between
_instances_ $X$ and associated _outcomes_ $Y$. In the standard setting of
supervised learning, training data consists of a set $\mathcal{D}$ of examples
in the form of tuples $(x,y)\in\mathcal{X}\times\mathcal{Y}$, with
$\mathcal{X}$ being the instance space and $\mathcal{Y}$ the set of possible
outcomes. Typically, these examples are assumed to be independent and
identically distributed (i.i.d.) according to some unknown probability measure
$P$ on $\mathcal{X}\times\mathcal{Y}$ (with mass or density function $p$),
which is the target of _generative_ learning. _Discriminative_ learning
proceeds from a _hypothesis space_
$\mathcal{H}\subset\mathcal{Y}^{\mathcal{X}}$, where each hypothesis is a map
(predictor) $\mathcal{X}\longrightarrow\mathcal{Y}$, and a loss function
$\ell:\,\mathcal{Y}\times\mathcal{Y}\longrightarrow\mathbb{R}$. The learner
then seeks the “best predictor”, _viz._ the expected loss minimizer (Bayes
predictor)

$h^{*}\in\operatorname*{argmin}_{h\in\mathcal{H}}\int_{\mathcal{X}\times\mathcal{Y}}\ell(h(x),y)\,d\,P\,,$
(1)

and leverages the data $\mathcal{D}$ to produce an approximation $h$ thereof.
The latter can then be used to make predictions $\hat{y}_{q}=h(x_{q})$ for new
query instances $x_{q}\in\mathcal{X}$.

A basic assumption in this setting is that $Y$ does (or will) exist
independently of the prediction. In other words, for every concrete instance
$X$, there is a “ground truth” $Y$, and hence a ground-truth (albeit non-
deterministic) dependence between instances and outcomes (specified by the map
$X\mapsto p(Y\,|\,X)$). For example, if the prediction is a weather forecast
for tomorrow, based on weather conditions in the last days ($X$), the “ground-
truth weather” $Y$ does (or will) indeed exist, independently of the
prediction, just like a disease exists independently of the diagnosis or a
handwritten digit independently of its image-based prediction.

This view of outcomes $Y$ or, more generally, a dependence between instances
and outcomes, as “ontic” entities that ought to be “discovered”, and hence of
ML as an _analytic_ task, is arguably less appropriate in the context of ADM,
where $Y$ is a decision made by the learner in a situation $X$. Here, the
problem of constructing a decision model is essentially of _synthetic_ nature.
Imagine, for example, that $Y$ is not a diagnosis but a drug or a therapy. In
this case, $Y$ is not an ontic entity and should be perceived as a
_prescription_ rather than a prediction. As a consequence, a real ground truth
may not exist in prescriptive ML. Moreover, as explained above, the notion of
“best decision” is not a good surrogate of a ground truth either, as it is
counterfactual and cannot be determined in practice, especially not
retrospectively. In fact, by taking actions in the real world, a prescriptive
learner is influencing and possibly changing the world through intervention,
in contrast to predictive modeling, where the learner is reasoning about the
world from an outside perspective (albeit an indirect influence is possible if
decisions are made by others based on the learner’s predictions [8, 18]).
Consequently, the goal of learning shifts from targeting a ground truth to
finding a “practicable” decision model, where practicability will typically
refer to several criteria simultaneously.

Thus, while predictive modeling is essentially an analytic task, prescriptive
modeling is of synthetic nature. Instead of hypothesizing about an underlying
ground truth, which may not exist or at least be unavailable for training and
evaluation, prescriptive ML is seeking practicable decision models that
balance value, complexity, cost, fairness, etc. in an appropriate manner.
These differences between predictive and prescriptive ML have implications for
the formalization of ML problems and pose several challenges for the
methodology needed to tackle them.

  * •

Challenge 1: Weak supervision. Training information in prescriptive ML will
typically be weak, biased, incomplete, and possibly context-dependent, thereby
posing severe challenges for learning. This is especially true for
observational data collected from human DMs acting as “teachers” [2, 20]. In
fact, decisions made in the past are presumably suboptimal and likely to be
biased toward the teacher’s preferences, so that, in contrast to supervised
learning, observed decisions cannot be considered as ground-truth examples
suggesting the best course of action in a given context. Besides, training
information will often be incomplete, lacking information about
“counterfactuals”, i.e., outcomes for those decisions that were not made, and
hence raising the need for counterfactual reasoning [30]. At best, a decision
is accompanied by some hint at how effective it was, or perhaps by a relative
comparison with another decision — for example, if one candidate $A$ failed to
accomplish a job while a second one $B$ succeeded, this could be captured in
terms of a preference $B\succ A$. Finally, decisions may exhibit various types
of dependencies. Often, for example, decisions are not made individually but
collectively, i.e., in the context of a set of other problems. For instance,
the decision of whether or not a candidate is accepted for a study program
with limited capacity depends on which other candidates also applied.
Obviously, standard assumptions on the data-generating process as commonly
made in the predictive modeling setup, such as i.i.d. observations, are
violated in such cases.

  * •

Challenge 2: Complex prescriptions and performance criteria. Decisions are not
always atomic units but may possess an internal structure, so that decision
spaces are structured spaces. For example, the structure can be combinatorial
if a (global) decision comprises several sub-decisions [14, 17], like in a
multi-step therapy, where the doctor may need to decide about the type of
medication, the dosages, and the time in-between the administrations.
Moreover, in cases where the learner is only supposed to make recommendations,
while leaving the final course of action to the user, the prescription of a
sole decision may not provide optimal support. Instead, a recommendation might
be a “pre-prescription” in the form of a _set_ of promising decisions, i.e., a
subset of candidates the user can choose from [5] — including complete
abstention in the extreme case. Or, the candidates could be further
prioritized by returning a ranking instead of an unsorted set. Regarding
prescriptive performance, the outcome of a decision might be rated in terms of
_multiple performance criteria_ (e.g., medicinal efficacy of a therapy, side-
effects, cost, etc.) at the same time, which may lead to incomparability (or
weaker comparability, for example, Pareto-dominance) and further increase the
complexity of the prescription task [11, 15].

  * •

Challenge 3: Uncertainty representation. The notion of uncertainty is of
utmost importance and plays a fundamental role in prescriptive modeling — see
e.g. [25] for a discussion from an ethical and epistemological perspective.
First of all, the learner’s uncertainty-awareness and the adequate
representation of its “prescriptive uncertainty” is a key prerequisite for
responsible and reliable decision making, including the possibility to delay
or abstain from a decision if uncertainty is too high. Deciding about the
right course of action will require a distinction between different types and
sources of uncertainty, notably between _aleatoric_ and _epistemic_
uncertainty, whence uncertainty representation should go beyond standard
probabilistic models [29]. Likewise, a proper representation of uncertainty is
essential for the two main components of a decision-theoretic framework as
advocated in Section 4, namely preference and uncertainty (about outcomes of
decisions in a given context), because these are learned from data and hence
hypothetical; again, a distinction between aleatoric and epistemic uncertainty
is important in this regard, with the latter essentially representing
uncertainty about uncertainty.

  * •

Challenge 4: Complexity and resource constraints. Facing a model construction
task of synthetic nature, one will naturally strive for a good compromise
between value and complexity of a decision model in prescriptive ML. Indeed,
if the learner has to make decisions under real-world conditions, constraints
regarding _time and computational resources_ as well as the availability of
information must be taken into account. In fact, depending on the situation
and application context, time and computational resources for collecting
information and applying decision models might be limited. Moreover, in the
context of prescriptive ML, the decision maker is not necessarily a powerful
computer but perhaps a human supported by a low-level computing device (e.g.,
a smartphone). The DM’s resources to collect, validate, and enter data might
be scarce, or decisions must be made quickly. In the extreme case, instead of
running an algorithm on a computer, the decision is made by the human himself
— imagine, e.g., a medical doctor who needs to make decisions in emergency
cases. In such situations, mental capacity will only allow for applying a
“rule of thumb”. In any case, the complexity of a decision model should be
well adapted to the experience of the DM, the information at hand, the
computational resources available, as well as possible time constraints to be
met.

  * •

Challenge 5: Ethics and fairness. Acting in the world comes with a certain
responsibility, especially if decisions have a direct influence on people.
Consequently, in addition to rationality, _ethical criteria_ such as fairness
have a major role to play in prescriptive ML. Broadly speaking, decisions
should not only be rational but also fair and ethical. Although the notion of
fairness in ML has been studied quite intensively in the recent past, there is
still very little work on combining fairness and rationality, for example, by
formalizing fairness in a decision-theoretic context [31, 21]. Moreover, ML
algorithms must be generalized so as to allow for learning fair decision
models. This will also influence the learning objectives and performance
measures. For example, the optimization of _average_ or _expected performance_
might be inappropriate in a social context, as it may suggest sacrificing
performance on an ethnic minority for the benefit of a bigger subpopulation
[52].

##  3 Positioning and Related Topics

The field of machine learning has grown quickly over the last decades, and a
wide repertoire of ML methodologies is now available. Prescriptive ML has
connections to many of these methodologies, including weakly supervised
learning [64], preference learning [22], structured output prediction [4, 36,
16], as well as uncertainty representation and quantification [50, 32, 1, 29].
Besides, as already mentioned, recent work on ethical, legal, and societal
aspects of AI have recently come to the fore. Especially relevant for
prescriptive ML is recent work on _fairness_ [51, 38, 56], which has mainly
been formalized in terms of constraints that the predictions of an ML model
should obey [65, 51, 42, 61].

In terms of methodology, a basic distinction (according to the type of
supervision) is often made between the paradigms of supervised, unsupervised,
and reinforcement learning [53]. Adhering to this rough categorization,
prescriptive ML is perhaps best positioned in-between the latter two. The type
of training data in prescriptive ML resembles the type of feedback in
reinforcement learning (RL) and related settings such as (contextual) multi-
armed bandits [34]. With the latter, it also shares the view of a learner as
an _acting_ rather than only predicting agent. However, prescriptive ML, at
least in its basic form, is more geared toward _one-time_ rather than
sequential decisions, _passive_ rather than active, and _offline_ (“batch”)
rather than online learning, which are properties it has in common with the
setting of supervised learning. Moreover, algorithms for RL or bandits
essentially compensate for the lack of counterfactuals through trial and error
strategies, often requiring thousands of trials before a good policy is found
[37] — the excessive exploration required for optimization (albeit to a lesser
extent for _satisficing_ [46]) is not only inappropriate from an ethical point
of view but also in conflict with the quest for per-instance guarantees [25].

More relevant from this perspective is the problem of offline (policy)
learning from data of the form $(x,a,u)$, where information about the outcome
or utility $u$ caused by a decision $a$ made in a specific context $x$ is
given, but lacking information about counterfactuals (i.e., outcomes for
decisions $a^{\prime}\neq a$). Variants of this problem have recently been
considered in different fields. In machine learning, it is also referred to as
_learning from logged bandit feedback_ or _batch learning from bandit
feedback_ [10]. This setting requires causal inference and counterfactual
estimation [40, 12], on the basis of which (supervised) learning principles
such as empirical risk minimization can be generalized. For example,
_counterfactual risk minimization_ has been proposed as a principle for
learning in this setting [54, 55]. _Regress-and-compare_ approaches [45, 8]
train regression models for estimating (counterfactual) outcomes, one per
decision, and then do “estimated utility maximization” [28]. As another
approach for dealing with partial (observational) data, the use of _propensity
scores_ (conditional probability of a treatment/decision given the context
features) has received increasing interest in the field of (medical)
statistics [47]. Combined with weighted classification, it is used for
learning individualized treatment rules in personalized medicine [58, 63].
Also, adaptations of classification trees for minimizing a prescription error
have recently been proposed [30, 9]. Related to this is work on causal trees
and forests [3, 57, 43]. The use of so-called _instrument variables_ for the
purpose of counterfactual prediction, in combination with deep learning
methods, is proposed in [27]. Although the problem can be seen as a special
case of prescriptive modeling, the setups considered are quite specific and
often coming with strong assumptions (e.g., binary treatment case or known
propensities/logging policy). Moreover, major challenges regarding
reliability, responsibility, and ethics are not addressed.

##  4 A Decision-theoretic Approach

The idea of machine learning as a means for automated data-driven decision
making suggests a methodology of prescriptive ML grounded in (prescriptive)
_decision theory_ as a formal foundation of rational decision making under
uncertainty. A decision-theoretic grounding promises to support the systematic
development of prescriptive ML methods in a principled (rather than ad hoc)
manner. Moreover, it involves the elucidation of underlying assumptions,
values, and norms, as well as the explication and (axiomatic) justification of
a decision-theoretic principle, thereby contributing to the transparency and
social acceptability of algorithmic decisions [24].

Figure 1: In standard statistical decision theory (above), the decision
principle of expected utility maximization takes probability distributions for
uncertainty and utility functions for preference representation as input. In
the envisioned prescriptive ML methodology (below), prescriptions are produced
in the form of bounded rational decisions (middle). The latter are made on the
basis of suitably represented preferences (right) and uncertainty (left) in
the current context, which are both obtained through machine learning. The
methodology shall also accommodate additional ethical criteria (such as
fairness), taken as another “input” for decision making.

Statistical decision theory [19, 6, 7] essentially builds upon the notions of
_preference_ and _uncertainty_ , which are formalized in terms of utility
functions and probability distributions, respectively (cf. Fig. 1, upper
panel). While this theory does provide a suitable basis for standard
(predictive) ML — for example, the principle of _expected loss minimization_
in ML is in direct correspondence with the celebrated (and axiomatically
justified) principle of expected utility maximization [39, 49] — utility
functions and probability distributions are not expressive enough in the
context of prescriptive ML. For example, preferences are not pre-specified but
learned from data, and hence uncertain and maybe not readily expressible in
numeric form. Moreover, the quest for a self-aware decision maker calls for an
expressive representation of uncertainty, notably a distinction between
epistemic (“I don’t know”) and aleatoric (“I know it’s a matter of chance”)
uncertainty. What is arguably needed, therefore, is a novel decision-theoretic
foundation of prescriptive modeling built upon recent work on generalized
uncertainty representation in ML [29] as well as preference modeling and
learning [22].

Further, classical decision theory does not account for complexity and
computational aspects, which are essential from an ML perspective and may
indeed prevent a DM from realizing “perfect rationality”. Therefore, in
addition to uncertainty and preference learning, prescriptive ML should
leverage the concept of _bounded rationality_ [48, 59], thereby connecting
three subfields of modern AI (cf. Fig. 1, lower panel). Last but not least, a
methodology of prescriptive ML shall also be capable of accommodating _social
and ethical_ aspects, especially fairness [51, 38, 56], as key concepts of
contemporary AI and responsible decision making.

##  5 Summary

We emphasized the need for a sound methodological foundation of prescriptive
modeling, that is, the use of machine learning to support automated decision
making, and argued that standard ML methodology geared toward predictive
modeling does not fully comply with the specific characteristics of
prescriptive ML and the challenges it implies. These challenges include the
need for learning from weak and possibly biased training information that
lacks “optimal” decisions provided as reference, the quest for reliable
prescriptions with individual guarantees, for responsibility and fairness, as
well as an increased need for an adequate representation of uncertainty and
for taking resource constraints into consideration.

Viewing prescriptions as (bounded) rational decisions under uncertainty, we
advocate a methodology grounded in decision theory and leveraging novel
approaches to handling uncertainty and preference in machine learning. A
methodology of that kind would support the systematic development of
prescriptive ML methods in a principled manner. Moreover, as a decision-
theoretic approach involves the elucidation of underlying assumptions, values,
and norms, it may contribute to the transparency and social acceptability of
algorithmic decisions. As such, it would provide a theoretical foundation for
the data-driven design of decision-making agents that act in a rational but at
the same time responsible manner, and thereby a profound basis of algorithmic
decision making.

## References

  * Abdar et al. [2021] M. Abdar, F. Pourpanah, S. Hussain, D. Rezazadegan, L. Liu, M. Ghavamzadeh, P. Fieguth, X. Cao, A. Khosravi, U.R. Acharya, V. Makarenkov, and S. Nahavandi.  A review of uncertainty quantification in deep learning: Techniques, applications and challenges.  _Information Fusion_ , 76:243–297, 2021.  doi: 10.1016/j.inffus.2021.05.008.
  * Atan et al. [2018] O. Atan, W.R. Zame, and M. van der Schaar.  Learning optimal policies from observational data.  In _Proc. ICML, 35th Int. Conf. on Machine Learning_ , Stockholm, Sweden, 2018.
  * Athey and Imbens [2016] S. Athey and G. Imbens.  Recursive partitioning for heterogeneous causal effects.  _Proc. of the National Academy of Science_ , 113(27):7353–7360, 2016.
  * Bakir et al. [2007] G. Bakir, T. Hofmann, B. Schölkopf, A. Smola, B. Taskar, and S. Vishwanathan, editors.  _Predicting Structured Data_.  MIT Press, 2007.
  * Bengs and Hüllermeier [2020] V. Bengs and E. Hüllermeier.  Preselection bandits.  In _Proc. ICML, 37th Int. Conf. on Machine Learning_ , Vienna, Austria, 2020.
  * Berger [1980] J.O. Berger.  _Statistical Decision Theory: Foundations, Concepts, and Methods_.  Springer-Verlag, 1980.
  * Berger [1985] J.O. Berger.  _Statistical Decision Theory and Bayesian Analysis_.  Springer-Verlag, 1985.
  * Bertsimas and Kallus [2020] D. Bertsimas and N. Kallus.  From predictive to prescriptive analytics.  _Management Science_ , 66(3):1025–1044, 2020\.
  * Bertsimas et al. [2019] D. Bertsimas, J. Dunn, and N. Mundru.  Optimal prescriptive trees.  _INFORMS Journal on Optimization_ , 1(2):164–183, 2019.  doi: 10.1287/ijoo.2018.0005.
  * Beygelzimer and Langford [2009] A. Beygelzimer and J. Langford.  The offset tree for learning with partial labels.  In _Proc. KDD, 15th ACM SIGKDD Int. Conf. on Knowledge Discovery and Data Mining_ , pages 129–138, 2009.
  * Blaszczynski et al. [2007] J. Blaszczynski, S. Greco, and R. Slowinski.  Multi-criteria classification – A new scheme for application of dominance-based decision rules.  _European Journal of Operational Research_ , 181(3):1030–1044, 2007.
  * Bottou et al. [2013] L. Bottou, J. Peters, J.Q. Candela, D.X. Charles, M. Chickering, E. Portugaly, D. Ray, P.Y. Simard, and E. Snelson.  Counterfactual reasoning and learning systems: the example of computational advertising.  _Journal of Machine Learning Research_ , 14(1):3207–3260, 2013.
  * Casteluccia and Métayer [2019] C. Casteluccia and D. Le Métayer.  Understanding algorithmic decision making: opportunities and challenges.  Technical Report PE 624.261, EPRS Study, European Parlament, 2019.
  * Chevaleyre et al. [2008] Y. Chevaleyre, U. Endriss, J. Lang, and N. Maudet.  Preference handling in combinatorial domains: From AI to social choice.  _AI Magazine_ , 29(4), 2008.
  * Corrente et al. [2013] S. Corrente, S. Greco, M. Kadzinski, and R. Slowinski.  Robust ordinal regression in preference learning and ranking.  _Machine Learning_ , 93(2-3):381–422, 2013.
  * Doppa et al. [2014] J.R. Doppa, A. Fern, and P. Tadepalli.  Structured prediction via output space search.  _Journal of Machine Learning Research_ , 15:1317–1350, 2014\.
  * Dragone et al. [2018] P. Dragone, S. Teso, and A. Passerini.  Constructive preference elicitation over hybrid combinatorial spaces.  In _Proc. AAAI Conf. on Artificial Intelligence_ , 2018.
  * Elmachtoub et al. [2020] A.N. Elmachtoub, J.C.M. Liang, and R. McNellis.  Decision trees for decision-making under the predict-then-optimize framework.  In _Proc. ICML, 37th Int. Conf. on Machine Learning_ , volume 119 of _PMLR_ , 2020.
  * Ferguson [1967] T. Ferguson.  _Mathematical Statistics: A Decision Theoretic Approach_.  Academic Press, New York, 1967.
  * Fernandez and Provost [2019] C. Fernandez and F. Provost.  Observational vs experimental data when making automated decisions using machine learning.  Technical report, NYU Stern School of Business, 2019.  URL SSRN:https://ssrn.com/abstract=3444678.
  * Friedler et al. [2021] S.A. Friedler, C. Scheidegger, and S. Venkatasubramanian.  The (im)possibility of fairness: Different value systems require different mechanisms for fair decision making.  _Communications of the ACM_ , 64(4):136–143, 2021\.  doi: 10.1145/3433949.
  * Fürnkranz and Hüllermeier [2010] J. Fürnkranz and E. Hüllermeier, editors.  _Preference Learning_.  Springer-Verlag, 2010.
  * Gambhir and Gupta [2017] M. Gambhir and V. Gupta.  Recent automatic text summarization techniques: a survey.  _Artificial Intelligence Review_ , 47:1–66, 2017.  doi: 10.1007/s10462-016-9475-9.
  * Goodman and Flaxman [2017] R. Goodman and S. Flaxman.  European Union regulations on algorithmic decision-making and a “right to explanation”.  _AI Magazine_ , 38(3):1–9, 2017.
  * Grote and Berens [2020] T. Grote and P. Berens.  On the ethics of algorithmic decision-making in healthcare.  _Journal of Medical Ethics_ , 46(3):205–211, 2020\.  doi: 10.1136/medethics-2019-105586.
  * Hansen [2020] K.B. Hansen.  The virtue of simplicity: On machine learning models in algorithmic trading.  _Big Data and Society_ , January 2020.  doi: 10.1177/2053951720926558.
  * Hartford et al. [2017] J. Hartford, G. Lewis, K. Leyton-Brown, and M. Taddy.  Deep IV: A flexible approach for counterfactual prediction.  In _Proc. ICML, 34th Int. Conf. on Machine Learning_ , number 70 in PMLR, Sydney, Australia, 2017.
  * Hüllermeier [2005] E. Hüllermeier.  Experience-based decision making: A satisficing decision tree approach.  _IEEE Transactions on Systems, Man, and Cybernetics – Part A: Systems and Humans_ , 35(5):641–653, 2005.
  * Hüllermeier and Waegeman [2021] E. Hüllermeier and W. Waegeman.  Aleatoric and epistemic uncertainty in machine learning: An introduction to concepts and methods.  _Machine Learning_ , 110(3):457–506, 2021.  doi: 10.1007/s10994-021-05946-3.
  * Kallus [2017] N. Kallus.  Recursive partitioning for personalization using observational data.  In _Proc. ICML, 34th Int. Conf. on Machine Learning_ , volume 70 of _PPMLR_ , pages 1789–1798, 2017.
  * Kamiran et al. [2012] F. Kamiran, A. Karim, and X. Zhang.  Decision theory for discrimination-aware classification.  In _Proc. IEEE Int. Conf. on Data Mining_ , pages 924–929, 2012\.  doi: 10.1109/ICDM.2012.45.
  * Kendall and Gal [2017] A. Kendall and Y. Gal.  What uncertainties do we need in Bayesian deep learning for computer vision?  In _Proc. NIPS, Advances in Neural Information Processing Systems_ , pages 5574–5584, 2017.
  * Kleinberg et al. [2018] J. Kleinberg, H. Lakkaraju, J. Leskovec, J. Ludwig, and S. Mullainathan.  Human decisions and machine predictions.  _The Quarterly Journal of Economics_ , 133(1):237–293, 2018.
  * Lattimore and Szepesvari [2020] T. Lattimore and C. Szepesvari.  _Bandit Algorithms_.  Cambridge University Press, 2020.
  * Lepri et al. [2018] B. Lepri, N. Oliver, E. Letouzé, A. Pentland, and P. Vinck.  Fair, transparent, and accountable algorithmic decision-making processes.  _Philosophy and Technology_ , 31:611–627, 2018.  doi: 10.1007/s13347-017-0279-x.
  * Maes et al. [2009] F. Maes, L. Denoyer, and P. Gallinari.  Structured prediction with reinforcement learning.  _Machine Learning_ , 77(2–3):271–301, 2009.
  * Mannor and Tsitsiklis [2004] S. Mannor and J.N. Tsitsiklis.  The sample complexity of exploration in the multi-armed bandit problem.  _Journal of Machine Learning Research_ , 5:623–648, 2004\.
  * Mehrabi et al. [2021] N. Mehrabi, F. Morstatter, N. Saxena, K. Lerman, and A. Galstyan.  A survey on bias and fairness in machine learning.  _ACM Computing Surveys_ , 54(6):1–15, 2021.
  * Neumann and Morgenstern [1953] J. Von Neumann and O. Morgenstern.  _Theory of Games and Economic Behavior_.  John Wiley and Sons, 1953.
  * Pearl [2000] J. Pearl.  _Causality: Models, Reasoning, and Inference_.  Cambridge University Press, 2000.
  * Pessach et al. [2020] D. Pessach, G. Singer, D. Avrahamia, H. Chalutz Ben-Gal, E. Shmueli, and I. Ben-Gala.  Employees recruitment: A prescriptive analytics approach via machine learning and mathematical programming.  _Decision Support Systems_ , 134, 2020.
  * Pitoura et al. [2021] E. Pitoura, K. Stefanidis, and G. Koutrika.  Fairness in rankings and recommendations: An overview.  _CoRR_ , abs/2104.05994v1, 2021.  URL http://arxiv.org/abs/2104.05994v1.
  * Powers et al. [2017] S. Powers, J. Qian, K. Jung, A. Schuler, H.S. Nigam, T. Hastie, and R. Tibshirani.  Some methods for heterogenous treatment effect estimation in high dimensions.  Technical report, Stanford University, Palo Alto, CA, USA, 2017.
  * Prosperi et al. [2020] M. Prosperi, Y. Guo, M. Sperrin, J.S. Koopman, J.S. Min, X. He, S. Rich, M. Wang, I.E. Buchan, and J. Bian.  Causal inference and counterfactual prediction in machine learning for actionable healthcare.  _Nature Machine Intelligence_ , 2:369–375, 2020.
  * Qian and Murphy [2011] M. Qian and S.A. Murphy.  Performance guarantees for individualized treatment rules.  _Annals of Statistics_ , 39(2), 2011.
  * Reverdy et al. [2017] P. Reverdy, V. Srivastava, and N.E. Leonard.  Satisficing in multi-armed bandit problems.  _IEEE Transactions on Automatic Control_ , 62(8):3788–3803, 2017.
  * Rosenbaum and Rubin [1983] P. Rosenbaum and D. Rubin.  The central role of the propensity score in observational studies for causal effects.  _Biometrika_ , 70(1):41–55, 1983.
  * Russell [2016] S.J. Russell.  Rationality and intelligence: A brief update.  In V. Müller, editor, _Fundamental Issues of Artificial Intelligence_. Springer, Cham., 2016.
  * Savage [1954] L.J. Savage.  _The Foundations of Statistics_.  John Wiley and Sons, Inc., New York, 1954.
  * Senge et al. [2014] R. Senge, S. Bösner, K. Dembczynski, J. Haasenritter, O. Hirsch, N. Donner-Banzhoff, and E. Hüllermeier.  Reliable classification: Learning classifiers that distinguish aleatoric and epistemic uncertainty.  _Information Sciences_ , 255:16–29, 2014.
  * Shrestha and Yang [2019] Y.R. Shrestha and Y. Yang.  Fairness in algorithmic decision-making: Applications in multi-winner voting, machine learning, and recommender systems.  _Algorithms_ , 12(199), 2019.  doi: 10.3390/a12090199.
  * Slowik and Bottou [2021] A. Slowik and L. Bottou.  Algorithmic bias and data bias: Understanding the relation between distributionally robust optimization and data curation.  _CoRR_ , abs/2106.09467v1, 2021.  URL http://arxiv.org/abs/2106.09467v1.
  * Sutton and Barto [2018] R.S. Sutton and A.G. Barto.  _Reinforcement Learning: An Introduction_.  Bradford Books, 2 edition, 2018.
  * Swaminathan and Joachims [2015a] A. Swaminathan and T. Joachims.  Counterfactual risk minimization: Learning from logged bandit feedback.  In _Proc. ICML, Int. Conf. on Machine Learning_ , pages 814–823, 2015a.
  * Swaminathan and Joachims [2015b] A. Swaminathan and T. Joachims.  The self-normalized estimator for counterfactual learning.  In _Proc. NIPS, Advances in Neural Information Processing Systems_ , volume 28, 2015b.
  * Tavakol [2020] M. Tavakol.  Fair classification with counterfactual learning.  In _Proc. SIGIR, 43rd Int. ACM Conf. on Research and Development in Information Retrieval_ , pages 2073–2076, Virtual Event, China, 2020. ACM.  doi: 10.1145/3397271.3401291.
  * Wager and Athey [2018] S. Wager and S. Athey.  Estimation and inference of heterogeneous treatment effects using random forests.  _Journal of the American Statistical Association_ , 113(523):1228–1242, 2018.
  * Westreich et al. [2010] D. Westreich, J. Lessler, and M.J. Funk.  Propensity score estimation: Machine learning and classification methods as alternatives to logistic regression.  _Journal of Clinical Epidemiology_ , 63(8), 2010.
  * Wheeler [2020] G. Wheeler.  Bounded rationality.  In E.N. Zalta, editor, _The Stanford Encyclopedia of Philosophy_. Metaphysics Research Lab, Stanford University, Fall 2020 edition, 2020.
  * Zagorecki et al. [2013] A.T. Zagorecki, D.E.A. Johnson, and J. Ristvej.  Data mining and machine learning in the context of disaster and crisis management.  _Int. Journal of Emergency Management_ , 9(4), 2013.
  * Zehlike et al. [2021] M. Zehlike, K. Yang, and J. Stoyanovich.  Fairness in ranking: A survey.  _CoRR_ , abs/2103.14000v2, 2021.  URL http://arxiv.org/abs/2103.14000v2.
  * Zhao et al. [2012] Y. Zhao, D. Zeng, A.J. Rush, and M.R. Kosorok.  Estimating individualized treatment rules using outcome weighted learning.  _Journal of the American Statistical Association_ , 107(499):1106–1118, 2012.  doi: 10.1080/01621459.2012.695674.
  * Zhou et al. [2017] X. Zhou, N. Mayer-Hamblett, U. Khan, and M.R. Kosorok.  Residual weighted learning for estimating individualized treatment rules.  _Journal of the American Statistical Association_ , 112(517):169–187, 2017.
  * Zhou [2018] Z.H. Zhou.  A brief introduction to weakly supervised learning.  _National Science Review_ , 5:44–53, 2018.
  * Zliobaite [2017] I. Zliobaite.  Measuring discrimination in algorithmic decision making.  _Data Mining and Knowledge Discovery_ , 31:1060–1089, 2017\.  doi: 10.1007/s10618-017-0506-1.


# Learning to Skip for Language Modeling

Dewen Zeng  Nan Du  Tao Wang  Yuanzhong Xu  Tao Lei  Zhifeng Chen  Claire Cui

###### Abstract

Overparameterized large-scale language models have impressive generalization
performance of in-context few-shot learning. However, most language models
allocate the same amount of parameters or computation to each token,
disregarding the complexity or importance of the input data. We argue that in
language model pretraining, a variable amount of computation should be
assigned to different tokens, and this can be efficiently achieved via a
simple routing mechanism. Different from conventional early stopping
techniques where tokens can early exit at only early layers, we propose a more
general method that dynamically skips the execution of a layer (or module) for
any input token with a binary router. In our extensive evaluation across 24
NLP tasks, we demonstrate that the proposed method can significantly improve
the 1-shot performance compared to other competitive baselines only at mild
extra cost for inference.

Machine Learning, ICML

##  1 Introduction

Transformer-based (Vaswani et al., 2017) large scale language models trained
with general corpus have shown tremendously improvement of generalization in
particular with in-context few-shot learning in recent years (Shoeybi et al.,
2019; Brown et al., 2020; Rae et al., 2021; Chowdhery et al., 2022; Hoffmann
et al., 2022). Despite of the impressive capability of text generation,
training and serving these giant models are non-trivial even with the recent
progress of hardware and software (Jouppi et al., 2017; Lepikhin et al., 2021;
Patterson et al., 2021). One of the major challenges is that the processing of
each input requires to activate all the parameters of a model, which often
leads to trillions of floating point operations (FLOPs) per prediction. This
imposes a big burden on both model training and inference since we have no
control over the amount of computation that can be assigned to each input
example.

In contrast, it is commonly believed that human cognition (Stanovich & West,
2000; Levy, 2008) uses varying cognitive efforts to operate and learn
depending on the ‘hardness’ of the input. Specifically, one may only need
small efforts (lower computational cost) to process ‘easy’ examples, like the
commonly used stop-words, punctuation, patches in the background of an image,
etc., but allow additional efforts (more computational cost) for ‘hard’
examples, e.g., a rare abstract concept for reasoning, when they are truly
needed. Therefore, allocating the same computational power of a large model
uniformly for processing all samples tend to be wasteful and less efficient.
Such issue might be even more exacerbated when training large models using
real-world data corpus in that the redundancy of trivial examples will be more
pronounced as more and more data are used.

Conditional computation (Bengio et al., 2013, 2015) is the paradigm where only
a small subset of the model parameters are activated based on the input
representation, thereby reducing the amount of computation needed per example.
However, due to the discreteness of the decisions based on each input,
training neural networks with conditionally activated components end-to-end
differentiably and efficiently is still challenging.

In this paper, we develop a simple framework, referred to as the SkipLayer,
which allows an input to skip any layer that can be wrapped inside it
conditioned on the contextual representation. More specifically, SkipLayer-
based models can be trained end-to-end differentiably while at the same time
the discrete decisions during the forward pass can still be respected, which
enables us to precisely control the performance-compute tradeoff through
external constraint. Moreover, because the discrete decisions can be preserved
during the forward pass, we also develop an efficient implementation so that
the additional computation can be further saved in both pretraining and
inference for the given target budget. We then apply SkipLayer to the
Transformer architecture (Vaswani et al., 2017) to demonstrate the potential
efficacy of the method for decoder-only language model pretraining and
decoding. Finally, we extensively validate our method on a suite of well
established NLP benchmarks ranging from open-domain QA tasks, reading
comprehension, common sense reasoning, to natural language inference tasks.
SkipLayer-based models have shown strong 1-shot performance with controllable
computation tradeoff between model quality and decoding efficiency compared to
a variety of competitive baselines.

##  2 Method

In this section, we elaborate on our proposed SkipLayer framework, its
efficient implementation, and the application to Transformer-based language
models.

###  2.1 SkipLayer

|
---|---

(a)

(b)

Figure 1: (a) Overview of our SkipLayer framework. The router can choose to
activate or skip the embeded layer logic based on the input context. (b)
Straight-Through Gumbel-Softmax is used for the router. In the forward pass,
binary variables are sampled. During the backward pass, gradients can be
backpropagated to update the router.

Let $X^{o}_{\text{[layer]}}=F_{[\text{layer}]}(X|\mathbb{W})$ denote a
parameterized layer (or module) of a neural network with input $X$ and output
$X^{o}$ given an optional set of weights represented by $\mathbb{W}$. For
instance, a plain FeedForward layer (FFN) can be denoted as
$X^{o}_{\text{FFN}}=F_{\text{FFN}}(X|\\{W^{i},W^{o}\\})$ where
$W^{i}\in\mathbb{R}^{m\times h}$ and $W^{o}\in\mathbb{R}^{h\times m}$ are the
input and output weight, respectively.

A SkipLayer $F_{\text{SL}}$ is designed to wrap an existing layer such that

$\displaystyle X^{o}_{\text{SL}}$
$\displaystyle=F_{\text{SL}}\left(F_{[\text{layer}]}(X)|W_{\text{G}}\right)$
(1) $\displaystyle=F_{[\text{layer}]}(X)\odot
G(X|W_{G})+X\odot(1-G(X|W_{G})),$

where $G(X|W_{G})\in\\{0,1\\}$ is a router function with the learnable weight
$W_{G}$. Figure 1(a) shows the overall framework.

Given a batch $X\in\mathbb{R}^{B\times T\times d}$ of $B$ sequences, each of
length $T$ and the embedding dimension $d$, for each token input
$X[b,t]\in\mathbb{R}^{d},b\leq B,t\leq T$, we have that

$\displaystyle
X^{o}_{\text{SL}}[b,t]=\begin{cases}F_{[\text{layer}]}\left(X[b,t]\right),&\text{if
$G(X[b,t])=1$}.\\\ X[b,t],&\text{otherwise}.\end{cases}$ (2)

Therefore, as shown in Figure 1(a), any existing layer applied to the input in
a pointwise way (e.g., FFN) can be easily embedded inside a SkipLayer. Based
on the context, if the router decides to skip, the input will be connected
directly to the output, otherwise it will go through the embedded layer logic.

#### Router Function.

Central to the SkipLayer is the router function $G(X|W_{G})$ which is learned
to assign only a subset of inputs to the embedded layer for the best model
performance under a given budget. For a batch of input tokens
$X\in\mathbb{R}^{B\times T\times d}$, the router outputs a binary mask matrix

$\displaystyle M=G(X|W_{G})\in\\{0,1\\}^{B\times
T},W_{G}\in\mathbb{R}^{d\times 2}.$ (3)

There are several choices for designing $G(X|W_{G})$. One choice is the
Sigmoid function $G(X|W_{G})=\sigma(XW_{G})$ that independently normalizes
each value to be within the continuous range $(0,1)$ as the soft approximation
to the binary masking. Although this approximation is easy to differentiate,
it needs an additional threshold to produce the binary decision rule of
Equation 2.

The second design choice is the Top-K $(K=1)$ routing which is widely used in
the works of (Du et al., 2022; Lepikhin et al., 2021), giving
$G(X|W_{G})=\text{Top-1}\left(XW_{G}\right)=\mathop{\mathrm{argmax}}{XW_{G}}$.
In order to address the indifferentiability of the $\mathop{\mathrm{argmax}}$
operator, as discussed in (Lepikhin et al., 2021), for each input token
$X[b,t]$, we first normalize the dot-product scores by
$g=\text{Softmax}(X[b,t]W_{G})\in\mathbb{R}^{2}$, and let

$\displaystyle X^{o}_{\text{SL}}[b,t]=\begin{cases}g[1]\cdot
F_{[\text{layer}]}\left(X[b,t]\right),&\text{if
$\mathop{\mathrm{argmax}}g=1$}.\\\ g[0]\cdot
X[b,t],&\text{otherwise},\end{cases}$ (4)

such that the gradients can be backpropgated through the coefficients $g$.
Unfortunately, in our experiments, we find that the Top-1 formulation cannot
precisely control the sparsity of the model (which is crucial to the
efficiency) since $g$ is still a soft approximation of the binary decision
rule.

We thus formulate the router with the Straight-Through Gumbel-Softmax trick
shown in Figure 1(b). In the forward pass, sampled binary values are returned
for the gating $G(X[b,t])$ as in Equation 2. In the backward pass, the soft
probabilities are used as $g$ in Equation 4 for the gradients to be propagated
back to update the router weights. Because during the forward pass we are able
to directly calculate the percentage of the tokens that are not skipped based
on the binary masking of Equation 3, we can better control the density of the
model.

#### Router Capacity.

The binary mask of the router output in Equation 3 is the assignment of a
subset of tokens in a batch to the embedded layer inside a SkipLayer. For
simplicity, suppose each sequence in a batch of size $B$ has the same sequence
length $T$. Then, the ratio $r=\frac{\sum_{i,j}M[i,j]}{B\times T}$ is the
percentage (or probability) that a token is assigned to the layer, which is
also referred to as the capacity. Consider $P$ as a global budget of how many
input tokens can be assigned to a layer. Following (Du et al., 2022; Lepikhin
et al., 2021), we introduce an auxiliary loss term
$\ell_{\text{aux}}=\sum_{i}^{L}(r_{i}-P)^{2}$ where $r_{i}$ is the capacity of
layer $i\leq L$, so that each layer will respect the budget constraint. The
overall loss function of the model will be
$\mathcal{L}=\ell_{\text{nll}}+\lambda\cdot\ell_{\text{aux}}$ where
$\ell_{\text{nll}}$ is the negative log-likelihood of predicting the next
token on average. By optimizing $\mathcal{L}$, on the one hand, the layer
capacity will be pushed to be closer to the target probability $P$. On the
other hand, the $\ell_{\text{aux}}$ term will continuously improve the model’s
predictive accuracy. Since the $\ell_{\text{aux}}$ term will enforce only $P$
percent of tokens in a batch to go with the layer, in order to reduce the
first term $\ell_{\text{nll}}$, ‘hard’ examples that lead to large marginal
reduction on average will be prioritized while ‘easy’ examples that already
achieve low perplexity will be skipped in order to save FLOPs.

The router capacity enables the flexibility of controlling the performance-
computation trade-off. Specifically, we can increase the number of layers in
total while at the same time reduce target probability $P$ to keep the average
number of activated layers roughly the same. This effectively separates the
increase of model capacity from the computation cost per prediction, and makes
it possible to trade off the increased model capacity for better prediction.
During serving, we are then able to load the model that can best utilize the
accelerators’ memory, lead to the highest prediction quality, and only mildly
increase the computation cost while still meeting the latency requirement
simultaneously.

Figure 2: Illustration of efficient SkipLayer implementation. We focus on the
sparse computation of FFN, non-skip tokens are gathered based on indices
generated by the router and then fed into the FFN as groups. Gsize is a hyper-
parameter that controls the group size. The results are scattered to the final
output. Figure 3: Overview of our SkipLayer for Transformer-based models,
example of a single layer. LN is the layer normalization layer, query, key and
value refers to the computation of query, key and value projections in the
self-attention layer. Attn is the attention computation. Residual connections
in the self attention layer and FFN layer are ignored for simplicity.

###  2.2 Efficient Implementation

The major advantage of SkipLayer-based models is that the number of inputs
computed by each layer is different across the entire stack of layers and
continuously varies during training. At the same time, this dynamic
characteristic is also challenging for implementation on TPU where
computations of tensors with static shapes are often preferred. The basic
implementation is to first apply the given layer logic in Figure 1(a) to the
entire batch and then multiply the output batch with the mask given by
Equation 3 so that the skipped tokens will not be used in the layer. However,
this masking mechanism is computational expensive since we should not spend
the same computations on the skipped inputs as those on the non-skipped ones
especially when the skip ratio is high.

We thus develop an efficient SkipLayer implementation based on dynamic gather
and scatter. The idea is illustrated in Figure 2 where we focus on the sparse
computation of a FFN layer since it is a widely used component and often
computationally intensive. The overall algorithm includes three major steps.

  1. 1.

All inputs are marked as skip or non-skip based on results of the router in
Equation 3.

  2. 2.

All non-skipped inputs (in the blue rectangles) are gathered and evenly
partitioned into groups. Although each group will be fed into the FFN for
computation sequentially, all the elements in the same group will be gathered,
computed, and scattered in parallel.

  3. 3.

The compute results from the non-skipped inputs will be scattered back to the
final outputs of this FFN layer, while the skipped inputs will be directly
written into the final outputs without any computation.

The group size (the number of inputs in a group), denoted as Gsize, is a
hyper-parameter that controls how many tokens will be processed by the FFN in
parallel. Because the number of non-skipped inputs in a batch is dynamic and
unknown in advance, Gsize affects the training efficiency. When Gsize is too
large, e.g., there is only one single group, this group may include too many
skipped inputs, leading to sub-optimal performance. When Gsize is too small,
it will produce too many groups of small size, and the computation will be
close to being sequential. Thus, there will be little parallelism, and the
overheads maybe even larger than the basic masking implementation. In
practice, we often set $\text{Gsize}\propto P\cdot BT$ where $P$ is the target
probability, $B$ is the batch size, and $T$ is the sequence length.

Algorithm 1 Forward pass of SkipLayer

Data: A batch of tokens $X\in\mathbb{R}^{B\times T\times d}$, target
probability $P$.

Get the mask $M$ by Equation 3.

Get the key and value projection $K\leftarrow F_{\text{key}}(X)$, $V\leftarrow
F_{\text{val}}(X)$.

for  _$b\leq B$ , $t\leq T$_ do

2 if  _$M[b,t]=1$_ then

3 Get the query projection $q\leftarrow F_{\text{query}}(X[b,t])$
$x^{\prime}\leftarrow
F_{\text{Attn}}\left(F_{\text{LN}}(X[b,t])|K,q,V\right)\cdot M[b,t]+X[b,t]$
$X_{\text{SL}}[b,t]\leftarrow
F_{\text{FFN}}\left(F_{\text{LN}}(x^{\prime})\right)+x^{\prime}$

4 else

5 $X_{\text{SL}}[b,t]\leftarrow X[b,t])\cdot(1-M[b,t])$

6 end if

7

8 end for

$\ell_{\text{aux}}\leftarrow(\sum_{b,t}M[b,t]/(B\cdot T)-P)^{2}$ return
_$X_{\text{SL}},\ell_{\text{aux}}$_

###  2.3 SkipLayer for Transformer-based Models

In this section, we focus in particular on applying SkipLayer to Transformer-
based decoder-only language models in the setup of in-context learning. A
Transformer layer mainly includes the self-attention, layer normalization, and
FFN as the sub-layers, and can be represented as

$\displaystyle{X^{l}}^{\prime}$
$\displaystyle=F_{\text{Attn}}\left(F_{\text{LN}}(X^{l})\right)+X^{l},$
$\displaystyle X^{l+1}$
$\displaystyle=F_{\text{FFN}}\left(F_{\text{LN}}({X^{l}}^{\prime})\right)+{X^{l}}^{\prime}.$
(5)

SkipLayer can be applied to a single Transformer layer shown in Figure 3. We
propose to wrap the entire Transformer layer into a SkipLayer to preserve the
atomicity of the self-attention ($F_{\text{Attn}}$) $\to$ FFN
($F_{\text{FFN}}$) structure. However, the $F_{\text{Attn}}$ layer and the
$F_{\text{FFN}}$ layer have slightly different skipping implementations. The
$F_{\text{FFN}}$ layer is often the most computationally intensive component
of a Transformer model, but can be applied to a batch of tokens in a point-
wise manner. Therefore, each input token of a batch can activate the
$F_{\text{FFN}}$ layer independently with the probability $P$. Because
$F_{\text{FFN}}$ consumes most of the computation in a Transformer layer, we
can thus have big savings in FLOPs when the activation probability $P$ is
small.

The self-attention layer $F_{\text{Attn}}$ consumes much less computation
relative to the $F_{\text{FFN}}$ layer. However, $F_{\text{Attn}}$ cannot
apply to a batch of tokens in the pointwise way because tokens need to attend
to each other to compute their own attention output. If most tokens in a batch
are skipped when $P$ is small, the left non-skipped tokens will lose most of
the context of the respective sequences they belong to. We also empirically
observe lower predictive quality when this simple skipping mechanism is
applied. Therefore, we propose the following partial skipping mechanism. As
shown in Figure 3, when the input tokens are skipped, their key and value
projections are still preserved (Line 2 in Algorithm 1) since they are part of
the context and are needed for the rest non-skipped tokens to further attend
to. However, because the skipped tokens do not require attention calculations
by themselves, we can still omit their query projections. It is worth to
mention that Line 3-11 in Algorithm 1 can be computed in parallel using our
efficient implementation in Section 2.2 to increase training speed.

Algorithm 2 shows the greedy decoding logic of one SkipLayer-based Transformer
layer. The router makes the skipping decision by picking the most likely
outcome. The key and value projections will be computed and saved in the
decoding cache (Line 2). Only when the router activates the current layer, the
query projection will be computed for current token, and then the decoding
cache $K$ and $V$ which contain the key and value projections of previous
decoding steps will be used to compute the self attention. Otherwise, there
will be no further computations from this layer.

Algorithm 2 SkipLayer per decoding step

Data: The current state $x\in\mathbb{R}^{d}$, key and value cache $K$, $V$

$m=G(x|W_{G})=\mathop{\mathrm{argmax}}{x^{\top}W_{G}}$

$K\leftarrow F_{\text{key}}(x)$, $V\leftarrow F_{\text{val}}(x)$

if  _$m=1$_ then

2 Get the query projection $q\leftarrow F_{\text{query}}(x)$
$x^{\prime}\leftarrow F_{\text{Attn}}\left(F_{\text{LN}}(x)|K,q,V\right)+x$
$x_{\text{SL}}\leftarrow
F_{\text{FFN}}\left(F_{\text{LN}}(x^{\prime})\right)+x^{\prime}$

3else

4 $x_{\text{SL}}\leftarrow x$

5 end if

return: $x_{\text{SL}}$

##  3 Experiment Setup

We focus on training decoder-only language models. This section elaborates our
training setup, hyperparameters, baselines, benchmarks, and evaluation
protocol.

Table 1: Architectures and sizes of the models trained in our experiments. All trained model share the same learning hyperparameters. Model | $n_{\text{params}}$ | $P$ | $L$ | Eff-L
---|---|---|---|---
Standard (6L) | 408M | 0 | 6 | 6
SkipLayer (12, 50%) | 766M | 50% | 12
SkipLayer (24, 25%) | 1.47B | 25% | 24
SkipLayer (48, 12.5%) | 2.92B | 12.5% | 48
Standard (12L) | 766M | 0 | 12 | 12
SkipLayer (24, 50%) | 1.47B | 50% | 24
SkipLayer (48, 25%) | 2.92B | 25% | 48
SkipLayer (96, 12.5%) | 5.79B | 12.5% | 96
Standard (24L) | 1.47B | 0 | 24 | 24
SkipLayer (48, 50%) | 2.92B | 50% | 48
SkipLayer (96, 25%) | 5.79B | 25% | 96

#### Dataset.

The pretraining dataset has 1.6 trillion tokens that are representative of a
wide range of natural language use cases. An in-house classifier is trained to
classify between a collection of curated text and other web pages so that we
are able to estimate the content quality of a webpage. A high-quality filtered
subset of webpages are combined with books, Wikipedia pages, conversations,
forums, and news to create the final dataset which is the same as (Du et al.,
2022) for training.

|  |  |
---|---|---|---

(a) Avg. 1-shot (NLG)

(b) Avg. 1-shot (NLU)

(c) Decoding time per token.

(d) Training speed.

Figure 4: Average 1-shot performance of different SkipLayer-based models for
comparable effective FLOPs per token prediction over the NLG tasks (a) and NLU
tasks (b). (c) Comparisons of the decoding time per token between the
SkipLayer-based models (SL) and the respective baseline models (STD). (d)
Comparisons of the training speed among models of different density.

#### Model Training.

We have trained several variants of SkipLayer-based models and baselines shown
in Table 1. The model dimension of all the models is 1,536 and the hidden
dimension of the FFN has 8$\times$ of the model dimension. The hidden dim of
each attention head is 64. $n_{\text{params}}$ is the total number of
trainable model parameters, $L$ is the total number of Transformer layers, $P$
is the probability of activating a layer, and Eff-L is the effective number of
layers activated on average. The sequence length is set to 1,024 tokens during
training, and the batch size includes 256 sequences. We set the learning rate
to 0.1 for the first 10K training steps and then decay following an inverse
square root schedule. We use Adafactor optimizer with first-moment decay
$\beta_{1}=0$ and second-moment decay $\beta_{2}=0.99$. The dropout rate is
set to 0 as the processed tokens are token from an extremely large training
corpus. We use the SentencePiece (Kudo & Richardson, 2018) subword tokenizer
with a vocabulary of size of 32K. During training, we use float32 for model
weights and bfloat16 for activations. Aside from the general negative log-
likelihood loss, we add the auxiliary loss discussed in Section 2.1 to control
the skip ratio of our SkipLayer model, the auxiliary loss weight $\lambda$ is
set to 0.1. Finally, Gsize is set to 1,024 for all SkipLayer-based models.

#### Model Evaluation.

To directly evaluate the effectiveness of SkipLayer-based models, we mainly
follow the 1-shot learning protocol suggested by (Radford et al., 2018), which
is widely used for evaluating the generalization quality of pre-trained
language models. We evaluate each example in the development set of a
benchmark. For each benchmark, only one example will be randomly drawn from
that task’s training set as the only demonstration and context, which will be
then concatenated with the evaluation example with two newlines in between,
and then fed into the model. We use exactly the same prompting format as
(Radford et al., 2018) for each downstream benchmark.

#### Benchmarks.

We use 24 datasets including four natural language generative (NLG) tasks and
20 natural language understanding (NLU) tasks for evaluations. For NLG tasks,
we compare the decoded sequence of tokens by the models to the ground-truth
and report the Exact Match (EM) accuracy. These tasks are TriviaQA, NQS,
WebQS, and SQuADv2. Greedy decoding is used for each task. All NLU tasks are
formulated into the form of selecting one correct answer from multiple
candidate options. The prediction is based on the maximum log-likelihood of
each option given the context $\log{P(\text{option}|\text{context})}$
normalized by the token length of each option. These NLU tasks include ANLI
(R1, R2, R3), ARC (Easy, Challenge), BoolQ, CB, COPA, Hellaswag, Openbookqa,
PIQA, Race (middle, high), ReCord, RTE, Storycloze, WIC, Winograd, Winogrande
and WSC273. Finally, we use the average of the scores across all datasets to
report the overall 1-shot performance of models on both NLG and NLU tasks.

#### Baselines.

We consider the following baselines to study the effectiveness of SkipLayer-
based models.

  * •

Standard base model (STD). The standard Transformer dense models without any
skipping operations.

  * •

WideFFN. Because the FFN often consumes the major computation and has big
impact to the predictive performance (Kocsis et al., 2022), WideFFN is thus
designed to further double the hidden dimension of the FFN component. We then
apply SkipLayer only to the FFN component without changing the total number of
layers. As a consequence, when the skipping probability $P=50\%$, the compute
FLOPs per prediction does not change due to the 2x FFN size.

  * •

HighwayNet (Srivastava et al., 2015) is among the first few works that propose
to learn a gating function that could help to train very deep networks
efficiently. We also apply this idea to the FFN component of the Transformer
layer as one baseline.

  * •

Random Gating (Rand). Random gating method is the baseline where the learned
gating function of a SkipLayer model is replaced by a pure random function
without learning. This baseline is designed to evaluate the importance of the
learned gating function to a model’s predictive performance.

|  |  |
---|---|---|---

(a) Avg. 1-shot (NLG, 6L)

(b) Avg. 1-shot (NLG, 12L)

(c) Avg. 1-shot (NLU, 6L)

(d) Avg. 1-shot (NLU, 12L)

Figure 5: Average 1-shot NLG and NLU performance of different methods with 6
(a-b) and 12 (c-d) effective number of activated layers, respectively.

|
---|---

(a) Avg. 1-shot (NLG)

(b) Avg. 1-shot (NLU)

Figure 6: Learned gating has significantly better performance than the
respective methods using random gating.

##  4 Full Results

We have conducted comprehensive evaluation of the effectiveness of SkipLayer
and report the quantitative results compared to different baselines in this
section.

#### How does SkipLayer perform in 1-shot learning?

SkipLayer allows each input to selectively activate a particular layer
depending on the context. By keeping the average number of activated layers
constant while increasing the total number of layers of a language model, we
expect that the increased model capacity can improve the predictive quality of
few-shot learning. Thus, in Figure 4(a-b), we first report the average 1-shot
performance of different SkipLayer models. In Figure 4(a), the y-axis is the
average 1-shot performance of all the NLG tasks, and the x-axis (log-scale) is
the compute FLOPs per token during the forward-pass. The black line first
shows the performance of the standard baseline models of 6, 12, 24 and 48
layers, respectively. For each baseline model of a particular number of
layers, e.g., 6L, we report the performance of the respective SkipLayer models
of comparable compute FLOPs. For example, the three yellow dots represent the
SkipLayer models of 12 layers (12L) with 50% density, 24 layers (24L) with 25%
density, and 48 layers (48L) with 12.5% density, respectively. Compared to the
6L baseline model, all these three SkipLayer models (in yellow) have the same
average number of activated layers of six (Effective number of layers) denoted
as Eff06. Similarly, the orange line (Eff12) and red line (Eff24) denote the
SkipLayer models with 12 and 24 average number of activated layers,
accordingly.

In all the cases, it first shows that as the SkipLayer models become deeper
and sparser (keeping the same activated number of layers), it keeps improving
the few-shot learning performance at the modest cost of increased FLOPs per
token prediction. For instance, SkipLayer (24L, 25%) has 51.5% performance
gain at the cost of 18.2% increased compute FLOPs compared to the 6L baseline.
Similarly, SkipLayer (48L, 25%) and SkipLayer (96L, 25%) have 28% and 22% at
the cost of 19.2% and 20% compared to the 12L and 24L baselines, respectively.
Moreover, in Figure 4(a) we can also observe that even though SkipLayer (48L,
12.5%) has less compute FLOPs compared to the 12L baseline, it has achieved
pretty close 1-shot performance. Similarly, SkipLayer (96L, 12.5%) has even
better 1-shot performance compared to the 24L baseline by using less compute
FLOPs. This verifies that we are able to trade off model capacity for better
predictive quality. Likewise, Figure 4(b) shows similar patterns that
increased model capacity can lead to better predictive quality across the NLU
tasks while at the cost of modest increased FLOPs.

#### Does SkipLayer decode and train efficiently?

As shown in Figure 4(a-b), deeper and sparser SkipLayer models have consistent
performance improvement in few-shot learning. We are also interested in
studying if they are able to decode and train fast. Figure 4(c) compares the
decoding time per token of different models using a single TPU v3 chip. It
shows that SkipLayer (12, 50%) has nearly the same per-step decoding time as
the baseline 6L. SkipLayer (24, 50%) also has similar speed as the baseline
12L. SkipLayer (48, 50%) has 8% 1-shot performance gain at the cost of 6%
increase in the per-step decoding time compared to the baseline 24L. As the
models become deeper, the per-step decoding time also increases. However, we
may find some good trade-off between quality and speed. For example, SkipLayer
(96, 25%) has achieved 20% decoding efficiency with a tiny quality loss of
only 0.5% compared to the baseline 48L. Figure 4(d) further shows the training
speed on a single TPU v4 of different SkipLayer models. In general, the
training speed decreases as more FLOPs are used per prediction. Compared to
the respective full dense baselines, SkipLayer (24, 50%) has 18% speed gains,
and SkipLayer (48, 12.5%) has 3x gains.

#### How does SkipLayer compare to the baselines?

Figure 5(a-b) first compare the average 1-shot performance of different
methods of effective 6 layers and 12 layers across all the NLG tasks. In
Figure 5(a) when all models are small, HighwayNet has similar performance as
the standard baseline 6L model. They both outperform the Random gating method
but marginally underperform the WideFFN method. When all the models become
deeper in Figure 5(b), the standard baseline 12L scales better than WideFFN,
HighwayNet and the Random gating method. However, in both cases, SkipLayer
based models have achieved the best performance compared to all the other
baselines. HighwayNet performs the worst among all the models we trained, one
possible reason could be Highway network is designed for FFN only
architectures, it works like a weighted residual branch added to the wrapped
layer, which helps the training stability of very deep FFN only networks.
However, residual branch is quite normal in today’s Transformer models, self-
attention and FFN both have residuals. In our implementation, we replace the
original residual with the highway network residual which leads to performance
degradation. Figure 5(c-d) further compare the average 1-shot performance of
different methods across all the NLU tasks. With a similar trend observed
before, WideFFN has a scaling performance closer to the standard baseline
models of 12 layers, both of which significantly outperform the HighwayNet and
the Random gating method, and SkipLayer based models lead to the best
performance overall when the models scale up.

#### Does the learned gating matter?

The gating function inside a SkipLayer enables each input example to activate
a subset of layers of the model based on the context. We have already observed
that this flexibility of switching model parameters improve the predictive
accuracy of the model when it becomes deeper and sparser. We wonder how much
gain the learned routing function can contribute to the predictive performance
of the SkipLayer-based models. We approach this question by comparing the
SkipLayer-based models with the respective Random gating baselines where the
gating is not learned by varying the density of the model side-by-side in
Figure 6. Because each token activates a layer randomly and independently,
Random gating baselines have the similar compute FLOPs per token prediction as
the SkipLayer-based models given the same model density. However, as shown in
Figure 6, for the same model density, the learned gating of SkipLayer-based
models performs significantly better than the respective methods using random
gating. Even though Random gating baselines also have the flexibility of
switching model parameters per token prediction, Figure 6 verifies that the
learned gating functions are much more effective to improve the prediction
accuracy.

#### How does the skipping behave during decoding?

Figure 7: Bubble chart of the skipping behaviour of tokens during greedy
decoding on TriviaQA dataset using our SkipLayer (12L, 50$\%$) model. Each dot
represents a token, larger dot size means more layers are skipped. Black/Red
texts show some tokens that skip the most/least layers.

To study the skipping behaviour of SkipLayer model and understand what kind of
tokens skip more layers than others during greedy decoding, we collected the
skipping statistics of the decoding results of 500 sampled questions in
TriviaQA using our SkipLayer (12L, 50$\%$) model. In Figure 7, we plotted the
bubble chart of 300 frequently used tokens in the decoding results according
to their averaged number of skipped layers (One token may have different
skipping patterns under different contexts). Larger dots in the figure
represents more layers are skipped. We can observe that tokens that skip the
most are mainly functional words like “and”, “to”, “ed” or “ing”. These tokens
can be easily inferred from the previous contexts and thus do not need much
computation to decode. While tokens that skip less are usually independent
words like “No” or “Paris”. This indeed shows that our SkipLayer model can
successfully identify such tokens and assign proper computation accordingly.

##  5 Related Work

#### Conditional Computation

(Bengio et al., 2013, 2015) is a paradigm where only a subset of model
parameters can be activated per input example. Early-exit is one kind of
implementation for conditional computation where external classifiers equipped
with confidence-based thresholding are used to exit early without going
through the whole stack of layers (Wang et al., 2017; Xin et al., 2020;
Schwartz et al., 2020; Liu et al., 2020; Dabre et al., 2020; Elbayad et al.,
2020; Schuster et al., 2022). Unlike these approaches where the computation is
activated for the bottom layers, SkipLayer-based models technically allows
each input to explore $2^{L}$ different compute paths for a model with $L$
stacked layers.

An alternative technique is to enable the model to ‘learn’ how to activate its
different sub-layers. Due to the discreteness of the activating decisions,
soft-approximations and RL-based implementations has been explored in the
vision (Srivastava et al., 2015; Wang et al., 2018) and NLP (Bapna et al.,
2020) community. Our approach is closer to the second learning approach but
differ in that SkipLayer does not require the soft-approximation during the
forward pass, giving computational savings not just in inference but also
during training. This difference is crucial since pretraining language models
is often time consuming and costly.

Additionally, concurrent work such as CODA (Lei et al., 2023) and CoLT5
(Ainslie et al., 2023) have applied similar token selection method to activate
Transformer layers. However, these work only apply conditional activation in
the encoder layers of an encoder-decoder model such as T5.

#### Mixture of Experts

have recently been proposed to improve model efficiency (Shazeer et al., 2017;
Gross et al., 2017; Lepikhin et al., 2021; Fedus et al., 2021; Roller et al.,
2021; Du et al., 2022; Artetxe et al., 2021; Lewis et al., 2021; Zhou et al.,
2022; Rajbhandari et al., 2022) by sparsely activating a subset of experts in
a MoE layer.

Our approach is orthogonal to MoE models in that an MoE layer can be easily
wrapped by the SkipLayer for additional efficiency. Moreover, SkipLayer can
apply conditional computation to both the self-attention and the Feed-Forward
(FFN) component of a Transformer layer, whereas MoE models mainly focus on
conditionally activating the FFN component in a MoE layer.

#### Structural Dropout

(Tompson et al., 2015; Ghiasi et al., 2018; Dai et al., 2019; Fan et al.,
2019; Zeng et al., 2021) randomly drops a group of weights, e.g., a layer (Fan
et al., 2019), during training to achieve better generalization and robustness
for pruning during inference. However, the amount of computation during
inference is still uniform per example. In contrast, SkipLayer-based models
learn the skipping patterns from the data which shows better performance than
the random skipping baseline, and potentially assign non-uniform amount of
computation to each example during inference.

##  6 Conclusions

We propose a new general method named SkipLayer for dynamically skipping the
execution of arbitrary layers based on the input context using a simple
routing algorithm. This method enables heterogeneous computation for tokens at
different complexity or importance so that more computation resources can be
used for improving the predictive quality of harder tokens. Our model
demonstrates significant 1-shot performance improvement across 24 NLP tasks
compared to other competitive baselines with only a small extra cost for
inference.

## References

  * Ainslie et al. (2023) Ainslie, J., Lei, T., de Jong, M., Ontañón, S., Brahma, S., Zemlyanskiy, Y., Uthus, D., Guo, M., Lee-Thorp, J., Tay, Y., et al.  Colt5: Faster long-range transformers with conditional computation.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , 2023.
  * Artetxe et al. (2021) Artetxe, M., Bhosale, S., Goyal, N., Mihaylov, T., Ott, M., Shleifer, S., Lin, X. V., Du, J., Iyer, S., Pasunuru, R., Anantharaman, G., Li, X., Chen, S., Akin, H., Baines, M., Martin, L., Zhou, X., Koura, P. S., O’Horo, B., Wang, J., Zettlemoyer, L., Diab, M., Kozareva, Z., and Stoyanov, V.  Efficient large scale language modeling with mixtures of experts.  2021\.  URL https://arxiv.org/abs/2112.10684.
  * Bapna et al. (2020) Bapna, A., Arivazhagan, N., and Firat, O.  Controlling computation versus quality for neural sequence models, 2020\.
  * Bengio et al. (2015) Bengio, E., Bacon, P., Pineau, J., and Precup, D.  Conditional computation in neural networks for faster models.  _CoRR_ , abs/1511.06297, 2015.  URL http://arxiv.org/abs/1511.06297.
  * Bengio et al. (2013) Bengio, Y., Léonard, N., and Courville, A. C.  Estimating or propagating gradients through stochastic neurons for conditional computation.  _CoRR_ , abs/1308.3432, 2013.  URL http://arxiv.org/abs/1308.3432.
  * Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D.  Language models are few-shot learners.  In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. F., and Lin, H. (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 1877–1901. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf.
  * Chowdhery et al. (2022) Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., Schuh, P., Shi, K., Tsvyashchenko, S., Maynez, J., Rao, A., Barnes, P., Tay, Y., Shazeer, N., Prabhakaran, V., Reif, E., Du, N., Hutchinson, B., Pope, R., Bradbury, J., Austin, J., Isard, M., Gur-Ari, G., Yin, P., Duke, T., Levskaya, A., Ghemawat, S., Dev, S., Michalewski, H., Garcia, X., Misra, V., Robinson, K., Fedus, L., Zhou, D., Ippolito, D., Luan, D., Lim, H., Zoph, B., Spiridonov, A., Sepassi, R., Dohan, D., Agrawal, S., Omernick, M., Dai, A. M., Pillai, T. S., Pellat, M., Lewkowycz, A., Moreira, E., Child, R., Polozov, O., Lee, K., Zhou, Z., Wang, X., Saeta, B., Diaz, M., Firat, O., Catasta, M., Wei, J., Meier-Hellstern, K., Eck, D., Dean, J., Petrov, S., and Fiedel, N.  Palm: Scaling language modeling with pathways.  2022\.  doi: 10.48550/ARXIV.2204.02311.  URL https://arxiv.org/abs/2204.02311.
  * Dabre et al. (2020) Dabre, R., Rubino, R., and Fujita, A.  Balancing cost and benefit with tied-multi transformers.  In Birch, A., Finch, A. M., Hayashi, H., Heafield, K., Junczys-Dowmunt, M., Konstas, I., Li, X., Neubig, G., and Oda, Y. (eds.), _Proceedings of the Fourth Workshop on Neural Generation and Translation, NGT@ACL 2020, Online, July 5-10, 2020_ , pp. 24–34. Association for Computational Linguistics, 2020.  doi: 10.18653/v1/2020.ngt-1.3.  URL https://doi.org/10.18653/v1/2020.ngt-1.3.
  * Dai et al. (2019) Dai, Z., Chen, M., Gu, X., Zhu, S., and Tan, P.  Batch dropblock network for person re-identification and beyond.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pp. 3691–3701, 2019.
  * Du et al. (2022) Du, N., Huang, Y., Dai, A. M., Tong, S., Lepikhin, D., Xu, Y., Krikun, M., Zhou, Y., Yu, A. W., Firat, O., Zoph, B., Fedus, L., Bosma, M. P., Zhou, Z., Wang, T., Wang, E., Webster, K., Pellat, M., Robinson, K., Meier-Hellstern, K., Duke, T., Dixon, L., Zhang, K., Le, Q., Wu, Y., Chen, Z., and Cui, C.  GLaM: Efficient scaling of language models with mixture-of-experts.  In _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 5547–5569. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/du22c.html.
  * Elbayad et al. (2020) Elbayad, M., Gu, J., Grave, E., and Auli, M.  Depth-adaptive transformer.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net, 2020.  URL https://openreview.net/forum?id=SJg7KhVKPH.
  * Fan et al. (2019) Fan, A., Grave, E., and Joulin, A.  Reducing transformer depth on demand with structured dropout.  _arXiv preprint arXiv:1909.11556_ , 2019.
  * Fedus et al. (2021) Fedus, W., Zoph, B., and Shazeer, N.  Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.  _CoRR_ , abs/2101.03961, 2021.  URL https://arxiv.org/abs/2101.03961.
  * Ghiasi et al. (2018) Ghiasi, G., Lin, T.-Y., and Le, Q. V.  Dropblock: A regularization method for convolutional networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Gross et al. (2017) Gross, S., Ranzato, M., and Szlam, A.  Hard mixtures of experts for large scale weakly supervised vision.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 6865–6873, 2017.
  * Hoffmann et al. (2022) Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Vinyals, O., Rae, J. W., and Sifre, L.  An empirical analysis of compute-optimal large language model training.  In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), _Advances in Neural Information Processing Systems_ , 2022.  URL https://openreview.net/forum?id=iBBcRUlOAPR.
  * Jouppi et al. (2017) Jouppi, N. P., Young, C., Patil, N., Patterson, D., Agrawal, G., Bajwa, R., Bates, S., Bhatia, S., Boden, N., Borchers, A., et al.  In-datacenter performance analysis of a tensor processing unit.  In _Proceedings of the 44th Annual International Symposium on Computer Architecture_ , pp. 1–12, 2017.
  * Kocsis et al. (2022) Kocsis, P., Súkeník, P., Brasó, G., Nießner, M., Leal-Taixé, L., and Elezi, I.  The unreasonable effectiveness of fully-connected layers for low-data regimes.  In _Proc. NeurIPS_ , 2022.
  * Kudo & Richardson (2018) Kudo, T. and Richardson, J.  Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing.  In _EMNLP_ , 2018.
  * Lei et al. (2023) Lei, T., Bai, J., Brahma, S., Ainslie, J., Lee, K., Zhou, Y., Du, N., Zhao, V. Y., Wu, Y., Li, B., Zhang, Y., and Chang, M.-W.  Conditional adapters: Parameter-efficient transfer learning with fast inference.  In _Advances in Neural Information Processing Systems_ , 2023.
  * Lepikhin et al. (2021) Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z.  GShard: Scaling giant models with conditional computation and automatic sharding.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=qrwe7XHTmYb.
  * Levy (2008) Levy, R.  Expectation-based syntactic comprehension.  _Cognition_ , 106(3):1126–1177, 2008.  doi: 10.1016/j.cognition.2007.05.006.
  * Lewis et al. (2021) Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettlemoyer, L.  Base layers: Simplifying training of large, sparse models.  In _International Conference on Machine Learning_ , 2021.
  * Liu et al. (2020) Liu, W., Zhou, P., Wang, Z., Zhao, Z., Deng, H., and Ju, Q.  FastBERT: a self-distilling BERT with adaptive inference time.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pp. 6035–6044, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.537.  URL https://aclanthology.org/2020.acl-main.537.
  * Patterson et al. (2021) Patterson, D., Gonzalez, J., Le, Q., Liang, C., Munguia, L.-M., Rothchild, D., So, D., Texier, M., and Dean, J.  Carbon emissions and large neural network training.  _arXiv preprint arXiv:2104.10350_ , 2021.
  * Radford et al. (2018) Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I.  Language models are unsupervised multitask learners.  2018\.  URL https://d4mucfpksywv.cloudfront.net/better-language-models/language-models.pdf.
  * Rae et al. (2021) Rae, J. W., Borgeaud, S., Cai, T., Millican, K., Hoffmann, J., Song, H. F., Aslanides, J., Henderson, S., Ring, R., Young, S., Rutherford, E., Hennigan, T., Menick, J., Cassirer, A., Powell, R., van den Driessche, G., Hendricks, L. A., Rauh, M., Huang, P., Glaese, A., Welbl, J., Dathathri, S., Huang, S., Uesato, J., Mellor, J., Higgins, I., Creswell, A., McAleese, N., Wu, A., Elsen, E., Jayakumar, S. M., Buchatskaya, E., Budden, D., Sutherland, E., Simonyan, K., Paganini, M., Sifre, L., Martens, L., Li, X. L., Kuncoro, A., Nematzadeh, A., Gribovskaya, E., Donato, D., Lazaridou, A., Mensch, A., Lespiau, J., Tsimpoukelli, M., Grigorev, N., Fritz, D., Sottiaux, T., Pajarskas, M., Pohlen, T., Gong, Z., Toyama, D., de Masson d’Autume, C., Li, Y., Terzi, T., Mikulik, V., Babuschkin, I., Clark, A., de Las Casas, D., Guy, A., Jones, C., Bradbury, J., Johnson, M., Hechtman, B. A., Weidinger, L., Gabriel, I., Isaac, W. S., Lockhart, E., Osindero, S., Rimell, L., Dyer, C., Vinyals, O., Ayoub, K., Stanway, J., Bennett, L., Hassabis, D., Kavukcuoglu, K., and Irving, G.  Scaling language models: Methods, analysis & insights from training gopher.  _CoRR_ , abs/2112.11446, 2021.
  * Rajbhandari et al. (2022) Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. Y., Awan, A. A., Rasley, J., and He, Y.  DeepSpeed-MoE: Advancing mixture-of-experts inference and training to power next-generation AI scale.  In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 18332–18346. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/rajbhandari22a.html.
  * Roller et al. (2021) Roller, S., Sukhbaatar, S., szlam, a., and Weston, J.  Hash layers for large sparse models.  In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), _Advances in Neural Information Processing Systems_ , volume 34, pp. 17555–17566. Curran Associates, Inc., 2021.  URL https://proceedings.neurips.cc/paper/2021/file/92bf5e6240737e0326ea59846a83e076-Paper.pdf.
  * Schuster et al. (2022) Schuster, T., Fisch, A., Gupta, J. P., Dehghani, M., Bahri, D., Tran, V. Q., Tay, Y., and Metzler, D.  Confident adaptive language modeling.  _CoRR_ , abs/2207.07061, 2022.  doi: 10.48550/arXiv.2207.07061.  URL https://doi.org/10.48550/arXiv.2207.07061.
  * Schwartz et al. (2020) Schwartz, R., Stanovsky, G., Swayamdipta, S., Dodge, J., and Smith, N. A.  The right tool for the job: Matching model and instance complexities.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pp. 6640–6651, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.593.  URL https://aclanthology.org/2020.acl-main.593.
  * Shazeer et al. (2017) Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. V., Hinton, G. E., and Dean, J.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  In _5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings_. OpenReview.net, 2017.  URL https://openreview.net/forum?id=B1ckMDqlg.
  * Shoeybi et al. (2019) Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B.  Megatron-lm: Training multi-billion parameter language models using gpu model parallelism.  _arXiv preprint arXiv:1909.08053_ , 2019.
  * Srivastava et al. (2015) Srivastava, R. K., Greff, K., and Schmidhuber, J.  Highway networks.  _arXiv preprint arXiv:1505.00387_ , 2015.
  * Stanovich & West (2000) Stanovich, K. E. and West, R. F.  Individual differences in reasoning: Implications for the rationality debate?  _Behavioral and Brain Sciences_ , 23(5):645–665, 2000.  doi: 10.1017/S0140525X00003435.
  * Tompson et al. (2015) Tompson, J., Goroshin, R., Jain, A., LeCun, Y., and Bregler, C.  Efficient object localization using convolutional networks.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 648–656, 2015.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I.  Attention is all you need.  In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
  * Wang et al. (2017) Wang, X., Luo, Y., Crankshaw, D., Tumanov, A., and Gonzalez, J. E.  IDK cascades: Fast deep learning by learning not to overthink.  _CoRR_ , abs/1706.00885, 2017.  URL http://arxiv.org/abs/1706.00885.
  * Wang et al. (2018) Wang, X., Yu, F., Dou, Z., Darrell, T., and Gonzalez, J. E.  Skipnet: Learning dynamic routing in convolutional networks.  In Ferrari, V., Hebert, M., Sminchisescu, C., and Weiss, Y. (eds.), _Computer Vision - ECCV 2018 - 15th European Conference, Munich, Germany, September 8-14, 2018, Proceedings, Part XIII_ , volume 11217 of _Lecture Notes in Computer Science_ , pp. 420–436. Springer, 2018.  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) Xin, J., Tang, R., Lee, J., Yu, Y., and Lin, J.  Deebert: Dynamic early exiting for accelerating bert inference.  _arXiv preprint arXiv:2004.12993_ , 2020.
  * Zeng et al. (2021) Zeng, Y., Dai, T., Chen, B., Xia, S.-T., and Lu, J.  Correlation-based structural dropout for convolutional neural networks.  _Pattern Recognition_ , 120:108117, 2021.
  * Zhou et al. (2022) Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A., Chen, Z., Le, Q., and Laudon, J.  Mixture-of-experts with expert choice routing.  _arXiv preprint arXiv:2202.09368_ , 2022.