Towards Fair and Robust Volumetric CT Classification via KL-Regularised Group Distributionally Robust Optimisation

TL;DR

This paper introduces a KL-regularised group distributionally robust optimization framework for fair and robust CT scan classification, addressing distribution shifts and subgroup disparities.

Contribution

It proposes a novel KL-regularised Group DRO method that stabilizes group weighting and enhances fairness and robustness in medical imaging tasks.

Findings

Achieved 0.835 F1 on COVID-19 classification, surpassing previous best by 5.9 points.

Attained 0.815 macro F1 across genders, outperforming previous methods by 11.1 points.

Significantly improved F1 scores for underrepresented subgroups like female Squamous cell carcinoma.

Abstract

Automated diagnosis from chest computed tomography (CT) scans faces two persistent challenges in clinical deployment: distribution shift across acquisition sites and performance disparity across demographic subgroups. We address both simultaneously across two complementary tasks: binary COVID-19 classification from multi-site CT volumes (Task 1) and four-class lung pathology recognition with gender-based fairness constraints (Task 2). Our framework combines a lightweight MobileViT-XXS slice encoder with a two-layer SliceTransformer aggregator for volumetric reasoning, and trains with a KL-regularised Group Distributionally Robust Optimisation (Group DRO) objective that adaptively upweights underperforming acquisition centres and demographic subgroups. Unlike standard Group DRO, the KL penalty prevents group weight collapse, providing a stable balance between worst-case protection and average performance. For Task 2, we define groups at the granularity of gender class, directly targeting severely underrepresented combinations such as female Squamous cell carcinoma. On Task 1, our best configuration achieves a challenge F1 of 0.835, surpassing the best published challenge entry by +5.9. On Task 2, Group DRO with {\alpha} = 0.5 achieves a mean per-gender macro F1 of 0.815, outperforming the best challenge entry by +11.1 pp and improving Female Squamous F1 by +17.4 over the Focal Loss baseline.