Controllable Latent Space Augmentation for Digital Pathology

TL;DR

This paper introduces HistAug, a controllable latent space augmentation method for digital pathology that efficiently generates realistic, semantically controlled image embeddings to improve multiple instance learning models, especially with limited data.

Contribution

HistAug is a novel generative model enabling efficient, controllable augmentation of pathology image embeddings, addressing limitations of traditional patch-level augmentation methods.

Findings

HistAug outperforms existing augmentation methods across multiple tasks.

It significantly improves model performance in low-data regimes.

Learned transformations outperform noise-based perturbations.

Abstract

Whole slide image (WSI) analysis in digital pathology presents unique challenges due to the gigapixel resolution of WSIs and the scarcity of dense supervision signals. While Multiple Instance Learning (MIL) is a natural fit for slide-level tasks, training robust models requires large and diverse datasets. Even though image augmentation techniques could be utilized to increase data variability and reduce overfitting, implementing them effectively is not a trivial task. Traditional patch-level augmentation is prohibitively expensive due to the large number of patches extracted from each WSI, and existing feature-level augmentation methods lack control over transformation semantics. We introduce HistAug, a fast and efficient generative model for controllable augmentations in the latent space for digital pathology. By conditioning on explicit patch-level transformations (e.g., hue, erosion), HistAug generates realistic augmented embeddings while preserving initial semantic information. Our method allows the processing of a large number of patches in a single forward pass efficiently, while at the same time consistently improving MIL model performance. Experiments across multiple slide-level tasks and diverse organs show that HistAug outperforms existing methods, particularly in low-data regimes. Ablation studies confirm the benefits of learned transformations over noise-based perturbations and highlight the importance of uniform WSI-wise augmentation. Code is available at https://github.com/MICS-Lab/HistAug.