PLM: Partial Label Masking for Imbalanced Multi-label Classification

TL;DR

This paper introduces Partial Label Masking (PLM), a novel training technique that improves multi-label classification on imbalanced datasets by adaptively balancing class ratios during training, enhancing minority class recall and overall performance.

Contribution

The paper presents a general, adaptable method for addressing long-tail imbalance in multi-label classification, extending beyond single-label approaches and compatible with various objectives and strategies.

Findings

PLM improves recall on minority classes in imbalanced datasets.

PLM enhances overall classification performance on multiple datasets.

The method is versatile and can be integrated with existing strategies.

Abstract

Neural networks trained on real-world datasets with long-tailed label distributions are biased towards frequent classes and perform poorly on infrequent classes. The imbalance in the ratio of positive and negative samples for each class skews network output probabilities further from ground-truth distributions. We propose a method, Partial Label Masking (PLM), which utilizes this ratio during training. By stochastically masking labels during loss computation, the method balances this ratio for each class, leading to improved recall on minority classes and improved precision on frequent classes. The ratio is estimated adaptively based on the network's performance by minimizing the KL divergence between predicted and ground-truth distributions. Whereas most existing approaches addressing data imbalance are mainly focused on single-label classification and do not generalize well to the multi-label case, this work proposes a general approach to solve the long-tail data imbalance issue for multi-label classification. PLM is versatile: it can be applied to most objective functions and it can be used alongside other strategies for class imbalance. Our method achieves strong performance when compared to existing methods on both multi-label (MultiMNIST and MSCOCO) and single-label (imbalanced CIFAR-10 and CIFAR-100) image classification datasets.