Provable Multi-instance Deep AUC Maximization with Stochastic Pooling

TL;DR

This paper introduces a novel stochastic pooling approach for deep AUC maximization in multi-instance learning, effectively handling large bag sizes with provable convergence and demonstrating superior performance on various datasets.

Contribution

It proposes a unified, provable multi-instance deep AUC maximization algorithm using stochastic pooling methods to address large bag size challenges in MIL.

Findings

The MIDAM algorithm converges at a rate comparable to state-of-the-art DAM methods.

Stochastic pooling methods improve computational efficiency for large bags.

Experimental results show superior performance on MIL and medical datasets.

Abstract

This paper considers a novel application of deep AUC maximization (DAM) for multi-instance learning (MIL), in which a single class label is assigned to a bag of instances (e.g., multiple 2D slices of a CT scan for a patient). We address a neglected yet non-negligible computational challenge of MIL in the context of DAM, i.e., bag size is too large to be loaded into {GPU} memory for backpropagation, which is required by the standard pooling methods of MIL. To tackle this challenge, we propose variance-reduced stochastic pooling methods in the spirit of stochastic optimization by formulating the loss function over the pooled prediction as a multi-level compositional function. By synthesizing techniques from stochastic compositional optimization and non-convex min-max optimization, we propose a unified and provable muli-instance DAM (MIDAM) algorithm with stochastic smoothed-max pooling or stochastic attention-based pooling, which only samples a few instances for each bag to compute a stochastic gradient estimator and to update the model parameter. We establish a similar convergence rate of the proposed MIDAM algorithm as the state-of-the-art DAM algorithms. Our extensive experiments on conventional MIL datasets and medical datasets demonstrate the superiority of our MIDAM algorithm.