In-hoc Concept Representations to Regularise Deep Learning in Medical Imaging
This work addresses the critical issue of model robustness and generalization in medical imaging, offering a lightweight, architecture-agnostic solution without dense concept supervision, though it is incremental as it builds on existing concept-based methods.
The paper tackles the problem of deep learning models in medical imaging relying on spurious correlations instead of clinically meaningful features, leading to poor generalization under distribution shifts. It introduces LCRReg, a regularization method using latent concept representations, which significantly improves robustness to spurious correlations and out-of-distribution generalization, as shown in tasks like diabetic retinopathy classification.
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