Revisiting End-to-End Learning with Slide-level Supervision in Computational Pathology

TL;DR

This paper revisits end-to-end learning in computational pathology, addressing optimization challenges with a novel MIL called ABMILX, leading to state-of-the-art results with efficient training.

Contribution

It introduces ABMILX, a new MIL method that mitigates E2E optimization issues, enabling efficient training and improved performance in computational pathology.

Findings

ABMILX outperforms state-of-the-art models on multiple benchmarks.

E2E training with ABMILX is computationally efficient (<10 RTX3090 hours).

The study highlights the potential of E2E learning in computational pathology.

Abstract

Pre-trained encoders for offline feature extraction followed by multiple instance learning (MIL) aggregators have become the dominant paradigm in computational pathology (CPath), benefiting cancer diagnosis and prognosis. However, performance limitations arise from the absence of encoder fine-tuning for downstream tasks and disjoint optimization with MIL. While slide-level supervised end-to-end (E2E) learning is an intuitive solution to this issue, it faces challenges such as high computational demands and suboptimal results. These limitations motivate us to revisit E2E learning. We argue that prior work neglects inherent E2E optimization challenges, leading to performance disparities compared to traditional two-stage methods. In this paper, we pioneer the elucidation of optimization challenge caused by sparse-attention MIL and propose a novel MIL called ABMILX. It mitigates this problem through global correlation-based attention refinement and multi-head mechanisms. With the efficient multi-scale random patch sampling strategy, an E2E trained ResNet with ABMILX surpasses SOTA foundation models under the two-stage paradigm across multiple challenging benchmarks, while remaining computationally efficient (<10 RTX3090 hours). We show the potential of E2E learning in CPath and calls for greater research focus in this area. The code is https://github.com/DearCaat/E2E-WSI-ABMILX.