Aligning (Medical) LLMs for (Counterfactual) Fairness

TL;DR

This paper addresses unfairness in medical LLMs by first rigorously profiling bias patterns across datasets, models, and prompting strategies via a counterfactual fairness framework. It then introduces a fairness-focused alignment method that leverages a teacher model and Preference Optimization within a knowledge-distillation pipeline, generating a bias-aware preference dataset through counterfactual prompts and semantic-similarity ranking. The proposed approach, implemented with LoRA-based fine-tuning and SimPO, demonstrates consistent reductions in bias across Q-Pain, Treatment Recommendation, and Triage tasks with minimal impact on overall task performance. The work highlights the importance of scalable fairness interventions in clinical AI, shows the value of model alignment to mitigate bias without sacrificing accuracy, and lays the groundwork for practical deployment in diverse medical settings.

Abstract

Large Language Models (LLMs) have emerged as promising solutions for a variety of medical and clinical decision support applications. However, LLMs are often subject to different types of biases, which can lead to unfair treatment of individuals, worsening health disparities, and reducing trust in AI-augmented medical tools. Aiming to address this important issue, in this study, we present a new model alignment approach for aligning LLMs using a preference optimization method within a knowledge distillation framework. Prior to presenting our proposed method, we first use an evaluation framework to conduct a comprehensive (largest to our knowledge) empirical evaluation to reveal the type and nature of existing biases in LLMs used for medical applications. We then offer a bias mitigation technique to reduce the unfair patterns in LLM outputs across different subgroups identified by the protected attributes. We show that our mitigation method is effective in significantly reducing observed biased patterns. Our code is publicly available at \url{https://github.com/healthylaife/FairAlignmentLLM}.

Figures (11)

• Figure 1: Visual description of the evaluation framework. By randomly rotating patient demographics within standardized clinical scenarios (red teaming), we assess the impact of sensitive attributes on LLM outputs across different clinical tasks, LLM types, and prompting techniques.
• Figure 2: Left panel: Proposed pipeline for fairness-aware model alignment in three steps. The non-shaded yellow-blocks area describes the dataset generation process, and the blue hatched area is the preference ranking step, where both candidate answers are compared to the reference answer. The red arrow denotes the alignment process through Preference Optimization (PO). Right panel: Example of the generation process for the preference dataset.
• Figure 3: Results on the Q-Pain dataset. The bars represent the average maximum difference probability of denying the pain treatment between two subgroups for each question. The error bars show the standard deviation. CNC: Chronic Non Cancer, ANC: Acute Non Cancer, Post Op: Postoperative
• Figure 4: Results on the Treatment Recommendation dataset with NEJM Healer vignettes.
• Figure 5: Results on the Triage dataset on a Likert Scale. The LLMs were presented with patient summaries and statements and were asked to rate their agreement with the statement. 1:Strongly disagree with the statement. 5:Strongly agree.