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

TL;DR

This study investigates stain variability as a potential shortcut in renal pathology AI, demonstrating that entropy regularization can mitigate stain bias without sacrificing lesion classification accuracy.

Contribution

The paper introduces a label-free entropy regularization method within a dual-head Bayesian architecture to prevent stain-related shortcut learning in multi-stain renal pathology datasets.

Findings

Stain identity is easily learnable, confirming a shortcut.

Stain supervision does not significantly affect lesion classification performance.

Entropy regularization maintains stain predictions at chance levels without harming accuracy.

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.