# FTFT: Efficient and Robust Fine-Tuning by TransFerring Training Dynamics

Yupei Du, Albert Gatt & Dong Nguyen
Utrecht University
Utrecht, the Netherlands
<EMAIL_ADDRESS>

###### Abstract

Despite the massive success of fine-tuning large Pre-trained Language Models
(PLMs) on a wide range of Natural Language Processing (NLP) tasks, they remain
susceptible to out-of-distribution (OOD) and adversarial inputs. Data map (DM)
is a simple yet effective dual-model approach that enhances the robustness of
fine-tuned PLMs, which involves fine-tuning a model on the original training
set (i.e. reference model), selecting a specified fraction of important
training examples according to the training dynamics of the reference model,
and fine-tuning the same model on these selected examples (i.e. main model).
However, it suffers from the drawback of requiring fine-tuning the same model
twice, which is computationally expensive for large models. In this paper, we
first show that 1) training dynamics are highly transferable across different
model sizes and different pre-training methods, and that 2) main models fine-
tuned using DM learn faster than when using conventional Empirical Risk
Minimization (ERM). Building on these observations, we propose a novel fine-
tuning approach based on the DM method: Fine-Tuning by transFerring Training
dynamics (FTFT). Compared with DM, FTFT uses more efficient reference models
and then fine-tunes more capable main models for fewer steps. Our experiments
show that FTFT achieves better generalization robustness than ERM while
spending less than half of the training cost.

##  1 Introduction

Current state-of-the-art performance in Natural Language Processing (NLP) is
dominated by large, pretrained language models (PLMs), which are typically
fine-tuned for downstream tasks. Scaling laws (Kaplan et al., 2020; Hoffmann
et al., 2022) suggest that better downstream performance is achieved with
larger pretrained language models. However, fine-tuning large PLMs is also
more expensive than fine-tuning small PLMs, in terms of both computational
resources and carbon emission (Strubell et al., 2019; Wu et al., 2022).

Moreover, despite making impressive progress on regular benchmarks, many
studies have shown that fine-tuned PLMs lack robustness against out-of-
distribution (OOD) and adversarial inputs. For instance, human annotators can
easily exploit the weaknesses of fine-tuned PLMs to trick these models to
yield incorrect predictions, on tasks such as Natural Language Inference (NLI)
(Nie et al., 2020) and Hate Speech Detection (HSD) (Vidgen et al., 2021b).

The problem of robustness can be mitigated using dual-model approaches. Such
methods first train a reference model to estimate the importance of each
training instance, and then train a main model based on the outputs of the
reference model (Nam et al., 2020; Utama et al., 2020; Sanh et al., 2021;
Karimi Mahabadi et al., 2020; Zhang et al., 2022; Liu et al., 2021). Among
these methods, the approach proposed by Swayamdipta et al. (2020) is
particularly attractive in view of its simplicity and the demonstration that
it consistently improves model performance on OOD test datasets. However, a
major drawback of this approach is that it improves robustness at the expense
of efficiency, which is especially problematic when using large PLMs. In
particular, this approach requires fine-tuning the same model twice. First, a
Data Map (DM) is constructed, based on the training dynamics from an initial
fine-tuning run of the model (i.e., the reference model) on the full dataset.
The DM divides the training data into three subsets which a model finds
ambiguous, hard-to-learn, and easy. The main model is then fine-tuned using
only either the ambiguous or the hard-to-learn subset of data. In the standard
DM approach, the reference model and the main model are the same PLM (e.g.,
$\text{DeBERTaV3}_{\text{Large}}$, He et al., 2023).

In this paper, we jointly address the issues of robustness and efficiency
without sacrificing the properties that make DMs attractive. We achieve this
by exploiting the transferability of training dynamics. Specifically, we study
whether training instances identified by more efficient reference models can
be used to improve the robustness of more capable — i.e. larger — main models.
We make three key observations. First, _training dynamics are highly
transferable_ across different model sizes (e.g.,
$\text{DeBERTaV3}_{\text{Small}}$ as the reference model and
$\text{DeBERTaV3}_{\text{Large}}$ as the main model) and different pre-
training methods (PTMs, e.g., ELECTRA (Clark et al., 2020) as the reference
model and DeBERTaV3 as the main model), Second, the condition for this
transfer to work well is that reference models should be reasonably strong.
Crucially, we identify a key property of effective reference models, namely
that they typically identify higher ratios of training instances as easy
cases, compared to less effective reference models. This observation can help
us inspect whether a reference model would work well without training the main
model. Third, the main model in the DM method learns faster than ERM,
achieving good performance using fewer steps.

Based on our observations, we propose Fine-Tuning by transFerring Training
dynamics (FTFT) to improve both the efficiency and the robustness of fine-
tuning. Concretely, FTFT improves the efficiency of the DM method in two ways.
First, FTFT uses more efficient reference models for identifying the subset of
ambiguous training instances. Second, when using this data subset to fine-tune
a more capable main model, FTFT uses substantially fewer training steps than
ERM fine-tuning. Experimental results on two tasks, NLI and HSD, using two
models, DeBERTaV3 (He et al., 2021) and ELECTRA (Clark et al., 2020), show
that FTFT can achieve better generalization robustness than ERM, while
lowering the training cost by a factor of more than two.

##  2 Background

###  2.1 Improving Model Robustness Using Dual-Model Approaches

Many previous studies have proposed dual-model approaches to improve model
robustness. Nam et al. (2020) first train a reference model using generalized
cross-entropy loss. Then, they train a main model while assigning higher
weights to training instances that are hard for the reference model. Utama et
al. (2020) use the same rationale, except for using a model that is trained on
a random subset of the full training data as the reference model, as well as
adding a confidence regularization mechanism. Sanh et al. (2021) use a
Product-of-Expert (PoE) approach, by first training a reference model with
limited capacity to capture dataset biases, and then training the main model
to avoid these biases. Karimi Mahabadi et al. (2020) also adopt a PoE loss,
but they train both the reference model and the main model in an end-to-end
fashion. Zhang et al. (2022) investigate the idea of regularizing hidden
representations by first training a reference model using ERM, and then using
contrastive learning to make the hidden representations of the same class and
different classes more similar and more separable. Liu et al. (2021) first
train a less capable reference model using heavy regularization and vanilla
SGD, and up-weigh the training instances that the reference model mis-predicts
when training the main model. The DM method proposed by Swayamdipta et al.
(2020) is a more nuanced approach to a similar idea. Instead of relying only
on the correctness of the reference model, they use training dynamics to
categorize training instances, and only use a subset of data to train the main
model. We turn to the details of this method below.

###  2.2 Data Map

Swayamdipta et al. (2020) propose a dual-model approach for improving model
robustness. First, a reference model is trained on the full original dataset.
Then, a Data Map (DM) is built based on the observed training dynamics, by
tracking the prediction probabilities of the true class ($p_{\mathrm{true}}$)
of each training instance across different epochs. Using the DM, training
instances can be categorized into three groups: ambiguous (i.e. the standard
deviation of $p_{\mathrm{true}}$ is in the top $q\%$ of all training
instances); hard-to-learn (i.e. the mean of $p_{\mathrm{true}}$ is at the
bottom $q\%$ of all training instances); and easy (i.e. neither ambiguous nor
hard-to-learn). The threshold $q\%$ is fixed and typically set to $33\%$. It
is worth noting that _ambiguous and hard-to-learn are not mutually exclusive_
: a training instance can be categorized as both hard-to-learn (its average of
$p_{\mathrm{true}}$ is at the bottom $q\%$) and ambiguous (its standard
deviation of $p_{\mathrm{true}}$ is among the top $q\%$). Finally, the main
model is only fine-tuned on the $q\%$ most ambiguous or hard-to-learn
datapoints. Swayamdipta et al. (2020) show that, with a slight loss of In-
Distribution (ID) performance, this approach improves model performance on
challenging Out-Of-Distribution (OOD) datasets. They also observe that
ambiguous data is more beneficial than hard-to-learn data for model training.
We therefore only experiment with ambiguous data.

Swayamdipta et al. (2020) uses the same PLM as both the reference and the main
model. In contrast, Sar-Shalom & Schwartz (2023) show that a DM constructed by
$\text{ELECTRA}_{\text{Large}}$ can be used to improve the robustness of
$\text{DeBERTaV3}_{\text{Large}}$ (He et al., 2021). However, they use a
different approach: instead of training only on the $q\%$ most ambiguous
datapoints, they add $k$ copies of this subset to the original training set;
also, they do not investigate DM transferability further. Inspired by their
observation, we study whether exploiting the transferability of training
dynamics can help improve the efficiency of the DM method, while retaining the
advantage of improved robustness.

##  3 Experimental Setup

In our experiments we study the transferability of training dynamics on two
tasks, Natural Language Inference (NLI) and Hate Speech Detection (HSD). We
consider both ID performance and OOD robustness. To study the transferability
across both model sizes and pretraining methods (PTMs), we mainly focus on two
PTMs, DeBERTa-V3 (He et al., 2023) and ELECTRA (Clark et al., 2020). Each of
these methods makes different model sizes available (small, base, and large).
To further investigate the impact of the capability of the reference model on
transferability, we also include TinyBERT (Turc et al., 2020), a very
efficient but relatively weak PLM. As a baseline, we also experiment with
random DM (i.e., randomly selecting $q\%$ of the training data). Unless
otherwise specified, we set the ratio $q\%$ of all DM methods to $33\%$.

#### Data

To study the impact of our method on model robustness, for each task, besides
the original train set and ID validation set, we also include a few
challenging OOD test sets.

For NLI, we use the MultiNLI dataset (Williams et al., 2018) as the train and
ID validation set, because of its wide coverage of diverse inputs (10 genres).
As OOD test sets, we use two challenging datasets designed to target
weaknesses of models trained on MultiNLI: WANLI (Liu et al., 2022) and
AdversarialNLI (Nie et al., 2020), which consists of three rounds of
adversarial data collection.

For HSD, we use CAD (Vidgen et al., 2021a) as the train and ID validation set.
CAD consists of Reddit posts covering diverse topics and writing styles,
annotated for multiple categories. We frame the task as a binary
classification task (hateful vs. non-hateful). Following Ramponi & Tonelli
(2022), we mark identity-related abuse as hateful and all other categories as
non-hateful. As OOD test sets, we use DynaHate (Vidgen et al., 2021b), which
also contains three rounds of adversarial data collection, as well as a
perturbed version for each round. Conceptualizations of hate speech vary
widely across hate speech datasets. We selected these two datasets, because
they use similar definitions and taxonomies concerning hate speech.

#### Models

We focus on two PLMs: DeBERTaV3 (He et al., 2023) and ELECTRA (Clark et al.,
2020), for two reasons. First, they have both been shown to perform well on
text classification tasks. Second, they are available in different sizes
(small, base, and large), allowing us to experiment with transfer between
different model sizes. We also experiment with TinyBERT to study the impact of
low reference model capacity. We provide the cost (measured in PFLOPs) for
fine-tuning different PLMs in Appendix A.1.111 We report PFLOPs rather than
GPU hours because on our setup, which uses NVIDIA A100 GPUs, we’ve noticed
occasional low GPU utilization, especially during the fine-tuning of smaller
models such as TinyBERT. In such cases, reporting GPU hours would not reflect
the computational costs accurately.

#### Training

For NLI, we train all models for 60k steps ($\sim 5$ epochs) and evaluate
every 4k steps. For TD, we train all models for 6k steps ($\sim 10$ epochs)
and evaluate every 400 steps. For all models, we use four different random
seeds and report their average performance. Full training details (e.g.,
optimization, software and hardware) are included in Appendix A.2.

##  4 Transferability of Training Dynamics

In this section, we study the transferability of training dynamics in the DM
method, i.e., whether we can use different reference and main models while
maintaining the robustness of the main model. Specifically, we study whether
training dynamics are transferable across different model sizes (§4.1, e.g.,
from $\text{DeBERTaV3}_{\text{Small}}$ to $\text{DeBERTaV3}_{\text{Large}}$)
and pretraining methods (§4.2, e.g., from $\text{ELECTRA}_{\text{Large}}$ to
$\text{DeBERTaV3}_{\text{Large}}$), for two reasons. First, transferability
across model sizes can help improve the efficiency of constructing a DM, by
enabling the use of more efficient reference models. Second, transferability
across pretraining methods can help achieve such an efficiency gain, when more
efficient versions of the main model pretraining method are either not
available, or are not able to model the task sufficiently well. Moreover, in
the long run, transferability across different models can help us gain a
better understanding of data ambiguity. Indeed, if the same subset of training
instances is consistently identified as ambiguous by reference models of
different sizes or pretraining methods, this would suggest that ambiguity is
to some extent intrinsic to the data instances rather than being completely
model-dependent.

Our results show training dynamics to be transferable across different model
sizes and pretraining methods, with a few exceptions. To better understand the
conditions for successful transfers, we analyze these few failure cases. We
observe that successful transfers require reference models to be reasonably
strong in identifying relatively easy training instances (§4.3). This
observation offers us a guideline for choosing reference models without
training the main model, which can be computationally prohibitive.

###  4.1 Transferability Across Different Model Sizes

In this section, we study whether smaller and more efficient models can be
used as reference models for training larger main models, without compromising
robustness. For example, when using $\text{DeBERTaV3}_{\text{Large}}$ as the
main model, we investigate whether there is a comparable performance when
using $\text{DeBERTaV3}_{\text{Small}}$ versus using
$\text{DeBERTaV3}_{\text{Large}}$ itself as reference models. Successful
transfers of this type can improve the efficiency of the DM method by enabling
the use of more efficient reference models.

Figure 1:  Consistency across different sizes of DeBERTaV3 on NLI. The numbers
are the fractions (0-1) of the ambiguous training instances shared by two
models. Training dynamics are transferable across different sizes: the
fractions of shared ambiguous instances between models of different sizes are
only slightly smaller than those between models of the same size but different
random seeds (shown as superscript).

We show the results for DeBERTaV3 in Table 1 (across different model sizes),
and make three observations. First, performance is better with larger PLMs as
main models, even when we use ERM. In support of this claim, we note that
$\text{DeBERTaV3}_{\text{Large}}$ performs better than
$\text{DeBERTaV3}_{\text{Small}}$ and $\text{DeBERTaV3}_{\text{Base}}$ when
fine-tuned with ERM or DM. This is the case for all test datasets, except when
the reference model is trained on a random $33\%$ subset of the train data.
Second, by comparing $\text{DeBERTaV3}_{\text{Large}}$ fine-tuned with ERM and
DM, we observe that ERM performs better on ID datasets, while DM outperforms
ERM on OOD datasets. This is consistent with Swayamdipta et al. (2020). Third
and most importantly, _training dynamics are transferable across different
model sizes_ : When using $\text{DeBERTaV3}_{\text{Large}}$ as the main model,
changing the reference model from $\text{DeBERTaV3}_{\text{Large}}$ to
$\text{DeBERTaV3}_{\text{Small}}$ or $\text{DeBERTaV3}_{\text{Base}}$ yields
comparable performance.

To understand this transferability, we analyze to which extent reference
models of different sizes are consistent in identifying ambiguous training
instances. Figure 1 displays the fraction of the ambiguous instances shared by
reference models of different sizes, or of the same size but different random
seeds. Consistent with observation regarding transferability, the fractions of
shared ambiguous instances between different sizes are only slightly smaller
than those between the same size but different random seeds (for comparison,
the fraction between random DMs is expected to be 0.33).

###  4.2 Transferability Across Different Pretraining Methods

Mode | Main Model | Ref. Model | MultiNLI | WANLI | AdversarialNLI
---|---|---|---|---|---
|  |  |  |  | R1 | R2 | R3
ERM | $\text{DeBERTaV3}_{\text{Small}}$ | - | 87.57 | 61.61 | 33.55 | 30.63 | 32.40
ERM | $\text{DeBERTaV3}_{\text{Base}}$ | - | 90.00 | 64.61 | 43.55 | 33.23 | 34.14
ERM | $\text{DeBERTaV3}_{\text{Large}}$ | - | 91.06 | 66.46 | 58.17 | 45.57 | 41.34
DM | $\text{DeBERTaV3}_{\text{Large}}$ | Random | 90.74 | 65.31 | 53.30 | 42.02 | 38.60
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{DeBERTaV3}_{\text{Large}}$ | 90.75 | 66.33 | 59.75 | 45.60 | 41.94
Across different model sizes
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{DeBERTaV3}_{\text{Small}}$ | 90.51 | 66.80 | 59.60 | 45.62 | 42.04
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{DeBERTaV3}_{\text{Base}}$ | 90.74 | 66.61 | 61.42 | 46.73 | 41.58
Across different pretraining methods
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Small}}$ | 90.91 | 62.09 | 49.62 | 38.5 | 35.98
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Base}}$ | 90.63 | 66.58 | 59.77 | 46.25 | 42.29
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Large}}$ | 90.80 | 66.42 | 58.95 | 44.57 | 41.52

Mode Main Model Ref. Model CAD DynaHate-Original DynaHate-Perturb R2 R3 R4 R2
R3 R4 ERM $\text{DeBERTaV3}_{\text{Small}}$ - 76.57 56.89 59.29 63.48 59.55
66.59 61.48 ERM $\text{DeBERTaV3}_{\text{Base}}$ - 78.64 60.53 64.28 68.89
60.81 69.48 63.12 ERM $\text{DeBERTaV3}_{\text{Large}}$ - 81.69 75.44 73.32
76.12 70.62 77.41 68.89 DM $\text{DeBERTaV3}_{\text{Large}}$ Random 76.22
63.38 61.59 71.21 64.05 72.10 62.88 DM $\text{DeBERTaV3}_{\text{Large}}$
$\text{DeBERTaV3}_{\text{Large}}$ 81.58 79.17 76.87 77.73 73.34 76.63 67.54
Across different model sizes DM $\text{DeBERTaV3}_{\text{Large}}$
$\text{DeBERTaV3}_{\text{Small}}$ 81.15 80.68 79.56 79.86 76.47 78.03 70.32 DM
$\text{DeBERTaV3}_{\text{Large}}$ $\text{DeBERTaV3}_{\text{Base}}$ 80.12 80.12
76.34 78.82 74.60 77.81 68.67 Across different pretraining methods DM
$\text{DeBERTaV3}_{\text{Large}}$ $\text{ELECTRA}_{\text{Small}}$ 79.74 78.09
77.40 78.75 75.26 76.79 70.05 DM $\text{DeBERTaV3}_{\text{Large}}$
$\text{ELECTRA}_{\text{Base}}$ 80.37 81.47 78.16 78.38 76.99 78.51 71.01 DM
$\text{DeBERTaV3}_{\text{Large}}$ $\text{ELECTRA}_{\text{Large}}$ 79.48 76.71
75.97 78.58 73.55 77.80 69.81

Table 1:  Transferability across different model sizes and pretraining
methods, using DeBERTaV3 as the main model, on NLI (top, accuracy) and HSD
(bottom, macro-F1). We compare the performance of 1) DeBERTaV3 of different
sizes fine-tuned using ERM, 2) $\text{DeBERTaV3}_{\text{Large}}$ as the main
model, using random DM (random $33\%$ training instances), and
$\text{DeBERTaV3}_{\text{Large}}$ as the reference model (Ref. Model) to
construct DM (original DM), 3) $\text{DeBERTaV3}_{\text{Large}}$ as the main
model, using $\text{DeBERTaV3}_{\text{Small}}$ and
$\text{DeBERTaV3}_{\text{Base}}$ as reference models to construct DM, 4)
$\text{DeBERTaV3}_{\text{Large}}$ as the main model, using ELECTRA of
different sizes as reference models to construct DM. R1-R4 in AdversarialNLI
and DynaHate refer to different rounds of collected data. Training dynamics
are transferrable across different sizes and pretraining methods: DM methods
using different reference model sizes and pretraining methods show comparable
performance.

In this section, we study the transferability of training dynamics across
different pretraining methods, by comparing performance when the reference and
main models are trained using the same pretraining method, or using different
pretraining methods. If transfers of this type are successful, we can improve
the efficiency of DM methods in case there is no version of the main model
pretraining method which is both efficient and effective as the reference
model.

Table 1 displays the results when using $\text{DeBERTaV3}_{\text{Large}}$ as
the main model with different reference models (across different pretraining
methods). _In most settings, training dynamics are transferable across
different pretraining methods_. We note that in most cases
$\text{DeBERTaV3}_{\text{Large}}$ achieves comparable performance when trained
by transferring a DM from reference models of varying sizes of DeBERTaV3 and
ELECTRA. However, such transfers are not always successful. When using
$\text{ELECTRA}_{\text{Small}}$ as the reference model, the performance is
clearly worse on the OOD datasets of NLI than using other PLMs. We hypothesize
that $\text{ELECTRA}_{\text{Small}}$ is not strong enough for constructing
effective DMs for MultiNLI, and will analyze this further in §4.3.

###  4.3 How Efficient Can We Be?

We have shown that training dynamics are transferable across different model
sizes and pretraining methods. In this section, we study the conditions for
successful transfers. Specifically, we focus on answering two questions: 1)
whether we can use very efficient but weak models as reference models; and 2)
what are the differences between effective and ineffective reference models.
Answers to these questions can serve as guidelines for selecting efficient yet
effective reference models.

#### The Use of Efficient but Weak Models

To answer this question, we compare the performance of a wide range of methods
of three types. First, we consider four models fine-tuned with ERM: the small,
base, and large versions of ELECTRA, and TinyBERT (Turc et al., 2020), which
is a very weak but efficient PLM. We use these four models because they are of
different sizes and capabilities
($\text{ELECTRA}_{\text{Large}}>\text{ELECTRA}_{\text{Base}}>\text{ELECTRA}_{\text{Small}}>\text{TinyBERT}$).
Second, we use these ERM models as reference models in the DM method to fine-
tune $\text{DeBERTaV3}_{\text{Large}}$. By using reference models of different
capabilities to fine-tune the same main model, we can inspect the impact of
reference model capability on transferability. Third, we also include the
results with $\text{ELECTRA}_{\text{Large}}$ as the main model, and different
sizes of ELECTRA as reference models. By comparing results with different main
models, we can better understand whether successful transfer of training
dynamics is due to the compatibility between reference and main models, or due
to the capability of the reference model itself. We also include the results
for $\text{ELECTRA}_{\text{Large}}$ fine-tuned with a random DM as a baseline.
We show our results for NLI in Table 2, and include the results for HSD in
Appendix B. We make two observations.

First, poorly performing reference models, such as TinyBERT and
$\text{ELECTRA}_{\text{Small}}$ fine-tuned with ERM, lead to failed transfers.
Also, the worse the reference models perform, the worse the OOD performance of
their main models are, e.g., TinyBERT leads to worse main model performance
than $\text{ELECTRA}_{\text{Small}}$. Moreover, transferring from such models
to main models of neither the same nor different pretraining method is
successful, e.g., transferring from $\text{ELECTRA}_{\text{Small}}$ to
$\text{ELECTRA}_{\text{Large}}$ and $\text{DeBERTaV3}_{\text{Large}}$. This
suggests that the success of transfers mostly depends on the reference model,
rather than the compatibility between the reference and the main model.
Second, using weak reference models for DM does not negatively affect ID
performance much. For instance, transferring from
$\text{ELECTRA}_{\text{Small}}$ to $\text{DeBERTaV3}_{\text{Large}}$ yields
the best accuracy on MultiNLI. We suspect the reason for this is that weak
models usually identify easy training instances as ambiguous data, which are
sufficient for obtaining satisfactory ID performance (Swayamdipta et al.,
2020).

#### Differences between Effective and Ineffective Reference Models

Because some reference models lead to failed transfers, it is important to
understand the differences between effective and ineffective reference models.
To answer this question, we start by considering the differences between a
weak and a reasonably strong reference model when categorizing training data
into ambiguous, hard-to-learn, and easy instances.

Assume we have a weak reference model which is only able to fit the easy
training instances but fails to fit hard training instances. As a result, this
weak reference model will assign increasing $p_{\text{true}}$ to easy training
instances across different epochs, while keeping $p_{\text{true}}$ for hard
training instances around the values expected in a random guessing scenario.
This reference model behavior means that $p_{\text{true}}$ will exhibit high
standard deviations on the subset of easy cases, which will therefore be
identified as ambiguous data; while hard training instances will have lower
means for $p_{\text{true}}$ and therefore be identified as hard-to-learn data.
In contrast, a reasonably capable reference model can fit the easy part of the
training instances during the early stage of training, which makes these
instances have both high means and low standard deviations for
$p_{\text{true}}$. In contrast, $p_{\text{true}}$ for harder instances will
gradually increase across different epochs, making these instances yielding
relatively low means and high standard deviations for $p_{\text{true}}$. As a
result, such cases will be identified as both ambiguous and hard-to-learn
(i.e. we expect a large overlap in these subsets). Because we select a fixed
percentage $q\%$ of instances as ambiguous or hard-to-learn, the larger
overlap results in a larger percentage of data instances to be classified as
easy when using strong reference models.

This reasoning is supported by our previous observation on the ID performance
of unsuccessful transfers (Table 2). Concretely, the fact that weaker
reference models (e.g. TinyBERT) identify easy training instances as
ambiguous, and that using only easy training instances for fine-tuning will
produce comparable ID performance but degraded OOD performance (Swayamdipta et
al., 2020), together explain our results for weak reference models.

To further validate our reasoning, we compute the percentages of easy training
instances from different reference models and show the results in Table 3.
Besides the default $q\%=33\%$, we also experiment with $q\%=25\%$ and
$q\%=50\%$. Consistent with our reasoning, ineffective reference models (i.e.
$\text{ELECTRA}_{\text{Small}}$ on NLI and TinyBERT on both datasets) indeed
identify fewer data points as easy compared with other models. Furthermore,
the overlap between hard-to-learn and ambiguous instances in successful
transfers is usually very large. For example, with $q\%=50\%$, all effective
reference models identify more than 45% of the training data as easy, (the
maximum is 50%, when ambiguous and hard-to-learn data align perfectly).

Mode | Main Model | Ref. Model | MultiNLI | WANLI | AdversarialNLI
---|---|---|---|---|---
|  |  |  |  | R1 | R2 | R3
ERM | TinyBERT | - | 67.32 | 43.40 | 23.30 | 28.10 | 30.94
ERM | $\text{ELECTRA}_{\text{Small}}$ | - | 81.98 | 54.11 | 23.38 | 28.57 | 30.25
ERM | $\text{ELECTRA}_{\text{Base}}$ | - | 88.53 | 63.06 | 34.58 | 30.73 | 31.29
ERM | $\text{ELECTRA}_{\text{Large}}$ | - | 90.75 | 65.85 | 54.20 | 39.38 | 36.10
DM | $\text{DeBERTaV3}_{\text{Large}}$ | TinyBERT | 89.17 | 60.02 | 41.83 | 34.58 | 34.54
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Small}}$ | 90.91 | 62.09 | 49.62 | 38.50 | 35.98
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Base}}$ | 90.63 | 66.58 | 59.77 | 46.25 | 42.29
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Large}}$ | 90.80 | 66.42 | 58.95 | 44.57 | 41.52
DM | $\text{ELECTRA}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Small}}$ | 89.88 | 61.53 | 45.90 | 36.20 | 31.89
DM | $\text{ELECTRA}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Base}}$ | 90.40 | 66.09 | 54.10 | 40.97 | 37.31
DM | $\text{ELECTRA}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Large}}$ | 90.33 | 65.37 | 53.73 | 39.67 | 36.17
DM | $\text{ELECTRA}_{\text{Large}}$ | Random | 89.99 | 65.03 | 51.25 | 39.02 | 34.98
Table 2:  Performance of different models on NLI with conventional fine-tuning (ERM) and DM using the $33\%$ most ambiguous data identified with different reference models. Random in Ref. Model means randomly selecting $33\%$ of the train data. The rows marked in gray are the results for DM training where the transfer was not successful, and the corresponding reference models. Successful transfer requires the reference model to be reasonably strong: reference models of clearly worse performance lead to degraded OOD performance for the main models.  | NLI: MultiNLI | HSD: CAD
---|---|---
| 25% | 33% | 50% | 25% | 33% | 50%
$\text{ELECTRA}_{\text{Small}}$ | 55.91% | 47.77% | 35.84% | 68.72% | 62.10% | 46.07%
$\text{ELECTRA}_{\text{Base}}$ | 69.10% | 62.88% | 46.62% | 73.36% | 65.54% | 47.56%
$\text{ELECTRA}_{\text{Large}}$ | 66.91% | 59.08% | 45.00% | 67.80% | 60.03% | 45.01%
$\text{DeBERTaV3}_{\text{Small}}$ | 67.08% | 61.31% | 46.08% | 72.05% | 65.00% | 47.94%
$\text{DeBERTaV3}_{\text{Base}}$ | 71.39% | 64.02% | 46.90% | 73.20% | 65.80% | 48.04%
$\text{DeBERTaV3}_{\text{Large}}$ | 72.97% | 64.84% | 46.97% | 74.06% | 65.89% | 47.72%
TinyBERT | 51.29% | 38.44% | 20.01% | 63.58% | 54.45% | 40.20%
Table 3:  Differences between effective and ineffective reference models. We
show the ratios of easy data identified by the data maps of different models;
cells marked in gray are the ratios of easy data identified by ineffective
reference models. The column names indicate different thresholds for $q\%$ in
the DM method. The key difference between effective and ineffective reference
models lies in their ability to identify easy cases: TinyBERT identifies less
easy data than other models on both datasets, and
$\text{ELECTRA}_{\text{Small}}$ identifies less easy data on NLI.

Figure 2:  Performance on NLI (top) and HSD (bottom) when training the main
model (ELECTRA) with fewer training steps. ERM is standard ERM fine-tuning on
the full training set. DM-* refers to fine-tuning
$\text{ELECTRA}_{\text{Large}}$ with the DM method, using reference model
$\text{ELECTRA}_{\text{*}}$. The X-axis is the percentile of full training
steps used, ranging from 1/15 to 7/15 of the total number of training steps.
The Y-axis is the percentile of performance compared with the full training
steps. Fine-tuning the main model using data maps is much faster than ERM: the
models achieve close-to-$100\%$ performance using only 1/3 of the training
steps.

##  5 FTFT: efficient and robust Fine-Tuning by transFerring Training dynamics

Given the transferability of training dynamics across model sizes and
pretraining methods (§4), we can improve the efficiency of the DM method by
using more efficient reference models. However, after training the reference
model, DM method still requires fine-tuning the main model in the same way as
ERM, making it less efficient than ERM. In this section, we investigate
whether we can overcome this limitation. We show that models trained using DM
learn faster than ERM, and fine-tuning the main model in DM for much fewer
than the full training steps of ERM is already sufficient. Based on these
observations, we propose a novel approach, Fine-Tuning by transFerring
Training dynamics (FTFT). We show that FTFT consistently offers both better
efficiency and better robustness over ERM fine-tuning.

#### DM Learns Faster Than ERM

Figure 2 shows the OOD test performance of $\text{ELECTRA}_{\text{Large}}$
fine-tuned with fewer steps (i.e. from 1/15 to 7/15 of the full training
steps). We compare DM with different reference model sizes against ERM. To
better show the different learning speeds, instead of the absolute performance
scores, we show the percentiles compared with using the full training steps
for each method. Clearly DM methods learns much faster than ERM on all OOD
datasets: the performance percentile curves of DM methods are almost always
above those of ERM. We also observe that DM methods using only 1/3 of the full
training steps already achieve comparable performance as using all training
steps. This result suggests that we can further improve the efficiency of the
DM method by training fewer steps, while maintaining its robustness advantage
over ERM. We show the results for $\text{DeBERTaV3}_{\text{Large}}$ in
Appendix B, and observe similar trends.

#### FTFT: Achieving both Efficiency and Robustness

Mode | Main Model | Ref. Model | Cost | MultiNLI | WANLI | AdversarialNLI
---|---|---|---|---|---|---
|  |  |  |  |  | R1 | R2 | R3
ERM | $\text{DeBERTaV3}_{\text{Large}}$ | - | 32.0 | 91.06 | 66.46 | 58.17 | 45.57 | 41.34
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{DeBERTaV3}_{\text{Large}}$ | 64.0 | 90.75 | 66.33 | 59.75 | 45.60 | 41.94
FTFT | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{DeBERTaV3}_{\text{Small}}$ | 15.2 | 90.12 | 66.42 | 60.30 | 45.75 | 43.66
FTFT | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{DeBERTaV3}_{\text{Base}}$ | 19.7 | 90.14 | 66.47 | 59.77 | 46.65 | 42.71

Mode Main Model Ref. Model CAD DynaHate-Original DynaHate-Perturb R2 R3 R4 R2
R3 R4 ERM $\text{DeBERTaV3}_{\text{Large}}$ - 81.69 75.44 73.32 76.12 70.62
77.41 68.89 DM $\text{DeBERTaV3}_{\text{Large}}$
$\text{DeBERTaV3}_{\text{Large}}$ 81.58 79.17 76.87 77.73 73.34 76.63 67.54
FTFT $\text{DeBERTaV3}_{\text{Large}}$ $\text{DeBERTaV3}_{\text{Small}}$ 80.73
78.77 77.53 79.48 76.19 77.53 69.31 FTFT $\text{DeBERTaV3}_{\text{Large}}$
$\text{DeBERTaV3}_{\text{Base}}$ 79.76 82.05 76.77 78.62 75.22 78.43 71.00

Table 4:  Comparison between FTFT and ERM/original DM fine-tuning. Performance
of DeBERTaV3 on NLI (top, accuracy) and HSD (bottom, macro-F1). ERM is
conventional ERM fine-tuning, and FTFT refers to using the $33\%$ most
ambiguous data identified by different reference models (i.e., Ref. Model)
_and only training for 1/3 of the total steps_. Cost refers to the fine-tuning
cost, with the cost of fine-tuning ELECTRA-Small with ERM as the unit. FTFT
yields both better efficiency and better robustness compared to both ERM fine-
tuning and the original DM method.

Fine-Tuning by transFerring Training dynamics (FTFT), involves two crucial
changes to the standard DM method. First, FTFT uses more efficient PLMs as
reference models. Second, to improve the efficiency of the main model, FTFT
uses only 1/3 of the training steps used for ERM fine-tuning. We use 1/3 of
the total training steps because we select 33% of the ambiguous train data,
and this choice means that we fine-tune on each training instance for the same
number of times as ERM fine-tuning on the full training dataset. Nevertheless,
we recommend to determine the number of training steps by monitoring model
performance. Table 4 summarizes the performance of FTFT using
$\text{DeBERTaV3}_{\text{Large}}$ as the main model, and
$\text{DeBERTaV3}_{\text{Small}}$ and $\text{DeBERTaV3}_{\text{Base}}$ as
reference models. We also include the performance of
$\text{DeBERTaV3}_{\text{Large}}$ fine-tuned with ERM and the original DM
method for comparison. We also show the cost for each method for NLI, using
the cost of fine-tuning $\text{ELECTRA}_{\text{Small}}$ for this task as the
unit. The relative cost for each method in HSD is the same as NLI.

We make two observations. First, FTFT methods show better robustness (i.e.
better performance on OOD datasets) than both ERM and the original DM method.
Second, FTFT is also much more efficient than ERM and the original DM method.
For example, FTFT with $\text{DeBERTaV3}_{\text{Base}}$ and
$\text{DeBERTaV3}_{\text{Small}}$ as the reference models only cost $19.7$ and
$15.2$ units of computation: compared with the 32/64 units of computation of
ERM/DM fine-tuning, these two FTFT choices are respectively 1.63/3.26 times
and 2.11/4.22 times cheaper.

##  6 Conclusions

Fine-tuning PLMs has been shown vulnerable to OOD and adversarial inputs. The
DM method has been shown to improve model robustness (Swayamdipta et al.,
2020), however, it is computationally expensive. In this paper, we have
presented a novel approach for fine-tuning PLMs, FTFT, which yields both
better efficiency and better robustness over conventional ERM fine-tuning
(§5). FTFT is built on the DM method, based on two observations: 1) in DM
methods, reference model training dynamics are highly transferable across
different model sizes (§4.1) and pretraining methods (§4.2), and 2) models
trained using DM learn faster than when using conventional ERM fine-tuning. We
have also discussed the conditions for successful FTFT runs (§4.3). We believe
FTFT will be an important tool for future researchers and practitioners to
perform efficient PLM fine-tuning, especially in situations where robustness
against out-of-distribution or adversarial inputs is essential.

#### Limitations

Nevertheless, our work has limitations. First, despite that we have observed
that effective reference models will identify more easy instances, we did not
conduct controlled experiments to validate whether this feature is the only
condition to ensure the success of transferring training dynamics. Future
researchers can perform more exhaustive empirical studies to investigate the
prediction of reference model effectiveness. Second, we have only tested FTFT
on two classification tasks. Future studies can extend our work to tasks of
other types, such as generative tasks (e.g., question answering), to examine
the generalizability of FTFT.

## References

  * Clark et al. (2020) Kevin Clark, Minh-Thang Luong, Quoc V. Le, and Christopher D. Manning.  ELECTRA: Pre-training text encoders as discriminators rather than generators.  In _ICLR_ , 2020.  URL https://openreview.net/pdf?id=r1xMH1BtvB.
  * He et al. (2021) Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen.  Deberta: Decoding-enhanced bert with disentangled attention.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=XPZIaotutsD.
  * He et al. (2023) Pengcheng He, Jianfeng Gao, and Weizhu Chen.  DeBERTav3: Improving deBERTa using ELECTRA-style pre-training with gradient-disentangled embedding sharing.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=sE7-XhLxHA.
  * Hoffmann et al. (2022) Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al.  Training compute-optimal large language models.  _arXiv preprint arXiv:2203.15556_ , 2022.
  * Kaplan et al. (2020) Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_ , 2020.
  * Karimi Mahabadi et al. (2020) Rabeeh Karimi Mahabadi, Yonatan Belinkov, and James Henderson.  End-to-end bias mitigation by modelling biases in corpora.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pp. 8706–8716, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.769.  URL https://aclanthology.org/2020.acl-main.769.
  * Lhoest et al. (2021) Quentin Lhoest, Albert Villanova del Moral, Yacine Jernite, Abhishek Thakur, Patrick von Platen, Suraj Patil, Julien Chaumond, Mariama Drame, Julien Plu, Lewis Tunstall, Joe Davison, Mario Šaško, Gunjan Chhablani, Bhavitvya Malik, Simon Brandeis, Teven Le Scao, Victor Sanh, Canwen Xu, Nicolas Patry, Angelina McMillan-Major, Philipp Schmid, Sylvain Gugger, Clément Delangue, Théo Matussière, Lysandre Debut, Stas Bekman, Pierric Cistac, Thibault Goehringer, Victor Mustar, François Lagunas, Alexander Rush, and Thomas Wolf.  Datasets: A community library for natural language processing.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_ , pp. 175–184, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.  URL https://aclanthology.org/2021.emnlp-demo.21.
  * Liu et al. (2022) Alisa Liu, Swabha Swayamdipta, Noah A. Smith, and Yejin Choi.  WANLI: Worker and AI collaboration for natural language inference dataset creation.  In _Findings of the Association for Computational Linguistics: EMNLP 2022_ , pp. 6826–6847, Abu Dhabi, United Arab Emirates, December 2022\. Association for Computational Linguistics.  URL https://aclanthology.org/2022.findings-emnlp.508.
  * Liu et al. (2021) Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn.  Just train twice: Improving group robustness without training group information.  In Marina Meila and Tong Zhang (eds.), _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 6781–6792. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/liu21f.html.
  * Loshchilov & Hutter (2019) Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=Bkg6RiCqY7.
  * Mosbach et al. (2021) Marius Mosbach, Maksym Andriushchenko, and Dietrich Klakow.  On the stability of fine-tuning {bert}: Misconceptions, explanations, and strong baselines.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=nzpLWnVAyah.
  * Nam et al. (2020) Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin.  Learning from failure: Training debiased classifier from biased classifier.  In _Advances in Neural Information Processing Systems_ , 2020.
  * Nie et al. (2020) Yixin Nie, Adina Williams, Emily Dinan, Mohit Bansal, Jason Weston, and Douwe Kiela.  Adversarial NLI: A new benchmark for natural language understanding.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_. Association for Computational Linguistics, 2020\.
  * Ramponi & Tonelli (2022) Alan Ramponi and Sara Tonelli.  Features or spurious artifacts? data-centric baselines for fair and robust hate speech detection.  In _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pp. 3027–3040, Seattle, United States, July 2022. Association for Computational Linguistics.  doi: 10.18653/v1/2022.naacl-main.221.  URL https://aclanthology.org/2022.naacl-main.221.
  * Sanh et al. (2021) Victor Sanh, Thomas Wolf, Yonatan Belinkov, and Alexander M Rush.  Learning from others’ mistakes: Avoiding dataset biases without modeling them.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=Hf3qXoiNkR.
  * Sar-Shalom & Schwartz (2023) Aviad Sar-Shalom and Roy Schwartz.  Curating datasets for better performance with example training dynamics.  In _Findings of the Association for Computational Linguistics: ACL 2023_ , pp. 10597–10608, Toronto, Canada, July 2023. Association for Computational Linguistics.  URL https://aclanthology.org/2023.findings-acl.674.
  * Strubell et al. (2019) Emma Strubell, Ananya Ganesh, and Andrew McCallum.  Energy and policy considerations for deep learning in NLP.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pp. 3645–3650, Florence, Italy, July 2019. Association for Computational Linguistics.  doi: 10.18653/v1/P19-1355.  URL https://aclanthology.org/P19-1355.
  * Swayamdipta et al. (2020) Swabha Swayamdipta, Roy Schwartz, Nicholas Lourie, Yizhong Wang, Hannaneh Hajishirzi, Noah A. Smith, and Yejin Choi.  Dataset cartography: Mapping and diagnosing datasets with training dynamics.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pp. 9275–9293, Online, November 2020\. Association for Computational Linguistics.  doi: 10.18653/v1/2020.emnlp-main.746.  URL https://aclanthology.org/2020.emnlp-main.746.
  * Turc et al. (2020) Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Well-read students learn better: On the importance of pre-training compact models, 2020.  URL https://openreview.net/forum?id=BJg7x1HFvB.
  * Utama et al. (2020) Prasetya Ajie Utama, Nafise Sadat Moosavi, and Iryna Gurevych.  Mind the trade-off: Debiasing NLU models without degrading the in-distribution performance.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pp. 8717–8729, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.770.  URL https://aclanthology.org/2020.acl-main.770.
  * Vidgen et al. (2021a) Bertie Vidgen, Dong Nguyen, Helen Margetts, Patricia Rossini, and Rebekah Tromble.  Introducing CAD: the contextual abuse dataset.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pp. 2289–2303, Online, June 2021a. Association for Computational Linguistics.  doi: 10.18653/v1/2021.naacl-main.182.  URL https://aclanthology.org/2021.naacl-main.182.
  * Vidgen et al. (2021b) Bertie Vidgen, Tristan Thrush, Zeerak Waseem, and Douwe Kiela.  Learning from the worst: Dynamically generated datasets to improve online hate detection.  In _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)_ , pp. 1667–1682, Online, August 2021b. Association for Computational Linguistics.  doi: 10.18653/v1/2021.acl-long.132.  URL https://aclanthology.org/2021.acl-long.132.
  * Williams et al. (2018) Adina Williams, Nikita Nangia, and Samuel Bowman.  A broad-coverage challenge corpus for sentence understanding through inference.  In _Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers)_ , pp. 1112–1122, New Orleans, Louisiana, June 2018. Association for Computational Linguistics.  doi: 10.18653/v1/N18-1101.  URL https://aclanthology.org/N18-1101.
  * Wolf et al. (2020) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_ , pp. 38–45, Online, October 2020. Association for Computational Linguistics.  URL https://www.aclweb.org/anthology/2020.emnlp-demos.6.
  * Wu et al. (2022) Carole-Jean Wu, Ramya Raghavendra, Udit Gupta, Bilge Acun, Newsha Ardalani, Kiwan Maeng, Gloria Chang, Fiona Aga, Jinshi Huang, Charles Bai, Michael Gschwind, Anurag Gupta, Myle Ott, Anastasia Melnikov, Salvatore Candido, David Brooks, Geeta Chauhan, Benjamin Lee, Hsien-Hsin Lee, Bugra Akyildiz, Maximilian Balandat, Joe Spisak, Ravi Jain, Mike Rabbat, and Kim Hazelwood.  Sustainable ai: Environmental implications, challenges and opportunities.  In D. Marculescu, Y. Chi, and C. Wu (eds.), _Proceedings of Machine Learning and Systems_ , volume 4, pp. 795–813, 2022.  URL https://proceedings.mlsys.org/paper_files/paper/2022/file/462211f67c7d858f663355eff93b745e-Paper.pdf.
  * Zhang et al. (2022) Michael Zhang, Nimit S Sohoni, Hongyang R Zhang, Chelsea Finn, and Christopher Re.  Correct-n-contrast: a contrastive approach for improving robustness to spurious correlations.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 26484–26516. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/zhang22z.html.

##  Appendix A Training Specifications

###  A.1 Experimental Setup

For training all models, we use AdamW (Loshchilov & Hutter, 2019) as the
optimizer with a batch size of 32. We also use a linear learning rate
scheduler with 10% warmup. For fine-tuning the small and base versions of both
DeBERTa-V3 and ELECTRA, as well as TinyBERT, we use a learning rate of 2e-5.
For DeBERTa-V3-Large and ELECTRA-Large, we respectively use 1e-5 and 2e-6 as
the learning rates. We also experimented with 1e-5 and 5e-6 for ELECTRA-Large.
However, we observed high ratios of failed runs, i.e., the training fails to
converge and produces a worse than majority class baseline (Mosbach et al.,
2021). We therefore follow the suggestions from Mosbach et al. (2021) to adopt
a lower learning rate. We use Python 3.9 and PyTorch 2.0 for all experiments.
For training PLMs, we use HuggingFace Transformers 4.32 (Wolf et al., 2020),
Accelerate 0.22, and Datasets 2.14 (Lhoest et al., 2021). All experiments are
performed on two NVIDIA A100 GPUs. Training all models takes approximately
seven GPU days.

###  A.2 Comparison of Training Costs

Here we exhibit the training costs of different models (PFLOPs), of a single
forward epoch. When taking back-propagation, optimization, and multiple epochs
into consideration, these numbers should scale proportionally.

| NLI: MultiNLI | HSD: CAD
---|---|---
$\text{DeBERTaV3}_{\text{Small}}$ | 3116.93 | 312.35
$\text{DeBERTaV3}_{\text{Base}}$ | 6233.34 | 624.68
$\text{DeBERTaV3}_{\text{Large}}$ | 22160.84 | 2220.9
$\text{ELECTRA}_{\text{Small}}$ | 694.61 | 69.98
$\text{ELECTRA}_{\text{Base}}$ | 6227.14 | 627.37
$\text{ELECTRA}_{\text{Large}}$ | 22138.79 | 2230.5
Table 5:  Comparison of training costs, measured in PFLOPs. Here we only
calculate the forward cost for a single epoch.

##  Appendix B Additional Results

Mode | Main Model | Ref. Model | CAD | DynaHate-Original | DynaHate-Perturb
---|---|---|---|---|---
|  |  |  | R2 | R3 | R4 | R2 | R3 | R4
ERM | TinyBERT | - | 71.72 | 43.19 | 49.6 | 52.88 | 43.21 | 57.5 | 52.91
ERM | $\text{ELECTRA}_{\text{Small}}$ | - | 74.15 | 48.93 | 56.93 | 58.93 | 55.75 | 61.84 | 59.52
ERM | $\text{ELECTRA}_{\text{Base}}$ | - | 76.91 | 62.13 | 60.71 | 61.72 | 55.42 | 64.64 | 62.05
ERM | $\text{ELECTRA}_{\text{Large}}$ | - | 68.07 | 70.07 | 62.56 | 70.27 | 63.91 | 70.45 | 67.02
DM | $\text{DeBERTaV3}_{\text{Large}}$ | TinyBERT | 78.91 | 71.68 | 71.69 | 76.52 | 71.22 | 75.92 | 68.23
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Small}}$ | 79.74 | 78.09 | 77.4 | 78.75 | 75.26 | 76.79 | 70.05
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Base}}$ | 80.37 | 81.47 | 78.16 | 78.38 | 76.99 | 78.51 | 71.01
DM | $\text{DeBERTaV3}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Large}}$ | 79.48 | 76.71 | 75.97 | 78.58 | 73.55 | 77.8 | 69.81
DM | $\text{ELECTRA}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Small}}$ | 74.35 | 75.62 | 67.92 | 74.06 | 68.43 | 71.62 | 67.3
DM | $\text{ELECTRA}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Base}}$ | 75.84 | 75.52 | 65 | 72.89 | 67.94 | 70.34 | 67.94
DM | $\text{ELECTRA}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Large}}$ | 71.93 | 65.25 | 62.78 | 72.45 | 62.92 | 71.17 | 68.23
DM | $\text{ELECTRA}_{\text{Large}}$ | Random | 74.14 | 62.14 | 58.13 | 68.23 | 63.3 | 67.93 | 65.22
Table 6:  Performance of different models on HSD with conventional fine-tuning
(ERM) and DM using the $33\%$ most ambiguous data identified with different
reference models. Random in Ref. Model means randomly selecting $33\%$ of the
train data. The rows marked in gray are the results for DM training where the
transfer was not successful, and the corresponding reference models.
Successful transfer requires the reference model to be reasonably strong:
reference models of clearly worse performance lead to degraded OOD performance
for the main models. Note that although using TinyBERT and
$\text{DeBERTaV3}_{\text{Large}}$ respectively as the reference and the main
model yields better performance than ELECTRA, that is due to the better
performance of $\text{DeBERTaV3}_{\text{Large}}$ than
$\text{ELECTRA}_{\text{Large}}$: it is still worse than other methods using
$\text{ELECTRA}_{\text{Large}}$ as the main model, especially on OOD test
sets.

Figure 3:  Performance on NLI (top) and HSD (bottom) when training the main
model (DeBERTaV3) with fewer training steps. ERM is standard ERM fine-tuning
on the full training set. DM-* refers to fine-tuning
$\text{ELECTRA}_{\text{Large}}$ with the DM method, using reference model
$\text{ELECTRA}_{\text{*}}$. The X-axis is the percentile of full training
steps used, ranging from 1/15 to 7/15 of the total number of training steps.
The Y-axis is the percentile of performance compared with the full training
steps. Fine-tuning the main model using data maps is much faster than ERM: the
models achieve close-to-$100\%$ performance using only 1/3 of the training
steps.

Mode | Main Model | Ref. Model | Cost | MultiNLI | WANLI | AdversarialNLI
---|---|---|---|---|---|---
|  |  |  |  |  | R1 | R2 | R3
ERM | $\text{ELECTRA}_{\text{Large}}$ | - | 32.0 | 90.75 | 65.85 | 54.2 | 39.38 | 36.1
DM | $\text{ELECTRA}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Large}}$ | 64.0 | 90.33 | 65.37 | 53.73 | 39.67 | 36.17
FTFT | $\text{ELECTRA}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Small}}$ | 11.7 | 89.46 | 60.87 | 46.05 | 35.82 | 32.6
FTFT | $\text{ELECTRA}_{\text{Large}}$ | $\text{ELECTRA}_{\text{Base}}$ | 19.7 | 90.06 | 65.85 | 54.7 | 40.23 | 37.58

Mode Main Model Ref. Model CAD DynaHate-Original DynaHate-Perturb R2 R3 R4 R2
R3 R4 ERM $\text{ELECTRA}_{\text{Large}}$ - 68.07 70.07 62.56 70.27 63.91
70.45 67.02 DM $\text{ELECTRA}_{\text{Large}}$ $\text{ELECTRA}_{\text{Large}}$
71.93 65.25 62.78 72.45 62.92 71.17 68.23 FTFT $\text{ELECTRA}_{\text{Large}}$
$\text{ELECTRA}_{\text{Small}}$ 75.52 73.45 66.54 73.79 67.08 72.05 68.11 FTFT
$\text{ELECTRA}_{\text{Large}}$ $\text{ELECTRA}_{\text{Base}}$ 77.15 74.17
65.53 73.42 68.77 70.6 68.7

Table 7:  Comparison between FTFT and ERM/original DM fine-tuning. Performance
of ELECTRA on NLI (top, accuracy) and HSD (bottom, macro-F1). ERM is
conventional ERM fine-tuning, and FTFT refers to using the $33\%$ most
ambiguous data identified by different reference models (i.e., Ref. Model).
Cost refers to the fine-tuning cost, with the cost of fine-tuning ELECTRA-
Small with ERM as the unit. FTFT yields both better efficiency and better
robustness compared to both ERM fine-tuning and the original DM method. For
NLI, we only train the main model for 1/3 of the total steps. For HSD, we
observe that ELECTRA converges a bit slower, see Figure 2. We therefore use
7/15 of the total steps for FTFT to obtain better performance (still less than
1/2), making the cost 15.9 and 23.9 for small and base reference models
respectively. FTFT yields both better efficiency and better robustness,
compared to both ERM fine-tuning and the original DM method, except for using
$\text{ELECTRA}_{\text{Small}}$ as the reference model on NLI.


# Multiphasic stochastic epidemic models

Petros Barmpounakis and Nikolaos Demiris
Department of Statistics
Athens University of Economics and Business
Greece
{barmpounakis<EMAIL_ADDRESS>

###### Abstract

At the onset of the Covid-19 pandemic, a number of non-pharmaceutical
interventions have been implemented in order to reduce transmission, thus
leading to multiple phases of transmission. The disease reproduction number
$R_{t}$, a way of quantifying transmissibility, has been a key part in
assessing the impact of such interventions. We discuss the distinct types of
transmission models used and how they are linked. We consider a hierarchical
stochastic epidemic model with piece-wise constant $R_{t}$, appropriate for
modelling the distinct phases of the epidemic and quantifying the true disease
magnitude. The location and scale of $R_{t}$ changes are inferred directly
from data while the number of transmissibility phases is allowed to vary. We
determine the model complexity via appropriate Poisson point process and
Dirichlet process-type modelling components. The models are evaluated using
synthetic data sets and the methods are applied to freely available data from
California and New York states as well as the United Kingdom and Greece. We
estimate the true infected cases and the corresponding $R_{t}$, among other
quantities, and independently validate the proposed approach using a large
seroprevalence study.

_K_ eywords First keyword $\cdot$ Second keyword $\cdot$ More

The emergence on early 2020 of Covid-19, an infectious disease caused by the
virus SARS-CoV2, has placed health systems around the globe under immense
pressure. On March 2020, the World Health Organization declared Covid-19 as a
global pandemic, and as of the end of September 2022 more than 6.5 million
have died due to illness or complications of it. At the beginning of the
pandemic in the absence of available vaccines or suitable medication the
majority of governments around the globe resorted to Non-Pharmaceutical-
Interventions (NPIs) in an attempt to stop the exponential spreading of the
virus and reduce transmissibility. Such NPIs involved measures like work-from-
home policies, school and university closures, stay-at-home guidance for
people in high-risk groups and full lockdowns.

These measures had an effect on reducing the transmissibility and resulted in
spreading trajectories that could not be properly described by the standard
epidemic models due to the resulting multiphasic nature of transmission. The
first systematic technique to assess these interventions was due to Flaxman et
al., (2020) who proposed a renewal equation model whose infection dynamics
were modelled through a multilevel framework incorporating NPIs. We amend this
model by inferring the points in time that the transmissibility changes as
well as the magnitude of infectiousness in a data-driven manner. We determine
the model complexity by using appropriate stochastic processes based upon
variations of the Poisson process (PP) and Dirichlet process (DP)-based priors
via their stick-breaking constructions (Miller and Harrison, (2018);
Sethuraman, (1994)).

Several models have been proposed in the literature for the estimation of
multiphasic infectious diseases, particularly Covid-19. Briefly, a stochastic
Susceptible-Exposed-Infectious-Removed (SEIR) model with a regression
framework for the effect of the NPIs on transmissibility is used in Knock et
al., (2021) while Birrell et al., (2021), Li et al., (2021) and Chatzilena et
al., (2022) use stochastic SEIR models where the transmission mechanism is
described by a system of non-linear ordinary differential equations and the
transmission rate is modelled by a diffusion process. Modelling the
transmission rate as a random walk facilitates gradual and smooth changes in
time. A piecewise linear quantile trend model was proposed by Jiang et al.,
(2021), a kernel-based SIR model distinguishing the different phases of the
transmissibility in space was developed by Geng et al., (2021) while Wistuba
et al., (2022) incorporated splines to estimate the reproduction number in
Germany.

Simpler forms of deterministic and stochastic multiphasic epidemic models have
been considered before. In the context of modelling SARS-CoV2 transmission
Flaxman et al., (2020) used an approach with fixed number, location and scale
of the $R_{t}$ change. Related work based upon variations of Dirichlet process
mixtures is presented in Hu and Geng, (2021) and Creswell et al., (2023). In
the former, the authors used a Mixture of finite mixtures (MFM) model on a
Susceptible-Infected-Recovered-Susceptible model, while in the latter the
authors used a suitably modified Pitman-Yor process but only for the scenario
of fitting to the observed cases, thus dispensing with the effort to estimate
the complete epidemic burden and the suitable adjustment for the reproduction
number. The main advantage of the proposed methodology is the intuitive
characterization of the epidemic in terms of multiple phases of
transmissibility. The number and magnitude of the distinct phases are
determined purely by data without explicitly using information about policy
changes and NPIs. This approach should be central to a retrospective
assessment of the NPIs: an evidence-based method for estimating the timing and
effect of those interventions, minimising the risk of introducing several
types of bias.

The paper is organized as follows. In section 1 we define the proposed
compartmental process, elucidate its equivalence with renewal process-based
models and describe the observation regimes of the data. In section 2 we
complete the model definition by characterising the complexity regimes.
Section 3 assesses the proposed models via simulation experiments while
section 4 contains the application to data from California and New York state,
the United Kingdom and Greece. The paper concludes with discussion.

##  1 Modelling Disease Transmission

###  1.1 Model Definition and Related Characterisations

The methodology for modelling the time-varying disease transmissibility has
been implemented under two distinct but equivalent models, the compartmental
Susceptible-Infectious-Removed (SIR) model and the seemingly simpler time-
since-infection model with population susceptibility reduction. Here we define
both models and delineate their equivalence.

The model assumes that the population has size $n$, is closed (demographic
changes during the course of the epidemic are ignored) homogeneous and
homogeneously mixing. In the stochastic SIR model, an infected individual
makes contact with any other individual on day $t$ at the points of a time-
homogeneous Poisson process with time-varying intensity
$\frac{\lambda_{t}}{n}$. This scaling is commonly adopted as it makes the
contact process independent of the size of the population (e.g., Andersson and
Britton,, 2000). If these (close) contacts of an infected individual occur
with a susceptible they result in an infection. Each individual remains
infectious for a random time period $Y$. All Poisson processes in this
construction are assumed to be independent. The disease reproduction number is
defined as $R_{t}=\lambda_{t}*E[Y]$, $t=1,\dots,T$ where $T$ is the time
horizon of the study.

For this model the expected number of new infections $c_{t+1}$ at day $t+1$ is
given by:

$\displaystyle\begin{split}E[c_{t+1}]&=S_{t}*\frac{\lambda_{t}}{n}*I_{t}*\Delta_{t+1-t},\\\
\end{split}$ (1)

with $I_{t}$ denoting the active set of infectives:

$\displaystyle\begin{split}I_{t}&=\sum_{s=0}^{t}\sum_{j=1}^{c_{s}}P(Y_{j}>t-s)\\\
\end{split}$ (2)

and $P(Y_{j}>t-s)$ the probability that individual $j$ infected on day $s$
remains infectious on day $t$. This probability is implicitly determined by
the disease characteristics. Then (1) can be rewritten as

$\displaystyle\begin{split}E[c_{t+1}]&=S_{t}*\frac{R_{t}}{n}*\frac{\sum_{s=0}^{t}\sum_{j=1}^{c_{s}}P(Y_{j}>t-s)}{E[Y]}=\frac{S_{t}}{n}*R_{t}*\sum_{s=0}^{t}c_{s}*g_{s}(t),\\\
\end{split}$ (3)

where $g_{s}(t)=\frac{P(Y>t-s)}{E[Y]}$ is called the generation interval which
defines the time from infection of an individual until the first infection
they generate, see for example Åke Svensson, (2015), Åke Svensson, (2007) and
Champredon et al., (2018). Note that equation (3) is used in the commonly
adopted technique of Cori et al., (2013) for estimating the instantaneous
reproduction number. In that approach, the term $\frac{S_{t}}{n}$ which
accounts for the depletion of the susceptible population is ignored since the
aim is somewhat different.

One should also consider potential ‘superspreading’ events when certain
individuals infected unusually large numbers of secondary cases (Shen et al.,,
2004; Lipsitch et al.,, 2003). We account for this variability assuming that
the individual reproduction number is gamma distributed with mean $R_{t}$ and
dispersion parameter $k$, yielding $c_{t}\sim NegativeBinomial(E[c_{t}],k)$
(Lloyd-Smith et al.,, 2005).

####  1.1.1 The Disease Reproduction Number

The reproduction number $R_{t}$ is of great practical interest as it is used
to assess if the epidemic is growing or shrinking. Here we consider two
distinct instances of reproduction number. The effective reproduction number
$R_{e}(t)=S_{t}*R_{t}$ describes the expected number of secondary cases
generated by an infectious individual. Then $R_{e}(t)>1$ and $R_{e}(t)<1$
indicate that the epidemic is growing or shrinking respectively and reducing
$R_{e}(t)$ below unity is the typical target of public health authorities. In
contrast, $R_{t}$ quantifies contacts that may not always result in new
infections, due to mixing with the immune proportion of the population.
Therefore, $R_{t}>1$ does not necessarily mean that the epidemic is growing. A
detailed discussion about reproduction numbers can be found in Pellis et al.,
(2022).

###  1.2 Observation Regimes

We consider two distinct observation regimes, one where the observed number of
cases corresponds to the total number of infections, explained below, and
whence the total number of infections is indirectly estimated, outlined in
1.2.2.

####  1.2.1 Observed Infections

The regime where the total number of infections are observed may be of
interest in its own right but may also be used for certain transmissible
diseases, for example in the analysis of influenza like illness data when
seroprevalence study information is available. Epidemic models are attractive
for analysing such data and are naturally defined in terms of infector-
infectee pair and the timing of such events. In reality however this type of
data is rarely available. Disease monitoring is based on the daily reported
infections, which are known to be susceptible to multiple problems, including
a time lag between the timing of infection and symptom onset or testing
positive.

In the case of Covid-19 a large proportion of the population experiences
asymptomatic or mild disease (Ward et al.,, 2021) leading to severe under-
reporting. Inference about the reproduction number can be robust when the
reported cases are used if depletion of the susceptible population is
accounted for, or if the observed proportion of cases remains constant over
time. One way to validate this assumption is by sequentially performing
seroprevalence studies to estimate the true disease prevalence and the
proportion of unreported incidences. However, regular such information was not
available in most countries. In the following subsection, we describe an
alternative approach that dispenses with the need for this assumption.

####  1.2.2 Unobserved Cases

The case where infections may not be directly observed has been studied in a
different context by Demiris et al., (2014). In the case of the pandemic, it
became immediately apparent that the observed number of infections only
partially accounts for the complete epidemic burden. An alternative technique
was proposed by Flaxman et al., (2020) where the true cases were estimated by
back-calculating infections from the daily reported deaths which are likely
less prone to under-reporting. This method has the additional advantage of
yielding an estimate of $S_{t}$. We adopt this approach for the second level
of our model and the daily deaths are linked with the true cases via:

$\displaystyle\begin{split}d_{t}&\sim NegativeBinomial(E[d_{t}],k)\\\
E[d_{t}]&=IFR*\sum_{i=0}^{t-1}c_{t}*\pi(i)\\\ \end{split}$ (4)

Accurate estimates of the infection fatality ratio ($IFR$) and time-from-
infection-to-death distribution ($\pi(i)$) are necessary for estimating
incidence, treated here as a latent parameter. The $IFR$ and $\pi(i)$
parameters may be calculated independently from external data or in a single
stage, leveraging additional evidence from seroprevalence studies as
illustrated in 4.

##  2 Epidemic Complexity Determination

The number of phases may be treated as a fixed but unknown integer or as a
random quantity to be modelled and estimated from data. We describe two such
models in the following two subsections.

###  2.1 Deterministic Number of Phases

For the models described above ‘model complexity’ refers to the number of
epidemic phases. In Flaxman et al., (2020) the number of phases was a-priori
selected and the times when the reproduction number $R_{t}$ changed were also
predefined. The locations of these points were informed by the NPIs
implemented by each government leading to a piece-wise constant reproduction
number $R_{t}$, effectively assuming immediate effect of those NPIs. We also
consider that $R_{t}$ is a piece-wise constant function and we amend this
transmission mechanism by inferring the location and magnitude of $R_{t}$
changes directly from the data. The number, $K$, of epidemic phases is
investigated using models with different $K$ values and the best model is
selected using the Watanabe–Akaike information criterion (WAIC) (Watanabe,,
2013) and Leave-one-out cross-validation (LOO) (Vehtari et al.,, 2017). The
model is defined as follows:

$R_{t}=\left\\{\begin{array}[]{ll}r_{1},\quad t\leq T_{1}\\\ ...\\\
r_{j+1},\quad T_{j}<t\leq T_{j+1}\\\ ...\\\ r_{K},\quad T_{K-1}<t\leq T\\\
\end{array}\right.$ $\displaystyle\begin{split}r_{j}&\sim f(\cdot),\quad
r_{j}\in(0,\infty),\quad j=1,...,K\\\ T_{i+1}&=T_{i}+e_{i}\\\
T_{1}&\sim\operatorname{Uniform}\left({3,T}\right)\\\
e_{i}&\sim\operatorname{Uniform}\left({0,100}\right),\quad i=1,...,K-1\\\
\end{split}$ (5)

###  2.2 Stochastic Number of Phases

Under the Bayesian paradigm, a natural but not trivial way is to treat the
model complexity, here the number of epidemic phases K, as a parameter and
learn its posterior distribution. The ‘reversible jump’ algorithm (e.g.,
Richardson and Green,, 1997) could be used to explore the joint space of K and
within-K models. Here we adopt a different approach and model $K$ as a
characteristic of two stochastic models, the Poisson process (PP) and
variations of the Dirichlet process (DP) (Ferguson,, 1973). For both
processes, we use the stick-breaking representation, see Miller and Harrison,
(2013) and Sethuraman, (1994) for the PP and DP respectively, facilitating
inference for $K$. The directed acyclic graph (Figure 1) represents the
general structure of our modelling framework.

Figure 1: Directed acyclic graph of the model. Ellipses denote parameters to
be learned by the model. The number of phases K is estimated by the DP/PP
model or via model selection criteria.

Estimating the number of phases of the epidemic and the associated location
and magnitude of the $R_{t}$ changes can lead to identifiability problems for
$R_{t}$ and its generative quantities, notably the total number of infections.
In order to overcome such issues we explore both a single and a multi-stage
modelling procedure (e.g., Bhatt et al.,, 2020). In the latter, at the first
stage, the latent disease cases are estimated using a Gaussian Process (GP)
model and then the medians of these latent cases are treated as data with
likelihood given in (3). The GP for the estimation of cases is presented in
the supplementary material.

####  2.2.1 Poisson Point Process-based Model

We consider that the arrival of new phases in the time horizon (0,T] is driven
by a time-homogeneous Poisson process with rate $\lambda$, with K growing
linearly with time. Hence, following the first epidemic phase, the number,
K-1, of new phases follows a Poisson distribution with rate $\lambda*T$ while
the duration of each phase a-priori follows an Exponential distribution with
rate $\lambda$. We follow Miller and Harrison, (2013) and use the
representation:

$\displaystyle\begin{split}R_{t}&=r_{z_{t}}\\\ r_{j}&\sim f(\cdot),\quad
r_{j}\in(0,\infty),\quad j=1,...,K\\\
z_{t}&\sim\operatorname{Categorical}\left({\pi_{1:K}}\right),\quad
t=1,...,T\\\ \pi_{K}&=1-\sum_{k=1}^{K}\pi_{k},\quad
K=min\\{j:\sum_{i=1}^{j}T_{i}\geq T\\}\\\ \pi_{k}&=\frac{T_{k}}{T},\quad
k=1,...,K-1\\\
T_{i}&\sim\operatorname{Exponential}\left({\lambda}\right),\quad
i=1,...,K_{max}\\\ \lambda&\sim\operatorname{Gamma}\left({0.02,1}\right)\\\
\end{split}$ (6)

truncating K at $Kmax=100$, far higher than data-supported estimates.

####  2.2.2 Dirichlet Process-based Model

An alternative model for the number of phases is based on the DP and its
stick-breaking construction:

$\displaystyle\begin{split}R_{t}&=r_{z_{t}}\\\ r_{j}&\sim f(\cdot),\quad
r_{j}\in(0,\infty),\quad j=1,...,L\\\
z_{t}&\sim\operatorname{Categorical}\left({w_{1:L}}\right),\quad t=1,...,T\\\
w_{L}&=\prod_{k<L}(1-v_{k}),\quad K=\sum_{k=1}^{L}I\\{w_{k}\geq 0\\}\\\
w_{l}&=v_{l}*\prod_{j=1}^{l-1}(1-v_{j}),\quad l=2,3,...,L-1\\\
w_{1}&=v_{1},\quad v_{i}\sim\operatorname{Beta}\left({1,\theta}\right),\quad
i=1,...,L-1\\\ \ \theta&\sim\operatorname{Gamma}\left({1,1}\right)\\\
\end{split}$ (7)

where $L$ is the truncation point of the DP, set here to 36. Here K is
increasing with the scaling parameter $\theta$.

##  3 Simulation Experiments

Simultaneously learning the parameters and the dimension of a model is
typically a challenging statistical task. Here we adopt a simulation-based
approach to inference whose details are given in the supplement. We assess the
performance of our methods by simulating epidemics of various characteristics
for 250 days. The epidemic model defined in (1) was used for simulating daily
infections and deaths. The population size was set at $10^{8}$ with $IFR=2\%$.
The discretized infectious period and the infection-to-death interval are
described in the supplementary material. The epidemic was simulated with 5
distinct increasing/decreasing phases resembling the observed Covid 19
outbreaks. The time-varying reproduction number was set as follows:

$R_{t}=\left\\{\begin{array}[]{ll}1.5,&t\leq 60\\\ 0.95,&60<t\leq 100\\\
1.35,&100<t\leq 150\\\ 0.8,&150<t\leq 200\\\ 1.8,&200<t\leq
250\end{array}\right.$

Using the model in (5) and the daily deaths as data the lowest WAIC and LOO
selected 5 changepoints. Models with varying (3, 4, 5 and 6) number of
changepoints incorrectly identified the first 10 days of the simulation as a
distinct phase (Figure 2). This can be attributed to the lack of information
at the start, a common issue in epidemic models. Following this period the
model with 5 changepoints correctly identifies the different epidemic phases,
including their timing and magnitude of change. The total daily infections
(Figure 2) are also accurately recovered. Inference was initiated the day that
10 cumulative deaths were observed. Plots for the other models may be found in
the supplementary material.

(a) Simulated (triangles) and estimated daily infections with 95% Cr.I.
(line).

(b) Real (solid line) and estimated reproduction number $R_{t}$ with 95% Cr.I.
(dashed line).

Figure 2: Simulation and estimates based on observing deaths

(a) Dirichlet process model

(b) Poisson Process model

Figure 3: True (solid line) and estimated reproduction number $R_{t}$ with 95%
Cr.I. (dashed line) based on observing infections

In addition to the findings that the models correctly select the right
complexity, it is interesting to summarise the model behaviour when
investigating model misspecification. Broadly, these findings may be
summarised as follows; when we fix the number of phases to be smaller than the
true one then the model is correctly recovering the early ones while it is
averaging the final phases leading to poorly fitted models. In contrast, when
fixing $K$ to be larger than the true one then we essentially recover the true
patterns and get a good fit. Hence, slightly overestimating model complexity
is not materially affecting the recovery of the true signal. A list of
detailed results is outlined in the supplement.

When fitting the models with a stochastic number of phases to daily
infections, both the PP and DP models are precisely estimating the number of
epidemic phases, the time of change and the true $R_{t}$ value (Figure 3). The
model was run for 100000 iterations and 8 chains. The analysis based on
observing deaths is included in the supplementary material. Briefly, the
intermediate phases of the epidemic are well estimated while the first and
final phases are recovered with noise. The level of smoothing introduced by
the cubic spline affects the noisy estimation of the cases; the lower the
degrees of freedom the smoother the estimation of cases and subsequently the
reproduction number.

##  4 Real-data Application

###  4.1 Data Description and Prepocessing

The models were fitted to daily reported deaths from two US states, California
and New York and two European countries, the United Kingdom and Greece. The
data are accessible from John Hopkins University and ECDC and the time horizon
ran to the end of June 2021 when many NPIs were lifted. Due to a lack of data
availability, the model does not account for reinfections. The age-
standardized $IFR$ for each country was informed by the meta-analysis from
COVID-19 Forecasting Team, (2022) accounting for time, geography and
population characteristics. We allowed the $IFR$ to vary over time, accounting
for the age structure of those infected, the burden of health systems and
amendments in treating the disease. The infection-to-death time and generation
interval were given a Gamma distribution with (mean, standard deviation) set
to (19, 8.5) and (6.5, 4.4) days respectively.

###  4.2 Analyses and Results

(a) California state

(b) New York state

(c) The United Kingdom

(d) Greece

Figure 4: Estimation of Effective Reproduction Number $R_{e}(t)$ with 50%
Cr.I. (solid and dashed lines) based on observing deaths, fixed number of
phases model.

_California_ was one of the first US states to report cases on the 26th of
January, 2020. A state of emergency was declared on March 4, 2020, and
mass/social gatherings were banned while a mandatory statewide stay-at-home
order was issued on March 19, 2020. We fitted the model to daily deaths and
using WAIC/LOO selected 6 changepoints. Figures 4 and 5 suggest that
$R_{e}(t)$ was reduced after imposing restrictions and fell below the critical
value of 1 after April 2020 when school closure was decided for the remainder
of the 2019–2020 academic year. The epidemic remained under control until the
summer of 2020 when $R_{e}(t)$ jumped slightly above 1 following a gradual
relaxation of measures. On August 31, 2020, a new set of measures called
‘Blueprint for a Safer Economy’ was applied and all models show that they were
effective, alongside the gained immunity of the population, at reducing the
effective reproduction number below one and keeping the epidemic under control
until the first half of October 2020. All models estimate a sharp increase in
$R_{e}(t)$, which resulted in an increase in the daily reported cases and
deaths between November 2020 and January 2021. Nighttime curfew and regional
stay-at-home orders were announced at the start of December 2020 whence
$R_{e}(t)$ remained stable and began declining. The initiation of the
vaccination program on early 2021 brought the epidemic under control with
$R_{e}(t)$ remaining below 1.

_New York_ state had, by April 10 2020, more confirmed cases than any country
outside the US and was heavily affected at the start of the pandemic, with
daily recorded deaths reaching a thousand in April. On March 15 all New York
City schools were closed and on March 20 state-wide stay-at-home order was
declared. As a result, the models show a drop of $R_{e}(t)$ below 1 from mid-
March 2020 until August 2020 (Figures 4 and 5). The best-performing models
based on WAIC and LOO had 7 changepoints (8 distinct phases). This model
estimates that after the summer of 2020, $R_{e}(t)$ remained above 1 up until
the start of 2021 with a small increase during November and the holiday
season. The DP and PP models show similar estimates for $R_{e}(t)$ (Figure 5).

(a) California state - DP model

(b) California state - PP model

(c) New York state - DP model

(d) New York state - PP

Figure 5: Estimation of Effective Reproduction Number $R_{e}(t)$ with 95%
Cr.I. (solid and dashed lines) based on observing deaths, multi-stage
approach.

For the _United Kingdom_ a model with 8 changepoints was selected by WAIC and
LOO. Until early March 2020, when a lockdown was imposed we estimate that
$R_{t}\approx 3.5$ (Figure 4). These measures were lifted in early June and
during the lockdown $R_{e}(t)$ remained below 1, and therefore under control.
After the summer $R_{e}(t)$ increased above 1 and the so-called rule of six
was imposed while on November 5, 2020, the second lockdown was announced. The
number of reported deaths was reduced after the initiation of the vaccination
program on January 4 2021. Virtually identical estimates for the UK $R_{e}(t)$
are inferred by the DP and PP models (Figures and additional details in the
supplementary material).

Figure 6:  Cumulative sum of estimated daily infections with 95% Cr.I. (dashed
lines) and the estimation of REACT-2 with 95% C.I. (solid lines) for the
United Kingdom

(a) California state

(b) New York state

(c) The United Kingdom

(d) Greece

Figure 7: Reported (triangles) and estimated deaths with 50% Cr.I. (solid and
dashed lines) based on observing deaths, fixed number of phases model.

We conducted an independent (or ‘external’) validation of the model
performance based upon REACT-2, an antibody prevalence study conducted in the
UK with the participation of more than 100000 adults (Ward et al.,, 2021).
This is a unique opportunity as it took place on early July 2020 when waning
immunity was unlikely and provides a reasonable estimate of the total disease
burden up to that time. The estimated prevalence for the adult population
(children were excluded) was 6.0% (95% CI: 5.8, 6.1) and our estimate for the
whole population is 7.5% (95% Cr.I.: 5.7, 10.) (Figure 6) well compatible with
that independent estimate.

For _Greece_ WAIC and LOO selected the 7-changepoint model. At the starting
phase, we estimate $R_{e}(t)=3.36$ $(sd=0.88)$ and a decrease below 1 on the
first half of March 2020 (Figure 4). On March 10 the government suspended most
activities, including educational, shopping and recreational while a week
later all nonessential movement was restricted. The $R_{e}(t)$ estimate
remained below 1 until early June 2020 when it increased following the lifting
of restrictions. During summer $R_{e}(t)$ remained over 1 until November 2020
since a case spike on October led to new measures. Similar estimates for the
$R_{e}(t)$ are obtained by the DP and PP models (Supplementary material).

The computation time was similar for the PP and DP models with the DP being
faster. More importantly, we get valuable insights on the effectiveness of the
measures imposed by the governments. For New York and the UK it appears that
the NPIs predate the reductions in transmissibility. California and Greece
adopted the measures before a large first wave, like other EU countries and US
states. All regions were similar when these measures were relaxed: multiple
epidemic waves emerged and the estimated $R_{e}(t)$ remained above 1.

The results of our simulation experiments corroborate the findings of the
application to real data from different areas. The time-ordering of the data
facilitates avoiding label-switching problems typically encountered when
fitting mixture models. By selecting the number of phases we capture mortality
changes in all the real-world examples (Figure 7). The DP and PP models can
infer a slightly higher number of phases but the conclusions are not
materially affected. This observation is in line with Rousseau and Mengersen,
(2011) who show a generally stable behaviour of such so-called overfitted
mixture models, theoretically verifying the robust behaviour of the developed
models.

##  5 Discussion

In this article, we propose 3 models for the transmission mechanism of
infectious diseases with multiple epidemic phases. We use freely available
data to estimate the points in time when transmissibility changes and the
realised magnitude of tho NPI effects. We adopt this approach since many of
these interventions coexist or overlap and identifiability issues can arise
when disentangling individual effects and the associated time lags.
Essentially, one may retrospectively assess the effect of the NPIs by
comparing the changes in the reproduction number with the dates that these
measures were imposed. Selecting the number of phases requires multiple runs
and the computation time can be an issue when nowcasting is essential for
decision-making. Estimating model complexity via the DP and PP models
represents an alternative approach that is computationally efficient and
statistically robust.

The DP and PP models can estimate more epidemic phases and this issue is
discussed in detail in Rousseau and Mengersen, (2011) and Miller and Harrison,
(2013). In our setting, this effect essentially relates to the start and end
of the epidemic and the inherent challenges of limited information. At the
start of the epidemic, such uncertainty dictates that estimates should be
interpreted with caution. In the end, this is less of an issue and is mostly
due to the time lag between cases and deaths. When one is working with the
observed infections these issues are largely removed and inference is
typically accurate throughout the duration of the data as indicated by our
simulation experiments.

The models developed in this work are assuming a homogeneous and homogeneously
mixing population, like most of the work studying SARS-CoV2 transmission. This
may be appropriate for large populations such as working at the state or
country level since functional central limit theorems can reasonably be
thought of as applicable (e.g., Andersson and Britton,, 2000). Our models can
naturally be extended when more detailed information is available and this is
the subject of current research.

## Acknowledgments

This article is part of the first author’s doctoral thesis, co-financed by
Greece and the European Union (European Social Fund-ESF) through the
Operational Programme ‘Human Resources Development, Education and Lifelong
Learning’ in the context of the Act ‘Enhancing Human Resources Research
Potential by undertaking a Doctoral Research’ Sub-action 2: IKY Scholarship
Programme for PhD candidates in the Greek Universities.

The authors are grateful to Kostas Kalogeropoulos and Petros Dellaportas for
useful comments on an earlier version of this article.

## References

  * Andersson and Britton, (2000) Andersson, H. and Britton, T. (2000).  Stochastic epidemic models and their statistical analysis.  Lecture notes in statistics. Springer, New York, NY, 2000 edition.
  * Bhatt et al., (2020) Bhatt, S., Ferguson, N., Flaxman, S., Gandy, A., Mishra, S., and Scott, J. A. (2020).  Semi-mechanistic bayesian modeling of covid-19 with renewal processes.
  * Birrell et al., (2021) Birrell, P., Blake, J., van Leeuwen, E., Gent, N., and De Angelis, D. (2021).  Real-time nowcasting and forecasting of covid-19 dynamics in england: the first wave.  Philosophical Transactions of the Royal Society B: Biological Sciences, 376(1829):20200279.
  * Champredon et al., (2018) Champredon, D., Dushoff, J., and Earn, D. J. D. (2018).  Equivalence of the erlang-distributed seir epidemic model and the renewal equation.  SIAM Journal on Applied Mathematics, 78(6):3258–3278.
  * Chatzilena et al., (2022) Chatzilena, A., Demiris, N., and Kalogeropoulos, K. (2022).  A modelling framework for the analysis of the transmission of sars-cov2.
  * Cori et al., (2013) Cori, A., Ferguson, N. M., Fraser, C., and Cauchemez, S. (2013).  A new framework and software to estimate time-varying reproduction numbers during epidemics.  Am J Epidemiol, 178(9):1505–1512.
  * COVID-19 Forecasting Team, (2022) COVID-19 Forecasting Team (2022).  Variation in the COVID-19 infection-fatality ratio by age, time, and geography during the pre-vaccine era: a systematic analysis.  Lancet, 399(10334):1469–1488.
  * Creswell et al., (2023) Creswell, R., Robinson, M., Gavaghan, D., Parag, K. V., Lei, C. L., and Lambert, B. (2023).  A bayesian nonparametric method for detecting rapid changes in disease transmission.  Journal of Theoretical Biology, 558:111351.
  * Demiris et al., (2014) Demiris, N., Kypraios, T., and Smith, L. V. (2014).  On the epidemic of financial crises.  Journal of the Royal Statistical Society. Series A (Statistics in Society), 177(3):697–723.
  * Ferguson, (1973) Ferguson, T. S. (1973).  A Bayesian Analysis of Some Nonparametric Problems.  The Annals of Statistics, 1(2):209 – 230.
  * Flaxman et al., (2020) Flaxman, S., Mishra, S., Gandy, A., Unwin, H. J. T., Mellan, T. A., Coupland, H., Whittaker, C., Zhu, H., Berah, T., Eaton, J. W., Monod, M., Perez-Guzman, P. N., Schmit, N., Cilloni, L., Ainslie, K. E. C., Baguelin, M., Boonyasiri, A., Boyd, O., Cattarino, L., Cooper, L. V., Cucunubá, Z., Cuomo-Dannenburg, G., Dighe, A., Djaafara, B., Dorigatti, I., van Elsland, S. L., FitzJohn, R. G., Gaythorpe, K. A. M., Geidelberg, L., Grassly, N. C., Green, W. D., Hallett, T., Hamlet, A., Hinsley, W., Jeffrey, B., Knock, E., Laydon, D. J., Nedjati-Gilani, G., Nouvellet, P., Parag, K. V., Siveroni, I., Thompson, H. A., Verity, R., Volz, E., Walters, C. E., Wang, H., Wang, Y., Watson, O. J., Winskill, P., Xi, X., Walker, P. G. T., Ghani, A. C., Donnelly, C. A., Riley, S., Vollmer, M. A. C., Ferguson, N. M., Okell, L. C., Bhatt, S., and Team, I. C. C.-. R. (2020).  Estimating the effects of non-pharmaceutical interventions on covid-19 in europe.  Nature, 584(7820):257–261.
  * Geng et al., (2021) Geng, X., Katul, G. G., Gerges, F., Bou-Zeid, E., Nassif, H., and Boufadel, M. C. (2021).  A kernel-modulated sir model for covid-19 contagious spread from county to continent.  Proceedings of the National Academy of Sciences, 118(21):e2023321118.
  * Hu and Geng, (2021) Hu, G. and Geng, J. (2021).  Heterogeneity learning for sirs model: an application to the covid-19.  Statistics and Its Interface, 14:73–81.
  * Jiang et al., (2021) Jiang, F., Zhao, Z., and Shao, X. (2021).  Modelling the covid-19 infection trajectory: A piecewise linear quantile trend model.  Journal of the Royal Statistical Society. Series B: Statistical Methodology.  Publisher Copyright: © 2021 Royal Statistical Society.
  * Knock et al., (2021) Knock, E. S., Whittles, L. K., Lees, J. A., Perez-Guzman, P. N., Verity, R., FitzJohn, R. G., Gaythorpe, K. A. M., Imai, N., Hinsley, W., Okell, L. C., Rosello, A., Kantas, N., Walters, C. E., Bhatia, S., Watson, O. J., Whittaker, C., Cattarino, L., Boonyasiri, A., Djaafara, B. A., Fraser, K., Fu, H., Wang, H., Xi, X., Donnelly, C. A., Jauneikaite, E., Laydon, D. J., White, P. J., Ghani, A. C., Ferguson, N. M., Cori, A., and Baguelin, M. (2021).  Key epidemiological drivers and impact of interventions in the 2020 sars-cov-2 epidemic in england.  Science Translational Medicine, 13(602):eabg4262.
  * Li et al., (2021) Li, Y. I., Turk, G., Rohrbach, P. B., Pietzonka, P., Kappler, J., Singh, R., Dolezal, J., Ekeh, T., Kikuchi, L., Peterson, J., Bolitho, A., Kobayashi, H., Cates, M. E., Adhikari, R., and Jack, R. L. (2021).  Efficient bayesian inference of fully stochastic epidemiological models with applications to COVID-19.  R Soc Open Sci, 8(8):211065.
  * Lipsitch et al., (2003) Lipsitch, M., Cohen, T., Cooper, B., Robins, J. M., Ma, S., James, L., Gopalakrishna, G., Chew, S. K., Tan, C. C., Samore, M. H., Fisman, D., and Murray, M. (2003).  Transmission dynamics and control of severe acute respiratory syndrome.  Science, 300(5627):1966–1970.
  * Lloyd-Smith et al., (2005) Lloyd-Smith, J. O., Schreiber, S. J., Kopp, P. E., and Getz, W. M. (2005).  Superspreading and the effect of individual variation on disease emergence.  Nature, 438(7066):355–359.
  * Miller and Harrison, (2013) Miller, J. W. and Harrison, M. T. (2013).  Inconsistency of Pitman-Yor process mixtures for the number of components.
  * Miller and Harrison, (2018) Miller, J. W. and Harrison, M. T. (2018).  Mixture models with a prior on the number of components.  Journal of the American Statistical Association, 113(521):340–356.  PMID: 29983475.
  * Pellis et al., (2022) Pellis, L., Birrell, P. J., Blake, J., Overton, C. E., Scarabel, F., Stage, H. B., Brooks-Pollock, E., Danon, L., Hall, I., House, T. A., Keeling, M. J., Read, J. M., and and, D. D. A. (2022).  Estimation of reproduction numbers in real time: Conceptual and statistical challenges.  Journal of the Royal Statistical Society: Series A (Statistics in Society), 185(S1).
  * Richardson and Green, (1997) Richardson, S. and Green, P. J. (1997).  On bayesian analysis of mixtures with an unknown number of components (with discussion).  Journal of the Royal Statistical Society: Series B (Statistical Methodology), 59(4):731–792.
  * Rousseau and Mengersen, (2011) Rousseau, J. and Mengersen, K. (2011).  Asymptotic behaviour of the posterior distribution in overfitted mixture models.  Journal of the Royal Statistical Society: Series B (Statistical Methodology), 73(5):689–710.
  * Sethuraman, (1994) Sethuraman, J. (1994).  A constructive definition of dirichlet priors.  Statistica Sinica, 4(2):639–650.
  * Shen et al., (2004) Shen, Z., Ning, F., Zhou, W., He, X., Lin, C., Chin, D. P., Zhu, Z., and Schuchat, A. (2004).  Superspreading SARS events, beijing, 2003.  Emerg Infect Dis, 10(2):256–260.
  * Vehtari et al., (2017) Vehtari, A., Gelman, A., and Gabry, J. (2017).  Practical bayesian model evaluation using leave-one-out cross-validation and waic.  Statistics and Computing, 27(5):1413–1432.
  * Ward et al., (2021) Ward, H., Atchison, C., Whitaker, M., Ainslie, K. E. C., Elliott, J., Okell, L., Redd, R., Ashby, D., Donnelly, C. A., Barclay, W., Darzi, A., Cooke, G., Riley, S., and Elliott, P. (2021).  Sars-cov-2 antibody prevalence in england following the first peak of the pandemic.  Nature Communications, 12(1):905.
  * Watanabe, (2013) Watanabe, S. (2013).  A widely applicable bayesian information criterion.  J. Mach. Learn. Res., 14(1):867–897.
  * Wistuba et al., (2022) Wistuba, T., Mayr, A., and Staerk, C. (2022).  Estimating the course of the covid-19 pandemic in germany via spline-based hierarchical modelling of death counts.  Scientific Reports, 12(1):9784.
  * Åke Svensson, (2007) Åke Svensson (2007).  A note on generation times in epidemic models.  Mathematical Biosciences, 208(1):300–311.
  * Åke Svensson, (2015) Åke Svensson (2015).  The influence of assumptions on generation time distributions in epidemic models.  Mathematical Biosciences, 270:81–89.