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

2025-08-19

AI Summary

Deep learning models for medical imaging are highly susceptible to distribution shifts and often rely on spurious correlations rather than clinically meaningful features. To address this, we propose Latent Concept Regularization (LCRReg): a lightweight, multi-concept–multi-class compatible regularization method that operates in the CNN latent space. LCRReg requires no concept labels from the primary training set; instead, it synthesizes high-fidelity, disentangled concept samples using only a small auxiliary dataset. It constructs concept-attribute vector (CAV)-based subspace constraints to enforce alignment between latent representations and semantically grounded concepts. Experiments demonstrate that LCRReg significantly improves model robustness against spurious correlations and out-of-distribution inputs—outperforming multitask learning, linear probing, and other baselines on both synthetic and real-world medical tasks. By bridging interpretability and robustness without sacrificing scalability, LCRReg establishes a new paradigm for trustworthy, clinically aligned AI in medicine.

Technology Category

• Computer Vision: Large Vision Models
• Machine Learning: Adversarial Learning & Robustness
• Natural Language Processing: Safety and Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data

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

Problem

• Improves generalization of medical imaging models under distribution shifts
• Reduces reliance on spurious correlations in deep learning
• Guides models toward clinically meaningful feature representations

Innovation

• Leverages Latent Concept Representations for regularization
• Uses auxiliary dataset to synthesize concept examples
• Guides CNN activation within concept-associated latent subspaces