End-to-end Multiple Instance Learning with Gradient Accumulation

TL;DR

This paper introduces a gradient accumulation training strategy for attention-based deep multiple instance learning, enabling end-to-end training on large histopathological images without GPU memory limitations, achieving comparable performance to traditional methods.

Contribution

The paper presents a novel gradient accumulation method that allows direct end-to-end training of ABMIL models on large-scale data without memory constraints.

Findings

Memory-efficient training matches baseline performance.

Gradient accumulation enables end-to-end learning on gigapixel images.

Training time is longer but results are comparable to existing methods.

Abstract

Being able to learn on weakly labeled data, and provide interpretability, are two of the main reasons why attention-based deep multiple instance learning (ABMIL) methods have become particularly popular for classification of histopathological images. Such image data usually come in the form of gigapixel-sized whole-slide-images (WSI) that are cropped into smaller patches (instances). However, the sheer size of the data makes training of ABMIL models challenging. All the instances from one WSI cannot be processed at once by conventional GPUs. Existing solutions compromise training by relying on pre-trained models, strategic sampling or selection of instances, or self-supervised learning. We propose a training strategy based on gradient accumulation that enables direct end-to-end training of ABMIL models without being limited by GPU memory. We conduct experiments on both QMNIST and Imagenette to investigate the performance and training time, and compare with the conventional memory-expensive baseline and a recent sampled-based approach. This memory-efficient approach, although slower, reaches performance indistinguishable from the memory-expensive baseline.