A Penalty Approach for Normalizing Feature Distributions to Build Confounder-Free Models

TL;DR

This paper introduces PDMN, a penalty-based extension of MDN, enabling stable, sample-wide estimation of confounder effects in feature normalization, improving model accuracy and independence from confounders in clinical data applications.

Contribution

The paper proposes a novel penalty approach for MDN, allowing trainable, sample-wide confounder normalization adaptable to various architectures.

Findings

PDMN improves model accuracy over MDN.

PDMN enhances independence from confounders.

Effective on MRI datasets and synthetic data.

Abstract

Translating machine learning algorithms into clinical applications requires addressing challenges related to interpretability, such as accounting for the effect of confounding variables (or metadata). Confounding variables affect the relationship between input training data and target outputs. When we train a model on such data, confounding variables will bias the distribution of the learned features. A recent promising solution, MetaData Normalization (MDN), estimates the linear relationship between the metadata and each feature based on a non-trainable closed-form solution. However, this estimation is confined by the sample size of a mini-batch and thereby may cause the approach to be unstable during training. In this paper, we extend the MDN method by applying a Penalty approach (referred to as PDMN). We cast the problem into a bi-level nested optimization problem. We then approximate this optimization problem using a penalty method so that the linear parameters within the MDN layer are trainable and learned on all samples. This enables PMDN to be plugged into any architectures, even those unfit to run batch-level operations, such as transformers and recurrent models. We show improvement in model accuracy and greater independence from confounders using PMDN over MDN in a synthetic experiment and a multi-label, multi-site dataset of magnetic resonance images (MRIs).