In-hoc Concept Representations to Regularise Deep Learning in Medical Imaging

TL;DR

LCRReg is a novel regularisation method that uses latent concept representations to improve the robustness and generalisation of deep learning models in medical imaging, especially under distribution shifts.

Contribution

It introduces LCRReg, a concept-guided regularisation technique that does not need concept labels in main training, enhancing model robustness in medical imaging tasks.

Findings

LCRReg improves robustness to spurious correlations in synthetic datasets.

LCRReg enhances out-of-distribution generalisation in diabetic retinopathy classification.

The method is lightweight and architecture-agnostic, outperforming several baselines.

Abstract

Deep learning models in medical imaging often achieve strong in-distribution performance but struggle to generalise under distribution shifts, frequently relying on spurious correlations instead of clinically meaningful features. We introduce LCRReg, a novel regularisation approach that leverages Latent Concept Representations (LCRs) (e.g., Concept Activation Vectors (CAVs)) to guide models toward semantically grounded representations. LCRReg requires no concept labels in the main training set and instead uses a small auxiliary dataset to synthesise high-quality, disentangled concept examples. We extract LCRs for predefined relevant features, and incorporate a regularisation term that guides a Convolutional Neural Network (CNN) to activate within latent subspaces associated with those concepts. We evaluate LCRReg across synthetic and real-world medical tasks. On a controlled toy dataset, it significantly improves robustness to injected spurious correlations and remains effective even in multi-concept and multiclass settings. On the diabetic retinopathy binary classification task, LCRReg enhances performance under both synthetic spurious perturbations and out-of-distribution (OOD) generalisation. Compared to baselines, including multitask learning, linear probing, and post-hoc concept-based models, LCRReg offers a lightweight, architecture-agnostic strategy for improving model robustness without requiring dense concept supervision. Code is available at the following link: https://github.com/Trustworthy-AI-UU-NKI/lcr\_regularization