# Fair Representation Learning for Heterogeneous Information Networks

Ziqian Zeng2, Rashidul Islam1, Kamrun Naher Keya1, James Foulds1, Yangqiu
Song2, Shimei Pan1

###### Abstract

Recently, much attention has been paid to the societal impact of AI,
especially concerns regarding its fairness. A growing body of research has
identified unfair AI systems and proposed methods to debias them, yet many
challenges remain. Representation learning for Heterogeneous Information
Networks (HINs), a fundamental building block used in complex network mining,
has socially consequential applications such as automated career counseling,
but there have been few attempts to ensure that it will not encode or amplify
harmful biases, e.g. sexism in the job market. To address this gap, in this
paper we propose a comprehensive set of de-biasing methods for fair HINs
representation learning, including sampling-based, projection-based, and graph
neural networks (GNNs)-based techniques. We systematically study the behavior
of these algorithms, especially their capability in balancing the trade-off
between fairness and prediction accuracy. We evaluate the performance of the
proposed methods in an automated career counseling application where we
mitigate gender bias in career recommendation. Based on the evaluation results
on two datasets, we identify the most effective fair HINs representation
learning techniques under different conditions.

## Introduction

Many real-world social and information networks such as social networks,
bibliographic networks, and biological networks are heterogeneous in nature
Liu et al. (2016); Zhou et al. (2007); Xiong et al. (2019). We can model such
networks as Heterogeneous Information Networks (HINs) which contain diverse
types of nodes and/or relationships. For example, we can model the Twitter
social network as an HIN where the nodes are users or tweets, and the links
are the follower/following relations between users and the
authoring/retweeting relations between a user and a tweet.

To support large-scale heterogeneous information networks mining, much
attention has recently been paid to representation learning for HINs where
each node in a network is mapped to a dense vector in a low dimensional
embedding space which preserves the relationships between nodes and important
structural characteristics of the original networks Fu, Lee, and Lei (2017);
Dong, Chawla, and Swami (2017). Due to its robustness, flexibility, and
scalability, representation learning methods for HINs have been widely used to
support diverse network mining tasks such as node classification Dong, Chawla,
and Swami (2017), link prediction Wang et al. (2018), community detection
Cavallari et al. (2017), and recommender systems Shi et al. (2019).

Despite its popularity, little attention has yet been paid to the
understanding and mitigation of biases toward certain demographics in HINs
representation, such as gender and racial biases. As demonstrated in a wide
range of recent discoveries, machine learning systems trained with human-
generated content (e.g., social media data) frequently inherit or even amplify
human biases in the data Angwin et al. (2016); Dastin (2018); Noble (2020).
For example, word embedding models which inspired some of the early work on
network embedding such as DeepWalk Perozzi, Al-Rfou, and Skiena (2014), were
shown to exhibit female/male gender stereotypes to a disturbing extent (e.g.,
“man is to computer programmer as woman is to homemaker”) Bolukbasi et al.
(2016). To overcome this, there has been a concentrated recent effort in the
natural language processing community Bolukbasi et al. (2016); Caliskan,
Bryson, and Narayanan (2017); Gonen and Goldberg (2019) on understanding and
mitigating the biases in word embeddings. In the network mining community,
however, not much attention has been paid to the biases in HINs
representation. Since HINs representation may encode harmful societal
prejudice, it may cause unintended bias or unfairness in downstream
applications. Therefore, it is important to ensure HINs representation is
unbiased and applications incorporating these embeddings are fair and will not
negatively impact vulnerable people and marginalized communities in our
society.

In this paper, we propose a range of fair HINs representation learning
algorithms, including sampling-based, projection-based, and graph neural
network (GNNs)-based methods to mitigate demographic bias in HINs
representation. We systematically study the behavior of these algorithms,
especially their capability in balancing the trade-off between prediction
accuracy and fairness.

To evaluate our algorithms, we applied our fair HINs representation learning
techniques to automated fair career recommendation. Career counseling plays an
important role in many people’s lives. Unbiased career advice based on an
accurate assessment of one’s interests, skills, and personality can help them
make proper career choices. Good career choice in turn can boost their
economic success, social standing, and quality of life. Biased career
counseling however may restrict occupational opportunities and stunt the
career development of disadvantaged populations (e.g., girls and minorities)
Alshabani and Soto (2020).

The main contributions of this work include:

  * •

To the best of our knowledge, this is the first systematic investigation on
measuring and mitigating demographic bias in heterogeneous information
networks representation. Although prior work has studied mitigating bias in
network embedding Rahman et al. (2019), they focused on homogeneous networks
instead of heterogeneous information networks.

  * •

We propose a comprehensive suite of de-biasing algorithms ranging from
sampling-based, projection-based, to graph neural networks (GNNs)-based
techniques to mitigate demographic biases in HINs representation.

  * •

We demonstrate the effectiveness of the proposed methods in mitigating gender
bias in automated career counseling on two real-world datasets. Our results
illuminate the prediction accuracy vs. fairness trade-off behavior of these
algorithms, providing guidance to practitioners.

Our code is available at https://github.com/HKUST-KnowComp/Fair˙HIN.

The rest of the paper is organized as follows. We start with a brief survey of
related work, followed by a problem statement and the details of the proposed
fair HINs representation learning algorithms. In the next section, we discuss
the experiments designed to evaluate the effectiveness of the proposed methods
under various conditions. We conclude the paper by summarizing the main
findings and pointing out some future directions.

## Related Work

In this section, we survey the research areas relevant to our paper. First, we
summarize the main representation learning methods for heterogeneous
information networks since they are the basis of our de-biasing algorithms.
Next, we survey the general area of fair machine learning with a special focus
on two most relevant subareas: fair word embeddings and fair homogeneous
network embeddings.

### HIN Representation Learning

We categorize the typical HIN representation learning methods into two main
types. Task-agnostic HIN representation learning methods focus on training
general-purpose network embeddings that can be used to support different
network mining tasks. In contrast, task-specific HIN representation learning
methods learn representation via optimizing a particular downstream task.

#### Task-agnostic HIN Representation Learning

Task-agnostic HIN representation learning methods normally employ a meta-path
guided sampling method to generate meta-paths. After obtaining meta-paths,
typical word embedding algorithms such as word2vec Mikolov et al. (2013a, b)
are applied on the instantiated meta-path sequences to learn network
embeddings. A meta-path is a path consisting of a sequence of relations
defined between different node types, which is either specified manually or
derived from additional supervision Fu, Lee, and Lei (2017); Dong, Chawla, and
Swami (2017); Shi et al. (2018). For example, a meta-path author
$\xrightarrow{\texttt{write}}$ paper $\xrightarrow{\texttt{written by}}$
author in a bibliographic network represents a co-authorship relation.

#### Task-specific HIN Representation Learning

GNNs-based methods are end-to-end supervised approaches to learn network
embeddings. GNNs have the ability to aggregate local features, and to learn
highly expressive representations. A comprehensive survey on GNNs Wu et al.
(2021) shows that the majority of the current GNNs are designed for
homogeneous networks. Recently, Zhang et al. (2018) proposed a GNN method to
explore heterogeneous networks. They translated an HIN into multiple
homogeneous networks and applied a GNN to each homogeneous network and then
combined them at the final layer.

### Fair Machine Learning: General Approaches

There is a rising awareness that bias and fairness issues in machine learning
(ML) algorithms can cause substantial societal harm Angwin et al. (2016);
Buolamwini and Gebru (2018). A concentrated effort in the machine learning
community aims to address this problem. Existing methods on fair machine
learning can be summarized into three general strategies.

The first strategy employs fairness-aware pre-processing techniques to adapt
the training data, including (a) modifying the values of sensitive attributes
and class labels; (b) mitigating the dependencies between sensitive attributes
and class labels by mapping the training data to another space Dwork et al.
(2012); Feldman et al. (2015); Islam et al. (2019), and (c) learning a fair
representation of the training data that is independent of the protected
attribute Zemel et al. (2013); Xie et al. (2017).

The second strategy focuses on employing a fairness-guided optimization in
model training. A fairness metric-based penalty is added as a constraint or a
regularization term to the existing objective to enforce fairness Zafar et al.
(2017b, a); Calders and Verwer (2010); Islam et al. (2021). For example, Zafar
et al. (2017b) minimizes the covariance between the sensitive attributes and
the (signed) distance between feature vectors and the decision boundary of a
classifier.

The third strategy modifies posteriors to satisfy fairness constraints. For
example, Hardt, Price, and Srebro (2016) selected a threshold such that the
true positive rates of different groups are equal.

Furthermore, there are existing works using post-processing techniques to
correct the bias depending on the protected attribute. For example,
d’Alessandro, O’Neil, and LaGatta (2017) proposed to audit fairness of the
model on a held-out set to mitigate bias by altering some predicted outputs.

### Fair Word Embeddings

Word embedding models Mikolov et al. (2013b, a) learn a mapping of each word
in the vocabulary to a vector in an embedding space to encode semantic meaning
and syntactic structures of natural languages. These models are typically
trained using a neural network-based representation learning algorithm on word
co-occurrence data computed from massive text corpora. Since text data may
contain societal stereotypes such as racism or sexism, word embeddings also
typically inherit or amplify biases present in the data Bolukbasi et al.
(2016); Caliskan, Bryson, and Narayanan (2017); Papakyriakopoulos et al.
(2020).

The most popular method of de-biasing word embeddings is to project each word
embedding orthogonally to the bias direction Bolukbasi et al. (2016) followed
by a crowd-sourcing based correction. Instead of relying on the crowd to find
gendered words, alternatively, the demographic bias direction can simply be
computed using gendered person names Dev and Phillips (2019).

### Fair Homogeneous Network Embedding

Fair network embedding is a relatively new area that is now beginning to be
addressed. In an approach related to ours, Rahman et al. (2019) modified
random walks to learn fairness-aware embeddings. At each step, the algorithm
partitions neighbors into different groups based on the values of the
sensitive attribute. The system tries to give each group the same probability
of being selected regardless of its size. The method is applicable only to
homogeneous networks, while our work addresses heterogeneous information
networks. Bose and Hamilton (2019) proposed an adversarial framework to learn
fair graph embeddings. They trained an encoder to optimize a downstream task
(i.e., link prediction task in knowledge graphs) as well as to filter out
information about sensitive attributes. They also trained a discriminator to
predict the sensitive attribute from the graph embeddings. This algorithm can
be applied to both homogeneous and heterogeneous networks. But they did not
utilize the property of heterogeneous information networks to learn graph
embeddings and to enforce fairness.

## Problem Statement

We first describe our problem setting. We assume that our dataset is a
heterogeneous information network $G=(V,E,T,R)$ with $|T|>1$, where $V$
denotes a set of nodes, $E$ denotes a set of edges, $T$ denotes a set of node
types (e.g., user, career, and item), and $R$ denotes a set of edge relations
(e.g., a user likes a Facebook page, a user rates a movie, and a user chooses
a career). We also assume a binary protected attribute $a$, which pertains to
at least one type of nodes(e.g. gender for the users). An illustrative example
of a heterogeneous information network (Facebook career network) is shown in
Figure 1.

We focus on fairness in the general task of link prediction in heterogeneous
information networks, which we illustrate with an application to career
counseling. Our goal is to learn node embeddings $e_{v}$ for all nodes $v\in
V$ such that their use for link prediction on held-out edges leads to accurate
performance, and the system behaves in a fair manner with regard to the
protected demographics $a$. Fairness is a complex socio-technical issue. So
far many fairness definitions have been proposed in the AI community. Here we
adopt two of the most widely used AI fairness definitions and consider their
applications to link prediction or node classification in heterogeneous
information networks, e.g. for career counseling. Career counseling aims to
automatically provide a career choice for users based on their interests
(e.g., clicking “like” on Facebook pages).

Figure 1: An illustrative example of a Facebook career network. Different
shapes (colors) indicate different types of nodes (relations). A square
represents a page. A circle represents a user. A hexagon represents a career.
Orange line represents “like” relation. Blue line represents “choose”
relation.

Demographic Parity Demographic Parity is one of the most well-known criteria
for fairness in machine learning classifiers Dwork et al. (2012). It is
defined as:

$P(\hat{y}=k|a=0)=P(\hat{y}=k|a=1)\;\forall k,$

where $\hat{y}$ is a predicted label, $y$ is the ground truth class label, $a$
is the protected attribute such as gender (here we only consider binary
protected attributes), and $k$ represents the possible values of the class
label $y$. The extent to which the demographic parity criterion is violated is
measured by a distance metric between the two conditional distributions,
typically chosen to be the total variation distance (see Eq. (9)).

In the case of link prediction in heterogeneous information networks, we adapt
the definition by letting $P(\hat{y}=k,a=i)$ be the empirical probability of
predicting a node as class label $k$ (e.g. a career such as computer science)
and the node is of protected type $i$ (e.g. a female user). Based on this
definition, in career counseling, to achieve perfect demographic parity, the
probability of recommending each career (e.g. computer science) given male
users should be the same as that given female users.

Equal Opportunity Another popular fairness definition is equal opportunity
Hardt, Price, and Srebro (2016). It is defined as:

$P(\hat{y}=k|a=0,y=k)=P(\hat{y}=k|a=1,y=k)\;\forall k\mbox{ .}$

We adapt this to link prediction in heterogeneous information networks similar
to demographic parity. we adapt the definition by letting
$P(\hat{y}=k,a=i,y=k)$ be the empirical probability of predicting a node as
class label $k$ (e.g. a career such as computer science) and the ground truth
is class label $k$ and the node is of protected type $i$ (e.g. a female user).
Based on this definition, in career counseling, to achieve perfect equal
opportunity, for a subgroup of people who indeed have chosen a particular
career (e.g., computer science), the probability of the system to recommend
this career (e.g., computer science) given male users should be the same as
that given female users.

## Methodology

In this section, we present three classes of fair HIN representation learning
methods: (1) fairness-aware sampling which extends existing sampling methods
for HIN embedding to mitigate bias, (2) embedding projection which employs a
vector-space projection operation to reduce biases in the embeddings, and (3)
graph neural networks (GNNs) with a fairness-aware loss, which trains a GNN to
optimize both prediction accuracy and fairness.

To illustrate these algorithms, we use fair career counseling as an example
application. We formulate fair career counseling as a fair link prediction in
HINs. Here, the nodes in the HIN includes a user (e.g., a Facebook user or a
movie reviewer), an item (e.g., a Facebook page or a movie) and a career
(e.g., the career concentration or occupation declared by a user). A user node
is linked to an item node via a “like” link, indicating a user’s preference
for the item. A user node can also be linked to a “career” node via a “choose”
link. The goal of automated career recommendation is to predict whether there
should be a link between a user and a career based on the interests/likes of a
user. In the following, we propose a range of methods for fair HIN
representation learning.

### Fairness-aware Sampling

To learn embeddings in a HIN, we often use meta-paths to guide a random walker
to generate traversal paths. A meta-path is a path consisting of a sequence of
relations defined between different node types. For example, user
$\xrightarrow{\texttt{choose}}$ career $\xrightarrow{\texttt{chosen by}}$ user
represents two users choosing the same career. After obtaining meta-paths, we
can use a meta-path based HIN representation learning algorithm such as
metapath2vec (M2V) to learn the node embeddings. M2V incorporates the
heterogeneous network structures into the skip-gram model Mikolov et al.
(2013b, a). Specifically, the objective function of M2V is defined as follows,

$\mathop{\rm argmax}_{\theta}\sum_{v\in V}\sum_{t\in T}\sum_{c_{t}\in
N_{t}(v)}\log p(c_{t}|v;\theta),$ (1)

where $N_{t}(v)$ denotes node $v$’s neighbors which are of type $t$, and
$p(c_{t}|v;\theta)$ is commonly defined as a softmax function. The objective
function of M2V is similar to that of the skip-gram model except that M2V is
sensitive to the type of node.

After learning user embeddings and career embeddings, we then train a
multilayer perceptron (MLP) to predict the career, where the input is a user
embedding and the career embeddings learned by M2V.

#### Traditional Meta-path Generation

A meta-path $\mathcal{P}$ is a path defined on the graph $G=(V,E,T,R)$, and is
denoted in the form of
$V_{1}\xrightarrow{R_{1}}V_{2}\xrightarrow{R_{2}}\cdots\xrightarrow{R_{l}}V_{l+1}$.
This meta-path defines a composite relation $R=R_{1}\circ
R_{2}\circ\cdots\circ R_{l}$ between type $1$ and type $l+1$, where $\circ$
denotes the composition operator on relations, and $V_{t}$ denotes all nodes
with type $t$.

Here we show how to use meta-paths to guide heterogeneous random walkers. The
transition probability at step $i$ is defined as follows,

$P(v^{i+1}|v^{i}_{t})=\begin{cases}\frac{1}{|N_{t+1}(v^{i}_{t})|},&(v^{i+1},v^{i}_{t})\in
E,\phi(\cdot)=t+1\\\ 0,&(v^{i+1},v^{i}_{t})\in E,\phi(\cdot)\neq t+1\\\
0,&(v^{i+1},v^{i}_{t})\not\in E\end{cases},$ (2)

where $v^{i}_{t}\in V_{t}$, and $N_{t+1}(v^{i}_{t})$ denotes $v^{i}_{t}$’s
neighbors which are of type $t+1$, and $\phi(\cdot)$ is a function which
returns the type of the node $v^{i+1}$. The transition probability function
shows that the random walker only goes to a certain type of nodes and the type
of next step is guided by meta-paths.

We manually define two meta-paths for our career counseling application.
career $\xrightarrow{\texttt{chosen by}}$ user $\xrightarrow{\texttt{like}}$
item $\xrightarrow{\texttt{liked by}}$ user $\xrightarrow{\texttt{choose}}$
career represents two careers chosen by two different users who like the same
item (e.g., a Facebook page, or a movie). user $\xrightarrow{\texttt{like}}$
item $\xrightarrow{\texttt{liked by}}$ user represents two users like the same
item.

#### Fair Meta-path Generation

Assume $a$ is a protected binary attribute and $i$ is a value of $a$, $g_{i}$
represents a group of users satisfying $a=i$. If $a$ is a binary attribute, we
define an advantaged group if it contains a larger number of training examples
than the other group (the disadvantaged group). Fair meta-path generation aims
to up-sample the disadvantaged group with higher probability while down-sample
the advantaged group. The sampling probability is inversely proportional to
the number of users in each group. As the number of users in the disadvantaged
group is less than that in the advantaged group, the disadvantaged group has a
higher probability to be sampled. As a result, the system is less likely to
neglect the disadvantaged group.

Next, we present the details of the algorithm. Fair meta-path sampling occurs
at the step where the current node type is career ($C$) and the next node type
is user ($U$). Assume $a$ is a protected attribute of a user (e.g., gender or
race). To simplify our explanation, we assume $a$ is binary. Based on the
values of $a$, we can cluster all the users into different groups (e.g.,
$g_{0}$ and $g_{1}$). The unnormalized transition probability at step $i$
where $t=C$ and $t+1=U$ is defined as follows:

$P(v^{i+1}|v^{i}_{C})=\begin{cases}\frac{r}{|N(v^{i}_{C},U,g_{0})|},&cond,\pi(v^{i+1})=g_{0}\\\
\frac{r}{|N(v^{i}_{C},U,g_{1})|},&cond,\pi(v^{i+1})=g_{1}\\\
0,&(v^{i+1},v^{i}_{C})\in E,\phi(\cdot)\neq U\\\ 0,&(v^{i+1},v^{i}_{C})\not\in
E\end{cases},$ (3)

where $\phi(\cdot)$ is a function which returns the type of node $v^{i+1}$,
and $\pi(\cdot)$ is a function which returns the value of the protected
attribute(e.g., the gender of a user) of node $v^{i+1}$; $cond$ is the
condition where $(v^{i+1},v^{i}_{C})\in E,\phi(\cdot)=U$, which means the
current node is a career node and the next node $v^{i+1}$ is a user node, and
they are connected; $N(v^{i}_{C},U,g_{0})$ denotes the neighbors of
$v^{i}_{C}$ which are of type $U$ and belong to group $g_{0}$, and
$N(v^{i}_{C},U,g_{1})$ denotes the neighbors of $v^{i}_{C}$ which are of type
$U$ and belong to group $g_{1}$; and $r$ is a hyper-parameter that can
determine to what extend the random walker selects the disadvantaged group in
order to correct the bias. Note that if $|N(v^{i}_{C},U,g_{0})|=0$ or
$|N(v^{i}_{C},U,g_{1})|=0$, then the corresponding unnormalized probability
should be zero. The transition probability in the second and third conditions
remains the same as that in Eq. (2).

### Embedding Projection

The second debiasing method is inspired by recent work on attenuating bias in
word vectors Dev and Phillips (2019). It is often used as a post-processing
step. We adapt this method to de-bias HINs embeddings. Basically, after
obtaining HINs embeddings, we can reduce their unfairness by eliminating the
effect of a protected attribute (e.g., gender) in learned user embeddings via
vector projection. The main difference between our method and the projection
method used in Dev and Phillips (2019) is the bias direction in the embedding
space.

Dev and Phillips (2019) computed a “bias direction” for word embeddings based
on the difference in the average embeddings of male and female names. In our
case, let $v_{g_{i}}$ be the direction of group $g_{i}$. We compute
$v_{g_{i}}$ by averaging all the embeddings of the users in group $g_{i}$,

$v_{g_{i}}=\frac{e_{u_{1}}+e_{u_{2}}+\cdots+e_{u_{n_{i}}}}{\|e_{u_{1}}+e_{u_{2}}+\cdots+e_{u_{n_{i}}}\|},$
(4)

where $e_{u_{1}}$, $e_{u_{2}}$, $e_{u_{n_{i}}}$ are the user embeddings in
$g_{i}$, and there are $n_{i}$ users in group $g_{i}$. Assuming the protected
attribute is binary, we can compute the bias direction $v_{b}$ using

$v_{b}=\frac{v_{g_{0}}-v_{g_{1}}}{\|v_{g_{0}}-v_{g_{1}}\|}\mbox{ .}$ (5)

To reduce bias in user embeddings, we project each user’s vector $e_{u}$
orthogonally to the bias direction $v_{b}$ as follows,

$e_{u}^{\prime}=e_{u}-<e_{u},v_{b}>v_{b},$ (6)

where $v_{b}$ is the bias direction, $<\cdot,\cdot>$ is the inner product
operation, and $e_{u}^{\prime}$ is the resulting de-biased user embedding.

### Fairness-aware Graph Neural Networks (GNNs)-based Representation Learning

In this section, we explore task-specific representation learning methods,
i.e., graph neural networks (GNNs). GNNs learn the node representation via
optimizing a downstream task (e.g., node classification) in a supervised
manner. In the career counseling application, task-agnostic representation
learning methods (e.g., M2V) consider that a career is a node and users nodes
are linked to career nodes. While GNNs methods consider that a career is the
label of a user node, so there are only two types of nodes (users and items)
in the input for GNNs models. Task-agnostic representation learning methods
consider career recommendation as a link prediction task, GNNs methods
formulate it as the node classification task. To reduce bias, we directly
incorporate a fairness-aware loss in addition to an accuracy-based objective.
In the following, we describe our method to make GNNs representation fair.

#### Graph Neural Networks

Graph Neural Networks (GNNs) Kipf and Welling (2017); Hamilton, Ying, and
Leskovec (2017b); Velickovic et al. (2018) are powerful tools for graph
mining. It has gained increasing popularity in various applications, including
social networks, knowledge graphs, recommender systems, and biomedical
research.

Let $G=(V,E)$ denote a graph with node features $X_{v}$ for $v\in V$. Given a
set of nodes $\\{v_{1},...,v_{n}\\}$ and their labels $\\{y_{1},...,y_{n}\\}$,
the task of GNNs is to learn a representation vector $h_{v}$ for node $v$ that
helps to predict the label of the node $v$. The prediction is denoted as
$\hat{y}$. In GNNs model, the representation of node $v$ is iteratively
updated by aggregating the representations of $v$’s neighboring nodes. After
$l$ iterations of aggregation, $v$’s representation captures the structural
characteristics of its $l$-hop network neighborhood. Formally, the
representation of node $v$ in $l$-th layer in a GNN is formulated as follows,

$h^{l}_{v}=Combine\big{(}h^{l-1}_{v},Agg(h^{l-1}_{v},h^{l-1}_{u})\big{)},$ (7)

where $h^{l}_{v}$ is the representation of node $v$ at the $l$-th
iteration/layer, and $u$ are neighboring nodes of $v$, and $Combine$ is a
function to combine the information of node $v$ and its neighboring nodes in
$l-1$ th layer, and $Agg$ is an aggregation operation which aggregates
neighbors of $v$. Commonly used aggregator operations are mean aggregator,
LSTM aggregator, and pooling aggregator. We initialize $h^{0}_{v}=X_{v}$. The
accuracy loss function of GNNs is:

${\mathcal{L}}_{acc}=-\sum_{i}^{N}{\log P(\hat{y_{i}}=y_{i})}\mbox{ .}$ (8)

#### Fairness-aware Loss

There are many fairness definitions. Here we only consider demographic parity
Dwork et al. (2012) and equal opportunity Hardt, Price, and Srebro (2016).

We introduce some notations first. $N^{i}$ denotes the number of users in
protected group $g_{i}$. Here, we only consider binary protected attributes,
hence, $i\in\\{0,1\\}$. $N^{i}_{k}$ denotes the number of samples which are
predicted to be $k$ and are in the protected group $g_{i}$. $N^{i}_{k,k}$
denotes the number of samples which are predicted correctly and are in
protected group $g_{i}$. The computation of empirical demographic parity and
equal opportunity is:

$\begin{split}\textit{diff}_{dp}&=\sum_{k}{\bigg{|}\frac{N^{0}_{k}}{N^{0}}-\frac{N^{1}_{k}}{N^{1}}\bigg{|}}\end{split},$
(9)

$\begin{split}\textit{diff}_{eo}&=\sum_{k}{\bigg{|}\frac{N^{0}_{k,k}}{N^{0}_{k}}-\frac{N^{1}_{k,k}}{N^{1}_{k}}\bigg{|}}\end{split}.$
(10)

The above computation is based on counting, which cannot be used as a loss
function, because the gradient cannot be back propagated. Hence we use the
probabilistic prediction $p(\hat{y}|x)$ of the model to replace the hard
count. $x$ is a sample in the dataset. Specifically, in career counseling
application, $x$ is a user. Hence, we define the fairness-aware loss for GNN
as follows,

$\begin{split}{\mathcal{L}}_{dp}&=\sum_{k}\bigg{(}\frac{\sum_{x\in
D:a=0}{P(\hat{y}=k|x)}}{N^{0}}\\\ &\quad\quad\quad-\frac{\sum_{x\in
D:a=1}{P(\hat{y}=k|x)}}{N^{1}}\bigg{)}^{2},\end{split}$ (11)

$\begin{split}{\mathcal{L}}_{eo}&=\sum_{k}\bigg{(}\frac{\sum_{x\in
D:a=0,y=k}{P(\hat{y}=k|x)}}{N^{0}_{k}}\\\ &\quad\quad\quad-\frac{\sum_{x\in
D:a=1,y=k}{P(\hat{y}=k|x)}}{N^{1}_{k}}\bigg{)}^{2},\end{split}$ (12)

where $x$ denotes a sample; $D:a=i$ is the set of samples where their values
of the attribute are $i$; $D:a=i,y=k$ is the set of samples where their values
of the attribute is $i$ and ground truth is $k$. We compute two final loss
functions in our GNN models. The first one is a demographic parity based
fairness-aware loss function:

${\mathcal{L}}_{acc}+\lambda_{dp}\cdot{\mathcal{L}}_{dp},$ (13)

where $\lambda_{dp}$ is a trade-off hyper-parameter.

The second one is an equal opportunity based fairness-aware loss function:

${\mathcal{L}}_{acc}+\lambda_{eo}\cdot{\mathcal{L}}_{dp},$ (14)

where $\lambda_{eo}$ is a trade-off hyper-parameter.

## Experiments

In this section, we study the behavior of the proposed de-biasing methods for
HINs representation and their applications in automated career counseling. We
tested the their performance on two datasets: a Facebook dataset and a
MovieLens dataset.

### Datasets

Dataset | Facebook | MovieLens
---|---|---
# careers | 48 | 14
# items (FB pages or movies) | 99,756 | 3,677
# users | 7,069 | 4,920
# male users | 2,721 | 3,558
# female users | 4,332 | 1,362
avg users per career | 62.81 | 93.47
avg users per item | 14.04 | 222.40
avg items per user | 198.15 | 166.21
Table 1: Statistics of Facebook and MovieLens datesets.

The Facebook dataset used in the study was collected as a part of the
myPersonality project Kosinski et al. (2015). The data was collected with an
explicit opt-in consent for reuse for research purposes. To protect privacy,
the data were anonymized.

The Facebook dataset contains rich information about a Facebook user such as
his/her demographics (e.g., gender), the Facebook pages he/she likes, and
his/her declared career concentrations (e.g., music, computer science,
psychology). Since Facebook “likes” contain rich information about a person’s
interests and preferences of a wide range of items/topics (e.g., books, music,
celebrities, brands, and TV shows), they can be used to suggest possible
career concentrations for a new user.

As the myPersonality dataset is no longer publicly available, to facilitate
research reproducibility, we employed a second dataset, the movieLens dataset,
which is publically available. It contains similar information to the Facebook
dataset such as the gender of a movie reviewer, the movies he/she likes, and
his/her occupation.

Table 1 shows the statistics of each dataset after we clean the data. The
Facebook network created from the myPersonality dataset consists of $7,069$
user nodes, $99,756$ item nodes (a.k.a FB page), and $48$ career nodes. If a
user likes a page, there is an edge between the user and the page. If a user
declared a career concentration on Facebook, there is an edge between the user
and the career.

The MovieLens-1M dataset originally contains $6,040$ users, $3,900$ movies,
and $19$ careers. After we remove some careers that our system should not
recommend such as “retired,” “unemployed,” and “K-12 student,” the resulting
MovieLens network consists of $4,920$ user nodes, $3,677$ movie nodes, and
$14$ career nodes. Similarly, there is an edge between the user and the movie
if a user rated a movie. If a user has declared an occupation, there is an
edge between the user and the occupation.

In both datasets, each user only has one career. In the Facebook dataset, we
filtered out users for which binary gender information was not provided by the
user. All users in MovieLens had binary gender information provided. We were
unable to study our method’s impact on non-binary identifying users as non-
binary gender information was not provided.

Since our datasets are not very large, to fully utilize our ground truth data,
we employed nested cross validation for model training, testing, and hyper-
parameter tuning. Specifically, we first split the whole dataset into two
parts, where 40% of the data were used for embeddings training, and the
remaining 60% were used for career prediction. For the data reserved for
career prediction, we split it further into three folds, with one of the folds
for testing, and the other two folds for training and validation. We further
split the training and validation data into four folds with one of the folds
as a validation set, and the other three folds for training. The proportion of
embeddings training, career prediction training, career prediction validation,
and career prediction test was 4:3:1:2.

For GNN based methods, to facilitate result comparison, the test sets were the
same as those used in the sampling-based methods. Since GNNs do not need any
data to train embeddings, the data for embeddings training (i.e., 40% of the
whole data) was added into the career prediction training data, so the
proportion of training, validation, and test for career prediction was 7:1:2.

### Evaluation Metrics

We use one prediction accuracy measure and two fairness measures to assess the
performance of all the methods.

Accuracy Measure: we use Mean Reciprocal Rank (MRR) as the measure of
prediction accuracy. It is a statistical measure for evaluating an ordered
list of items. MRR is computed as the mean of the multiplicative inverse of
the rank of the correct answers. MRR is frequently used in information
retrieval and recommender systems to assess the system output. To compute MRR,
let $rank_{i}$ denote the rank of the ground truth (i.e, a user’s career
choice), and $N$ denote the number of samples. Then
$MRR=\frac{1}{N}\sum_{i}^{N}{1/rank_{i}}$.

Fairness Measures: we use two widely used fairness measures in our evaluation:
Demographic Parity and Equal Opportunity, which are defined in Eq. (9) and Eq.
(10).

### Compared Methods

We have implemented a total of eight different methods. Among them, one is
traditional HINs embedding methods without gender de-biasing (M2V), and one is
graph neural network methods without gender de-biasing (GNN). The rest are
fairness-aware methods. Among the seven fairness-aware methods, one is a naive
baseline (called balance-data). One is the adversarial baseline proposed by
Bose and Hamilton (2019). The rest of the methods employ a single or a
combination of the de-biasing methods we proposed. The details of each method
are described below.

balance-data: For each career, we randomly remove users from the advantaged
group so that we have an equal number of male and female users in the
embedding training data. In this way, we deliberately create a balanced
dataset to remove gender bias. Note that only the embeddings training set is
balanced. The career prediction training/validation/test sets remain the same
to facilitate fair comparison. After creating a balanced dataset, we use the
metapath2vec algorithm Dong, Chawla, and Swami (2017) to train embeddings.

adversarial: Bose and Hamilton (2019) proposed an adversarial framework to
learn fair graph embeddings. They trained an encoder to learn graph embeddings
as well as to filter out information about sensitive attributes. They also
trained a discriminator to predict the sensitive attribute from the graph
embeddings. The encoder we used in the experiment is the GraphSaint model Zeng
et al. (2019). The discriminator we used in the experiment is an MLP which has
the same architecture in Bose and Hamilton (2019). Similar to GNNs-based
method, we formulate the career prediction task to a node classification task.

M2V: we use the traditional meta-path generation algorithm in Eq. (2) to
generate meta-paths, and learn user embeddings and career embeddings using
metapath2vec algorithm Dong, Chawla, and Swami (2017).

M2V + fair sampling: we use fair meta-path generation algorithm shown in Eq.
(3) to generate meta-paths, and the rest of the processes are the same as M2V.

projection: we calculate the bias direction using Eq. (5) and project all the
user embeddings to the direction that is orthogonality to the bias direction.
All methods which learn embeddings explicitly can use the projection method as
a post-processing step.

GNN: we use the GraphSaint algorithm Zeng et al. (2019), a state-of-the-art
GNN algorithm to predict career. We have also tried GraphSAGE Hamilton, Ying,
and Leskovec (2017a). However, the results were much worse than GraphSaint. So
in this paper, we only report the results using GraphSaint. We formulate the
career prediction task to a node classification task, where the career choice
of a user is the label of the user node.

GNN-demographic-parity: we use Eq. (13) which combines the demographic-parity
based fairness-aware loss and the accuracy loss to train the GraphSaint model
Zeng et al. (2019).

GNN-equal-opportunity: we use Eq. (14) which combines the equal-opportunity
based fairness-aware loss with the accuracy loss to train the GraphSaint model
Zeng et al. (2019).

For the sampling-based based methods above, after we learn the embeddings, we
train a multilayer perceptron (MLP) to predict the career. The inputs are the
user embeddings and the career embeddings learned by sampling-based methods.
Since the output of the softmax layer in MLP is a probability distribution, we
rank all career candidates based on their corresponding probabilities, and use
mean reciprocal rank as the measure of prediction accuracy.

For all sampling-based methods, we used the same hyper-parameters listed
below. The dimension of embeddings was $128$; the size of negative samples was
$5$; the context window size was $5$. For all GNNs-based methods and the
encoder in the adversarial method, we used two convolutional layers. The
dimension of the hidden layer is $128$. The features of nodes are generated
using the following procedure. Each Facebook page or movie is associated with
a title or description. We average the word embeddings of all the words in the
title or description as its features. For each user, we average the features
of all the linked items (e.g., FB page or movie) as his/her features. For each
career, we average the features of all linked users as its features. The
dimension of the features was $50$.

(a) MRR vs demographic parity on Facebook

(b) MRR vs equal opportunity on Facebook

(c) MRR vs demographic parity on MovieLens

(d) MRR vs equal opportunity on MovieLens

Figure 2: The Pareto frontier of each method to demonstrate the accuracy and
fairness tradeoff.

### Model Selection

Many of the methods used in our experiment have hyper-parameters that we need
to tune. For M2V, we tune the number of walks and the length of each walk. For
fair-sampling based methods, we tune the number of walks, the length of a
walk, and the ratio $r$ in Eq. (3). For the projection method, the hyper-
parameters are the same as M2V or M2V + fair-sampling. For GNN-demographic-
parity and GNN-equal-opportunity, the hyperparameters are $\lambda_{dp}$ and
$\lambda_{eo}$ respectively. For the adversarial method, the hyperparameter is
the trade-off coefficient $\gamma$ between the accuracy loss and fairness loss
in the encoder.

We used grid search to tune hyper-parameters for sampling-based methods, where
the range of the number of walks and the length of each walk were
$\\{10,50,100,200\\}$, $r\in\\{1,2,\cdots,10\\}$ and the search step size was
$1$. We used grid search to tune hyper-parameters for GNNs-based methods,
where $\lambda_{dp},\lambda_{eo}\in\\{10,20,\cdots,100\\}$ and the search step
size was $10$. We used grid search to tune hyper-parameters for the
adversarial method, where $\gamma\in\\{0.1,0.2,\cdots,1.0\\}$ and the search
step size was $0.1$.

### Experimental Results

To study the performance of different fair representation learning algorithms
for HINs, especially their ability to balance the trade-off between prediction
accuracy and fairness, we plot the results on the test datasets in Figure
2(a)-2(d). The Y-axis of each chart represents a fairness measure (either
demographic parity in Figure 2(a) and 2(c) or equal opportunity in Figure 2(b)
and 2(d)), while the X-axis represents the prediction accuracy (MRR). Since
some algorithms such as sampling-based methods and fairness-aware GNNs-based
approaches have hyper-parameters that can be tuned to achieve different
fairness and accuracy tradeoffs, we show the Pareto frontier of each method
which consists of models that are not dominated by other alternatives from
that method. We say that a model $A$ dominates an alternative model $B$ if
model $A$ outscores model $B$ regardless of the trade-off between fairness and
accuracy – that is, if $A$ is better than $B$ in both fairness and accuracy.

As shown in these charts, the traditional GNN method (represented as the
_black star_) has relatively good MRR but poor fairness as it only tries to
optimize prediction accuracy. The second baseline method that employs a naive
data balancing technique to remove bias (represented as the _black square_)
achieved moderate prediction accuracy as well as moderate fairness. The
balance-data model performed better on the MovieLens dataset than on the
Facebook dataset because the Facebook dataset is more biased and thus a lot of
data has to be removed from the advantaged group to achieve balance. Among the
rest of the models, some methods such as the traditional sampling-based method
(represented by the green crosses) and the fairness-aware sampling-based
method (represented by red dots) are capable of achieving high MRR at the cost
of low fairness (for both demographic parity and equal opportunity, the lower
the value is, the fairer it is). Other methods, such as GNN-demographic
parity/equal opportunity (represented by the blue dots/blue diamonds in the
charts) and a combination of fair sampling-based and projection-based methods
can achieve the best fairness and a reasonably good prediction accuracy.

Method | $dp_{LF}$ | $dp_{MF}$ | $dp_{HF}$ | $eo_{LF}$ | $eo_{MF}$ | $eo_{HF}$
---|---|---|---|---|---|---
balance-data | 0.2973 | 0.2973 | 0.2951 | 0.2898 | 0.2661 | 0.2513
adversarial | 0.3158 | 0.3158 | 0.3158 | 0.3155 | 0.3149 | 0.2960
M2V | 0.3428 | 0.3120 | 0.2951 | 0.3171 | 0.2750 | 0.2636
M2V + fair-sampling | 0.3466 | 0.3303 | 0.2856 | 0.3250 | 0.2916 | 0.2780
M2V + projection | 0.3177 | 0.3177 | 0.3177 | 0.3217 | 0.3088 | 0.2647
M2V + fair-sampling + projection | 0.3354 | 0.3351 | 0.3267 | 0.3284 | 0.3023 | 0.2918
GNN | 0.3341 | – | – | 0.3341 | – | –
GNN-demographic-parity | 0.3400 | 0.3279 | 0.3277 | – | – | –
GNN-equal-opportunity | – | – | – | 0.3382 | 0.3266 | 0.3125
Table 2: Comparison of different HINs representation learning methods for career prediction on the Facebook dataset. Mean Reciprocal Rank (MRR) is reported under different demographic parity (dp)/equal opportunity (eo) constraints. We report mean of 5 runs. Here, LF, MF, and HF represent the low fairness, medium fairness and high fairness conditions.  Method | $dp_{LF}$ | $dp_{MF}$ | $dp_{HF}$ | $eo_{LF}$ | $eo_{MF}$ | $eo_{HF}$
---|---|---|---|---|---|---
balance-data | 0.4022 | 0.3780 | 0.3279 | 0.4088 | 0.3868 | 0.3119
adversarial | 0.3929 | 0.3922 | 0.3849 | 0.3914 | 0.3860 | 0.3734
M2V | 0.4061 | 0.3840 | 0.3209 | 0.4104 | 0.3901 | 0.3160
M2V + fair-sampling | 0.4097 | 0.3946 | 0.3265 | 0.4114 | 0.3912 | 0.3151
M2V + projection | 0.4064 | 0.4056 | 0.3857 | 0.4066 | 0.3977 | 0.3356
M2V + fair-sampling + projection | 0.4143 | 0.3967 | 0.3371 | 0.4152 | 0.3949 | 0.3268
GNN | 0.4105 | – | – | 0.4105 | – | –
GNN-demographic-parity | 0.3932 | 0.3925 | 0.3881 | – | – | –
GNN-equal-opportunity | – | – | – | 0.3915 | 0.3886 | 0.3795
Table 3: Comparison of different HINs representation learning methods for
career prediction on the MovieLens dataset. Mean Reciprocal Rank (MRR) is
reported under different demographic parity (dp)/equal opportunity (eo)
constraints. We report mean of 5 runs. Here, LF, MF, and HF represent the low
fairness, medium fairness and high fairness conditions.

To further illustrate each model’s ability in balancing the trade-off between
accuracy and fairness and facilitate model comparison, we fixed the fairness
dimension and only compare the prediction accuracy of these models. We choose
three reference fairness thresholds to illustrate model performance under
three conditions: high fairness (HF), medium fairness (MF), and low fairness
(LF). The HF condition simulates a real-world scenario where a system is used
in making consequential decisions (e.g., sentencing). Thus under this
condition, we may want to choose a model with high fairness. The low fairness
condition simulates a real-world scenario where the fairness of a system is
not consequential (e.g., for entertainment).

The reference fairness thresholds were based on the baseline GNN method. Since
traditional GNN performed poorly on the fairness dimension, we use its
fairness performance as the threshold to simulate the LF condition. The MF
condition is 75% of demographic parity or equal opportunity value achieved by
traditional GNN. The HF condition is 50% of demographic parity or equal
opportunity value achieved by traditional GNN. Once we select the three
fairness thresholds, we report the performance of the model with the highest
MRR among all the models satisfying the fairness constraint. We report the
mean MRR of 5 runs.

As shown in Table 2, on the Facebook dataset, under the LF condition and the
demographic parity measure, traditional methods without any de-biasing such as
M2V and GNN perform quite well. This is not surprising since the LF condition
puts relatively few fairness constraints on the systems. Thus systems that
totally ignore the fairness constraints (e.g., M2V and GNN) performed quite
well. M2V + fair-sampling even outperforms traditional methods slightly. It
demonstrates that M2V + fair-sampling does not harm accuracy under a LF
condition. Under the MF condition and the demographic parity measure, M2V +
fair-sampling and M2V + fair-sampling + projection outperform two baselines
(i.e., M2V and adversarial) by 2%. It demonstrates that our proposed fair-
sampling methods are effective under a MF condition. Under the HF condition
and the demographic parity measure, GNN-demographic-parity outperforms M2V and
adversarial by 3% and 1% respectively. It demonstrates that our proposed
fairness-aware loss is effective under a HF condition. Under the equal
opportunity measure, GNN-equal-opportunity is the best under all conditions.
It outperforms baselines significantly. It demonstrates the effectiveness of
fairness-aware loss under the equal opportunity measure. Under all fairness
conditions and the equal opportunity measure, M2V + fair-sampling outperforms
M2V significantly. It shows that the fair-sampling technique can improve M2V
under all fairness conditions if the measure is equal opportunity. In summary,
under all fairness conditions, M2V + fair-sampling + projection perform well
on demographic parity measure, and GNN-equal-opportunity performed well on
equal opportunity measure.

The results on the MovieLens dataset are slightly different. First, all
methods perform quite well under the LF condition. M2V + fair-sampling +
projection outperforms M2V and adversarial by 1% and 2% respectively. Under
the MF condition, M2V + projection outperforms M2V and adversarial by 2% and
1% respectively if the measure is demographic parity. Under the MF condition,
M2V + projection outperforms M2V and adversarial by 1% if the measure is equal
opportunity. Under the HF condition, GNN-demographic-parity and GNN-equal-
opportunity outperform M2V by 6% on demographic parity measure and equal
opportunity measure respectively. They are comparable with the adversarial
method. In summary, under LF and MF fairness conditions, projection-based
methods perform well. Under HF fairness conditions, GNN-demographic-parity and
GNN-equal-opportunity methods perform well.

To summarize, based on results on Facebook and MovieLens dataset, M2V, GNN,
and fair-sampling based methods can be used if under the LF condition. Under
the HF and MF condition, project-based methods(e.g., M2V + projection and M2V
+ fair-sampling + projection), and GNN with fairness-aware loss methods (e.g.,
GNN-demographic-parity and GNN-equal-opportunity) consistently performed well.

## Conclusion

In this paper, we systematically investigated a wide range of de-biasing
algorithms for HINs representation, ranging from sampling-based, projection-
based, to graph neural networks (GNNs)-based approaches. We evaluated the
effectiveness of these methods in an automated career counseling task where we
mitigate harmful gender bias in career recommendation. Based on the evaluation
results on two datasets, we identified different algorithms that are effective
under different conditions, which provides valuable guidance to practitioners.
One possible downstream application for fair HINs representation learning is
communities recommendation in e-commerce platforms and social platforms.
Unbiased communities recommendation relies on users’ interests rather than
users’ sensitive attributes. We leave it to future work.

## References

  * Alshabani and Soto (2020) Alshabani, N.; and Soto, S. 2020.  Early 20th‐Century Career Counseling for Women: Contemporary Practice and Research Implications.  _Career Development Quarterly_ 68(1): 78 – 93.
  * Angwin et al. (2016) Angwin, J.; Larson, J.; Mattu, S.; and Kirchner, L. 2016.  Machine bias: There’s software used across the country to predict future criminals. and it’s biased against blacks.  _ProPublica, May_ 23\.
  * Bolukbasi et al. (2016) Bolukbasi, T.; Chang, K.-W.; Zou, J. Y.; Saligrama, V.; and Kalai, A. T. 2016.  Man is to computer programmer as woman is to homemaker? debiasing word embeddings.  In _NeurIPS_ , 4349–4357.
  * Bose and Hamilton (2019) Bose, A.; and Hamilton, W. 2019.  Compositional fairness constraints for graph embeddings.  In _ICML_ , 715–724.
  * Buolamwini and Gebru (2018) Buolamwini, J.; and Gebru, T. 2018.  Gender shades: Intersectional accuracy disparities in commercial gender classification.  In _FAT_ , 77–91.
  * Calders and Verwer (2010) Calders, T.; and Verwer, S. 2010.  Three naive Bayes approaches for discrimination-free classification.  _Data Mining and Knowledge Discovery_ 21(2): 277–292.
  * Caliskan, Bryson, and Narayanan (2017) Caliskan, A.; Bryson, J. J.; and Narayanan, A. 2017.  Semantics derived automatically from language corpora contain human-like biases.  _Science_ 356(6334): 183–186.
  * Cavallari et al. (2017) Cavallari, S.; Zheng, V. W.; Cai, H.; Chang, K. C.-C.; and Cambria, E. 2017.  Learning community embedding with community detection and node embedding on graphs.  In _CIKM_ , 377–386.
  * d’Alessandro, O’Neil, and LaGatta (2017) d’Alessandro, B.; O’Neil, C.; and LaGatta, T. 2017.  Conscientious classification: A data scientist’s guide to discrimination-aware classification.  _Big data_ 5(2): 120–134.
  * Dastin (2018) Dastin, J. 2018.  Amazon scraps secret AI recruiting tool that showed bias against women.  _Reuters_ .
  * Dev and Phillips (2019) Dev, S.; and Phillips, J. 2019.  Attenuating bias in word vectors.  In _AISTATS_ , 879–887.
  * Dong, Chawla, and Swami (2017) Dong, Y.; Chawla, N. V.; and Swami, A. 2017.  metapath2vec: Scalable representation learning for heterogeneous networks.  In _KDD_ , 135–144.
  * Dwork et al. (2012) Dwork, C.; Hardt, M.; Pitassi, T.; Reingold, O.; and Zemel, R. 2012.  Fairness through awareness.  In _ITCS_ , 214–226.
  * Feldman et al. (2015) Feldman, M.; Friedler, S. A.; Moeller, J.; Scheidegger, C.; and Venkatasubramanian, S. 2015.  Certifying and removing disparate impact.  In _KDD_ , 259–268.
  * Fu, Lee, and Lei (2017) Fu, T.-y.; Lee, W.-C.; and Lei, Z. 2017.  Hin2vec: Explore meta-paths in heterogeneous information networks for representation learning.  In _CIKM_ , 1797–1806. ACM.
  * Gonen and Goldberg (2019) Gonen, H.; and Goldberg, Y. 2019.  Lipstick on a pig: Debiasing methods cover up systematic gender biases in word embeddings but do not remove them.  In _NAACL-HLT_ , 609–614.
  * Hamilton, Ying, and Leskovec (2017a) Hamilton, W.; Ying, Z.; and Leskovec, J. 2017a.  Inductive representation learning on large graphs.  In _NeurIPS_ , 1024–1034.
  * Hamilton, Ying, and Leskovec (2017b) Hamilton, W. L.; Ying, R.; and Leskovec, J. 2017b.  Representation Learning on Graphs: Methods and Applications.  _IEEE Data Engineering Bulletin_ 40(3): 52–74.
  * Hardt, Price, and Srebro (2016) Hardt, M.; Price, E.; and Srebro, N. 2016.  Equality of opportunity in supervised learning.  In _NeurIPS_ , 3315–3323.
  * Islam et al. (2021) Islam, R.; Keya, K.; Zeng, Z.; Pan, S.; and Foulds, J. 2021.  Debiasing Career Recommendations with Neural Fair Collaborative Filtering.  In _WWW_.
  * Islam et al. (2019) Islam, R.; Keya, K. N.; Pan, S.; and Foulds, J. 2019.  Mitigating Demographic Biases in Social Media-Based Recommender Systems.  In _KDD Workshop on Social Impact Track_.
  * Kipf and Welling (2017) Kipf, T. N.; and Welling, M. 2017.  Semi-supervised classification with graph convolutional networks.  In _ICLR_.
  * Kosinski et al. (2015) Kosinski, M.; Matz, S. C.; Gosling, S. D.; Popov, V.; and Stillwell, D. 2015.  Facebook as a research tool for the social sciences: Opportunities, challenges, ethical considerations, and practical guidelines.  _American Psychologist_ 70(6): 543.
  * Liu et al. (2016) Liu, M.; Liu, S.; Zhu, X.; Liao, Q.; Wei, F.; and Pan, S. 2016.  An Uncertainty-Aware Approach for Exploratory Microblog Retrieval.  _IEEE Transactions on Visualization and Computer Graphics_ 22(1): 250–259.
  * Mikolov et al. (2013a) Mikolov, T.; Chen, K.; Corrado, G.; and Dean, J. 2013a.  Efficient estimation of word representations in vector space.  _arXiv preprint arXiv:1301.3781_ .
  * Mikolov et al. (2013b) Mikolov, T.; Sutskever, I.; Chen, K.; Corrado, G. S.; and Dean, J. 2013b.  Distributed representations of words and phrases and their compositionality.  In _NeurIPS_ , 3111–3119.
  * Noble (2020) Noble, S. 2020.  _Algorithms of Oppression: How Search Engines Reinforce Racism_ , volume 22.
  * Papakyriakopoulos et al. (2020) Papakyriakopoulos, O.; Hegelich, S.; Serrano, J. C. M.; and Marco, F. 2020.  Bias in word embeddings.  In _FAT_ , 446–457.
  * Perozzi, Al-Rfou, and Skiena (2014) Perozzi, B.; Al-Rfou, R.; and Skiena, S. 2014.  Deepwalk: Online learning of social representations.  In _KDD_ , 701–710.
  * Rahman et al. (2019) Rahman, T.; Surma, B.; Backes, M.; and Zhang, Y. 2019.  Fairwalk: towards fair graph embedding.  In _IJCAI_ , 3289–3295.
  * Shi et al. (2019) Shi, C.; Hu, B.; Zhao, W. X.; and Yu, P. S. 2019.  Heterogeneous Information Network Embedding for Recommendation.  _IEEE Transactions on Knowledge and Data Engineering_ 31(2): 357–370.
  * Shi et al. (2018) Shi, Y.; Gui, H.; Zhu, Q.; Kaplan, L.; and Han, J. 2018.  Aspem: Embedding learning by aspects in heterogeneous information networks.  In _ICDM_ , 144–152.
  * Velickovic et al. (2018) Velickovic, P.; Cucurull, G.; Casanova, A.; Romero, A.; Liò, P.; and Bengio, Y. 2018.  Graph Attention Networks.  In _ICLR_.
  * Wang et al. (2018) Wang, H.; Zhang, F.; Hou, M.; Xie, X.; Guo, M.; and Liu, Q. 2018.  Shine: Signed heterogeneous information network embedding for sentiment link prediction.  In _ICWSM_ , 592–600.
  * Wu et al. (2021) Wu, Z.; Pan, S.; Chen, F.; Long, G.; Zhang, C.; and Philip, S. Y. 2021.  A comprehensive survey on graph neural networks.  _IEEE Transactions on Neural Networks and Learning Systems_ 32(1): 4–24.
  * Xie et al. (2017) Xie, Q.; Dai, Z.; Du, Y.; Hovy, E.; and Neubig, G. 2017.  Controllable invariance through adversarial feature learning.  In _NeurIPS_ , 585–596.
  * Xiong et al. (2019) Xiong, Y.; Guo, M.; Ruan, L.; Kong, X.; Tang, C.; Zhu, Y.; and Wang, W. 2019.  Heterogeneous network embedding enabling accurate disease association predictions.  _BMC Medical Genomics_ 12(10): 186.
  * Zafar et al. (2017a) Zafar, M. B.; Valera, I.; Gomez Rodriguez, M.; and Gummadi, K. P. 2017a.  Fairness beyond disparate treatment & disparate impact: Learning classification without disparate mistreatment.  In _WWW_ , 1171–1180.
  * Zafar et al. (2017b) Zafar, M. B.; Valera, I.; Rodriguez, M. G.; and Gummadi, K. P. 2017b.  Fairness constraints: Mechanisms for fair classification.  In _AISTATS_ , 962–970.
  * Zemel et al. (2013) Zemel, R.; Wu, Y.; Swersky, K.; Pitassi, T.; and Dwork, C. 2013.  Learning fair representations.  In _ICML_ , 325–333.
  * Zeng et al. (2019) Zeng, H.; Zhou, H.; Srivastava, A.; Kannan, R.; and Prasanna, V. 2019.  Graphsaint: Graph sampling based inductive learning method.  In _ICLR_.
  * Zhang et al. (2018) Zhang, Y.; Xiong, Y.; Kong, X.; Li, S.; Mi, J.; and Zhu, Y. 2018.  Deep collective classification in heterogeneous information networks.  In _WWW_ , 399–408.
  * Zhou et al. (2007) Zhou, D.; Orshanskiy, S. A.; Zha, H.; and Giles, C. L. 2007.  Co-ranking Authors and Documents in a Heterogeneous Network.  In _ICDM_ , 739–744.


# Stop Wasting My Time!
Saving Days of ImageNet and BERT Training
with Latest Weight Averaging

Jean Kaddour
Centre for Artificial Intelligence
University College London
<EMAIL_ADDRESS>

###### Abstract

Training vision or language models on large datasets can take days, if not
weeks. We show that averaging the weights of the $k$ latest checkpoints, each
collected at the end of an epoch, can speed up the training progression in
terms of loss and accuracy by dozens of epochs, corresponding to time savings
up to ~$68$ and ~$30$ GPU hours when training a ResNet50 on ImageNet and
RoBERTa-Base model on WikiText-103, respectively. We also provide the code and
model checkpoint trajectory to reproduce the results and facilitate research
on reusing historical weights for faster
convergence111https://github.com/jeankaddour/lawa.

##  1 Introduction

The arsenal of deep learning methods (e.g., architectures, regularizers, pre-
trainers, etc.) has been growing rapidly; for the last decade, thousands of
them have been proposed yearly. Arguably, many are brittle and not as
universally effective as initially claimed [24]. One way to filter what really
works is by testing methods on large datasets. For example, for vision tasks,
methods that demonstrated success on ImageNet have often proven to be
successful in other tasks [1].

Large datasets, however, require access to expensive multi-GPU machines to
enable data parallelism and reasonable training durations. Less well-funded
researchers do not have access to supercomputers, and lengthy training runs
make quick, iterative experimentation of research ideas difficult. Simple
task-, model-, and optimizer-agnostic methods that can be easily added to
existing training pipelines and speed up training time have the potential to
make deep learning research more accessible.

In the 90s, Polyak & Juditsky [33] studied how to accelerate the convergence
speed of stochastic gradient descent in the convex loss function regime. They
proved that the running average of the model weights iterates
$\bar{\bm{\theta}}=\frac{1}{t}\sum_{i=1}^{t}\bm{\theta}_{i}$ converges to the
loss minimizer $\bm{\theta}^{*}$ asymptotically with the highest possible
rate. When visualizing a convex loss function, the geometric intuition is
simple: whenever the optimizer oscillates around a minimum, the average of the
iterates will be closer to it.

However, in deep learning, loss functions are highly non-convex [8]. Weight
averaging has been mainly used to improve the model’s generalization
performance at the end of or after training [14, 16].

#### Contribution

We revisit weight averaging applied to neural networks from a convergence
speed perspective. Inspired by Li et al. [22], we focus on the _middle_ stage
of training: after the dramatic changes of the local loss landscape during the
very first training steps [10, 7] but _before_ the optimizer converges.
Because the weights still undergo substantial change in that middle phase,
averaging _all_ collected models, e.g., by maintaining a moving average [14,
16], can be sub-optimal. Therefore, we propose to average the $k$ latest
checkpoints (each collected at the end of an epoch) throughout training, which
we refer to as _LAtest Weight Averaging_ (LAWA).

Figure 1: LAWA speeds up convergence on ImageNet and (Ro)BERT(a) training. We
highlight the GPU hours saved by LAWA as the longest time until the baseline
optimizer (SGD or Adam) matches LAWA ’s performance. We plot the losses and
top-5 accuracies in Appendix A.

##  2 LAtest Weight Averaging (LAWA)

Algorithm 1 Pseudocode in PyTorch style

⬇

# k: number of latest checkpoints

ckpts = []

lawa_model = copy.deepcopy(model)

for epoch in range(num_epochs):

for (x,y) in train_loader:

train_step(x, y, model, optimizer)

ckpts.append(get_params(model))

if epoch + 1 >= k:

update_lawa_model(lawa_model, ckpts)

del ckpts[0]

def update_lawa_model(lawa_model, ckpts):

for p, avg_p in zip(list(lawa_model.parameters()),torch.mean(ckpts)):

p.copy_(avg_p)

The key idea is to collect model checkpoints once at the end of each epoch in
a queue. LAWA ’s solution at the end of epoch $E$ is
$\bm{\theta}^{\text{LAWA\,}}_{E}:=\frac{1}{k}\sum_{i=E-k+1}^{E}\bm{\theta}_{i}$.

Requirements include few training loop modifications, as shown in Algorithm 1,
and additional memory. In practice, we store the checkpoints in RAM or on disk
and only transfer them to the GPU once we want to evaluate
$\bm{\theta}^{\text{LAWA\,}}$. To improve the time complexity of the averaging
operation, one can use a circular queue 222Coding interview preparers might
remember this Leetcode problem..

The number of latest weights $k$ is a hyper-parameter, and we achieve good
results across both experiments with default value $k=6$, as shown in Figure
2. However, we observe that averaging too many checkpoints ($k>16$) results in
worse performance.

(a) ResNet50 / ImageNet.

(b) RoBERTa-Base / WikiText103.

Figure 2: LAWA is fairly insensitive w.r.t. $k$, but high $k$ can hurt its
performance early on.

The averaging coefficients can also follow a different pattern, e.g., in
Appendix C, we experiment with an exponential moving average (assigning higher
weights to the more recent checkpoints) and observe that this works worse than
uniform coefficients. One can also learn the averaging coefficients [38, 22],
but for simplicity and to avoid additional computational costs, we do not do
so.

The checkpoint saving frequency might be thought of as a hyper-parameter;
however, in this work, we always set it to one epoch. When there is so much
data that we are in a sub-one-epoch training regime [18, 13], we may collect
checkpoints every $\nu$ steps and need to tune $\nu$. Another heuristic might
be to collect a checkpoint whenever the validation loss has not improved [28].

If the network includes batch norm layers, then their statistics for
$\bm{\theta}^{\text{LAWA\,}}_{E}$ are unknown. Prior work [14] has suggested
computing them by an inference pass through the training dataset. We do not
observe a large effect of doing so compared to simply copying
$\bm{\theta}_{E}$’s statistics, possibly because we only average the $k$
latest weights instead of keeping one running average over many epochs [14].

##  3 Results

We run all experiments on a machine with 4x NVIDIA 3090s and report its wall-
clock time.

###  3.1 Image Classification: ResNet50 on ImageNet

We consider the ImageNet $1000$-classes classification task [4], which
includes $1.28$M training images and $50$k validation images. To train a
ResNet50 [12], we use the official PyTorch implementation [32] and train for
$90$ epochs using SGD with a momentum value of $0.9$ and a cosine learning
rate schedule. Our 4-GPU machine takes ~$26$min for one epoch. For
$\bm{\theta}^{\text{LAWA\,}}$, we re-compute the batch norm layer statistics
with a full inference pass through the training dataset before evaluating
$\bm{\theta}^{\text{LAWA\,}}$.

In Figure 1(a), we observe that LAWA reaches a high accuracy dramatically
faster, e.g., validation accuracy of around ~$66$% (the final accuracy is
~$76$%) is reached ~$40$ epochs (~$68$ hours) earlier than the baseline
optimizer (SGD). However, we also note that its head start decreases towards
the end of the training, and the highest reached accuracy is not reached much
earlier. This observation raises the question of whether we can use LAWA to
“jump forward” and continue the training from $\bm{\theta}^{\text{LAWA\,}}$ to
reach the optimal accuracy faster, which we further discuss in Section 4.

###  3.2 Masked Language Modeling: RoBERTa-Base on WikiText-103

Next, we pre-train a (Ro)BERT(a)-Base [5, 25] model with masked language
modeling (MLM) objective on the WikiText-103 dataset [27] with $103$M and
$218$k tokens for training and validation set, respectively. We follow the
training recipe provided by fairseq [31]: We train with Adam [20] for $200$
epochs, using a batch size of $2048$, a polynomial learning rate decay with
$10$k warmup steps and a peak learning rate of $0.0005$. Our 4-GPU machine
needs ~$10$min for one training epoch.

In Figure 1(b), we report the training and validation MLM cross-entropy (CE)
losses as a function of the number of training steps (as typically done in
NLP). We observe that LAWA consistently improves the losses, and it reaches
Adam’s final best validation loss ~$45$ epochs ahead, saving ~$30$ GPU hours.
Interestingly, $\bm{\theta}^{\text{LAWA\,}}$’s final validation performance is
noticeably better than $\bm{\theta}^{\text{ADAM}}$’s, confirming previous
results on improved generalization obtained with weight averaging [14, 16].

##  4 Future Work

#### Continuing training from $\bm{\theta}^{\text{LAWA\,}}$.

It is tempting to think that we may “jump forward” training by applying the
LAWA procedure and then continue training from there if some target accuracy
has not been reached yet. One issue is that we would need to adjust the
learning rate each time we “jump”. In practice, we may not know by how much
(if at all) we accelerated the training progression. Hence, it remains unclear
how to adjust a learning rate scheduler or the state variables of an adaptive
optimizer.

####  $k$ scheduler.

In Figure 2, we observe that at different times, different $k$ values perform
better; e.g., during the end of the training, higher $k$ performs better;
motivating a scheduler for $k$.

#### Accelerating training from the very beginning.

We focus on speeding up training during the middle stage of training: after
the first training steps but long before the optimizer converges. The reason
for that is that in the very early training phase, the gradient typically
moves with large magnitudes until it converges to a smaller subspace of the
loss function’s Hessian, in which it then remains over long periods of
training [10, 7] (middle stage). We empirically confirm that averaging during
the early phase worsens the baseline’s performance, as can be seen in Figure
4.

#### Combining LAWA with other acceleration techniques.

As we will discuss in the next section, there are several other techniques
available to accelerate neural network training. For example, the SAM
optimizer [6] can accelerate training too [37], and Kaddour et al. [16] show
that SAM combined with weight averaging can further boost the final test
performance.

#### Relationships between LAWA and optimization hyper-parameters.

For example, SGD becomes unstable for certain learning rates [39, 15]; can we
similarly characterize when LAWA is effective?

#### Applying other operations to a set of checkpoints.

For example, by learning a hyper-network [11] that takes in one or more
checkpoints and predicts the model parameters at later training stages.

#### When does it not work?

LAWA may not always cause speed-ups because Kaddour et al. [16] reported some
negative results on using weight averaging to improve the model’s final
performance.

##  5 Related work

The idea of weight averaging is not novel; it has been studied widely in
linear settings [33, 30, 21].

Szegedy et al. [35] used weight averaging to create the GoogLeNet model,
which, at that time, set a new state of the art in the ImageNet 2014 challenge
[4]. Izmailov et al. [14] introduce _Stochastic Weight Averaging_ (SWA), a
weight averaging strategy starting from pre-trained models to move them to
better-generalizing regions in the same loss basin. Kaddour et al. [16]
extensively study SWA’s effectiveness, including non-typical domains like
graph-structured data, and suggest combining it with SAM [6] to boost its
final performance further. Wortsman et al. [38] propose to average weights of
multiple models with different hyper-parameter configurations. All three works
(i) average weights toward the end or even after convergence, (ii) focus on
the models’ final test performances, and (iii) incorporate one moving-averaged
model, while we show in Figure 2 that too large $k$ can result in suboptimal
results, especially at earlier training times.

This work is heavily inspired by Li et al. [22]’s _Trainable Weight Averaging_
(TWA), who propose to learn averaging coefficients for training speed-ups.
Concurrently, Guo et al. [9] observe that running the SWA procedure multiple
times accelerates convergence. In some sense, LAWA generalizes their procedure
by keeping an average of the $k$ latest checkpoints instead of running SWA
sequentially. Another related optimizer utilizing an auxiliary set of “fast
weights” before updating the weights of interest is the _Lookahead_ (LA)
optimizer [41]. We compare LAWA and LA in Appendix B.

Another line of work has shown that training data re-weighting can speed up
training. For example, some re-weighting methods focus on proxy models [2,
29], importance sampling [3, 19] or removing spurious correlations [17, 36].

## Acknowledgements

I thank Matt J. Kusner and Mingtian Zhang for feedback and fruitful
discussions. I acknowledge support from the Engineering and Physical Sciences
Research Council with grant number EP/S021566/1.

## References

  * Beyer et al. [2020] Beyer, L., Hénaff, O. J., Kolesnikov, A., Zhai, X., and Oord, A. v. d.  Are we done with imagenet?  _arXiv preprint arXiv:2006.07159_ , 2020.
  * Coleman et al. [2019] Coleman, C., Yeh, C., Mussmann, S., Mirzasoleiman, B., Bailis, P., Liang, P., Leskovec, J., and Zaharia, M.  Selection via proxy: Efficient data selection for deep learning.  _arXiv preprint arXiv:1906.11829_ , 2019.
  * Csiba & Richtárik [2018] Csiba, D. and Richtárik, P.  Importance sampling for minibatches.  _The Journal of Machine Learning Research_ , 19(1):962–982, 2018.
  * Deng et al. [2009] Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pp. 248–255. Ieee, 2009.
  * Devlin et al. [2018] Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Foret et al. [2021] Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B.  Sharpness-aware minimization for efficiently improving generalization.  In _9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021_. OpenReview.net, 2021.  URL https://openreview.net/forum?id=6Tm1mposlrM.
  * Frankle et al. [2020] Frankle, J., Schwab, D. J., and Morcos, A. S.  The early phase of neural network training.  _arXiv preprint arXiv:2002.10365_ , 2020.
  * Goodfellow et al. [2016] Goodfellow, I., Bengio, Y., and Courville, A.  _Deep Learning_.  MIT Press, 2016.  http://www.deeplearningbook.org.
  * Guo et al. [2022] Guo, H., Jin, J., and Liu, B.  Stochastic weight averaging revisited.  _arXiv preprint arXiv:2201.00519_ , 2022.
  * Gur-Ari et al. [2018] Gur-Ari, G., Roberts, D. A., and Dyer, E.  Gradient descent happens in a tiny subspace.  _arXiv preprint arXiv:1812.04754_ , 2018.
  * Ha et al. [2016] Ha, D., Dai, A., and Le, Q. V.  Hypernetworks.  _arXiv preprint arXiv:1609.09106_ , 2016.
  * He et al. [2016] He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 770–778, 2016.
  * Hoffmann et al. [2022] Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al.  Training compute-optimal large language models.  _arXiv preprint arXiv:2203.15556_ , 2022.
  * Izmailov et al. [2018] Izmailov, P., Podoprikhin, D., Garipov, T., Vetrov, D. P., and Wilson, A. G.  Averaging weights leads to wider optima and better generalization.  In Globerson, A. and Silva, R. (eds.), _Proceedings of the Thirty-Fourth Conference on Uncertainty in Artificial Intelligence, UAI 2018, Monterey, California, USA, August 6-10, 2018_ , pp. 876–885. AUAI Press, 2018.  URL http://auai.org/uai2018/proceedings/papers/313.pdf.
  * Jastrzebski et al. [2020] Jastrzebski, S., Szymczak, M., Fort, S., Arpit, D., Tabor, J., Cho*, K., and Geras*, K.  The break-even point on optimization trajectories of deep neural networks.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=r1g87C4KwB.
  * Kaddour et al. [2022a] Kaddour, J., Liu, L., Silva, R., and Kusner, M. J.  A fair comparison of two popular flat minima optimizers: Stochastic weight averaging vs. sharpness-aware minimization.  _arXiv preprint arXiv:2202.00661_ , 2022a.  doi: 10.48550/ARXIV.2202.00661.  URL https://arxiv.org/abs/2202.00661.
  * Kaddour et al. [2022b] Kaddour, J., Lynch, A., Liu, Q., Kusner, M. J., and Silva, R.  Causal machine learning: A survey and open problems.  _arXiv preprint arXiv:2206.15475_ , 2022b.  URL https://arxiv.org/abs/2206.15475.
  * Kaplan et al. [2020] Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_ , 2020.
  * Katharopoulos & Fleuret [2018] Katharopoulos, A. and Fleuret, F.  Not all samples are created equal: Deep learning with importance sampling.  In _International conference on machine learning_ , pp. 2525–2534. PMLR, 2018.
  * Kingma & Ba [2014] Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Lakshminarayanan & Szepesvari [2018] Lakshminarayanan, C. and Szepesvari, C.  Linear stochastic approximation: How far does constant step-size and iterate averaging go?  In _International Conference on Artificial Intelligence and Statistics_ , pp. 1347–1355. PMLR, 2018.
  * Li et al. [2022] Li, T., Huang, Z., Tao, Q., Wu, Y., and Huang, X.  Trainable weight averaging for fast convergence and better generalization, 2022.  URL https://arxiv.org/abs/2205.13104.
  * Lin et al. [2017] Lin, T.-Y., Dollár, P., Girshick, R., He, K., Hariharan, B., and Belongie, S.  Feature pyramid networks for object detection.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 2117–2125, 2017.
  * Lipton & Steinhardt [2018] Lipton, Z. C. and Steinhardt, J.  Troubling trends in machine learning scholarship.  _arXiv preprint arXiv:1807.03341_ , 2018.
  * Liu et al. [2019] Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V.  Roberta: A robustly optimized bert pretraining approach.  _arXiv preprint arXiv:1907.11692_ , 2019.
  * Loshchilov & Hutter [2017] Loshchilov, I. and Hutter, F.  SGDR: Stochastic gradient descent with warm restarts.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=Skq89Scxx.
  * Merity et al. [2016] Merity, S., Xiong, C., Bradbury, J., and Socher, R.  Pointer sentinel mixture models, 2016.  URL https://arxiv.org/abs/1609.07843.
  * Merity et al. [2017] Merity, S., Keskar, N. S., and Socher, R.  Regularizing and optimizing lstm language models.  _arXiv preprint arXiv:1708.02182_ , 2017.
  * Mindermann et al. [2022] Mindermann, S., Brauner, J. M., Razzak, M. T., Sharma, M., Kirsch, A., Xu, W., Höltgen, B., Gomez, A. N., Morisot, A., Farquhar, S., et al.  Prioritized training on points that are learnable, worth learning, and not yet learnt.  In _International Conference on Machine Learning_ , pp. 15630–15649. PMLR, 2022.
  * Neu & Rosasco [2018] Neu, G. and Rosasco, L.  Iterate averaging as regularization for stochastic gradient descent.  In _Conference On Learning Theory_ , pp. 3222–3242. PMLR, 2018\.
  * Ott et al. [2019] Ott, M., Edunov, S., Baevski, A., Fan, A., Gross, S., Ng, N., Grangier, D., and Auli, M.  fairseq: A fast, extensible toolkit for sequence modeling.  _arXiv preprint arXiv:1904.01038_ , 2019.
  * Paszke et al. [2019] Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Köpf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S.  Pytorch: An imperative style, high-performance deep learning library.  In Wallach, H. M., Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pp. 8024–8035, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/bdbca288fee7f92f2bfa9f7012727740-Abstract.html.
  * Polyak & Juditsky [1992] Polyak, B. T. and Juditsky, A. B.  Acceleration of stochastic approximation by averaging.  _SIAM journal on control and optimization_ , 30(4):838–855, 1992.
  * Sagun et al. [2017] Sagun, L., Evci, U., Guney, V. U., Dauphin, Y., and Bottou, L.  Empirical analysis of the hessian of over-parametrized neural networks.  _arXiv preprint arXiv:1706.04454_ , 2017.
  * Szegedy et al. [2015] Szegedy, C., Liu, W., Jia, Y., Sermanet, P., Reed, S., Anguelov, D., Erhan, D., Vanhoucke, V., and Rabinovich, A.  Going deeper with convolutions.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 1–9, 2015.
  * Tang et al. [2020] Tang, K., Huang, J., and Zhang, H.  Long-tailed classification by keeping the good and removing the bad momentum causal effect.  _Advances in Neural Information Processing Systems_ , 33:1513–1524, 2020.
  * Team [2021] Team, T. M. M.  composer.  https://github.com/mosaicml/composer/, 2021.
  * Wortsman et al. [2022] Wortsman, M., Ilharco, G., Gadre, S. Y., Roelofs, R., Gontijo-Lopes, R., Morcos, A. S., Namkoong, H., Farhadi, A., Carmon, Y., Kornblith, S., et al.  Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time.  In _International Conference on Machine Learning_ , pp. 23965–23998. PMLR, 2022.
  * Wu et al. [2018] Wu, L., Ma, C., et al.  How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Zagoruyko & Komodakis [2016] Zagoruyko, S. and Komodakis, N.  Wide residual networks.  _arXiv preprint arXiv:1605.07146_ , 2016.
  * Zhang et al. [2019] Zhang, M. R., Lucas, J., Ba, J., and Hinton, G. E.  Lookahead optimizer: k steps forward, 1 step back.  In Wallach, H. M., Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pp. 9593–9604, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/90fd4f88f588ae64038134f1eeaa023f-Abstract.html.

##  Appendix A Losses and Top-5 Accuracies

For completeness, we also plot the training and validation losses for both
experiments and the top-5 accuracies for the ImageNet experiment.

Figure 3 shows similar speed up trends of LAWA over SGD as discussed in the
main body (Figure 1).

Figure 4 shows the training and validation losses for RoBERTa-Base trained on
WikiText103. Here, we also include losses during earlier stages of training
and point out that during these more fluctuant phases, LAWA performs worse. We
expect this behavior because previous works pointed out that the network
undergoes dramatic changes in early phases [34, 10, 7].

Figure 3: Additional ResNet50 / ImageNet Metrics: training/validation top-5
accuracies and losses. Figure 4: Starting point of LAWA being effective: on
RoBERTa-Base / WikiText103.

##  Appendix B LAWA vs. Lookahead

We compare LAWA ($k=10$) against Lookahead [41] on the moderately-sized
CIFAR-100 dataset (50k training images) and train a ResNet34 [12]. We follow
commonly used hyper-parameters (see e.g., [40, 23]), and train with SGD for
$200$ epochs, using a batch size of $256$, a momentum value of $0.9$ and a
cosine learning rate scheduler [26, 6] with initial learning rate $\eta=0.1$.
For LA, we use $\alpha=0.8$ and $k_{LA}=5$, as suggested by the authors for
this particular CIFAR100 dataset.

Figure 5: LAWA ($k=10$) outperforms Lookahead [41]. We plot the mean and
standard deviation of ResNet34/CIFAR-100 experiments across three random
seeds.

Figure 5 shows the training/test accuracy/loss as a function of the number of
epochs. LAWA reaches high test accuracy around $90$ epochs earlier than
SGD/LA.

Initially, we started experimenting with this learning task before scaling up
to larger datasets. Since we only observed slight but not dramatic
improvements in LA over the baseline, we did not evaluate LA in the larger-
scale ImageNet and BERT experiments. However, note that we apply LAWA to the
SGD checkpoints; an interesting future direction can be to combine LAWA with
LA, i.e., to average over checkpoints obtained with LA.

##  Appendix C Uniform vs. Exponentially Decayed Averaging Coefficients

We compare uniform (UNI, corresponding to $\bm{\theta}^{\text{LAWA\,}}$ by
default) and exponentially-decaying (EXP) weight coefficients. We follow the
same ResNet34 / CIFAR100 setup as in the previous section.

For EXP, we set $\alpha=0.9$ and compute

$\bm{\theta}^{\text{EXP}}_{0}=\bm{\theta}_{0},\quad\forall
E>0:\bm{\theta}^{\text{EXP}}_{E}:=\alpha\bm{\theta}_{E}+(1-\alpha)\bm{\theta}^{\text{EXP}}_{E-1}.$
(1)

We set $k=10$ for both strategies. Figure 6 shows that UNI slightly
outperforms EXP; however, the difference is not large.

Figure 6: Uniform coefficients (i.e., LAWA by default) slightly outperform
exponentially-decaying ones for $k=10$. We plot the mean and standard
deviation of ResNet34/CIFAR-100 experiments across three random seeds.