Copyright for this paper by its authors. Use permitted under Creative Commons
License Attribution 4.0 International (CC BY 4.0).

SemREC’21: Semantic Reasoning Evaluation Challenge, ISWC’21, Oct 24 – 28,
Albany, NY

[orcid=0000-0002-4354-9138<EMAIL_ADDRESS>url=http://www.uni-mannheim.de/dws/people/researchers/phd-students/nicolas-
heist/ ] [orcid=0000-0002-4354-9138<EMAIL_ADDRESS>url=http://www.heikopaulheim.com/ ]

# The CaLiGraph Ontology as a Challenge for OWL Reasoners

Nicolas Heist  Heiko Paulheim  Data and Web Science Group, University of
Mannheim, Germany

(2021)

###### Abstract

CaLiGraph is a large-scale cross-domain knowledge graph generated from
Wikipedia by exploiting the category system, list pages, and other list
structures in Wikipedia, containing more than 15 million typed entities and
around 10 million relation assertions. Other than knowledge graphs such as
DBpedia and YAGO, whose ontologies are comparably simplistic, CaLiGraph also
has a rich ontology, comprising more than 200,000 class restrictions. Those
two properties – a large A-box and a rich ontology – make it an interesting
challenge for benchmarking reasoners. In this paper, we show that a reasoning
task which is particularly relevant for CaLiGraph, i.e., the materialization
of owl:hasValue constraints into assertions between individuals and between
individuals and literals, is insufficiently supported by available reasoning
systems. We provide differently sized benchmark subsets of CaLiGraph, which
can be used for performance analysis of reasoning systems.

###### keywords:

Knowledge Graph Reasoning Scalability owl:hasValue Restrictions
Materialization

##  1 Introduction

In the recent decade, open cross-domain knowledge graphs have been recognized
as an interesting ingredient to build intelligent systems. This gave rise to
the development of different such open knowledge graphs, such as DBpedia [1],
YAGO [2], and Wikidata [3]. While those knowledge graphs come with large-scale
A-boxes comprising millions of entities, their T-boxes are usually not very
expressive, defining mainly a class hierarchy and relations with domains and
ranges, but not using complex class constructors [4].

CaLiGraph is a comparatively new knowledge graph, which is constructed from
categories, list pages, and other lists in Wikipedia. It uses DBpedia as a
training set to derive interpretations of lists and categories [5]. This
allows to extract more entities from list pages [6] and other listings [7],
making the A-box comprise more than twice as many entities as DBpedia. At the
same time, the definitions of the derived classes are also maintained in the
T-box, making it more expressive than the ontologies of other knowledge
graphs.

With those characteristics, CaLiGraph creates some interesting tasks to
reasoners, since it requires the materialization of individual and literal
assertions, and, at the same time, poses high scalability requirements when
asking a reasoner to materialize 10 million A-box assertions. Since those
requirements are hardly met by currently available reasoners, the
materialization of the CaLiGraph A-box is carried out by custom code, not by
running a reasoner, albeit being a standard materialization task which an OWL2
EL reasoner should be capable of performing.

In this paper, we report on experiments with three reasoning systems on the
materialization of CaLiGraph. Furthermore, we introduce differently sized
subsets of CaLiGraph, which allow for performing scalability experiments.

The rest of this paper is structured as follows. Section 2 gives a short
introduction into CaLiGraph. In section 3, we describe a set of experiments we
conducted to investigate the fitness of existing reasoners for processing
CaLiGraph. We close with a short summary and an outlook on future work.

Figure 1: Example of a Category in Wikipedia

##  2 The CaLiGraph Ontology

As described above, CaLiGraph processes categories and list pages in
Wikipedia, and derives classes and instances from those. For example, when
processing the Wikipedia category Swedish Rock
Musicians,111https://en.wikipedia.org/wiki/Category:Swedish_rock_musicians
shown in Fig. 1, it will create a class clgo:Swedish_rock_musician,222The
following name space conventions are used throughout this paper:
clgo=http://caligraph.org/ontology/
clgr=http://caligraph.org/resource/ and heuristically assign the following
subclass axioms to it:

clgo:Swedish_rock_musician a owl:Class ; rdfs:subClassOf clgo:Rock_musician ,
clgo:Swedish_musician , [ a owl:Restriction ; owl:onProperty clgo:birthPlace ;
owl:hasValue clgr:Sweden ] , [ a owl:Restriction ; owl:onProperty clgo:genre ;
owl:hasValue clgr:Rock_music ] , [ a owl:Restriction ; owl:onProperty
clgo:occupation ; owl:hasValue clgr:Musician ] . clgo:Swedish_rock_guitarist
rdfs:subClassOf clgo:Swedish_rock_musician . clgo:Swedish_rock_singer
rdfs:subClassOf clgo:Swedish_rock_musician .

The subclass axioms are derived from the subcategory relations in Wikipedia,
the restrictions are created using a mix of statistic and textual heuristics
[5, 6].

As shown in this example, restrictions using owl:hasValue are quite frequently
used in CaLiGraph. Overall, the current release has more than 200,000 of such
restrictions, most of them using a resource as their object, but some also
using literals (e.g., restrictions on birth years).

When populating the A-box, CaLiGraph creates entities for Wikipedia pages, as
well as entities found on list pages. The only axioms created for those
entities are the class memberships for the classes derived from the
corresponding categories and list pages; the remaining axioms are materialized
according to the class definitions. In the example above, given that the
entity Dennis Lyxzén is found in the category Swedish rock singers (which is a
subcategory of the above mentioned Swedish rock musicians, CaLiGraph will
produce the following axiom:

    clgr:Dennis_Lyxzn a clgo:Swedish_rock_singer .

Given the T-box above, the derived axioms according to the definition above
should include:

    clgr:Dennis_Lyxzn a clgo:Swedish_rock_musician .
    clgr:Dennis_Lyxzn a clgo:Rock_musician .
    clgr:Dennis_Lyxzn a clgo:Swedish_musician .
    clgr:Dennis_Lyxzn clgo:birthPlace clgr:Sweden .
    clgr:Dennis_Lyxzn clgo:genre clgr:Rock_music .
    clgr:Dennis_Lyxzn clgo:occupation clgr:Musician .

The first three axioms are obtained by materializing the subclass relations,
the latter three by materializing the restrictions based on owl:hasValue. Note
that the example only depicts a subset of the information on the instance
clgr:Dennis_Lyxzén in CaLiGraph, which also contains information on his
birthplace and birth year.333http://caligraph.org/resource/Dennis_Lyxzén

While CaLiGraph uses a custom implementation to materialize those A-box
axioms, this is a task which should, in principle, be possible to carry out
for an OWL DL reasoner. Moreover, since restrictions with owl:hasValue are
possible in OWL2 EL and OWL2 RL [8], we would expect that reasoners supporting
those fragments should be capable of carrying out that materialization as well
[9].

The characteristics of the most recent version 2.1.0 of CaLiGraph are shown in
table 1. The table shows that a reasoner, provided with the 14 million
individuals and their 50M direct type assertions, should in principle be able
to derive 139M transitive type relations, as well as 10M relation and literal
assertions.

No. of classes | 1,061,598
---|---
No. of relations | 338
No. of subclass assertions | 1,711,270
No. of restrictions |
…with individuals | 214,248
…with literals | 12,925
No. of classes with direct restrictions | 223,552
No. of classes with direct or inherited restrictions | 488,364
No. of instances | 14,452,393
No. of direct type assertions | 50,169,052
No. of transitive type assertions | 138,713,499
No. of relation assertions |
…with individuals | 9,637,354
…with literals | 969,022
Table 1: Statistics of the CaLiGraph ontology

This kind of reasoning – at least at that scale – is rather underrepresented
in current reasoning benchmarks, and often not evaluated at all. For example,
in the ORE2015 benchmark, the task of materialization is defined as _the
computation of all instances for all named classes in the ontology_ [10]. In
other words: in that benchmark, the task is to infer all class assertions for
an instance given all its literal and individual assertions, while for
CaLiGraph, we require the reverse direction.

Among the 1,920 test cases in the ORE2015 benchmarks [10], there are only 311
which use owl:hasValue constraints with an individual (max. 6,455 in a single
file), and 279 which use owl:hasValue with a literal (max. 1,104 in a single
file). The more recent OWL2Bench dataset [11], proposed in 2020, contains
exactly one owl:hasValue restriction, mainly for checking whether or not a
reasoner supports that construct, but does not allow for evaluating that
support at scale. In contrast, the CaLiGraph ontology makes much heavier use
of those constructs, as shown above, and therefore imposes more challenging
scalability requirements with respect to that reasoning task.

##  3 Experiments

To validate how far reasoning systems actually meet those challenges, we
conducted a series of experiments, both with CaLiGraph as a whole and with
differently-sized subsets of CaLiGraph.

###  3.1 Reasoning Systems

In our experiments, we chose three commonly known reasoners: ELK [12], HermiT
[13], and Pellet [14]. All of them claim to be able to process OWL2 EL
ontologies. Specifically, we used the latest non-commercial versions, i.e.:

  * •

ELK version 0.4.3

  * •

Pellet version 2.3.6

  * •

HermiT version 1.4.5.456

The reasoning systems were called through OWL
API444https://github.com/owlcs/owlapi using the respective connectors to the
reasoners. As the reasoners rely on different versions of OWL API, we provide
an individual dependency configuration for each of them in our implementation.
Note that since OWL API does not support restrictions with explicit URIs and
labels, we transformed them to blank nodes (as shown in Fig. 2) before
importing them with OWL API.

###  3.2 Sandbox Example

As a first sanity check, we extracted a simple sandbox example from CaLiGraph,
which we coin clg_10 in the following. The sandbox example contains two
classes with one owl:hasValue restriction each, one using an individual, one
using a literal. For each of the two classes, there is an individual, so that
an OWL2EL reasoner should be capable of inferring one individual and one
literal assertion. The sandbox example is shown in Fig. 2.

    @prefix xsd: <http://www.w3.org/2001/XMLSchema#> .
    @prefix rdfs: <http://www.w3.org/2000/01/rdf-schema#> .
    @prefix owl: <http://www.w3.org/2002/07/owl#> .
    @prefix clgo: <http://caligraph.org/ontology/> .
    @prefix clgr: <http://caligraph.org/resource/> .

    clgo:Agent a owl:Class;
      rdfs:subClassOf owl:Thing .

    clgo:Organization a owl:Class;
      rdfs:subClassOf clgo:Agent .

    clgo:Person a owl:Class;
      rdfs:subClassOf clgo:Agent;
      owl:disjointWith clgo:Organization .

    clgo:International_organization a owl:Class;
      rdfs:subClassOf clgo:Organization .

    clgo:Organization_based_in_Asia a owl:Class;
      rdfs:subClassOf clgo:Organization .

    clgo:Organization_based_in_China a owl:Class;
      rdfs:subClassOf clgo:Organization_based_in_Asia .

    clgo:headquarter a owl:ObjectProperty .

    clgr:China a owl:NamedIndividual .

    _:b1 a owl:Restriction;
      owl:onProperty clgo:headquarter;
      owl:hasValue clgr:China .

    clgo:International_organization_based_in_China a owl:Class;
      rdfs:subClassOf clgo:Organization_based_in_China, clgo:International_organization, _:b1 .

    clgr:International_Center_on_Small_Hydro_Power a owl:NamedIndividual,
                                                     clgo:International_organization_based_in_China .

    clgo:Organization_disestablished_in_1939 a owl:Class;
      rdfs:subClassOf clgo:Organization .

    clgo:activeYearsEndYear a owl:DatatypeProperty .

    _:b2 a owl:Restriction;
      owl:onProperty clgo:activeYearsEndYear;
      owl:hasValue "1939"^^xsd:integer .

    clgo:Military_unit_or_formation_disestablished_in_1939 a owl:Class;
      rdfs:subClassOf clgo:Organization_disestablished_in_1939, _:b2 .

    clgr:46th_Mixed_Brigade a owl:NamedIndividual, clgo:Military_unit_or_formation_disestablished_in_1939 .

Figure 2: Sandbox example for testing the reasoners’ capabilities of
processing restrictions (with individuals and literals) and disjointness
axioms

From our three reasoners, only Pellet is capable of inferring both assertions.
HermiT can infer the resource-valued assertion, but not the literal-valued
one, while ELK does not infer any of the two. Table 2 summarizes the results
of the sandbox experiment. Here, _Disjointness materialization_ refers to the
addition of disjointness axioms to subclasses. In the example above, if
`clgo:Swedish_Rock_Musician` was defined as a subclass of `clgo:Person`, and
`clgo:Person` is disjoint with `clgo:Organization`, the reasoner would also
infer that `clgo:Swedish_Rock_Musician` is disjoint with `clgo:Organization`
(and its subclasses). However, we just inspected that out of curiosity, and
removed disjointess axioms from the larger scale datasets.

Reasoner | ELK | HermiT | Pellet
---|---|---|---
Subclass materialization | X | X | X
Disjointness materialization |  | X | X
Restrictions with individuals |  | X | X
Restrictions with literals |  |  | X
Table 2: Reasoner capabilities as evaluated in the sandbox experiments

We tried a couple of other reasoners, which were not capable of completely
processing the Sandbox example, including

  * •

jcel [15] version 0.24.121

  * •

Snorocket [16] version 4.0.17

  * •

Konclude [17] version 0.7.0

Those did either not produce any of the desired statements, or reported by
themselves that they do not support inference based on owl:hasValue
constructs. Thus, we decided to only keep ELK as a representative of those
reasoners which cannot produce the inferences sought to compare its
performance to those which can, and discarded the rest.

###  3.3 Scalability Experiments

To conduct experiments for scalability, we extracted differently sized subsets
of CaLiGraph, each of which should have around $10^{n}$ classes for different
values of $n$. The resulting datasets are depicted in table 3. The datasets
were designed so that they create meaningful subtrees of the overall hierarchy
in CaLiGraph. The corresponding SPARQL query is depicted in figure 3. The
different datasets are generated by varying the inner SELECT statement which
retrieves the leaf classes. Datasets up to clg_10e4 are drawn from subclasses
of `clgo:Organization`, while clg_10e5 is drawn from the complete ontology.
The relatively low number of transitive types of the latter dataset is
explained by the fact that the average type depth of subclasses of
`clgo:Organization` is much higher than the overall average type depth.

Dataset | Triples | Classes | Restrictions | Instances | Inferrable Assertions
---|---|---|---|---|---
|  |  |  |  | Trans. Types | Individuals | Literals
clg_10 | 35 | 9 | 2 | 3 | 10 | 1 | 1
clg_10e2 | 510 | 101 | 3 | 70 | 425 | 20 | 1
clg_10e3 | 43,120 | 1,000 | 79 | 18,072 | 4,500 | 514 | 5,706
clg_10e4 | 297,266 | 10,006 | 1,299 | 115,501 | 73,683 | 12,085 | 10,369
clg_10e5 | 4,641,400 | 99,923 | 12,147 | 1,968,282 | 42,820 | 187,929 | 7,120
clg_full | 54,914,982 | 1,061,598 | 227,173 | 14,452,393 | 138,713,499 | 9,637,354 | 969,022
Table 3: Overview of the different CaLiGraph subsets used for the evaluation

    PREFIX owl: <http://www.w3.org/2002/07/owl#>
    PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>
    PREFIX clgo: <http://caligraph.org/ontology/>
    CONSTRUCT {
        ?leaf ?leafPred ?leafObj .
        ?super ?superPred ?superObj .
        ?ind a owl:NamedIndividual, ?leaf .
        ?resPred ?resPredPred ?resPredObj .
        ?resObj ?resObjPred ?resObjObj .
    } WHERE {
        # Collect set of leaf classes to bootstrap sample from
        # Sample subgraph can be changed by picking something else than clgo:Organization as basis
        # Use a specific LIMIT for N to restrict the number of leaf classes in the sample
        {
            SELECT ?leaf {
                ?leaf rdfs:subClassOf+ clgo:Organization .
                FILTER NOT EXISTS {?sub rdfs:subClassOf ?leaf}
            }
            LIMIT N
        }
        # Collect all relevant superclasses of the leaves
        ?leaf rdfs:subClassOf+ ?super ;
              ?leafPred ?leafObj .
        ?super ?superPred ?superObj .
        # Collect individuals for the leaves
        ?ind a ?leaf .
        # Add restrictions, if available
        OPTIONAL {
            ?leaf rdfs:subClassOf+ ?res .
            ?res a owl:Restriction ;
                 owl:onProperty ?resPred ;
                 owl:hasValue ?resObj .
            # Add information about the predicates and objects of restrictions to the sample
            ?resPred ?resPredPred ?resPredObj .
            OPTIONAL {
                ?resObj ?resObjPred ?resObjObj .
            }
        }
        # Discard disjointnesses as this would blow up the sample size
        FILTER (?superPred != owl:disjointWith)
        FILTER (?leafPred != owl:disjointWith)
    }

Figure 3: SPARQL query to generate datasets for scalability experiments

We conducted experiments on a system with 32 Intel Xeon CPUs and 200GB RAM,
and we set a timeout of 72h. Overall, ELK was the only system capable of
processing all datasets (but not returning any of the inferrable individual
and literal assertions, as discussed above). The processing of the largest
dataset with ELK took about 3.5 hours.

HermiT and Pellet, on the other hand, could only process clg_10e2 and
clg_10e3, respectively, and timed out on the larger ones (with Pellet running
out of memory on clg_full. This shows that those reasoning systems are still
3-4 orders of magnitude away from the reasoning capabilities required for
processing CaLiGraph.

Figure 4: Scalability of the Reasoning Systems investigated

##  4 Conclusion and Outlook

In this paper, we investigated three state of the art OWL2EL reasoners, i.e.,
ELK, HermiT, and Pellet. The task for those reasoners was to materialize the
A-box of the CaLiGraph knowledge graph, particularly the creation of around
10M A-box assertions from owl:hasValue definitions. As discussed above, this
kind of reasoning at scale is currently not evaluated in benchmarks such as
ORE2015 and OWL2Bench.

The outcome of our experiments is that the reasoners evaluated are either (a)
not capable of performing that inference at all, as in the case of ELK, or (b)
not scalable enough to process knowledge graphs with 10k instances or more –
which is some orders of magnitude below the size of popular knowledge graphs
[4].

With this work, we have also provided a benchmark suite of differently sized
test sets.555The datasets, as well as the code for running the experiments,
are available at https://github.com/nheist/CaLiGraph-for-SemREC This allows
for running scalability experiments also with respect to materializing a large
amount of individual and literal A-box assertions. Moreover, the benchmark is
suitable for running evaluations with approximate reasoning system, since the
recall and precision of A-box assertions generated can be quantified.
Therefore, the dataset is also suitable for measuring the trade-off of
scalability and accuracy of approximate reasoning systems.

## References

  * Lehmann et al. [2015] J. Lehmann, R. Isele, M. Jakob, A. Jentzsch, D. Kontokostas, P. N. Mendes, S. Hellmann, M. Morsey, P. Van Kleef, S. Auer, et al.,  Dbpedia–a large-scale, multilingual knowledge base extracted from wikipedia,  Semantic web 6 (2015) 167–195.
  * Suchanek et al. [2007] F. M. Suchanek, G. Kasneci, G. Weikum,  Yago: a core of semantic knowledge,  in: Proceedings of the 16th international conference on World Wide Web, 2007, pp. 697–706.
  * Vrandečić [2012] D. Vrandečić,  Wikidata: A new platform for collaborative data collection,  in: Proceedings of the 21st international conference on world wide web, 2012, pp. 1063–1064.
  * Heist et al. [2020] N. Heist, S. Hertling, D. Ringler, H. Paulheim,  Knowledge graphs on the web-an overview,  in: P. H. Ilaria Tiddi, Freddy Lécué (Ed.), Knowledge Graphs for eXplainable Artificial Intelligence: Foundations, Applications and Challenges, volume 47 of Studies on the Semantic Web, IOS Press/AKA, 2020, pp. 3–22.
  * Heist and Paulheim [2019] N. Heist, H. Paulheim,  Uncovering the semantics of wikipedia categories,  in: International semantic web conference, Springer, 2019, pp. 219–236.
  * Heist and Paulheim [2020] N. Heist, H. Paulheim,  Entity extraction from wikipedia list pages,  in: European Semantic Web Conference, Springer, 2020, pp. 327–342.
  * Heist and Paulheim [2021] N. Heist, H. Paulheim,  Information extraction from co-occurring similar entities,  in: Proceedings of the Web Conference 2021, 2021, pp. 3999–4009.
  * Motik et al. [2009] B. Motik, B. C. Grau, I. Horrocks, Z. Wu, A. Fokoue, C. Lutz, et al., Owl 2 web ontology language profiles, 2009. URL: https://www.w3.org/TR/owl2-profiles/.
  * Antoniou et al. [2018] G. Antoniou, S. Batsakis, R. Mutharaju, J. Z. Pan, G. Qi, I. Tachmazidis, J. Urbani, Z. Zhou,  A survey of large-scale reasoning on the web of data,  The Knowledge Engineering Review 33 (2018).
  * Parsia et al. [2017] B. Parsia, N. Matentzoglu, R. S. Gonçalves, B. Glimm, A. Steigmiller,  The owl reasoner evaluation (ore) 2015 competition report,  Journal of Automated Reasoning 59 (2017) 455–482.
  * Singh et al. [2020] G. Singh, S. Bhatia, R. Mutharaju,  Owl2bench: A benchmark for owl 2 reasoners,  in: International Semantic Web Conference, Springer, 2020, pp. 81–96.
  * Kazakov et al. [2014] Y. Kazakov, M. Krötzsch, F. Simančík,  The incredible elk,  Journal of automated reasoning 53 (2014) 1–61.
  * Glimm et al. [2014] B. Glimm, I. Horrocks, B. Motik, G. Stoilos, Z. Wang,  Hermit: an owl 2 reasoner,  Journal of Automated Reasoning 53 (2014) 245–269.
  * Sirin et al. [2007] E. Sirin, B. Parsia, B. C. Grau, A. Kalyanpur, Y. Katz,  Pellet: A practical owl-dl reasoner,  Journal of Web Semantics 5 (2007) 51–53.
  * Mendez [2012] J. Mendez,  jcel: A modular rule-based reasoner.,  in: ORE, 2012.
  * Lawley and Bousquet [2010] M. J. Lawley, C. Bousquet,  Fast classification in protégé: Snorocket as an owl 2 el reasoner,  in: Proc. 6th Australasian Ontology Workshop (IAOA’10). Conferences in Research and Practice in Information Technology, volume 122, 2010, pp. 45–49.
  * Steigmiller et al. [2014] A. Steigmiller, T. Liebig, B. Glimm,  Konclude: system description,  Journal of Web Semantics 27 (2014) 78–85.


# Understanding the Training Speedup from Sampling with Approximate Losses

Rudrajit Das Correspondence to Rudrajit Das<EMAIL_ADDRESS>Xi Chen
<EMAIL_ADDRESS>and Sujay Sanghavi<EMAIL_ADDRESS>UT
Austin  Xi Chen∗ Amazon  Bertram Ieong  Amazon  Parikshit Bansal  UT Austin
Sujay Sanghavi∗ UT Austin  Amazon

###### Abstract

It is well known that selecting samples with large losses/gradients can
significantly reduce the number of training steps. However, the selection
overhead is often too high to yield any meaningful gains in terms of overall
training time. In this work, we focus on the greedy approach of selecting
samples with large approximate losses instead of exact losses in order to
reduce the selection overhead. For smooth convex losses, we show that such a
greedy strategy can converge to a constant factor of the minimum value of the
average loss in fewer iterations than the standard approach of random
selection. We also theoretically quantify the effect of the approximation
level. We then develop SIFT which uses early exiting to obtain approximate
losses with an intermediate layer’s representations for sample selection. We
evaluate SIFT on the task of training a 110M parameter 12 layer BERT base
model, and show significant gains (in terms of training hours and number of
backpropagation steps) without any optimized implementation over vanilla
training. For e.g., to reach 64% validation accuracy, SIFT with exit at the
first layer takes $\sim$ 43 hours compared to $\sim$ 57 hours of vanilla
training.

##  1 Introduction

Stochastic Gradient Descent (SGD) and its variants are the algorithms of
choice for solving large-scale optimization problems that arise in training
machine learning models. These are problems of the form
$\min_{\bm{w}\in\mathbb{R}^{d}}F(\bm{w})$, where $F(.)$ is the expected
population loss and $\bm{w}\in\mathbb{R}^{d}$ is the vector of model
parameters. More specifically, if $f(\bm{w},.)$ is the per-sample loss and
$\mathcal{D}$ is the data distribution, then
$F(\bm{w})=\mathbb{E}_{\bm{x}\sim\mathcal{D}}[f(\bm{w},\bm{x})]$. The standard
SGD update rule at $\bm{w}$ with step-size/learning rate $\eta$ is:

$\bm{w}_{+}=\bm{w}-\eta\nabla f(\bm{w},\bm{x}),$ (1)

where the sample/example $\bm{x}$ is drawn from $\mathcal{D}$. In practice,
the gradient over a single sample is replaced by the average gradient over a
mini-batch of samples; for our theoretical results, we shall stick to batch-
size = 1 (as in eq. 1).

The choice of the sample $\bm{x}$ in eq. 1 significantly impacts the speed of
convergence. There is copious work on selecting “important” samples for
speeding up convergence. In this paradigm, the seminal idea is that of
importance sampling which proposes to sample the examples such that resulting
stochastic gradient is unbiased while having the minimum possible variance
(Zhao and Zhang,, 2015; Alain et al.,, 2015; Needell et al.,, 2014). It turns
out that the optimal solution is to sample the examples with probability
proportional to their per-sample gradient norms. While this has a clean
theoretical solution, it is completely infeasible from the practical
standpoint because of the high cost of computing the per-sample gradient
norms. To address this shortcoming, several approximations to this exact
solution have been proposed; we discuss these and related ideas in Section 2.
In this space, one high-level idea is to use the per-sample loss value as a
proxy to the per-sample gradient norm, favoring samples with high losses for
performing the update (Loshchilov and Hutter,, 2015; Shrivastava et al.,,
2016; Katharopoulos and Fleuret,, 2017; Kawaguchi and Lu,, 2019; Zhang et
al.,, 2019). We discuss the important works based on this idea in detail in
Section 2. Since exact loss computation is itself expensive especially for
large models, some works rely on approximate loss values for sample selection;
for instance, using an auxiliary model to predict loss values of the actual
model being trained (Katharopoulos and Fleuret,, 2017; Zhang et al.,, 2019).
So in summary, several approximations to the core idea of exact importance
sampling have been proposed to make it practically usable.

Despite the myriad sample selection approaches, there is a paucity of
theoretical results quantifying how much speed-up we can obtain over random
sampling and their limitations – especially for approaches involving the use
of approximate quantities. In this work, we focus on the approach of choosing
the sample with the highest loss for performing the gradient-based update.
More specifically, suppose we are given $R>1$ i.i.d. samples drawn from
$\mathcal{D}$ and their approximate loss values with the constraint that we
can pick only one sample by inspecting the approximate losses to perform the
update. We consider the greedy approach of selecting the sample with the
largest approximate loss, and we call this greedy SGD (GSGD). Also, we will
refer to the default approach of picking a sample uniformly at random as
(vanilla) SGD. In Section 5, we theoretically compare the convergence rates of
GSGD (with approximate losses) against SGD, characterizing the benefits and
limitations of the former. We would like to clarify that we are not claiming
that GSGD is a novel algorithm; in fact, it is very similar to OSGD (Kawaguchi
and Lu,, 2019) in spirit111We do not call it OSGD because Kawaguchi and Lu,
(2019) use exact loss values whereas we use approximate loss values. Moreover,
Kawaguchi and Lu, (2019) consider the finite-sum (ERM) setting whereas we
focus on the stochastic setting.. Rather, we are claiming that our theoretical
characterization of its benefits and limitations – with approximate losses –
is the first of its kind, to the best of our knowledge.

On the applied side, we propose to use early exiting (Teerapittayanon et al.,,
2016; Schwartz et al.,, 2020) as a light-weight way to obtain approximate
losses for sample selection in training large ML models (e.g., transformers).
To be clear, our novelty is in the light-weight filtering process for training
via early exiting. Empirical results on a 110M parameter BERT base model show
that early exit-based sample selection yields significant improvements.

We will now elaborate on our main contributions.
(a) In Section 5, we theoretically compare greedy SGD (GSGD) (i.e., pick the
sample with the highest approximate loss) against vanilla SGD (i.e., pick a
sample at random).

  * •

Theorem 1 provides a convergence bound for GSGD on smooth convex losses
assuming that the argmax of the approximate losses is equal to the argmax of
the actual losses (5). In Section 5.1, we consider a setting wherein 5 no
longer holds. We quantify the degradation in the performance of GSGD in this
setting; see Theorems 3 and 4.

  * •

A key insight is that GSGD can converge to a (problem-dependent) constant
factor of $F^{*}=\min_{\bm{w}}F(\bm{w})$ in fewer iterations than SGD; see
Remark 2. This is of interest in training large ML models on extremely large
datasets, where converging exactly to $F^{*}$ is infeasible and instead
converging faster to $\mathcal{O}(F^{*})$ is desirable. On the negative side,
our result indicates that GSGD may not converge to $F^{*}$ asymptotically and
hence it can be worse than SGD asymptotically (Remark 1).

(b) In Section 6, we propose early exiting as a way to cheaply obtain
approximate losses with an intermediate layer’s representations for sample
selection in training large ML models such as transformers. Specifically, we
propose to do backpropagation on only 50% of the samples in a batch with the
highest approximate losses obtained via early exiting. To our knowledge, early
exiting has not been studied for accelerating training. We call this early-
exit based sample “sifting” process SIFT. This can be seamlessly integrated
with other sample selection schemes; in particular, we also try large entropy-
based sample selection with early exiting.

  * •

In Section 7, we show the efficacy of SIFT in training a 12 layer BERT base
model with 110M parameters from scratch. Specifically, to achieve 64%
validation accuracy, loss-based and entropy-based SIFT with exit at the first
layer take roughly 43 and 40 hours, respectively, compared to roughly 57 hours
of vanilla training involving no sample selection.

  * •

In Section 6.1, we quantify the probability of correctly selecting the sample
with the largest actual loss using early exiting for feedforward linear neural
networks as a function of the model’s weights.

##  2 Related Work

There is a large body of work proposing several kinds of sample-selection
schemes for accelerated training. These include importance sampling methods,
sample reordering based approaches and algorithms focusing on samples with
higher losses. Our approach falls in the last category.

Optimal importance sampling. Importance sampling asks how should the training
examples be sampled so as to obtain an unbiased stochastic gradient with the
minimum possible variance. Zhao and Zhang, (2015); Alain et al., (2015);
Needell et al., (2014) show that the optimal solution to this problem is to
sample the examples with probability proportional to their per-sample gradient
norms. Unfortunately, this is completely infeasible in practice because of the
high cost of computing the per-sample gradient norms.

Approximate importance sampling. Several papers attempt to approximate the
optimal importance sampling procedure so as to make it feasible. For convex
settings, Zhao and Zhang, (2015); Needell et al., (2014) propose sampling with
probability proportional to the smoothness constant of the per-sample losses,
while Borsos et al., (2018); Stich et al., (2017) propose adaptive sampling
strategies. Going beyond convex settings, Alain et al., (2015) present a
distributed approach for importance sampling, while Katharopoulos and Fleuret,
(2017) approximate the importance weights with loss values which are predicted
by a smaller network. Katharopoulos and Fleuret, (2018) derive an upper bound
on the per-sample gradient norms for deep-learning networks which take
essentially the same time to compute as the per-sample loss values, and
propose using the upper bounds for approximating the importance weights.
Johnson and Guestrin, (2018) approximate the true sampling distribution by
solving a robust optimization problem.

Sample reordering. Another related line of work attempts to improve the
way/order in which samples are presented while training. Bengio et al., (2009)
propose curriculum learning wherein the key idea is to present easier examples
before harder ones but this requires prior information about the training set.
Several improved modifications of this idea are out there (Tsvetkov et al.,,
2016; Kumar et al.,, 2010; Jiang et al.,, 2017; Kim and Choi,, 2018; Zhang et
al.,, 2019; Jiang et al.,, 2019). There is also a long line of papers that
theoretically analyze conventional sample ordering schemes (such as shuffle
once, random reshuffling, etc.) as well as improved sample ordering schemes
(Bertsekas et al.,, 2011; Recht and Ré,, 2012; Gurbuzbalaban et al.,, 2019,
2022; Haochen and Sra,, 2019; Safran and Shamir,, 2020; Mishchenko et al.,,
2020; Lu et al.,, 2022; Mohtashami et al.,, 2022).

Selection of samples with large losses. As mentioned in Section 1, GSGD is not
a novel algorithm and there are prior approaches with the same underlying
principle as GSGD, i.e., focus on the samples with the largest loss values.
The two algorithms closest to GSGD in spirit are online hard example mining
(OHEM) proposed by Shrivastava et al., (2016) and ordered SGD (OSGD) proposed
by Kawaguchi and Lu, (2019) – except that both algorithms use exact losses for
sample selection. OHEM was proposed for training object detectors (in computer
vision) and it involves back-propagating using the gradients of only the
samples with the $k$ largest losses in batch of size $b$ with $k<b$. However,
Shrivastava et al., (2016) do not provide any theoretical guarantees. OSGD
(Kawaguchi and Lu,, 2019) is essentially the same algorithm as OHEM
(Shrivastava et al.,, 2016), but with convergence guarantees and
generalization bounds. However, unlike our theoretical results for GSGD,
Kawaguchi and Lu, (2019) do not show how/when/to what extent ordered SGD is
better than plain SGD. So even though GSGD is similar to OSGD in spirit, our
theoretical results (for GSGD) are much more comprehensive. Loshchilov and
Hutter, (2015) propose a sampling strategy which favors picking examples with
larger losses. Fan et al., (2017) introduce the average top-$k$ loss and
advocate minimizing this loss rather than the empirical average over all the
samples.

Early Exiting. Early exiting (Teerapittayanon et al.,, 2016; Schwartz et al.,,
2020) is a promising approach to decreasing the computational cost of
multilayered neural architectures by approximating the output of a model
through its intermediate feature representations. This saves computational
costs by dynamically deciding the number of layers/modules (attention-blocks
in Transformers) to use during inference by exiting based on some metric
computed on the intermediate representations themselves. Initially used for
ResNets (Teerapittayanon et al.,, 2016), early exiting is now widely popular
even for transformer models, especially language models (Schwartz et al.,,
2020; Xin et al.,, 2020; Schuster et al.,, 2022; Rotem et al.,, 2023; Xin et
al.,, 2021; Zhu,, 2021). Recent work around early exiting for large language
models also explores accelerated decoding (Schuster et al.,, 2022) and
improving factuality (Chuang et al.,, 2023). In this work, we solely use early
exiting as a light-weight way to obtain approximate losses with the
intermediate layer representations for sample selection. To our knowledge,
early exiting has not been used in prior work for speeding up training.

##  3 Notation

Vectors and matrices are in bold font. For a natural number $n$, we sometimes
denote the set $\\{1,\ldots,n\\}$ by $[n]$. The (Gauss) error function and
complementary error function are defined as:

$\text{erf}(t):=\frac{2}{\sqrt{\pi}}\int_{0}^{t}e^{-z^{2}}{d}z\text{ and
}\text{erfc}(t):=1-\text{erf}(t).$ (2)

Note that $\lim_{t\to\infty}\text{erf}(t)=1$. For any $z\in\mathbb{R}$, we
define the sigmoid function as $\text{sig}(z)=\frac{1}{1+e^{z}}$.

##  4 Problem Setting

We briefly recap the optimization problem introduced in Section 1. Given
access to a first-order optimization oracle, we would like to minimize
$F(\bm{w}):=\mathbb{E}_{\bm{x}\sim\mathcal{D}}[f(\bm{w},\bm{x})]$, where
$\bm{w}\in\mathbb{R}^{d}$ and $\mathcal{D}$ is the data distribution.

Standard First-Order Stochastic Optimization Oracle. A query at $\bm{w}$
returns $\nabla_{\bm{w}}f(\bm{w},\bm{x})$, where $\bm{x}\sim\mathcal{D}$.

We consider a more generous oracle. However, it can also perform one gradient
evaluation per query (same as the standard oracle).

######  Definition 1 (Proposed First-Order Stochastic Optimization Oracle).

A query at $\bm{w}$ first returns a set of $R>1$ samples
$\\{\bm{x}^{(1)},\ldots,\bm{x}^{(R)}\\}=\mathcal{S}^{(R)}$ drawn i.i.d. from
$\mathcal{D}$ and their approximate function222Throughout this paper, we
interchange “loss” value and “function” value freely. In most cases, we use
“function” value in the context of optimization-based discussions and “loss”
value in the context of ML-based discussions. values
$\\{\widetilde{f}(\bm{w},\bm{x}^{(1)}),\ldots,\widetilde{f}(\bm{w},\bm{x}^{(R)})\\}$.
The user can then pick one $\widehat{\bm{x}}$ from $\mathcal{S}^{(R)}$, and
the oracle will return $\nabla_{\bm{w}}f(\bm{w},\hat{\bm{x}})$.

Later, we shall make an assumption on the relation between the approximate
function value $\widetilde{f}(\bm{w},\bm{x})$ and the actual function value
${f}(\bm{w},\bm{x})$. Now that we have introduced the oracle that we consider
in this work, we state the $\widehat{\bm{x}}$ chosen by vanilla SGD and greedy
SGD (GSGD).
SGD choice: Pick $\widehat{\bm{x}}$ uniformly at random from
$\mathcal{S}^{(R)}$.
GSGD choice: Pick $\widehat{\bm{x}}=\text{arg
max}_{\bm{x}\in\mathcal{S}^{(R)}}\widetilde{f}(\bm{w},\bm{x})$.
We state the GSGD update rule in more detail next.

###  4.1 Greedy SGD (GSGD) Algorithm

In the $k^{\text{th}}$ iteration, we observe a set of $R$ i.i.d. samples drawn
from $\mathcal{D}$, say
$\mathcal{S}_{k}^{(R)}=\\{\bm{x}_{k}^{(1)},\ldots,\bm{x}_{k}^{(R)}\\}$. We
pick:

$\widehat{\bm{x}}_{k}=\text{arg
max}_{\bm{x}\in\mathcal{S}_{k}^{(R)}}\widetilde{f}(\bm{w}_{k},\bm{x}).$ (3)

The update of Greedy SGD (GSGD) with step-size $\eta_{k}$ is:

$\bm{w}_{k+1}=\bm{w}_{k}-\eta_{k}\nabla f(\bm{w}_{k},\widehat{\bm{x}}_{k}).$
(4)

Vanilla SGD’s update is the same as eq. 4, except with $\widehat{\bm{x}}_{k}$
being a random sample from $\mathcal{S}_{k}^{(R)}$.

##  5 GSGD vs. SGD for Smooth Convex Objectives

We begin by stating some assumptions and definitions.

######  Assumption 1 (Continuity).

For any $\bm{x}\sim\mathcal{D}$, $f(\bm{w},\bm{x})$ is continuous w.r.t.
$\bm{w}$.

######  Assumption 2 (Convexity).

For any $\bm{x}\sim\mathcal{D}$, $f(\bm{w},\bm{x})$ is convex w.r.t. $\bm{w}$.

######  Assumption 3 (Smoothness).

For any $\bm{x}\sim\mathcal{D}$, $f(\bm{w},\bm{x})$ is $L$-smooth w.r.t.
$\bm{w}$.

######  Assumption 4.

For any $\bm{x}\sim\mathcal{D}$,
$\min_{\bm{w}\in\mathbb{R}^{d}}f(\bm{w},\bm{x})=0$.

Let $\Phi_{F}:=\text{arg min}_{\bm{w}\in\mathbb{R}^{d}}F(\bm{w})$ and
$F^{*}:=\min_{\bm{w}\in\mathbb{R}^{d}}F(\bm{w})$. We restrict our attention to
the case of $\Phi_{F}$ being closed and compact.

Let us first consider the case where the argmax of the approximate function
values is the same as the argmax of the exact function values; we will relax
this assumption later in Section 5.1.

######  Assumption 5 (Approximate function values preserve argmax).

In the setting of Definition 1, $\widetilde{f}$ satisfies:

$\textup{arg
max}_{\bm{x}\in\mathcal{S}^{(R)}}\widetilde{f}(\bm{w},\bm{x})=\textup{arg
max}_{\bm{x}\in\mathcal{S}^{(R)}}{f}(\bm{w},\bm{x}).$

Under 5, $\widehat{\bm{x}}_{k}$ in eq. 3 becomes the same as $\text{arg
max}_{\bm{x}\in\mathcal{S}_{k}^{(R)}}{f}(\bm{w}_{k},\bm{x})$.

######  Definition 2.

For $R>1$, let

$\widehat{F}_{R}(\bm{w}):=\mathbb{E}_{\\{\bm{x}^{(j)}\\}_{j=1}^{R}\text{
}\underset{\textup{iid}}{\sim}\text{ }\mathcal{D}}\Big{[}\max_{\text{
}\bm{x}\in\\{\bm{x}^{(j)}\\}_{j=1}^{R}}f(\bm{w},\bm{x})\Big{]}.$

Define

$\rho_{R}(\bm{w}):=\frac{\widehat{F}_{R}(\bm{w})}{F(\bm{w})}\text{ and
}\rho_{R}^{*}:=\inf_{\bm{w}\notin\Phi_{F}}\rho_{R}(\bm{w}).$

Also, suppose

$\sup_{\bm{w}^{*}\in\Phi_{F}}\widehat{F}_{R}(\bm{w}^{*})\leq\Delta_{R}.$

Except for the trivial case of $f(\bm{w},\bm{x}_{1})=f(\bm{w},\bm{x}_{2})$
$\forall$ $\bm{x}_{1},\bm{x}_{2}$ which we disregard, $\rho_{R}^{*}$ is
strictly bigger than 1. Further, consider a point
$\widehat{\bm{w}}^{*}\notin\Phi_{F}$ that is $\epsilon$-close to some
$\bm{w}^{*}\in\Phi_{F}$ (i.e.,
$\|\widehat{\bm{w}}^{*}-\bm{w}^{*}\|_{2}\leq\epsilon$) and let $\epsilon\to
0$. In that case, under Assumption 1 (and because the max operation preserves
continuity) and Definition 2,
$\widehat{F}_{R}(\widehat{\bm{w}}^{*})\to\widehat{F}_{R}({\bm{w}}^{*})\leq\Delta_{R}$.
Also, $F(\widehat{\bm{w}}^{*})\geq F^{*}$. Thus, by definition,
$\rho_{R}^{*}\leq{\widehat{F}_{R}(\widehat{\bm{w}}^{*})}/{F(\widehat{\bm{w}}^{*})}\leq({\Delta_{R}}/{F^{*}})$.
Hence, we have that:

$1<\rho_{R}^{*}\leq\frac{\Delta_{R}}{F^{*}},\text{ }\forall\text{ }R>1.$ (5)

In Section 5.2, we quantify $\rho_{R}(\bm{w})$ and $\rho_{R}^{*}$ for fitting
a model with the squared loss.

For our convergence results, we assume that our initialization is $\bm{w}_{0}$
and let:

$D_{0}:=\min_{\bm{w}^{*}\in\Phi_{F}}\|\bm{w}_{0}-\bm{w}^{*}\|.$

We are now ready to state our convergence results.

######  Theorem 1 (GSGD).

Suppose Assumptions 1, 2, 3, 4 and 5 hold. Set $\eta_{k}=\eta<\frac{1}{L}$ for
all $k$. Then, GSGD has the following convergence guarantee after $K$
iterations:

$\mathbb{E}\Bigg{[}F\Bigg{(}\frac{1}{K}\sum_{k=0}^{K-1}\bm{w}_{k}\Bigg{)}\Bigg{]}\leq\frac{D_{0}^{2}}{2\rho_{R}^{*}\eta\big{(}1-\eta
L\big{)}K}+\frac{\Delta_{R}}{\rho_{R}^{*}\big{(}1-\eta L\big{)}}.$

The proof of Theorem 1 can be found in Appendix B. We now state a
corresponding folklore result for SGD.

######  Theorem 2 (SGD).

Suppose Assumptions 2, 3 and 4 hold. Set $\eta_{k}=\eta<\frac{1}{L}$ for all
$k$. Then, SGD has the following convergence guarantee after $K$ iterations:

$\mathbb{E}\Bigg{[}F\Bigg{(}\frac{1}{K}\sum_{k=0}^{K-1}\bm{w}_{k}\Bigg{)}\Bigg{]}\leq\frac{D_{0}^{2}}{2\eta\big{(}1-\eta
L\big{)}K}+\frac{\eta LF^{*}}{1-\eta L}+F^{*}.$

From Theorem 1, observe that in general, GSGD may not converge to the minimum
value $F^{*}$ asymptotically (i.e., with $K\to\infty$). At best, we can show
that GSGD converges to $\frac{\Delta_{R}}{\rho_{R}^{*}}$ which is $\geq F^{*}$
(this follows from eq. 5). But SGD can indeed converge to $F^{*}$
asymptotically by setting $\eta=\frac{1}{\eta L\sqrt{K}}$ for example in
Theorem 2. Based on this, we make the following remark.

######  Remark 1.

GSGD may be worse than SGD asymptotically.

However, GSGD is better than SGD early on as $\rho_{R}^{*}>1$ and assuming
$\Delta_{R}=\mathcal{O}(F^{*})$ (and $F^{*}\neq 0$), GSGD can converge to a
constant factor of $F^{*}$ in fewer iterations than SGD. We formalize this
next.

######  Corollary 1 (Up to what point is GSGD better than SGD?).

Suppose we run GSGD and SGD with constant step-size $\eta<\frac{1}{L}$. In
that case, until
$K=\frac{D_{0}^{2}(\rho_{R}^{*}-1)}{2\eta(\Delta_{R}-\rho_{R}^{*}F^{*})}$
iterations, the convergence bound of GSGD in Theorem 1 is better than that of
SGD in Theorem 2. Thus, GSGD can converge to $\frac{\Delta_{R}-F^{*}}{(1-\eta
L)(\rho_{R}^{*}-1)}$ function value in fewer iterations than SGD.

######  Remark 2.

Based on Corollary 1, when $\Delta_{R}=\mathcal{O}(F^{*})$, GSGD can converge
to $\mathcal{O}(F^{*})$ function value in fewer iterations than SGD. This is
of particular interest in training large ML models such as transformers on
extremely large datasets, where minimizing the training loss exactly is
infeasible and converging faster to a constant factor of the minimum loss
value is desirable.

###  5.1 Beyond Argmax-Preserving Approximate Function Values

Our previous results were under 5, i.e., the argmax of the approximate
function values ($\widetilde{f}$) is always equal to the argmax of the actual
function values (${f}$). Here we relax this assumption by instead modeling
$\widetilde{f}$ as a noisy version of $f$, and provide a convergence result
for GSGD in such a setting. Modeling $\widetilde{f}$ as a noisy version of $f$
is analogous to modeling the stochastic gradients as a noisy version of the
actual gradient in vanilla stochastic optimization. Specifically, we make the
following assumption.

######  Assumption 6 (Approximate function values).

There exists $\mu(\bm{w})\in\mathbb{R}$ and $\sigma\geq 0$ such that:

$\widetilde{f}(\bm{w},\bm{x})={f}(\bm{w},\bm{x})\exp\Big{(}\mu(\bm{w})+\sigma{\zeta}(\bm{w},\bm{x})\Big{)},$

where ${\zeta}(\bm{w},\bm{x})$ is i.i.d. random noise with mean 0 and variance
1.

Thus, the approximate function value is the actual function values times the
exponential of random noise (i.e.,
$\exp\big{(}\sigma{\zeta}(\bm{w},\bm{x})\big{)}$) times some other scaling
(i.e., $\exp(\mu(\bm{w}))$). 6 is pretty mild as it does not involve any
particular distributional assumptions on ${\zeta}(\bm{w},\bm{x})$ (such as
Gaussian, etc.). The important thing to note is that under 6,
$\widehat{\bm{x}}_{k}$ in eq. 3 is not always equal to $\text{arg
max}_{\bm{x}\in\mathcal{S}_{k}^{(R)}}{f}(\bm{w}_{k},\bm{x})$ – unlike 5.

######  Definition 3.

Let

$\displaystyle\widehat{F}_{R,\textup{approx}}(\bm{w}):=\mathbb{E}_{\\{\bm{x}^{(j)}\\}_{j=1}^{R}\text{
}\underset{\textup{iid}}{\sim}\text{
}\mathcal{D},\zeta}\Big{[}f\big{(}\bm{w},\bm{x}^{(j^{*})}\big{)}\Big{|}j^{*}=\textup{arg
max}_{j\in[R]}\widetilde{f}(\bm{w},\bm{x}^{(j)})\Big{]}.$

Define

$\rho_{R,\textup{approx}}(\bm{w}):=\frac{\widehat{F}_{R,\textup{approx}}(\bm{w})}{F(\bm{w})}\text{
and
}\rho_{R,\textup{approx}}^{*}:=\inf_{\bm{w}\notin\Phi_{F}}\rho_{R,\textup{approx}}(\bm{w}).$

Clearly, $\widehat{F}_{R,\text{approx}}(\bm{w})\leq\widehat{F}_{R}(\bm{w})$
(as defined in Definition 2). So we also have:

$\sup_{\bm{w}^{*}\in\Phi_{F}}\widehat{F}_{R,\text{approx}}(\bm{w}^{*})\leq\sup_{\bm{w}^{*}\in\Phi_{F}}\widehat{F}_{R}(\bm{w}^{*})\leq\Delta_{R}.$

It is easy to extend the proof of Theorem 1 to obtain the following result for
GSGD under 6.

######  Theorem 3 (GSGD).

Suppose Assumptions 1, 2, 3, 4 and 6 hold. Set $\eta_{k}=\eta<\frac{1}{L}$ for
all $k$. Then, GSGD has the following convergence guarantee after $K$
iterations:

$\mathbb{E}\Bigg{[}F\Bigg{(}\frac{1}{K}\sum_{k=0}^{K-1}\bm{w}_{k}\Bigg{)}\Bigg{]}\leq\frac{D_{0}^{2}}{2\rho_{R,\textup{approx}}^{*}\eta\big{(}1-\eta
L\big{)}K}+\frac{\Delta_{R}}{\rho_{R,\textup{approx}}^{*}\big{(}1-\eta
L\big{)}}.$

The expectation in Theorem 3 also includes the randomness due to $\zeta(.)$
(i.e., the noise in the approximate function values). The proof of Theorem 3
is almost identical to the proof of Theorem 1 (see Appendix B) and is obtained
by replacing $\widehat{F}_{R}(\bm{w})$ with
$\widehat{F}_{R,\text{approx}}(\bm{w})$ and $\rho_{R}^{*}$ with
$\rho_{R,\text{approx}}^{*}$.

For the subsequent results in this subsection, we shall consider $R=2$.
Specifically, we will provide a lower bound for
$\widehat{F}_{2,\text{approx}}(\bm{w}),\rho_{2,\text{approx}}(\bm{w})$ and
$\rho_{2,\text{approx}}^{*}$ in terms of
$\widehat{F}_{2}(\bm{w}),\rho_{2}(\bm{w})$ and $\rho_{2}^{*}$, respectively,
as a function of the noise level $\sigma$.

######  Theorem 4.

Suppose 6 holds with $\sigma\leq\frac{1}{2\sqrt{2}}$. Define
$p(\sigma):=\big{(}1-0.72\big{(}1-e^{-\sqrt{2}\sigma}\big{)}\big{)}$. Then:

$\widehat{F}_{2,\textup{approx}}(\bm{w})\geq
p(\sigma)\widehat{F}_{2}(\bm{w}),\text{ }\rho_{2,\textup{approx}}(\bm{w})\geq
p(\sigma)\rho_{2}(\bm{w})\text{ and }\rho_{2,\textup{approx}}^{*}\geq
p(\sigma)\rho_{2}^{*}.$

The proof of Theorem 4 is in Appendix C. It is worth pointing out that our
result is independent of $\mu(\bm{w})$; intuitively, this is because a
constant scaling (w.r.t. the samples) does not change the argmax. Regarding
the dependence w.r.t. $\sigma$, notice that $p(\sigma)$ is a decreasing
function of $\sigma$. So as the noise level $\sigma$ increases, the lower
bound becomes worse; this happens because the quality of the approximate
function value worsens as $\sigma$ increases. Also, as a quick sanity check,
observe that $p(0)=1$. This makes sense because $\sigma=0$ means no effective
noise, i.e., the approximate function values match the actual function values
modulo the constant scaling $\exp(\mu(\bm{w}))$.

In the following corollary, we specify a bound for the noise level $\sigma$
below which $\rho_{2,\textup{approx}}^{*}>1$ and thus, we can obtain a speed-
up over SGD.

######  Corollary 2.

As long as

$\sigma<\sigma_{\textup{max}}:=\frac{1}{\sqrt{2}}\log\Big{(}\frac{18\rho_{2}^{*}}{25-7\rho_{2}^{*}}\Big{)},$

$p(\sigma)\rho_{2}^{*}>1$ and therefore, $\rho_{2,\textup{approx}}^{*}>1$.
Hence, for $\sigma<\sigma_{\textup{max}}$, GSGD with $R=2$ using approximate
function values for sample selection can be faster than SGD.

###  5.2 Quantifying $\rho_{R}(\bm{w})$

Here we shall quantify $\rho_{R}(\bm{w})$ for a particular case. Let us
consider the problem of learning a parameterized model
$\mathcal{M}(\bm{w}^{*},.)$ with the squared loss. In this case, our per-
sample objective function $f$ is:

$f(\bm{w},\bm{x})=\Big{(}\mathcal{M}(\bm{w},\bm{x})-\mathcal{M}(\bm{w}^{*},\bm{x})\Big{)}^{2}.$
(6)

We make the following assumption.

######  Assumption 7.

For $\bm{w}\neq\bm{w}^{*}$,
$\mathcal{M}(\bm{w},\bm{x})-\mathcal{M}(\bm{w}^{*},\bm{x})\underset{\textup{iid}}{\sim}\mathcal{N}\big{(}\varepsilon(\bm{w}),\delta^{2}(\bm{w})\big{)}$,
i.e., the per-sample prediction error is a Gaussian random variable.

Modeling the per-sample prediction error as a Gaussian random variable is
fairly reasonable and a similar assumption has been made in Pennington and
Bahri, (2017). We provide a lower bound for $\rho(\bm{w})$ in this setting.

######  Theorem 5.

Suppose 7 holds. Then:

$\displaystyle\frac{\Big{(}\varepsilon^{2}(\bm{w})+\Big{(}{\frac{\pi}{2}\log\frac{R}{4\log
R}}\Big{)}\delta^{2}(\bm{w})+\sqrt{2\pi\log\frac{R}{4\log
R}}\varepsilon(\bm{w})\delta(\bm{w})\Big{)}\Big{(}1-\frac{1}{R}\Big{)}}{\varepsilon^{2}(\bm{w})+\delta^{2}(\bm{w})}.$

The proof of Theorem 5 is in Appendix D.

######  Corollary 3.

In the setting of Theorem 5, if
$\varepsilon(\bm{w})\leq\mathcal{O}(\delta(\bm{w}))$ for all $\bm{w}$ in our
region of optimization, then:

$\rho_{R}^{*}\geq\Bigg{(}1+\Omega\Bigg{(}\log\frac{R}{\log
R}+\sqrt{\log\frac{R}{\log R}}\Bigg{)}\Bigg{)}\Big{(}1-\frac{1}{R}\Big{)}.$

##  6 Approximate Losses via Early Exiting in Neural Networks

For sample selection in GSGD-like algorithms, we propose to use an
approximation of the actual loss which is computed on the “early” predictions
obtained by applying the linear head after the final layer to an intermediate
layer’s representation (instead of the final layer’s representation).
Specifically, suppose $\bm{\theta}$ is the linear head and for a sample
$\bm{x}$ with label $y$, $\mathcal{R}(\bm{x})$ and
$\widetilde{\mathcal{R}}(\bm{x})$ are the final layer’s and some intermediate
layer’s representation with the same dimension as $\bm{\theta}$, respectively.
Then the approximate and actual losses with the cross-entropy loss function
denoted by $\ell$ are
$\ell(y,\text{softmax}(\bm{\theta}^{\top}\widetilde{\mathcal{R}}(\bm{x})))$333For
a vector $\bm{v}=[v_{1},\ldots,v_{d}]$,
$\text{softmax}(\bm{v})=[\hat{v}_{1},\ldots,\hat{v}_{d}]$ with
$\hat{v}_{i}={\exp({v_{i}})}/{\sum_{j=1}^{d}\exp({v_{j}})}$. and
$\ell(y,\text{softmax}(\bm{\theta}^{\top}{\mathcal{R}}(\bm{x})))$,
respectively; we use the former to select samples for performing the gradient-
based update.

We shall now consider a feedforward linear neural network (Kawaguchi,, 2016)
and quantify the probability of correctly picking the sample with the largest
actual loss when using early exiting.

###  6.1 Probability of Correctly Selecting the Sample with the Largest Actual
Loss

We consider a binary classification problem, where each sample
$\bm{x}\sim\mathcal{N}(\vec{0}_{d},\bm{\text{I}}_{d})$ has a binary label
$y\in\\{0,1\\}$. Our model is a $k$-layer linear feed-forward network
parameterized by $\\{\bm{W}_{i}\\}_{i=1}^{d}\in\mathbb{R}^{d\times d}$ and
$\bm{\theta}\in\mathbb{R}^{d}$, where the soft prediction $\hat{y}$ for
$\bm{x}$ is:

$\hat{y}=\text{sig}\Big{(}\bm{\theta}^{\top}\big{(}\bm{W}_{k}\times\bm{W}_{k-1}\times\ldots\times\bm{W}_{1}\big{)}\bm{x}\Big{)}.$
(7)

In eq. 7, $\text{sig}(.)$ is the sigmoid function444For binary classification
problems, the sigmoid function essentially plays the role of the softmax
function mentioned earlier. as defined in Section 3. For any $j\in[k]$, let us
define:

$\bm{A}_{j}:=\bm{W}_{j}\times\ldots\times\bm{W}_{1}\text{ and
}\bm{B}_{j}:=\bm{W}_{k}\times\ldots\times\bm{W}_{j+1}.$ (8)

Then,

$\hat{y}_{j}=\text{sig}\big{(}\bm{\theta}^{\top}\bm{A}_{j}\bm{x}\big{)}$ (9)

is the “early prediction” for $\bm{x}$ at the $j^{\text{th}}$ layer. Note that
$\hat{y}=\hat{y}_{k}$, where:

$\hat{y}_{k}=\text{sig}\big{(}\bm{\theta}^{\top}\bm{B}_{j}\bm{A}_{j}\bm{x}\big{)}.$
(10)

In the context of GSGD, we once again focus on the case of $R=2$. For any two
i.i.d. samples $\bm{x}^{(1)}$ and $\bm{x}^{(2)}$, let $y^{(1)}$ and $y^{(2)}$
$\in\\{0,1\\}$ be the corresponding ground truth labels and let
$\hat{y}_{j}^{(1)}$ and $\hat{y}_{j}^{(2)}$ be the corresponding early
predictions at the $j^{\text{th}}$ layer. Further, let $\ell_{j}^{(1)}$ and
$\ell_{j}^{(2)}$ be the corresponding cross-entropy losses of the early
predictions at the $j^{\text{th}}$ layer, i.e.,

$\ell_{j}^{(i)}=-y^{(i)}\log\big{(}\hat{y}_{j}^{(i)}\big{)}-(1-y^{(i)})\log\big{(}1-\hat{y}_{j}^{(i)}\big{)},$
(11)

for $i\in\\{1,2\\}$. In the context of GSGD, $\ell_{j}^{(1)}$ and
$\ell_{j}^{(2)}$ are the approximate function values, whereas $\ell_{k}^{(1)}$
and $\ell_{k}^{(2)}$ are the actual function values. In this section, we are
interested in quantifying the probability (over the randomness of data) that
early exiting at the $j^{\text{th}}$ layer preserves the argmax; specifically,
we wish to quantify

$p_{j}:=\mathbb{P}_{\bm{x}^{(1)},\bm{x}^{(2)}}\Big{(}\text{arg
max}_{i\in[1,2]}\ell_{j}^{(i)}=\text{arg
max}_{i\in[1,2]}\ell_{k}^{(i)}\Big{)}.$ (12)

Note that $p_{k}=1$. We have the following result on this.

######  Theorem 6.

Define

$\beta_{j}:=\frac{\langle\bm{A}_{j}^{\top}\bm{\theta},\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}\rangle}{\|\bm{A}_{j}^{\top}\bm{\theta}\|_{2}\|\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}\|_{2}}.$

We restrict our attention to the case of $\beta_{j}\geq 0$. Then:

$p_{j}=1-\frac{1}{2\sqrt{\pi}}\int_{0}^{\infty}\exp\Big{(}-\frac{y^{2}}{4}\Big{)}\textup{erfc}\Bigg{(}\frac{\beta_{j}y}{2\sqrt{1-\beta_{j}^{2}}}\Bigg{)}dy,$
(13)

where $\textup{erfc}(.)$ is as defined in eq. 2. We also have the following
lower bound for $p_{j}$:

$p_{j}\geq 1-\sqrt{\frac{2-2\beta_{j}^{2}}{2-\beta_{j}^{2}}}.$ (14)

The proof of Theorem 6 can be found in Appendix E. Note that $\beta_{j}$ is
the (normalized) correlation between $\bm{A}_{j}^{\top}\bm{\theta}$ and
$\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}$. As per Theorem 6, $p_{j}=1$
when $\beta_{j}=1$; this makes sense because $\beta_{j}=1$ implies
$\bm{A}_{j}^{\top}\bm{\theta}$ and
$\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}$ are parallel and so the argmax
operation is unaffected if we use $\hat{y}_{j}$ instead of $\hat{y}_{k}$. It
is worth mentioning that the lower bound for $p_{j}$ in eq. 14 is loose for
small $\beta_{j}$, but we believe it is tight up to constant factors for
$\beta_{j}\approx 1$; in particular, it is exact in the case of $\beta_{j}=1$.
Specifically, we have the following simple corollary for $\beta_{j}\approx 1$.

######  Corollary 4.

In the setting of Theorem 6, suppose $\beta_{j}=1-\tau_{j}$ where $\tau_{j}\to
0$. Then, $p_{j}\geq 1-\mathcal{O}(\sqrt{\tau_{j}})$.

In simple words, if $\bm{A}_{j}^{\top}\bm{\theta}$ is strongly correlated with
$\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}$, then early exiting at the
$j^{\text{th}}$ layer preserves the argmax with high probability.

In the next section, we empirically demonstrate the efficacy of early exiting
in a transformer model.

##  7 Empirical Evaluation

We demonstrate the efficacy of early exiting in accelerating the training of a
BERT base model from scratch with the masked language modeling (MLM) loss.
Specifically, we apply early exiting for selecting samples in a mini-batch
version of greedy AdamW which is a simple extension of the greedy SGD idea to
AdamW. We refer to this practical early exit-based sample selection or
“sifting” strategy as SIFT. We describe SIFT in more detail later but at a
high level, we propose to backpropagate on only 50% of the examples in a batch
with the highest approximate losses obtained via early exiting. We would like
to emphasize that our novelty is in the light-weight sifting process for
training via early exiting and not the idea of backpropagating on the samples
with large losses555This is the same idea as OSGD (Kawaguchi,, 2016) and OHEM
(Shrivastava et al.,, 2016)..

The BERT base model used for our experiments here consists of 12 layers with a
hidden dimension of 768; the total number of parameters is 110M. We train on
BookCorpus (Zhu et al.,, 2015) and English Wikipedia which are two diverse and
extensive standard corpora. The validation set for assessing the model’s
performance is derived from the development partition of the training corpus.
This partition ensures that the validation set represents a diverse and
unbiased subset of the overall data. In our training, we use 512 as the
maximum sequence length. To process the input data, we use the “bert-base-
uncased” tokenizer from the Hugging Face model repository. The masking is
applied after tokenization with a uniform masking rate of 15%. Our experiments
were conducted on AWS p4d.24xlarge instances (8 NVIDIA A100 Tensor core GPUs).

Our baseline algorithm is the vanilla approach with no kind of sample
filtering. The micro-batch per GPU is set to 32 sequences (so there are 32
sequences per GPU * 8 GPUs * 512 tokens = 131072 tokens per batch). We will
now elaborate on SIFT which does greedy selection based on early exit loss.

Loss-based SIFT: We select 50% sequences per batch (for training) with the
largest MLM losses computed on the early predictions obtained from an
intermediate layer in the way described in the beginning of Section 6.
Specifically, we show results for the first, second, third, sixth and last
(i.e., twelfth) layer. The micro-batch per GPU is set to 64 sequences. Since
we select half of the sequences for training, the effective batch size per GPU
is 32 which is the same as that of baseline.

The early-exit based filtering idea is pretty general in the sense that it can
be seamlessly integrated with other sample selection schemes. In particular,
we also tried greedy selection based on early exit entropy (instead of loss).
We describe it below.

Entropy-based SIFT: Everything is the same as loss-based SIFT except that we
select 50% sequences per batch whose early predictions obtained from an
intermediate layer have the largest entropies instead of MLM losses.

Prediction entropy has been used a measure of model uncertainty for active
learning; see for e.g., Ren et al., (2021); Gal et al., (2017). So our
proposal of entropy-based SIFT may also be of interest in active learning with
large-scale models.

It is worth mentioning that in our implementation, the forward propagation for
selecting samples in SIFT is implemented naively before full forward and
backward propagation for the update step; one could potentially make this more
efficient. However, even without any optimization of the initial forward
propagation for sample selection, we obtain significant gains as we discuss
later. For the experiments with SIFT, we do not perform any sample filtering
for the first 20K steps; this is to provide an initial warm-up period and is
aligned with the suggestion of Kawaguchi and Lu, (2019) for GSGD. The total
number of update steps for each algorithm is 800K. We defer some other
experimental details to Appendix F.

We would like to point out that the goal of our experiments is to only show
the efficacy of early exiting for approximately selecting samples with large
losses/entropies and not propose a SOTA sample selection scheme.

Results. Figures 1 and 2 show the performance of SIFT with exits at the first,
sixth and last layers666We do not show the plots for exits at the second and
third layers here to avoid congestion; we report the validation accuracies of
all these layers in Table 1. and baseline w.r.t. the number of samples used
for backpropagation (i.e., backpropagation sample complexity) and the number
of training hours, respectively. We report the validation accuracy at the last
step with early exit at all the layers we considered in Table 1. Please see
the figure and table captions for detailed discussion but the key takeaways
are as follows:

  * •

SIFT is better than the baseline in terms of backpropagation sample complexity
as well as wall-clock time.

  * •

Entropy-based SIFT does better than loss-based SIFT.

  * •

For entropy-based SIFT, exit at the sixth layer has the best performance while
for loss-based SIFT, exit at the last layer has the best performance.

In Table 2 (Appendix F), we report the total time spent in the forward pass
for selecting samples in SIFT as a function of the layer number.

Figure 1: Validation accuracy vs. backpropagation sample complexity. In terms of performance, entropy-based SIFT $>$ loss-based SIFT $>$ baseline. All the accuracies are listed in Table 1 but for quick reference, loss-based and entropy-based SIFT with exit at the first layer are better than the baseline by 1.33% and 1.75%, respectively. For loss-based SIFT, exit at the last layer has the best performance while for entropy-based SIFT, exit at the sixth layer has the best performance.  Figure 2: Validation accuracy vs. wall-clock time. The observations and trends are the same as Figure 1; please see the discussion therein. Algorithm | SIFT Layer # | Val. Accuracy
---|---|---
Baseline | N/A | 0.6466
Loss-based SIFT | 1 | 0.6599
Loss-based SIFT | 2 | 0.6597
Loss-based SIFT | 3 | 0.6604
Loss-based SIFT | 6 | 0.6607
Loss-based SIFT | 12 | 0.6640
Entropy-based SIFT | 1 | 0.6641
Entropy-based SIFT | 2 | 0.6672
Entropy-based SIFT | 3 | 0.6678
Entropy-based SIFT | 6 | 0.6680
Entropy-based SIFT | 12 | 0.6664
Table 1: Validation accuracy at the last step. For loss-based SIFT, exit at
the last layer has the best performance while the performance of other layers
is nearly the same. However for entropy-based SIFT, exit at the sixth layer
has the best performance.

##  8 Conclusion and Limitations

In this work, we theoretically characterized the benefits (as well as
limitations) of the greedy approach of selecting samples with large
approximate losses instead of exact losses. We also showed the promise of
early exiting in speeding up training of a transformer model.

We will mention some limitations in our current work which we hope to explore
and address in future work. As we mentioned, we did not pipeline the early
exit forward propagation step for sample selection with the full forward and
backward propagation steps for model update in our current implementation;
doing so can yield bigger gains. In our current work, we have only shown the
efficacy of SIFT on BERT. In the future, we hope to test SIFT on larger
transformer models and other kinds of models where early exiting is feasible
(such as ResNets). On the theory side, our current convergence result is for
convex functions; we would like to derive a similar result for non-convex
functions too.

## References

  * Alain et al., (2015) Alain, G., Lamb, A., Sankar, C., Courville, A., and Bengio, Y. (2015).  Variance reduction in sgd by distributed importance sampling.  arXiv preprint arXiv:1511.06481.
  * Bengio et al., (2009) Bengio, Y., Louradour, J., Collobert, R., and Weston, J. (2009).  Curriculum learning.  In Proceedings of the 26th annual international conference on machine learning, pages 41–48. ACM.
  * Bertsekas et al., (2011) Bertsekas, D. P. et al. (2011).  Incremental gradient, subgradient, and proximal methods for convex optimization: A survey.  Optimization for Machine Learning, 2010(1-38):3.
  * Borsos et al., (2018) Borsos, Z., Krause, A., and Levy, K. Y. (2018).  Online variance reduction for stochastic optimization.  arXiv preprint arXiv:1802.04715.
  * Bulatov, (2022) Bulatov, Y. (2022).  Mathematics stack exchange: Proving $1-\exp(-4x^{2}/\pi)\geq\text{erf}(x)^{2}$.  https://math.stackexchange.com/questions/6908/proving-1-exp-4x2-pi-ge-texterfx2.
  * Chuang et al., (2023) Chuang, Y.-S., Xie, Y., Luo, H., Kim, Y., Glass, J., and He, P. (2023).  Dola: Decoding by contrasting layers improves factuality in large language models.  arXiv preprint arXiv: 2309.03883.
  * Fan et al., (2017) Fan, Y., Lyu, S., Ying, Y., and Hu, B. (2017).  Learning with average top-k loss.  In Advances in neural information processing systems, pages 497–505.
  * Gal et al., (2017) Gal, Y., Islam, R., and Ghahramani, Z. (2017).  Deep bayesian active learning with image data.  In International conference on machine learning, pages 1183–1192. PMLR.
  * Gurbuzbalaban et al., (2019) Gurbuzbalaban, M., Ozdaglar, A., and Parrilo, P. A. (2019).  Convergence rate of incremental gradient and incremental newton methods.  SIAM Journal on Optimization, 29(4):2542–2565.
  * Gurbuzbalaban et al., (2022) Gurbuzbalaban, M., Ozdaglar, A., and Parrilo, P. A. (2022).  Why random reshuffling beats stochastic gradient descent.  arXiv preprint arXiv:1510.08560.
  * Haochen and Sra, (2019) Haochen, J. and Sra, S. (2019).  Random shuffling beats sgd after finite epochs.  In International Conference on Machine Learning, pages 2624–2633. PMLR.
  * Jiang et al., (2019) Jiang, A. H., Wong, D. L.-K., Zhou, G., Andersen, D. G., Dean, J., Ganger, G. R., Joshi, G., Kaminksy, M., Kozuch, M., Lipton, Z. C., et al. (2019).  Accelerating deep learning by focusing on the biggest losers.  arXiv preprint arXiv:1910.00762.
  * Jiang et al., (2017) Jiang, L., Zhou, Z., Leung, T., Li, L.-J., and Fei-Fei, L. (2017).  Mentornet: Learning data-driven curriculum for very deep neural networks on corrupted labels.  arXiv preprint arXiv:1712.05055.
  * Johnson and Guestrin, (2018) Johnson, T. B. and Guestrin, C. (2018).  Training deep models faster with robust, approximate importance sampling.  In Advances in Neural Information Processing Systems, pages 7265–7275.
  * Katharopoulos and Fleuret, (2017) Katharopoulos, A. and Fleuret, F. (2017).  Biased importance sampling for deep neural network training.  arXiv preprint arXiv:1706.00043.
  * Katharopoulos and Fleuret, (2018) Katharopoulos, A. and Fleuret, F. (2018).  Not all samples are created equal: Deep learning with importance sampling.  arXiv preprint arXiv:1803.00942.
  * Kawaguchi, (2016) Kawaguchi, K. (2016).  Deep learning without poor local minima.  Advances in neural information processing systems, 29.
  * Kawaguchi and Lu, (2019) Kawaguchi, K. and Lu, H. (2019).  Ordered sgd: A new stochastic optimization framework for empirical risk minimization.  arXiv preprint arXiv:1907.04371.
  * Kim and Choi, (2018) Kim, T.-H. and Choi, J. (2018).  Screenernet: Learning self-paced curriculum for deep neural networks.  arXiv preprint arXiv:1801.00904.
  * Kumar et al., (2010) Kumar, M. P., Packer, B., and Koller, D. (2010).  Self-paced learning for latent variable models.  In Advances in Neural Information Processing Systems, pages 1189–1197.
  * Loshchilov and Hutter, (2015) Loshchilov, I. and Hutter, F. (2015).  Online batch selection for faster training of neural networks.  arXiv preprint arXiv:1511.06343.
  * Lu et al., (2022) Lu, Y., Meng, S. Y., and De Sa, C. (2022).  A general analysis of example-selection for stochastic gradient descent.  In International Conference on Learning Representations (ICLR), volume 10.
  * Mishchenko et al., (2020) Mishchenko, K., Khaled, A., and Richtárik, P. (2020).  Random reshuffling: Simple analysis with vast improvements.  Advances in Neural Information Processing Systems, 33:17309–17320.
  * Mohtashami et al., (2022) Mohtashami, A., Stich, S., and Jaggi, M. (2022).  Characterizing & finding good data orderings for fast convergence of sequential gradient methods.  arXiv preprint arXiv:2202.01838.
  * Needell et al., (2014) Needell, D., Ward, R., and Srebro, N. (2014).  Stochastic gradient descent, weighted sampling, and the randomized kaczmarz algorithm.  In Advances in neural information processing systems, pages 1017–1025.
  * Nesterov et al., (2018) Nesterov, Y. et al. (2018).  Lectures on convex optimization, volume 137.  Springer.
  * Pennington and Bahri, (2017) Pennington, J. and Bahri, Y. (2017).  Geometry of neural network loss surfaces via random matrix theory.  In International conference on machine learning, pages 2798–2806. PMLR.
  * Recht and Ré, (2012) Recht, B. and Ré, C. (2012).  Toward a noncommutative arithmetic-geometric mean inequality: Conjectures, case-studies, and consequences.  In Conference on Learning Theory, pages 11–1. JMLR Workshop and Conference Proceedings.
  * Ren et al., (2021) Ren, P., Xiao, Y., Chang, X., Huang, P.-Y., Li, Z., Gupta, B. B., Chen, X., and Wang, X. (2021).  A survey of deep active learning.  ACM computing surveys (CSUR), 54(9):1–40.
  * Rotem et al., (2023) Rotem, D., Hassid, M., Mamou, J., and Schwartz, R. (2023).  Finding the sweet spot: Analysis and improvement of adaptive inference in low resource settings.  Annual Meeting of the Association for Computational Linguistics.
  * Safran and Shamir, (2020) Safran, I. and Shamir, O. (2020).  How good is sgd with random shuffling?  In Conference on Learning Theory, pages 3250–3284. PMLR.
  * Schuster et al., (2022) Schuster, T., Fisch, A., Gupta, J., Dehghani, M., Bahri, D., Tran, V., Tay, Y., and Metzler, D. (2022).  Confident adaptive language modeling.  Advances in Neural Information Processing Systems, 35:17456–17472.
  * Schwartz et al., (2020) Schwartz, R., Stanovsky, G., Swayamdipta, S., Dodge, J., and Smith, N. A. (2020).  The right tool for the job: Matching model and instance complexities.  In Jurafsky, D., Chai, J., Schluter, N., and Tetreault, J., editors, Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6640–6651, Online. Association for Computational Linguistics.
  * Shrivastava et al., (2016) Shrivastava, A., Gupta, A., and Girshick, R. (2016).  Training region-based object detectors with online hard example mining.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 761–769.
  * Stich et al., (2017) Stich, S. U., Raj, A., and Jaggi, M. (2017).  Safe adaptive importance sampling.  In Advances in Neural Information Processing Systems, pages 4381–4391.
  * Teerapittayanon et al., (2016) Teerapittayanon, S., McDanel, B., and Kung, H. (2016).  Branchynet: Fast inference via early exiting from deep neural networks.  In 2016 23rd International Conference on Pattern Recognition (ICPR), pages 2464–2469.
  * Tsvetkov et al., (2016) Tsvetkov, Y., Faruqui, M., Ling, W., MacWhinney, B., and Dyer, C. (2016).  Learning the curriculum with bayesian optimization for task-specific word representation learning.  arXiv preprint arXiv:1605.03852.
  * Xin et al., (2020) Xin, J., Tang, R., Lee, J., Yu, Y., and Lin, J. (2020).  DeeBERT: Dynamic early exiting for accelerating BERT inference.  In Jurafsky, D., Chai, J., Schluter, N., and Tetreault, J., editors, Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2246–2251, Online. Association for Computational Linguistics.
  * Xin et al., (2021) Xin, J., Tang, R., Yu, Y., and Lin, J. (2021).  BERxiT: Early exiting for BERT with better fine-tuning and extension to regression.  In Merlo, P., Tiedemann, J., and Tsarfaty, R., editors, Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, pages 91–104, Online. Association for Computational Linguistics.
  * Zhang et al., (2019) Zhang, J., Yu, H.-f., and Dhillon, I. S. (2019).  Autoassist: A framework to accelerate training of deep neural networks.  arXiv preprint arXiv:1905.03381.
  * Zhao and Zhang, (2015) Zhao, P. and Zhang, T. (2015).  Stochastic optimization with importance sampling for regularized loss minimization.  In international conference on machine learning, pages 1–9.
  * Zhu, (2021) Zhu, W. (2021).  LeeBERT: Learned early exit for BERT with cross-level optimization.  In Zong, C., Xia, F., Li, W., and Navigli, R., 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. Association for Computational Linguistics.
  * Zhu et al., (2015) Zhu, Y., Kiros, R., Zemel, R., Salakhutdinov, R., Urtasun, R., Torralba, A., and Fidler, S. (2015).  Aligning books and movies: Towards story-like visual explanations by watching movies and reading books.  In Proceedings of the IEEE international conference on computer vision, pages 19–27.

Appendix

##  Appendix A Some Preliminaries for the Proofs

######  Fact 1 (Nesterov et al., (2018)).

For an $L$-smooth function $g:\mathbb{R}^{p}\xrightarrow{}\mathbb{R}$,
$\|\nabla g(\bm{y})\|_{2}^{2}\leq
2L(g(\bm{y})-\min_{\bm{z}\in\mathbb{R}^{p}}g(\bm{z}))$ $\forall$
$\bm{y}\in\mathbb{R}^{p}$.

######  Fact 2 (Cantelli’s inequality).

For any random variable $Y$ and any $t>0$:

$\mathbb{P}\Big{(}Y-\mathbb{E}[Y]\geq
t\Big{)}\leq\frac{\textup{Var}(Y)}{\textup{Var}(Y)+t^{2}}.$

######  Fact 3 (Bulatov, (2022)).

We have:

$\textup{erf}(t)\leq\sqrt{1-\exp\Big{(}-\frac{4t^{2}}{\pi}\Big{)}}.$

######  Fact 4.

It holds that:

$\textup{erfc}(t)\leq 2\exp\Big{(}-\frac{t^{2}}{2}\Big{)}.$

###### Proof.

Let $Z\sim\mathcal{N}(0,1)$. It can be checked that
$\text{erf}(t)=2\mathbb{P}(Z\in(0,t))$. Thus,

$\text{erfc}(t)=1-\text{erf}(t)=2\Big{(}\frac{1}{2}-\mathbb{P}\big{(}Z\in(0,t)\big{)}\Big{)}=2\mathbb{P}(Z\geq
t).$ (15)

But using the Gaussian tail bound, we have $\mathbb{P}(Z\geq t)\leq
e^{-t^{2}/2}$. Using this in eq. 15 gives us the desired result. ∎

######  Fact 5.

For any $t>0$:

$\int_{0}^{\infty}\exp\Big{(}-\frac{y^{2}}{2t^{2}}\Big{)}dy=\sqrt{\frac{\pi}{2}}t.$

###### Proof.

Let $Z\sim\mathcal{N}(0,t^{2})$. We have:

$\int_{0}^{\infty}\exp\Big{(}-\frac{y^{2}}{2t^{2}}\Big{)}dy=\sqrt{2\pi
t^{2}}\underbrace{\Bigg{(}\frac{1}{\sqrt{2\pi
t^{2}}}\int_{0}^{\infty}\exp\Big{(}-\frac{y^{2}}{2t^{2}}\Big{)}dy\Bigg{)}}_{=\mathbb{P}(Z>0)=\frac{1}{2}}=\sqrt{\frac{\pi}{2}}t.$
(16)

∎

##  Appendix B Proof of Theorem 1

###### Proof.

Consider some $\bm{w}^{*}\in\Phi_{F}$. With a constant step-size, say $\eta$,
we have for any iteration $k$:

$\mathbb{E}_{\mathcal{S}_{k}^{(R)}}\big{[}\|\bm{w}_{k+1}-\bm{w}^{*}\|^{2}\big{]}=\|\bm{w}_{k}-\bm{w}^{*}\|^{2}-2\eta\big{\langle}\mathbb{E}_{\mathcal{S}_{k}^{(R)}}\big{[}\nabla
f(\bm{w}_{k},\widehat{\bm{x}}_{k})\big{]},\bm{w}_{k}-\bm{w}^{*}\big{\rangle}+\eta^{2}\mathbb{E}_{\mathcal{S}_{k}^{(R)}}\big{[}{\|\nabla
f(\bm{w}_{k},\widehat{\bm{x}}_{k})\|^{2}}\big{]}.$ (17)

Using Definition 2 in eq. 17, we get:

$\mathbb{E}_{\mathcal{S}_{k}^{(R)}}\big{[}\|\bm{w}_{k+1}-\bm{w}^{*}\|^{2}\big{]}=\|\bm{w}_{k}-\bm{w}^{*}\|^{2}-2\eta\langle\nabla\widehat{F}_{R}(\bm{w}_{k}),\bm{w}_{k}-\bm{w}^{*}\rangle+\eta^{2}\mathbb{E}_{\mathcal{S}_{k}^{(R)}}\big{[}{\|\nabla
f(\bm{w}_{k},\widehat{\bm{x}}_{k})\|^{2}}\big{]}.$ (18)

Using the convexity of $f(\bm{w},\bm{x})$ w.r.t. $\bm{w}$ (2) and the fact
that pointwise maximum of convex functions is also convex, we conclude that
$\widehat{F}_{R}(\bm{w})$ is convex. Thus:

$\langle\nabla\widehat{F}_{R}(\bm{w}_{k}),\bm{w}_{k}-\bm{w}^{*}\rangle\geq\widehat{F}_{R}(\bm{w}_{k})-\widehat{F}_{R}(\bm{w}^{*})\geq\widehat{F}_{R}(\bm{w}_{k})-\Delta_{R},$
(19)

where the last step follows from the fact that
$\sup_{\bm{w}^{*}\in\Phi_{F}}\widehat{F}_{R}(\bm{w}^{*})\leq\Delta_{R}$.

Further, using 3, 1 and 4, we get:

$\mathbb{E}_{\mathcal{S}_{k}^{(R)}}\big{[}\|\nabla
f(\bm{w}_{k},\widehat{\bm{x}}_{k})\|^{2}\big{]}\leq
2L\mathbb{E}_{\mathcal{S}_{k}^{(R)}}\big{[}f(\bm{w}_{k},\widehat{\bm{x}}_{k})\big{]}\leq
2L\widehat{F}_{R}(\bm{w}_{k}).$ (20)

Plugging eq. 19 and eq. 20 into eq. 17, we get:

$\mathbb{E}_{\mathcal{S}_{k}^{(R)}}\big{[}\|\bm{w}_{k+1}-\bm{w}^{*}\|^{2}\big{]}=\|\bm{w}_{k}-\bm{w}^{*}\|^{2}-2\eta\big{(}1-\eta
L\big{)}\widehat{F}_{R}(\bm{w}_{k})+2\eta\Delta_{R}.$ (21)

Let us impose $\eta<\frac{1}{L}$ so that $1-\eta L>0$. From Definition 2, we
have $\widehat{F}_{R}(\bm{w}_{k})\geq\rho_{R}^{*}F(\bm{w}_{k})$. Using this
above, we get:

$\mathbb{E}_{\mathcal{S}_{k}^{(R)}}\big{[}\|\bm{w}_{k+1}-\bm{w}^{*}\|^{2}\big{]}=\|\bm{w}_{k}-\bm{w}^{*}\|^{2}-2\rho_{R}^{*}\eta\big{(}1-\eta
L\big{)}F(\bm{w}_{k})+2\eta\Delta_{R}.$ (22)

After some rearrangement, we get:

$F(\bm{w}_{k})\leq\frac{\|\bm{w}_{k}-\bm{w}^{*}\|^{2}-\mathbb{E}_{\mathcal{S}_{k}^{(R)}}\big{[}\|\bm{w}_{k+1}-\bm{w}^{*}\|^{2}\big{]}}{2\rho_{R}^{*}\eta\big{(}1-\eta
L\big{)}}+\frac{\Delta_{R}}{\rho_{R}^{*}\big{(}1-\eta L\big{)}}.$ (23)

Next, we sum the above for all $k\in\\{0,\ldots,K-1\\}$ while taking
expectation throughout. After dividing both sides of the resultant inequality
by $\frac{1}{K}$, we get:

$\frac{1}{K}\sum_{k=0}^{K-1}\mathbb{E}[F(\bm{w}_{k})]\leq\frac{\|\bm{w}_{0}-\bm{w}^{*}\|^{2}}{2\rho_{R}^{*}\eta\big{(}1-\eta
L\big{)}K}+\frac{\Delta_{R}}{\rho_{R}^{*}\big{(}1-\eta L\big{)}}.$ (24)

Let $\overline{\bm{w}}_{K}=\frac{1}{K}\sum_{k=0}^{K-1}\bm{w}_{k}$. Applying
Jensen’s inequality, we get
$\mathbb{E}[F(\overline{\bm{w}}_{K})]\leq\frac{1}{K}\sum_{k=0}^{K-1}\mathbb{E}[F(\bm{w}_{k})]$.
Using this in eq. 24 gives us:

$\mathbb{E}\big{[}F\big{(}\overline{\bm{w}}_{K}\big{)}\big{]}\leq\frac{\|\bm{w}_{0}-\bm{w}^{*}\|^{2}}{2\rho_{R}^{*}\eta\big{(}1-\eta
L\big{)}K}+\frac{\Delta_{R}}{\rho_{R}^{*}\big{(}1-\eta L\big{)}}.$ (25)

Observe that this analysis holds for any $\bm{w}^{*}\in\Phi_{F}$, including
the one that is closest to $\bm{w}_{0}$. Using this in the above bound gives
us:

$\mathbb{E}\big{[}F\big{(}\overline{\bm{w}}_{K}\big{)}\big{]}\leq\frac{D_{0}^{2}}{2\rho_{R}^{*}\eta\big{(}1-\eta
L\big{)}K}+\frac{\Delta_{R}}{\rho_{R}^{*}\big{(}1-\eta L\big{)}},$ (26)

where $D_{0}:=\min_{\bm{w}^{*}\in\Phi_{F}}\|\bm{w}_{0}-\bm{w}^{*}\|$. ∎

##  Appendix C Proof of Theorem 4

###### Proof.

For $j\in\\{1,2\\}$, we have
$\widetilde{f}(\bm{w},\bm{x}^{(j)})={f}(\bm{w},\bm{x}^{(j)})\exp\Big{(}\mu(\bm{w})+\sigma{\zeta}(\bm{w},\bm{x}^{(j)})\Big{)}$,
where ${\zeta}(\bm{w},\bm{x}^{(j)})$ is i.i.d. random noise with mean 0 and
variance 1. For brevity of notation, let:

$\mu=\mu(\bm{w})\text{ and
}\widetilde{f}^{(j)}=\widetilde{f}(\bm{w},\bm{x}^{(j)}),\text{
}f^{(j)}={f}(\bm{w},\bm{x}^{(j)})\text{ and
}{\zeta}^{(j)}={\zeta}(\bm{w},\bm{x}^{(j)})\text{ for }j\in\\{1,2\\}.$

Without loss of generality, let $f^{(1)}\geq f^{(2)}$. Let $j^{*}=\text{arg
max}_{j\in[2]}\widetilde{f}^{(j)}$. Let us consider the case of $j^{*}=1$.
This happens when $f^{(2)}\exp(\mu+\sigma{\zeta}^{(2)})\leq
f^{(1)}\exp(\mu+\sigma{\zeta}^{(1)})$; this is equivalent to:

$\sigma({\zeta}^{(2)}-{\zeta}^{(1)})\leq\log\frac{f^{(1)}}{f^{(2)}}.$ (27)

Let $Z={\zeta}^{(2)}-{\zeta}^{(1)}$. Note that $Z$ is random variable with
mean 0 and variance 2. As per the above discussion, we have:

$\mathbb{P}\big{(}j^{*}=1\big{)}=\mathbb{P}\Bigg{(}Z\leq\frac{1}{\sigma}\log\frac{f^{(1)}}{f^{(2)}}\Bigg{)}=1-\mathbb{P}\Bigg{(}Z>\frac{1}{\sigma}\log\frac{f^{(1)}}{f^{(2)}}\Bigg{)}.$
(28)

Thus,

$\mathbb{P}\big{(}j^{*}=2\big{)}=\mathbb{P}\Bigg{(}Z>\frac{1}{\sigma}\log\frac{f^{(1)}}{f^{(2)}}\Bigg{)}.$
(29)

Then, we have:

$\displaystyle\mathbb{E}_{\\{\zeta^{(1)},\zeta^{(2)}\\}}\Big{[}f^{(j^{*})}\Big{]}$
$\displaystyle=f^{(1)}\mathbb{P}\big{(}j^{*}=1\big{)}+f^{(2)}\mathbb{P}\big{(}j^{*}=2\big{)}$
(30)
$\displaystyle=f^{(1)}\Bigg{(}1-\mathbb{P}\Bigg{(}Z>\frac{1}{\sigma}\log\frac{f^{(1)}}{f^{(2)}}\Bigg{)}\Bigg{)}+f^{(2)}\mathbb{P}\Bigg{(}Z>\frac{1}{\sigma}\log\frac{f^{(1)}}{f^{(2)}}\Bigg{)}$
(31)
$\displaystyle=f^{(1)}-\big{(}f^{(1)}-f^{(2)}\big{)}\mathbb{P}\Bigg{(}Z>\frac{1}{\sigma}\log\frac{f^{(1)}}{f^{(2)}}\Bigg{)}.$
(32)

Since $\mathbb{E}[Z]=0$ and $\text{Var}(Z)=2$, using Cantelli’s inequality
(2), we have:

$\mathbb{P}\Bigg{(}Z>\frac{1}{\sigma}\log\frac{f^{(1)}}{f^{(2)}}\Bigg{)}\leq\frac{2}{2+\frac{1}{\sigma^{2}}\log^{2}\frac{f^{(1)}}{f^{(2)}}}.$
(33)

Using this in eq. 32, we get:

$\mathbb{E}_{\\{\zeta^{(1)},\zeta^{(2)}\\}}\Big{[}f^{(j^{*})}\Big{]}\geq
f^{(1)}-\big{(}f^{(1)}-f^{(2)}\big{)}\Bigg{(}\frac{2}{2+\frac{1}{\sigma^{2}}\log^{2}\frac{f^{(1)}}{f^{(2)}}}\Bigg{)}.$
(34)

Let us define the function
$h(t;\sigma):=\frac{2(1-e^{-t})}{2+\frac{t^{2}}{\sigma^{2}}}$ for $t\geq 0$.
Then, note that:

$\big{(}f^{(1)}-f^{(2)}\big{)}\Bigg{(}\frac{2}{2+\frac{1}{\sigma^{2}}\log^{2}\frac{f^{(1)}}{f^{(2)}}}\Bigg{)}=f^{(1)}\times
h\Bigg{(}\log\frac{f^{(1)}}{f^{(2)}};\sigma\Bigg{)}.$

From Lemma 1, we have $h(t;\sigma)\leq
0.72\big{(}1-e^{-\sqrt{2}\sigma}\big{)}$ for all $t\geq 0$, when
$\sigma\leq\frac{1}{2\sqrt{2}}$. Using this above gives us:

$\big{(}f^{(1)}-f^{(2)}\big{)}\Bigg{(}\frac{2}{2+\frac{1}{\sigma^{2}}\log^{2}\frac{f^{(1)}}{f^{(2)}}}\Bigg{)}\leq
0.72\Big{(}1-e^{-\sqrt{2}\sigma}\Big{)}f^{(1)},$ (35)

when $\sigma\leq\frac{1}{2\sqrt{2}}$. Plugging this into eq. 34 gives us:

$\mathbb{E}_{\\{\zeta^{(1)},\zeta^{(2)}\\}}\Big{[}f^{(j^{*})}\Big{]}\geq\Big{(}1-0.72\Big{(}1-e^{-\sqrt{2}\sigma}\Big{)}\Big{)}f^{(1)}.$
(36)

Rewriting the above in terms of the full notation, we obtain:

$\mathbb{E}_{\\{{\zeta}(\bm{w},\bm{x}^{(j)})\\}_{j=1}^{2}}\Big{[}f\big{(}\bm{w},\bm{x}^{(j^{*})}\big{)}\Big{|}j^{*}=\text{arg
max}_{j\in[2]}\widetilde{f}(\bm{w},\bm{x}^{(j)})\Big{]}\geq\Big{(}1-0.72\Big{(}1-e^{-\sqrt{2}\sigma}\Big{)}\Big{)}\max_{j\in[2]}{f}(\bm{w},\bm{x}^{(j)}).$
(37)

Thus:

$\widehat{F}_{2,\text{approx}}(\bm{w})\geq\Big{(}1-0.72\Big{(}1-e^{-\sqrt{2}\sigma}\Big{)}\Big{)}\mathbb{E}_{\\{\bm{x}^{(j)}\\}_{j=1}^{2}}\Big{[}\max_{j\in[2]}{f}(\bm{w},\bm{x}^{(j)})\Big{]}=\Big{(}1-0.72\Big{(}1-e^{-\sqrt{2}\sigma}\Big{)}\Big{)}\widehat{F}_{2}(\bm{w}).$
(38)

Hence, we also have:

$\rho_{2,\text{approx}}(\bm{w})=\frac{\widehat{F}_{2,\text{approx}}(\bm{w})}{F(\bm{w})}\geq\Big{(}1-0.72\Big{(}1-e^{-\sqrt{2}\sigma}\Big{)}\Big{)}\frac{\widehat{F}_{2}(\bm{w})}{F(\bm{w})}=\Big{(}1-0.72\Big{(}1-e^{-\sqrt{2}\sigma}\Big{)}\Big{)}\rho_{2}(\bm{w}),$
(39)

and:

$\rho_{2,\text{approx}}^{*}=\inf_{\bm{w}\notin\Phi_{F}}\rho_{2,\text{approx}}(\bm{w})\geq\Big{(}1-0.72\Big{(}1-e^{-\sqrt{2}\sigma}\Big{)}\Big{)}\inf_{\bm{w}\notin\Phi_{F}}\rho_{2}(\bm{w})=\Big{(}1-0.72\Big{(}1-e^{-\sqrt{2}\sigma}\Big{)}\Big{)}\rho_{2}^{*}.$
(40)

∎

######  Lemma 1.

Consider the function
$h(t;\sigma):=\frac{2(1-e^{-t})}{\big{(}2+\frac{t^{2}}{\sigma^{2}}\big{)}}$
for $t\geq 0$ and $\sigma\leq\frac{1}{2\sqrt{2}}$. Then, we have:

$\max_{t\geq 0}h(t;\sigma)\leq 0.72\Big{(}1-e^{-\sqrt{2}\sigma}\Big{)}.$

###### Proof.

Let $t^{*}=\text{arg max}_{t\geq 0}h(t;\sigma)$. Setting
$\frac{{d}h}{{d}t}\big{|}_{t=t^{*}}=0$, we obtain the following equation:

$e^{t^{*}}=1+\frac{\sigma^{2}}{t^{*}}+\frac{t^{*}}{2}.$ (41)

Using the series expansion of the exponential function above, we get:

$1+\sum_{j=1}^{\infty}\frac{(t^{*})^{j}}{j!}=1+\frac{\sigma^{2}}{t^{*}}+\frac{t^{*}}{2}\implies(t^{*})^{2}\underbrace{\Big{(}1+2\sum_{j=1}^{\infty}\frac{(t^{*})^{j}}{(j+1)!}\Big{)}}_{:=\nu(t^{*})}=2\sigma^{2}.$
(42)

Since $\nu(t^{*})\geq 1$, we conclude that:

$t^{*}\leq\sqrt{2}\sigma.$ (43)

Further, since $\sigma\leq\frac{1}{2\sqrt{2}}$, we also have
$t^{*}\leq\frac{1}{2}$. Note that $\nu(t^{*})$ is an increasing function of
$t^{*}$. So $\nu(t^{*})\leq\nu\big{(}\frac{1}{2}\big{)}$. Also, by using the
series expansion of the exponential function, it can be verified that
$\nu(t^{*})=\frac{2(e^{t^{*}}-1)}{t^{*}}-1$. Thus, we have:

$\nu(t^{*})\leq\nu\Big{(}\frac{1}{2}\Big{)}\leq 4e^{1/2}-5.$ (44)

Using this in eq. 42, we get:

$t^{*}\geq\frac{\sqrt{2}\sigma}{\sqrt{4e^{1/2}-5}}\geq
0.62\big{(}\sqrt{2}\sigma\big{)}.$ (45)

Combining the bounds of eq. 43 and eq. 45, we deduce that
$t^{*}=c\big{(}\sqrt{2}\sigma\big{)}$, where $c\in[0.62,1]$. Thus,
$(1-e^{-t^{*}})\leq 1-e^{-\sqrt{2}\sigma}$ and
$\frac{2}{2+\frac{t^{2}}{\sigma^{2}}}\leq\frac{2}{2(1+0.62^{2})}\leq 0.72$.
Using all of this, we obtain:

$\max_{t\geq 0}h(t;\sigma)=h(t^{*};\sigma)\leq
0.72\big{(}1-e^{-\sqrt{2}\sigma}\big{)}.$ (46)

∎

##  Appendix D Proof of Theorem 5

###### Proof.

Per 7,
$\mathcal{M}(\bm{w},\bm{x})-\mathcal{M}(\bm{w}^{*},\bm{x})\underset{\text{iid}}{\sim}\mathcal{N}\big{(}\varepsilon(\bm{w}),\delta^{2}(\bm{w})\big{)}$
for $\bm{w}\neq\bm{w}^{*}$. Clearly,

$F(\bm{w})=\mathbb{E}_{\bm{x}}\big{[}f(\bm{w},\bm{x})\big{]}=\varepsilon^{2}(\bm{w})+\delta^{2}(\bm{w}),$
(47)

and

$\widehat{F}_{R}(\bm{w})=\mathbb{E}_{\\{\bm{x}^{(1)},\ldots,\bm{x}^{(R)}\\}}\Bigg{[}\max_{\bm{x}\in\\{\bm{x}^{(1)},\ldots,\bm{x}^{(R)}\\}}f(\bm{w},\bm{x})\Bigg{]}.$
(48)

For conciseness, we shall denote $\varepsilon(\bm{w})$ and $\delta(\bm{w})$ by
just $\varepsilon$ and $\delta$, respectively. Also, let
$Z_{i}=\mathcal{M}(\bm{w},\bm{x}^{(i)})-\mathcal{M}(\bm{w}^{*},\bm{x}^{(i)})$
for $i\in[R]$. As per our concise notation, note that each
$Z_{i}\underset{\text{iid}}{\sim}\mathcal{N}(\varepsilon,\delta^{2})$. Hence:

$\widehat{F}_{R}(\bm{w})=\mathbb{E}\Big{[}\max_{i\in[R]}Z_{i}^{2}\Big{]}.$
(49)

We shall obtain a lower bound for
$\mathbb{E}\big{[}\max_{i\in[R]}Z_{i}^{2}\big{]}$. Let
$Y\sim\mathcal{N}(0,1)$. For $t>0$, we have:

$\displaystyle\mathbb{P}\Big{(}\max_{i\in[R]}(Z_{i}-\varepsilon)\leq
t\Big{)}=\mathbb{P}\Big{(}\cap_{i\in[R]}(Z_{i}-\varepsilon)\leq t\Big{)}$
$\displaystyle=\Big{(}\mathbb{P}\big{(}Y\leq{t}/{\delta}\big{)}\Big{)}^{R}$
(50)
$\displaystyle=\Big{(}\frac{1}{2}+\frac{1}{2}\text{erf}\Big{(}\frac{t}{\sqrt{2}\delta}\Big{)}\Big{)}^{R}$
(51)
$\displaystyle\leq\Bigg{(}\frac{1}{2}+\frac{1}{2}\sqrt{1-\exp\Big{(}-\frac{2t^{2}}{\pi\delta^{2}}\Big{)}}\Bigg{)}^{R}$
(52)
$\displaystyle\leq\Bigg{(}\frac{1}{2}+\frac{1}{2}\Bigg{(}1-\frac{1}{2}\exp\Big{(}-\frac{2t^{2}}{\pi\delta^{2}}\Big{)}\Bigg{)}\Bigg{)}^{R}$
(53)
$\displaystyle=\Bigg{(}1-\frac{1}{4}\exp\Big{(}-\frac{2t^{2}}{\pi\delta^{2}}\Big{)}\Bigg{)}^{R}$
(54)
$\displaystyle\leq\exp\Big{(}-\frac{R}{4}\exp\Big{(}-\frac{2t^{2}}{\pi\delta^{2}}\Big{)}\Big{)}.$
(55)

In eq. 51, the error function is as defined in eq. 2. Equation 52 follows from
3. Equation 53 is obtained by using the fact that $\sqrt{1-a}\leq
1-\frac{a}{2}$ for all $a\in[0,1]$. Equation 55 follows from the fact that
$1-a\leq e^{-a}$ for all $a\in\mathbb{R}$. So,
$\max_{i\in[R]}(Z_{i}-\varepsilon)\leq t$ with a probability of at most
$\exp\Big{(}-\frac{R}{4}\exp\Big{(}-\frac{2t^{2}}{\pi\delta^{2}}\Big{)}\Big{)}$.
Let us choose $t=\delta\sqrt{\frac{\pi}{2}\log\frac{R}{4\log R}}$. With this
choice, we get:

$\max_{i\in[R]}Z_{i}\geq\varepsilon+\delta\sqrt{\frac{\pi}{2}\log\frac{R}{4\log
R}}\text{ }\text{ w.p. }\geq 1-\frac{1}{R}.$

Thus,

$\mathbb{E}\Big{[}\max_{i\in[R]}Z_{i}^{2}\Big{]}\geq\Bigg{(}\varepsilon+\delta\sqrt{\frac{\pi}{2}\log\frac{R}{4\log
R}}\Bigg{)}^{2}\Big{(}1-\frac{1}{R}\Big{)}\\\
=\Bigg{(}\varepsilon^{2}+{\frac{\pi}{2}\log\frac{R}{4\log
R}}\delta^{2}+\sqrt{2\pi\log\frac{R}{4\log
R}}\varepsilon\delta\Bigg{)}\Big{(}1-\frac{1}{R}\Big{)}.$ (56)

Plugging this into eq. 49 and replacing $\varepsilon$ and $\delta$ with their
complete notations, i.e., $\varepsilon(\bm{w})$ and $\delta(\bm{w})$, we get:

$\widehat{F}_{R}(\bm{w})\geq\Bigg{(}\varepsilon^{2}(\bm{w})+\Big{(}{\frac{\pi}{2}\log\frac{R}{4\log
R}}\Big{)}\delta^{2}(\bm{w})+\sqrt{2\pi\log\frac{R}{4\log
R}}\varepsilon(\bm{w})\delta(\bm{w})\Bigg{)}\Big{(}1-\frac{1}{R}\Big{)}.$ (57)

Hence:

$\rho_{R}(\bm{w})=\frac{\widehat{F}_{R}(\bm{w})}{F(\bm{w})}\geq\frac{\Big{(}\varepsilon^{2}(\bm{w})+\Big{(}{\frac{\pi}{2}\log\frac{R}{4\log
R}}\Big{)}\delta^{2}(\bm{w})+\sqrt{2\pi\log\frac{R}{4\log
R}}\varepsilon(\bm{w})\delta(\bm{w})\Big{)}\Big{(}1-\frac{1}{R}\Big{)}}{\varepsilon^{2}(\bm{w})+\delta^{2}(\bm{w})}.$
(58)

∎

##  Appendix E Proof of Theorem 6

###### Proof.

For $i\in\\{1,2\\}$, we have:

$\hat{y}_{j}^{(i)}=\text{sig}\big{(}\bm{\theta}^{\top}\bm{A}_{j}\bm{x}^{(i)}\big{)}\text{
and
}\hat{y}^{(i)}=\hat{y}_{k}^{(i)}=\text{sig}\big{(}\bm{\theta}^{\top}\bm{B}_{j}\bm{A}_{j}\bm{x}^{(i)}\big{)}.$
(59)

For the i.i.d. samples $\bm{x}^{(1)}$ and $\bm{x}^{(2)}$ with ground truth
labels $y^{(1)}$ and $y^{(2)}$, recall that $\ell_{j}^{(1)}$ and
$\ell_{j}^{(2)}$ are the corresponding cross-entropy losses of the early
predictions at the $j^{\text{th}}$ layer, i.e.,

$\ell_{j}^{(i)}=-y^{(i)}\log\big{(}\hat{y}_{j}^{(i)}\big{)}-(1-y^{(i)})\log\big{(}1-\hat{y}_{j}^{(i)}\big{)}\text{
for }i\in\\{1,2\\}.$ (60)

We are interested in:

$p_{j}:=\mathbb{P}_{\bm{x}^{(1)},\bm{x}^{(2)}}\Big{(}\text{arg
max}_{i\in[1,2]}\ell_{j}^{(i)}=\text{arg
max}_{i\in[1,2]}\ell_{k}^{(i)}\Big{)}.$

Using symmetry, we get:

$p_{j}:=\mathbb{P}_{\bm{x}^{(1)},\bm{x}^{(2)}}\Big{(}\ell_{j}^{(1)}\geq\ell_{j}^{(2)}\Big{|}\ell_{k}^{(1)}\geq\ell_{k}^{(2)}\Big{)}.$
(61)

Note that $p_{k}=1$. Henceforth, we shall drop the subscript
$\bm{x}^{(1)},\bm{x}^{(2)}$ for conciseness. Let $\bar{y}^{(1)}=2y^{(1)}-1$
and $\bar{y}^{(2)}=2y^{(2)}-1$ be the centered ground truth labels (i.e.,
$\in\\{-1,1\\}$) of $\bm{x}^{(1)}$ and $\bm{x}^{(2)}$, respectively. It can be
verified that:

$\ell_{j}^{(1)}\geq\ell_{j}^{(2)}\iff\bm{\theta}^{\top}\bm{A}_{j}\Big{(}\bar{y}^{(1)}\bm{x}^{(1)}-\bar{y}^{(2)}\bm{x}^{(2)}\Big{)}\geq
0\text{ and
}\ell_{k}^{(1)}\geq\ell_{k}^{(2)}\iff\bm{\theta}^{\top}\bm{B}_{j}\bm{A}_{j}\Big{(}\bar{y}^{(1)}\bm{x}^{(1)}-\bar{y}^{(2)}\bm{x}^{(2)}\Big{)}\geq
0.$ (62)

Let $\bm{z}:=\bar{y}^{(1)}\bm{x}^{(1)}-\bar{y}^{(2)}\bm{x}^{(2)}$. Since
$\bm{x}^{(1)},\bm{x}^{(2)}\underset{\text{iid}}{\sim}\mathcal{N}(\vec{0}_{d},\bm{\text{I}}_{d})$,
we have that $\bm{z}\sim\mathcal{N}(\vec{0}_{d},2\bm{\text{I}}_{d})$.

Using eq. 62 and the definition of $\bm{z}$ followed by the application of
Bayes’ theorem in Equation 61, we obtain:

$p_{j}=\mathbb{P}\Big{(}\bm{\theta}^{\top}\bm{A}_{j}\bm{z}\geq
0\Big{|}\bm{\theta}^{\top}\bm{B}_{j}\bm{A}_{j}\bm{z}\geq
0\Big{)}=\frac{\mathbb{P}\big{(}\bm{\theta}^{\top}\bm{A}_{j}\bm{z}\geq
0,\bm{\theta}^{\top}\bm{B}_{j}\bm{A}_{j}\bm{z}\geq
0\big{)}}{\mathbb{P}\big{(}\bm{\theta}^{\top}\bm{B}_{j}\bm{A}_{j}\bm{z}\geq
0\big{)}}.$ (63)

Since $\bm{z}\sim\mathcal{N}(\vec{0}_{d},2\bm{\text{I}}_{d})$,
$\mathbb{P}\big{(}\bm{\theta}^{\top}\bm{B}_{j}\bm{A}_{j}\bm{z}\geq
0\big{)}=\frac{1}{2}$. Using this above, we get:

$p_{j}=2\mathbb{P}\Big{(}\bm{\theta}^{\top}\bm{A}_{j}\bm{z}\geq
0,\bm{\theta}^{\top}\bm{B}_{j}\bm{A}_{j}\bm{z}\geq 0\Big{)}.$ (64)

For ease of notation, let $u_{1}=\bm{\theta}^{\top}\bm{A}_{j}\bm{z}$ and
$u_{2}=\bm{\theta}^{\top}\bm{B}_{j}\bm{A}_{j}\bm{z}$ and
$\bm{u}=\begin{bmatrix}u_{1}\\\ u_{2}\end{bmatrix}$. Note $\bm{u}$ is a
multivariate Gaussian random variable with
$\mathbb{E}[\bm{u}]=\begin{bmatrix}0\\\ 0\end{bmatrix}$ and

$\mathbb{E}[\bm{u}\bm{u}^{\top}]:=\bm{\Sigma}=2\begin{bmatrix}\|\bm{A}_{j}^{\top}\bm{\theta}\|_{2}^{2}&\big{\langle}\bm{A}_{j}^{\top}\bm{\theta},\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}\big{\rangle}\\\
\big{\langle}\bm{A}_{j}^{\top}\bm{\theta},\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}\big{\rangle}&\|\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}\|_{2}^{2}\end{bmatrix}.$
(65)

Let $\alpha_{1}=\|\bm{A}_{j}^{\top}\bm{\theta}\|_{2}$,
$\alpha_{2}=\|\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}\|_{2}$ and

$\beta=\frac{\langle\bm{A}_{j}^{\top}\bm{\theta},\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}\rangle}{\|\bm{A}_{j}^{\top}\bm{\theta}\|_{2}\|\bm{A}_{j}^{\top}\bm{B}_{j}^{\top}\bm{\theta}\|_{2}}.$

In the theorem statement, we shall denote $\beta$ by $\beta_{j}$ to indicate
the dependence on the layer number $j$; we drop the subscript $j$ here for
conciseness. With this notation, we have:

$\bm{\Sigma}=2\begin{bmatrix}\alpha_{1}^{2}&\beta\alpha_{1}\alpha_{2}\\\
\beta\alpha_{1}\alpha_{2}&\alpha_{2}^{2}\end{bmatrix}.$ (66)

The pdf of $\bm{u}$ is therefore:

$\displaystyle\phi(\bm{u})$
$\displaystyle=\frac{1}{2\pi\sqrt{\text{det}(\bm{\Sigma})}}\exp\Bigg{(}-\frac{1}{2}\bm{u}^{\top}\bm{\Sigma}^{-1}\bm{u}\Bigg{)}$
(67)
$\displaystyle=\frac{1}{4\pi\alpha_{1}\alpha_{2}\sqrt{1-\beta^{2}}}\exp\Bigg{(}-\frac{1}{4(1-\beta^{2})}\Big{(}\frac{u_{1}^{2}}{\alpha_{1}^{2}}-\frac{2\beta
u_{1}u_{2}}{\alpha_{1}\alpha_{2}}+\frac{u_{2}^{2}}{\alpha_{2}^{2}}\Big{)}\Bigg{)}.$
(68)

Using this in eq. 64, we get:

$\displaystyle p_{j}$
$\displaystyle=2\int_{0}^{\infty}\int_{0}^{\infty}\phi(\bm{u}){d}u_{1}{d}u_{2}$
(69)
$\displaystyle=\frac{1}{2\pi\alpha_{1}\alpha_{2}\sqrt{1-\beta^{2}}}\int_{0}^{\infty}\int_{0}^{\infty}\exp\Bigg{(}-\frac{1}{4(1-\beta^{2})}\Big{(}\frac{u_{1}^{2}}{\alpha_{1}^{2}}-\frac{2\beta
u_{1}u_{2}}{\alpha_{1}\alpha_{2}}+\frac{u_{2}^{2}}{\alpha_{2}^{2}}\Big{)}\Bigg{)}{d}u_{1}{d}u_{2}$
(70)
$\displaystyle=\frac{1}{2\pi\alpha_{1}\alpha_{2}\sqrt{1-\beta^{2}}}\int_{0}^{\infty}\exp\Big{(}-\frac{u_{1}^{2}}{4\alpha_{1}^{2}}\Big{)}\Bigg{(}\int_{0}^{\infty}\exp\Bigg{(}-\frac{\big{(}u_{2}-\big{(}\frac{\beta\alpha_{2}}{\alpha_{1}}\big{)}u_{1}\big{)}^{2}}{4\alpha_{2}^{2}(1-\beta^{2})}\Bigg{)}du_{2}\Bigg{)}du_{1}.$
(71)

With some simple change of variables and some simplification, the above
equation becomes:

$\displaystyle p_{j}$
$\displaystyle=\frac{1}{\pi\alpha_{1}}\int_{0}^{\infty}\exp\Big{(}-\frac{u_{1}^{2}}{4\alpha_{1}^{2}}\Big{)}\Bigg{(}\int_{-\frac{\beta
u_{1}}{2\sqrt{1-\beta^{2}}\alpha_{1}}}^{\infty}\exp(-t^{2})dt\Bigg{)}du_{1}$
(72)
$\displaystyle=\frac{1}{2\sqrt{\pi}\alpha_{1}}\int_{0}^{\infty}\exp\Big{(}-\frac{u_{1}^{2}}{4\alpha_{1}^{2}}\Big{)}\Bigg{(}\frac{2}{\sqrt{\pi}}\int_{0}^{\frac{\beta
u_{1}}{2\sqrt{1-\beta^{2}}\alpha_{1}}}\exp(-t^{2})dt+\frac{2}{\sqrt{\pi}}\int_{0}^{\infty}\exp(-t^{2})dt\Bigg{)}du_{1}$
(73)
$\displaystyle=\frac{1}{2\sqrt{\pi}\alpha_{1}}\int_{0}^{\infty}\exp\Big{(}-\frac{u_{1}^{2}}{4\alpha_{1}^{2}}\Big{)}\Bigg{(}\text{erf}\Big{(}\frac{\beta
u_{1}}{2\sqrt{1-\beta^{2}}\alpha_{1}}\Big{)}+\underbrace{\lim_{y\to\infty}\text{erf}(y)}_{=1}\Bigg{)}du_{1}$
(74)
$\displaystyle=\frac{1}{2\sqrt{\pi}\alpha_{1}}\int_{0}^{\infty}\exp\Big{(}-\frac{u_{1}^{2}}{4\alpha_{1}^{2}}\Big{)}\text{erf}\Big{(}\frac{\beta
u_{1}}{2\sqrt{1-\beta^{2}}\alpha_{1}}\Big{)}du_{1}+\underbrace{\frac{1}{2\sqrt{\pi}\alpha_{1}}\int_{0}^{\infty}\exp\Big{(}-\frac{u_{1}^{2}}{4\alpha_{1}^{2}}\Big{)}du_{1}}_{=\frac{1}{2}\text{
using \lx@cref{creftypecap~refnum}{fact5-jan30}}}$ (75)
$\displaystyle=\frac{1}{2\sqrt{\pi}\alpha_{1}}\int_{0}^{\infty}\exp\Big{(}-\frac{u_{1}^{2}}{4\alpha_{1}^{2}}\Big{)}\text{erf}\Big{(}\frac{\beta
u_{1}}{2\sqrt{1-\beta^{2}}\alpha_{1}}\Big{)}du_{1}+\frac{1}{2}.$ (76)

Replacing $({u_{1}}/{\alpha_{1}})$ by $y$ above gives us:

$\displaystyle p_{j}$
$\displaystyle=\frac{1}{2\sqrt{\pi}}\int_{0}^{\infty}\exp\Big{(}-\frac{y^{2}}{4}\Big{)}\text{erf}\Big{(}\frac{\beta
y}{2\sqrt{1-\beta^{2}}}\Big{)}dy+\frac{1}{2}$ (77)
$\displaystyle=\underbrace{\frac{1}{2\sqrt{\pi}}\int_{0}^{\infty}\exp\Big{(}-\frac{y^{2}}{4}\Big{)}dy}_{=\frac{1}{2}\text{
using
\lx@cref{creftypecap~refnum}{fact5-jan30}}}-\frac{1}{2\sqrt{\pi}}\int_{0}^{\infty}\exp\Big{(}-\frac{y^{2}}{4}\Big{)}\text{erfc}\Bigg{(}\frac{\beta
y}{2\sqrt{1-\beta^{2}}}\Bigg{)}dy+\frac{1}{2}$ (78)
$\displaystyle=1-\frac{1}{2\sqrt{\pi}}\int_{0}^{\infty}\exp\Big{(}-\frac{y^{2}}{4}\Big{)}\text{erfc}\Bigg{(}\frac{\beta
y}{2\sqrt{1-\beta^{2}}}\Bigg{)}dy.$ (79)

This is the exact final expression for $p_{j}$.

Using 4 fact above yields:

$\displaystyle p_{j}$ $\displaystyle\geq
1-\frac{1}{\sqrt{\pi}}\int_{0}^{\infty}\exp\Big{(}-\frac{y^{2}}{4}\Big{)}\exp\Big{(}-\frac{\beta^{2}y^{2}}{8{(1-\beta^{2})}}\Big{)}dy$
(80)
$\displaystyle=1-\frac{1}{\sqrt{\pi}}\int_{0}^{\infty}\exp\Bigg{(}-\frac{y^{2}}{8}\Big{(}\frac{2-\beta^{2}}{1-\beta^{2}}\Big{)}\Bigg{)}dy$
(81) $\displaystyle=1-\sqrt{\frac{2-2\beta^{2}}{2-\beta^{2}}}.$ (82)

Equation 82 follows from 5.

Finally, the theorem statement follows by adding the subscript $j$ to $\beta$
in Equation 79 and Equation 82 (recall that we omitted the subscript $j$
earlier for conciseness). ∎

###  E.1 Proof of Corollary 4

###### Proof.

Plugging in $\beta_{j}=1-\tau_{j}$ into the lower bound for $p_{j}$ in Theorem
6, we get:

$\displaystyle p_{j}\geq
1-\sqrt{\frac{2-2(1-\tau_{j})^{2}}{2-(1-\tau_{j})^{2}}}=1-\sqrt{\frac{2\tau_{j}(2-\tau_{j})}{1+\tau_{j}(2-\tau_{j})}}=1-\mathcal{O}(\sqrt{\tau_{j}}),$

for $\tau_{j}\to 0$. ∎

##  Appendix F Remaining Experimental Details

For AdamW, we used the following hyper-parameter values: learning rate = 1e-4,
$\ell_{2}$ weight decay = 0.01, $\beta_{1}$ = 0.9 and $\beta_{2}$ = 0.999. The
learning rate warmup was over the first 0.2% of total steps followed by linear
decay. We used the GELU activation and a dropout probability of 0.1 on all the
layers. The training loss is the sum of the mean masked LM likelihood and the
mean next sentence prediction likelihood.

SIFT Criterion | SIFT Layer # | Sampling Time(hrs)
---|---|---
Loss-based | 1 | 0.8452
Loss-based | 2 | 1.1662
Loss-based | 3 | 1.4764
Loss-based | 6 | 2.5901
Loss-based | 12 | 4.5086
Entropy-based | 1 | 0.9121
Entropy-based | 2 | 1.2181
Entropy-based | 3 | 1.5398
Entropy-based | 6 | 2.6195
Entropy-based | 12 | 4.5148
Table 2: Sampling time of SIFT as a function of early exit layer. For
comparison, the full forward and backward propagation times (for the update)
in all the cases is $\sim$ 3.8 hours and $\sim$ 7.9 hours, respectively.