PRIME: Deep Imbalanced Regression with Proxies

Jongin Lim, Sucheol Lee, Daeho Um, Sung-Un Park, Jinwoo Shin

TL;DR: This paper introduces PRIME, a novel proxy-based representation learning scheme for imbalanced regression.

Abstract: Data imbalance remains a fundamental challenge in real-world machine learning. However, most existing work has focused on classification, leaving imbalanced regression underexplored despite its importance in many applications. To address this gap, we propose PRIME, a framework that leverages learnable proxies to construct a balanced and well-ordered feature space for imbalanced regression. At its core, PRIME arranges proxies to be uniformly distributed in the feature space while preserving the ordinal structure of regression targets, and then aligns each sample feature to its corresponding proxy. By using proxies as reference points, PRIME induces the desired structure of learned representations, promoting better generalization, especially in underrepresented target regions. Moreover, since proxy-based alignment resembles classification, PRIME enables the seamless application of class imbalance techniques to regression, facilitating more balanced feature learning. Extensive experiments demonstrate the effectiveness and broad applicability of PRIME, achieving state-of-the-art performance on four real-world regression benchmark datasets across diverse target domains.

Lay Summary:

In many real-world situations, artificial intelligence (AI) systems struggle when certain types of data are rare—for example, predicting rare events or extreme values. While researchers have made progress on this problem for classification tasks, less attention has been paid to similar challenges in predicting continuous values, known as regression.

To bridge this gap, we developed a new method called PRIME. PRIME uses “proxies”, or reference points, to help the AI model learn in a more balanced and organized way, even when the training data is uneven. These proxies are arranged to reflect the order of target values, and each data point is trained to align with its appropriate proxy.

Our work provides a flexible and unified framework that brings class imbalance solutions into regression tasks, paving the way for a new paradigm in handling imbalanced regression.

Primary Area: Deep Learning->Other Representation Learning

Keywords: Imbalanced Regression, Representation Learning, Proxy

Submission Number: 6759