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

Thanks for your detailed review and valuable feedback. Below is our concise response:

[A1] The explanation is as follows:

• Your interpretation of our method is correct. We first learn score to model the complex data distribution. A subsequent score-based reweighting ensures balanced representation, improving learning for under-represented regions.
• Regarding "worst group accuracy", this metric assesses performance based on groups defined by a non-causal sensitive attribute. The group with the lowest accuracy is termed the "worst group", reflecting model robustness when the chosen attribute does not causally determine the target variable. For example, when predicting income levels, city of residence could serve as the grouping attribute. It does not directly determine an individual's income, but different cities may correlate with varying income distributions. Therefore, such correlations may induce bias and imbalance, adversely affecting accuracy for certain groups. Specifically, data imbalance significantly impacts model robustness, as samples from low-density regions are naturally under-represented and form the "worst group". Achieving robust model performance requires maintaining high accuracy across all groups. Our method addresses this issue by reweighting samples based on our score proxy, ensuring balanced representation independent of inherent data distribution biases.

In summary, worst group accuracy serves as a measure of model robustness but can be negatively affected by data imbalance. Our proposed method mitigates this issue by ensuring balanced learning across all samples, resulting in enhanced worst group accuracy and robust model performance.

[A2] In our paper, "distribution shift" refers to the differences between imbalanced training data and test data—the latter implicitly treated as balanced under the robustness measure. Please refer to A3 for detailed explanations. Regarding how our dataset presents train-test shifts (i.e., imbalance in training data), please see Table 10, which shows the distribution of samples across different groups.

[A3] This is a valuable insight. Indeed, our method remains beneficial even in scenarios without explicit train-test shifts. The core intuition behind our approach involves modeling data distribution and balancing the original data via our score-based proxy. Train-test shifts do not directly affect either our method or evaluation criteria.

In our paper, we emphasize "distribution shifts" primarily to highlight the distinction between the imbalanced training distribution and the test distribution—the latter being implicitly approximated as balanced under our robustness measure (which computes worst-group accuracies across all sensitive attributes). Specifically, the "shift" we discuss stresses the imbalance present in the training data with respect to either covariates $x$ or target class $y$. To measure model robustness, the ideal evaluation criterion is the sum of worst-group accuracies across all potential non-causal attributes. Under this criterion, an ideal model should perform well in every region of the data space, implicitly corresponding to a balanced test distribution. Therefore, the primary cause of degraded model performance under this testing condition is the imbalance present in training data, creating an implicit train-test shift highlighted in our main text. To mitigate this imbalance, we model the training distribution with our score-based method, thereby effectively enhancing robustness. We will further clarify this point in our final manuscript.

[A4] We deeply appreciate your insightful question. Indeed, our score-based method is compatible with tree-based ensembles. To demonstrate this, we conducted an additional experiment employing three tree-based models to assess whether our sample weights improve models' robustness beyond neural networks. The results are presented at here. For a comprehensive comparison, we also evaluated JTT, a boundary-based method discussed in Section 1, using these tree-based models. Results indicate that our method consistently enhances worst-group accuracy across base models. In contrast, JTT performs well only on certain models. This finding supports our earlier observation of boundary-based methods—they rely solely on decision boundaries, which are disconnected from training distributions, leading to unstable robustness improvements. Conversely, our strategy effectively captures distributional imbalance, generating reliable weights to enhance worst-group accuracy. These results further confirm our method could work well with tree-based models.

We would like to thank you once again for raising these concerns. We believe that these discussions significantly enhance the rigor of our paper. Please feel free to reach out with any questions—we are more than happy to engage further.

We sincerely appreciate your detailed reviews. Below are our responses to your concerns (R) and questions (A):

[R1] Due to readability considerations, we report standard deviations in Appendix A.3. For detailed dataset statistics, please refer to Appendix A.5, which lists the attributes selected for evaluation, and Table 10, which provides sample statistics. In addition, Table 7 clarifies how our hyperparameters affect runtime. For all these suggestions, the revisions are currently underway and we believe our next version will meet your requirements.

[R2] We appreciate your suggestions on the undiscussed studies and will incorporate them in our manuscript.

[A1] The primary role of the autoencoder in our method is to provide latent representations for subsequent score-based modeling. It serves solely as a module for semantic compression and the architecture design does not directly affect subsequent score-based modeling. In our implementation, we follow the practice from TabSyn [1]. Specifically, a tokenizer first converts both numerical and categorical features into an embedding space of dimension $d \times (N _{\text{num}} + \sum _{i=1}^{N _{\text{cat}}}C _i)$, where $N _{\text{num}}, N _{\text{cat}}, C _i$ denote the counts of numerical features, categorical features, and categories within the $i$-th categorical feature, respectively. Then a transformer serves as encoder and decoder, and its output is detokenized to reconstruct the original features. Only the tokenizer and encoder are used for further score-based modeling.

[A2] Sure. We have conducted additional experiments on three datasets from [3] and [4], as you suggested. The results are listed in here. Our method consistently delivers the best worst-group accuracy, clearly enhancing robustness. While mean accuracy sometimes reflects a trade-off, our method often matches or exceeds baseline performance, confirming that our robustness is improved through effective score-based balancing rather than sacrificing overall accuracy. For more details about datasets, please check A5.

[A3] We would like to first clarify that our main contribution is a novel score-based method designed to enhance robustness, rather than providing a new model architecture. Therefore, our approach (and those boundary-based baselines) is complementary to model architectures (e.g., neural networks, tree-based models), making direct comparisons inappropriate.

However, our reweighting strategy can indeed integrate with tree-based models. To address your concern, we conducted extra experiments using 3 tree-based models on 4 datasets. Details are provided in our A4 to Reviewer eedu. These results confirm that our score-based weights could work well with tree models to enhance robustness.

[A4] Nearly all datasets used in our experiments contain categorical features. How we preprocess these features into latent embeddings is detailed in A1.

[A5] Validating performance across diverse datasets is crucial, and our experiments have already been designed with this principle in mind. Given the primary objective of our paper—to enhance robustness under various data imbalances—we selected datasets reflecting different imbalance scenarios. Specifically, datasets in columns 2–9 of Table 1 primarily exhibit class-label $y$ imbalance. To evaluate robustness under covariate $x$ imbalance, we adopted the ACS dataset with state-level validation (columns 2–7 in Table 2). Thus, we believe the selected datasets comprehensively cover key data imbalance challenges.

Regarding your references, we kindly clarify that [3] is the source repository for several datasets we used (e.g., Adult, Bank). :) Given time and resource constraints, evaluating all datasets from [3, 4] is infeasible. However, our original experiment design has tested our method on datasets used in prior studies [1, 2]. Moreover, we also appended extra experiments on 3 new datasets from [3, 4], as reported in A2. Thus, we believe our current validation sufficiently reveals the effectiveness of our method across diverse imbalance scenarios.

We deeply appreciate your suggestions and will incorporate these clarifications into our revised paper. Please feel free to reach out with any new concerns about our method—we are more than happy to engage further.

[1] Zhang, Hengrui, et al. "Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Space." The Twelfth International Conference on Learning Representations.

[2] Liu, Jiashuo, et al. "On the need for a language describing distribution shifts: Illustrations on tabular datasets." Advances in Neural Information Processing Systems 36 (2023).

[3] Asuncion, Arthur, and David Newman. "UCI machine learning repository." Nov. 2007.

[4] Gijsbers, Pieter, et al. "Amlb: an automl benchmark." Journal of Machine Learning Research 25.101 (2024).

We would like to thank you for providing helpful comments and positive feedbacks. Below are our responses to your concerns.

Lack of details about noise preconditioning factor $\sigma$ and hyper-parameter $\tau$

[A1] We appreciate this observation. Regarding the network preconditioning factor $\sigma$, we adopt the methodology proposed in EDM [1] and TabSyn [2] (see Appendix A.2 and Table 4 for details). $\tau$ controls the strength of reweighting based on $\text{SimDiff(·)}$. Its sensitivity analysis is presented in Figure 4 and discussed in Section 4.7.

Lack of discussion for essential references

[A2] We thank the reviewer for highlighting these previous studies. We will incorporate a detailed discussion of these works in the final manuscript.

Revise the captions in Figure 1

[A3] We appreciate this suggestion. Figure 1 displays the results from our synthetic experiment, illustrating how boundary-based methods erroneously upweight samples based on the model's prediction boundary rather than the implicit distribution. In contrast, our score-based proxy effectively models the distribution without requiring prior knowledge, thereby yielding more reliable sample weights for robust classification. We will update the caption in the subsequent version of the manuscript.

We would like to express our sincere thanks once again for your efforts in reviewing our paper. Should you have any further comments or queries, we would greatly appreciate the opportunity to address them promptly.

[1] Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models." Advances in neural information processing systems 35 (2022): 26565-26577.

[2] Zhang, Hengrui, et al. "Mixed-type tabular data synthesis with score-based diffusion in latent space." arXiv preprint arXiv:2310.09656 (2023).