FairTune: Optimizing Parameter Efficient Fine Tuning for Fairness in Medical Image Analysis

TL;DR

This work considers the highly effective workflow of adapting pre-trained models to downstream medical imaging tasks using parameter-efficient fine-tuning (PEFT) techniques and proposes FairTune, a framework to optimise the choice of PEFT parameters with respect to fairness.

Abstract

Training models with robust group fairness properties is crucial in ethically sensitive application areas such as medical diagnosis. Despite the growing body of work aiming to minimise demographic bias in AI, this problem remains challenging. A key reason for this challenge is the fairness generalisation gap: High-capacity deep learning models can fit all training data nearly perfectly, and thus also exhibit perfect fairness during training. In this case, bias emerges only during testing when generalisation performance differs across subgroups. This motivates us to take a bi-level optimisation perspective on fair learning: Optimising the learning strategy based on validation fairness. Specifically, we consider the highly effective workflow of adapting pre-trained models to downstream medical imaging tasks using parameter-efficient fine-tuning (PEFT) techniques. There is a trade-off between updating more parameters, enabling a better fit to the task of interest vs. fewer parameters, potentially reducing the generalisation gap. To manage this tradeoff, we propose FairTune, a framework to optimise the choice of PEFT parameters with respect to fairness. We demonstrate empirically that FairTune leads to improved fairness on a range of medical imaging datasets. The code is available at https://github.com/Raman1121/FairTune

Figures (5)

• Figure 1: Bias arises during train-test generalisation. Left (Training AUROC): High-capacity deep models can exhibit perfect group fairness during training because they can classify all the training data perfectly. Right (Validation AUROC): Bias arises because the disadvantaged subgroup has worse generalisation error than the privileged subgroup. Fine-tuning ViT-Base on the Papila dataset.
• Figure 2: Illustration that shows how our approach optimises the structure of PEFT with respect to fairness. Hyperparameter optimisation (HPO) selects a mask that decides which components of a pre-trained model $\theta$ are fine-tuned using PEFT. For each sampled mask, the fine-tuned model is evaluated on the validation set to compute the fairness loss $\mathcal{L}^{fair}$, which is then reported to the HPO algorithm that decides what masks to sample and which is the final best option.
• Figure 3: FairTune leads to stable fine-tuning with reduced differences between the best and worst performing subgroups compared to conventional fine-tuning from Figure \ref{['fig:teaser']}.
• Figure 4: Frequency of selecting a specific component for fine-tuning across different scenarios.
• Figure 5: Optimization trajectory for the outer loop PEFT mask search. The plots shown here are for the 12-bit Attention Tuning PEFT search space. The saturation in the objective value demonstrates that the objective is saturated within 200 outer loop iterations trials.