Invariant Shape Representation Learning For Image Classification
This addresses robustness issues in medical image classification (e.g., brain and cardiovascular MRI) where correlations can vary across environments like age groups, though it is incremental as it builds on existing invariant risk minimization methods.
The paper tackles the problem of spurious correlations between shape features and target variables in image classification by introducing invariant shape representation learning (ISRL), which improves prediction accuracy by learning invariant features across multiple environments, achieving state-of-the-art results on simulated 2D images and real 3D MRI datasets.
Geometric shape features have been widely used as strong predictors for image classification. Nevertheless, most existing classifiers such as deep neural networks (DNNs) directly leverage the statistical correlations between these shape features and target variables. However, these correlations can often be spurious and unstable across different environments (e.g., in different age groups, certain types of brain changes have unstable relations with neurodegenerative disease); hence leading to biased or inaccurate predictions. In this paper, we introduce a novel framework that for the first time develops invariant shape representation learning (ISRL) to further strengthen the robustness of image classifiers. In contrast to existing approaches that mainly derive features in the image space, our model ISRL is designed to jointly capture invariant features in latent shape spaces parameterized by deformable transformations. To achieve this goal, we develop a new learning paradigm based on invariant risk minimization (IRM) to learn invariant representations of image and shape features across multiple training distributions/environments. By embedding the features that are invariant with regard to target variables in different environments, our model consistently offers more accurate predictions. We validate our method by performing classification tasks on both simulated 2D images, real 3D brain and cine cardiovascular magnetic resonance images (MRIs). Our code is publicly available at https://github.com/tonmoy-hossain/ISRL.