$f$-FERM: A Scalable Framework for Robust Fair Empirical Risk Minimization

Machine learning models are increasingly deployed in critical applications ranging from healthcare (Ahmad et al., 2018) to image processing (Krizhevsky et al., 2017), education to job recruitment (Boselli et al., 2018), and social networking to cybersecurity (Xin et al., 2018). Machine learning practitioners have adopted learning algorithms to fathom inherently difficult and crucial problems. However, naïve deployment of these models may lead to serious shortcomings such as biased predictions against minority groups (Angwin et al., 2016; Buolamwini & Gebru, 2018), vulnerability to adversarial attacks (Madry et al., 2017; Carlini & Wagner, 2017; Baharlouei et al., 2023), or lack of generalizability (Arjovsky et al., 2019). Consequently, it is of utmost importance to have reliable and trustworthy models that are, in particular, fair and comply with equality norms and provisions worldwide (Act, 1964; Elford, 2023).

With the increasing concern for the trustworthiness of unchecked machine learning algorithms, a broad class of paradigms has been proposed to counteract and mitigate both the cause and effects of model unreliability. Imposing statistical independence between model output and particular input features is of interest in various domains, especially when the generalization of a trained model is based on a collection of spurious features present in the training dataset (Dwork et al., 2012; Hardt et al., 2016; Yan et al., 2017). These could be sensitive features like gender, race, age, and/or income in the context of fairness or confounding factors like environmental artifacts in the context of image classification (Arjovsky et al., 2019). Existing literature on imposing statistical independence between selected input features and model outputs is directed into three approaches: pre-processing, post-processing, and in-processing methods.

Pre-processing methods entail upstream changes made in datasets to mask sensitive features or reduce the dependency of output variables on sensitive features through transforming data in a stage before the training phase (Kamiran & Calders, 2012; Zemel et al., 2013; Ustun et al., 2019). Post-processing methods involve model-specific adjustments to the model’s output to ensure the independence of predictions and sensitive attributes (Hardt et al., 2016; Alghamdi et al., 2022). While pre-processing and post-processing methods do not affect the training procedure, they fail to exploit underlying training mechanisms for the best achievable accuracy-fairness tradeoffs. Unsurprisingly enough, optimizing accuracy and fairness jointly (in-processing) leads to better tradeoffs than sequentially optimizing fairness and accuracy in a pre-processing or post-processing fashion.

In-processing methods alternatively add fairness constraints or regularizers, penalizing dependence between sensitive attributes and output variables. (Zafar et al., 2017) utilizes covariance as the measure of independence between the sensitive attributes and the predictions. While such a measure is amenable to stochastic updates, it fails to capture correlations beyond linear. Alternatively, several non-linear measures such as Rényi correlation (Baharlouei et al., 2020), $\chi^{2}$ divergence (Lowy et al., 2022), $L_{\infty}$ distance (Donini et al., 2018), and Maximum Mean Discrepancy (MMD) (Prost et al., 2019) are proposed in the literature to establish the independence of the predictors and sensitive attributes. In-processing techniques can be model-specific (Wan et al., 2021; Aghaei et al., 2019) or generalizable to different training algorithms (Baharlouei et al., 2020; Lowy et al., 2022).

In the spirit of in-processing methods, input data-driven constraints or regularization terms are used to modify training objectives of problems like learning generalizable models to new environments, invariant learning, and learning in the presence of distribution shifts (Arjovsky et al., 2019; Mary et al., 2019; Baharlouei et al., 2020). Such constrained/regularized reformulations are prevalent in learning robust classifiers against adversarial attacks (Sinha et al., 2018), meta-learning (Balaji et al., 2018), federated learning (Deng et al., 2023), and alternative learning paradigms such as learning distributionally robust optimization (DRO) models (Kuhn et al., 2019; Levy et al., 2020), tilted empirical risk minimization (TERM) (Li et al., 2020), and Squared-root Lasso (Belloni et al., 2011).

While in-processing techniques outperform pre-processing and post-processing approaches, they are not scalable to large datasets because of a lack of adaptability to stochastic optimization (Mary et al., 2019; Lowy et al., 2022). All aforementioned examples consist of regularization terms in their objective functions where the gradient cannot be described as a linear combination of data point functions. As a result, applying stochastic gradient descent or other stochastic first-order methods on the objective functions of such problems might not converge, especially for small batch sizes.

Motivated by this, (Lowy et al., 2022) proposes a stochastic optimization framework for Exponential Rényi Mutual Information as the measure of independence. More recently Zhong & Tandon (2023) use $f$-divergences as regularization terms to establish the independence between sensitive attributes and predictions. They estimate the $f$-divergence regularizers offline through multi-layer neural networks to avoid the computational challenges of devising scalable stochastic methods for nonconvex min-max problems. Our approach, on the other hand, directly solves the variational formulation for both full-batch and stochastic settings with convergence guarantees to non-spurious solutions. In Section 2, using the variational representation of $f$-divergences, we present a convergent stochastic optimization framework for fair learning via $f$-divergences. (Lowy et al., 2022) is a special case of $f$-divergences where $f(t)=t^{2}-1$ ($\chi^{2}$ divergence). Aside from $\chi^{2}$, all other divergences listed in Table 1 are not introduced in the literature to the best of our knowledge.

Designing convergent stochastic algorithms for fair empirical risk minimization can be further explored in scenarios involving changes in the data distribution from the source to the target domain. Detection and mitigation of biases against protected groups in the presence of distribution shifts have been extensively studied in recent years. Lechner et al. (2021) theoretically shows that learning fair representations (pre-processing) is nearly impossible for the popular notions of fairness, such as demographic parity in the presence of the distribution shift. Ding et al. (2021), on the other hand, experimentally demonstrates that applying post-processing fairness techniques (Hardt et al., 2016) to learn fair predictors of income concerning race, gender, and age fails to transfer from one US state (training domain) to another state. Overlooking distribution shifts can lead to catastrophic decisions threatening the well-being of human subjects when deploying a trained model in certain hospitals to other hospitals (Schrouff et al., 2022). The current literature for handling distribution shifts with in-processing methods relies on certain assumptions on the type of distribution shift (demographic shift (Fang et al., 2020; Du & Wu, 2021; Maity et al., 2021; Giguere et al., 2021), label shift (Dai & Brown, 2020), and/or covariate shift (Rezaei et al., 2021; Singh et al., 2021)) or explicit access to the causal graph (Mishler & Dalmasso, 2022; Schrouff et al., 2022) of predictors, sensitive attributes, and target variables. As a result, they face practical limitations and cannot cope with most real-world problems involving complex shifts that cannot be categorized in the ones assumed in their works.

Alternatively, Taskesen et al. (2020) provides convex objective functions for imposing fairness on logistic regression using constraint optimization. Staib & Jegelka (2019) use MMD for defining uncertainty sets around training distribution, whereas Husain (2020) use Integral Probability Measure (IPM) to mitigate the distribution shift. The main limitation of these approaches is their reliance on the convexity of the underlying learning model and lack of scalability due to incompatibility with stochastic optimization algorithms. Wang et al. (2023) uses the Maximum Mean Discrepancy (MMD) distance between the spectral norm of the Hessian matrix at advantaged and disadvantaged data points. However, they do not provide convergence guarantees for their proposed algorithm to any notion of optimality. In addition, the method is not necessarily amenable to stochastic updates. While we naturally define the uncertainty set directly on the joint distribution of sensitive attributes and predictions, they use the curvature of the obtained solution quantified by the norm of the Hessian matrix as a heuristic for promoting the robustness of the fair solution.

Contributions: This paper establishes a scalable (stochastic) fair empirical risk minimization framework through regularization via $f$-divergences ($f$-FERM) for both standard and distributed shift settings. $f$-FERM presents a unified methodology based on the Legendre-Fenchel transformation, enabling us to develop theoretically convergent first-order stochastic algorithms when only small batches of data are available at each iteration. Further, we have presented the first distributionally robust optimization framework under $\ell_{p}$ norms uncertainty sets covering nonconvex losses such as neural networks. The presented framework for fair inference in the presence of distribution shifts does not rely on the causal graph describing the causal interaction of input features, sensitive attributes, and target variables, which is rarely available in practical problems.

Paper Organization: We structure our response towards designing scalable, robust, and fair algorithms into two sections. Section 2 motivates the design of unbiased gradient estimators of objectives with information-theoretic $f$-divergence regularizers. In Section 3, we present our approach for fair inference in the presence of the distribution shift in detail. Our experiments provide an extensive examination of various $f$-divergences and their suitability as regularizers and also show the consistency of our method across all batch sizes in contrast to existing benchmarks. Similar experiments are carried out for robust training on varying amounts of distributional shifts in data.

A widely studied problem in algorithmic fairness is promoting a notion of group fairness, such as demographic parity, equalized odds, equality of opportunity, or sufficiency through an in-processing method. For these notions, we aim to establish a [conditional] statistical independence between the predictions (e.g., the creditworthiness of the individual) and the sensitive attributes (e.g., gender, race). For simplicity of presentation, we formulate all problems under the demographic parity notion, which requires statistical independence between the prediction and the sensitive attribute. Without loss of generality, all formulations and methods are generalizable to other aforementioned notions of group fairness by considering conditional random variables (see Appendix A). A popular in-processing approach for training fair (classification) models under the demographic parity notion is to regularize the empirical risk minimization:

where ${\bm{\theta}}$ is the learning parameters (e.g., weights of the neural network); ${\mathbf{x}}_{i}\in\mathbb{R}^{d}$ is the $i$-th input feature vector; $y_{i}$ is the actual label/class for sample $i$; $\hat{y}_{{\bm{\theta}}}({\mathbf{x}}_{i})$ is the prediction of the model for sample $i$; and $\ell(\hat{y}_{{\bm{\theta}}}({\mathbf{x}}_{i}),y_{i})$ is the loss function measuring the “goodness-of-fit" for sample $i$. Here, $\mathcal{D}$ is a divergence between the joint probability distribution of the predictions and sensitive attributes and the Kronecker product of their marginal distributions. Recall that $\hat{y}_{{\bm{\theta}}}$ and $s$ are statistically independent iff $\mathbb{P}({\hat{y}_{{\bm{\theta}}}({\mathbf{x}}),s})$ follows $\mathbb{P}(\hat{y}_{{\bm{\theta}}}({\mathbf{x}}))\otimes\mathbb{P}(s)$. Thus, the second term in (1) is zero iff $\hat{y}_{{\bm{\theta}}}$ and $s$ are statistically independent (complete fairness under the demographic parity notion).

This section studies the fair empirical risk minimization regularized by a broad class of $f$-divergence measures. Let $\mathbb{P}$ and $\mathbb{Q}$ be two discrete probability measures taking values in $\mathcal{P}=\{1,\dots,m\}$. The $f$-divergence between $\mathbb{P}$ and $\mathbb{Q}$ is defined as (Polyanskiy & Wu, 2022, Def 4.9)(see Appendix B for the general continuous case):

The above definition, which is also known as $f$-mutual information (Lu et al., 2023; Csiszár, 1967), covers many known divergence measures used for imposing fairness, such as KL-divergence for the choice of $f(t)=t\log(t)$ (Shui et al., 2022), or $\chi^{2}$ divergence when $f(t)=(t-1)^{2}$ (Lowy et al., 2022). As shown in Appendix C, $\mathcal{D}_{f}$ in (1) is zero if and only if the probability distribution of $s$ and $\hat{y}_{{\bm{\theta}}}$ are statistically independent for the choices of $f$ listed in Table 1. In addition, we prove that these $f$-divergences either cover or provide upper bounds for the popular notions of fairness violations in the literature, such as $\ell_{p}$ distances, Rényi correlation (Baharlouei et al., 2020), and demographic parity (equalized odds) violation. This means that by minimizing these regularizers, we are minimizing an upper bound of (other) popular fairness violation measures, and thus we are controlling them implicitly. Further, unlike Rényi correlation (Baharlouei et al., 2020; Grari et al., 2020), we can utilize Legendre-Fenchel duality (and variational representation) to develop (provably) convergent algorithms with stochastic (mini-batch) updates. This formulation and the resulting stochastic optimization algorithm are described in the next subsection.

In the previous section, we assumed that the training and test domains have the same distribution. However, this assumption is not necessarily valid in certain applications (Fang et al., 2020). In particular, a model that behaves fairly on the training data distribution may have an unfair performance in the test phase. To address this issue, this section develops stochastic algorithms for fair empirical risk minimization via $f$-divergences in the presence of the distribution shifts.

Assume that $\hat{\mathbb{P}}_{s,y}(s,\hat{y})$ is the joint distribution of sensitive attributes and predictions on the training data. The distributionally robust fair empirical risk minimization via $f$-divergences is formulated as:

$\mathcal{B}=\mathcal{B}(\hat{\mathbb{P}},\delta)$ is the distributional uncertainty set defined as a certain ball around the training distribution $\hat{\mathbb{P}}$ with radius $\delta$. This formulation guarantees that the model fairness is preserved (up to a violence of f-divergence less than $\kappa$) even when the test distribution slightly changes. With a slight change of notation, $\hat{\mathbb{P}}$ refers to the training distribution, whereas $\mathbb{P}$ is the optimization parameter.

One can define the uncertainty set through an $\epsilon$ neighborhood around the joint distribution of the training data characterized by a distance measure such as $\ell_{p}$ norms, Wasserstein distance, or MMD distance. While these distributionally robust uncertainty sets are thoroughly analyzed for empirical risk minimization (ERM) (Kuhn et al., 2019; Blanchet et al., 2019; Levy et al., 2020), the DRO formulation for ERM is limited to the Wasserstein distance for the fair logistic regression (Taskesen et al., 2020) and MMD distance (Wang et al., 2023) on the distribution curvature as a heuristic for robustness. Unfortunately, none of these approaches offer a convergent algorithm with stochastic updates. Further, some of these approaches are limited to special loss functions and heuristics. On the other hand, we study imposing the distributionally robust fairness via $f$-divergences for a general loss function where the uncertainty set is characterized by $\ell_{p}$ norms (Section 3.1) or $f$-divergences (Section 3.2). Our results show that the former approach is more suitable when lower levels of robustness for fairness are required, and the latter works better for handling larger distribution shifts.

We use three popular notions of group fairness: demographic parity, equalized odds, and equality of opportunity violations (see Appendix A for definitions) to measure the fairness of trained models. To run Algorithm 1, we set $\eta_{{\bm{\theta}}}$ and $\eta_{\alpha}$ to $10^{-5}$ and $10^{-6}$ respectively in all experiments. Further, by changing $\lambda$, we get different points in the trade-off curve between accuracy and fairness. The range of $\lambda$ depends on the $f$-divergence (see Appendix J for more information on tuning hyper-parameters). In the inference phase of our experiments, we use the standard maximum likelihood decoding based on the output of the softmax layer, i.e., the predicted label is the label with the highest logit value.

As we will see in this section, several $f$-divergence measures lead to reasonable fairness/accuracy tradeoffs and can outperform existing benchmarks. However, no single $f$-divergence measure uniformly outperforms other measures in all the experiments. Thus, we believe in applications, the choice of the $f$-divergence can be viewed as a hyperparameter that can be tuned by cross-validation.

This paper presented a unified stochastic framework for fair empirical risk minimization via $f$-divergences ($f$-FERM). The key idea is to reformulate the objective function as a min-max optimization problem using Legendre-Fenchel duality of $f$-divergence. This enables us to develop an unbiased gradient estimator and a convergent stochastic first-order algorithm. Furthermore, we robustified $f$-FERM using $\ell_{p}$ norm balls as the uncertainty set against distributional changes. While our empirical investigation delves into the performance and fairness distinctions among various $f$-divergences, a more comprehensive analysis is warranted to determine the optimal $f$-divergence concerning the tradeoff between performance and fairness, faster convergence, and asymptotic behaviors. Furthermore, the distributionally robust formulation of fair empirical risk minimization and the advantages of each formulation can be explored beyond $f$-divergences as the measure of fairness violation and $\ell_{p}$ norm balls as uncertainty sets.

This work was supported by the NSF CAREER Award CCF2144985, the AFOSR Young Investigator Program Award FA9550-22-1-0192, a gift from the USC-Meta Center for Research and Education in AI, and a gift from Google.