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

Authors: Guangkun Nie, Gongzheng Tang, Shenda Hong
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Deep imbalanced regression, sparse data region optimization

Abstract:
Imbalanced data distributions are prevalent in real-world scenarios, presenting significant challenges in both classification and regression tasks. This imbalance often causes deep learning models to overfit in regions with abundant data (manyshot regions) while underperforming in regions with sparse data (few-shot regions). Such characteristics limit the applicability of deep learning models across various domains, notably in healthcare, where rare cases often carry greater clinical significance. While recent studies have highlighted the benefits of incorporating distributional information in imbalanced classification tasks, similar strategies have been largely unexplored in imbalanced regression. To address this gap, we propose Dist Loss, a novel loss function that integrates distributional information into model training by jointly optimizing the distribution distance between model predictions and target labels, alongside sample-wise prediction errors. This dual-objective approach encourages the model to balance its predictions across different label regions, leading to significant improvements in accuracy in fewshot regions. We conduct extensive experiments across three datasets spanning computer vision and healthcare: IMDB-WIKI-DIR, AgeDB-DIR, and ECG-KDIR. The results demonstrate that Dist Loss effectively mitigates the impact of imbalanced data distributions, achieving state-of-the-art performance in few-shot regions. Furthermore, Dist Loss is easy to integrate and complements existing methods. To facilitate further research, we provide our implementation at https://github.com/Ngk03/DIR-Dist-Loss.

Decision: accept (poster)
Reviewer scores: 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730436821823

thread_type:
official_review

<thread_object>
creator: reviewer_zurs

<content>
summary:
The paper proposes a new approach, called Dist Loss, for imbalanced deep regression. Specifically, the method generates pseudo-labels taking into consideration label distribution information, and at the same time generates pseudo-predictions. The paper is novel and well-written, but has a number of issues and can be strengthened.

strengths:
The paper proposes an interesting and concise idea by generating pseudo-labels and pseudo-predictions for imbalanced deep regression. The paper is clearly written with elucidatory illustrations for easier understanding.

weaknesses:
All datasets are image datasets and all models are ResNet-like models. The method has not been evaluated on more datasets like text data or tabular data. 

The experiment setting may not be representative enough. For all three real-world datasets, their imbalance rates are fixed. However, the input data could be manually tweaked to generate datasets covering a wider range of imbalance rates, and to investigate the algorithm's performance in scenarios from mildly to extremely heavy-tailed distributions. 

There is no theoretical analysis of the algorithm.

questions:
In Figure 2, what is the exact form of the loss function L(.)? This is not clear from both the figure and the main text. Specifically, in Figure 2, for the two sequences [1, 3, 3, 4, 4, 4, 6] and [1, 2, 2, 3, 5, 6, 7], how is the loss computed?

The first step of the algorithm is discretization which maps the continuous label space Y into B bins. In such a way, the regression problem is transformed into a classification problem. Could you also show classification metrics such as overall accuracy and few-shot accuracy, along with the regression metrics, so that we have a more thorough understanding of the result? Again, what are the correlation metrics (Pearson's or Spearman's) of the results? Could you please also present some actual-by-predicted correlation plots of the prediction?

How did you determine the number of bins B?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730675364249

thread_type:
official_review

<thread_object>
creator: reviewer_7cfx

<content>
summary:
The paper addresses the challenge of imbalanced data distributions in regression tasks, especially the difficulty deep learning models face in accurately predicting few-shot settings. The authors introduce a new loss function, called Dist Loss, designed to align the distribution of model predictions with the label distribution using kernel density estimation and pseudo-predictions. The authors also contribute a differentiable relaxation for their loss function to make the model end-to-end training. The paper conducts experiments on the regression of three datasets, compares with five related works, and has demonstrated significant improvements.

strengths:
The paper has the following strengths:

(i) The problem of imbalanced regression tasks is important to study but less well investigated in related works. The authors contributed a new loss function to train the model and showed good results compared to other baselines in two aspects: using the new loss function integrated into the baseline and combining the existing training loss with their proposed one. The running time is as fast as the mean squared error and significantly faster than other baselines. 

(ii) The idea of using a pseudo-label to define the distance between prediction and label distribution is interesting. Also, the authors applied different sorting techniques to make their Dist Loss ($L(S_P, S_L)$) trainable.

weaknesses:
While the paper achieved good results, its presentation lacks details and information to help the reviewer understand core ideas and how the method efficiently addresses imbalanced problems in regression tasks.

(i) In Figure 2, what is the role of **"Labels"** and **"Predictions"** terms (inside **Dist loss calculation**) while the Loss function already received **Pseudo-labels** and **Sorted Prediction**?

(ii) In Section 3.2.1, authors jumped into detail about how they compute pseudo-labels and pseudo-predictions but they have not addressed two key questions:
+ Why do we need to compute "pseudo-here"? How is it important to deal with the imbalanced dataset?
+ Section 3.2.1 is based on the Empirical label distribution (KDE), but no details about KDE are discussed (either background and how it is used), which makes the readers unable to follow the details later. For e.g., is the KDE step applied at a batch size level or at the whole training dataset?

Another unclear problem is the standard regression formulation, which transforms continuous label space Y into discretized B bins (**Section 3.1**). How is the standard learning scheme trained with these discretized labels? 

(iii) The descriptions in Equation (1) and (2) and the definition of $N_{L'}$ are also hard to imagine without a figure to illustrate core ideas beyond.

Overall, the paper's methodology section is not well presented, with information presented sporadically, and does not address core concerns about imbalanced settings.

questions:
The reviewer has the following questions:


(a) can you explain how the KDE is used in your formulation? What is input data? Do we need to use it in the testing part or only in the training steps?

(b) can you explain why the pseudo-label is important to overcome imbalances in regression?

(c) The ablation study in Table 4 is difficult to capture. In particular, the authors mentioned, *"We demonstrate the effects of using different functions, considering the probability-based inversely weighed MAE and MSE losses."* There are two questions here. First, what is *probability-based inversely* (INV) ? and why $INV - L_1$ or $ INV - L_2$ shows benefit of your proposed one?

(d) because dist_loss works with discrete values in sorted orders, how can you make the training process stable and avoid exploding gradients given large values from the loss function?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731070533298

thread_type:
official_review

<thread_object>
creator: reviewer_c9ay

<content>
summary:
This paper tackles the issue of imbalanced regression, which often results in poor performance in regions with few samples, by proposing a new loss function termed Dist Loss. Dist Loss is designed to reduce the distribution gap between model predictions and actual labels, thereby promoting closer alignment with the target distribution, particularly in sparse data regions. The approach has shown enhanced performance in few-shot scenarios across several datasets, such as IMDB-WIKI-DIR, AgeDB-DIR, and ECG-Ka-DIR, and complements existing techniques well. Through extensive experimentation, Dist Loss demonstrates state-of-the-art accuracy for rare instances, highlighting its relevance for critical fields like healthcare.

strengths:
1. The paper introduces Dist Loss, a new loss function specifically designed to tackle the problem of imbalanced regression, improving model performance in low-sample regions.

2. Dist Loss is easy to incorporate with current techniques and enhances their performance in few-shot areas, making it a practical addition to existing approaches.

3. The proposed method is evaluated on different datasets.

weaknesses:
1. According to the problem setting, the distribution has been discretized using fixed bin size, which can also be adapted into classification setting if one have a look at Figure 1. However the authors claimed in section 2.2 that the methods developed for imbalanced image classification can not be well adapted into imbalanced regression task. From my point of view, if bin is chosen then classification based method can also be used in regression task. More explanation should be provided by the authors.


2. The performance gain of the proposed loss function is not promising compared with Balanced MSE in Table I. It even underperforms BalancedMSE on IMDB/WIKI/DIR dataset using GM metric.

3. The authors are also suggested to compare with some other existing related works, e.g., a.

In a and b, the task is evaluated in four different settings, which are all, many, medium, few, which is also mentioned by the authors. However it seems in the provided Tables there are only results for few shot regions. I find the original different shot settings more convincing since it can showcase the generalizability of the proposed loss function among different shot settings.

Another question, why the reported performance of ConR is different compared with the performance reported by their paper for AgeDB DIR dataset? Did the authors follow the same experimental setting as introduced in ConR and other existing works?

a. Wang, Z., & Wang, H. (2024). Variational imbalanced regression: Fair uncertainty quantification via probabilistic smoothing. Advances in Neural Information Processing Systems, 36.

b. Keramati, M., Meng, L., & Evans, R. D. (2023). Conr: Contrastive regularizer for deep imbalanced regression. ICLR 2024.

4. In Table 3, it seems that balanced MSE is better than the proposed Dist Loss regarding the training time consumption. The benefits from the training efficiency pespective is now well highlighted.


5. I have some concerns with Table 4. From the paper of ConR as mentioned in b, LDS with ConR can achieve better performance on AgeDB DIR dataset, with 9.62 for MAE and 6.87 for AgeDB DIR when we considering the few shot setting. The authors are suggested to compare with all other existing works in Table 4.

questions:
1. In Section 2.2, the authors claim that imbalanced image classification methods cannot be well adapted for imbalanced regression. However, with the chosen discretized distribution using fixed bin sizes, wouldn't classification-based methods be applicable? Could the authors elaborate on why classification-based methods might not work well in this setting and clarify the limitations or challenges?

2. The proposed loss function does not show a strong performance improvement over Balanced MSE, as shown in Table I, and it even underperforms Balanced MSE on the IMDB/WIKI/DIR dataset using the GM metric. Can the authors explain why the proposed loss function struggles to outperform Balanced MSE in these cases and discuss any specific scenarios where it demonstrates advantages?

3.Could the authors compare their method with other existing works, such as Wang & Wang (2024) and Keramati et al. (2023)? These works evaluate the task across different settings (all, many, medium, few), which seems like a more convincing framework to assess generalizability. Why did the authors choose not to include similar shot-based evaluations? Would this additional analysis provide deeper insights into the model’s robustness?


4.The performance of ConR reported in the paper differs from what is reported in the original ConR paper for the AgeDB DIR dataset. Did the authors follow the same experimental settings as in these related works? If there were any differences in the setup, could they clarify these and explain how they might impact the results?

5.Table 3 shows that Balanced MSE is more efficient in terms of training time compared to the proposed Dist Loss. Could the authors discuss the trade-offs between training time efficiency and performance gain? Given the additional time cost, what are the practical advantages of using Dist Loss?
Performance on AgeDB DIR Dataset in Few-Shot Settings:

6. According to Keramati et al. (2023), ConR combined with LDS achieved better results on the AgeDB DIR dataset in the few-shot setting (MAE of 9.62 and 6.87). Why does the proposed method seem to underperform compared to this? Did the authors consider integrating ConR or LDS in their approach to enhance few-shot performance, and if not, what were the reasons?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732475177373

thread_type:
official_comment

<thread_object>
creator: area_chair_qczk

<content>
title:
discussion

comment:
Dear reviewers,

Thank you for your contribution. Soon the discussion period is about to end. Please go over the response from the authors and initiate discussion.

regards

AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
Dear authors, 

This draft received overall positive ratings. Please take in consideration input from the reviewers to update the draft. 
Please not that reviewers have given good rating with the condition of draft to be updated e.g. "I would emphasize that in case this paper is accepted by AC decision, please take into account the above explanations to make major improvements in the method section, either writing part or coming up with a better figure for the main method if possible." 
Other reviewer has asked the code to be released. 

I concur their requests. Draft should be updated and full implementation  (covering all aspects, experiment settings, configs, etc..) should be made publicly available before the camera-ready version is submitted.   

A side note: 
* Abstract should be revised to be improve readability and clarity. 
* ”vanilla model” appears to be defined in the Caption of Fig.1 (Here, the term ”vanilla model” refers to a model that employs
* no specialized techniques to address imbalanced data.). Since the term is being used later in the paper too, such definitions should be in the main draft. 
* Suggestion: Remove "bold" format from "Thus, it is crucial to explore ...."
* Suggestion: Dist Loss is quite generic term, changing to something more specific will help remove any confusion. 




regards

AC