Philipp Ratz<EMAIL_ADDRESS>
Université du Québec à Montréal and François Hu<EMAIL_ADDRESS>
Université de Montréal and Arthur Charpentier<EMAIL_ADDRESS>
Université du Québec à Montréal

# Fairness Explainability using Optimal Transport
with Applications in Image Classification

###### Abstract

Ensuring trust and accountability in Artificial Intelligence systems demands
explainability of its outcomes. Despite significant progress in Explainable
AI, human biases still taint a substantial portion of its training data,
raising concerns about unfairness or discriminatory tendencies. Current
approaches in the field of Algorithmic Fairness focus on mitigating such
biases in the outcomes of a model, but few attempts have been made to try to
explain _why_ a model is biased. To bridge this gap between the two fields, we
propose a comprehensive approach that uses optimal transport theory to uncover
the causes of discrimination in Machine Learning applications, with a
particular emphasis on image classification. We leverage Wasserstein
barycenters to achieve fair predictions and introduce an extension to pinpoint
bias-associated regions. This allows us to derive a cohesive system which uses
the enforced fairness to measure each features influence _on_ the bias. Taking
advantage of this interplay of enforcing and explaining fairness, our method
hold significant implications for the development of trustworthy and unbiased
AI systems, fostering transparency, accountability, and fairness in critical
decision-making scenarios across diverse domains.

###### keywords:

Algorithmic Fairness, Explainable Artificial Intelligence, Image
Classification

##  1 Introduction

Machine Learning (ML) algorithms are widely used in critical domains ranging
from recruiting and law enforcement to personalized medicine [Berk(2012),
García-Peñalvo et al.(2018)García-Peñalvo, Cruz-Benito, Martín-González,
Vázquez-Ingelmo, Sánchez-Prieto, and Therón, Esteva et al.(2021)Esteva, Chou,
Yeung, Naik, Madani, Mottaghi, Liu, Topol, Dean, and Socher]. Their usage is
not beyond debate however, as fairness related issues remain poorly
understood. Further, ML algorithms can perpetuate societal stereotypes and
discriminatory practices by associating protected attributes, such as gender
and race, with predictions - even if the attribute in question is not directly
used in the modelling process [Pedreshi et al.(2008)Pedreshi, Ruggieri, and
Turini, Noiret et al.(2021)Noiret, Lumetzberger, and Kampel, Mehrabi et
al.(2021)Mehrabi, Morstatter, Saxena, Lerman, and Galstyan]. This can lead to
discriminatory behavior towards certain subgroups, where examples include
sexist recruiting algorithm, facial recognition systems that perform poorly
for females with darker skin and challenges in recognizing specific subgroups
in self-driving cars [Goodall(2014), Nyholm and Smids(2016), Dastin(2018)].
None of these algorithms were designed with explicit malice, but nevertheless
delivered biased results. As [Kearns and Roth(2019)] put it, “machine learning
won’t give you anything like gender neutrality ‘for free’ that you didn’t
explicitely ask for”. Whereas there has been notable progress in the
elimination of biases from black box models, challenges persist in identifying
the source of biases and explaining _why_ unfair outcomes materialized [Ali et
al.(2023)Ali, Abuhmed, El-Sappagh, Muhammad, Alonso-Moral, Confalonieri,
Guidotti, Del Ser, Díaz-Rodríguez, and Herrera]. Especially in fields where
complex black-box models are employed, explaining unfairness is often reliant
on testing hypotheses one-by-one, which can quickly become infeasible in the
era of big-data. A key reason for this is that existing methods aim to explain
how a given _score_ was constructed not how a _bias_ was introduced.
Additionally, criticism has been raised regarding the potential misuse of
standard explainable AI tools, which can result in misleading explanations
that validate incorrect models [Alvarez Melis and Jaakkola(2018),
Rudin(2019)]. This issue, often referred to as ”fairwashing” [Aïvodji et
al.(2019)Aïvodji, Arai, Fortineau, Gambs, Hara, and Tapp, Aïvodji et
al.(2021)Aïvodji, Arai, Gambs, and Hara], underscores the importance of
exercising caution in the application of such tools.

To address both the fairness and explainablity concerns, we take a two-fold
approach in this article. As a First Step (A), we use existing research on
mitigating the impact of sensitive attributes within a fairness-aware
framework and extend this explicitly to pre-trained models using optimal
transport theory [Villani(2021)]. We then turn our attention to explainablity
and issues arising from fairwashing in a second step (B). Here, aim to model
the algorithmic bias directly, providing valuable insights into the root
causes behind the biases. Importantly, our work fulfills both local (specific
model outputs) and global (model insights from data) explanation requirements,
as highlighted in [Arrieta et al.(2020)Arrieta, Díaz-Rodríguez, Del Ser,
Bennetot, Tabik, Barbado, García, Gil-López, Molina, Benjamins, et al.]. The
combination of these two steps, (A) and (B), has the advantage that
algorithmic fairness can be enforced and steps to mitigate the source of the
biases can be put into place. By filling this important gap, we contribute to
enhancing fairness, transparency and accountability in Artificial Intelligence
(AI) systems. Note that throughout this article, we focus on applications
involving images, rather than traditional tabular data, as it allows us to
showcase the effectiveness of our approach. Specifically, with images, a more
direct interpretation is possible, even without domain knowledge. However, the
techniques presented here can easily be applied to standard tabular data sets
as well.

###  1.1 Scoping and Definitions

The field of Algorithmic Fairness considers different metrics for distinct
goals. Here we opt to consider fairness at the distributional level,
specifically, we focus on the Demographic Parity (DP) notion of fairness
[Calders et al.(2009)Calders, Kamiran, and Pechenizkiy]. This strict
definition aims to achieve independence between sensitive attributes and
predictions without relying on labels. Formally, let $(\boldsymbol{X},S,Y)$ be
a random tuple with distribution $\mathbb{P}$, where
$\boldsymbol{X}\in\mathcal{X}\subset\mathbb{R}^{d}$ represents the features,
$S\in\mathcal{S}\subset\mathbb{N}$ a sensitive feature, considered discrete,
across which we would like to impose fairness and $Y\in\mathcal{Y}:=\\{0,1\\}$
represents the task to be estimated. As an illustration, consider Figure 1.
Our study primarily focuses on binary classification tasks; however, the
methodologies and techniques discussed can be readily extended and generalized
to regression tasks or multi-task problems. Also note that we include the
sensitive variable $S$ in the model, which is a somewhat paradoxical feature
in Algorithmic Fairness. However, both empirical studies [Lipton et
al.(2018)Lipton, McAuley, and Chouldechova] and theoretical research [Gaucher
et al.(2023)Gaucher, Schreuder, and Chzhen] have consistently shown that de-
biasing algorithms that do not consider the sensitive attribute, referred to
as fairness-unaware, exhibit inferior fairness compared to fairness-aware
approaches [Dwork et al.(2012)Dwork, Hardt, Pitassi, Reingold, and Zemel] that
leverage the sensitive feature.

$\\{\text{Young}\\}$$\\{\text{Female}\\}$Image Tensor$(\boldsymbol{X},\ S,\
Y)$ Figure 1:  Attributes of an Image, here we consider the image to be the
features (that is, $\boldsymbol{X}$) and the sensitive attribute (here,
gender) as well as the label (young) a categorical variable.

In binary classification, we aim to determine the probability response to get
$1$ in $\mathcal{Z}:=[0,1]$ for classifying $(\boldsymbol{X},S)$, also known
as a probabilistic classifier (or soft classifier). To achieve strong DP-
fairness for a given predictor $f$, the objective is to ensure that $S$ and
$f(\boldsymbol{X},S)$ are independent. This definition is more flexible than
its commonly used weak counterpart. Weak DP-fairness, aims to establish the
independence of $S$ and the hard classifier
$c_{f}(\boldsymbol{X},S):=\mathds{1}\\{f(\boldsymbol{X},S)\geq 0.5\\}$ but
restricts the user to a given threshold. Throughout our study, our goal is to
satisfy the strong DP-fairness, formally defined as the following conditions
for all $s,s^{\prime}\in\mathcal{S}$ and $u\in\mathbb{R}$:

$\mathbb{P}\left(f(\boldsymbol{X},S)\leq
u|S=s\right)=\mathbb{P}\left(f(\boldsymbol{X},S)\leq
u|S=s^{\prime}\right)\enspace.$ (1)

Further, we define $\mathcal{F}$ as the set of soft classifiers of the form
$f:\mathcal{X}\times\mathcal{S}\to\mathcal{Z}$. Also, given $s\in\mathcal{S}$
we denote,

  * •

$\nu_{f}$ (resp. $\nu_{f|s}$) the probability measure of $f(\boldsymbol{X},S)$
(resp. $f(\boldsymbol{X},S)|S=s$);

  * •

$F_{f|s}(u):=\mathbb{P}\left(f(\boldsymbol{X},S)\leq u|S=s\right)$ its
cumulative distribution function (CDF);

  * •

$Q_{f|s}(v):=\inf\\{u\in\mathbb{R}:F_{f|s}(u)\geq v\\}$ the associated
quantile function.

Throughout the remainder of the paper, we assume that for any
$f\in\mathcal{F}$, both the measures $\nu_{f}$ and $\nu_{f|s}$ have finite
second-order moments and their densities exist. With these notations, the DP
notion of fairness defined in Equation eq:IntroDP is rewritten as
$F_{f|s}(u)=F_{f|s^{\prime}}(u)$ for all $s,s^{\prime}\in\mathcal{S}$ and
$u\in\mathbb{R}$.

###  1.2 Related Work

####  1.2.1 Algorithmic Fairness

In recent years, research on algorithmic fairness has grown significantly. The
most common approaches can broadly be categorized into pre-processing, in-
processing, and post-processing methods. Pre-processing ensures fairness in
the input data by removing biases before applying ML models [Park et
al.(2021)Park, Hwang, Kim, and Byun, Qiang et al.(2022)Qiang, Li, Brocanelli,
and Zhu]. In-processing methods incorporate fairness constraints during model
training [Wang et al.(2020b)Wang, Qinami, Karakozis, Genova, Nair, Hata, and
Russakovsky, Joo and Kärkkäinen(2020)], where fairness constraints usually
modify the loss landscape and prevent the model from learning an unfair
solution. Whereas these two approaches focus on the model parameters itself,
post-processing techniques aim to achieve fairness through modifications of
the final scores [Karako and Manggala(2018), Kim et al.(2019)Kim, Ghorbani,
and Zou]. This has the advantage that it works with any kind of estimator,
including (partially) pre-trained models. Especially in computationally
intensive fields, transfer learning or partial fine-tuning are prevalent to
reduce training time and improve generalization. Due to this, post-processing
methods are easily integrable into a standard workflow at low computational
cost, hence we focus our work to this latter category. Of particular
importance is the literature using Optimal Transport, a mathematical framework
for measuring distributional differences. Intuitively, the goal is to
_transport_ unfair scores to fair ones while minimizing the effects of this
intervention to maintain predictive accuracy. In regression, methods like
[Chzhen et al.(2020)Chzhen, Denis, Hebiri, Oneto, and Pontil] and [Gouic et
al.(2020)Gouic, Loubes, and Rigollet] minimize Wasserstein distance to reduce
discrimination. Similarly, in classification, [Chiappa et al.(2020)Chiappa,
Jiang, Stepleton, Pacchiano, Jiang, and Aslanides] and [Gaucher et
al.(2023)Gaucher, Schreuder, and Chzhen] leverage optimal transport to achieve
fair scores. Recent work by [Hu et al.(2023a)Hu, Ratz, and Charpentier] also
achieves fairness in multi-task learning through joint optimization. Though
still a nascent field, applications thereof become more common [Zehlike et
al.(2020)Zehlike, Hacker, and Wiedemann, Charpentier et al.(2023)Charpentier,
Hu, and Ratz]. However, despite the extensive use of optimal transport theory
in algorithmic fairness, there is only limited research that delves into
applications that go beyond the standard case of tabular data, a first
shortcoming we address in this article.

####  1.2.2 Explainable AI

We focus on creating a simple fairness explanation method for computer vision,
narrowing down our exploration to two Explainable AI (XAI) subfields. Among
the most widely known methods are _model-agnostic_ techniques, such as LIME
[Ribeiro et al.(2016)Ribeiro, Singh, and Guestrin, Garreau and Mardaoui(2021)]
and SHAP [Shapley(1997), Lundberg and Lee(2017)], that do not depend on
specific assumptions about the model’s architecture. Though these approaches
can be extended to the analysis of images, a range of XAI methods have been
developed more specifically for the use of deep neural networks. These methods
commonly focus on local explainability, which often involves highlighting
important pixels (referred to as _attention maps_) for individual task
predictions. Global explainablity can then be achieved through the
identification of significant regions across the whole prediction analysis. As
these approaches leverage the specific architecture of neural networks they
are referred to as _model-specific_ approaches. Notable examples include Grad-
CAM [Selvaraju et al.(2017)Selvaraju, Cogswell, Das, Vedantam, Parikh, and
Batra] and its various variants, like Grad-CAM++ [Chattopadhay et
al.(2018)Chattopadhay, Sarkar, Howlader, and Balasubramanian] and Score-CAM
[Wang et al.(2020a)Wang, Wang, Du, Yang, Zhang, Ding, Mardziel, and Hu].
Recent work by [Franco et al.(2021)Franco, Oneto, Navarin, and Anguita] for
fair and explainable systems, generating two attention maps for local insight.
Global explainability is achieved through t-SNE representations, but explicit
discrimination explanations are often lacking, raising potential _fairwashing_
concerns [Alikhademi et al.(2021)Alikhademi, Richardson, Drobina, and
Gilbert]. Our approach sets itself apart from the aforementioned methods by
directly generating attention maps that specifically describe the model’s
unfair decisions, offering a clearer and more focused explanation for
discriminatory outcomes.

###  1.3 Contributions and Outline

In summary, we extend the literature of fairness-aware algorithms based on
optimal transport theory through the following points:

  * •

Fair decision-making: We adapt a post-processing model using optimal transport
theory for computer vision tasks, bringing the theory closer to the community.
We ensure fair and unbiased outcomes and show that the solution is optimal
with respect to the relative rankings, and independent of the bias.

  * •

Explainable artificial intelligence: Our main contribution is to use the
optimal transport plan to develop an XAI approach for identifying changes in
the data, describing unfair results. In computer vision applications, our
method directly highlights the regions most responsible for the stated bias,
facilitating direct identification of discrimination. The method is also
easily extendable to tabular data sets.

The remainder of this article is structured as follows. First, we provide a
brief background on optimal transport theory and establish its connection to
algorithmic fairness. Then, we turn our focus to our XAI methodologies to
uncover the causes of discrimination. Finally, we showcase their performance
through numerical experiments on the CelebA dataset.

##  2 Background on Optimal Transport

In this section, we present the fundamental concepts from optimal
transportation theory. Specifically, we focus on the Wasserstein distance and
give a brief overview of notable results in optimal transport theory with one-
dimensional measures, where all the main results can be found in [Villani et
al.(2009), Santambrogio(2015), Villani(2021)].

###  2.1 Wasserstein Distances

Let $\nu_{f_{1}}$ and $\nu_{f_{2}}$ be two probability measures on
$\mathcal{Z}$. The squared Wasserstein distance (cf. [Santambrogio(2015)],
definition §5.5.1) between $\nu_{f_{1}}$ and $\nu_{f_{2}}$ is defined as

$\mathcal{W}_{2}^{2}(\nu_{f_{1}},\nu_{f_{2}})=\inf_{\pi\in\Pi(\nu_{f_{1}},\nu_{f_{2}})}\mathbb{E}_{(Z_{1},Z_{2})\sim\pi}\left(Z_{2}-Z_{1}\right)^{2}\enspace,$

where $\Pi(\nu_{f_{1}},\nu_{f_{2}})$ is the set of distributions on
$\mathcal{Z}\times\mathcal{Z}$ having $\nu_{f_{1}}$ and $\nu_{f_{2}}$ as
marginals. If the infimum is achieved, the resulting coupling is referred to
as the optimal coupling between $\nu_{f_{1}}$ and $\nu_{f_{2}}$. If either of
the predictors belongs to $\mathcal{F}$, the optimal coupling can be
determined (refer to [Villani(2021)], Thm 2.12) as follows: if
$Z_{1}\sim\nu_{f_{1}}$ and $Z_{2}\sim\nu_{f_{2}}$, where
$f_{2}\in\mathcal{F}$, there exists a mapping $T:\mathbb{R}\to\mathbb{R}$ such
that

$\mathcal{W}_{2}^{2}(\nu_{f_{1}},\nu_{f_{2}})=\mathbb{E}\left(Z_{2}-T(Z_{2})\right)^{2}\enspace,$

with $T(Z_{2})\sim\nu_{f_{1}}$. We call $T$ the optimal transport map from
$\nu_{f_{2}}$ to $\nu_{f_{1}}$. Moreover, in the univariate setting, a closed-
form solution is explicitly provided as: $T(\cdot)=Q_{f_{1}}\circ
F_{f_{2}}(\cdot)$.

###  2.2 Wasserstein Barycenters

Throughout this article, we will frequently make use of Wasserstein
Barycenters. It can be defined for a given family of $K$ measures
$(\nu_{f_{1}},\dots,\nu_{f_{K}})$ and weights
$\boldsymbol{w}=(w_{1},\dots,w_{K})\in\mathbb{R}_{+}^{K}$ such that
$\sum_{s=1}^{K}w_{s}=1$. Then $Bar(w_{s},\nu_{f_{s}})_{s=1}^{K}$ represents
the Wasserstein barycenter of these measures which is the minimizer of

$Bar(w_{s},\nu_{f_{s}})_{s=1}^{K}=\underset{\nu}{\operatorname{arg}\\!\operatorname{min}}\;\sum_{s=1}^{K}w_{s}\cdot\mathcal{W}_{2}^{2}\left(\nu_{f_{s}},\nu\right)\enspace.$

The work in [Agueh and Carlier(2011)] shows that in our configuration, with
$f_{s}\in\mathcal{F}$, this barycenter exists and is unique. Put into words,
the Wasserstein barycenter can be used to find a representative distribution
that lies between multiple given distributions. For our applications, this
will ensure that the predictive distributions given any values $s$ will
coincide. The optimal transport problem then aims to minimize the total amount
of changes required to achieve this.

##  3 Fairness Projection using Optimal Transport

Optimal transport theory provides a means to ensure specific forms of
algorithmic fairness. We provide a short summary of the some necessary
concepts, all the main results can be found in [Chiappa et al.(2020)Chiappa,
Jiang, Stepleton, Pacchiano, Jiang, and Aslanides, Chzhen et al.(2020)Chzhen,
Denis, Hebiri, Oneto, and Pontil] and [Gouic et al.(2020)Gouic, Loubes, and
Rigollet].

###  3.1 Unfairness and Risk: a Warm-up

Paraphrasing the quote above that machine learning will not give a fair
classifier for free, we first need to define both the objective of the
classification and the concept of DP in a unified manner. DP will be used to
determine the fairness of a classifier.

######  Definition 3.1 (Demographic Parity).

Given a soft classifier $f$, its unfairness is quantified by

$\mathcal{U}(f)=\max_{s,s^{\prime}\in\mathcal{S}}\sup_{u\in\mathcal{Z}}\left|F_{f|s}(u)-F_{f|s^{\prime}}(u)\right|\enspace,$
(2)

and $f$ is called (DP-)fair if and only if $\mathcal{U}(f)=0$.

Consider $f^{*}(\boldsymbol{X},S):=\mathbb{E}\left[Y|\boldsymbol{X},S\right]$,
the Bayes rule that minimizes the following squared risk

$\mathcal{R}(f):=\mathbb{E}\left(Y-f(\boldsymbol{X},S)\right)^{2}\enspace.$

The associated hard classifier $c_{f^{*}}(\boldsymbol{X},S)$ has the property
of minimizing the risk of misclassification, which makes the squared risk
applicable for both regression and classification (for more details, see
[Gaucher et al.(2023)Gaucher, Schreuder, and Chzhen]). In line with this, we
adopt a popular approach to algorithmic fairness by incorporating DP-fairness
principles into risk minimization [Chzhen et al.(2020)Chzhen, Denis, Hebiri,
Oneto, and Pontil, Gouic et al.(2020)Gouic, Loubes, and Rigollet], that is:

$\min_{f\in\mathcal{F}}\left\\{\mathcal{R}(f):\mathcal{U}(f)=0\right\\}\enspace.$

By construction, this optimization effectively balances both risk
$\mathcal{R}$ and unfairness $\mathcal{U}$, leading to improved predictive
performance, reduced biases, and mitigation of potentially offensive or
discriminatory errors.

###  3.2 Optimal Fair Projection: Theoretical and Empirical Estimators

The two objectives, fairness and predictive accuracy are often in conflict
with one another. Most of the recent work in algorithmic fairness has
therefore focused on finding either a precise joint solution or an optimal
trade-off between the two. When starting from the best possible predictor
without any constraints, we refer to its optimal _fair_ counterpart as the
_optimal fair projection_. This estimator should minimize the unfairness
across the sensitive variables, while maintaining the best predictive accuracy
under this constraint. Much work has been done within the field to achieve
univariate optimal fair projections. Recall $p_{s}=\mathbb{P}(S=s)$ and let
$f_{B}\in\mathcal{F}$, where its measure is the Wasserstein barycenter
$\nu_{f_{B}}:=Bar(p_{s},\nu_{f^{*}|s})_{s\in\mathcal{S}}$. Then, studies
conducted by [Chzhen et al.(2020)Chzhen, Denis, Hebiri, Oneto, and Pontil] and
[Gouic et al.(2020)Gouic, Loubes, and Rigollet] demonstrate that

$f_{B}=\underset{f\in\mathcal{F}}{\operatorname{arg}\\!\operatorname{min}}\;\\{\mathcal{R}(f):\mathcal{U}(f)=0\\}\enspace.$

Therefore, $f_{B}$ represents the optimal fair predictor in terms of
minimizing unfairness-risk. Previous studies also offer a precise closed-form
solution: for all $(\boldsymbol{x},s)\in\mathcal{X}\times\mathcal{S}$,

$f_{B}(\boldsymbol{x},s)=\left(\sum_{s^{\prime}=1}^{K}p_{s^{\prime}}Q_{f^{*}|s^{\prime}}\right)\circ
F_{f^{*}|s}\left(f^{*}(\boldsymbol{x},s)\right)\enspace.$ (3)

To employ the results on real data the plug-in estimator of the Bayes rule
$f^{*}$ is given by $\hat{f}$, which corresponds to any DP-unconstrained ML
model trained on a training set
$\\{(\boldsymbol{x}_{i},s_{i},y_{i})\\}_{i=1}^{n}$ $n$ i.i.d. realizations of
$(\boldsymbol{X},S,Y)$. The empirical counterpart is then defined as:

$\widehat{f_{B}}(\boldsymbol{x},s)=\left(\sum_{s^{\prime}=1}^{K}\Hat{p}_{s^{\prime}}\Hat{Q}_{\Hat{f}|s^{\prime}}\right)\circ\Hat{F}_{\Hat{f}|s}\left(\Hat{f}(\boldsymbol{x},s)\right)\enspace,$
(4)

where $\Hat{p}_{s}$, $\Hat{F}_{\hat{f}|s}$ and $\Hat{Q}_{\hat{f}|s}$
corresponds to the empirical counterparts of $p_{s}$, $F_{f^{*}|s}$ and
$Q_{f^{*}|s}$. Note that, with the exception of $\hat{f}$, the remaining
quantities can be constructed using an unlabeled _calibration_ dataset,
denoted as
$\mathcal{D}^{\text{pool}}:=\\{(\boldsymbol{X}_{i},S_{i})\\}_{i=1}^{N}$, which
consists of N i.i.d. copies of $(\boldsymbol{X},S)$. The pseudo-code of this
approach is provided in Algorithm 1. We also visualize a possible model flow
in Figure 2, where the calibration layer corresponds to the inner workings of
Algorithm 1 and specifically Equation 4. [Chzhen et al.(2020)Chzhen, Denis,
Hebiri, Oneto, and Pontil] show that if the estimator $\Hat{f}$ is a good
proxy for $f^{*}$, then under mild assumptions on the distribution
$\mathbb{P}$, the calibrated post-processing approach $\widehat{f_{B}}$ is a
good estimator of $f_{B}$, enabling accurate and fair estimation of the
instances. [Gaucher et al.(2023)Gaucher, Schreuder, and Chzhen] demonstrates
that the hard classifier $c_{f_{B}}$ maximizes accuracy under DP-constraint,
and the classifier $c_{\widehat{f_{B}}}$ is proven to be a good estimator.

Algorithm 1 Fairness projection.

Input: instance $(\boldsymbol{x},{s})$, base estimator $\Hat{f}$, unlabeled
data $\mathcal{D}^{\text{pool}}=\\{(\boldsymbol{x}_{i},s_{i})\\}_{i=1}^{N}$.

Step 0. Split $\mathcal{D}^{\text{pool}}$ to construct the group-wise sample
$\\{\boldsymbol{x}_{i}^{s}\\}_{i=1}^{N_{s}}\sim\mathbb{P}_{\boldsymbol{X}|S=s}\quad\text{forany}s\in\mathcal{S}\enspace;$
with $N_{s}$ the number of images corresponding to $S=s$

Step 1. Compute the frequencies $(\widehat{p}_{s})_{s}$ from
$\\{s_{i}\\}_{i=1}^{N}$;

Step 2. Estimate $(\hat{F}_{\hat{f}|s})_{s}$ and $(\hat{Q}_{\hat{f}|s})_{s}$
from $\\{\boldsymbol{x}_{i}^{s}\\}_{i=1}^{N_{s}}$;

Step 3. Compute $\widehat{f}_{B}$ according to Eq. eq:plugin;

Output: fair predictor $\widehat{f}_{B}(\boldsymbol{x},s)$ at point
$(\boldsymbol{x},s)$.

##  4 Explainable AI using the Transportation Plan

Whereas the results from above enable the correction of given scores, they can
be considered a treatment of the symptoms rather than the cause of unfairness.
To this end, we extend the fairness procedures from above to explicitly
pinpoint these sources bias. The key idea that we pursue is that the transport
map used in Equation eq:plugin can be used to construct group-wise
counterfactual estimates. This approach clarifies differences between pre- and
post-processed scores, enabling the use of established XAI methods to uncover
the underlying causes of unfairness.

###  4.1 Local explainability

To isolate the features responsible for discrimination, we extend the fair
projection method by introducing an auxiliary learning task to detect the
source of biases directly. The binary task, denoted
$\Tilde{Y}\in\Tilde{\mathcal{Y}}:=\\{0,1\\}$, is estimated using the
distributions of the unfair predictor $f^{*}(\boldsymbol{X},S)$ and the DP-
fair predictor $f_{B}^{*}(\boldsymbol{X},S)$. This task can be chosen
according to specific goals, and we provide a sample of possibilities in Table
1. The formulation

$d_{B}(\boldsymbol{X},S):=f_{B}^{*}(\boldsymbol{X},S)-f^{*}(\boldsymbol{X},S)\enspace,$

then offers an intuitive explanation of unfairness. As an example, in the
context of a wage model, a positive value of $d_{B}(\boldsymbol{X},S)$
indicates a group discrimination for an individual $(\boldsymbol{X},S)$ while
a negative value might indicate favoritism. Its magnitude
$|d_{B}(\boldsymbol{X},S)|$ can be interpreted as the “degree” of
discrimination or favoritism, and its squared version an indicator for
extremes. The proposition below provides an interpretation of the quantity
$d_{B}(\boldsymbol{X},S)$ within the probabilistic framework, specifically in
the context of a binary sensitive attribute scenario.

######  Proposition 1 (Bias detection characterization).

Suppose $\mathcal{S}=\\{1,2\\}$ a binary sensitive feature scenario. Given
$(\boldsymbol{x},s)\in\mathcal{X}\times S$ and
$\Bar{s}\in\mathcal{S}-\\{s\\}$, there exists an optimal transport map from
$\nu_{f^{*}|s}$ to $\nu_{f^{*}|\Bar{s}}$, denoted
$T_{s\to\Bar{s}}:\mathcal{Z}\to\mathcal{Z}$, such that
$d_{B}(\boldsymbol{x},s)=f^{*}_{B}(\boldsymbol{x},s)-f^{*}(\boldsymbol{x},s)$
can be rewritten as,

$d_{B}(\boldsymbol{x},s)=p_{\Bar{s}}\cdot\left(\ T_{s\to\Bar{s}}\circ
f^{*}(\boldsymbol{x},s)-f^{*}(\boldsymbol{x},s)\ \right)\enspace,$ (5)

where $T_{s\to\Bar{s}}\circ f^{*}(\boldsymbol{X},s)\sim\nu_{f^{*}|\Bar{s}}$.

######  Proof 4.1.

Given $(\boldsymbol{x},s)\in\mathcal{X}\times S$, we are interested in the
quantity
$d_{B}(\boldsymbol{x},s)=f^{*}_{B}(\boldsymbol{x},s)-f^{*}(\boldsymbol{x},s)$.
For simplicity, we denote
$u_{s}(\boldsymbol{x})=F_{f^{*}|s}(f^{*}(\boldsymbol{x},s))$. Then, given
$\Bar{s}\in\mathcal{S}-\\{s\\}$, we have, f^*_B(x, s) - f^*(x, s) & = (
∑_s’=1, 2 p_s’Q_f^*—s’)∘F_f^*—s( f^*(x, s) ) - f^*(x, s)
= p_s⋅Q_f^*—s (u_s(x)) - p_s⋅Q_f^*—s (u_s(x)) + Q_f^*—s (u_s(x))
= p_s⋅Q_f^*—s (u_s(x)) - p_s⋅Q_f^*—s (u_s(x))
= p_s⋅( Q_f^*—s (u_s(x)) - f^*(x, s)) . Since $f^{*}\in\mathcal{F}$, by
definition of the Wasserstein distance, there exists a transport map
$T_{s\to\Bar{s}}:\mathcal{Z}\to\mathcal{Z}$ such that
$T_{s\to\Bar{s}}(\cdot)=Q_{f^{*}|\Bar{s}}\circ F_{f^{*}|s}(\cdot)$ with
$T_{s\to\Bar{s}}\circ f^{*}(\boldsymbol{X},s)\sim\nu_{f^{*}|\Bar{s}}$, which
concludes the proof.

In a binary sensitive framework, Proposition 1 asserts that
$|d_{B}(\boldsymbol{X},s)|$ depends on how much the DP-unconstrained
prediction for $(\boldsymbol{X},s)$ deviates from its projection onto
$\nu_{f^{*}|\Bar{s}}$. In other words, the r.h.s. of Equation
eq:BiasDetectChara measures the disparity between the initial prediction and
the projected prediction, where features $\boldsymbol{X}|S=s$ are aligned with
$\boldsymbol{X}|S=\Bar{s}$. As a concrete example, changing a male
individuals’ gender to female, the projection also modifies related attributes
(such as height or weight) to match a female counterpart, ensuring
comparability when these attributes naturally differ on a group level. Note
that $p_{\Bar{s}}$ enhances this bias for over-represented $\Bar{s}$ groups,
but reduces its significance for under-represented ones.

Task description | Probabilistic framework | Empirical framework
---|---|---
Discrimination | $\mathds{1}\\{f_{B}^{*}(\boldsymbol{X},S)-f^{*}(\boldsymbol{X},S)\geq 0\\}$ | $\mathds{1}\\{\widehat{f_{B}}(\boldsymbol{X},S)-\widehat{f}(\boldsymbol{X},S)\geq 0\\}$
Bias size | $\mathds{1}\\{|f_{B}^{*}(\boldsymbol{X},S)-f^{*}(\boldsymbol{X},S)|\geq\tau\\}$ | $\mathds{1}\\{|\widehat{f_{B}}(\boldsymbol{X},S)-\widehat{f}(\boldsymbol{X},S)|\geq\tau\\}$
Outliers | $\mathds{1}\\{(f_{B}^{*}(\boldsymbol{X},S)-f^{*}(\boldsymbol{X},S))^{2}\geq\tau\\}$ | $\mathds{1}\\{(\widehat{f_{B}}(\boldsymbol{X},S)-\widehat{f}(\boldsymbol{X},S))^{2}\geq\tau\\}$
Table 1: Bias / discrimination detection task

Alternatively, the new task described in Equation eq:BiasDetectChara can be
viewed as a decomposition of biases into implicit and explicit components.
Indeed, applying triangle inequality we have —d_B(x, s)— & ≤p_s⋅( —T_s→s
∘f^*(x, s) - f^*(x, s)— + —f^*(x, s) - f^*(x, s)— ) . In this context,
implicit bias, which refers to the hidden influence of $s$, is measured as the
difference between two values $T_{s\to\Bar{s}}\circ
f^{*}(\boldsymbol{X},s)-f^{*}(\boldsymbol{X},\Bar{s})$ both follow the same
distribution $\nu_{f^{*}|\Bar{s}}$ (although not independent). This implicit
bias represents the variation in predicted outcomes when the features
$\boldsymbol{X}$ under $S=s$ are aligned with those under $S=\Bar{s}$, in
contrast to unconstrained predictions where $s$ is simply replaced with
$\Bar{s}$ without aligning the features. Explicit bias, conversely, is
simplistically expressed as
$f^{*}(\boldsymbol{X},\Bar{s})-f^{*}(\boldsymbol{X},s)$. This measurement
aligns to the principle of _ceteris paribus_ , meaning “ _all other things
being equal_ ”. However, when considered in isolation, this condition can lead
to unrealistic situations. We provide a visual explanation in the Appendix A.

##### Data-driven procedure

In real datasets, we use plug-in estimators from Section 3.2 to estimate
$f^{*}$ and $f_{B}^{*}$, producing
$\widehat{d_{B}}=\widehat{f_{B}}-\widehat{f}$, the empirical counterpart of
$d_{B}$. Our goal is to train an estimator
$g:\mathcal{X}\to\Tilde{\mathcal{Y}}$, where $\Tilde{Y}\in\Tilde{\mathcal{Y}}$
represents the new target task outlined in Table 1’s last column. XAI methods
are then used to pinpoint areas causing observed model unfairness in the
initial ML model. For image classification, popular techniques like Grad-CAM
([Selvaraju et al.(2017)Selvaraju, Cogswell, Das, Vedantam, Parikh, and
Batra]) create attention maps highlighting these biased areas. The pseudo-code
for this approach is provided in Algorithm 2 with
$\Tilde{Y}^{\tau}:=\mathds{1}\\{|\widehat{f_{B}}(\boldsymbol{X},S)-\widehat{f}(\boldsymbol{X},S)|\geq\tau\\}$
as the desired XAI task.

######  Remark 2 (Impact of the parameter $\tau$ on the bias detection).

In Table 1, various tasks require establishing a threshold $\tau>0$ to
identify essential bias-contributing regions. We suggest determining $\tau$ at
a specific quantile $\alpha\in(0,1)$ within the sample
$\\{|\widehat{d_{B}}(\boldsymbol{x_{i}},s_{i})|\\}_{1\leq i\leq N}$, denoted
as $\widehat{Q}_{|\widehat{d_{B}}|}(\alpha)$. In particular, the choice of
$\alpha$ significantly influences the behavior of the bias detector. Indeed, a
larger $\alpha$ emphasizes causes with very high biases, while a smaller value
identifies all possible causes the bias detector can identify. Opting for
$\alpha=0.75$ might be a good choice since it results in a more balanced
dataset, with approximately equal occurrences of $\Tilde{Y}^{\tau}=0$ and
$\Tilde{Y}^{\tau}=1$, while also distinguishing significant areas for
unfairness.

Algorithm 2 Bias Detection

Input: new instance $(\boldsymbol{x},s)$, base estimator $\Hat{f}$, fair
estimator $\widehat{f_{B}}$, unlabeled sample
$\mathcal{D}^{\text{pool}}=\\{(\boldsymbol{x}_{i},s_{i})\\}_{i=1}^{N}$.

Step 0. Generate $(\Tilde{y}^{\tau}_{i})_{1\leq i\leq N}$, where
$\Tilde{y}^{\tau}$ can be constructed according to Table 1 or according to the
objective

Step 1. Train a learning model $g$ on
$\\{(\boldsymbol{x}_{i},\Tilde{y}^{\tau}_{i})\\}_{i=1}^{N}$;

Step 2. Use Grad-CAM (or other XAI method) to generate attention maps;

Output: Attention map of $g$ on instance $(\boldsymbol{x},s)$.

##  5 Experiments

We opt to showcase our method on image data, rather than tabular data, as it
helps to highlight an important practical aspect. Computer vision tasks are
often performed on (partially) pre-trained models and compute time presents a
major issue. The post-processing approach outlined in Equation eq:plugin is
particularly attractive in these circumstances. As an example, rendering the
scores fair throughout the applications in the experimental section took less
than 0.1 seconds on average. Further, the XAI approach outlined in the
previous section also works on pre-trained models as the transportation plan
only depends on the produced scores. We first present how the standard
approach outlined in Section 3 can be extended to standard computer vision
architectures and then show how the bias detection task can help identify the
regions associated with the bias.

###  5.1 Extension to Image Classification

An important detail from the above section is that in order to eliminate a
bias from the predictions, the sensitive feature $S$ must be included in the
modelling process to satisfy the assumptions of the optimal transport theory.
This is due to the fact that simply excluding the sensitive information might
lead the model to proxy for the sensitive variable which leads back to the
initial problem of the biased predictions. For image classification, we
consider a standard split into a feature (or embedding) and classification
block, where we use a pre-trained embedding model and keep its parameters
fixed throughout the whole procedure. The sensitive feature can then be added
through a simple layer concatenation of the output from the embedding before
it is fed into a classification block. As this will still result in biased
scores, a supplementary calibration layer is added, which implements Equation
eq:plugin and needs to be trained separately from the main model on a
calibration data set. This indicates the need to split the data into three
separate parts, the train set to fit the classification block to the specific
data, a calibration set (corresponding to $\mathcal{D}^{\text{pool}}$ in
Algorithm 1, that does _not_ need to be labeled), and a standard test set. The
architecture is visualized in Figure 2.

Classification Block (Pretrained) Embedding
${\textstyle\hat{f}(\boldsymbol{X},\
S)}$${\textstyle\hat{f}_{B}(\boldsymbol{X},\ S)}$ Calibration Layer Unfair
Scores Fair Scores $S$$\boldsymbol{X}$ Figure 2: Standard fairness approach,
the sensitive variable $S$ must be included the latest in the classification
block. The calibration layer contains an operation which makes use of the
Wasserstein Barycenter. The illustration above _Unfair Scores_ represent the
marginal score distributions induced by a binary sensitive variable which are
then transported to a common distribution for the _Fair Scores_.

###  5.2 Dataset and model

For our illustrations, we use the CelebA [Liu et al.(2015)Liu, Luo, Wang, and
Tang] data set, containing more than 200,000 celebrity images, each annotated
with 40 attributes. For the visualisations we selected a subset of 5,000
images with well-aligned facial features to obtain averaged predictions.
Following the setup from Figure 2, we choose the pre-trained version of
torchvision’s IMAGENET1K_V2 trained on the imagenet database as our embedding
model. With it, we use the built-in preprocessing steps on each image but
consider the parameters fixed. On top of the embedding, we then add a
classification block, similar in spirit to what [Wang et al.(2020b)Wang,
Qinami, Karakozis, Genova, Nair, Hata, and Russakovsky] proposed. This block
contains, before the output layer, three layers of size $[512,256,32]$ which
take as input the vector of size 2048 from the embedding block and the
sensitive feature of size 1. Each intermediate layer has a ReLU activation
function applied to it and uses a 0.1 dropout. We split the data into 64%
training, 16% calibration and 20% test data. The model is then trained task
wise for 10 epochs using a binary cross entropy loss and average results over
10 runs111All code can be found at:
https://github.com/FairInterpret/fair_imagesgithub.com/FairInterpret/fair_images.

###  5.3 Metrics and Prediction Tasks

As the approach is valid for soft classifiers, we use the Area under the ROC
curve (AUC) as the performance metric on the test set (denoted $A(f)$). We
also measure the unfairness $\widehat{\mathcal{U}}(f)$ on the test-set as the
empirical counterpart of Equation eq:Unf, based on the Kolmogorov–Smirnov
test,

$\widehat{\mathcal{U}}(f):=\max_{s,s^{\prime}\in\mathcal{S}}\sup_{t\in\mathcal{Z}}\left|\Hat{F}_{f|s}(t)-\Hat{F}_{f|s^{\prime}}(t)\right|\enspace,$

where $\Hat{F}_{f|s}$ is the empirical CDF of $f(\boldsymbol{X},S)|S=s$.

We consider three different binary prediction tasks from the data set (the
variables _Attractive_ , _Beard_ and _Young_) and consider _Gender_ as the
sensitive variable. As bias identification usually requires substantial domain
expertise, we demonstrate how our method works when we isolate the _Beard_
prediction task from influences of _Gender_. As the positive labels for the
task are almost exclusively present for male instances, we would expect the
bias to be the largest for this task. Further, the task also has a well
defined region of the image that should be used by the model to predict the
label, making it suitable for visualisations with Grad-CAM.

###  5.4 Results

| Uncalibrated | Fairness-aware
---|---|---
Metric | $A(f)$ | $\hat{\mathcal{U}}(f)$ | $A(f)$ | $\hat{\mathcal{U}}(f)$
Attractive |  0.855$\pm$0.002 |  0.447$\pm$ 0.028 |  0.769$\pm$ 0.002 | 0.011$\pm$ 0.001
Beard |  0.941$\pm$ 0.002 |  0.896 $\pm$ 0.009 |  0.731 $\pm$ 0.004 | 0.010$\pm$ 0.002
Young |  0.858$\pm$0.003 |  0.323$\pm$ 0.036 |  0.81$\pm$ 0.003 | 0.013$\pm$ 0.003
Table 2: AUC & Unfairness over 10 repetitions. Colored values highlight the
achieved fairness.

The numerical results are summarized in Table 2. The _Uncalibrated_ columns
present the results for a standard model that does not have a Calibration
Layer, the _Fairness-aware_ columns represent a model of the form of Figure 2.
As expected, all uncalibrated models present a significant level of
unfairness, indicating the model learned to use gender in its predictions. The
fairness-aware architecture manages to eliminate the bias, as indicated in the
highlighted column in the results table, though it also results in a lower
predictive accuracy as suggested by the theoretical analysis.

We then apply our XAI methodology to the data. We use the _Bias size_ task
from Table 1, re-fit the model to the new task $\Tilde{Y}^{\tau}$ and evaluate
the attention maps using the Grad-CAM algorithm of [Selvaraju et
al.(2017)Selvaraju, Cogswell, Das, Vedantam, Parikh, and Batra]. To see how
this can effectively prevent _fairwashing_ and help establish a causal
relation, we compare the results to attention maps obtained from the initial
model (that is, from the model that modelled the _Beard_ task). Results are
visualized in Figure 3. The left three columns are averages of 5,000 well-
aligned raw images, split by the _Gender_ variable and the Grad-CAM attention
map of the initial model. Most of the attention is indeed focused around the
region where one would expect the relevant characteristics of a beard to lie,
indicating a small bias (second and third column). However, this would stand
in stark contrast with the numerical results in Table 2. With the help of our
methodology we can instead isolate the regions that contribute to the bias
which we observe (fourth column). We can clearly see that the neural network
extensively focuses on features associated with gender (such as a receding
hairline, twice as likely for male individuals or blond hair, around 10 times
more likely for female individuals, or earrings around 20 times more likely
for females), but not the area where a beard would be expected. Such ratios
are easily identifiable and can be investigated further and our method can
effectively help identify relevant features.

Figure 3: GradCAM activations with gender as a sensitive feature when
estimating “Beard” task.

##  6 Conclusion

We investigated the use of optimal transport in both fairness and
explainability applications for machine learning methods. Though both fields
have produced significant advanced in recent years, the strong ties between
the two fields remain underexplored. We showed through our applications that a
standard optimal transport method for tabular data can readily be extended to
computer vision tasks, by adapting the model architecture slightly. This
resulted in an extremely efficient routine that can easily be integrated into
standard workflows. Given that the method is based on optimal transport, we
were then able to derive an XAI method based on the optimal transportation
plan, which helps to identify the sources of a bias, permitting a more
targeted investigation into the causes rather than the symptoms. This method
also enables researchers to adopt a more holistic approach to the choice of
sensitive variables, alleviating concerns of _fairwashing_ [Aïvodji et
al.(2019)Aïvodji, Arai, Fortineau, Gambs, Hara, and Tapp] and indeed opens up
a more objective discussion about emerging issues related to intersectional -
[Kong(2022)] or sequential [Hu et al.(2023b)Hu, Ratz, and Charpentier]
fairness.

This research was enabled in part by support provided by Calcul Québec
(calculquebec.ca) and the Digital Research Alliance of Canada (alliancecan.ca)

## References

  * [Agueh and Carlier(2011)] Martial Agueh and Guillaume Carlier.  Barycenters in the wasserstein space.  _SIAM Journal on Mathematical Analysis_ , 43(2):904–924, 2011.
  * [Aïvodji et al.(2019)Aïvodji, Arai, Fortineau, Gambs, Hara, and Tapp] Ulrich Aïvodji, Hiromi Arai, Olivier Fortineau, Sébastien Gambs, Satoshi Hara, and Alain Tapp.  Fairwashing: the risk of rationalization.  In _International Conference on Machine Learning_ , pages 161–170. PMLR, 2019.
  * [Aïvodji et al.(2021)Aïvodji, Arai, Gambs, and Hara] Ulrich Aïvodji, Hiromi Arai, Sébastien Gambs, and Satoshi Hara.  Characterizing the risk of fairwashing.  _Advances in Neural Information Processing Systems_ , 34:14822–14834, 2021.
  * [Ali et al.(2023)Ali, Abuhmed, El-Sappagh, Muhammad, Alonso-Moral, Confalonieri, Guidotti, Del Ser, Díaz-Rodríguez, and Herrera] Sajid Ali, Tamer Abuhmed, Shaker El-Sappagh, Khan Muhammad, Jose M Alonso-Moral, Roberto Confalonieri, Riccardo Guidotti, Javier Del Ser, Natalia Díaz-Rodríguez, and Francisco Herrera.  Explainable artificial intelligence (xai): What we know and what is left to attain trustworthy artificial intelligence.  _Information Fusion_ , page 101805, 2023.
  * [Alikhademi et al.(2021)Alikhademi, Richardson, Drobina, and Gilbert] Kiana Alikhademi, Brianna Richardson, Emma Drobina, and Juan E Gilbert.  Can explainable ai explain unfairness? a framework for evaluating explainable ai.  _arXiv preprint arXiv:2106.07483_ , 2021.
  * [Alvarez Melis and Jaakkola(2018)] David Alvarez Melis and Tommi Jaakkola.  Towards robust interpretability with self-explaining neural networks.  _Advances in neural information processing systems_ , 31, 2018.
  * [Arrieta et al.(2020)Arrieta, Díaz-Rodríguez, Del Ser, Bennetot, Tabik, Barbado, García, Gil-López, Molina, Benjamins, et al.] Alejandro Barredo Arrieta, Natalia Díaz-Rodríguez, Javier Del Ser, Adrien Bennetot, Siham Tabik, Alberto Barbado, Salvador García, Sergio Gil-López, Daniel Molina, Richard Benjamins, et al.  Explainable artificial intelligence (xai): Concepts, taxonomies, opportunities and challenges toward responsible ai.  _Information fusion_ , 58:82–115, 2020.
  * [Berk(2012)] Richard Berk.  _Criminal justice forecasts of risk: A machine learning approach_.  Springer Science & Business Media, 2012.
  * [Calders et al.(2009)Calders, Kamiran, and Pechenizkiy] T. Calders, F. Kamiran, and M. Pechenizkiy.  Building classifiers with independency constraints.  In _IEEE international conference on Data mining_ , 2009.
  * [Charpentier et al.(2023)Charpentier, Hu, and Ratz] Arthur Charpentier, François Hu, and Philipp Ratz.  Mitigating discrimination in insurance with wasserstein barycenters.  _arXiv preprint arXiv:2306.12912_ , 2023.
  * [Chattopadhay et al.(2018)Chattopadhay, Sarkar, Howlader, and Balasubramanian] Aditya Chattopadhay, Anirban Sarkar, Prantik Howlader, and Vineeth N Balasubramanian.  Grad-cam++: Generalized gradient-based visual explanations for deep convolutional networks.  In _2018 IEEE winter conference on applications of computer vision (WACV)_ , pages 839–847. IEEE, 2018.
  * [Chiappa et al.(2020)Chiappa, Jiang, Stepleton, Pacchiano, Jiang, and Aslanides] S. Chiappa, R. Jiang, T. Stepleton, A. Pacchiano, H. Jiang, and J. Aslanides.  A general approach to fairness with optimal transport.  In _AAAI_ , 2020.
  * [Chzhen et al.(2020)Chzhen, Denis, Hebiri, Oneto, and Pontil] E. Chzhen, C. Denis, M. Hebiri, L. Oneto, and M. Pontil.  Fair regression with wasserstein barycenters.  In _Advances in Neural Information Processing Systems_ , 2020.
  * [Dastin(2018)] Jeffrey Dastin.  Amazon scraps secret ai recruiting tool that showed bias against women, Oct 2018.  URL https://www.reuters.com/article/us-amazon-com-jobs-automation-insight-idUSKCN1MK08G.
  * [Dwork et al.(2012)Dwork, Hardt, Pitassi, Reingold, and Zemel] 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.
  * [Esteva et al.(2021)Esteva, Chou, Yeung, Naik, Madani, Mottaghi, Liu, Topol, Dean, and Socher] Andre Esteva, Katherine Chou, Serena Yeung, Nikhil Naik, Ali Madani, Ali Mottaghi, Yun Liu, Eric Topol, Jeff Dean, and Richard Socher.  Deep learning-enabled medical computer vision.  _NPJ digital medicine_ , 4(1):5, 2021.
  * [Franco et al.(2021)Franco, Oneto, Navarin, and Anguita] Danilo Franco, Luca Oneto, Nicolò Navarin, and Davide Anguita.  Toward learning trustworthily from data combining privacy, fairness, and explainability: an application to face recognition.  _Entropy_ , 23(8):1047, 2021.
  * [García-Peñalvo et al.(2018)García-Peñalvo, Cruz-Benito, Martín-González, Vázquez-Ingelmo, Sánchez-Prieto, and Therón] Francisco García-Peñalvo, Juan Cruz-Benito, Martín Martín-González, Andrea Vázquez-Ingelmo, José Carlos Sánchez-Prieto, and Roberto Therón.  Proposing a machine learning approach to analyze and predict employment and its factors.  2018\.
  * [Garreau and Mardaoui(2021)] Damien Garreau and Dina Mardaoui.  What does lime really see in images?  In _International conference on machine learning_ , pages 3620–3629. PMLR, 2021.
  * [Gaucher et al.(2023)Gaucher, Schreuder, and Chzhen] Solenne Gaucher, Nicolas Schreuder, and Evgenii Chzhen.  Fair learning with wasserstein barycenters for non-decomposable performance measures.  In _International Conference on Artificial Intelligence and Statistics_ , pages 2436–2459. PMLR, 2023.
  * [Goodall(2014)] Noah J Goodall.  Ethical decision making during automated vehicle crashes.  _Transportation Research Record_ , 2424(1):58–65, 2014.
  * [Gouic et al.(2020)Gouic, Loubes, and Rigollet] Thibaut Le Gouic, Jean-Michel Loubes, and Philippe Rigollet.  Projection to fairness in statistical learning.  _arXiv preprint arXiv:2005.11720_ , 2020.
  * [Hu et al.(2023a)Hu, Ratz, and Charpentier] François Hu, Philipp Ratz, and Arthur Charpentier.  Fairness in multi-task learning via wasserstein barycenters.  In Danai Koutra, Claudia Plant, Manuel Gomez Rodriguez, Elena Baralis, and Francesco Bonchi, editors, _Machine Learning and Knowledge Discovery in Databases: Research Track_ , pages 295–312, Cham, 2023a. Springer Nature Switzerland.  ISBN 978-3-031-43415-0.
  * [Hu et al.(2023b)Hu, Ratz, and Charpentier] François Hu, Philipp Ratz, and Arthur Charpentier.  A sequentially fair mechanism for multiple sensitive attributes.  _arXiv preprint arXiv:2309.06627_ , 2023b.
  * [Joo and Kärkkäinen(2020)] Jungseock Joo and Kimmo Kärkkäinen.  Gender slopes: Counterfactual fairness for computer vision models by attribute manipulation.  In _Proceedings of the 2nd international workshop on fairness, accountability, transparency and ethics in multimedia_ , pages 1–5, 2020.
  * [Karako and Manggala(2018)] Chen Karako and Putra Manggala.  Using image fairness representations in diversity-based re-ranking for recommendations.  In _Adjunct Publication of the 26th Conference on User Modeling, Adaptation and Personalization_ , pages 23–28, 2018.
  * [Kearns and Roth(2019)] M. Kearns and A. Roth.  _The ethical algorithm: The science of socially aware algorithm design_.  Oxford University Press, 2019.
  * [Kim et al.(2019)Kim, Ghorbani, and Zou] Michael P Kim, Amirata Ghorbani, and James Zou.  Multiaccuracy: Black-box post-processing for fairness in classification.  In _Proceedings of the 2019 AAAI/ACM Conference on AI, Ethics, and Society_ , pages 247–254, 2019.
  * [Kong(2022)] Youjin Kong.  Are “intersectionally fair” ai algorithms really fair to women of color? a philosophical analysis.  In _Proceedings of the 2022 ACM Conference on Fairness, Accountability, and Transparency_ , pages 485–494, 2022.
  * [Lipton et al.(2018)Lipton, McAuley, and Chouldechova] Zachary Lipton, Julian McAuley, and Alexandra Chouldechova.  Does mitigating ml’s impact disparity require treatment disparity?  _Advances in neural information processing systems_ , 31, 2018.
  * [Liu et al.(2015)Liu, Luo, Wang, and Tang] 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.
  * [Lundberg and Lee(2017)] Scott M Lundberg and Su-In Lee.  A unified approach to interpreting model predictions.  _Advances in neural information processing systems_ , 30, 2017.
  * [Mehrabi et al.(2021)Mehrabi, Morstatter, Saxena, Lerman, and Galstyan] Ninareh Mehrabi, Fred Morstatter, Nripsuta Saxena, Kristina Lerman, and Aram Galstyan.  A survey on bias and fairness in machine learning.  _ACM Computing Surveys (CSUR)_ , 54(6):1–35, 2021.
  * [Noiret et al.(2021)Noiret, Lumetzberger, and Kampel] Sophie Noiret, Jennifer Lumetzberger, and Martin Kampel.  Bias and fairness in computer vision applications of the criminal justice system.  In _2021 IEEE Symposium Series on Computational Intelligence (SSCI)_ , pages 1–8. IEEE, 2021.
  * [Nyholm and Smids(2016)] Sven Nyholm and Jilles Smids.  The ethics of accident-algorithms for self-driving cars: An applied trolley problem?  _Ethical theory and moral practice_ , 19(5):1275–1289, 2016.
  * [Park et al.(2021)Park, Hwang, Kim, and Byun] Sungho Park, Sunhee Hwang, Dohyung Kim, and Hyeran Byun.  Learning disentangled representation for fair facial attribute classification via fairness-aware information alignment.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pages 2403–2411, 2021.
  * [Pedreshi et al.(2008)Pedreshi, Ruggieri, and Turini] Dino Pedreshi, Salvatore Ruggieri, and Franco Turini.  Discrimination-aware data mining.  In _Proceedings of the 14th ACM SIGKDD international conference on Knowledge discovery and data mining_ , pages 560–568, 2008.
  * [Qiang et al.(2022)Qiang, Li, Brocanelli, and Zhu] Yao Qiang, Chengyin Li, Marco Brocanelli, and Dongxiao Zhu.  Counterfactual interpolation augmentation (cia): A unified approach to enhance fairness and explainability of dnn.  In _Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence, IJCAI_ , pages 732–739, 2022.
  * [Ribeiro et al.(2016)Ribeiro, Singh, and Guestrin] Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin.  ” why should i trust you?” explaining the predictions of any classifier.  In _Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining_ , pages 1135–1144, 2016.
  * [Rudin(2019)] Cynthia Rudin.  Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead.  _Nature machine intelligence_ , 1(5):206–215, 2019.
  * [Santambrogio(2015)] Filippo Santambrogio.  Optimal transport for applied mathematicians.  _Birkäuser, NY_ , 55(58-63):94, 2015.
  * [Selvaraju et al.(2017)Selvaraju, Cogswell, Das, Vedantam, Parikh, and Batra] Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra.  Grad-cam: Visual explanations from deep networks via gradient-based localization.  In _Proceedings of the IEEE international conference on computer vision_ , pages 618–626, 2017.
  * [Shapley(1997)] L. S Shapley.  A value for n-person games.  _Classics in game theory_ , 69, 1997.
  * [Villani(2021)] Cédric Villani.  _Topics in optimal transportation_ , volume 58.  American Mathematical Soc., 2021.
  * [Villani et al.(2009)] Cédric Villani et al.  _Optimal transport: old and new_ , volume 338.  Springer, 2009.
  * [Wang et al.(2020a)Wang, Wang, Du, Yang, Zhang, Ding, Mardziel, and Hu] Haofan Wang, Zifan Wang, Mengnan Du, Fan Yang, Zijian Zhang, Sirui Ding, Piotr Mardziel, and Xia Hu.  Score-cam: Score-weighted visual explanations for convolutional neural networks.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops_ , pages 24–25, 2020a.
  * [Wang et al.(2020b)Wang, Qinami, Karakozis, Genova, Nair, Hata, and Russakovsky] Zeyu Wang, Klint Qinami, Ioannis Christos Karakozis, Kyle Genova, Prem Nair, Kenji Hata, and Olga Russakovsky.  Towards fairness in visual recognition: Effective strategies for bias mitigation.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 8919–8928, 2020b.
  * [Zehlike et al.(2020)Zehlike, Hacker, and Wiedemann] Meike Zehlike, Philipp Hacker, and Emil Wiedemann.  Matching code and law: achieving algorithmic fairness with optimal transport.  _Data Mining and Knowledge Discovery_ , 34(1):163–200, 2020.

##  Appendix A Visual representation

To help with the interpretation of the different quantities in Section 4, we
created a graph below. As a simple working example, consider a wage model and
we would like to analyse the distributions induced by a binary sensitive
feature. In Figure 4, we visualise these two marginal distributions, using
their CDF. The CDF has the advantage that it scales both marginal
distributions to the same range.

Consider predictions run on the x-axis (from $f^{*}(x)$, where the sensitive
feature is implicitly included in the features). As explained in the text, we
would like to compare two individuals across their relative position within
their subgroup (recall that, intuitively, the barycenter we apply to achieve
fairness recreates a common distribution - but keeps the relative within group
ordering constant to minimize the transportation cost). This can be done via
the projection, that is, for a value on the x-axis, say the value obtained for
$f(x_{0},s=1)$, we calculate the value on the CDF for the group $s=1$ (the
blue curve) and project the point to the CDF of the group $s=2$ (the orange
dotted line). This then gives us the counterfactual based on optimal transport
on the $x-axis$. Compare this to the ceteris paribus prediction (that is
$f(x_{0},s=2)$), which does not take into account the group level changes in
$x_{0}$. This allows a very natural interpretation of the distance we are
using for our auxiliary task. In a sense, we would like to explain the
features that are relevant to explain large differences, be it directly
related (that is the distance $|f(x_{0},s=1)-f(x_{0},s=2)|$) or indirectly
related (the remainder of the total distance) to the group variable $s$.

$F(x)$$f^{*}(x)$$Q_{f^{*}|s=2}(F_{f^{*}|s=1}(f(x_{0},s=1)))$$f(x_{0},s=1)$$F_{f^{*}|s=1}(f(x_{0},s=1))$$f(x_{0},s=2)$
Figure 4:  Illustration of different points and counterfactuals.


11institutetext: CERN, 1 Esplanade des Particules, Geneva, Switzerland

# Accelerating GAN training using highly parallel hardware on public cloud

Renato Cardoso 11<EMAIL_ADDRESS>Dejan Golubovic 11 Ignacio Peluaga
Lozada 11 Ricardo Rocha 11 JoÃ£o Fernandes 11 Sofia Vallecorsa 11

###### Abstract

With the increasing number of Machine and Deep Learning applications in High
Energy Physics, easy access to dedicated infrastructure represents a
requirement for fast and efficient R&D. This work explores different types of
cloud services to train a Generative Adversarial Network (GAN) in a parallel
environment, using Tensorflow data parallel strategy. More specifically, we
parallelize the training process on multiple GPUs and Google Tensor Processing
Units (TPU) and we compare two algorithms: the TensorFlow built-in logic and a
custom loop, optimised to have higher control of the elements assigned to each
GPU worker or TPU core. The quality of the generated data is compared to Monte
Carlo simulation. Linear speed-up of the training process is obtained, while
retaining most of the performance in terms of physics results. Additionally,
we benchmark the aforementioned approaches, at scale, over multiple GPU nodes,
deploying the training process on different public cloud providers, seeking
for overall efficiency and cost-effectiveness. The combination of data
science, cloud deployment options and associated economics allows to burst out
heterogeneously, exploring the full potential of cloud-based services.

##  1 Introduction

In recent years, several studies have demonstrated the benefits of using Deep
Learning (DL) to solve different tasks related to data processing in High
Energy Physics (HEP): for example, generative models are being tested as fast
alternatives to Monte Carlo based simulation and anomaly detection algorithms
are being explored to design searches for rare new-physics processes.

Training of such models has been made tractable with the improvement of
optimization methods and the advent of dedicated infrastructure. High
Performance Computing (HPC) and storage technologies are often required by
these types of projects, together with the provision of high processing
capacity on multi-architectures, ranging from large multi-core systems to
hardware accelerators like GPUs or TPUs. These requirements directly translate
into the necessity to extend the HEP computing model that, so far, has relied
on on-premise resources and grid computing. A hybrid model has the potential
to seamlessly complement and extend local infrastructure, using widely
supported techniques, allowing flexible and efficient resource bursting to
public cloud for CERN research teams and science beyond. In this context, our
work explores the deployment of HEP DL models on native cloud-based
frameworks, with well-established industry support and large communities,
whilst, at the same time, integrate it with the ongoing work to transparently
and generically burst CERN on-premises infrastructure to public clouds. For
this reason, in addition to the optimisation of the specific physics use case,
we provide relevant feedback on the most efficient models to develop an end-
to-end integrated strategy that includes cost prediction and optimization when
consuming public cloud resources at scale, for scientific research. The
development and optimisation of deep generative models as fast simulation
solutions is one of the most important research directions as far as DL for
HEP is concerned. HEP experiments rely heavily on simulations in order to
model complex processes and describe detectors response: the classical Monte
Carlo approach can reproduce theoretical expectations with a high level of
precision, but it is extremely demanding in terms of computing resources
hsf2019 . For this reason, different strategies trading some physics accuracy
for speed are being developed as fast simulation techniques and deep
generative models are today among the most promising solutions GAN1 ; GAN2 ;
WGAN_Thorben ; Salamani2018 .
This work focuses on 3DGAN GAN_Gulrukh : a three-dimensional convolutional
Generative Adversarial Network (GAN) that is being optimised for the
simulation of electromagnetic calorimeters. This work summarises the specific
development required by the 3DGAN architecture and the adversarial training
process for efficient data parallel training on Google TPUs and multiple GPUs
setups. It also details the effect of distributed and mixed-precision training
(on TPUs) on 3DGAN physics performance: this aspect is essential to ensure
that the performance of complex simulations tasks (such as the GAN generation
process) is preserved at an acceptable level. In addition, we describe the
3DGAN training deployment on public cloud providers (Google Cloud Platform
(GCP) and Microsoft Azure) comparing different approaches, based on
Kubernetes, Kubeflow and the Azure Machine Learning service. Our deployments
are among the first examples of establishing an innovative hybrid platform
successfully deploying scientific deep learning workloads on public clouds at
large scale kube .

##  2 Generative Adversarial Networks in High Energy Physics

Generative models represent a fundamental part of deep learning and GANs
2014Googfellow , in particular, can generate sharp and realistic images with
high resolution wgan ; karras2017progressive . In HEP, most GAN applications
address the problem of calorimeters simulation. The LaGAN de2017learning  and
CaloGAN paganini2017calogan  models first introduced GAN for simulating
simplified detectors; since then, multiple prototypes have been developed
Salamani2018 ; WGAN_Thorben ; icip18_3dgan . While distributed deep neural
networks training is a well-established field, with multiple approaches
thoroughly investigated and state-of-the-art solutions widely spread, the
specific GAN use case is still a relatively new research subject:
mustafa2019cosmogan  represents a notable example at large scale, while
additional work can be found in IXPUG_3DGAN ; blind3_3dgan . Low precision
computation (int8, float16, bfloat16) is another active research field, aimed
at reducing computing resources, mostly targeted to speed-up inference: models
are trained in float32 precision, on multi-node clusters, quantized to lower
precision and then loaded on light-weight devices such as edge-devices. Lower
precision formats can be used at training time as well, in combination with
full precision float32 in order to maintain the level of accuracy (mixed
precision training) MixedPrecision2 ; IntelMixedPrecision .

Details about the 3DGAN model can be found in GAN_Gulrukh : it represents the
first proof of concept on using 3D convolutional GANs to simulate high
granularity calorimeters. The 3DGAN model uses an ACGAN-like acgan  approach
to introduce physics constraints and improve image fidelity: in particular,
two quantities, representing the particle energy, $E_{P}$, and its trajectory,
measured by the angle $\theta$, are used as input to the generation step. In
general terms, simulating samples using generative models is much faster than
using a Monte Carlo approach. In the case of 3DGAN, a several orders of
magnitude speed-up is observed with respect to the HEP Monte Carlo state-of-
the-art geant4 : a 67000x speed-up has been recently achieved using int8
quantized inference on Intel Xeon Cascade Lake icipram . The training process
is however very time consuming: an entire week is needed in order to train the
model to convergence using a single NVIDIA V100 GPU and therefore a
distributed training approach is essential CHEP2018_distributed .

##  3 The adversarial training process acceleration

The Tensorflow 2.3 distribute.Strategy API runs the training process across a
single node with multiple GPUs, multiple nodes, as well as multiple TPU cores,
while ensuring portability with minimal code changes; the downside is that
only graph mode is available on TPUs, whereas time penalties are introduced in
eager mode on GPUs. This work relies on the mirrored strategy, that implements
synchronous parallel training over multiple GPUs on a single node; the multi-
worker mirrored strategy, which extends it across multiple nodes and the TPU
strategy, that implements synchronous training on TPUs, using its own,
optimised, All_reduce and collective operations tensorflow2015-whitepaper .

The 3DGAN adversarial training, shown in Algorithm 1, trains, alternately, the
discriminator and generator networks an equal number of times. It includes
several steps to initialise the generation process, train the discriminator
network on real and fake data and, finally, the generator itself.

Algorithm 1 The 3DGAN training process

for $N_{epochs}$ do

for $N_{batches}$ do

get 3D image batch

get labels: $E_{p}$, $E_{CAL}$, $\theta$

generator input: sample latent noise batch and concatenate(noise, $E_{p}$,
$\theta$)

generate fake 3D image batch

calculate fake $E_{CAL}$ batch

train discriminator on real batch

train discriminator on fake batch

for 2 do

get labels: $E_{p}$, $E_{CAL}$, $\theta$

generator input: sample latent noise batch and concatenate(noise, $E_{p}$,
$\theta$)

train the generator

end for

end for

end for

Algorithm 1 can be implemented using standard training APIs, such as
keras.train_on_batch or custom training loops. In the first case, the
distribution logic is built-in and only the batch size is manually updated
according to the number of parallel replicas. Taking a look at the pseudo-code
in Algorithm 1, it is evident that the generator input initialization steps
are not distributed by keras.train_on_batch: they are sequentially run and end
up being a bottleneck when the number of replicas increases. This effect is
shown in Figure 1: the discriminator training time (per batch) stays the same
while increasing the number of GPUs, with a small increase due to reduce and
broadcast operations. On the other hand, the generator initialisation time
increases linearly with the number of GPUs, since it is sequentially run.

(a)

(b)

Figure 1: (left) Discriminator training times divided between processing, real
data training and fake data training. (right) Generator training times. Times
are quoted for one data batch.

In order to reduce this bottleneck, we rewrite Algorithm 1 as a custom
training loop, using Tensorflow tf.function to explicitly redefine the forward
step in training mode. In addition, the tf.function includes all previously
sequential steps. At this point, the remaining bottleneck is represented by
the data pre-processing and distribution steps. In order to address this
issue, we rely on some important advantages related to using a customised
training loop: loading data directly to the accelerators (GPUs and TPUs) and
running data preparation (batching and shuffling) on the CPU host while the
GPUs/TPUs are training. More specifically, we convert the original HDF5 data
to the native tensor format, TF Records, and introduce iterators instead of
manually instantiating batches on the CPU host. This final optimisation step
yields results we discuss in the next section for TPUs and in section 5 for a
larger number of GPUs on Google Cloud Platform (GCP) and Microsoft Azure
clouds.

##  4 Training on Google Tensor Processing Units

Application-specific integrated circuits (ASICs), such as Google Tensor
Processing Units (TPUs) dedicated to ML/DL workloads, are being developed with
the objective of accelerating matrix multiplication. In order to have
efficient communication across TPUs, Google implements a 2D Torus topology by
using Inter-Core Interconnect (ICI) links that enable a direct connection
between chips and simplify the rack-level deployment. In addition, the v2 and
v3 TPUs use two cores per chip, further decreasing the latency with respect to
single-core chips. In particular, each TPU core contains a 128x128 matrix
multiplier unit (MXU) which uses systolic arrays.

Google TPUs use bfloat16 for MXU calculations, in order to accelerate data
processing while, at the same time, maintaining a high level of accuracy. To
fully exploit bfloat16, the Tensorflow XLA compiler maps TPU dependent
instructions while also targeting CPU and GPUs instructions. The XLA compiler
analyses the code and optimises the data transfer among different TPU
elements, maps concurrent instructions in different parts of the TPU hardware
and increases the ICI communication capabilities between TPU cores. All these
optimisation steps occur during the first tf.function iteration, which in
consequence, is slower. For this reason we quote average results excluding the
first batch.

###  4.1 Results

Our initial test is run on 8 cores TPUs, using a batch size $BS=128$ in order
to match the design MXU dimension. Results for v2 and v3 TPUs are compared in
figure 2: using the double core v3 TPU we obtain 2x speed-up on the epoch
training time. Figure 2 also shows, on the center panel, the batch size impact
on training time on 8 cores v3 TPU: a batch size $BS=64$ is under utilising
the MXU capacity and, therefore, it is processed in the same amount of time as
a 128 batch; on the other hand, $BS=256$, which is double the optimal size, is
processed using twice the time.

(a)

(b)

(c)

Figure 2: (left) Time per epoch (in seconds) on 8 cores v2 and v3 TPU,
$BS=128$. (center) Time per epoch (in seconds) on 8 cores v3 TPU for batch
sizes $BS=64,128,256$. (right) Measured training time per epoch (in seconds)
compared to the theoretical linear speed-up.

Figure 2, right panel, shows weak scaling results, obtained using a $BS=128$
batch size on v3 TPUs with a number of cores ranging from 8 to 128: the time
per epoch scales linearly. Running on a 128 cores TPU, reduces the training
time per epoch to about 30 seconds, allowing 3DGAN to train to convergence in
a little more than one hour. As mentioned above, preserving the fidelity of
the synthetic images is critical, when accelerating training. In order to
verify this constraint, we validate the quality of the generated images
against Monte Carlo simulation: the left panel on figure 3 shows the
calorimeter energy response along the longitudinal plane. It can be noticed,
that while the GAN model seems to slightly distort the response toward the
left side of the distribution, the agreement remains overall very good.

(a)  (b)  (c)

Figure 3: Calorimeter energy response. The GAN prediction is in blue, Monte
Carlo in red. (left) Along the detector longitudinal plane. (center) Along the
detector transverse plane in linear scale and (right) logarithmic scale.

The energy response along the transverse plane, shown in the center panel of
figure 3, is also optimally reproduced, both in the central calorimeter cells
(the distribution bulk highlighted in the linear scale plot) and at the volume
edges (shown in the logarithmic scale plot).

To further analyse the results obtained, a profiling tool was employed using
tensorboard. The program is mostly compute bound, in fact the TPUs are only
idle 0.7% of the time, with 38% of the time being used for the forward-pass
(represented as the functional models), 61% for back-propagation (the gradient
tapes) and 0.2% for AllReduce.

On the other hand, analysis on the GPUs shows an increase on the idle time to
around 2.9%, for a higher number of GPUs, this is due to an increase in the
time it takes to do the AllReduce operation, which occurs after every gradient
tape operation shown in the profile.

##  5 Automated deployment using public cloud services

Training deep generative models for simulation is a workload that fits well
the public cloud for multiple reasons. First of all, training is done
periodically and can greatly benefit from the use of on-demand resources when
a longer-term investment on-premises cannot be justified. Secondly the
availability of a large number of GPUs, often a scarce resource on data
centers on-premises, as well as access to the most recent architectures.
Finally, the access to specialized accelerators such as Google TPUs and
Graphcore IPUs, that are vendor specific, allowing for a generic unlocked
deployment strategy, whilst at the same time, benefiting from a side-by-side
comparison.

The deployment can take different approaches, including the low level
management of virtual machines which keeps the burden of a significant
fraction of the infrastructure management on the end user; on the other hand,
one can also use open platforms such as Kubernetes, which offer APIs
abstracting the infrastructure and delegating most of the operations to the
cloud service itself, being available in all major public clouds; another
possibility is to use MLaaS provided by the vendor, that normally consists in
a PaaS optimised and dedicated to machine learning, such as the Azure Machine
Learning service. Below, we describe and present results for the last two
approaches.

###  5.1 Google Cloud Platform with Kubeflow

Our first setup relies on Kubeflow kubeflow . Our existing on-premises
deployment and configuration was reused and directed to the Google Kubernetes
Engine (GKE), Google Cloud’s managed Kubernetes service offering, benefiting
from the resources under CERN IT’s Cloud Broker Pilot project (CloudBank
EU)CloudBank . To monitor resource usage during training we rely on Prometheus
and custom Grafana dashboards showing GPU power usage, memory usage,
temperature and overall utilization.

The Kubernetes cluster includes multiple node groups, each of them with a
profile where the number of GPUs per node varies from 1 to 8. This allows the
evaluation of the penalty introduced from spreading the training across a
larger number of nodes, as well as the penalty of increasing the density in a
single node. Instances are of the n1-standard family, from n1-standard-4 for
single GPU nodes to n1-standard-32 for 8 GPU nodes. All tests rely on
preemptible instances, which are low SLA resources living for a maximum of 24
hours that can be deleted at any point if capacity is reclaimed - but are over
3 times cheaper compared to reserved instances for V100 GPUs. No reclaim was
experienced during the full set of tests performed.

####  5.1.1 Experiments and Results

In the initial set of runs, the goal is to determine an optimal batch size for
training the neural network. Higher batch sizes mean less steps to complete
the entire data set, which results in faster completion of the training.
Though smaller batch sizes can yield to better generalization performance, in
this case we tolerate a small loss of generalization performance (as discussed
in section 5.2). From the hardware perspective, higher batch sizes offer
better utilization of GPUs by exploiting advantages of the hardware. For the
purpose of the paper, as models are being trained on enterprise machine
learning platforms, it is important to provide maximum utilization of
resources in order to minimize the cost for training jobs. This experiment is
conducted by fixing the number of nodes (4), GPUs per node (8) and the number
of epochs (5). The only hyper-parameter varying is the batch size. Batch sizes
are selected to be multiples of 2 due to the GPU Single Instruction Multiple
Data (SIMD) organization. Minimum batch size is 16, since it is sufficiently
small to provide good generalization performance. The maximum batch size (96)
is selected to fit the GPU memory. As expected, the training times are reduced
when using larger batches. The increase in performance is prominent between
batch size of 16 and 32. With batch sizes of 64 and 80 the difference is
almost negligible.

As described in the infrastructure chapter, the cluster being used can be
configured with a different number of nodes and GPUs per node. In this
experiment, we select the optimal layout with respect to the TensorFlow
distributed training setup, including the number of GPUs per node, TensorFlow
workers and number of GPUs in each worker. The objective is to understand the
impact of communication overhead in the TensorFlow distributed setup. The
experiment has two sets of runs using a total of 32 GPUs. In the first set the
number of nodes (4) and the number of GPUs per node (8) are fixed. Different
combinations of workers are selected: 32 workers with 1 GPU each, 16 workers
with 2 GPUs each, 8 workers with 4 GPUs each and 4 workers with 8 GPUs each.
In the second set the number of nodes and GPUs per node varies: 32 nodes with
1 GPU each, 16 nodes with 2 GPUs each, 8 nodes with 4 GPUs each and 4 nodes
with 8 GPUs each. The workers configuration follows the nodes configuration
and is the same as in the previous run: 32 workers with 1 GPUs each, 16
workers with 2 GPUs each, 8 workers with 4 GPUs each and 4 workers with 8 GPUs
each. Essentially, the same set of worker configurations is run on a different
hardware layout. The results in figure 4, on the right, show that the impact
of different hardware configurations is not huge across the runs. The last
three runs from both sets show times that differ by less than 3 seconds. The
first run in both sets (32 workers with 1 GPU each) shows a significant
increase in time, meaning the optimal configuration should not strive for less
GPUs per node and less GPUs per worker. The best results are achieved when the
number of workers matches the number of nodes and the number of GPUs per
worker matches the number of GPUs per node. In subsequent experiments this
setup will be used.

(a)

(b)

(c)

Figure 4: (left) Impact of batch size on time to perform one epoch of
training. (center) Impact of various combinations of hardware layout and TF
worker configuration on time to perform one epoch of training. Orange
represents the first set of runs, where hardware configuration is fixed with 4
nodes with 8 GPUs each, while number of workers and GPUs per worker change.
Blue represents the second set of runs where number of workers and GPUs per
worker correspond to number of nodes and GPUs per node. (right) Batch
processing times for the second epoch of training using hardware layout with 4
nodes with 8 GPUs per node. Orange represents the batch processing times when
number of workers is 32 and each worker has 1 GPU, Blue represents the batch
processing times when number of workers is 4 and each worker has 8 GPUs.

It was also observed that a higher number of workers introduces instability in
the processing time per one batch, leading to instability in the processing
time per one epoch and a less predictable time for the entire run completion.
With a run with 32 workers, the variation in batch processing time is
significant compared to the run with 4 workers. This goes in favor of avoiding
higher number of workers and focusing on using nodes with as many GPUs as
possible.

We next determine how the performance changes with the increased number of
GPUs available to workers. This set of runs strives for maximum utilization of
the maximum number of GPUs available. Building on the previous tests a batch
size of 96 and one worker per node are selected, varying only the number of
nodes and GPUs per node.

(a)

(b)

Figure 5: (left) The impact of increasing number of total GPUs assigned to
workers on time to process one epoch of training. (right) Cost of running one
epoch on the Google Cloud Platform. The cost per epoch remains similar for all
GPU combinations, though the training time improves almost linearly.

The performance improvement is close to linear with a slight drop when moving
to 128 GPUs, which are positive results regarding TensorFlow’s distributed
training capabilities to scale workloads horizontally (see figure 5). This is
particularly true when combined with Kubeflow and the ease of configuration
and automation of the training process in large infrastructures.

####  5.1.2 Cost Analysis Overview

An important point when using public cloud resources is to understand upfront
the costs involved in order to reach optimal service settings and
configurations and, consequently, maximize cost effectiveness to reach a
sustainable position. The infrastructure costs during the tests on Google
Cloud Platform are driven by the cost of the GPUs, with virtual machines
accounting for less than 5% of the total. Below, we present the numbers for
using 2 to 128 GPUs and for a subset of TPU configurations. The calculation is
done taking the training time of one epoch against the hourly GPU costs in
region europe-west4  and include both reserved and preemptible options. It
should be noted that when using reserved instances it is possible to apply for
sustained or committed-use discounts which reduce the difference to
preemptible instances costs. A few points are worth noting (figure 5, right
panel): the cost per epoch remains similar when increasing the number of GPUs,
while the training time is reduced up to 52 times for 128 GPUs (almost linear
up to 64 GPUs); using preemptible GPUs can allow for significant cost savings
and should be used whenever possible; finally, the best results are achieved
using preemptible TPU v3-8, which are 2.4 times cheaper than their GPU
equivalent, considering the training time. This is an interesting result that
positions TPUs as a cost-effective option combined with their potential to
reduce the overall training time. Preemptible TPUs are only available up to 8
cores, with the reserved instance costs being less attractive. As an example,
the overall training cost with 128 GPUs is the same as TPU v3-32, but it takes
half the time. Future work will include a comparison with a larger number of
TPU cores.

###  5.2 Microsoft Azure with Azure’s Machine Learning Service

In the case of Microsoft Azure cloud, the 3DGAN deployment makes use of the
Azure Machine Learning service azure  offered through CERN openlab’s Azure
cloud research enrollment. The service allows the user to be abstracted from
the underlying infrastructure detail, since hardware is a managed component of
the service stack. This means, essentially, that the job needs to be
configured while the provisioning of the compute cluster is entirely
operationalised by the Azure Machine Learning service, as part of the created
deployment workflow. The interaction with the Azure ML service can be done
using different ways such as Web interface, CLI and SDK. In this case, the SDK
approach is preferred because of its wider support. It provides a space,
called Workspace, where the user can manage the different run submissions. A
detailed description of the general architecture can be found in azure . For
our tests, we use GPU-powered nodes with 24 vCPU cores, 448 GiB memory and 4
V100 GPU. The 3DGAN training time scales very close to linear up to a total of
16 nodes (64 GPUs). Azure ML automatically optimises the data set management,
in aspects such as caching, pre-fetching and parallel data loading according
to the hardware infrastructure setup. In figure 6, right panel, we verify that
manually modifying caching, pre-fetching and auto-tuning features in
Tensorflow does not improve on the train time performance.

(a)

(b)

Figure 6: (left) Microsoft Azure training time per epoch for a $BS=64$ batch
size. (right) Azure ML service automatic optimisation performance compared to
manual cache and prefetch implementation. All times are measured in seconds.

In terms of image fidelity, the work in IXPUG_3DGAN  has shown that data
parallel training can affect the quality of the images especially at the
border of the detector sensitive region. Indeed, in figure 7, on the left, the
calorimeter energy response in the detector transverse plane shows a similar
performance degradation at the edges of the energy distribution. The amount of
energy deposited in this region is orders of magnitudes smaller than in the
central pixels; it is therefore much harder for the GAN to correctly predict
it. It should be noted, however, that this behaviour does not critically
affect physics applications, since such small energy deposits are usually
neglected. On the other hand, the slight shift of the longitudinal energy
distribution, on the right in figure 7, is significant and additional
optimisation is needed in order to reduce this effect.

(a)

(b)

Figure 7:  Calorimeter energy response as predicted by 64 GPUs GAN (blue) and
Monte Carlo (red). (left) Distribution along the detector transverse plane in
logarithmic scale.(right) Distribution along the calorimeter depth in linear
scale.

The experiments on The CERN openlab Azure enrollment were performed using an
Azure credits grant, currently evolving to provide capabilities for realistic
analysis and overall resource costing optimisation. These capabilities were
not yet ready at the time this paper was produced and for this reason we defer
any cost analysis to future studies.

##  6 Conclusions and Future Plans

With Deep Learning models in the HEP domain becoming more complex,
computational requirements increase, triggering the need to consider HPCaaS
and MLaaS offers in the public cloud. This work presents results of the first
deployment of a three-dimensional convolutional GAN for detector simulation on
TPUs, demonstrating an efficient parallelization of the adversarial training
process: the 3DGAN training time is brought down from about a week to around
one hour. This result enables large architecture hyper-parameter scans to run
in just a few days (instead of weeks) and, therefore, it greatly extends the
range of detector geometries than can be simulated by the 3DGAN model.
Validation against Monte Carlo shows a slight over-estimation of the energy
distribution along the edges of the detector sensitive volume for a large
number of nodes (above 64 GPUs): simple strategies, typically employed to ease
the effect of increasing global batch sizes are not sufficient to improve
convergence at scale. We are currently investigating several different
directions, such as the role played by the batch normalisation layer, which is
one of the components greatly affecting the 3DGAN model performance
GAN_Gulrukh . In addition to the physics result, this work demonstrates the
deployment of scientific DL workloads using public cloud services,
complementing the on-premises infrastructure offers in use by research
organisations. This is achieved by exploring innovative, open hybrid
provisioning and orchestration models, profiting from the technological
potential of each cloud provider, unlocked from any particular commercial
vendor or provisioning model. Although public cloud services are offered in
the general market as commodity services, they must be validated for research
use cases technically but also in terms of cost optimisation requirements.
These elements combined, as part of a global strategy when scaling out on-
demand heterogeneously, can profit from the full potential of cloud-based
services that are evolving from basic elastic provisioning of virtual
resources, to a transparent and adaptive smart continuum. These possibilities
must be assessed with close alignment across performance, usability and
economics, considering complexities such as type of resources, data movement
patterns, optimisation of architectures adapted to the scientific workloads,
provisioning lock-in risks in and also data processing aspects to achieve a
long-term sustainable position. Future work will continue to explore models of
resources combination from specific vendors and explore seamless orchestration
of resources federated across multiple cloud providers, combining technical
and cost-effectiveness requirements. In addition, data governance aspects will
also be considered, ensure the presence of a set of technical measures to
guarantee adequacy with current legislation (e.g. GDPR and Free Flow of data).

## References

  * (1) J. Albrecht et al., Computing and Software for Big Science 3, 7 (2019)
  * (2) M. Erdmann et al., arXiv preprint arXiv:1807.01954 (2018)
  * (3) V. Chekalina et al., arXiv preprint arXiv:1812.01319 (2018)
  * (4) Erdmann et al., Comput Softw Big Sci 3 (2019)
  * (5) A. Ghosh (ATLAS Collaboration), Tech. rep., CERN, Geneva (2019)
  * (6) G. Khattak et al., in _18th IEEE ICMLA Conference_ (2019)
  * (7) R. Rocha, L. Heinrich, Kubecon - CloudNativeCon Europe 2019 (2019)
  * (8) I.J. Goodfellow et al., arXiv preprint arXiv:1406.2661 (2014)
  * (9) M. Arjovsky et al. (2017), 1701.07875
  * (10) T. Karras et al., arXiv preprint arXiv:1710.10196 (2017)
  * (11) L. de Oliveira et al., arXiv preprint arXiv:1701.05927 (2017)
  * (12) M. Paganini et al., arXiv preprint arXiv:1705.02355 (2017)
  * (13) G.R. Khattak et al., in _25th IEEE ICIP conference_ (2018)
  * (14) M. Mustafa et al., Computational Astrophysics and Cosmology 6, 1 (2019)
  * (15) S. Vallecorsa et al., in _High Performance Computing_ (2018), Vol. 11203
  * (16) A. Author1, _Title3: removed because of ISC review policy_ (Address2, 2019)
  * (17) S.R. Nandakumar et al., Frontiers in Neuroscience 14 (2020)
  * (18) J. Osorio et al. (2020)
  * (19) A. Odena et al., arXiv preprints arXiv:1610.09585 (2016)
  * (20) _Geant_ (2020), accessed Dec 20, 2020, http://geant.cern.ch/
  * (21) F. Rehm et al., in _Proceedings of the 10th ICPRAM conference_ (2021)
  * (22) J.R. Vlimant et al., in _EPJ Web of Conferences_ (2019), Vol. 214
  * (23) M. Abadi et al., accessed Dec 20, 2020, https://www.tensorflow.org/
  * (24) _Kubeflow_ (2020), [Online; accessed Dec 21 2020], https://www.kubeflow.org/
  * (25) _Cloud Broker Project (CloudBank EU)_ (2020), accessed Dec 20, 2020, https://indico.cern.ch/event/919839/contributions/4147766/attachments/2162052/3648772/Cloud_Broker_Pilot.pdf
  * (26) _Microsoft Azure Machine Learning service_ (2020), accessed Dec 20, 2020, https://azure.microsoft.com/en-us/services/machine-learning/