Balancing the Scales: Reinforcement Learning for Fair Classification

Research on mitigating bias can be divided in those that tackle the training data Wang et al. (2019), those that attempt to remove bias from representations Ravfogel et al. (2020); Haghighatkhah et al. (2022), and those that adjust the learning process Elazar and Goldberg (2018); Han et al. (2021). Within approaches that adjust the learning processes, we distinguish two main categories: those that add adversarial learners to ignore protected attributes Wadsworth et al. (2018), and more closely to our work, approaches that adjust the loss function to emphasize performance on minority classes.

Prior work that modified the training setup to increase fairness used methods such as down/upsampling Wang et al. (2019) and reweighting the loss function Höfler et al. (2005); Lahoti et al. (2020). Han et al. (2022a) evaluate both down-sampling and loss reweighting on two datasets for fair text classification. Both techniques are applied to align training with different definitions of fairness. Downsampling using the Equal Opportunity fairness metric demonstrated impressive results. In this paper, we take the first step to explore whether reward scaling in reinforcement learning can improve fairness in classification.

Literature on RL applications for classification predominantly considers the following two theoretical frameworks: Markov Decision Process (MDP) and the Contextual Multi-Armed Bandit (CMAB).

Early work by Wiering et al. (2011) casts classification as a sequential decision-making task, by introducing a classification variant of the MDP. In their setup agents manipulate memory cells to encode information by applying an action sequence on a single sample. They demonstrated competitive performance, but, remained limited to small tasks due to the computational complexity. Lin et al. (2020) extended this work, by introducing a variant of the classification MDP and applying a Deep Q-learning Network (DQN) to binary classification of images and texts. They focused on mitigating bias arising from class imbalance by scaling the rewards inversely proportional to the class frequency. However, in their setup the sequential component was taken over multiple data points, which assumes sequential dependency among data samples in the classification task.

The RL framework CMAB offers a promising alternative because it considers the input as a sequence of independent states. We formalize our classification task as a CMAB problem, because this is consistent with the independence of data points in the commonly shuffled datasets. Dudík et al. (2014) use CMAB agents by modifying K-class classification as a K-armed bandit problem, where the agent receives a reward of 1 for correct and 0 for incorrect classification. Dimakopoulou et al. (2019) use this framework and modify different CMAB algorithms to balance exploration and exploitation and compare the original and modified agents on 300 classification datasets. However, their analysis focused on datasets with either limited classes, features, or observations. To the best of our knowledge, we are the first to extend reward scaling for fair multi-class classification or to apply reward scaling for classification with CMAB.

We formalize the multi-class classification task as a finite contextual multi-armed bandit (CMAB) problem. In each round $t$, an agent is presented with a context vector $x_{t}\in\mathbb{R}^{d}$. The agent chooses an action $a_{t}\in A$ from a fixed set of arms, based on the policy $a_{t}\sim\pi(x_{t})$. After the action is taken, the environment returns a reward: $r_{t}\sim\mathcal{R}$. In a multi-class classification framework, the action space is the set of all possible classes, while the context vector is a representation of the input, e.g. a contextual text embedding (see Section 4.1 for more information). Within a finite number of rounds, the agent aims to learn the optimal policy to maximize the total reward. In other words, given a set of testing data, we aim to learn the optimal policy to maximize the selection of correct classes.

We extend the CMAB framework for fair classification by constructing a reward function that counters data imbalances. We assign a reward scale for each sensitive state $(a,g)$, comprising the desired class $a$ (e.g. occupation) and protected attribute $g$ (e.g. gender). The total reward for a given prediction is calculated as $\mathcal{R}(a,a_{pred},g)=\mathcal{A}cc(a,a_{pred})\cdot\mathcal{W}(a,g)$. It comprises an accuracy term $\mathcal{A}cc$, and a reward scale matrix $\mathcal{W}$. Unlike previous work Dudík et al. (2014), which defines the accuracy term as $\mathcal{A}cc\in\{0,1\}$, we define it as $\mathcal{A}cc\in\{-1,1\}$. Which allows us to scale the reward for both correct (+1) and incorrect classifications (-1). We use the term reward scale to indicate that this approach adjusts the magnitude but not the sign of the reward. Section 3.3 presents various designs of the reward scale.

We select three different RL algorithms and adapt them to learn optimal policies for fair classification in the formalized CMAB problem. These algorithms include one classical CMAB algorithm that addresses the linear relationship between the expected reward and the context, as well as two popular deep RL algorithms for MDP problems, Deep Q-Network (DQN) and Proximal Policy Optimization (PPO), which allow us to leverage non-linear approximations. The two deep RL algorithms are selected as they are representative of the two key types of deep RL approaches: value-based methods and policy gradient methods. By employing these three algorithms, we aim to investigate the application of diverse RL methods.

Lin et al. (2020) implement reward scaling on an imbalanced binary classification dataset. Inspired by this, we propose different ways of reward scales for multi-class classification with imbalances of protected attributes. For context, we use the profession classification dataset, BiasBios, where reward scaling tackles the sub-class imbalance of the protected attribute gender. To illustrate the influence of various reward scales Figure 2 shows the scales of a balanced (Professor) and an imbalanced (Nurse) class for the protected groups with attribute gender.

For the first method, we extend the implementation of Lin et al. (2020) into the multi-classification setting and reduce the reward for the majority by scaling it with the imbalanced ratio $\rho^{a}_{imb}=\frac{|D^{a}_{min}|}{|D^{a}_{maj}|}$, which is the ratio between the number of samples of the minority and majority class in class $a$.

Figure 2 demonstrates that $\mathcal{W}_{\rho+}(x)$ scales with a reverse of the bias within a class, however, compared to a balanced class, the reward scale of the majority is very low. Therefore, we propose a second design that keeps the scales of the majority group in the imbalanced class equal to the scales of the balanced class. Thus we set the majority value at 1 and only increase the minority value, based on the inversed imbalanced ratio.

The third implementation adopts the Equal Opportunity (EO) formalization used by Han et al. (2022a). Contrary to the previous two methods it ensures the average weights per class remain equal, providing an improved theoretical fairness among classes. The EO objective is achieved by aggregating the loss per sensitive state and then scaling it. However, our work scales per instance, thus we convert the EO objective to instance-specific weights (see Appendix B.3) and obtain:

Lastly, we also employ the Inverse Probability Weighting (IPW) technique Höfler et al. (2005). Full fairness across classes and protected groups is obtained by scaling with the joint probability, resulting in:

Parallel to reward scaling in RL is (instance) reweighing in supervised learning Han et al. (2022a); Lahoti et al. (2020), here loss reweighting for clarity. Loss reweighting has been a popular technique for imbalanced datasets, where the loss of each data sample is scaled to mitigate the class imbalance, traditionally using the IPW Höfler et al. (2005). The weighted cross-entropy loss using the true probability $p$, predicted probability $q$:

We implement supervised learning with loss reweighing to serve as a strong baseline and highlight the connection between loss reweighing and reward scaling.

The BiasBios De-Arteaga et al. (2019) consists of 393,423 biographies labeled with one of 28 professions, and a binary gender label. Following De-Arteaga et al. (2019), the data is randomly split according to 65% training, 25% testing, and 10% for validating. The dataset contains two imbalances: varying frequencies of the professions and a difference in gender percentage for each class.

Following Ravfogel et al. (2020) and Han et al. (2022a) we also evaluate on the Emoji (Elazar and Goldberg, 2018) sentiment analysis task of Twitter data Blodgett et al. (2016). The task involves binary sentiment classification evaluation with race as the protected attribute, approximated through the provided labels Standard American English (SAE) and African American English (AAE). As per Han et al. (2021), the dataset is composed of Happy (40% AAE, 10% SAE), and Sad: (10% AAE, 40% SAE). We use the same train, dev, and test splits of 100k/8k/8k instances, respectively.

Following prior work, we evaluate performance using accuracy and fairness using the True Positive Ratio (TPR) gap De-Arteaga et al. (2019); Ravfogel et al. (2020). The TPR gap of a class $a\in A$ is calculated as: $TPR_{gap}^{a}=TPR_{g}^{a}-TPR_{\sim g}^{a}$, where $g$ and $\sim g$ represent the two options for the sensitive states. The global TPR metric, GAP, is then calculated as the root mean square of the individual metrics:

To quantify performance and fairness as a single metric we use the Distance To the Optimum (DTO) introduced in Han et al. (2022a). DTO combines the metrics (accuracy, 1-GAP) as dimensions of evaluation space and computes the Euclidean distance between the achieved and Utopian point. The smaller the distance to the Utopian point (lower DTO), the better. We report the DTO with the Utopian accuracy and GAP as the best values across all evaluated models.

While accuracy measures the overall performance and GAP the disparity among protected groups within a class, these metrics do not capture imbalance performance across classes. Therefore we also evaluate our algorithms using the macro-averaged F1 metric to detect if minority classes are ignored. All metrics are scaled by 100 for ease of reading and all metrics are represented in the tables as the mean $\pm$ std over 5 random seeds, except DTO which is taken over the mean.

Each algorithm uses the same classifier architecture, except LinUCB, which has a custom set of learnable parameters. The classifier has one hidden layer MLP. All models are trained for 10 epochs, except LinUCB, which achieved optimal performance within 2 epochs. All models are evaluated on the validation set after 50k iterations to account for different convergence speeds of models. The best model throughout training and across hyperparameters is selected using DTO. We apply hyperparameter optimization on both datasets for each of the algorithms, for details see Appendix A.3.

Besides the supervised implementation in Section 3.4, abbreviated to Sup, we also compare our models against various existing debiasing methods. INLP Ravfogel et al. (2020) debiases embeddings by iteratively training classifiers to predict the protected attribute, it then removes this information from the embedding using a projection of the classifier’s nullspace. MP Haghighatkhah et al. (2022) simplifies the INLP setup by using a single Mean Projection (MP) between the representation of each class’s protected groups. DAdv, Han et al. (2021) removes sensitive information from the embeddings by applying adversarial training using diverse adversaries. Lastly, BTEO Han et al. (2022a) subsamples the dataset to establish equal opportunity. We implement these existing methods with the same training settings for fair comparison. Notably, we highlight how Supervised $\mathcal{W}^{EO}$ is theoretically equal to instance reweighing in Han et al. (2022a), but our implementation achieves significantly higher performance.

We first evaluate the effect of the different reward scales discussed in Section 3.3 by providing the results for the implementations of PPO and Sup, see Table 1, for full table see Appendix D.4.

The results presented in the table demonstrate that the imbalance ratio $\rho$ yields substantial gains in fairness and accuracy when applied to increase the reward for the minority class ($\mathcal{W}_{\rho-}$) as opposed to decreasing the reward for the majority class ($\mathcal{W}_{\rho+}$). Especially the accuracy of PPO is sensitive to this, suggesting that PPO might not work very well for states with low reward scales.

As hypothesized, scaling with the joint probability of class and protected attribute as in $\mathcal{W}_{IPW}$ proves to be too unstable. It results in the worst accuracy for both algorithms, with a minor improvement in fairness for supervised learning. The overall difference between $\mathcal{W}_{EO}$ and $\mathcal{W}_{\rho-}$ is minimal as expected from their similar reward scales depicted in Figure 2. We use EO in our experiments because of its better theoretical foundation.

The main results of our experiments are summarized in Table 2.

On the BiasBios dataset, our DQN and PPO implementations achieve strong results, with PPO outperforming DQN in fairness, as PPO’s lower GAP shows. Our supervised implementation, Sup $\mathcal{W}^{EO}$, surpasses all other baselines. LinUCB performs significantly worse on this task, obtaining one of the worst DTO scores. In contrast, on the Emoji dataset, LinUCB achieves one of the best performance-fairness trade-offs, as indicated by the low DTO score, and is only surpassed by BTEO, which has a slightly higher accuracy. On the other hand, DQN obtains a lower accuracy and PPO lower fairness compared to other metrics.

These results suggest that the classical CMAB algorithm LinUCB excels on binary datasets, while the deep RL implementations, DQN and PPO, perform better in multi-class settings. Notably, although DQN and PPO obtain competitive results on the BiasBios, their F1 score is considerably lower than the baselines. Further analysis of per class metrics (see Appendix D.2) reveals that while the F1 for most classes was on par with the supervised setup, both deep RL algorithms failed to recall two of the very sparse classes.

Compared to baseline methods such as BTEO and DAdv, our DQN and PPO implementations demonstrate competitive performance on the BiasBios dataset. Table 2 also shows that, contrary to previous work Han et al. (2022a),³³3The EO scaled supervised implementation of Han et al. (2022a) achieves an Accuracy of 75.7 and GAP of 13.9 loss-scaling for supervised learning (Sup $\mathcal{W}^{EO}$) achieves superior overall performance to downsampling (BTEO). Downsampling only seems to outperform scaling when enough data is present, as demonstrated by its lower GAP for the Emoji dataset.

We investigate the influence of reward scaling on our models by training them with and without scaling. Table 3 presents the results on BiasBios as the mean performance without scaling and the change in metrics when EO scaling is applied.

Without reward scaling the three RL algorithms achieve similar accuracy to the supervised approach but at the cost of a lower F1 score. As mentioned above, the RL algorithms fail on two very sparse classes, which explains the drop in GAP and F1. Failing to classify any instances of a class correctly results in a TPR gap of 0 for that class, since the result is "fair" among both genders.

The EO reward scale significantly reduces the GAP of all implementations, at the cost of a slight decrease in Accuracy and F1 for most models. However, on LinUCB the scaling causes a large performance reduction with only a small GAP reduction, suggesting that scaling hinders the performance more than it improves the fairness.

To inspect the weak effect of reward scaling on LinUCB, we analyze the TPR gap per profession against the gender imbalance for both cases in Figure 3. Without scaling, LinUCB’s performance follows a predictable positive correlation with gender imbalance. For instance, in the Nurse class ($\sim$ 90% women), the model performs better for the majority group (women), resulting in a positive TPR gap. However, after reward scaling this correlation is inverted, causing the model to perform worse for the majority group and better for the minority group. In case of the Nurse class, the model obtains a negative TPR gap. This suggests LinUCB is oversensitive to scaling on the BiasBios dataset, causing it to overcompensate and penalize the majority group.

For a more in-depth analysis of each model’s sensitivity to subclass imbalance, we train them on the Emoji dataset under various stereotyping ratios. A stereotyping ratio represents the proportion of the AAE and SAE samples in each class. For example, a stereotyping ratio of 0.2 means the data is distributed as Happy (20% AAE, 80% SAE), Sad (80% AAE, 20% SAE).

Figure 4 reveals a strong inverse relationship between LinUCB’s fairness and the stereotyping ratio. Although the stereotypical ratios are symmetric at the value of 0.5 the fairness of LinUCB is asymmetric at this value. Thus there is a residual representation bias in the data that is not addressed by the reward scaling. In contrast, supervised learning maintains a mostly stable fairness, except for the most extreme ratios. Interestingly, LinUCB reveals a reverse pattern in best and worst fairness.

The relatively low accuracy of DQN and poor performance on fairness of PPO are consistent across ratios. However, PPO does have the most constant fairness and performance across stereotyping ratios, indicating good training stability.

We now examine how the strength of the protected information affects the efficacy of reward scaling. We focus on two scenarios that modify the gender signal in the representations: 1) adding explicit gender information, thus increasing the gender signal strength 2) debiasing the embeddings using MP, which reduces it. Table 4 presents the mean and relative difference compared to the results without the specified modification.⁴⁴4Full tables in Appendix D.5

Providing the model with gender information increases the overall accuracy. However, the impact on fairness, as indicated by the GAP score varies among algorithms. The GAP score increases for the two algorithms with the lowest GAPs (Sup, PPO) and decreases for the two with the highest GAP (LinUCB, DQN). Only the algorithms that perform less well on fairness benefited from access to protected attribute.

Removing the bias with MP reduces the test accuracy for nearly all algorithms, indicating some useful information is removed. Again, the modification increased relatively low GAP scores, and decreased relatively high scores. As such, changing to a representation with relatively low bias helps LinUCB and DQN, whereas Sup and PPO that already achieved better fairness mainly see their overall performance hindered.

Notably, the differences in Table 4 are relatively small and hardly ever surpass the standard deviation provided in Table 2. This suggests that while the strength of the protected information influences performance and fairness, the impact might be less pronounced than the choice algorithmic design. Moreover, all methods reduced the GAP score compared to only applying MP (which yields a GAP of 13.9, see Table 2).