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

W2: 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.

[R2] Thanks for raising this question. We will address this issue step by step.

First, we want to stress that training a score model could be exactly viewed as training a density estimator just like traditional density estimation method. In fact, given a trained score model, the exact probability density could be estimated solely based on score. The detailed explanation is stated in Section 3.3.1. To be specific, score itself, i.e., $\nabla_{\mathbf{x}_t} \text{log}p(\mathbf{x}_t)$ denotes the gradient of the log-likelihood of the original data distribution $p(z)$ at a given noise scale $t$. The diffusion process is modeled as a stochastic differential equation (SDE). As introduced in [1], the corresponding probability flow ordinary differential equation (ODE) is given by: $\mathrm{d}\mathbf{x} = \underbrace{\left[\mathbf{f}(\mathbf{x}, t) - \frac{1}{2}g(t)^2 \mathbf{s} _{\theta}(\mathbf{x}, t) \right]} _{=: \tilde{\mathbf{f}} _\mathbf{\theta}(\mathbf{x}, t)} \mathrm{d}\bar{t}$. By solving this ODE and using the instantaneous change of variables formula, the log-likelihood of the original data distribution could be computed as $\log p _0(\mathbf{x}(0)) = \log p _T(\mathbf{x}(T)) + \int _0^T \nabla \cdot \tilde{\mathbf{f}} _\mathbf{\theta}(\mathbf{x}(t), t) \mathrm{d} t$, which has the same effect as traditional statistical density estimation method, while leveraging the flexibility and efficiency of score-based modeling.

Second, we want to clarify that our proxy follows the same trend as the probability likelihood computed via score model's ODE, which is exactly the estimated density. However, our proxy actually offers numerically more stable sample-level weights compared to the raw likelihood of the probability density, which makes our approach more suitable for sample-level classification tasks. This is the key question we have discussed in Section 3.3.1 and Section 4.6. Note that our final aim is to balance the datasets globally at the sample level. It requires that the weights of those samples should be within a numerically stable range. While samples can have varying importances, each valid sample must still have an opportunity to contribute to the final classification model. However, the probability density fails to meet this requirement. Extreme differences in log-likelihood values lead to highly imbalanced sample weights, failing to reflect the relative magnitudes of densities among different samples globally. Because the log-likelihood strictly follows the functional form of the original data distribution, it leads to explicit numerical differences. When only a few data points have significantly high probability densities, their log-likelihoods become much higher than those of other samples. Using these likelihoods to compute sample weights causes high-density samples to disproportionately overshadow others, whether those have low or moderately high densities. This imbalance diminishes distinctions among samples outside the highest-density regions, reducing the ultimate training on the remaining samples to an unweighted process without meaningful density differentiation.

To validate this point, we conduct the experiment in Section 4.6. In Figure 3a, we could observe that only a small number of points located at cluster centers exhibit significantly higher log-likelihoods. Reweighting based on these estimates could lead to disproportionate emphasis on a small number of extreme high-density points, potentially obscuring the distinction between other high- and low-probability regions. In contrast, our score similarity in Figure 3b alleviates this problem. Our similarity measure not only differentiates density clearly but also maintains a consistent and stable numerical range, making it suitable for sample reweighting.

The key difference between our proxy and exact likelihood lies in the interaction between models as defined in Eq. 12. The subtraction operation between outputs of two score models captures the relative difficulty of training models on specific samples. Exact likelihood, on the other hand, reflects only the absolute probability, making it well-suited for generative tasks but less appropriate for sample-level classification tasks. The numerical values of probability density have a scale mismatch with the requirements of sample-level weights for classification, whereas our proxy operates on the same scale as the score model outputs, ensuring consistency and stability in sample reweighting.

[1] Song Y, Sohl-Dickstein J, Kingma D P, et al. Score-based generative modeling through stochastic differential equations[J]. arXiv preprint arXiv:2011.13456, 2020.

Q: 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.

[A] Addressing the issue of distribution shift in tabular data is crucial. Previous studies [1] have highlighted that tabular datasets are often subject to diverse types of distribution shifts. Furthermore, given that each column in tabular data represents specific semantics, ensuring robustness and fairness across sensitive attributes is particularly important. These considerations motivate the development of our method that can maintain robustness under potential covariate and label shifts in tabular datasets.

To address your request, we conducted experiments on the ColorMNIST dataset. Following the original setup described in [2], we excluded the random label-flipping process from [2] to preserve the underlying predictive mechanism. Robustness was evaluated using the worst-group accuracy based on the non-causal feature, i.e., the color of the digits. Each experiment was repeated three times, and we report the mean and standard deviation of the classification accuracies. The results are summarized as follows:

The ColorMNIST dataset includes only one non-causal attribute—namely, the color of the digits—which facilitates validation. Our objective is to improve the worst-group accuracy while maintaining the mean accuracy as much as possible, a goal that our method successfully achieves as demonstrated in the table above. This experiment demonstrates that our approach can be adapted to other types of datasets, provided that the non-causal attributes are identifiable and corresponding validations are feasible. By modeling the original data distribution through score-based methods, our approach enhances robustness without relying on prior information. In summary, our method shows potential for generalization to other data domains. However, this paper focuses on addressing challenges in tabular data, and we leave the exploration of broader generalization to future work.

[1] Liu J, Wang T, Cui P, et al. On the need for a language describing distribution shifts: Illustrations on tabular datasets[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2] Arjovsky M, Bottou L, Gulrajani I, et al. Invariant risk minimization[J]. arXiv preprint arXiv:1907.02893, 2019.

W3: 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.

[R3] Thanks for proposing this concern. Our method use score-based model to model the original data distribution, whose aim is similar to density estimation methods like Kernel Density Estimation (KDE). Traditional density estimation methods, such as KDE, often require a large amount of background data. However, thanks to the advantages of diffusion models, our score-based density estimation is not sensitive to sample size, as demonstrated in our subsequent experimental results.

Here we investigate the influence of training samples' size on our score-based proxy. Subsets of varying sizes were selected from the original training dataset, and score models were trained on these subsets. The final classification model was then trained using the weights generated from these score models but tested on the original test data. We denote $R$ as the ratio of the subset size relative to the original dataset. The experimental results are listed in the following table. It demonstrates that the weights from our score-based proxy are largely insensitive to the sample size. The robustness of our method arises from score model's ability to construct the latent score field. As stated in [1], neural networks can memorize the score field even when the number of training samples is finite. This property ensures that our score-based weights faithfully approximate the original probability density, enabling robust classification model training even with reduced subsets ($R=0.5$ or $R=0.8$). However, when $R=0.2$, the number of training samples becomes insufficient to construct an accurate score field, which makes the estimated weights less accurate. However, even when we used only one-fifth of the data for training, the performance of our model still outperformed the results of the other baselines on the full dataset in Table 1 of our paper.

In summary, our method could generate more effective weights due to the process of modeling implicit score field, which only requires training a neural network. The process of predicting score for new samples does not require to calculate the interaction with other given training samples like statistical density estimation methods. Therefore, the number of training samples has minimal impact on this process, provided the dataset size remains relatively reasonable.

[1] Kadkhodaie Z, Guth F, Simoncelli E P, et al. Generalization in diffusion models arises from geometry-adaptive harmonic representations[C]//The Twelfth International Conference on Learning Representations.

Due to word count limitations, we are unable to include the complete table in our "Response to Reviewer 4w6D (1.1/4)". This response provides the experimental results referred in R1.

[1] Shen Z, Cui P, Zhang T, et al. Stable learning via sample reweighting[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2020, 34(04): 5692-5699.

[2] Kuang K, Cui P, Athey S, et al. Stable prediction across unknown environments[C]//proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining. 2018: 1617-1626.

[3] Zhang X, Cui P, Xu R, et al. Deep stable learning for out-of-distribution generalization[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021: 5372-5382.

We would like to thank you for taking the time to review our paper and for providing insightful comments. Below are responses to the concerns raised by you. Please let us know if you require any further information, or if anything is unclear.

W1: 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.

[R1] Thanks for proposing this concern. The idea of stable learning-based method is to make features decorrelated to achieve a balanced data distribution in the sample level, which is similar to our method. However, our method offers three significant advantages over stable learning, as detailed below:

1. Ability to handle potential $Y$-shift problems: Our method use a similarity difference measure to address potential $Y$-shift as discussed in Section 3.3.2. In contrast, stable learning focuses solely on the decorrelation of covariates without differentiating the information carried by labels $y$.

2. Balanced Distribution via Original Distribution Modeling vs. Feature Independence: Our method obtains a balanced distribution by modeling the original distribution while stable learning achieves balance through feature independence. The reweighting process in stable learning relies on feature decorrelation under a linear assumption. However, it is important to note that a balanced distribution does not equate to feature independence, and feature decorrelation does not necessarily achieve balance, particularly in non-linear cases.

   Consider an example where feature decorrelation fails at the sample level but our method succeeds. Suppose there are two features $x_0$ and $x_1$, with samples distributed such that $0 < x_0 < 1$ and $x_0 < x_1 < x_0+1$. Suppose that original data distribution is imbalanced, e.g., samples with $x_0 > 0.5$ all share a same higher density than those with $x_0 < 0.5$. Under this circumstance, our method can transform the original data distribution $p(x_0, x_1)$ into a uniform data distribution $U(x_0, x_1)$ at the sample level easily. This is achievable because the score model captures and recovers the original data distribution $p(x_0, x_1)$.

   In contrast, stable learning, which enforces feature decorrelation to achieve independence, cannot produce a uniform balanced distribution at the sample level while maintaining independence between $x_0$ and $x_1$. The geometry of the sample space—a parallelogram region—dictates that $x_0$ and $x_1$ cannot be independent while maintaining a uniform sample-level distribution, i.e., $U(x_0, x_1) \neq p(x_0) \cap p(x_1)$. Thus, stable learning fails to balance the data distribution at the sample level while simultaneously decorrelating features.

   Decorrelating features to achieve feature independence is a good idea, but it is insufficient to ensure a balanced dataset without the essential assumptions. In contrast, our method does not depend on the presence or absence of feature correlations. Instead, it estimates the implicit score field to model the original distribution, enabling a robust balancing operation based on the estimated distribution.

3. Enhanced Experimental Outcomes Relative to Stable Learning: To facilitate a quantitative comparison between our method and stable learning, we conducted supplementary experiments and updated the results in Tables 1 and 2 of our revised manuscript. Our approach shows a minimum improvement of 5% in the worst group accuracy compared to the SRDO baseline across both tables. For clarity, SRDO’s performance is highlighted in blue in the revised paper. The superiority of our score-based reweighting over stable learning stems from the accuracy of the weights derived from our score model compared to the predicted probabilities generated by stable learning predictors. Stable learning methods typically train a predictor to estimate probabilities, which involves the generation of synthetic samples. For instance, SRDO [1] employs empty vectors to create samples, whereas StableNet [3] uses a Random Fourier Transformation. The quality of these synthetic samples crucially affects the training process of the predictor, thereby making the estimated probabilities potentially unreliable. In contrast, the score in our method represents the gradient of the log-likelihood of the original distribution $p(x)$. This ensures that our reweighting process remains rigorously faithful to the actual data distribution.

In a word, our method does not rely on any assumptions about the original data distribution. It utilizes the score, i.e., the gradient of estimated log-likelihood of original data distribution, to perform sample reweighting, which is flexible and easy to conduct. The additional experiments further confirm the effectiveness of our method.

We would like to thank you for taking the time to review our paper and for providing comments. Below are responses to the concerns raised by you. Please let us know if you require any further information, or if anything is unclear.

Q1: 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?

[A1] As noted in your comments, addressing the issue of distribution shift in tabular data is crucial. Previous studies [1] have highlighted that tabular datasets are often subject to diverse types of distribution shifts. Furthermore, given that each column in tabular data represents specific semantics, ensuring robustness and fairness across sensitive attributes is particularly important. These considerations motivate the development of our method that can maintain robustness under potential covariate and label shifts in tabular datasets.

To address your request, we conducted experiments on the ColorMNIST dataset. Following the original setup described in [2], we excluded the random label-flipping process from [2] to preserve the underlying predictive mechanism. Robustness was evaluated using the worst-group accuracy based on the non-causal feature, i.e., the color of the digits. Each experiment was repeated three times, and we report the mean and standard deviation of the classification accuracies. The results are summarized as follows:

The ColorMNIST dataset includes only one non-causal attribute—namely, the color of the digits—which facilitates validation. Our objective is to improve the worst-group accuracy while maintaining the mean accuracy as much as possible, a goal that our method successfully achieves as demonstrated in the table above. This experiment demonstrates that our approach can be adapted to other types of datasets, provided that the non-causal attributes are identifiable and corresponding validations are feasible. By modeling the original data distribution through score-based methods, our approach enhances robustness without relying on prior information. In summary, our method shows potential for generalization to other data domains. However, this paper focuses on addressing challenges in tabular data, and we leave the exploration of broader generalization to future work.

[1] Liu J, Wang T, Cui P, et al. On the need for a language describing distribution shifts: Illustrations on tabular datasets[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2] Arjovsky M, Bottou L, Gulrajani I, et al. Invariant risk minimization[J]. arXiv preprint arXiv:1907.02893, 2019.

Q5: The experimental results 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

[A5] For the additional experiments, please refer to A1. For comparisons with related baselines, particularly stable learning, please see A4. We want to emphasize that our primary goal is to enhance overall robustness, as measured by the average of worst-group accuracies, while maintaining the mean accuracy as much as possible. Notably, our method improves the worst-group accuracy by nearly 3% while maintaining minimal impact on the mean accuracy—a remarkable achievement compared to baseline methods.

Due to word count limitations, we are unable to include the complete table in our "Response to Reviewer f1qS (4.1/5)". This response provides the experimental results referred in A4.

[1] Shen Z, Cui P, Zhang T, et al. Stable learning via sample reweighting[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2020, 34(04): 5692-5699.

[2] Kuang K, Cui P, Athey S, et al. Stable prediction across unknown environments[C]//proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining. 2018: 1617-1626.

[3] Zhang X, Cui P, Xu R, et al. Deep stable learning for out-of-distribution generalization[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021: 5372-5382.

Q4: If the model could only solve the P(X) problem, the stable learning methods could also handle this problem. What is your major advantange?

[A4] First, we would like to clarify that our method addresses not only $X$-shift but also potential $Y$-shift, as detailed in A2 and mentioned in Section 3.3.2 of the main text.

Second, since you referred to stable learning-based methods [1, 2, 3], we provide a comprehensive empirical comparison between our method and stable learning along with analysis experiments. Stable learning aims to decorrelate features to achieve a uniform and balanced data distribution, which is similar to our approach. However, our method offers three significant advantages over stable learning, as detailed below:

1. Ability to handle potential $Y$-shift problems: Our method use a similarity difference measure to address implicit $Y$-shift as discussed in A2. In contrast, stable learning focuses solely on the decorrelation of covariates without differentiating the information carried by labels $y$.

2. Balanced Distribution via Original Distribution Modeling vs. Feature Independence: Our method obtains a balanced distribution by modeling the original distribution while stable learning achieves balance through feature independence. The reweighting process in stable learning relies on feature decorrelation under a linear assumption. However, it is important to note that a balanced distribution does not equate to feature independence, and feature decorrelation does not necessarily achieve balance, particularly in non-linear cases.

   Consider an example where feature decorrelation fails at the sample level but our method succeeds. Suppose there are two features $x_0$ and $x_1$, with samples distributed such that $0 < x_0 < 1$ and $x_0 < x_1 < x_0+1$. Suppose that original data distribution is imbalanced, e.g., samples with $x_0 > 0.5$ all share a same higher density than those with $x_0 < 0.5$. Under this circumstance, our method can transform the original data distribution $p(x_0, x_1)$ into a uniform data distribution $U(x_0, x_1)$ at the sample level easily. This is achievable because the score model captures and recovers the original data distribution $p(x_0, x_1)$.

   In contrast, stable learning, which enforces feature decorrelation to achieve independence, cannot produce a uniform balanced distribution at the sample level while maintaining independence between $x_0$ and $x_1$. The geometry of the sample space—a parallelogram region—dictates that $x_0$ and $x_1$ cannot be independent while maintaining a uniform sample-level distribution, i.e., $U(x_0, x_1) \neq p(x_0) \cap p(x_1)$. Thus, stable learning fails to balance the data distribution at the sample level while simultaneously decorrelating features.

   Decorrelating features to achieve feature independence is a good idea, but it is insufficient to ensure a balanced dataset without the essential assumptions. In contrast, our method does not depend on the presence or absence of feature correlations. Instead, it estimates the implicit score field to model the original distribution, enabling a robust balancing operation based on the estimated distribution.

3. Enhanced Experimental Outcomes Relative to Stable Learning: To facilitate a quantitative comparison between our method and stable learning, we conducted supplementary experiments and updated the results in Tables 1 and 2 of our revised manuscript. Our approach shows a minimum improvement of 5% in the worst group accuracy compared to the SRDO baseline across both tables. For clarity, SRDO’s performance is highlighted in blue in the revised paper. The superiority of our score-based reweighting over stable learning stems from the accuracy of the weights derived from our score model compared to the predicted probabilities generated by stable learning predictors. Stable learning methods typically train a predictor to estimate probabilities, which involves the generation of synthetic samples. For instance, SRDO [1] employs empty vectors to create samples, whereas StableNet [3] uses a Random Fourier Transformation. The quality of these synthetic samples crucially affects the training process of the predictor, thereby making the estimated probabilities potentially unreliable. In contrast, the score in our method represents the gradient of the log-likelihood of the original distribution $p(x)$. This ensures that our reweighting process remains rigorously faithful to the actual data distribution.

In a word, our method does not rely on any assumptions about the original data distribution. It utilizes the score, i.e., the gradient of estimated log-likelihood of original data distribution, to perform sample reweighting, which is flexible and easy to conduct. The additional experiments further confirm the effectiveness of our method. Please view our revised paper or "Response to Reviewer f1qs (4.2/5)" for experimental results.

Q3: 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.

[A3] Thanks for this comment. The "prior information" in this section actually has two distinct meanings:

• Pure Prior Information: This refers to explicit prior knowledge that can enhance the training of predictive models, such as domain labels.

• Predefined Prior Assumption: This represents implicit assumptions about the training data distribution, which underpin the design of the method.

We want to emphasize that while the works you referenced may not explicitly introduce additional covariates, their methodologies are implicitly based on various assumptions about the data distribution. This aligns with the second meaning of "prior information" as defined in our text. Detailed explanations are provided below:

• The optimization process of the second stage in HRM implicitly assumes that samples with the same domain label follow a uniform distribution. HRM [1,2] implicitly assumes that the samples could be clustered only by data's heteogeneity. Therefore, it generates domain labels to divide subgroups. Since domain labels are discrete, samples with the same domain label are treated equivalently during traning final predictive models. Therefore, the underlying assumption in the final optimization process is that the samples within each cluster conform to a uniform distribution.
• Stable learning explicitly assumes feature independence could lead to balanced distribution, which is fulfilled in linear situations. The framework of stable learning methods [2] is inherently built on linear assumptions, targeting specific linear prediction problems. A detailed comparison between stable learning and our method is provided in A4.

In contrast to these methods, our approach neither assumes a specific mathematical form for the original data distribution nor introduces additional covariates. The score-based reweighting mechanism remains entirely faithful to the original given data distribution. We will include these comparisons in Appendix A.9 of our revised paper and provide a detailed explanation of the term "prior information" as requested.

[1] Liu J, Hu Z, Cui P, et al. Heterogeneous risk minimization[C]//International Conference on Machine Learning. PMLR, 2021: 6804-6814.

[2] Shen Z, Cui P, Zhang T, et al. Stable learning via sample reweighting[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2020, 34(04): 5692-5699.

Q2: 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?

[A2] We want to first clarify the statement in this answer to make our following explanations more clear: In our paper, $P(X)$, $P(Y)$, $P(Y|X)$, and $P(X, Y)$ refer to the data distribution. As for the specific types of distribution shift, we commonly use $X$-shift to indicate covariate shift problem, which arises primarily from imbalanced data distributions of the covariates $x$. $Y$-shift refers to the imbalances in the outcome labels $y$.

Returning to your question: Based on your text in Q4, we guess your question as asking whether our method, since the score model is trained solely on covariates, addresses only the $X$-shift problem. The answer is no. While the $X$-shift problem is the primary focus of our method, it is also capable of addressing the potentia $Y$-shift problem. The reasons are as follows:

• The types of distribution shift our method tackles: Our method primarily focuses on addressing $X$-shift at the level of all non-causal covariates, as well as potential $Y$-shift in class labels. The objective is to ensure that the model consistently makes robust predictions across all non-causal covariates rather than optimizing for robustness with respect to a single covariate.
• The reason why boundary-based methods may fail: Our paper claim that previous methods fail to achieve such global robustness because they rely on a pretrained boundary $P(Y|X)$. Here $P(Y|X)$ denotes conditional probability, not the $Y|X$-shift referenced in your text. While under $X$-shift, a balanced training data distribution is required when pursuing a global robust model just like the example in Figure 1. Previous boundary-based methods may fail because they could not model the original training distribution $P(X)$ effectively. In turn, they are unable to balance the training data.
• The detailed explanations about how our method utilize the information of $y$ and handle potential $Y$-shift problem: We guess what you want to ask is that training score model is based on covariates $x$ and how can we model the joint data distribution $P(X, Y)$. Our approach approximates the joint distribution indirectly using the similarity difference measure defined in Eq. 12. As described in Section 3.3.2, we partition the training data by class and train separate score models for each class. The similarity function $\text{Sim}(\cdot)$ in Eq. 11 estimates the probability density within a specific class. For each sample, we compute the similarity difference between score models trained on samples from different classes. The subtraction operation in Eq. 12 ensures that samples from minority classes exhibit lower $\text{SimDiff}$ values compared to those from majority classes, effectively capturing $P(Y)$ indirectly.

In summary, we train score models separately based on the given $y$ labels in the training data and use $\text{SimDiff}$ to approximate $P(X, Y)$. Our similarity difference measure allows us to represent probability density at the sample level. Consequently, our method effectively addresses both $X$-shift at the level of all non-causal covariates and potential $Y$-shift in class labels. If you have any further questions, please do not hesitate to reach out. We are committed to addressing all your concerns.

Dear Reviewer f1qS,

Thanks for your valuable efforts in reviewing our paper. Based on your questions, we have revised the paper in the following aspects:

• For W1, we have deployed our method on image data and provided the detailed results in A1.
• For W2, we have clarified the exact types of distribution shifts our method can address in A2.
• For W3, we have explained the meaning of "prior information" in our setting and revised Section 2 for better understanding.
• For W4, we have illustrated our method's advantages over stable learning both analytically and experimentally in Appendix A.9.
• For W5, we have added an experimental comparison with the proposed baseline and updated Table 1 and 2 accordingly.

We sincerely hope that our previous responses could adequately address your concerns. With the deadline for authors to upload the revised PDF approaching, we would like to kindly inquire if you have further concerns or suggestions. We greatly value this opportunity for dialogue and are always more than happy to engage in further discussions. Provided that your concerns have been well-addressed, we would greatly appreciate it if you could reevaluate our paper and thus consider upgrading your score.

Best regards,

Authors

Dear Reviewer f1qS,

We would like to express our sincere thanks once again for your efforts in reviewing our paper. As there are only two days left for discussion, we would kindly inquire if you have any further suggestions or concerns. We believe we have addressed all the remaining concerns raised by the other reviewers, and we hope that our previous responses have adequatetly addressed your concerns as well. Should you have any further comments or queries, we would greatly appreciate the opportunity to address them promptly.

Best regards,

Authors

Dear Reviewer f1qS,

We sincerely appreciate your efforts in reviewing our paper. As the Author-Reviewer discussion period is drawing to a close, we would like to kindly inquire if you have any additional suggestions. We are more than willing to engage in further discussions to address any potential concerns.

Best regards,

Authors

Rebuttal-Q2-2: Moreover, I have concerns on why a uniform training distribution might offer better generalization performance. A uniform distribution can cause the dataset to lose its original information. Further justification is needed regarding the correctness of towarding a "uniform distribution". This is a more critical question.

[Rebuttal-R2-2] Thanks for raising this concern. We would like to address it point by point.

The meaning of "uniform" in our paper.

[Rebuttal-R2-2-1]

• The "uniform" in our paper denotes sample-level balance.

[Rebuttal-R2-2-1-1] We want to first clarify that the "uniform" in our paper refers to the sample level, which aims to achieve a balanced joint distribution $P(X)$. Please note that this is different from the "uniform" at the feature level, which requires every feature to fulfill a uniform marginal distribution $U(x_i)$. We want to emphasize that these two concepts are fundamentally different. Sample-level uniformity is critical for ensuring robustness, and we will explain the rationale in Rebuttal-R2-2-2. In contrast, a uniform distribution at the feature level would indeed cause the dataset to lose its original significant information for generalization and we will discuss it in Rebuttal-R2-2-3.

• A balanced, sample-level-uniformly-distributed dataset is not the same as a dataset where all features follow a uniform marginal distribution.

[Rebuttal-R2-2-1-2] We would like to point out that a reweighted dataset whose joint distribution is balanced does not necessarily imply that all its features will follow a uniform distribution separately. Moreover, uniform-distributed marginal distribution does not always indicate a balanced joint distribution. Here we will provide an example, please refer to this link for a detailed illustration (try downloading the PDF if the webpage can't exhibit correctly). Here we have a dataset which contains six samples in the left figure, and four of them are totally the same. If we aim for a balanced dataset at the sample level, we can just enhance the weights for those two minority sample as four times large as majority samples just like the situation in the right figure. Note that in this case, the dataset becomes balanced through reweighting, while $p(x_0)$ obviously does not fulfill a uniform distribution.

• Achieving a balanced, sample-level-uniformly-distributed dataset does not require feature independence, and in fact, feature independence may not promote a balanced dataset.

[Rebuttal-R2-2-1-3] We want to clarify that deploying reweighting operation to achieve a balanced dataset does not require feature independence. In addition, some operations, such as decorrelating features and promoting feature independence just like stable learning-based methods you have referred in Q4, could not assure a balanced dataset. Here we would like to detail the example where feature decorrelation fails at the sample level but our method succeeds. The corresponding illustration is in this link. Suppose there are two features $x _0$ and $x _1$, with samples distributed such that $0<x _0<1$ and $x _0<x _1<x _0+1$. Suppose that original data is imbalanced, e.g., samples with $x _0>0.5$ all share a same higher density than those with $x _0<0.5$. Under this circumstance, our method can transform the original data distribution $p(x _0,x _1)$ into a uniform joint distribution $U(x _0,x _1)$ at the sample level easily. This is achievable because the score model captures and recovers the original data distribution $p(x _0,x _1)$. In contrast, the geometry of the sample space—a parallelogram region—dictates that $x _0$ and $x _1$ cannot be independent while maintaining a uniform sample-level distribution, i.e., $U(x _0,x _1) \neq p(x _0)\cap p(x _1)$. In addition, if $x _0$ and $x _1$ both fulfills a uniform distribution, the joint distribution $p(x _0,x _1)$ will not be balanced, which also confirms our Rebuttal-R2-2-1-2.

In summary, our "uniformity" refers to the balanced joint distribution, which is different from the feature-level "uniformity". Therefore, the drawbacks brought by feature-level uniformity do not match with our method, which will be detailed in Rebuttal-R2-2-3. Understanding the distinction between the two meanings of "uniform" in this context is critical. We will include a discussion of the different interpretations of "uniform" in the final version of our paper.

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.

[A2] This is an excellent concern. As you mentioned, selecting non-causal attributes is crucial for evaluating robustness globally. To address this, we present the Pearson correlation coefficients between the attributes and the target variable in both the training and test datasets in the following table.

The Pearson correlation coefficients between our selected sensitive attributes (shown in regular font) and the target variable in both training and test datasets are presented in the following table. For comparison, we also include some typical attributes (shown in italic font) that were not selected for evaluation. Sensitive attributes are chosen for evaluation if they meet at least one of the following conditions:

1. Weak Linear Correlation with the Target Variable: The selected covariates exhibit weak correlations with the target variable. For instance, in the Adult dataset, our selected attributes—marital status, race, and sex—show relatively smaller Pearson coefficients compared to attributes such as native country and workclass.
2. Differences in Correlation Statistics Between Training and Test Datasets: Selected attributes exhibit varying correlation coefficients across training and test datasets, suggesting that these attributes are not direct causes of the target variable. For example, in the Taxi dataset, the selected attributes demonstrate significant differences in correlation coefficients between training and test datasets, reinforcing their selection as non-causal in this context.

These details have been added to Appendix A.5 of our revised paper.

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.

[R3] This is an excellent question. As requested, we have added t-SNE visualizations in a new Section A.8. Please refer to Figure 5 in our revised paper. Specifically, we provide a straightforward visualization on the Default dataset to demonstrate that our method effectively balances real-world datasets.

In the last table in R1, the Default dataset primarily exhibits label ($y$) shift. Thus, we divide the representations into a majority group ($y=0$) and a minority group ($y=1$). In Figure 5.a, points from the minority group ($y=1$) are shaded in dark green, while the majority group ($y=0$) is colored brown. We observe that minority samples are often situated in the marginal regions of clusters, indicating that they have low probability densities. Correspondingly, we visualize the score-based weights of these samples in Figure 5.b, with points assigned higher weights represented by warmer colors (e.g., purple).

We expect the minority group points in Figure 5.a to be assigned higher weights in Figure 5.b. Comparing the two figures, we observe that our method significantly increases the weights for samples in the minority group. Almost all dark green points in Figure 5.a are shaded closer to purple in Figure 5.b, validating the effectiveness of our score-based weighting approach.

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.

[A1] Thanks for proposing this question and your observation is very insightful! As discussed in Section 4.2, our primary evaluation metric is the worst-group accuracy, as our goal is to achieve robust predictions across all non-causal covariates. It is widely acknowledged in the research community that there is often a trade-off between worst-group accuracy and mean accuracy. Including mean accuracy in our tables serves to demonstrate that our method does not achieve higher worst-group accuracy at the cost of a substantial drop in mean accuracy. As shown, our method achieves the best worst-group accuracy across nearly all datasets, though it does not always achieve the highest mean accuracy. This explains why, in Table 1, some mean accuracy results for our method are not the best.

To address your question regarding the mean accuracy results in Table 2, we specifically highlight the ACS dataset because it demonstrates strong selection bias, as identified in the WhyShift study [1]. This study revealed that the ACS datasets exhibit differing causal mechanisms based on spatial location, meaning the predictive relationship between features and labels is not consistent across states. In this case, the instability in mean accuracy is influenced not only by the trade-off between mean and worst-group accuracy but also by shifts in causal relationships. Therefore, we report this factor in our paper to distinguish the results in columns 11–19 of Table 2 from those in Table 1 and columns 2–10 of Table 2.

[1] Liu J, Wang T, Cui P, et al. On the need for a language describing distribution shifts: Illustrations on tabular datasets[J]. Advances in Neural Information Processing Systems, 2024, 36.

We would like to thank you for providing detailed and helpful comments. In addition, we are pleased that you like the idea of this paper. Please see our response to your concerns:

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.

[R1] Thanks for proposing this question. Our similarity measure is designed to generate sample-level weights which could help balance the datasets with distribution shift. To intuitively show that our score-based weights effectively promote balance and facilitate robust model training, we have added a new section A.7 in our revised paper. The following content provides the detailed explanation.

To verify the effectiveness of our score-based weights, we take a straightforward approach: dividing samples by their sensitive attributes and target labels, then computing the sum of weights for each group. Ignoring the challenges of training models due to group variance, we expect the weights to achieve balance. Specifically, the majority group’s weighted sum is expected to decrease relative to its unweighted sum, while the minority group’s weighted sum should increase.

The results for four groups after applying our score-based weights are presented in the following table. In almost all cases, the weighted sum for the majority group decreases, while the minority group’s weighted sum increases. This confirms that our weights effectively balance the original distribution shift, offering a clear explanation of how our method performs on real-world datasets. These findings suggest that our approach is suitable for addressing imbalances in non-synthetic datasets, and the score-based similarity measure can reliably balance datasets at the sample level. We are willing to provide more experimental results if you have any other concern.

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.

[R2] Thank you for your comments. As you correctly noted, we have reported all standard deviations for our experiments in the initial version of the paper. However, due to page constraints and concerns about visual clarity, we initially presented these tables separately. Combining them into a single table at that time would have resulted in a cluttered layout and excessively small font sizes. However, we have carefully considered your feedback, and revisions are currently underway. We are confident that our next version will address your concerns and fulfill your requirements.

Q3: Is there an intuitive way to extend this approach to multiclass classification? Would it be computationally feasible?

[A3] Thanks for raising the concern. For multiclass classification problems, suppose there are $C$ classes denoted by $Y=\{y_0, y_1, ..., y_{C-1}\}$, we could still use $\text{SimDiff}$ to differentiate densities. Similar to the binary classification scenario, we train $C$ separate score models, one for each class. Then for a data point $z_i$ with label $y_i$, $\text{Sim}(z_i;F_{y=y_i})$ is computed using the score model $F_{y=y_i}$ corresponding to its class.

Next, $\text{Sim}$ is calculated for the other $C-1$ score models, and the mean of these $C-1$ values is used to represent $\text{Sim}(z_i;F_{y \neq y_i})$. As a result, Eq.12 is modified as follows:

$\text{SimDiff}(z_i) = \frac{ \sum_{y' \in Y \setminus \{y_{z_i}\}} \text{Sim}(z_i; F_{y = y'})}{C-1} - \text{Sim}(z_i; F_{y = y_{z_i}})$.

For binary classification problems ($C=2$), this equation naturally reduces to the original Eq. 12 in the main text.

It is important to note that training a score model for a specific label $y_m$ requires only the samples belonging to that class. Therefore, the total computational cost for training the score models scales linearly with the total number of training samples, ensuring no additional computational overhead.

We will update these details in the revised manuscript and leave further discussions for future works.

I appreciate the efforts the authors have put in for clarifying my doubt, the mixed optimization objective is indeed interesting. While this could be a whole research direction in itself, the interplay between mean and worst case accuracy is an interesting discussion and going beyond what the authors presented seems a bit out of scope.

I'm happy to recommend acceptance for the paper :)

P.S. Have updated the rating to accept accordingly.

We would like to thank you for providing helpful comments and positive feedbacks. Following are point to point response to address your concern.

Q1: 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.

[A1] Thanks for proposing this concern. We would like to clarify that our method does not directly model $P(X, Y)$. Instead, we approximate it with $\text{SimDiff}$, as referenced in your question. $\text{SimDiff}$ integrates the information from $y$ by differentiating samples with different labels and training separate score models for each $y$. As stated in Section 3.3, we approximate $P(X)$ within each class with the similarity measure $\text{Sim}$. The subtraction operation in Eq. 12 ensures that samples from minority classes exhibit lower $\text{SimDiff}$ values compared to those from majority classes, indirectly capturing $P(Y)$. Therefore, the label information $y$ is incorporated by computing the similarity difference, $\text{SimDiff}$, across different class labels. In summary, while $\text{SimDiff}$ serves as a proxy for the joint distribution $P(X, Y)$, it effectively captures and reflects the relative relationships among labels $y$. As a result, it provides sufficient information to generate sensible weights which help balance the imbalanced dataset.

Q2: 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?

[A2] Thanks for proposing this question. Our similarity measure is designed to generate sample-level weights which could help balance the datasets with distribution shift. To intuitively show that our score-based weights effectively promote balance and facilitate robust model training, we have added a new section A.7 in our revised paper. The following content provides the detailed explanation.

To verify the effectiveness of our score-based weights, we take a straightforward approach: dividing samples by their sensitive attributes and target labels, then computing the sum of weights for each group. Ignoring the challenges of training models due to group variance, we expect the weights to achieve balance. Specifically, the majority group’s weighted sum is expected to decrease relative to its unweighted sum, while the minority group’s weighted sum should increase.

The results for four groups after applying our score-based weights are presented in the following table. In almost all cases, the weighted sum for the majority group decreases, while the minority group’s weighted sum increases. This confirms that our weights effectively balance the original distribution shift, offering a clear explanation of how our method performs on real-world datasets. These findings suggest that our approach is suitable for addressing imbalances in non-synthetic datasets, and the score-based similarity measure can reliably balance datasets at the sample level. We are willing to provide more experimental results if you have any other concern.