Probabilistic Bilevel Coreset Selection

TL;DR

This paper introduces a novel probabilistic bilevel approach for coreset selection that learns sample weights to efficiently select representative data subsets, outperforming existing methods especially in noisy and imbalanced data scenarios.

Contribution

It proposes the first continuous probabilistic bilevel formulation for coreset selection, with an efficient solver using unbiased policy gradient, improving performance over traditional greedy methods.

Findings

Outperforms existing coreset selection methods in various tasks.

Effective in challenging label-noise and class-imbalance scenarios.

Provides convergence guarantees for the proposed training procedure.

Abstract

The goal of coreset selection in supervised learning is to produce a weighted subset of data, so that training only on the subset achieves similar performance as training on the entire dataset. Existing methods achieved promising results in resource-constrained scenarios such as continual learning and streaming. However, most of the existing algorithms are limited to traditional machine learning models. A few algorithms that can handle large models adopt greedy search approaches due to the difficulty in solving the discrete subset selection problem, which is computationally costly when coreset becomes larger and often produces suboptimal results. In this work, for the first time we propose a continuous probabilistic bilevel formulation of coreset selection by learning a probablistic weight for each training sample. The overall objective is posed as a bilevel optimization problem, where 1) the inner loop samples coresets and train the model to convergence and 2) the outer loop updates the sample probability progressively according to the model's performance. Importantly, we develop an efficient solver to the bilevel optimization problem via unbiased policy gradient without trouble of implicit differentiation. We provide the convergence property of our training procedure and demonstrate the superiority of our algorithm against various coreset selection methods in various tasks, especially in more challenging label-noise and class-imbalance scenarios.