A Multi-Component Reward Function with Policy Gradient for Automated Feature Selection with Dynamic Regularization and Bias Mitigation

TL;DR

This work addresses bias arising from hidden dependencies in high-dimensional feature spaces by introducing a reinforcement learning-based feature selection framework. It defines a multi-component reward that couples predictive performance (AUC) with direct bias penalties and indirect penalties derived from a causal-graph representation, plus sparsity incentives, all within a policy-gradient (REINFORCE) learning loop. An MLP-based policy sequentially selects feature subsets, while a non-differentiable decision-tree predictor evaluates subset quality on masked data, guiding the agent via the reward signal. Applied to credit scoring, the approach demonstrates improved discrimination (ROC AUC) and reduced bias relative to baselines, with case studies illustrating how direct and indirect penalties influence feature choice and convergence. The work highlights a principled path toward fairness-aware, data-driven feature selection in high-stakes domains, while noting potential limitations in capturing complex non-linear dependencies and suggesting directions for deeper causal integration.

Abstract

Static feature exclusion strategies often fail to prevent bias when hidden dependencies influence the model predictions. To address this issue, we explore a reinforcement learning (RL) framework that integrates bias mitigation and automated feature selection within a single learning process. Unlike traditional heuristic-driven filter or wrapper approaches, our RL agent adaptively selects features using a reward signal that explicitly integrates predictive performance with fairness considerations. This dynamic formulation allows the model to balance generalization, accuracy, and equity throughout the training process, rather than rely exclusively on pre-processing adjustments or post hoc correction mechanisms. In this paper, we describe the construction of a multi-component reward function, the specification of the agents action space over feature subsets, and the integration of this system with ensemble learning. We aim to provide a flexible and generalizable way to select features in environments where predictors are correlated and biases can inadvertently re-emerge.

Figures (7)

• Figure 1: The model architecture consists of the standard reinforcement learning components: agent, actions, environment, state, and reward. The agent is parameterized as a multi-layer perceptron (MLP) with ReLU activations and a softmax output layer to approximate the action probability distribution.
• Figure 2: ROC curve comparing the performance of the RL policy with feature selection against baseline random forest and logistic regression models. This figure illustrates that our proposed model outperforms the baseline random forest and logistic regression models.
• Figure 3: This figure illustrates the relationship between AUC and bias scores across different models. The comparison accounts for both direct and indirect sources of bias, providing insight into each model’s balance between predictive accuracy and fairness.
• Figure 4: Indirect penalty over episodes, showing steady reduction and stabilization, indicating minimized feature redundancy and convergence of the learning policy.
• Figure 5: Total rewards over episodes, the plot shows improved policy performance compared to the random baseline.