SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhimeng Jiang $^{1*}$ , Xiaotian Han $^{1*}$ , Hongye Jin $^{1}$ , Guanchu Wang $^{2}$ , Rui Chen $^{3}$ , Na Zou $^{1}$ , Xia Hu $^{2}$ $^{1}$ Texas A&M University,  $^{2}$ Rice University,  $^{3}$ Samsung Electronics America
Fairness in machine learning has attracted increasing attention in recent years. The fairness methods improving algorithmic fairness for in-distribution data may not perform well under distribution shifts. In this paper, we first theoretically demonstrate the inherent connection between distribution shift, data perturbation, and model weight perturbation. Subsequently, we analyze the sufficient conditions to guarantee fairness (i.e., low demographic parity) for the target dataset, including fairness for the source dataset, and low prediction difference between the source and target datasets for each sensitive attribute group. Motivated by these sufficient conditions, we propose robust fairness regularization (RFR) by considering the worst case within the model weight perturbation ball for each sensitive attribute group. We evaluate the effectiveness of our proposed RFR algorithm on synthetic and real distribution shifts across various datasets. Experimental results demonstrate that RFR achieves better fairness-accuracy trade-off performance compared with several baselines. The source code is available at https://github.com/zhimengj0326/RFR_NeurIPS23
Previous research [1-4] has shown that a classifier trained on a specific source distribution will perform worse when testing on a different target distribution, due to distribution shift. Recently, many studies have focused on investigating the impact of distribution shift on machine learning models, where fairness performance degradation is even more significant than that of prediction performance [5]. The sensitivity of fairness over distribution shift challenges machine learning models in high stake applications, such as criminal justice [6], healthcare [7], and job marketing [8]. Thus the transferability of the fairness performance under distribution shift is a crucial consideration for real-world applications.
To achieve the fairness of the model (already achieved fairness on the source dataset) on the target dataset, we first reveal that distribution shift is equivalent to model weight perturbation, and then seek to achieve fairness under distribution shift via model weight perturbation. Specifically,  $i$ ) we reveal the inherent connection between distribution shift, data perturbation, and model weight perturbation. We theoretically demonstrate that any distribution shift can be equivalent to data perturbation and model weight perturbation in terms of loss value. In other words, the effect of distribution shift on model training can be attributed to data or model perturbation.  $ii$ ) Given the established connection between the distribution shift and model weight perturbation, we next tackle fairness under the distribution shift problem. We first investigate demographic parity relation between source and target datasets. Achieving fair prediction (e.g., low demographic parity) in the source dataset is insufficient for fair prediction in the target dataset. More importantly, the average prediction difference between source and target datasets with the same sensitive attribute also matters for achieving fairness in target datasets.
(b)
Figure 1: The overview of distribution shift understanding. The left part demonstrates distribution shift can be transformed as data perturbation, while the right part shows that data perturbation and model weight perturbation are equivalent.
Motivated by the established equivalence connection, we propose robust fairness regularization (RFR) to enforce average prediction robustness over model weight. In this way, the well-trained model can tackle the distribution shift problem in terms of demographic parity. Considering the expensive computation for the inner maximization problem in RFR, we accelerate RFR by obtaining the approximate closed-form model weight perturbation using first-order Taylor expansion. In turn, an efficient RFR algorithm, trading the inner maximization problem into two forward and backward propagations for each model update, is proposed to achieve robust fairness under distribution shift. Our contributions are highlighted as follows:
- We theoretically reveal the inherent connection between distribution shift, data perturbation, and model weight perturbation. In other words, distribution shift and model perturbation are equivalent in terms of loss value for any model architecture and loss function.
- We first analyze the sufficient conditions to guarantee fairness transferability under distribution shift. Based on the established connection, we propose RFR explicitly pursuing sufficient conditions for robust fairness by considering the worst case of model perturbation.
- We evaluate the effectiveness of RFR on various real-world datasets with synthetic and real distribution shifts. Experiments results demonstrate that RFR can mostly achieve better fairness-accuracy tradeoff with both synthetic and real distribution shifts.
2 Understanding Distribution Shift
In this section, we first provide the notations used in this paper. And then, we theoretically understand the relations between distribution shift, data perturbation, and model weight perturbation, as shown in Figure 1.
We consider source dataset  $\mathcal{D}_S$  and target dataset  $\mathcal{D}_{\mathcal{T}}$ , defined as a probability distribution  $\mathcal{P}_S$  and  $\mathcal{P}_{\mathcal{T}}$  for samples  $S \in S$  and  $T \in T$ , respectively. Each sample defines values for three random variables: features  $X$  with arbitrary domain  $\mathcal{X}$ , binary sensitive attribute  $A \in \mathcal{A} = \{0,1\}$ , and label  $Y$  arbitrary domain  $\mathcal{Y}$ , i.e.,  $\mathcal{S} = \mathcal{T} = \mathcal{X} \times \mathcal{A} \times \mathcal{Y}$ . We denote  $\delta$  as data perturbation on source domain  $\mathcal{S}$ , where  $\delta_X(X)$  and  $\delta_Y(Y)$  data perturbation of features and labels. Using  $\mathcal{P}(\cdot)$  to denote the space of probability distributions over some domain, we denote the space of distributions over examples as  $\mathcal{P}(\mathcal{S})$ . We use  $||\cdot||_p$  as  $L_p$  norm. Let  $f_{\theta}(x)$  be the output of the neural networks parameterized with  $\theta$  to approximate the true label  $y$ . We define  $l(f_{\theta}(x), y)$  as the loss function for neural network training, and the optimal model parameters  $\theta^*$  training on source data  $\mathcal{S}$  is given by  $\theta^* = \arg \min_{\theta} \mathcal{R}_S$ , where  $\mathcal{R}_S = \mathbb{E}_{(X,Y) \sim \mathcal{P}_S}[l(f_{\theta}(X), Y)]$ . Since the target distribution  $\mathcal{P}_{\mathcal{T}}$  may be different with source distribution  $\mathcal{P}_S$ , the well-trained model  $f_{\theta^*}(\cdot)$  trained on source dataset  $\mathcal{S}$  typically does not perform well in target dataset  $\mathcal{T}$ .
2.2 Distribution Shift is Data Perturbation
Deep neural networks are immensely challenged by data quality or dynamic environments [9, 10]. The well-trained deep neural network model may not perform well if existing feature/label noise in training data or distribution shifts among training and target environments. In this subsection, we reveal the inherent equivalence between distribution shift and data perturbation for any neural networks  $f_{\theta}(\cdot)$ , source distribution  $\mathcal{P}_S$ , and target distribution  $\mathcal{P}_{\mathcal{T}}$  using optimal transport. We first provide the formal definition of optimal transport as follows:
Definition 1 (Optimal Transport [II]). Considering two distributions with probability distribution  $P_{S}$  and  $P_{\mathcal{T}}$ , and cost function moving from  $s$  to  $t$  as  $c(s,t)$ , optimal transport between probability distribution  $P_{S}$  and  $P_{\mathcal{T}}$  is given by
$$
\gamma^ {*} (s, t) = \arg \inf  _ {\gamma \in \Gamma (P _ {S}, P _ {T})} \iint c (s, t) \gamma (s, t) d s d t, \tag {1}
$$
where the distribution set  $\Gamma(P_{\mathcal{S}}, P_{\mathcal{T}})$  is the collection of all possible transportation plans and given by  $\Gamma(P_{\mathcal{S}}, P_{\mathcal{T}}) = \left\{\gamma(s, t) > 0, \int \gamma(s, t) \mathrm{d}t = P_{\mathcal{S}}(s), \int \gamma(s, t) \mathrm{d}s = P_{\mathcal{T}}(t)\right\}$ . In other words, the distribution set consists of all possible joint distributions with margin distribution  $P_{\mathcal{S}}$  and  $P_{\mathcal{T}}$ .
Based on the definition of optimal transport, we demonstrate that distribution shift is equivalent to data perturbation in the following theorem:
Theorem 2.1. For any two different distributions with probability distribution  $P_{S}$  and  $P_{\mathcal{T}}$ , adding data perturbation  $\delta$  on source data  $S$  can make perturbated source data and target data with the same distribution, where the distribution of data perturbation  $\delta$  is given by
$$
\mathcal {P} (\delta) = \int_ {\mathcal {S}} \gamma^ {*} (s, s + \delta) \mathrm {d} s. \tag {2}
$$
Additionally, for any positive  $p > 0$ , data perturbation  $\delta$  with minimal power  $\mathbb{E}[||\delta ||_p^p]$  is given by Eq.(2) if optimal transport plan  $\gamma^{*}(\cdot ,\cdot)$  is calculated based on Eq. (1) with cost function  $c(s,t) = ||s - t||_p^p$ .
Proof sketch. Given two different distributions, there are many possible perturbations to move one distribution to another. Optimal transport can select the optimal feasible perturbations or transportation plan in terms of specific objectives (e.g., perturbations power). Given the optimal transportation plan, we can derive the distribution of perturbations based on basic probability theory.
Theorem 2.1 demonstrates the equivalence between distribution shift and data perturbation, where data perturbation distribution is dependent on optimal transport between source and target distribution. The intuition is that such distribution shift can be achieved via optimal transportation (i.e., data perturbation). Based on Theorem 2.1 we have the following Corollary 2.2 on the equivalent of neural network behavior for source and target data:
Corollary 2.2. Given source and target datasets with probability distribution  $\mathcal{P}_S$  and  $\mathcal{P}_{\mathcal{T}}$ , there exists data perturbation  $\delta$  so that the training loss of any neural network  $f_{\theta}(\cdot)$  for target distribution equals that for source distribution with data perturbation  $\delta$ , i.e.,
$$
\mathbb {E} _ {(X, Y) \sim \mathcal {P} _ {\mathcal {T}}} [ l (f _ {\theta} (X), Y) ] = \mathbb {E} _ {\delta_ {X} (X), \delta_ {Y} (Y)} \mathbb {E} _ {(X, Y) \sim \mathcal {P} _ {\mathcal {S}}} [ l (f _ {\theta} (X + \delta_ {X} (X)), Y + \delta_ {Y} (Y)) ]. \qquad (3)
$$
In other words, for the model trained with loss minimization on source data, the deteriorated performance on the target dataset stems from the perturbation of features and labels. The proof sketch is based on Theorem 2.1 since adding a perturbation on the source dataset can be consistent with the target dataset distribution. Therefore, the conclusion can hold for any loss function that is only dependent on the dataset.
2.3 Data Perturbation Equals Model Weight Perturbation
Although we understand that the distribution shift can be attributed to the data perturbation of features and labels in source data, it is still unclear how to tackle the distribution shift issue. A natural solution
2 Data perturbation  $\delta$  includes the perturbation for features  $\delta_X(X)$  and labels  $\delta_Y(Y)$ .
is to adopt adversarial training to force the well-trained model to be robust over data perturbation. However, it is complicated to generate data perturbation on features and labels simultaneously, and many well-developed adversarial training methods are mainly designed for adversarial feature perturbation. Fortunately, we show that model weight perturbation is equivalent to data perturbation by Theorem 2.3.
Theorem 2.3. Considering the source dataset with distribution  $\mathcal{P}_S$ , suppose the source dataset is perturbed with data perturbation  $\delta$ , and the neural network is given by  $f_{\theta}(\cdot)$ , for general case, there exists model weight perturbation  $\Delta \theta$  so that the training loss on perturbed source dataset is the same with that for model weight perturbation  $\Delta \theta$  on source distribution:
$$
\mathbb {E} _ {\delta_ {X} (X), \delta_ {Y} (Y)} \mathbb {E} _ {(X, Y) \sim \mathcal {P} _ {\mathcal {S}}} [ l (f _ {\theta} (X + \delta_ {X} (X)), Y + \delta_ {Y} (Y)) ] = \mathbb {E} _ {(X, Y) \sim \mathcal {P} _ {\mathcal {S}}} [ l (f _ {\theta + \Delta \theta} (X), Y) ]. \tag {4}
$$
Proof sketch. The loss under data perturbation and weight perturbation can be analyzed using first-order Tayler expansion. The critical step is to find the condition of the equivalence for the first order term of data perturbation and weight perturbation. Fortunately, the existence of such conditions can be easily proved using linear algebra.
Theorem 2.3 demonstrates that the training loss on perturbed data distribution (but fixed model weight) equals the training loss on perturbed model weight (but original data distribution). In other words, the training loss fluctuation from data perturbation can equal model weight perturbation. Furthermore, we conclude that chasing a robust model over data perturbation can be achieved via model weight perturbation, i.e., finding a "flattened" local minimum in terms of the target objective is sufficient for robustness.
In this section, we first analyze the sufficient condition to achieve robust fairness over distribution shift in terms of demographic parity. Based on the analysis, we propose a simple yet effective robust fairness regularization via explicit adopting group model weight perturbation. Note that the proposed robust fairness regularization involves a computation-expensive maximization problem, we further accelerate model training via first-order Taylor expansion.
3.1 Robust Fairness Analysis
We consider binary sensitive attribute  $A \in \{0,1\}$  and demographic parity as fairness metric, i.e., the average prediction gap of model  $f_{\theta}(\cdot)$  for different sensitive attribute groups in target dataset  $\Delta DP_{\mathcal{T}} = |\mathbb{E}_{\mathcal{T}_0}[f_\theta (\mathbf{x})] - \mathbb{E}_{\mathcal{T}_1}[f_\theta (\mathbf{x})]|$ , where  $\mathcal{T}_0$  and  $\mathcal{T}_1$  represent the target datasets for sensitive attribute  $A = 0$  and  $A = 1$ , respectively. However, only source dataset  $\mathcal{S}$  is available for neural network training. In other words, even though demographic parity on source dataset  $\Delta DP_{\mathcal{S}} = |\mathbb{E}_{\mathcal{S}_0}[f_\theta (\mathbf{x})] - \mathbb{E}_{\mathcal{S}_1}[f_\theta (\mathbf{x})]|$  is low, demographic parity on target dataset  $\Delta DP_{\mathcal{T}}$  may not be guaranteed to be low due to distribution shift.
To investigate robust fairness over distribution shift, we try to reveal the connection between demographic parity for source and target datasets. We bound the demographic parity difference for source and target datasets as follows:
$$
\begin{array}{l} D P _ {\mathcal {T}} \stackrel {(a)} {\leq} D P _ {\mathcal {S}} + \left| | \mathbb {E} _ {\mathcal {T} _ {0}} [ f _ {\theta} (\mathbf {x}) ] - \mathbb {E} _ {\mathcal {T} _ {1}} [ f _ {\theta} (\mathbf {x}) ] | - | \mathbb {E} _ {\mathcal {S} _ {0}} [ f _ {\theta} (\mathbf {x}) ] - \mathbb {E} _ {\mathcal {S} _ {1}} [ f _ {\theta} (\mathbf {x}) ] | \right| \\ \stackrel {(b)} {\leq} D P _ {\mathcal {S}} + \left| \mathbb {E} _ {\mathcal {S} _ {0}} \left[ f _ {\theta} (\mathbf {x}) \right] - \mathbb {E} _ {\mathcal {T} _ {0}} \left[ f _ {\theta} (\mathbf {x}) \right] \right| + \left| \mathbb {E} _ {\mathcal {S} _ {1}} \left[ f _ {\theta} (\mathbf {x}) \right] - \mathbb {E} _ {\mathcal {T} _ {1}} \left[ f _ {\theta} (\mathbf {x}) \right] \right|, \tag {5} \\ \end{array}
$$
where inequality (a) and (b) hold due to  $a - b \leq |a - b|$  and  $\left|\left|a - b\right| - \left|a' - b'\right|\right| \leq \left|a - a'\right| + \left|b - b'\right|$ , respectively, for any  $a, a', b, b'$ . In other words, in order to minimize demographic parity for target dataset  $DP_{\mathcal{T}}$ , the objective of  $DP_{S}$  minimization is insufficient. The minimization of prediction difference for source and target datasets given sensitive attribute groups  $A = 0$  and  $A = 1$ , defined as  $\Delta_0 = |\mathbb{E}_{\mathcal{S}_0}[f_\theta(\mathbf{x})] - \mathbb{E}_{\mathcal{T}_0}[f_\theta(\mathbf{x})]|$  and  $\Delta_1 = |\mathbb{E}_{\mathcal{S}_1}[f_\theta(\mathbf{x})] - \mathbb{E}_{\mathcal{T}_1}[f_\theta(\mathbf{x})]|$ , are also beneficial to achieve robust fairness over distribution shift in terms of demographic parity.
The bound in Eq. (5) is tight when both condition (a)  $DP_{S} \leq DP_{\mathcal{T}}$  and (b) maximum and minimum of value set  $\min \left\{\mathbb{E}_{\mathcal{S}_0}[f_\theta (\mathbf{x})],\mathbb{E}_{\mathcal{T}_0}[f_\theta (\mathbf{x})]\right\}$ ,  $\min \left\{\mathbb{E}_{\mathcal{S}_1}[f_\theta (\mathbf{x})],\mathbb{E}_{\mathcal{T}_1}[f_\theta (\mathbf{x})]\right\}$  both are from source or target distribution. Even though conditions (a) and (b) may not hold for the neural network model, we
would like to that our goal is not to obtain a tight bound for demographic parity on target distribution. Instead, we aim to find sufficient conditions to guarantee low demographic parity and such low demographic parity can be achieved in model training without any target distribution information. The proposed upper bound Eq. (5) actually reveals sufficient conditions, which can be achieved by our proposed RFR algorithm.
3.2 Robust Fairness Regularization
Motivated by Section 3.1 and distribution shift understanding in Section 2, we develop a robust fairness regularization to achieve robust fairness over distribution shift. Section 3.1 demonstrates that fair prediction on the target dataset requires fair prediction on the source dataset and low prediction difference between the source and target dataset for each sensitive attribute, i.e., low  $\Delta_0 = |\mathbb{E}_{\mathcal{S}_0}[f_\theta (\mathbf{x})] - \mathbb{E}_{\mathcal{T}_0}[f_\theta (\mathbf{x})]|$  and  $\Delta_{1} = |\mathbb{E}_{\mathcal{S}_{1}}[f_{\theta}(\mathbf{x})] - \mathbb{E}_{\mathcal{T}_{1}}[f_{\theta}(\mathbf{x})]|$ . Based on Theorem 2.3, there exists  $\epsilon_0$  so that the following equality holds:
$$
\Delta_ {0} = \left| \mathbb {E} _ {\mathcal {S} _ {0}} \left[ f _ {\theta} (\mathbf {x}) \right] - \mathbb {E} _ {\epsilon_ {0}} \mathbb {E} _ {\mathcal {S} _ {0}} \left[ f _ {\theta + \epsilon_ {0}} (\mathbf {x}) \right] \right|. \tag {6}
$$
Note that the distribution shift is unknown, we consider the worst case for model weight perturbation  $\epsilon_0$  within  $L_{p}$ -norm perturbation ball with radius  $\rho$  as follows:
$$
\Delta_ {0} = \left| \mathbb {E} _ {\mathcal {S} _ {0}} \left[ f _ {\theta} (\mathbf {x}) \right] - \mathbb {E} _ {\epsilon_ {0}} \mathbb {E} _ {\mathcal {S} _ {0}} \left[ f _ {\theta + \epsilon_ {0}} (\mathbf {x}) \right] \right| \leq \max  _ {\| \epsilon_ {0} \| _ {p} \leq \rho} \left| \mathbb {E} _ {\mathcal {S} _ {0}} \left[ f _ {\theta + \epsilon_ {0}} (\mathbf {x}) \right] - \mathbb {E} _ {\mathcal {S} _ {0}} \left[ f _ {\theta} (\mathbf {x}) \right] \right|, \tag {7}
$$
where  $\| \cdot \|_p$  represents  $L_p$  norm,  $\rho$  and  $p$  are hyperparameters. Note that the feasible region of model weight perturbation  $\epsilon_0$  is symmetric, and the neural network prediction is locally linear around parameter  $\theta$ , we have  $\mathbb{E}_{S_0}[f_{\theta + \epsilon_0}(\mathbf{x})] - \mathbb{E}_{S_0}[f_{\theta}(\mathbf{x})] \approx -\left(\mathbb{E}_{S_0}[f_{\theta - \epsilon_0}(\mathbf{x})] - \mathbb{E}_{S_0}[f_{\theta}(\mathbf{x})]\right)$  due to the local linearity. In other words, the absolute operation can be removed if we consider the maximation problem in a symmetric feasible region since there are always non-negative values for the pair perturbation  $\epsilon_0$  and  $-\epsilon_0$ . Therefore, we can further bound  $\Delta_0$  as
$$
\begin{array}{l} {\Delta_ {0}} \leq {\max _ {\| \epsilon_ {0} \| _ {p} \leq \rho} | \mathbb {E} _ {\mathcal {S} _ {0}} [ f _ {\theta + \epsilon_ {0}} (\mathbf {x}) ] - \mathbb {E} _ {\mathcal {S} _ {0}} [ f _ {\theta} (\mathbf {x}) ] |} \\ \approx \max  _ {\| \epsilon_ {0} \| _ {p} \leq \rho} \mathbb {E} _ {\mathcal {S} _ {0}} [ f _ {\theta + \epsilon_ {0}} (\mathbf {x}) ] - \mathbb {E} _ {\mathcal {S} _ {0}} [ f _ {\theta} (\mathbf {x}) ] \triangleq \mathcal {L} _ {R F R, \mathcal {S} _ {0}}, \tag {8} \\ \end{array}
$$
Similarly, we can bound the prediction difference for source and target distribution with sensitive group  $A = 1$  as follows:
$$
\Delta_ {1} \leq \max  _ {\| \epsilon_ {1} \| _ {p} \leq \rho} \mathbb {E} _ {\mathcal {S} _ {1}} \left[ f _ {\theta + \epsilon_ {1}} (\mathbf {x}) \right] - \mathbb {E} _ {\mathcal {S} _ {1}} \left[ f _ {\theta} (\mathbf {x}) \right] \triangleq \mathcal {L} _ {R F R, \mathcal {S} _ {1}}. \tag {9}
$$
Therefore, demographic parity for source and target distribution relation is given by  $DP_{\mathcal{T}} \leq DP_{S} + \mathcal{L}_{RFR,S_0} + \mathcal{L}_{RFR,S_1}$ , we propose robust fairness regularization (RFR) to achieve robust fairness as
$$
\mathcal {L} _ {R F R} = \mathcal {L} _ {R F R, S _ {0}} + \mathcal {L} _ {R F R, S _ {1}}. \tag {10}
$$
It is worth noting that our proposed RFR is agnostic to the training loss function and model architectures. Additionally, we follow [12] to accelerate model training using sharpness-aware minimization via trading maximization problem can be simplified as two forward and two backward propagations. More details on training acceleration are in Appendix E.
3.3 The Proposed Method
In this subsection, we introduce how to use the proposed RFR to achieve robust fairness, i.e., the fair model (low  $\Delta DP_{S}$ ) trained on the source dataset will also be fair on the target dataset (low  $\Delta DP_{\mathcal{T}}$ ). Considering binary sensitive attribute  $A \in \{0,1\}$  and binary classification problem  $Y \in \{0,1\}$ , the classification loss is denoted as
$$
\mathcal {L} _ {C L F} = \mathbb {E} _ {\mathcal {S}} [ - Y f _ {\theta} (X) - (1 - Y) (1 - f _ {\theta} (X)) ]. \tag {11}
$$
To achieve fairness, we consider demographic parity as fairness regularization, i.e.,
$$
\mathcal {L} _ {D P} = \left| \mathbb {E} _ {\mathcal {S} _ {0}} \left[ f _ {\theta} (\mathbf {x}) \right] - \mathbb {E} _ {\mathcal {S} _ {1}} \left[ f _ {\theta} (\mathbf {x}) \right] \right|, \tag {12}
$$