[Paper Review] EnD: Entangling and Disentangling deep representations for bias correction

Enzo Tartaglione, Carlo Alberto Barbano+1 more|arXiv (Cornell University)|Mar 2, 2021

Adversarial Robustness in Machine Learning45 references75 citations

TL;DR

This paper proposes EnD, a regularization method that corrects bias in deep neural networks by entangling features from the same target class while disentangling features sharing the same bias label, using an information bottleneck at a network layer. The method improves generalization on unbiased test sets without additional training complexity, achieving state-of-the-art performance on biased datasets like IMDB and CORDA for COVID-19 CXR classification.

ABSTRACT

Artificial neural networks perform state-of-the-art in an ever-growing number of tasks, and nowadays they are used to solve an incredibly large variety of tasks. There are problems, like the presence of biases in the training data, which question the generalization capability of these models. In this work we propose EnD, a regularization strategy whose aim is to prevent deep models from learning unwanted biases. In particular, we insert an "information bottleneck" at a certain point of the deep neural network, where we disentangle the information about the bias, still letting the useful information for the training task forward-propagating in the rest of the model. One big advantage of EnD is that we do not require additional training complexity (like decoders or extra layers in the model), since it is a regularizer directly applied on the trained model. Our experiments show that EnD effectively improves the generalization on unbiased test sets, and it can be effectively applied on real-case scenarios, like removing hidden biases in the COVID-19 detection from radiographic images.

Motivation & Objective

• Address the problem of model generalization failure due to hidden biases in training data, especially in real-world applications like medical imaging.
• Develop a regularization strategy that corrects bias without modifying data or adding trainable components to the model.
• Enable deep models to learn task-relevant features by discouraging reliance on spurious correlations linked to bias factors such as data origin or demographic attributes.
• Demonstrate effectiveness on highly biased datasets and real-world scenarios, including COVID-19 diagnosis from chest X-rays.
• Provide a post-hoc, model-agnostic regularization technique that can be applied after training to mitigate bias propagation.

Proposed method

• Apply a regularizer at a specific layer Γ in a trained deep neural network to create an information bottleneck.
• Entangle deep feature representations from samples of the same target class to preserve task-relevant information.
• Disentangle deep feature representations from samples sharing the same bias label (e.g., data origin or age group) to suppress spurious correlations.
• Optimize the regularizer via standard back-propagation, without requiring additional training or auxiliary models.
• Use the regularizer to minimize mutual information between biased features and the final prediction, favoring unbiased representations.
• Integrate the regularizer into the loss function during training, allowing the model to learn to avoid biased features while maintaining performance on the main task.

Experimental results

Research questions

• RQ1 Can a regularizer be designed to disentangle bias-related features while preserving task-relevant representations in deep models?
• RQ2 Does the proposed EnD method improve model generalization on unbiased test sets without requiring retraining or data modification?
• RQ3 How effective is EnD in correcting bias when the bias is not explicitly known (e.g., data origin or image artifacts in medical imaging)?
• RQ4 Can EnD outperform existing debiasing techniques such as adversarial training or feature inversion in real-world, highly biased scenarios?
• RQ5 Does EnD lead to more robust and interpretable attention maps in vision models, such as Grad-CAM, by focusing on relevant anatomical regions?

Key findings

• On the IMDB Face dataset, EnD improved gender prediction accuracy from 57.84% to 69.40% on the EB1 test set and from 59.86% to 65.49% on EB2 when trained on EB1, outperforming baseline and prior methods.
• For age prediction, EnD increased performance from 48.91% to 74.25% on the EB1 test set when trained on EB2, demonstrating strong mitigation of gender-based bias.
• On the CORDA-CDSS dataset, EnD improved the true negative rate from 59.26% to 76.30% and the balanced accuracy from 64.63% to 72.22%, reducing data-source bias.
• On CORDA-SLG, EnD increased the true positive rate from 52.14% to 68.37% and balanced accuracy from 69.88% to 75.94%, showing improved robustness across data sources.
• Grad-CAM visualization confirmed that EnD-regularized models focus on relevant lung regions, while vanilla models attend to irrelevant artifacts, indicating improved interpretability.
• EnD achieved state-of-the-art performance compared to methods like BlindEye and Kim et al. [16], especially on highly biased settings (e.g., ρ = 0.999 in Biased MNIST and IMDB).