Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data

Authors: Yunze Tong, Fengda Zhang, Zihao Tang, Kaifeng Gao, Kai Huang, Pengfei Lyu, Jun Xiao
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: robustness, score model, reweighting

Abstract:
Machine learning models often perform well on tabular data by optimizing average prediction accuracy. However, they may underperform on specific subsets due to inherent biases and spurious correlations in the training data, such as associations with non-causal features like demographic information. These biases lead to critical robustness issues as models may inherit or amplify them, resulting in poor performance where such misleading correlations do not hold. Existing mitigation methods have significant limitations: some require prior group labels, which are often unavailable, while others focus solely on the conditional distribution $P(Y|X)$, upweighting misclassified samples without effectively balancing the overall data distribution $P(X)$. To address these shortcomings, we propose a latent score-based reweighting framework. It leverages score-based models to capture the joint data distribution $P(X, Y)$ without relying on additional prior information. By estimating sample density through the similarity of score vectors with neighboring data points, our method identifies underrepresented regions and upweights samples accordingly. This approach directly tackles inherent data imbalances, enhancing robustness by ensuring a more uniform dataset representation. Experiments on various tabular datasets under distribution shifts demonstrate that our method effectively improves performance on imbalanced data.

Decision: reject
Reviewer scores: 6, 8, 5, 6

Reviews and discussion:
<thread>
creation date:
1730101902259

thread_type:
official_review

<thread_object>
creator: reviewer_4w6d

<content>
summary:
This work introduces a latent score-based reweighting framework to mitigate biases and spurious correlations in machine learning models on tabular data. The method overcomes limitations of existing techniques by capturing the joint data distribution without needing prior group labels and by focusing on balancing the overall data distribution, not just the conditional distribution. It identifies and upweights underrepresented samples based on the similarity of score vectors, thereby enhancing robustness and ensuring a more uniform dataset representation. The approach has been shown to effectively improve performance on imbalanced data across various datasets, even under distribution shifts.

strengths:
1) This paper captures the joint data distribution without the need for prior group labels. It enhances model robustness by ensuring a more uniform dataset representation. Some cases are provided to explain their motivation.  
2) The experimental results on various tabular datasets have demonstrated the effectiveness of the proposed method.

weaknesses:
1) Besides boundary-based methods that are mainly compared in this paper, other approaches, like stable learning, are also proposed to balance the overall data distribution. These methods should be also discussed and compared.
2) The design of the method is primarily driven by motivation. The effective generalization bounds require theoretical analysis and proof. In particular, the impact of the error from using proxies instead of density estimates on the model’s generalization capability needs to be explained, along with its limitations.
3) One of the challenges in estimating density estimates is the requirement for a large amount of data; that is, a substantial number of samples are needed to achieve accurate density estimates. It is suggested to investigate whether the proposed proxy of density is independent of the sample size. Relevant experiments, like testing the method on subsets of the datasets with varying sizes and comparing the results,  are recommended to test the effectiveness of the method from this perspective.

questions:
Why is the method proposed in this paper limited to tabular data? Please give the reasons why it is not applicable to other data types, such as images or natural language.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730364587080

thread_type:
official_review

<thread_object>
creator: reviewer_f1qs

<content>
summary:
The paper studies the distribution shift problem on tabular data. A diffusion model is used to estimate the data density. Then a sample weigthed is computed based on the density score to make the dataset more balanced. The proposed method is novel. Experiments on several tabular datasets with distribution shifts validate the effectivenss of the proposed method.

strengths:
1. The distribution shift problem on tabular data is important.
2. Applying diffusion model to estimate the sample weight to reweight the sample is novel.
3. Experiments results on several classical tabular datasets validate the proposed method could relieve the distribution shift problem.

weaknesses:
1. Motivation is not clear. The title includes the term "tabular data". However, the proposed model does not have any specific design on the tabular data. Moreover, why cannot the model be used in image data? There a lot of image benchmarks can be used to evaluate the problem method, such as WILDS, colored  MNIST, domainbed, waterbirds, celebA, etc.
2. Concept is confused. The manuscript that previous method mainly solves P(Y|X), then the proposed method solve P(X,Y) shift problem (i.e., both P(X) and P(Y|X) problem). To my understanding, the score is estimated only based on P(X), so it could only solve P(X) problem. How the model could handle P(Y|X) problem?
3. Related work has factual error. For the "Achieving robustness with prior information" paragraph, a lots of work do not require the prior information, such as HRM, stable learning, etc. The authors should carefully discuss the difference with them.
4. If the model could only solve the P(X) problem, the stable learning methods could also handle this problem. What is your major advantange?
5. The experimental resutls are weak. The improvement over baselines are marginal. More common used datasets on image should be tested (see W1). More related baselines should be compared, such as methods in [1].
[1] Change is Hard: A Closer Look at Subpopulation Shift

questions:
Please reponse the questions in weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730621803604

thread_type:
official_review

<thread_object>
creator: reviewer_ftyk

<content>
summary:
The paper attempts to improve robustness by modeling the joint distribution and taking a global perspective on re-weighing as opposed to the popular methods which mostly focus on rectifying the samples near the decision boundary. The observation of most current robustness methods are limited by the pre-trained classification boundaries $P(Y \mid X)$ adds a fresh perspective to the robustness literature. The paper has solid experiments and ablations to support its claims and it is well written.

strengths:
- [S1] Well motivated paper, clearly highlighting the issue in current literature that the methodology is attempting to mitigate.
- [S2] The demo code seems well written and simple enough to run and adapt.
- [S3] I really like the idea proposed by the authors as it intuitively seems to overcome quite a few problems, the experiments and ablations provide good support to the approach.

weaknesses:
- [W1] While I like the motivation presented in terms of the sinusoidal function and the multi-modal gaussian, it would be interesting to look at a more real world scenario where one can expect the proposed method to suit the setup better than traditional methods. That would further help strengthen the case.
- [W2] Table 1 does not include variance values in the main paper, it would be helpful to contrast the effectiveness of the method across various runs, given the general discussion about variance in other parts of the paper as well. Effectively combining Table 5 and Table 6 in the Table 1 and Table 2 or referencing the variation with some reformatting could be possible in the given space. Even moving Algorithm 1, potentially at the cost of Related Work might be a good idea.
- [W3] While the visualization in Figure 3 helps gain insight into how the method works as a proxy, it would be interesting to see some analytical reasoning behind why this might work at scale for larger datasets.

questions:
- [Q1] I’m unsure of how, despite being 2nd best on the mean performance in 2 columns in Table 2 (Columns 11 to 19) the authors claim the method is not best in this scenario due to variations in the predictive mechanism as opposed to Table 1, where the method is not the best (or even top 2) on 3 datasets. So the story doesn’t translate as well, do the authors have justifications other than the intuition as to why they made this claim? It is not obvious how the change in the causal mechanism is a potential cause for variation in the predictive mechanism; more details on this would help clarify the statement.
- [Q2] I am also slightly curious about the formalism behind Section A.4, is there prior literature investigating or stating these are the non-causal attributes? If so are all other attributes not selected necessary causal? The selection of these attributes plays a big role in the methodology, so including correlation statistics or defining a proper heuristic for selection of these attributes would help make a stronger case.
- [Q3] Is there an intuitive way to extend this approach to multiclass classification? Would it be computationally feasible?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730684700672

thread_type:
official_review

<thread_object>
creator: reviewer_bage

<content>
summary:
This paper proposes a latent score-based reweighting framework for improving robustness of classification accuracy under distribution shift. Score-based model (diffusion model) on latent representations of data distribution is used to estimate the joint data distribution P(X,Y), so that low density regions are upweighted to ensure a more balanced representation during training. The proposed method differs from previous methods since it does not require prior information and it models the joint data distribution without making assumptions on the pre-trained classification boundary of P(Y|X). Experiments on various tabular datasets under distribution shift demonstrate that the proposed method improves worse case classification accuracy as compared to baseline methods.

strengths:
This paper provides a novel method that solves an important task and works well in practice by performing experiments across synthetic data and several real datasets.

weaknesses:
1. Diffusion model is used to model latent covariates X. SimDiff addresses y shift by taking the difference between between two similarity metrics across two different classes. However, the paper claims that it directly estimates the joint data distribution P(X,Y). It would be better to make the connection clearer. 
2. Average similarity measure is proposed based on an intuition derived from a simulation example with mixtures of two Gaussians. Is the intution generalizable to more general settings? If the intuition is not satisfied in all settings, how would it affect the performance of the method?

questions:
See weakness section above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732287489986

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Responses to All Reviewers

comment:
We sincerely thank all reviewers for their thoughtful feedback and for recognizing our paper as well-written (Reviewers BAGe, fTYk), driven by a novel motivation (Reviewers BAGe, fTYk, f1qS), and validated effectively on real-world datasets (All Reviewers). We are also particularly pleased that our examples and illustrations helped most reviewers grasp the key ideas of our work. 

We acknowledge all reviewers' questions. The main requests and concerns can be summarized as follows:

1. Provide more specific details on how the experimental validation was conducted.
2. Investigate whether our score-based method is sensitive to sample size, as traditional density estimation methods often are.
3. Explain how our intuition generalizes to real-world cases.
4. Compare our method with stable learning, another typical reweighting-based approach.
5. Adapt our method to other types of data formats.

We have carefully addressed these concerns in our revised paper. The major revisions are as follows:

- **For (1):** We have added essential details on how sensitive attributes were selected, as outlined in Appendix A.5.
- **For (2):** We conducted new experiments and reviewed past studies, as presented in Appendix A.6. These results demonstrate that our method is inherently robust to sample size variations, a property stemming from score-based modeling.
- **For (3):** We provided both quantitative analyses (Appendix A.7) and visual illustrations (Appendix A.8) to explain how our method generalizes to real-world scenarios.
- **For (4):** We highlighted three main advantages of our approach over stable learning-based methods in Appendix A.9. Additionally, we applied a typical stable learning method SRDO to all datasets and updated the corresponding experimental results in the revised paper.
- **For (5):** We conducted additional experiments on an image dataset and shared the results with corresponding reviewers. Our method outperformed all baselines in these tests. However, it is important to emphasize that the focus of this paper is on tabular data, which represents a significant and practical domain of application. Thus, we will reserve the discussion of results on other data types for future work.

To facilitate review, we have color-coded all revisions in the updated manuscript. If any reviewer has additional concerns or suggestions, please feel free to reach out—we are always willing to engage in further discussion.

Best regards,

The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732588479430

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Request for Further Discussions

comment:
We sincerely appreciate the insightful questions and helpful suggestions from all the reviewers. We believe our responses have addressed all the concerns raised so far, but we are more than happy to provide any further clarifications if needed. As the discussion period has already passed the halfway point, we kindly invite reviewers to share any additional questions or suggestions. We deeply value this opportunity for dialogue and are eager to engage further to ensure all concerns are thoroughly addressed.

Best regards,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper describes a new approach to handle distribution shift by reweighing the framework. The proposed approach is motivated by the need to mitigate biases and spurious correlations in machine learning models on tabular data. The method identifies and upweights underrepresented samples based on the similarity of score vectors to ensure an uniform dataset representation. The approach has been shown to effectively improve robustness over tabular datasets.

### Recommendation

Reviewers engaged with this submission and identified strengths in its topic and approach. Following the rebuttal and discussion, the overall sentiment among reviewers was "weakly" positive. Given the borderline evaluation, I also reviewed the submission as well -- considering the original submission, the supplementary materials, the reviews, and the responses. 

Overall, my recommendation is to reject the paper at this time. My recommendation is primarily based on a number of weaknesses that affect clarity, significance, and potentially soundness. Some of these issues could be addressed through a round of revisions but should be validated through an additional round of peer review. 

### Weakness 1: Soundness/Significance/Limitations 

The proposed approach is primarily validated through experiments on real-world data where the authors examine accuracy. The paper demonstrates gains. However, it does not sufficiently explore the underlying drivers or the required conditions. This is concerning in several respects.

1. Significance - i.e., readers who wish to adopt this method may not be convinced, especially given that other approaches rest on firm theoretical foundations

2. Limitations - i.e., practitioners may adopt the proposed method without understanding the potential risk. This seems concerning as the effectiveness arises under causal assumptions which can be misspecified.

3. Validity - i.e., in the worst case, we may be observing an artifact. This is a minor point but one that is entirely possible given that we are examining summary statistics on a handful of experiments.

One way to address this would be to perform a synthetic study where we compare the method to alternatives on a simple 2D problem where we observe ground truth. This study can build intuition and highlight potential limitations. Stepping back, reweighing is practical but has several limitations (e.g., potential for misspecification,  and impact on other metrics).

### Weakness 2: Claims/Clarity

The discussion highlights several key points of confusion among reviewers. Many of these issues pertain to the central claims of the paper – e.g., which kinds of distribution shifts does this handle, why would other methods not work, why is this limited to tabular datasets? Some of these issues were resolved through discussion, but should be integrated into an updated paper. Overall, the claims should be measured and supported through evidence.