Dist Loss: Enhancing Regression in Few-Shot Region through Distribution Distance Constraint

TL;DR

The paper tackles imbalanced regression by formulating Dist Loss, a differentiable loss that minimizes distribution distance between predictions and true labels. It achieves this by estimating the label distribution with kernel density estimation, constructing pseudo-labels and pseudo-predictions, and employing fast differentiable sorting to compare distributions through a differentiable objective $L(\mathcal{S}_P,\mathcal{S}_L)$ alongside the standard regression loss. Empirical results on IMDB-WIKI-DIR, AgeDB-DIR, and ECG-K-DIR show state-of-the-art performance in few-shot and competitive results in median-shot regions, with Dist Loss offering complementary gains to existing methods. The approach demonstrates robustness across diverse domains and suggests a practical path to improving regression in highly skewed data settings, including healthcare applications. All mathematical notation in the paper is expressed with proper delimitation, e.g., discretization into $B$ bins, KDE-based label distributions $p_i$, and the distribution distance $L(\mathcal{S}_P,\mathcal{S}_L)$.

Abstract

Imbalanced data distributions are prevalent in real-world scenarios, posing significant challenges in both imbalanced classification and imbalanced regression tasks. They often cause deep learning models to overfit in areas of high sample density (many-shot regions) while underperforming in areas of low sample density (few-shot regions). This characteristic restricts the utility of deep learning models in various sectors, notably healthcare, where areas with few-shot data hold greater clinical relevance. While recent studies have shown the benefits of incorporating distribution information in imbalanced classification tasks, such strategies are rarely explored in imbalanced regression. In this paper, we address this issue by introducing a novel loss function, termed Dist Loss, designed to minimize the distribution distance between the model's predictions and the target labels in a differentiable manner, effectively integrating distribution information into model training. Dist Loss enables deep learning models to regularize their output distribution during training, effectively enhancing their focus on few-shot regions. We have conducted extensive experiments across three datasets spanning computer vision and healthcare: IMDB-WIKI-DIR, AgeDB-DIR, and ECG-Ka-DIR. The results demonstrate that Dist Loss effectively mitigates the negative impact of imbalanced data distribution on model performance, achieving state-of-the-art results in sparse data regions. Furthermore, Dist Loss is easy to integrate, complementing existing methods.

Figures (7)

• Figure 1: A real-world healthcare task of potassium (K$^+$) concentration regression from ECGs. (a) Both hyperkalemia (high K$^+$) and hypokalemia (low K$^+$) are predominantly found in the few-shot region, with normal K$^+$ are located in the many-shot region. Hyperkalemia and hypokalemia are life-threatening conditions that can lead to cardiac arrest and ventricular fibrillation, necessitating accurate and timely detection. Conversely, normal K$^+$ concentrations (the many-shot region) are of little concern, as inaccurate and untimely detection of these samples has minimal impact. Here, we follow pmlr-v139-yang21m to define the few-, median-, many-shot regions. (b) illustrates the significant distribution discrepancy between the vanilla model's predictions and the labels, stemming from the imbalanced data distribution. Here, the term "vanilla model" refers to a model that employs no specialized techniques to address imbalanced data. The orange histogram represents the label distribution, while the blue histogram depicts the prediction distribution from the vanilla model. It is evident that the model's predictions are heavily concentrated in the many-shot region and seldom fall into the few-shot region. (c) demonstrates the effectiveness of Dist Loss in reducing the distribution discrepancy. The orange histogram indicates the label distribution, and the blue histogram shows the prediction distribution from the model enhanced with Dist Loss. It is clear that the distribution discrepancy is significantly reduced.
• Figure 2: The presence of imbalanced data distributions introduces a noticeable distribution discrepancy between the model’s predictions and labels. Dist Loss mitigates this imbalance by simultaneously minimizing this discrepancy and sample-wise prediction errors. Initially, KDE is applied to estimate the label distribution and compute the expected frequency of each label within a batch, thereby generating pseudo-labels that incorporate label distribution information. For example, given the labels [1, 3, 4, 6] and their computed expected frequencies [1, 2, 3, 1], the resulting pseudo-labels would be [1, 3, 3, 4, 4, 4, 6], where each label appears according to its expected frequency. Subsequently, the model’s predictions within a batch are sorted to obtain an ordered sequence that captures the prediction distribution. For instance, if the model’s initial predictions are [5, 2, 6, 3, 2, 7, 1], sorting them yields [1, 2, 2, 3, 5, 6, 7], preserving the distributional characteristics of the predictions. Measuring the distance between these pseudo-labels and pseudo-predictions, which both encapsulate distribution information, provides an approximation of the distributional discrepancy. By optimizing both the distribution distance and sample-level prediction errors during training, the model effectively alleviates the adverse effects of imbalanced data, significantly enhancing accuracy, particularly in few-shot regions.
• Figure 3: To illustrate the core concept behind Dist Loss, the figure simplifies its computation by assuming that the batch size equals the total number of training samples.
• Figure 4: Overview of label distributions in the training sets for the IMDB-WIKI-DIR, AgeDB-DIR, and ECG-K-DIR datasets. The classification of shot types for IMDB-WIKI-DIR and AgeDB-DIR follows the definitions provided in pmlr-v139-yang21m.
• Figure 5: Data distribution diagrams for the eight datasets derived from the ECG-K-DIR dataset with varying imbalance ratios.