# Ranking & Reweighting Improves Group Distributional Robustness

Yachuan Liu, Bohan Zhang, Qiaozhu Mei, Paramveer Dhillon

(University of Michigan, Ann Arbor, MI, USA.
<EMAIL_ADDRESS>
###### Abstract

Recent work has shown that standard training via empirical risk minimization
(ERM) can produce models that achieve high accuracy on average but low
accuracy on underrepresented groups due to the prevalence of spurious
features. A predominant approach to tackle this group robustness problem
minimizes the worst group error (akin to a minimax strategy) on the training
data hoping that it will generalize well on the testing data. However, this is
often suboptimal, especially when the out-of-distribution (OOD) test data
contains previously unseen groups. Inspired by ideas from the information
retrieval and learning-to-rank literature, this paper first proposes to use
Discounted Cumulative Gain (DCG) as a metric of model quality for facilitating
better hyperparameter tuning and model selection. Being a ranking-based
metric, DCG weights multiple poorly-performing groups (instead of considering
just the group with the worst performance). As a natural next step, we build
on our results to propose a ranking-based training method called Discounted
Rank Upweighting (DRU) which differentially reweights a ranked list of poorly-
performing groups in the training data to learn models that exhibit strong OOD
performance on the test data. Results on several synthetic and real-world
datasets highlight the superior generalization ability of our group-ranking-
based (akin to soft-minimax) approach in selecting and learning models that
are robust to group distributional shifts.

##  1 Introduction

Text data are naturally split into groups in many machine learning contexts,
e.g., sentiment classification with reviews from different users or a
personalized dialogue system. In both these examples, a group corresponds to a
user. In other contexts, such as online toxicity detection, the groups might
be implicit, e.g., certain user demographics, and need annotation. More
broadly, consider the scenario where training examples are stratified non-
uniformly into groups. Our goal is to build a model for this scenario that
generalizes to all groups by providing comparable classification accuracies—a
key objective of deploying robust and fair machine learning practices [1, 2,
3].

Recent work on robust and equitable machine learning has shown that the
traditional approach of minimizing the average training error, also known as
empirical risk minimization (ERM), can be suboptimal for this grouped data
setting. ERM produces models that achieve low test error on average but incur
high errors on underrepresented groups in the data, which raises serious
ethical and fairness concerns. One of the main reasons ERM conceals poor
performance on minority groups behind a vastly superior average accuracy is
its reliance on spurious relationships between labels and some features in the
majority groups to achieve high average accuracy [4, 5, 6, 7, 8]. Such
correlations between labels and features are nonexistent or present with an
opposite sign in the minority (or new) groups. This leads to ERM severely
underperforming on these groups while overfitting to the majority groups.

Prior research tackles spurious correlation by building models with low worst-
group error on the training dataset. One such prominent model, Group
Distributional Robust Optimization (Group DRO), seeks to minimize the worst
group’s training loss [9]. While Group DRO has shown promising performance
compared to ERM on some benchmark datasets, it is known to perform poorly when
the different groups contain varying amounts of predictive signal [10]. Group
DRO assumes the test groups are all seen during training and each group has a
distribution that is invariant between training and test. Following this
assumption, the worst group in test is also likely the worst in training.
Thus, Group DRO extends the minimax distributional robust optimization (DRO)
framework [11] to groups. However, this assumption does not always hold in
reality, and it is especially problematic in domain-generalization scenarios
where the test data contains previously unseen out-of-distribution (OOD)
groups that do not overlap with the training or validation data.111It turns
out that the online algorithm that implements Group DRO in the paper [9] does
consider multiple groups, unlike the theory proposed in the paper, but it
makes strong parametric assumptions and weights the different groups
exponentially, which often leads to suboptimal performance on the test
dataset.

Moreover, unlike datasets with few groups and clear identification of spurious
features by construction (analogous to a controlled experiment), e.g.,
WaterBirds [12], spurious features can be hard to locate in naturally grouped
datasets, as they usually are not present exclusively in a subset of groups.
For instance, in tasks such as sentiment classification of user reviews,
potential spurious features such as the writing style can be present in all
the groups to varying degrees. We conjecture that differentially reweighting
the various groups will help mitigate the impact of spurious features and help
us identify robust predictive patterns in the data.

In this paper we draw on ideas from the learning-to-rank literature to provide
a more effective solution to the group distributional robustness problem.
Specifically, we rank and reweight different groups based on their training
errors (as opposed to considering just the worst-performing group as done by
Group DRO or DRO). To summarize, we make two key contributions in this paper:
We develop methods that reweight groups based on the (reverse) ranking of
their classification accuracy to 1) choose hyperparameters and perform model
selection, and 2) train the model.

First, we use Discounted Cumulative Gain (DCG) [13] metric from the
information retrieval and learning-to-rank literature to rank and then
reweight several poorly performing groups to inform model selection. DCG
allows us to consider the validation performance across more than one group
while choosing hyperparameters, thus lowering the risk of overfitting.
Further, the DCG metric is less prone to having ties in hyperparameter
choices, leading to statistically identified models. Next, we turn to the task
of developing a novel training method for group distributional robustness.
Borrowing intuition from our use of DCG for model selection, we propose a new
robust training method called Discounted Rank Upweighting (DRU). DRU
iteratively upweights groups during each training epoch based on that group’s
classification accuracy ranking.

At a high level, our proposed approach can be seen as a soft-minimax strategy,
which smooths the predictive signal from multiple poorly-performing groups by
weighting them based on their accuracy-based ranked order. As we show later,
both DCG-based model selection and DRU-based model training outperform
multiple state-of-the-art methods for group distributional robustness on
several synthetic and real-world benchmark datasets.

The rest of the paper is organized as follows. Section 2 discusses the related
work. Section 3 describes the preliminaries, including problem setup and
baseline methods. Section 4 describes the methods for group distributional
robustness compromising DCG for model selection and ranking and reweighting as
a novel model training method, the DRU. Then, in sections 5 and 6, we present
the experiment results for DCG-based model selection and DRU for model
training, respectively. We conclude in Section 7.

##  2 Related Work

This paper focuses on group distributional robustness, i.e., training models
generalize well across groups. There are other notions of robustness in
machine learning, e.g., adversarial robustness or the study of long-tailed
distributions, but they are beyond the scope of this paper.

###  2.1 Group Distributional Shifts

There are a couple of ways to split data into groups based on prior work.
First, groups can occur in data organically based on the data collection
procedure. For example, all reviews by a given user can be assembled into a
group, or all the images taken from a particular camera can constitute a
group. All the items in one group share similar characteristics and are
assumed to follow the same data-generating process. These organic groups can
be further divided into sub-populations based on meta-information about each
group. For instance, the user groups can be divided into sub-groups based on
the demographic constitution of those groups. Similarly, images taken from the
same camera can be divided into sub-groups based on the photographer’s
identity. The data can also be split into groups based on the interaction
between the output label and a spurious feature, e.g., Waterbirds [12], CelebA
[14], and MultiNLI [15] datasets.

Given the importance and prevalence of the grouped data setting, several
algorithms have been developed for removing disparity in performance across
the different groups. Some popular algorithms include Group Distributionally
Robust Optimization (Group DRO), which directly minimizes the worst group’s
regularized error during model training [16, 9]. Invariant Risk Minimization
(IRM) penalizes the distributions of learned representations with different
optimal linear classifiers [17]. Both Group DRO and IRM require group
annotation at training time. Recently, an approach called Just Train Twice
(JTT) has been proposed that does not require group information at training
time [18]. JTT instead just upweights misclassified examples and retrains the
model. It has been demonstrated to provide superior performance to Group DRO
or IRM.

###  2.2 Learning to Rank

It is a subfield of the information retrieval literature which aims to build
systems that can accurately retrieve top $k$ documents from a document
database. Essentially, it involves ranking the documents in a database based
on their content. The common evaluation measures used in this literature
include Mean Average Precision (MAP), Discounted Cumulative Gain (DCG), and
(Normalized) Discounted Cumulative Gain ((N)DCG), and (N)DCG at $k$ [13].

DCG at $k$ simply adds up the scores earned at each position with inverse
logarithm weights up to the $k^{th}$ document, i.e.,

$\text{DCG@k}=\sum_{i=1}^{k}\frac{\text{Score}_{(i)}}{\log_{2}(i+1)}$ (1)

While our approach is inspired by learning-to-rank, the major difference is
that in information retrieval literature, higher weights are assigned to the
higher-ranked items (e.g., most relevant documents), while in our setting,
higher importance weights are given to lower-ranked groups (i.e., worst
performing groups). To the best of our knowledge, this is the first work that
uses a ranking-based approach to facilitate a soft-minimax strategy of
training machine learning models with group distributional robustness.

##  3 Preliminaries

We consider the standard supervised learning setup of classifying an input
$x\in\mathcal{X}$ as a label $y\in\mathcal{Y}$. We assume that the training
data comprises of $m_{train}$ groups from a set $\mathcal{G}$ where each group
$g\in\mathcal{G}$ consists of $n_{g}$ data points from a probability
distribution $P_{g}(\mathcal{X},\mathcal{Y})$. In addition to the feature
$x_{j}$ and label $y_{j}$, each training example $j$ is also annotated with
the subpopulation/group $g_{j}\in\mathcal{G}$ that it belongs to. To
summarize, the training dataset contains $n_{train}$ samples with group
annotations in the format
$\\{(x_{1},y_{1},g_{1}),\ldots,(x_{n_{train}},y_{n_{train}},g_{n_{train}})\\}$.
Our goal is to learn a model
$f_{\theta}:\mathcal{X}\times\mathcal{G}\rightarrow\mathcal{Y}$ parameterized
by $\theta\in\Theta$. The group loss for group $g$ is the average loss over
all examples in $g$, and we denote it as
$l_{g}(\theta)=\mathop{\mathbb{E}}_{(x,y)\sim
P_{g}(\mathcal{X},\mathcal{Y})}\mathcal{L}(x,y;f_{\theta})$, for a loss
function $\mathcal{L}$ and a machine learning model $f_{\theta}$.

This paper assumes no group overlap between OOD test and training/validation
sets. This is a more challenging setting than the alternative scenario
commonly considered in the prior literature, where the test set only contains
new proportions of groups but no previously unseen groups. The performance
evaluation metric for a robust model under group distribution shift is the OOD
test set accuracy. More concretely, it is preferable to have a model with high
worst-group accuracy on the OOD test data, but that does not sacrifice the
average accuracy significantly.

###  3.1 Baseline Methods

We compare our approach against several competitive baselines as described
below. All the methods (including our approach) use the same base learner
$f_{\theta}$—a finetuned DistilBERT model [19]. We describe the hyperparameter
choices and other technical details of the various methods later in the paper.

##### Empirical Risk Minimization (ERM):

This is the standard training method that trains models to minimize the
average training loss. The method doesn’t take any group information into
consideration while training the model.

##### Group Distributionally Robust Optimization (Group DRO):

Group DRO uses distributionally robust optimization to explicitly minimize the
loss on the worst-case domain (or group) during training. We operationalize
Group DRO by using the online algorithm provided by [9].

##### Just Train Twice (JTT):

As briefly described earlier, JTT is a recently proposed approach that
requires no group annotations and has shown superior performance over Group
DRO and its variations on several challenging benchmark applications [18]. JTT
involves a two-stage training approach which first trains a standard ERM model
for several epochs and then trains a second model that upweights the training
examples that the first model has misclassified.

##  4 Methods

As just described, the goal of group distributional robustness is to learn
models with superior worst-group accuracy on the OOD test dataset without
sacrificing average accuracy. To achieve this goal, a common surrogate
optimization objective function that past literature has employed is to learn
models to maximize worst-group accuracy on the OOD validation dataset (See
Equation 2; $G_{val}$ denotes groups in the validation dataset).

$\min_{\theta\in\Theta}\max_{g\in G_{val}}\ l_{g}(\theta)$ (2)

This minimax approach to robust model selection is suboptimal since it ignores
the predictive signal from other groups. It also simplistically assumes that
the worst-performing group on the validation dataset is distributionally
similar to the worst group on the OOD test set. So, instead of this hard
minimax approach, we propose a soft-minimax approach that weights the errors
from several poorly performing groups on the validation dataset to inform the
hyperparameter choices for model selection. Intuitively, our approach can be
seen as performing smoothing by borrowing statistical strength from several
groups instead of just the worst group. We leverage the information retrieval
and learning to rank literature to help us operationalize this soft group-
weighting. This literature contains several ranking-based metrics that provide
discounted importance to various items, e.g., the Discounted Cumulative Gain
(DCG) metric.

###  4.1 Discounted Cumulative Gain (DCG) for Model selection:

First, use any base learner model, e.g., ERM to get the classification losses
incurred by the different groups on the validation set. Next, sort all the
$m_{val}$ groups according to their loss $g(1),g(2),\ldots,g(m_{val})$ from
the largest (worst) to the smallest (best group). Then the composite DCG
metric with a cutoff $k$ becomes,

$\text{DCG@k}(\theta)=\sum_{i=1}^{k}\frac{l_{g(i)}(\theta)}{\log_{2}(i+1)}$
(3)

Equation 3 considers the $k$ groups with the highest OOD validation errors and
provides them increasing weights (higher weight for worse performing group).

Since the inverse logarithm function flattens fast as the number of groups
increases, one can use DCG at the quantile-level as opposed to group-level.
The quantile-level DCG takes in a list of quantiles
$\textbf{{q}}=[q_{1},\ldots,q_{k}]$ that corresponds to groups
$g^{(q_{1})},\ldots,g^{(q_{k})}$ at these quantiles, for example,
$\textbf{{q}}=[0,1,\ldots,k]$ are the groups at quantile 0 (worst-group),
quantile 1, up to quantile $k$. This leads to a slightly modified expression
for DCG as shown below:

$\text{DCG}_{\textbf{{q}}}@k(\theta)=\sum_{i=1}^{k}\frac{l_{g^{(q_{i})}}(\theta)}{\log_{2}(i+1)}$
(4)

####  4.1.1 Evaluation of the Model Selection metric

We just described a soft minimax-based model selection strategy, which
generalizes the hard minimax used previously in the literature. Recall that
the ultimate goal of effective model selection is to choose a model with
superior performance on the OOD test set. So, how do we evaluate the
effectiveness of our proposed metric over alternative model selection metrics?

An excellent way to think about this is how much concordance or agreement
exists between the models selected by a given metric on the validation and the
OOD test sets. A superior model selection metric should yield similar rankings
of candidate models on either validation or test datasets. Thus, the best
model on the validation set will also be the best model on the test data
leading to effective model selection. For instance, consider our soft minimax
metric; let’s assume it ranks three candidate models as S2$>$ S1$>$ S3 based
on validation set accuracy. Then, if the test set accuracies222Recall that on
test data we only care about worst group accuracy. of these three models are
also S2$>$ S1$>$ S3, we consider the metric a good model selection strategy,
and we can pick model S2 from this class of models. We use this intuition to
guide our evaluation strategy for the model selection metric.

Let $\text{r}_{val}(M)$ denote the ranked accuracy list of models based on
metric $M$, e.g., hard minimax, soft minimax, or average, on the validation
set. Next, let $\text{r}_{test}(\text{worst-group-accuracy})$ represent a
similar model accuracy list but based on worst group performance on the test
set. Then, the metric $M$’s concordance $C(M)$ can be defined as the
similarity between $\text{r}_{val}(M)$ and $\text{r}_{test}(\text{worst-group-
accuracy})$. Hence, a superior evaluation metric should have a high degree of
similarity between the two rankings.

$C(M)=\text{similarity}(\text{r}_{val}(M),\ r_{test}(\text{worst-group-
accuracy}))$ (5)

The similarity in above Equation 5 can be operationalized by a function such
as euclidean distance or cosine similarity.

###  4.2 Learning-To-Rank inspired novel method for Model Training

We just saw the use of DCG to select the best model from candidate models
efficiently. Next, we propose a new method for training a model. Inspired by
the recent success of the Just Train Twice (JTT) method [18] for group
distributional robustness, we propose a method that performs iterative
upweighting of training examples. However, unlike JTT, it leverages group
annotations at training time. Our novel approach, called Discounted Rank
Upweighting (DRU), iteratively ranks the groups by their accuracy and then
upweights poorly performing groups. The key idea is to upweight training
samples from the groups with the highest training errors and assign them
differential importance commensurate with their ranking.

At each epoch $t$ (excluding the first one) during the training process, a
sample $x$ with label $y$ in the group $g\in\mathcal{G}_{train}$ is upweighted
as

$w^{t}_{g}=\begin{cases}\frac{\log_{2}(C+2)}{\log_{2}(r^{t-1}_{g}+2)}&r^{t-1}_{g}<=C\\\
\quad\quad 1&r^{t-1}_{g}>C\end{cases}$ (6)

where $r^{t-1}_{g}$ is either the ranking index or the ranking quantile of the
group $g$ in the training set (ascending order of training accuracy) evaluated
from the previous (t-1) epoch. $C$ is a hyperparameter that controls the
cutoff for upweighting (akin to $k$ in $DCG@k$). If the group ranking is
greater than the cutoff, then the weight is one, that is, no upweighting.
Otherwise, if the group has lower accuracy, then, it will be weighted by the
discounted log function shown above. Note that the constant ‘2’ in this
function is used to have discounted factors for training groups that are
consistent with those proposed by the standard DCG metric [13]. If the
upweighting is applied to all samples of each group regardless of their
classification accuracy in the previous epoch, then the training objective of
each epoch for a model with parameters $\theta$ is

$J^{t}_{DRU}(\theta)=\sum_{g\in\mathcal{G}_{train}}\sum_{(x,y)\in
g}w^{t}_{g}*\mathcal{L}(x,y;f_{\theta})$ (7)

for $t\neq 0$. As one can infer, the first epoch is always the standard ERM
training.

The upweighting scheme shown in Equation 7 upweights all the samples from a
given group. One can also choose to upweight only the misclassified samples
from the previous epoch. Assuming the misclassified samples to constitute an
error set $E$, the modified training objective function becomes:

$J^{t}_{DRU}(\theta,E)=\sum\limits_{g\in\mathcal{G}_{train}}\left[\sum\limits_{(x,y)\in
g\cap E}w^{t}_{g}*\mathcal{L}(x,y;f_{\theta})+\sum\limits_{(x,y)\in g\setminus
E}\mathcal{L}(x,y;f_{\theta})\right]$

(8)

We compare both these objective functions in our empirical results.

##  5 Experimental Setup

###  5.1 Datasets

We use three real-world review sentiment classification datasets: AMAZON-WILDS
[10], IMDB Movie Review Dataset [20], and a variation of Yelp Open
Dataset333https://www.yelp.com/dataset/. The prediction task for all three
datasets is to classify the review text to its corresponding 1-to-5 star
rating. Each review is associated with a group, which corresponds to all
reviews written by the same user. Each dataset consists of an in-distribution
(ID) training set and out-of-distribution (OOD) validation and test sets. The
OOD validation and test sets comprise reviews from disjoint sets of users
(groups). The users in the training dataset are randomly split 50/50 to be in
the ID validation and ID test datasets. Table 1 provides the summary
statistics of each dataset. For more preprocessing and descriptive details of
the dataset, please refer to the Appendix A. Intuitively, the performance of a
ERM model should significantly downgrade on OOD validation and test sets than
on ID validation and test sets. Table 1 confirms the significant accuracy
drops from ID to OOD on all three datasets.

We also generate several synthetic datasets in which each observation of a
group is sampled from a ‘shared’ signal across all groups plus an
‘idiosyncratic’ signal which is group-dependent. The distribution and strength
of the two signals are different across different datasets. Please refer to
Appendix D for more details.

Dataset | # Groups | Group size | ID val | OOD val | ID test | OOD test
---|---|---|---|---|---|---
AMAZON | (1252, 1334, 1334) | 75 | (75.7, 58.7, 24.0) | (72.3, 54.7, 6.3) | (74.7, 57.3, 24.0) | (71.9, 53.3, 12.0)
IMDB | (666, 561, 560) | 25 | (64.7, 46.7, 26.7) | (62.6, 43.1, 15.6) | (65.4, 48.0, 20.0) | (63.2, 42.9, 15.0)
Yelp | (500, 523, 522) | 100 | (65.2, 54.9, 41.0) | (64.5, 54.0, 34.0) | (64.0, 55.9, 26.9) | (63.0, 52.0, 18.0)

Table 1: Dataset Details. Note: 1) Number of groups is provided in the format
(training, OOD validation, OOD test). 2) The ERM model accuracies are given in
the format (average, 10-th percentile, worst group). All three performance
metrics are lower on OOD val/test sets than on their ID counterparts. 3) 10-th
percentile group is one that has a lower accuracy than 90% of all groups.

###  5.2 Experiment Setup for Model Selection

We consider the following metrics that one can use to select models and
hyperparameters from a validation set: (1) worst-group: accuracy of the worst
group; (2) average: average across all groups; (3) 10th percentile: the
accuracy of the group at the 10th percentile (lower accuracy than 90% groups);
(4) gDCG@k: DCG at group level for the k percent worst-performing groups
(k=10, 50); (5) qDCG@k: DCG at quantile level for percentiles [0,1,…,k] (k=10,
50).

To compare the effectiveness of these metrics for model selection, we trained
DistilBERT base-learner models using JTT. We varied the two hyperparameters
(first stage step $T$ and upweighting factor $\lambda$) of JTT to generate 16
candidate models. In particular, we considered $T\in\\{1,2,3,5\\}$ and
$\lambda\in\\{2,3,5,10\\}$. Next, we rank these 16 models using the various
metrics on their OOD validation set accuracy and then rank all the 16 models
on worst group accuracy on the test OOD dataset. This process provides us with
a ranked list of 16 models for each model selection metric on the validation
set and another list ranking all the 16 models on their worst group accuracy
on the OOD test set. Finally, we can assess the model selection performance of
all the metrics by computing the similarity between their rankings of models
on validation set with the “ground-truth” ranking of models on the test set
(cf. Equation 5). In particular, we calculate the similarity of two ranked
lists using euclidean distance, cosine similarity, and the NDCG of the model
ranking on validation using the ranking on test as gold standard scores. The
metric(s) which ranks the models on validation into the most similar positions
with the true performance of the models on test set is(are) the most effective
for selection.

###  5.3 Experiment Setup for Model Training

Similar to the DCG-based model selection, the DRU-based upweighting can also
be performed by ranking the quantiles of the groups, i.e., qDRU, or simply
ranking the indices of the groups, i.e., gDRU. qDRU is preferable when the
number of groups is large since the logarithm function flattens out quickly in
such a case. These upweighting methods also have one hyperparameter $C$, which
controls the cutoff or the amount of smoothing.

Since we want to contrast with the hard minimax approach throughout this
paper, we train models by only upweighting samples of the worst-performing
group from the previous training epoch by a constant weight $\lambda$. It is
easy to see that this hard minimax approach is a special case of our soft
minimax approach in which the cutoff for DRU is the rank of the worst group
(0). We denote this boundary case as Worst in our results.

In their basic form, our upweighting methods qDRU, gDRU upweight all the
samples from certain groups (cf. Equation 7). A related upweighting strategy
can be to further zoom in to each group and only upweight the misclassified
examples from that group (cf. Equation 8). We experiment with this seemingly
more precise weighting strategy and denote it using the suffix “+M” in our
results. For clarity, the default strategy of upweighting all examples from a
group is suffixed “+G”. This leads to four different variants of our DRU
models, qDRU + M, qDRU + G, gDRU + M, and gDRU + G.

We further compare against another variant of the upweighting strategy that
upweights only the misclassified examples from the previous training epoch by
a constant factor $\lambda$. We call this approach Const in our results. This
variant will help us tease apart the impact of the ranking-based logarithmic
weighting since it is plausible that the improved accuracy might not be
sensitive to the differential upweighting. Note that the Const method can not
be applied to all the samples from each group (+G) since upweighting all the
examples by the same amount makes the weights useless.

In addition to these methods, we compare against baseline methods ERM, Group
DRO, and JTT described in Section §3.1. All methods except ERM and JTT have
one hyperparameter ($C$ for DRU-based methods, Const, and Worst, and stepsize
for Group GRO). We also considered IRM [17] as a potential baseline but could
not obtain comparable performance to other baselines (ERM, Group DRO, JTT). We
hypothesize that IRM does not fit our scenario where there are a large number
of groups. We therefore do not include IRM in the following experiments.

We follow the lead of the authors of the WILDS Distribution Shift Benchmark
Suite [10] and use a finetuned base uncased DistilBERT model as our base
learner in all our experiments. We used the following hyperparameters for
DistilBERT as also suggested by [10]: batch size 16; learning rate $1\times
10^{-5}$ for AdamW optimizer [21]; L2-regularization strength 0.01; 5 epochs
with early stopping; and a maximum number of 512 tokens. Next, for both qDRU
and gDRU, we performed a grid search to tune the cutoff hyperparameter
$C\in[10,20,50,100]$. $\lambda$ is selected from the list [2, 3, 5] for all
methods that performed constant upweighting. For JTT, $T\in\\{0,1,2\\}$ and
$\lambda\in\\{2,3,5\\}$. Finally, for Group DRO we fixed the step size as 0.01
following the best practice reported in [10]. For the DistilBERT, we use the
implementation of
HuggingFace444https://huggingface.co/docs/transformers/model_doc/distilbert.
All the experiments are run on the NVIDIA GEFORCE RTX 2080 Ti using the
PyTorch Framework.

##  6 Results and Discussion

###  6.1 Model Selection Results

Table 2 shows the model selection results for Amazon, IMDB, and Yelp datasets.
We report the similarity between the ranked list of the 16 candidate models on
OOD validation set corresponding to each model selection metric and the ranked
list of the models by their true accuracy on the worst group in OOD test.

As can be seen from the results, the “worst-group,” that is, the hard minimax
approach, usually performs worse compared to other metrics. This is surprising
since it has identical semantics to the ground-truth metric we used on the
test data (worst-group accuracy), confirming that the hard minimax metric has
poor generalizability when distribution shift is present. The quantile-level
DCG metrics (qDCGs) perform the best on all three datasets. Specifically,
qDCG@10 performs the best on Amazon and IMDB, and qDCG@50 performs the best on
the Yelp dataset.

The “10th percentile” metric works better than worst-group or average
accuracy, although not as good as the DCG-based metrics. This is
understandable as the “10th percentile” worst-group is a special case of rank-
based metric and it also smooths the hard minimax to a certain extent.

A hidden but practically important strength of the DCG-based metric that is
not visible in the result tables is its ability to break ties between
candidate models. In our experiments, 43.75% and 37.5% of models have
identical 10th percentile and worst-group accuracies on all three datasets,
respectively. This lack of model identification makes it hard to assess which
model is the best, and one has to resort to suboptimal heuristics such as
random tiebreaks to choose the best model. Ties are rare in the case of DCG-
based metrics since they evaluate models using discounted (logarithm weighted)
ranks of several poorly performing groups instead of just a single accuracy
number as done by the “worst-group” or “10th percentile” metrics. Thus, the
smoothing produced by our soft minimax metrics leads us to select better
models. However, our results are still a bit inconclusive regarding how much
smoothing is optimal (with regard to the cutoff threshold of DCG) since that
threshold hyperparameter varies over the datasets in our experiments.

| Amazon | IMDB | Yelp
---|---|---|---
Metric | ED | CS | NDCG | ED | CS | NDCG | ED | CS | NDCG
worst-group | 27.0 | .72 | .77 | 14.5 | .93 | .92 | 15.3 | .92 | .91
average | 27.0 | .74 | .78 | 15.9 | .92 | .88 | 12.9 | .94 | .97
10th percentile | 23.5 | .80 | .88 | 15.3 | .92 | .93 | 13.6 | .93 | .96
gDCG@10 | 21.4 | .84 | .86 | 12.6 | .95 | .95 | 14.0 | .94 | .94
gDCG@50 | 24.0 | .80 | .82 | 15.6 | .92 | .93 | 12.5 | .95 | .97
qDCG@10 | 20.4 | .85 | .86 | 12.6 | .95 | .95 | 13.2 | .94 | .94
qDCG@50 | 24.0 | .80 | .82 | 15.0 | .92 | 0.93 | 12.2 | .95 | .97

Table 2: Concordance between the ranked lists of the models on OOD validation
by different metrics and the ranked list by worst-group accuracies on OOD
test. Note: ED = Euclidean Distance (lower is better), CS = Cosine Similarity
(higher is better), NDCG = Normalized Discounted Cumulative Gain using test-
worst-group ranking list as the gold standard (higher is better).

###  6.2 Model Training Results

The results of the test OOD datasets for the various methods are shown in
Table 3. The tables report average and 10th percentile group accuracy for
completeness, but the worst-group accuracy on the test OOD dataset is the
target. Broadly, we see a trend that the DRU-based methods, which smoothly
upweight multiple poorly-performing groups, outperform ERM and other methods
with hard upweighting rules. Among DRU-based methods, qDRU+M provides the
highest worst group accuracy (target metric) on the OOD test dataset on
average, with a comfortable margin of improvement over JTT, Group DRO, Worst,
Const which are statistically significant under a bootstrapped t-test. The
10th percentile and average group performance of qDRU+M and other DRU-based
methods are also competitive or even better than the baselines, suggesting
that our soft-minimax-based methods improve OOD worst-group performance
without a perceptible sacrifice of accuracy on better-performing groups or
average accuracy. We also showcase similar improved performances of our
methods on several synthetic scenarios (Appendix D). In all these scenarios,
DRU-based methods still significantly outperform the baseline methods, and
qDRU+M is nearly consistently the best method on average, 10th and worst group
accuracy metrics.

Dataset ERM Group DRO JTT Worst + G Worst + M Amazon-WILDs 71.9/53.3/12.0
70.0/53.3/8.0 71.6/53.3/9.3 72.2/53.3/12.0 70.2/53.3/17.3 Yelp 63.0/52.0/18.0
59.2/49.0/27.0 61.7/51.0/19.0 62.8/52.0/20.0 62.7/53.0/25.0 IMDB
63.2/42.9/15.0 61.0/42.0/10.8 62.5/42.3/10.0 62.5/42.9/22.5 63.3/43.8/17.5

Dataset qDRU + G qDRU + M gDRU + G gDRU + M Const + M Amazon-WILDs
70.2/54.7/17.3 70.1/53.3/18.7 70.2/53.3/17.3 71.5/54.7/14.7 71.0/53.3/14.7
Yelp 62.6/53.0/23.0 62.5/52.0/21.0 62.8,53.0/24.0 62.1/52.0/27.0
63.0/53.0/21.0 IMDB 64.1/45.8/20.0 62.0/43.3/25.0 61.0/42.9/18.9
62.4/43.8/15.0 62.6/44.1/15.0

Table 3: Results on OOD Test. G: upweighting all samples of a group (Equation
7). M: upweighting misclassified samples of a group (Equation 8). qDRU:
upweighting according to ranking percentile. gDRU: upweighting according to
ranking index. The results are in the format of (average accuracy/10th
percentile accuracy/worst group accuracy). Bold: best OOD test performance
(bootstrapped t-test)

We observe a slightly different pattern on real-world datasets where,
interestingly, qDRU+G model consistently outperforms others on 10th percentile
accuracy. Finally, we would like to note that real-world scenarios are often
considerably more complex than the controlled synthetic setting since
individual groups and misclassified examples may be more affected by random
noises or even adversarial signals (e.g., spam or fake reviews). And in these
scenarios, the 10th percentile group accuracy may be a more reasonable target
for group robustness, and upweighting all examples in a group may be more
resilient to noise than only upweighting misclassified examples. To summarize
the three metrics into one for easier comparison, we also report the
t-statistics of group accuracy of each method in Appendix B, which shows our
methods did better in reducing the variance among groups without sacrificing
the average group performance.

##  7 Conclusion and Future Work

In conclusion, this paper highlights the weakness of the canonical approach in
group distributional robustness literature of focusing only on the worst group
accuracy for model selection and model training. We introduced a suite of
methods inspired by the information retrieval and learning-to-rank literature
for group robust model selection and model training. Essentially, our ranking-
based soft minimax approaches smooth the predictive signal learned at training
time by performing a discounted weighting which leads to improved
generalization performance on the OOD test dataset in the challenging domain
generalization [10] setting. Our theoretical intuition regarding the fit of
ranking-based methods for group robustness is backed by our methods’ equally
strong empirical performance on synthetic and several real-world benchmark
datasets. Group identities carry a strong predictive signal (even if they do
not overlap in training/test) since we observe that group-based approaches
perform better than those that ignore the group structure. Though, more
research needs to be done to investigate this deeply. The learning to rank
literature implicitly assumes orthogonality between the search results (or
group features in our case). So, as part of future work, it will be
interesting to study how the number of groups and their correlation structure
impact the performance of our ranking-based methods.

## Limitations

As discussed in the paper, one of this work’s limitations is that we do not
discuss the optimal cutoff threshold of either DCG-based model selection
metrics or DRU model training as they are data-dependent. We also need to
develop theoretical arguments to guarantee convergence of Equation 7, though
empirically, we didn’t face any issues with model convergence. Finally, like
most prior work, our methods require group labels which can be expensive to
annotate if they are not naturally observed. So, future work should also look
into unsupervised grouping strategies and apply our methods.

## Ethics Statement

We did not notice any immediate ethical issues in our work. Our proposed
methods improve the tail-group accuracy, so our approach ensures that we do
not adversely impact marginalized and disadvantaged groups. For the licenses,
the Amazon-WILDS dataset [10] is licensed under the MIT license. The IMDB
dataset [20] is licensed under CC-BY 4.0. The Yelp Open Dataset provides a
YELP DATASET TERMS OF USE with permission to use for academic purposes. Our
use of the three existing datasets for academic purposes is consistent with
their intended use. All usernames in these datasets will be anonymized into
hashing values when releasing the datasets.

## References

  * [1] Cynthia Dwork, Moritz Hardt, Toniann Pitassi, Omer Reingold, and Richard Zemel.  Fairness through awareness.  In Proceedings of the 3rd innovations in theoretical computer science conference, pages 214–226, 2012.
  * [2] Moritz Hardt, Eric Price, and Nati Srebro.  Equality of opportunity in supervised learning.  Advances in neural information processing systems, 29, 2016.
  * [3] Jon Kleinberg, Sendhil Mullainathan, and Manish Raghavan.  Inherent trade-offs in the fair determination of risk scores.  arXiv preprint arXiv:1609.05807, 2016.
  * [4] Dirk Hovy and Anders Søgaard.  Tagging performance correlates with author age.  In Proceedings of the 53rd annual meeting of the Association for Computational Linguistics and the 7th international joint conference on natural language processing (volume 2: Short papers), pages 483–488, 2015.
  * [5] Su Lin Blodgett, Lisa Green, and Brendan O’Connor.  Demographic Dialectal Variation in Social Media: A Case Study of African-American English.  In Proceedings of EMNLP, 2016.
  * [6] Rachael Tatman.  Gender and dialect bias in youtube’s automatic captions.  In Proceedings of the first ACL workshop on ethics in natural language processing, pages 53–59, 2017.
  * [7] Tatsunori Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang.  Fairness without demographics in repeated loss minimization.  In International Conference on Machine Learning, pages 1929–1938. PMLR, 2018.
  * [8] John C Duchi, Peter W Glynn, and Hongseok Namkoong.  Statistics of robust optimization: A generalized empirical likelihood approach.  Mathematics of Operations Research, 46(3):946–969, 2021.
  * [9] Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang.  Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization.  arXiv preprint arXiv:1911.08731, 2019.
  * [10] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al.  Wilds: A benchmark of in-the-wild distribution shifts.  In International Conference on Machine Learning, pages 5637–5664. PMLR, 2021.
  * [11] Hongseok Namkoong and John C Duchi.  Stochastic gradient methods for distributionally robust optimization with f-divergences.  Advances in neural information processing systems, 29, 2016.
  * [12] Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie.  The caltech-ucsd birds-200-2011 dataset.  2011\.
  * [13] Kalervo Järvelin and Jaana Kekäläinen.  Cumulated gain-based evaluation of ir techniques.  ACM Transactions on Information Systems (TOIS), 20(4):422–446, 2002\.
  * [14] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang.  Deep learning face attributes in the wild.  In Proceedings of the IEEE international conference on computer vision, pages 3730–3738, 2015.
  * [15] Adina Williams, Nikita Nangia, and Samuel R Bowman.  A broad-coverage challenge corpus for sentence understanding through inference.  arXiv preprint arXiv:1704.05426, 2017.
  * [16] Weihua Hu, Gang Niu, Issei Sato, and Masashi Sugiyama.  Does distributionally robust supervised learning give robust classifiers?  In International Conference on Machine Learning, pages 2029–2037. PMLR, 2018.
  * [17] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz.  Invariant risk minimization.  arXiv preprint arXiv:1907.02893, 2019.
  * [18] 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 International Conference on Machine Learning, pages 6781–6792. PMLR, 2021.
  * [19] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.  arXiv preprint arXiv:1910.01108, 2019.
  * [20] Aditya Pal, Abhilash Barigidad, and Abhijit Mustafi.  Identifying movie genre compositions using neural networks and introducing genrec-a recommender system based on audience genre perception.  In 5th International Conference on Computing, Communication and Security, ICCCS 2020, Patna, India, October 14-16, 2020, pages 1–7. IEEE, 2020.
  * [21] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  arXiv preprint arXiv:1711.05101, 2017.
  * [22] Udit Arora, William Huang, and He He.  Types of out-of-distribution texts and how to detect them.  arXiv preprint arXiv:2109.06827, 2021.
  * [23] Paramveer S Dhillon, Dean Foster, and Lyle H Ungar.  Minimum description length penalization for group and multi-task sparse learning.  The Journal of Machine Learning Research, 12:525–564, 2011.
  * [24] Paramveer S Dhillon, David A Wolk, Sandhitsu R Das, Lyle H Ungar, James C Gee, and Brian B Avants.  Subject-specific functional parcellation via prior based eigenanatomy.  NeuroImage, 99:14–27, 2014.
  * [25] Bing Xu, Naiyan Wang, Tianqi Chen, and Mu Li.  Empirical evaluation of rectified activations in convolutional network.  arXiv preprint arXiv:1505.00853, 2015.

## Appendix

##  Appendix A Dataset Descriptions

##### 1) AMAZON-WILDS:

Collected as part of the WILDS dataset suite [10], the Amazon-WILDS dataset
involves predicting star ratings from users’ reviews of Amazon products. The
training set has $245,502$ reviews from $1252$ users (at least 75 reviews per
user). The ID validation set consists of 46,950 reviews from 626 of the 1252
users in the training set. The ID test set is the same size as the ID
validation set and contains reviews from the remaining 626 users from the
training dataset. Finally, the OOD validation and the OOD test sets each have
$100,050$ reviews (75 reviews per user) from $1,334$ new users.

##### 2) IMDB Movie Reviews:

We downloaded the IMDB dataset from [20] and modified it to exhibit
considerable OOD performance drops on the validation and test sets. To
construct our dataset, we aggregate the data at the user level and split it
into training, validation, and test sets using K-means clustering to ensure a
significant distributional shift from ID to OOD sets. Specifically, we
calculate the average of pre-trained DistilBERT embeddings of each user’s
reviews and then cluster their embeddings (k=2). One cluster is randomly
selected as the ID set, and the other is the OOD set. Next, users in the OOD
set are randomly split into OOD validation and OOD test sets, and the users in
the ID set with at least 50 reviews are randomly divided into ID validation
and ID test sets. The final training set has $41,146$ reviews from $666$
users. The ID validation set consists of 20,070 reviews from 333 of the 666
users in the training set. Similarly, the ID test set contains 21,083 reviews
from the other half of the users in the training dataset. The OOD validation
and test sets include $42,703$ and $43,451$ reviews from $561$ and $560$
unseen users, respectively, with each user containing at least 25 reviews.

##### 3) Yelp Business Reviews:

WILDS dataset suite [10] contains a modified version of the Yelp Open Dataset;
however, there’s no accuracy drop from their ID set to OOD set. Thus, we
modify it by clustering at the user level in a similar fashion as we did for
the IMDB dataset. We set k=6 and select the two farthest clusters as the OOD
and ID sets to have a significant out-of-distribution performance drop. The
training set comprises $64,931$ reviews from $500$ users. The ID validation
and test sets consist of $20,070$ and $21,083$ reviews from 333 out of the 666
users in the training set, respectively. Finally, the OOD validation and test
sets include $52,200$ and $52,300$ reviews from $522$ and $523$ unseen users,
respectively.

##  Appendix B T-statistic of Real World Datasets

The t-statistics of each method on OOD datasets are shown in Table 4. For each
method, the t-statistic is computed by dividing the average group performance
by the standard error of all groups’ performances.

A robust model aims to have a higher worst group performance while not
sacrificing the average performance which results in a smaller standard error
of the groups’ performances, so the t-stat should be the higher the better.
The results show that our proposed methods are usually consistently better
than ERM, JTT and GroupDRO.

Table 4: T-stats of group performance on OOD Test (T-stats here calculated as
$\frac{\bar{p}}{SE(p)}$ where $\bar{p}$ is the average group performance,
$SE(p)$ is the standard error of all groups’ performance). A higher t-stat is
achieved by reducing the variance among group performance (e.g., improving
worst-group performance) while maintaining a high average accuracy. G:
upweighting all samples of a group. M: upweighting misclassified samples of a
group. qDRU: upweighting according to ranking percentile. gDRU: upweighting
according to ranking index. Bold: best OOD test performance under the same
setting.

Dataset ERM Group DRO JTT Worst + G Worst + M Amazon-WILDs 183.0 164.0 171.3
179.7 186.3 Yelp 164.0 169.2 165.1 169.2 176.1 IMDB 93.6 94.7 92.4 94.0 96.6

Dataset qDRU + G qDRU + M gDRU + G gDRU + M Const + M Amazon-WILDs 190.7 194.3
188.8 188.5 187.0 Yelp 173.3 169.2 177.7 165.6 172.0 IMDB 102.1 101.4 94.7
99.5 95.9

##  Appendix C Validation Performances of Real World Datasets

Results of validation sets of each method are shown in Table 5.

Table 5: Results on OOD Validation. G: upweighting all samples of a group. M:
upweighting misclassified samples of a group. qDRU: upweighting according to
ranking percentile. gDRU: upweighting according to ranking index. The results
are in the format of (average accuracy/10th percentile accuracy/worst group
accuracy). Bold: best OOD validation performance under the same setting.

Dataset ERM Group DRO JTT Worst + G Worst + M Amazon-WILDs (72.3/54.7/5.3)
(70.7/54.7/5.8) (72.5/53.3/5.3) (72.9/54.7/6.7) (71.0/53.3/8.0) Yelp
(64.5/54.0/34.0) (60.5/51.0/31.0) (63.2/52.0/32.0) (64.0/54.0/35.0)
(64.2/54.0/35.0) IMDB (62.6/43.1/15.6) (61.1/40.5/15.5) (62.0/43.0/15.4)
(61.5/41.9/17.6) (63.0/44.4/20.0)

Dataset qDRU + G qDRU + M gDRU + G gDRU + M Const + M Amazon-WILDs
(70.9/54.7/6.7) (71.1/54.7/6.7) (70.9/54.7/6.7) (72.1/56.0/8.0)
(71.8/54.7/8.0) Yelp (64.0/53.1/40.0) (64.3/54.3/40.0) (63.9/54.0/34.0)
(63.6/54.0/38.0) (64.4/53.0/37.0) IMDB (63.3/44.4/19.3) (61.9/46.0/15.7)
(61.0/44.0/15.4) (61.9/43.2/20.2) (62.5/44.4/17.6)

##  Appendix D Synthetic Data Experiments

We showcase its improved performance in a controlled environment where we
generate the data using a fixed data-generating process. Past work has also
used synthetic data to validate new methods for distribution shift [22, 17].

###  D.1 Synthetic Data Generation:

Algorithm 1 Synthetic Data Generation

$Q$: number of reviews per group; $U$: % of groups with idiosyncratic
predictive signal, $W$: a set of predefined Gaussian idiosyncratic predictive
signals; $p$: probability of each idiosyncratic signal; $M\sim
N(\mu_{s},\sigma_{s}^{2})$, the shared predictive signal.

$\mathcal{S}\leftarrow\\{\\}$

for $g\in\mathcal{G}$ do

$a_{g}\leftarrow Truncated\mathcal{N}(0.75,0.25,0,1)$;

$w_{g}\sim N(\mu_{g},\sigma_{g}^{2})\leftarrow random.choices(W,p)$;

$has\\_idiosyncratic\leftarrow\mathcal{U}(0,1)<=U$

for $i=$ 1,2,…,Q do

$shared_{g}^{i}\leftarrow\mathcal{N}(\mu_{s},\sigma_{s}^{2})$

$x_{g}^{i}\leftarrow shared_{g}^{i}$

if $has\\_idiosyncratic$ then

$idiosyncratic_{g}^{i}\leftarrow N(\mu_{g},\sigma_{g}^{2})$

$x_{g}^{i}\leftarrow x_{g}^{i}+a_{g}*idiosyncratic_{g}^{i}$

end if

$y_{g}^{i}=\mathbbm{1}_{\mathbb{R}_{+}}(f(x_{g}^{i})+\mathcal{N}(0,0.25))$

$\mathcal{S}\leftarrow\mathcal{S}\cup\\{(x_{g}^{i},y_{g}^{i})\\}$

end for

end for

return $\mathcal{S}$

The procedure for synthetic data generation is summarized in Algorithm 1. Our
data-generating process assumes that each observation is generated by
combining two predictive signals. The first signal is a “shared signal”
present across all the groups and easily captured by any model and the second
signal is an “idiosyncratic signal” that varies significantly across groups.
As part of our controlled simulation setup, we vary the percentage of groups
$U$ containing idiosyncratic signals in a dataset. Specifically, we model the
shared signal $M$ by a Gaussian distribution $\mathcal{N}(\mu_{s},\sigma_{s})$
and the idiosyncratic signal is operationalized by a set of $W$ Gaussian
distributions whose each element represents a unique idiosyncratic signal. For
each sample $i$ of a given group $g$, we sample a shared signal
$shared_{g}^{i}$ from the distribution $M$. Next, if the group $g$ contains
idiosyncratic signals (depends on $U$), one idiosyncratic signal distribution
$w_{g}$ is sampled from $W$ based on a prior distribution $p$. Then for each
sample $i$ of the group, an idiosyncratic signal $idiosyncratic_{g}^{i}$ is
sampled from $w_{g}$. Then, another simulation parameter $a\in[0,1]$ that
controls the strength of the idiosyncratic signal is sampled from a truncated
Gaussian distribution. Note that for a group without idiosyncratic signals,
each sample of the group will only have the shared signal. Finally, we get the
feature representation for a given sample $i$ of the group $g$ as
$x_{g}^{i}=shared_{g}^{i}+b*a*idiosyncratic_{g}^{i}$ where $b$ is a boolean
value indicating whether the group $g$ contains the idiosyncratic signal
(determined by $U$). The label of the sample is obtained as
$y=\mathbbm{1}_{\mathbb{R}_{+}}(f(x_{g}^{i})+\sigma)$, where
$\mathbbm{1}_{\mathbb{R}_{+}}$ is the indicator function, $f$ is a function
that transforms features into labels (e.g., sin function) and $\sigma$ is
random Gaussian noise.

Table 6: Four different synthetic dataset settings. $p$ is the prior
distribution of the four idiosyncratic signals (before normalizing). $U$ is
the portion of groups that have idiosyncratic signals in format of (training,
validation, test)

| $p_{train}$ | $p_{val}$ | $p_{test}$ | $U$
---|---|---|---|---
1 | (1, 1, 1, 1) | (1, 1, 1, 1) | (1, 5, 1, 5) | (0.8, 0.8, 0.8)
2 | (1, 1, 1, 1) | (1, 1, 1, 1) | (1, 5, 1, 5) | (0.2, 0.2, 0.8)
3 | (0, 1, 1, 1) | (1, 1, 1, 1) | (1, 5, 1, 5) | (0.2, 0.2, 0.8)
4 | (1, 1, 0, 0) | (1, 1, 0, 0) | (0, 0, 1, 1) | (0.2, 0.2, 0.8)

###  D.2 Experiment setup on Synthetic Data:

As shown in Algorithm 1, in our experiments, the dimension of all signals is
2. The shared predictive signal $M$ is $\mathcal{N}$([0, 0], $4I$) where $I$
is the two-dimensional identity matrix. There are four idiosyncratic signals
in $W$ and their mean and variance values are chosen as [([0.25, 0.25], $I$),
([0.25, -0.25], $I$), ([-0.25, 0.25], $I$), and ([-0.25, -0.25], $I$)]
respectively. $f$ is a sine function that takes the sum of all feature
dimensions as input. The strength factor $a$ is sampled from a truncated
Gaussian distribution $\mathcal{N}(0.75,0.25,0,1)$ and the random noise
$\sigma$ is assumed to be distributed $\mathcal{N}(0,0.25)$.

Table 7: Results on Synthetic Dataset. G: upweighting all samples of a group.
M: upweighting misclassified samples of a group. qDRU: upweighting according
to ranking percentile. gDRU: upweighting according to ranking index. The
results are in the format of (average accuracy, 10th percentile accuracy,
worst group accuracy). BOLD: best OOD test performance under the same setting.
Underline: second best performance.

Dataset ERM GroupDRO JTT Const+M Worst+G Worst+M Setting 1 (70.2, 62.7, 54.7)
(78.7, 72.0, 62.7) (76.5, 69.3, 62.7) (67.0, 60.0, 45.3) (76.6, 69.3, 58.7)
(77.8, 72.0, 62.7) Setting 2 (68.1, 61.3, 52.0) (76.0, 69.3, 57.3) (73.1,
66.7, 58.7) (67.6, 61.3, 52.0) (78.1, 72.0, 62.7) (76.9, 70.7, 58.7) Setting 3
(68.9, 61.3, 50.7) (75.7, 69.3, 60.0) (76.3, 70.7, 61.3) (70.3, 64.0, 56.0)
(77.4, 70.7, 64.0) (76.3, 69.3, 61.3) Setting 4 (66.9, 60.0, 50.7) (75.3,
68.0, 60.0) (75.2, 69.3, 60.0) (66.9, 60.0, 49.3) (76.5, 69.3, 61.3) (76.7,
70.7, 62.7)

Our approaches qDRU+G qDRU+M gDRU+G gDRU+M Setting 1 (77.2, 70.7, 61.3) (77.6,
70.7, 64.0) (75.9, 69.3, 60.0) (78.5, 72.0, 65.3) Setting 2 (77.7, 72.0, 61.3)
(79.0, 73.3, 65.3) (77.3, 70.7, 58.7) (79.5, 73.3, 65.3) Setting 3 (77.1,
70.7, 62.7) (81.0, 74.7, 69.3) (76.8, 69.3, 60.0) (76.4, 69.3, 62.7) Setting 4
(79.4, 73.3, 65.3) (80.4, 74.7, 68.0) (77.3, 70.7, 60.0) (76.9, 70.7, 62.7)

Using these parameters, we generate synthetic datasets under four different
settings as shown in Table 6. Setting 1 is the one that exhibits the slightest
distribution shift since 80% of the groups in the training, validation, and
test sets contain the same set of idiosyncratic signals. Setting 2 shows a
realistic real-world scenario where the training data uniformly
($p_{train}(w_{i})=1$) contains each of the idiosyncratic signals, but only
20% of the training groups have an idiosyncratic signal. The test data, on the
other hand, contains the idiosyncratic signal in 80% of the groups. The third
and fourth settings show substantial distribution shifts since they represent
the case where some of the idiosyncratic signals are altogether hidden from
the training dataset. This happens routinely in real-world scenarios when the
training dataset is not large enough to include all the unique signals
introduced by unseen groups in the OOD test data.

We generate 1000 training groups, 500 test OOD groups, and 500 validation OOD
groups for each of these settings. Each group contains 75 samples. The base
learner for training all these datasets is a three-layer feed-forward neural
network with a hidden state size of 128. Each layer is connected by LeakyReLU
[25], and a 0.5 dropout rate is applied. We performed a grid search
$C\in[5,10,20,50,100]$ to select the best cutoff for qDRU. For all the
upweighting methods with constant factors, i.e., Worst, Const, and JTT,
$\lambda$ is selected from the list [2, 3, 4, 5]. The step size for Group DRO
is chosen as 0.01, and the first and second training steps of JTT were 5.
Finally, we performed the model selection using qDCG@10 metric, which was the
best performing metric as we saw in §2.

###  D.3 Synthetic Data Results:

The results are shown in Table Table 7, and as can be seen, the DRU-based
methods significantly outperform the baseline methods in all simulation
settings. DRU variants significantly boost the worst group accuracy (the last
number in the accuracy lists (90,70,50) in Table 7) by up to 10% in some cases
compared to the baselines. Overall, qDRU+M is consistently the best DRU-
variant except for Setting 1, in which there is only a mild distribution
shift. Group DRO and JTT also perform as well as DRU-based methods in Setting
1, but their relative performance drops in settings with significant
distribution shifts. Interestingly, the constant upweighting method Const+M
performs even worse than ERM which doesn’t perform any weighting at all. When
significant distribution shifts are present (especially in setting 3 and 4
where unseen idiosyncratic signals are present in the test dataset), qDRU+M
not only improves the OOD worst group accuracy but also the 10th percentile of
worst groups, and it even introduces a considerable improvement (over 6%) in
average test accuracy compared to ERM, GroupDRO, and JTT baselines. The
performance of qDRU+G is notable in setting 4 where most distribution shift is
present, and its performance is only inferior to qDRU+M.


# Neural Network Architecture for Database Augmentation Using Shared Features

William C. Sleeman IV  Rishabh Kapoor  Preetam Ghosh  Virginia Commonwealth
University, Department of Radiation Oncology, Richmond, VA  Virginia
Commonwealth University, Department of Computer Science, Richmond, VA

###### Abstract

The popularity of learning from data with machine learning and neural networks
has lead to the creation of many new datasets for almost every problem domain.
However, even within a single domain, these datasets are often collected with
disparate features, sampled from different sub-populations, and recorded at
different time points. Even with the plethora of individual datasets, large
data science projects can be difficult as it is often not trivial to merge
these smaller datasets. Inherent challenges in some domains such as medicine
also makes it very difficult to create large single source datasets or multi-
source datasets with identical features. Instead of trying to merge these non-
matching datasets directly, we propose a neural network architecture that can
provide data augmentation using features common between these datasets. Our
results show that this style of data augmentation can work for both image and
tabular data.

###### keywords:

neural networks , machine learning , databases , autoencoders , classification

††journal: Journal of Biomedical Informatics

##  1 Introduction

Data analysis tools from machine learning and artificial intelligence are now
used to solve problems within almost every industry and problem domain.
Although more data is continuously being collected to support those advanced
methods, real world data science problems are often challenged with a limited
number of examples or missing features. For example, collecting medical data
is costly and time consuming as it requires patient consent, data security for
protecting privacy, and the need for subject matter experts. If these
challenges are addressed, collecting large patient cohorts still may not be
possible as inclusion for a given study may be restrictive or the medical
center may not treat enough such patients each year. Even outside of medicine,
many of the available datasets are still relatively small. Within the popular
UC Irvine Machine Learning Repository (UCI), half of the datasets contain
fewer than 1,600 examples [1].

To make the most of modern machine learning and deep learning algorithms,
quality data is needed for training and the more data available often produces
the best results. Very large, data-centric companies have the capability to
construct massive datasets but this is often not scalable to many problem
domains. Collecting, cleaning, and storing large datasets is expensive and
smaller companies or medical studies focusing on a single diagnosis simply may
have no way to generate large amounts of data. Aggregating across multiple
data sources is one solution to address small datasets but becomes non-trivial
when features do not match, some of the data is missing, or different feature
encoding is used. However, it is likely some of the features are present for
multiple datasets if they belong to the same problem domain. Medical data
often includes demographics information like age, sex, race, or diagnosis and
other domains may use industry standards resulting in feature overlap.

In addition to the common features, each dataset likely has other unique
features that make it difficult for direct aggregation. However, the shared
context between these datasets may provide enough information to transfer the
unique context between these non-matching datasets. Our approach uses
autoencoder networks to convert common features from a given dataset (A) to
its full complement of features. The same common features from a different
dataset (B) can then be passed through A’s network to generate synthetic
features for dataset A. The unique features from the new synthetic examples
can then be added to the existing features for database B, creating a new
augmented dataset.

The primary motivation for this work is to address the challenges faced when
trying to learn from relatively small medical datasets. As previously
mentioned, it is often impracticable to create large datasets focusing on
specific medical questions. There are many high quality medical datasets and
we hypothesized that predictive performance would be improved by including
information from other datasets based on the common features. In Section 5.3,
we discuss how this data augmentation method can be applied to real world
medical datasets.

In summary, we propose an autoencoder based approach for augmenting datasets
using common features. Our experiments include both CNN networks for images
and fully connected networks for tabular data, providing insight on how
feature sharing impacts performance. We also show that this method can improve
classifier performance, specifically with the use case of cancer datasets.

This work addresses the challenge of combining knowledge across disparate
datasets and our main contributions are the following:

  * 1.

Proposed architecture: We propose an autoencoder based solution for augmenting
dataset using common features. This provides a method to create synthetic
examples that are completely missing from a dataset which adds more context.

  * 2.

Study of the impact of feature sharing: This initial experimental study is
performed on both image and tabular based datasets and the impact of data
sharing is investigated.

  * 3.

Use case with clinical datasets: A real world example of cancer data is used
to show the proposed method’s performance.

  * 4.

Software: We provide a publicly available software package using Python with
Keras on GitHub for future research.

##  2 Related Works

One challenge with real world machine learning and deep learning projects is
the limited size and quality of training data. Some domains like medicine are
notorious for the difficulty of building large datasets with concerns of
privacy, data collection cost, and regulatory restrictions. Even when enough
data is present, issues like class imbalance can reduce model performance.

Data augmentation is often used to increase the size of a dataset or improve
data quality. One of the most common ways to augment image data is to apply
shifts, rotations, color adjustment, or zooming. This keeps the main concepts
of the image but introduces enough variation to help with overfitting while
providing an almost unlimited number of permutations. More recent advances in
image augmentation include patch cutout, blurring, and image mixing where two
training examples are blended using various methods [2, 3]. The XtremeAugment
method was also developed to generate new examples by adding one or more
training objects with adjusted color or viewpoint, all on existing or new
backgrounds [4].

Instead of perturbing existing data, completely new synthetic examples can
also be generated. Techniques like random oversampling, Synthetic Minority
Oversampling Technique (SMOTE) [5], including its many derived methods, can be
used to add more examples to the training data. Although those algorithms were
originally designed for traditional machine learning problems, DeepSMOTE [6]
was later created for deep learning problems. Unlike traditional data
augmentation and random oversampling, some of the SMOTE based methods can
create examples in specific portions of the feature space. This allows for
placing more focus on areas of interest like decision boundaries, safe
regions, or regions specifed by clustering algorithms [7, 8, 9]. Undersampling
can also be used for class balancing [10, 11] by reducing the size of the
majority class. Both over- and undersampling can be done within the same
algorithm shown by the Combined Synthetic Oversampling and Undersampling
Technique for Imbalanced Data Classification (CSMOUTE) method [12].

Another solution for creating larger datasets is to combine multiple smaller
dataset that include the same kind of information. However, these datasets may
not share the same features or have the same data ranges. Several works
proposed solutions to database merging including a method that combined
multiple image datasets using only their principal component analysis (PCA)
space and did not require the original training images to be kept [13]. PCA
was used for feature reduction which could aid in combining datasets with many
non-shared features [14]. Another work showed that singular value
decomposition (SVD) could be used to combine partially overlapping microarray
gene expression datasets [15].

The occurrence of missing feature values is a common challenge as one survey
suggested that over half of the UCI datasets had a missing feature rate of at
least 30% [16]. Imputation is often used to replace missing values which can
be based on methods like simple statistics, regression, hot-deck method,
clustering, and other machine learning algorithms [17]. Neural networks have
been used for imputation with both traditional autoencoders [18, 19] and GANs
[20].

##  3 Proposed Architecture

Applying traditional data augmentation methods may result in sub-optimal
results when faced with disparate datasets. Model generalization and class
imbalance can be partially addressed with synthetic examples generated from a
single source but this excludes any novel information present in other
relevant datasets. Feature reduction methods may remove critical relationships
within smaller classes or interesting regions of the feature space. To address
some of these limitations, we propose a neural network based method that
generates unseen features from common ones, thereby improving the predictive
power of learning algorithms.

The core component is an autoencoder network and we investigated the use of
both a traditional autoencoder (AE) and a variational autoencoder (VAE)
architecture. Figure 1 shows the general architecture used in the experiments
found in Section 4. Input data is passed through encoding layer and compressed
into the latent space. The primary difference between these two architectures
is the step right before the latent space layer as depicted in gray for AE and
blue for VAE. While the AE uses a layer like the prior encoder layers, the VAE
splits into mean ($\mu$) and variance ($\sigma^{2}$) sub-layers. A stochastic
sampling method is then used to produce the layer output when generating the
latent values. This also means that the same input data may result in
different outputs after the model is frozen unlike the traditional AE.

Figure 1: An example autoencoder network, broken down into the input, encoder,
decoder, and output sections. This architecture varies between the methods
used in the experimental study with AE shown in gray and VAE in blue.

If two datasets, referred to A and B, represent different aspects of the same
concept, the unique features from each dataset may be correlated. This
architecture uses the common features to generate synthetic replacements for
feature values that are only present in the other dataset. By combining the
existing and synthetic features, the new examples simulate the scenario where
all features came from a single complete data source. In the following
description of the processing steps, we are augmenting dataset A with the
unique features in dataset B. Below are the detailed steps of this process:

Identify Common Features: First, the features common between the two datasets
are identified. Min-max normalization is performed on the common features
across both datasets to ensure proper alignment and the same process is used
for the dataset specific columns resulting in a range of [0, 1]. Next, new
sub-datasets named CA and CB are extracted from the pre-processed A and B,
representing the common features of the examples present in each dataset.

Fit the Autoencoder Network with B: An autoencoder network is trained to map
CB data to the full compliment of features found in B.

Predict with B Data: The autoencoder is frozen and the CA is passed through
the CB-B network to predict the B features values.

Append Generated B Features: Database A is now augmented with the B synthetic
features, adding information from the similar, but independent, database B.

To better visualize this process, we show in Figure2 how the MNIST image data
was split by image columns to simulate the two dataset problem considered
above. Dataset A is given the left half and dataset B is given the right half
of the MNIST image columns.

Common features between these new datasets are simulated by reintroducing some
columns from the opposing dataset. To make the experiments more consistent,
the resulting training images are padded with zero columns so the final
dimensions are the same as in the original dataset. The example in Figure 2
shows this process with six common columns, with first row showing the dataset
A with the 14 left hand columns, the six common columns, and the result of
combining left hand and common columns with padding. The process is performed
for the right hand dataset B with a comparison to the original image in the
middle.

Figure 2: Example of how images could be split to simulate two different
datasets. Right hand (A) and its associated middle common columns (CA) go to
the masked database A image and the left hand (B) with its common columns (CB)
go to database B.

Figure 3 shows an entire diagram of the architecture. Using the split image
data described above, the common columns are used to train on the CB-B
network. Common columns from the test A dataset are encoded using the CB-B
network and decoded into synthetic B features. The original A data is then
combined with the synthetic features to create the augmented image. While the
resulting image is not perfect, it gets closer to replicating the original
image.

Figure 3: The proposed data augmentation network. Starting at the bottom,
common columns of the masked training example from dataset B are used to train
an autoencoder network. Test data from dataset A is processed using the CB-B
network to get synthetic B features. These new features are added to the A
masked test data to generate the augmented image.

##  4 Experimental Study

In this section, we apply the data augmentation process to both image and
tabular data using traditional and variational autoencoder networks. Source
code used for these experiments is provided at
https://github.com/fsleeman/database-augmentation-network which was written is
Python with Keras.

###  4.1 MNIST and CIFAR10

In these first experiments, we test the augmentation process using the MNIST
[21] and CIFAR10 [22] datasets. The image features were split using the same
process as shown in Figure 2 to simulate datasets with unavailable data. For
the purpose of these experiments, each image column was treated as if it was
as a single feature. The middle columns that overlap with the left and right
sides of the original images were marked as common features (shown in Figure 2
as CA and CB). An even number of common columns (n) were used to provided a
symmetry that allowed for the same experiments to be performed on both sides
of the images. We performed experiments between 2 and total columns – 2 for
both datasets, where MNIST had 28 total columns and CIFAR10 had 32. Future
experiments can investigate other combinations of feature sharing but the
small size of MNIST and CIFAR10 images limits the number of useful
combinations.

The impact of the data augmentation process was evaluated using a simple CNN
based classifier. This model used two CNN layers with max pooling, flattened,
a with dropout, a dense layer, one more dropout layer, and a final softmax
layer. Both MNIST and CIFAR10 datasets had ten classes and so argmax was used
to choose the predicted class. The only difference between how these datasets
were processed was that the MNIST networks used one color channel as included
gray scale images and the CIFAR10 used three color channels for its RGB color.
Twenty percent of the data was held out for testing and the remaining data was
used for training and validation. To better represent two completely
independent datasets, the training data was split, so database A got the first
half of the examples (by index) and database B got the second half. This meant
that there was no information overlap between those datasets including the
common features. Training was then performed using 5-fold cross validation and
the remaining test data was used to provide the results in Sections 5.1 and
5.2.

###  4.2 Lung Cancer Case Study

One of the major problems faced with clinically oriented data science projects
is the lack of data. While there are many high quality datasets in the field
of medicine, they often do not have the same features or formatting that would
allow them to be combined directly. One motivation of this work was to devise
an approach that would improve model performance with information from
independent medical datasets. In these experiments, we apply the proposed data
augmentation method on two disjointed lung cancer datasets.

The National Institutes of Health (NIH) has provided a number of cancer
related datasets including the Genomic Data Commons (GDC) [23] and the
Surveillance, Epidemiology, and End Results (SEER) Program [24]. Started in
2016, the GDC is a harmonized cancer dataset which combines data from several
modalities such as gene expression, mutations, pathology images, prescribed
drugs, and clinical outcomes across over eighty six thousand patients. SEER
has been collecting cancer data since 1973 from cancer registries across the
United States which currently has over fifteen million reported cases.

We have chosen to limit the data in our experiments to lung cancer because it
is one of the most common diagnoses and has a wide range of survival outcomes.
Unlike the SEER dataset which is in a tabular format, GDC is mostly made up of
other modalities which further reduced the amount of data used. After
filtering for lung cancer, samples with missing values were then removed.
There are many methods for addressing missing features but these examples were
removed because we wanted to experiment on the cleanest available data. Other
approaches in handling missing features should be further investigated.

The data used in the following experiments resulted in 522 examples for GDC
and 674,008 for SEER after the filtering and data cleaning. Seven common
features between the two cleaned datasets were then identified as shown in
Table 1: sex, year of diagnosis, age at diagnosis, race, International
Classification of Diseases version 10 (ICD-10) code, histology, and
laterality. Categorical features were split into multiple columns using one-
hot encoding and min-max scaling was performed across both datasets for
normalization. After one-hot encoding, there were a total of 27 individual
common features.

The unique features of the two datasets covered other relevant lung cancer
information that likely would add value to the learning process. GDC included
details such as smoking history, pathological staging, and ethnicity while
SEER included group staging, income, rural vs. urban locale, surgery type, the
number of tumors, among others. These unique features were also min-max
normalized, resulting in a total of 68 features for GDC and 78 for SEER.

SEER and NIH Common Features Feature Type Possible Values Sex Categorical
Female, Male Year of diagnosis Numerical 1957 to 2017 Age at diagnosis
Numerical 0 to 88 Race Categorical American Indian/Alaska Native, Asian or
Pacific Islander, Black, White, Unknown ICD-10 Code Categorical Main bronchus,
Upper lobe, Middle lobe, Lower lobe, Overlapping lesion, Lung (NOS) Histology
Categorical Adenocarcinoma, Bronchiolo-alveolar carcinoma, non-mucinou
Invasive mucinous adenocarcinoma, Adenocarcinoma with mixed subtypes Papillary
adenocarcinoma (NOS), Clear cell adenocarcinoma (NOS) Mucinous adenocarcinoma,
Signet ring cell carcinoma, Acinar cell carcinoma Laterality Categorical Left,
Right, Other

Table 1: A list of features common between the GDC and SEER datasets with
their data types and valid value ranges.

##  5 Discussion of Results

###  5.1 MNIST

The results in Table 2 show the $F_{1}$ classification performance for the
images with missing data and with both basic AE and VAE augmentation. Twenty
six total experiments were performed across thirteen common columns
combinations on both sides of the images with one of the augmentation methods
providing twenty of the best results. The basic AE usually did better with
fewer common columns compared with the VAE which may suggest that the more
complicated VAE network benefits more with additional data. While the overall
improvements of augmentation do not seem to be very significant, they still
represent a large portion of the remaining performance as most of the $F_{1}$
scores are already above 99 percent.

MNIST Dataset Common A Only A + B* AE A + B* VAE B Only A* + B AE A* + B VAE
Columns 2 97.16 97.17 97.08 96.50 96.42 96.30 4 97.70 97.66 97.62 97.48 97.42
97.31 6 98.02 98.12 98.04 98.10 98.18 98.04 8 98.42 98.46 98.36 98.57 98.73
98.58 10 98.61 98.70 98.64 98.87 98.93 98.92 12 98.86 98.86 98.93 99.07 99.11
98.99 14 98.96 99.02 98.97 99.09 99.14 99.08 16 99.11 99.08 99.09 99.15 99.15
99.17 18 99.11 99.08 99.09 99.11 99.14 99.17 20 99.10 99.05 99.15 99.11 99.14
99.20 22 99.11 99.03 99.10 99.13 99.15 99.16 24 99.07 99.05 99.13 99.15 99.13
99.17 26 99.08 99.10 99.08 99.09 99.10 99.16

Table 2: The $F_{1}$ scores for the MNIST classifier experiments with the left
and right hand data as the primary datasets. The $*$ symbol marks the
synthetic features that were generated with either the AE or VAE based
networks.

Figure 4 shows line graph of the performance of the MNIST images with missing
columns against the AE and VAE based augmentation methods. As expected,
including more common columns makes the problem easier for both the missing
data case and with augmentation but performance flattens out after enough data
is available. Because of how the handwritten digit were presented, there is
almost no useful information at the far left and right sides of the original
images and most of the information is in the middle. The first few common
columns are in this critical section which explains why adding just a few
columns makes a significant impact on the result. This is most apparent with
the AE plots where most of the gains were found within the first eight common
columns.

Figure 4: Plot showing the comparison of MNIST images with missing columns
against the AE and VAE based augmentation methods.

Figure 5 show examples of the generated digits for the A and B sides of the
images. In most of these cases, the AE and VAE algorithms did a reasonably
good job replacing the missing image columns. The largest discrepancy between
these methods was with the second to last image column for the number five.
This was a more difficult problem as the digit was poorly written and the
bottom semicircle was completely closed. The AE got much closer when
attempting to complete the right hand size of the image as the VAE may have
mistook the image as a partially completed eight. A similar problem occurred
when completing the left hand side as it could be seen as a six. While the VAE
did appear to make some of the clearest augmentations, its mistakes were more
pronounced but might be improved with larger training sets.

Figure 5: Comparison of augmented images for both the A and B side of the
MNIST dataset. In the augmentation results, A* and B* refers to their
corresponding synthetic representations. These images were generated using
eight common columns.

###  5.2 CIFAR10

Like with MNIST, Table 3 shows that the data augmentation process improved
$F_{1}$ scores for 23 of the 30 cases across the AE and VAE experiments. The
impact of augmentation was more pronounced with this dataset which may be
attributed to the increased complexity of the data. Many of the experiments
showed improvements in the range of 0.5 to 1.0%. The AE method gave most of
the top results, although the VAE outperformed the non-augmented data in the
majority of experiments. Augmentation tended to work better when there was
enough data to learn from and some critical information was still missing that
needed to be replaced. This ideal range was around 8 - 24 common columns for
CIFAR10.

CIFAR10 Dataset Common A Only A + B* AE A + B* VAE B Only A* + B AE A* + B VAE
Columns 2 64.50 64.38 64.13 64.15 64.07 63.67 4 65.95 65.52 65.96 65.37 65.15
64.83 6 66.14 66.01 66.33 65.30 65.83 65.70 8 66.80 66.99 67.48 66.02 66.75
66.43 10 67.32 67.01 67.80 66.47 66.89 66.68 12 67.68 68.07 67.54 67.19 67.51
67.36 14 68.07 68.74 68.14 67.23 67.82 67.69 16 68.94 68.88 68.62 67.98 68.46
68.80 18 68.58 68.95 69.10 68.63 68.60 68.31 20 68.69 69.26 69.13 68.59 69.02
68.58 22 69.29 69.50 69.18 68.96 69.44 69.04 24 69.39 70.03 69.92 69.09 70.11
69.38 26 69.80 69.49 70.38 69.66 69.51 69.77 28 69.81 70.14 70.10 69.99 69.66
69.92 30 70.12 70.32 70.38 69.90 69.77 69.79

Table 3: The $F_{1}$ scores for the CIFAR10 classifier experiments. The $*$
symbol marks the synthetic features that were generated with either the AE or
VAE based networks.

The value of data augmentation for CIFAR10 is more evident as shown in Figure
6. $F_{1}$ scores increase when more data is included but does not plateau
like with MNIST as there is useful data throughout the images. Both
augmentation methods improve scores in most cases and the relative performance
increase is present as more columns are added.

Figure 6: Plot showing the comparison of CIFAR10 images with missing columns
against the AE and VAE based augmentation methods.

Figure 7 show examples of the generated digits for the A and B sides of the
images. Although augmentation improved $F_{1}$ scores, the quality of the
generated image data was much worse than with MNIST. The CIFAR10 dataset is
not very large compared to recent deep learning image datasets and with its
size further reduced to create the non-overlapping A and B datasets. While
there may not be enough training data to produce high quality image data, it
was still enough to help the classifier. In some cases, such as the boats
depicted in columns two and tree of Figure 7, the new image data does look
like a very blurry version of the true image data. However, the frog images in
columns 5, 6, and 8 provide almost no information but may be attributed to the
diverse coloring, backgrounds, or orientations. Boats, on the other hand, tend
to look more similar as they often share common colors and backgrounds which
may require fewer training examples.

Figure 7: Comparison of augmented images for both the A and B side of the
CIFAR10 dataset. In the augmentation results, A* and B* refers to their
corresponding synthetic representations. These images were generated using
eight common columns.

###  5.3 GDC and SEER

As mentioned in Section 4.2, these experiments used 522 and 647k examples for
the GDC and SEER datasets respectively. Unlike the prior image based
experiments, the number of common columns were naturally presented as the two
datasets were already separated. Because the GDC dataset was much smaller, the
SEER dataset was randomly undersampled to evaluate the impact that the dataset
size has on the augmentation process. Classification was performed to predict
if a patient would survive 24 months or longer after lung cancer treatment.

Table 4 shows the $F_{1}$ scores for the experiments performed with different
level of SEER undersampling. Although using a high percentage of the SEER
dataset did not help the GDC classification, using 500 to 10k SEER examples
did. This resulted in both datasets being more balanced and using much more
SEER examples in the training process could make the latent space used for
decoding too cluttered to generate useful features. When using the 10k SEER
examples for training, the AE network increased $F_{1}$ score by almost 2.3%.

When augmenting the SEER dataset, the VAE outperformed the AE method in all
cases and improved the overall $F_{1}$ scores six out of eight times compared
to just using the original data. As expected, including more of the original
SEER data improved the non-augmented classification results. However,
augmentation started to consistently improve results when at least 10k SEER
examples were used.

These results show that cross-database feature augmentation can improve
classifier performance, especially when the larger dataset is the one being
augmented. One potential benefit from this kind of augmentation is that good
results may be achievable without the entire dataset. This can be more
important for cases where datasets are very large and are slow to train or
when it is difficult to acquire large training dataset, which is a common
challenge with medical data. Augmenting only 100k SEER examples with the small
GDC dataset provided better results than the entire SEER dataset without
augmentation.

The smaller GDC dataset got less benefit from augmentation but did see some
improvement with a smaller selection of SEER examples. Additional tuning of
the autoencoding network and feature engineering may result in a more
significant impact with even smaller datasets.

GDC and SEER Tabular Datasets SEER Count GDC GDC + SEER* AE GDC + SEER* VAE
SEER GDC* + SEER AE GDC* + SEER VAE 250 83.74 82.16 79.36 77.22 69.63 75.61
500 79.25 79.49 77.76 77.09 75.00 77.14 1k 82.62 83.00 79.61 77.36 73.41 75.82
10k 84.61 86.88 74.63 78.94 77.71 79.87 50k 89.29 89.58 82.45 81.83 81.67
82.96 100k 81.96 81.56 77.18 82.73 82.23 83.93 250k 86.01 84.36 83.38 82.21
82.00 83.61 500k 84.91 83.48 81.23 83.10 82.92 84.26 647k 84.83 80.33 83.27
83.42 83.11 84.60

Table 4: $F_{1}$ scores for the GDC and SEER experiments. Since the SEER
dataset is much larger than GDC, experiments were run with different subsets
with the total included examples in the SEER Count column. These two datasets
were augmented in the same manner for MNIST and CIFAR10. Figure 8: Plot
showing the comparison of augmentation with the GDC and SEER datasets using
the AE and VAE networks.

###  5.4 Future Work

In addition to the work presented in this paper, there are a number of other
topics that should be further investigated. The data cleaning process
performed for the GDC and SEER datasets removed any examples that had many
missing features and entire features that had a significant number of missing
values to ensure the cleanest possible data. However, some of this removed
data could be used if a more permissive imputation process was used such as
the traditional statistical based replacement or even within the autoencoder
process itself.

Class imbalance can negatively impact classifier performance as it can be
biased towards the majority class which can also affect deep learning problems
[25, 26, 27]. Although the MNIST and CIFAR10 datasets were mostly balanced,
the GDC and SEER datasets has a much higher imbalance with the two year
survival examples representing approximately 72% of the dataset. The proposed
feature generation system could also be extended to create entirely new
features for class balancing as previously done with a VAE based network [28].
The potential benefit of majority class undersampling, minority class
oversampling, or some combination of both could be further pursued using both
traditional machine learning and deep learning algorithms.

Instance level difficulty is another approach used to improve model
performance but was not considered in this work. The data augmentation method
may benefit by focusing more on specific types of examples. Most of the
current work on instance level difficulty has focused on a single dataset so
there is not much research on how the difficulty of examples between multiple
datasets would affect mutual encodings.

In this work, only one dataset was used to augment another but there may be
cases where the augmentation process could benefit from multiple datasets.
Each extra dataset could be used to generate different groups of features to
augment the target dataset. There could also be features common between
datasets A and B but different features common between B and C.

The simpler AE and VAE networks used in the experiments could be replaced with
deeper networks of the same type or more complicated architectures like
stacked autoencoders and GANs. These neural networks often perform better on
datasets with many more examples or features such as high resolution pictures,
3D medical images, or gene expression data. In addition to classification,
there are other types of machine learning algorithms that could benefit from
the proposed technique such as regression and clustering.

##  6 Conclusion

Within many problem domains, there are often many related datasets that do not
have matching features making the curation of large datasets difficult. In
this work, we have proposed an autoencoder based solution for data
augmentation using common features between these disparate datasets. As long
as these datasets are contextually similar, the information unique from each
dataset could be used to improve the performance of classifiers. We have shown
that this approach can work with both image and tabular data as well as
different types of autoencoder networks.

Although our experiments used relatively shallow autoencoders, they could be
replaced with much complicated or problem specific architectures such as GANs,
stacked autoencoders, and deep pre-trained models. In addition to the MNIST
and CIFAR10 examples, we showed that the data augmentation can work on real
world medical datasets. Even the small GDC dataset was able to provide a
benefit for classifying the much larger SEER dataset. This method is
especially useful for domains like medicine which often is limited to small
datasets collected as part of individual studies.

This initial work on data augmentation using common features has suggested
several new areas of research. This approach could be extended to include the
use of multiple datasets for augmentation, larger autoencoder style networks,
and different styles of imputation. Utilizing data properties like class
imbalance or instance level difficulty may provide further benefits.

## References

  * [1] D. Dua, C. Graff, UCI machine learning repository (2017).
URL http://archive.ics.uci.edu/ml

  * [2] D. Lewy, J. Mańdziuk, An overview of mixing augmentation methods and augmentation strategies, Artificial Intelligence Review (2022) 1–59.
  * [3] V. D. Jadhav, L. V. Patil, A study on medical image data augmentation using learning techniques, in: ICT Analysis and Applications, Springer, 2023, pp. 23–30.
  * [4] S. Nesteruk, S. Illarionova, T. Akhtyamov, D. Shadrin, A. Somov, M. Pukalchik, I. Oseledets, Xtremeaugment: Getting more from your data through combination of image collection and image augmentation, IEEE Access 10 (2022) 24010–24028.
  * [5] N. V. Chawla, K. W. Bowyer, L. O. Hall, W. P. Kegelmeyer, Smote: synthetic minority over-sampling technique, Journal of artificial intelligence research 16 (2002) 321–357.
  * [6] D. Dablain, B. Krawczyk, N. V. Chawla, Deepsmote: Fusing deep learning and smote for imbalanced data, IEEE Transactions on Neural Networks and Learning Systems (2022).
  * [7] H. Han, W.-Y. Wang, B.-H. Mao, Borderline-smote: a new over-sampling method in imbalanced data sets learning, in: International conference on intelligent computing, Springer, 2005, pp. 878–887.
  * [8] C. Bunkhumpornpat, K. Sinapiromsaran, C. Lursinsap, Safe-level-smote: Safe-level-synthetic minority over-sampling technique for handling the class imbalanced problem, in: Pacific-Asia conference on knowledge discovery and data mining, Springer, 2009, pp. 475–482.
  * [9] G. Douzas, F. Bacao, F. Last, Improving imbalanced learning through a heuristic oversampling method based on k-means and smote, Information Sciences 465 (2018) 1–20.
  * [10] W.-C. Lin, C.-F. Tsai, Y.-H. Hu, J.-S. Jhang, Clustering-based undersampling in class-imbalanced data, Information Sciences 409 (2017) 17–26.
  * [11] M. A. Arefeen, S. T. Nimi, M. S. Rahman, Neural network-based undersampling techniques, IEEE Transactions on Systems, Man, and Cybernetics: Systems (2020).
  * [12] M. Koziarski, Csmoute: Combined synthetic oversampling and undersampling technique for imbalanced data classification, in: 2021 International Joint Conference on Neural Networks (IJCNN), IEEE, 2021, pp. 1–8.
  * [13] G. N. Costache, P. Corcoran, P. Puslecki, Combining pca-based datasets without retraining of the basis vector set, Pattern Recognition Letters 30 (16) (2009) 1441–1447.
  * [14] T. Nguyen, R. Khadka, N. Phan, A. Yazidi, P. Halvorsen, M. A. Riegler, Combining datasets to increase the number of samples and improve model fitting, arXiv preprint arXiv:2210.05165 (2022).
  * [15] A. W. Schreiber, N. J. Shirley, R. A. Burton, G. B. Fincher, Combining transcriptional datasets using the generalized singular value decomposition, BMC bioinformatics 9 (1) (2008) 1–15.
  * [16] W.-C. Lin, C.-F. Tsai, Missing value imputation: a review and analysis of the literature (2006–2017), Artificial Intelligence Review 53 (2) (2020) 1487–1509.
  * [17] T. Emmanuel, T. Maupong, D. Mpoeleng, T. Semong, B. Mphago, O. Tabona, A survey on missing data in machine learning, Journal of Big Data 8 (1) (2021) 1–37.
  * [18] S. J. Choudhury, N. R. Pal, Imputation of missing data with neural networks for classification, Knowledge-Based Systems 182 (2019) 104838.
  * [19] B. K. Beaulieu-Jones, J. H. Moore, P. R. O.-A. A. C. T. CONSORTIUM, Missing data imputation in the electronic health record using deeply learned autoencoders, in: Pacific symposium on biocomputing 2017, World Scientific, 2017, pp. 207–218.
  * [20] C. Shang, A. Palmer, J. Sun, K.-S. Chen, J. Lu, J. Bi, Vigan: Missing view imputation with generative adversarial networks, in: 2017 IEEE International Conference on Big Data (Big Data), IEEE, 2017, pp. 766–775.
  * [21] Y. LeCun, The mnist database of handwritten digits, http://yann. lecun. com/exdb/mnist/ (1998).
  * [22] A. Krizhevsky, G. Hinton, et al., Learning multiple layers of features from tiny images (2009).
  * [23] R. L. Grossman, A. P. Heath, V. Ferretti, H. E. Varmus, D. R. Lowy, W. A. Kibbe, L. M. Staudt, Toward a shared vision for cancer genomic data, New England Journal of Medicine 375 (12) (2016) 1109–1112.
  * [24] National Cancer Institute, DCCPS, Surveillance Research Program, Surveillance, Epidemiology, and End Results (SEER) Program (www.seer.cancer.gov) SEER*Stat Database: Incidence - SEER Research Data, 8 Registries, Nov 2021 Sub (1975-2019) - Linked To County Attributes - Time Dependent (1990-2019) Income/Rurality, 1969-2020 Counties, released April 2022, based on the November 2021 submission.
  * [25] F. J. Pulgar, A. J. Rivera, F. Charte, M. J. del Jesus, On the impact of imbalanced data in convolutional neural networks performance, in: Hybrid Artificial Intelligent Systems: 12th International Conference, HAIS 2017, La Rioja, Spain, June 21-23, 2017, Proceedings 12, Springer, 2017, pp. 220–232.
  * [26] M. Buda, A. Maki, M. A. Mazurowski, A systematic study of the class imbalance problem in convolutional neural networks, Neural networks 106 (2018) 249–259.
  * [27] J. M. Johnson, T. M. Khoshgoftaar, Survey on deep learning with class imbalance, Journal of Big Data 6 (1) (2019) 1–54.
  * [28] Z. Wan, Y. Zhang, H. He, Variational autoencoder based synthetic data generation for imbalanced learning, in: 2017 IEEE symposium series on computational intelligence (SSCI), IEEE, 2017, pp. 1–7.