Shortcut Learning in Glomerular AI: Adversarial Penalties Hurt, Entropy Helps

2026-04-09Computer Vision and Pattern Recognition

Computer Vision and Pattern Recognition
AI summary

The authors studied whether AI models that classify kidney lesions in lupus patients accidentally rely on the type of stain used in microscope images instead of the actual disease features. They used a large dataset from multiple centers and stains to test different model setups, including one that tries to ignore stain information without having stain labels. They found that while models can easily learn stain type, their lesion predictions were not biased by it when trained properly. They also showed that adding a special method to discourage the model from depending on stain details did not reduce accuracy. This suggests that well-prepared data and certain model designs can prevent unwanted shortcuts related to staining in kidney pathology AI.

stain variabilityrenal pathologylupus nephritisglomerular lesionsBayesian CNNVision Transformer (ViT)Monte Carlo dropoutentropy regularizationdistribution shiftshortcut learning
Authors
Mohammad Daouk, Jan Ulrich Becker, Neeraja Kambham, Anthony Chang, Hien Nguyen, Chandra Mohan
Abstract
Stain variability is a pervasive source of distribution shift and potential shortcut learning in renal pathology AI. We ask whether lupus nephritis glomerular lesion classifiers exploit stain as a shortcut, and how to mitigate such bias without stain or site labels. We curate a multi-center, multi-stain dataset of 9{,}674 glomerular patches (224$\times$224) from 365 WSIs across three centers and four stains (PAS, H\&E, Jones, Trichrome), labeled as proliferative vs.\ non-proliferative. We evaluate Bayesian CNN and ViT backbones with Monte Carlo dropout in three settings: (1) stain-only classification; (2) a dual-head model jointly predicting lesion and stain with supervised stain loss; and (3) a dual-head model with label-free stain regularization via entropy maximization on the stain head. In (1), stain identity is trivially learnable, confirming a strong candidate shortcut. In (2), varying the strength and sign of stain supervision strongly modulates stain performance but leaves lesion metrics essentially unchanged, indicating no measurable stain-driven shortcut learning on this multi-stain, multi-center dataset, while overly adversarial stain penalties inflate predictive uncertainty. In (3), entropy-based regularization holds stain predictions near chance without degrading lesion accuracy or calibration. Overall, a carefully curated multi-stain dataset can be inherently robust to stain shortcuts, and a Bayesian dual-head architecture with label-free entropy regularization offers a simple, deployment-friendly safeguard against potential stain-related drift in glomerular AI.