Class-based Subset Selection for Transfer Learning under Extreme Label Shift

TL;DR

This paper introduces a novel class-based subset selection method using Wasserstein distance for transfer learning under extreme label shift, improving performance when source and target label spaces differ significantly.

Contribution

It proposes a new approach that selects and weights source classes to optimize transfer learning across divergent label distributions, with theoretical analysis and empirical validation.

Findings

Outperforms existing methods in extreme label shift scenarios

Effectively handles disjoint label spaces

Provides theoretical bounds for transfer performance

Abstract

Existing work within transfer learning often follows a two-step process -- pre-training over a large-scale source domain and then finetuning over limited samples from the target domain. Yet, despite its popularity, this methodology has been shown to suffer in the presence of distributional shift -- specifically when the output spaces diverge. Previous work has focused on increasing model performance within this setting by identifying and classifying only the shared output classes between distributions. However, these methods are inherently limited as they ignore classes outside the shared class set, disregarding potential information relevant to the model transfer. This paper proposes a new process for few-shot transfer learning that selects and weighs classes from the source domain to optimize the transfer between domains. More concretely, we use Wasserstein distance to choose a set of source classes and their weights that minimize the distance between the source and target domain. To justify our proposed algorithm, we provide a generalization analysis of the performance of the learned classifier over the target domain and show that our method corresponds to a bound minimization algorithm. We empirically demonstrate the effectiveness of our approach (WaSS) by experimenting on several different datasets and presenting superior performance within various label shift settings, including the extreme case where the label spaces are disjoint.