AugDiff: Diffusion-Based Feature Augmentation for Multiple Instance Learning in Whole Slide Image

Abstract

Multiple instance learning (MIL), a powerful strategy for weakly supervised learning, is able to perform various prediction tasks on gigapixel whole slide images (WSIs). However, the tens of thousands of patches in WSIs usually incur a vast computational burden for image augmentation, limiting the performance improvement in MIL. Currently, the feature augmentation-based MIL framework is a promising solution, while existing methods such as mixup often produce unrealistic features. To explore a more efficient and practical augmentation method, we introduce the diffusion model (DM) into MIL for the first time and propose a feature augmentation framework called AugDiff. The diverse generation capabilities of DM guarantee a various range of feature augmentations, while its iterative generation approach effectively preserves semantic integrity during these augmentations. We conduct extensive experiments over four distinct cancer datasets, two different feature extractors, and three prevalent MIL algorithms to evaluate the performance of AugDiff. Ablation study and visualization further verify the effectiveness. Moreover, we highlight AugDiff's higher quality augmented feature over image augmentation and its superiority over self-supervised learning. The generalization over external datasets indicates its broader applications. The code is open-sourced on https://github.com/szc19990412/AugDiff .