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
Kun Kuang
Venue: ICML 2025 poster (2025)
Keywords: robustness
score model
reweighting
Primary area: general_machine_learning->representation_learning
TL;DR: We propose using score-based reweighting to achieve the overall robustness under distribution shifts.

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.

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Yunze Tong" Fengda Zhang"
Zihao Tang'
Kaifeng Gao Kun Kuang'
Kai Huang?
Pengfei Lyu? Jun Xiao
Abstract
1. Introduction
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 in the training data, such as associations with non- causal features like demographic information. These biases lead to critical robustness issues models may inherit 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 X(Y]X), upweight- ing 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, without relying 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 more uniform dataset representation. Experiments on various tabular datasets under distribution shifts demonstrate that our method effectively improves perfor- mance on imbalanced data. Code is available at https::/gittii com wwomm/yuneeTogg/ latent secntt-score-basee www.rree-baeddd-eeeiitttttt..iia
Machine learning applied tabular data has achieved sig- nificant success in various practical domains (Tang al., 2024a; Yang et al., 2024; 2025b; Wu et al., 2025; Zhu et al., 2025; Lin et al., 2025). While models trained using empiri- cal risk minimization (ERM) often perform well on average, achieving low test errors overall, they can still exhibit high error rates on specific subsets data (Vapnik, 1999; Tang et al., 2024b; Lv al., 2024). This inconsistency highlights critical robustness issue, stemming from inherent biases in the data. For example, training data may exhibit spurious correlations where target labels are statistically linked to non-causal features like demographic information. As a result, models trained on such biased data may inherit or even amplify these biases, leading to poor performance on data subsets where these correlations do not apply.
To tackle this issue, researchers have proposed various meth- ods to mitigate biases in the training data. Some approaches utilize additional prior information, such as group labels, to reduce the impact of spurious correlations by resampling or adding regularization terms (Arjovsky et al., 2019; Sagawa* et al., 2020; Yang et al., 2025a). Unfortunately, in many practical situations, such prior information is incomplete or unavailable, limiting these methods' usefulness. To over- come this, other methods focus automatically generating proxy information for debiasing (Nam et al., 2020; Liu et al., 2021; Qiu et al., 2023). A common strategy involves pre- training a classification model on the training data and then upweighting samples that the model classifies incorrectly, Since these misclassifications are often determined by sam- ple's 's proximity to the model's decision boundary, we refer these boundary-basee based methods (Liu et al., 2021; LaBonte et al., 2023; Ghaznavi et al., 2024). While these methods can enhance robustness by focusing on under-represented training samples near the classification boundary, their effec- tiveness is constrained as they may fail to achieve a globally balanced distribution of training data.
Eqqual contribution Department of Computer Science and Technology, Zhejiang University. Hangzhou, China 3 *Ant Group, Hangzhou, China. Correspondence to: Kun Kuang <kunkuang@ www.edu..n.............
To illustrate the limitations of koundary-based methods, consider binary classification problem with two features, and where the true classification boundary follows sine curve (see Figure la). The training data is heavily
Proceedings of the 4rnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Thus Dosts
Uee Deady
Ler Beatt
(a) Unweighted(b) Weighted data(c) Weighted data by la- training data by JTT
tent score
(a) Data density visu-(b) Score field & ex-(c) Score similarity pected weights
alization
proxy
Figure The score fields of mixture of two Gaussian.
methods by directly targeting data imbalance, aiming to improve performance across diverse data subsets.
(d) Unweighted (c) Weighted (f) Weighted density data density (clas- density with JTT with latent score sifier accuracy at (classifier accu- (classifier accuracy 60.25%)
However, directly using density estimates from the score- based model poses challenges due to potential extreme val- ues, leading to "density explosions" for certain samples. These extremes can overemphasize a few extremely high- density points, making it difficult to distinguish between other high- and low-density samples. To overcome this, we propose using a proxy for density based on an important observation. We demonstrate this observation using data from a mixture of two Gaussians, as shown in Figure 2. Regions with higher probability densities are represented by warmer colors in Figure 2a. To achieve balanced data dis- tribution, these regions are expected to have lower weights, as shown in Figure 2b. In Figure 2c, we illustrate that in high-density regions (e.g., points B and D), the score vec- tors (pink arrows) of neighboring samples tend to align with the direction toward the high-density sample (e.g.. cyan arrows B,B, where i 1,). In contrast, samples in low-density regions (e.g., points A and C') do not exhibit this similarity. Essentially, if sample B has higher density than sample A, the similarity between the score vectors and B, will generally be greater than that for A, This direc- tional similarity serves as proxy for density. By using this proxy, we avoid instability from extreme values, enabling a more stable and effective data reweighting approach.
racy 69.50%) 75.75%)
The comparison of weighted density plots synthetic dataset. JTT upweights samples based on the model's prediction boundary, which is not assured to achieve balance. In contrast, weights could yield more balanced weighted density.
Figure 1.
concentrated in regions where E [ 4 O ('T 7r] resulting in a highly imbalanced dataset, as visualized in Figure 1d. Ideally, reweighting should correct this bias, yielding a more uniformly distributed dataset. Boundary- based methods like JTT (Liu et al., 2021) try to achieve this by assigning higher weights to samples with large classifica- tion errors (Figure 1b). However, many reweighted samples are still from high-density regions near the decision bound- ary, rather than the low-density regions where the imbalance is most severe. The resulting distribution (Figure le) shows that the boundary-basee based method does not fully balance the training data. Consequently, the accuracy improvement on the balanced test set modest, increasing from 60.25% to only 69.50%.
The fundamental limitation of boundary- bbssed methods is their exclusive focus on the pre- pre-trained classification bound- ary, X)) neglecting the overall data distribution. This oversight leads to ineffective data balancing. Motivated by this observation, we aim to develop new approaches that not only consider classification errors but also more effectively address inherent data imbalances, ensuring balanced and uniform distribution across the dataset.
Our method offers two key advantages. First, it requires additional prior information, such as group labels, mak- ing it applicable in scenarios where such information is unavailable incomplete. This flexibility allows broad adoption without relying on external data or assumptions. Second, our approach faithfully represents the joint data dis- tribution P(X, Y). By leveraging score- based models, we overcome the limitations of pre-trained classification bound- aries PYY| which often inherit biases from the training process. This ensures that the final classification model remains unbiased and accurately captures the underlying re- lationships between features and labels, leading to improved performance and robustness, as shown in Figure Ic and If.
To more effectively address inherent data imbalances. we propose leveraging score- score-based models, also known as diffu- sion models (Song & Ermon, 2019; Ho et al., 2020). These models can capture the joint data distribution P(X, Y)., pro- viding powerful tool for modeling the underlying data structure. Our key idea is to use a score-based model to estimate the density of samples in the dataset. By identi- fying low-density regions-areas areas where data is underrep- resented- we can upweight samples from these regions, ensuring a more balanced and unbiased representation dur- ing training. This approach moves beyond boundary-bast
In summary, our contributions are: (1) We recognize that improving robustness requires modeling the joint data distri-
2

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
bution, a gap in current boundary-based methods that do not capture the true data distribution.
shifts, we employ score-based methods faithfully capture the joint distribution. The framework follows three steps: (1) training diffusion models on latent representations to model the data distribution, (2) sampling several timesteps and estimating scores to align the probability density of training data, and (3) reweighting samples based on data density to ensure a balanced distribution. The detailed procedure provided in Algorithm
(2) We introduce new framework that accurately reflects the joint data distribution, detailed in Section 3. (3) We assess our method on vari- ous tabular datasets under distribution shifts, demonstrating through extensive experiments that our approach effectively enhances robustness.
2. Related Works
3.1. Preliminary
Achieving robustness with prior information. Commonly, researchers aim to train robust models towards distribution shift. Some works train models with the help of given or self- generated domain labels (Sagawa* et al., 2020; Arjovsky et al., 2019; Sun & Saenko, 2016; Tong et al., 2023; Zhang et al.. 2024b; Yang et al., 2025a). Generally they expect the model's 's performance in the worst domain acceptable. Therefore, they may add regularizer into training loss when treating domain labels as extra supervised information. Pre- defining the form of prior data distribution is also common approach (Shen et al., 2020; Duchi & Namkoong, 2021; Shen et al., 2023; Gu et al., 2024).
We first introduce the background on score models (Song & Ermon, 2019). Score models can model data distributions by learning score (i.e., the gradients of probability density), and have shown remarkable performance in generative tasks, Song et al. proposed a unified framework based on Itô stochastic differential equations (SDEs). Training a score model typically involves an iterative forward and backward process. In the forward pass, a complex data distribution is gradually transformed into a Gaussian distribution by progressively adding noise, described by the following SDE:
dx = f(x, t) + g(t)dw,
(I)
Achieving robustness without prior information. Apart from introducing the prior information, some other methods turn to use the performance of models to guide the final training process (Liu et al., 2021; Nam et al., 2020; Levy et al., 2020; Qiu et al., 2023; LaBonte et al., 2023; Ghaznavi et al., 2024). They typically contain two stages. In the first stage, they identify some biased samples which will make ERM-bbaed models make false predictions. In the second stage, they assign greater weights to these samples to develop a more robust model. However, notice that few of the following methods model the joint data distribution, which is not fit for seeking overall robustness.
where x E Rdd with Xp ~ Pa representing the data distribu- tion, & E [0, f: Idd [0, + 12 g : [0, - R. and w Rd is standard Wiener process. The backward pro- cess reconstructs the original data structure from noisy data. Song et al. also introduced the corresponding "probability flow" ordinary differential equation (ODE):
dx f(x, () -()x))))000000 px dī,
(2)
where 7 represents time flowing backward from T to O. We use a neural network, sp(x,,,) t), to estimate the score of the transformed data distribution at time t, Vx log px()). The training loss for se t) is defined through denoising mctre-matching:
Generative Models for Robustness. Through different approaches, generative models could be deployed in vari- ous scenarios with specific needs (Zhao et al., 2024; 2025; Hu et al., 2025). In addition, generative models can pro- duce novel and diverse data, enriching training datasets and leading to models with a deeper understanding of semantic content. Consequently, several studies have leveraged gener- ative models to enhance robustness, primarily through data augmentation techniques (Choi et al., 2021; Ilse et al., 2020; Dendorfer et al., 2021; Zhang et al.. 2022; 2024a). These methods have also demonstrated promising performance improvements.
[11so(xe, t) - Vxx log Pejo (x::o(x )10)]
(3)
where (t) represents the time variable distribution, and X(1) is a positive weighting function that stabilizes the time- dependent loss magnitude (Song et al., 2021). The diffusion process generally employs Gaussian transition kernels, lead- ing to Pijo (xx, (o)) (M,,,?1)
In summary, score models aim to compute score at each time scale with different noise, and finally reconstruct the clean sample with the guidance of score (Karras et al., 2022; Xu et al., 2022; 2023).
3. Method
We present an unbiased learning framework that operates without prior information or pre-trained classification bound- aries. The model is designed to provide robust predictions across range of biased covariates. To address distributional
3.2. Training Distribution Modeling
To ensure robustness, we first train score models to approx- imate the probability distribution. To enable training on
3

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
limited computational resources, we adopt the approach of latent diffusion (Rombach et al., 2022; Zhang al., 2024c).
3.3. Probability Density Proxy
At this stage, our goal is to estimate the probability den- sity of each sample using the score models trained in Sec- tion 3.2. We first address the limitations of computing exact log-likelihood via the probability flow ordinary differential equation (ODE) proposed in (Song et al., 2021). Based on the insights from Figure 2, we introduce an approach to estimate the relative probability density through a similarity measure. The difference of score similarities acts as proxy, preserving the relative magnitudes of probability densities among samples and offering controllable numerical range at the model level. This makes it well-suited for computing weights, as described in Section 3.4.
Specifically,, we first train a variational autoencoder (VAE),
Ø2 = PEn DDee, using 3-VAE (Higgins et al., 2017), where the coefficient / balances the reconstruction loss and the KL-divergence loss:
Løz = Erenon (x, x) + Blas-
(4)
We then obtain the latent representations using the trained encoder Fine. The subsequent diffusion process operates on these latent representations, z & (x), rather than the raw data. Leveraging the VAE allows us to model the latent joint distribution of meaningful semantics, rather than superficial features in the raw data.
After the VAE model trained, we use score- based methods to model the underlying distribution p(z) (Song al., 2021; Zhang al., 2024c):
3.3.1. EXACT LIG-LIKELDHOOD COMPUTATION VIA PROBABILITY FLOW ODE
Since our goal is to distinguish data based on high or low probability density, a natural approach is to compute the exact log- likelihodd using the trained score model from Sec- tion 3.2. Song et al. have proposed an estimation method for this. The forward diffusion process is represented by a stochastic differential equation (SDE) that gradually trans- forms a complex data distribution into a known prior distri- bution by injecting noise, as described in Eq. 1. The corre- sponding "probability flow" ordinary differential equation (ODE) is given in Eq. 2. By replacing the score Vx log Pr (x) with neural network (x,, f), the probability flow ODE takes the form:
2e- 20 + (t)e, NO,, I),
(Forward Process)
(5)
dz,- - (t) log p(z)) df
+ (t)
(Reverse Process)
(6)
where 2o is the initial embedding from the encoder, z is the diffused embedding at time t, and (t) is the noise level. In the reverse process (Eq. 6). 2z, p(zz) represents the score function of 2₁- Following Zhang et al., we set the noise scale (t) = f, making it linear with respect to time.
Following the approach of EDM (Karras et al., 2022), we train our neural network Fe to directly predict the output at given timestep rather than scaled unit variance term (4)e. The neural network is preconditioned with o-dependent skip connection, defined as:
dx f(x:t) 0)( sx t)
df.
(9)
Using the instantaneous change of variables formula (Chen et al., 2018), the log-likelihood of po()) can be computed as:
Do(2,0)= = (a)= + Coua(a) Fo) Croise (o)), (7) where (atsp(a) modulates the skip connection between timesteps, (r) and (o) scale the input and output magnitudes, and Cnoise(a) maps the noise level a into conditioning input for Fg. The details of these scaling fac- tors are listed in Appendix B. The final training loss is:
log (x)))(((0)) = log Pr (x(T)) + L.'' V (x(().)) dt,
(10)
where the random variable x(t) a function of is obtained by solving the ODE in Eq. 9.
C(Do, 0) = E,A(()ss=pte)) (2 + (()), ()- - ==1 (8)
However, this estimation method has a limitation: extreme differences in log- log-likelihood values lead to highly imbal- anced sample weights, failing to reflect the relative mag- nitudes of densities among different samples globally. Be- cause the log- ligeliiel strictly follows the functional form of the original data distribution, it leads to explicit numer- ical differences. When only a few data points have sig- nificantly high probability densities, their lig-likelihoods become much higher than those of other samples. Using these likelihoods to compute sample weights causes high- density samples to disproportionately overshadow others,
Here, X(t) is a positive weighting function to maintain the time-dependent loss at a consistent magnitude. The entire process uses the covariates x as input. After obtaining the latent representation z, we partition z by their corresponding label y. For binary classification, two models are trained separately on 2 and v using Eq. 8. This allows us to model the joint training distribution for each class through these score models. We demonstrate the advantages of class- wise data separation and independent score model training in Section 3.3.
4

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
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 clear density differentiation. To validate this, we visualize the probability density using log-likelihood on a synthetic dataset in Section 4.6. In summary, although log-likelihood faithfully reflects the original density, does not meet the requirement for global relative density comparison needed for sample reweighting.
within a specific class, where a larger Sim(-) indicates lower density. The SimDiff(-) measure further accounts for distribution shifts in the label y while preserving the properties of Sim(-) that reflect p(z). When the number of training samples for given class yk is small, Sim(: Fy=ve) increases because the score model Pyy=y. is trained on a narrower and less comprehensive data space. Consequently, Fy=y finds it more challenging to guide noisy points back to their original locations via score fields, leading higher estimation errors in Sim(:;Fy=so ). The subtraction oper- ation thus causes samples from the minority class to have a lower SimDDff(-) compared to those from the majority class. In this way, SimDiff captures both covariate distribu- tion shifts in p(x) (p(z)) via Sim(-) and label distribution shifts in p(y) through the subtraction operation, providing a means to indicate relative probability density.
3.3.2. DATA DENSITY ESTIMATION VIA SIMILARITY DIFFERENCE MEASURE
Because the exact log- ligeliiooi is unsuitable for sample reweighting, we adopt an alternative measure for global density modeling. As discussed in Section we estimate the relative probability density by measuring similarities between the scores of noisy points. According to the for- ward process in Eq. 5, a clean point 2g transforms into a noisy point ze by adding noise scaled by (t). This mirrors the scenario depicted in Figure 2, where point A moves to its noisy neighbor A/-, The vector AA; corresponds to (L)e, and its length |AA|]] represents the magnitude of the added noise [|a(t))ll2. To sample a neighborhood of noisy points around a data point z;, we select T' fixed timesteps to. tu.tt....,ty tr- and add noise as per Eq. 5. Repeating this process K times yields T K noisy points. We then es- timate the relative probability density using these points, employing the squared error as the similarity metric. For sample 2, and preconditioned network Fe. we compute the aggregated similarity, i.e., the average similarity across several noise scales, as:
In summary, lower SimDiff(.) indicates a lower probabil- ity density. Whether the imbalance arises from covariates x or labels SimDiff(-) consistently and faithfully reflects the relative probability density. The whole computation pro- cess requires no prior information or predefined boundaries.
3.4. Unbiased Learning on Distribution Data
Our ultimate goal is achieving the robustness for all sensi- tive covariates x. The implicit joint distribution on latent representation z has been modeled in Section 3.2 and Sec- tion 3.3 through score models. Therefore, we could conduct sample reweighting guided by Eq. 12 directly to acquire an overall unbiased training distribution. Each z; is assigned a weight wi through
exp( (S((-iDif((=:)/))
(13)
Σ N-1 exp( - Si)D)ff(ss)/))
Sim(z;; Po)
1 KL
E, (() - =1]]
where N is the number of all training samples and t de- notes a temperature which controls the scale of reweighting. Finally, we train an unbiased classification model &b as:
(11) The key difference between Eq. 8 and Eq. 11 is that Sim(z;;F) samples only from specific fixed and sparse timesteps, accounting for different noise scales. Moreover, we randomly sample these timesteps multiple times and average the results to obtain a robust similarity measure.
Lelassification
B( [w; v:)),
(14)
where l stands for cross entropy loss. When testing, we only use Eene and to to make predictions.
Using Eq. 11, we compute the aggregated similarity across all training data points to indicate their densities. How- ever, in addition to shifts in latent covariates x, y-shift also commonly occurs in real-world data. To address this, we model the data distribution for each class separately (see Section 3.2) and use the difference in aggregated similarities as a proxy for the final relative probability density. For ex- ample, for data point 2 with class label ye,, we compute (i)) as:
4. Experiment
To evaluate the effectiveness of our method, we conduct extensive experiments in various settings.
4.1. Datasets
To comprehensively validate our method, we selected six diverse datasets that exhibit various types of distribution shifts, including covariate shifts and concept shifts.
Si)Diff(z)) = Sim(z; F - Sim(z;; ) (12)
Sim(-) allows us to differentiate the densities of samples
The details of our used datasets are as follows:
5

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Adult (Becker& & Kohavi, 1996): The goal is to predict whether income exceeds $50K/yr based census data.
two components: Variational AutoEncoder (VAE) for gen- erating latent representations, and MultiLayer Perceptron (MLP) for classification. Both our method and the baseline models utilize the same architecture. For training the VAE, we followed the default setrings described in TabSyn (Zhang et al., 2024c).
Bank (Moro et al., 2012): The data is related with direct marketing campaigns of Portuguese banking institu- tion. The task is to predict if the client will subscribe a term deposit.
In addition, our method involves training score models. Fol- lowing the guidance from TabSyn, we use a 4-layer MLP architecture as the backbone for our score models. We also adopt the training paradigm from EDM (Karras et al., 2022), which mitigates the influence of varying noise scales on neural network training. The noise scale is set to increase linearly with time, i.e., a(t) = t, where follows a log- normal distribution, resulting in In (a)()) ~ (Prmann: Pala). The I'mman and Pad are set 1.2 and 1.2, respectively, When computing aggregated similarity in Eq. 11, we choose 0.002 and 80 as the minimum and maximum values of / and select T timesteps with equal intervals between them. Regarding hyperparameter selection, we set T to 10 and т to 3 for all experiments.
Default (Yeh, 2016): This dataset aims at the case of customers' default payments in Taiwan. The task to predict if the client will default payment next month.
Shoppers (Sakar & Kastro, 2018): The task is to predict if the user's session ends with the shopping behavior. The distribution shift mainly demonstrates as y-shift: About 84.5% samples did not end with shopping, and the rest were positive class samples.
Taxi (Navas, 2018): This dataset collects some informa- tion about the pickup and dropoff of taxi rides. The task is to predict whether the total ride duration time exceeds 30 minutes. We choose the data collected in Mexican City.
4.4. Compared Baselines
US-Wide ACS PUMS Data (Ding et al., 2021): This large dataset contains individual records from US Cen- sus sources. We choose the task where the outcome is whether an individual's income exceeds 50k. The perfor- mance is validated separately on three randomly chosen states.
As referred in Section 4.2, our goal is to enhance overall robustness, with the accuracy on the worst group serving as our main evaluation metric. Therefore, we compare our method with several robust machine learning techniques, including ERM (Vapnik, 1999), CVaR-DRO (Levy et al., 2020), x -DRO (Levy et al., 2020), KL-DRO (Duchi & Namkoong, 2021), JTT (Liu et al., 2021), EIII. (Creager et al., 2021), FAM (Petzka et al., 2021; Zou et al., 2024) and SRDO (Shen et al., 2020). Among these, CVaR- DDO, KL-DRO, JTT, and SRDO employ reweighting processes similar to ours, while EIIL and FAM enhance robustness by adding regularization terms.
Specifically, for the Taxi and ACS datasets, we followed the preprocessing guidelines outlined in the WhyShift benchmark (Liu et al., 2024).
4.2. Evaluation Metrics
Our objective is ensure that the model consistently makes robust predictions across all non-causal covariates, rather than focusing on a single one. For each dataset, we select several deterministic non-causal attributes based on prior knowledge- for example, sex and race in the ACS income dataset. We then record the worst-case prediction results for each feature and compute the average of the worst-group accuracies across these features. Details of the selected features are provided in Appendix B.2. Additionally, we track the mean accuracy for these features. We aim for our model to enhance the average worst-group performance without significantly compromising mean accuracy.
4.5. Classification Results
4.5.1. ROBUSTNESS ACROSS VARIOUS DISTRIBUTION SHIFTS
We first demonstrate the effectiveness of our method under various types of distribution shifts in Table 1. Our method consistently achieves the highest worst-case accuracy across all five datasets, with an average improvement of at least 3% in worst-case accuracy over all baseline methods. Addi- tionally, it attains the best mean accuracy on two of the five datasets, highlighting its ability to enhance robustness with- out significant reduction in overall prediction accuracy.
4.3. Implementation Details
In contrast, DRO-based methods like KL-DRO perform inconsistently on our datasets. For instance, while KI.-DRO achieves the second-best accuracy on the Shoppers dataset, its performance deteriorates significantly on Default and Taxi. This inconsistency stems from KL-DRO's s reliance on
We conducted each experiment three times with different random seeds and report the mean results in Table and Table 2. The corresponding standard deviations are pro- vided in Appendix C. Our prediction model consists of
6

---

Latent Ccooe-BBased Reweighting for Robust Classification on Imbalanced Tabular Data
Table 1. The classification results five datasets which exhibit various kinds of distribution shifts. The bold and underline denote the best and the second best results respectively.
Adult
Bank
Default
Shoppers
Taxi
Average
Methods
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
ERM
71.30
48.18
70.23
40.13
62.98
36.12
77.12
53.91
67.53
59.48
69.83
47.56
CVaR-DRO
71.95
49.02
68.93
38.71
62.43
34.60
76.65
51.26
65.55
54.51
69.10
45.62
x -DRO
71.85
50.09
70.88
40.04
62.03
34.16
76.95
52.40
68.27
62.43
70.00
47.83
KL--DR
74.33
49.17
69.95
39.89
59.03
19.39
79.53
59.74
62.30
47.93
69.03
43.22
EMI..
69.37
38.97
61.85
21.69
65.05
28.91
74.82
46.18
69.52
58.00
68. 12
38.75
JTT
71.46
49.93
68.78
37.77
62.47
35.51
78.10 10
52.59
67.47
60.01
69.65
47.16
FAM
72.85
49.59
71.03
41.09
62.50
37.12
76.60
53.51
67,90
62.24
70.18
48.71
SRDO
71.27
46.44
66,34
32.08
62.38
33.34
76.24
53.11
64.57
58.18
68.16
44.63
Ours
74.33
54.79
69.50
41.28
62.62
38.78
79.68
60.73
67.85
63.14
70.80
51.74
expanding the search space without adequately modeling the underlying data distribution.
robustness of our approach.
For the generalization experiment on different-souree data (Columns 11 to 19), our method does not achieve the highest mean accuracy. This is due to the change in causal mechanisms across states (Liu et al., 2024), leading to variations in the predictive mechanism. How- ever, our method still achieves strong worst-case accu- racy, thanks to our modeled relative probability density. Our training process ensures fair treatment of covariates through score-based similarity, which is more reliable than the boundaries employed by other methods.
Another popular invariant learning- based method, EIIL (Creager et al., 2021), achieves high mean accuracy. How- ever, its worst-case accuracy is relatively disappointing com- pared to other methods. EIIL divides samples into different groups based on generated environment labels, which cre- ates implicit boundaries based on data's inherent character- istics. While this boundary- -based approach may improve overall accuracy by exploiting data biases, it undermines the balance and robustness required to handle sensitive non- causal attributes.
Compared to the above methods which rely on either prior information or boundaries, our method could achieve con- sistent and stable robustness on these datasets. The results demonstrate the benefit of modeling relative probability density as referred in Section 1.
4.5.3. INTEGRATION WITH TREE-BASED MODELS
Though our score-modeling process relies on neural net- works, the resulting weights in Eq. 13 can be applied to various model architectures during the training of the final classification model. Prior studies have shown that tree- based models perform particularly well on tabular data, es- pecially in industrial settings (Grinsztajn et al., 2022: Liu et al., 2024). Consequently, we apply our score-derived weights to tree-based models to assess the generality of our approach. As shown in Table 3, we evaluate three widely used models CatBoost (Prokhorenkova et al., 2018), Light- GBM (Ke al., 2017), and XGBoost (Chen & Guestrin, 2016)-with our weighting scheme. For a thorough com- parison, we also integrate JTT, a boundary-based method in- troduced in Section 1, into 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 in Section 1-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 en- hance worst-group accuracy. These results further confirm our method could work well with tree-based models.
4.5.2. GENERALIZATION CAPABILITY UNDER SELECTION BIAS
We further validate the generalization capability of our method with the ACS Income dataset. The key objective is to assess whether our model can learn a robust predictive function from a single source environment that generalizes well to others. We randomly select three geographically distant states in the U.S. and train classification models separately on data from each state.
The second to ninth columns of Table 2 show the results when the model is tested on data from its source state. We then evaluate each model on data from another state in a round-robin manner, with the results recorded in the tenth to last column. The observations are as follows:
For the single-source data experiments (Columns 2 to 10), our method achieves both higher mean accuracy and worst-case accuracy compared to most baselines. This consistency with the results in Table 1 supports the
7

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Table 2. The performance on ACS Income task.
The bold and underline denote the best and the second hest results respectively.
Methods
AZ
MA
MI
Average
AZ
MA
MA
>>
MI
MI
> AZ
Average
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
ERM
76.95
64.07
78.18
74.05
73.70
62.42
75.83
66.04
73.05
57.82
74.28
67.78
72.38
56.77
73.23
60.79
CVaR-DRG
77.00
64.41
77.85
72.85
73.48
61.88
76.11
66.39
74.10
59.42
74.85
68.57
71.75
54.67
73.57
60.88
XRDDRO
76.95
64.07
77.55
71.82
73.13
59.03
75.88
64.97
74.53
61.18
74.23
68.53
71.33
53.05
73.36
60.92
KL-DRO
75.68
60.98
78.23
71.22
76,10
62.88
76.67
65.03
74.20
57.78
75.35
68.72
74.73
57.20
74,76
61.23
EIIL
74.98
55.22
78.18
68.42
75.60
63.77
76.25
62.47
75,35
56,80
76.92
65.62
74.78
57.32
75.68
59,91
JT
75.70
57.73
77.48
68.48
74.48
60.63
75.88
62.28
73.90
55.17
74.00
66.57
72.65
55.17
73.52
58,97
FAM
75.00
57.82
78.73
70.12
74.38
59.45
76.03
62.46
72.93
54.60
74.68
66.35
73.40
54.53
73.67
58.49
SRDO
75.17
60.03
77.85
69.52
73,30
54.63
75.44
61.39
74.00
58.37
74.78
61.77
73.83
60.05
74,39
60.07
Ours
76.28
66.42
78.35
73.33
75.53
67.10
76.72
68.95
75.00
62,85
74.75
68.75
74.73
64.10
74.83
65.23
4.6. Visualization of Score-similarity-asse Weights
depth represents the log- log-likelihood of the predicted proba- bility density. Points with higher density are shaded closer to violet, while those with lower density appear more blue. Interestingly, not all points in high- tigh-probability regions are assigned warm colors. In fact, most points are blue, similar to those in low-density regions. Only a few points located at cluster centers exhibit significantly higher log- logllliiihhoos. Reweighting based on these estimates could lead to dispro- portionate emphasis on small number of extreme high- density points, potentially obscuring the distinction between other high- and cow-probbbility regions.
In this section, provide the visualization results on synthetic dataset directly explain the reweighting process of our methods. In addition, we will explain why we use our aggregated similarity difference measure in Eq. 12 in- stead of running a sampler based on the probability flow ordinary differential equation that allow for exact likelihood computation proposed by Song et al..
Synthetic Data Generation. Our imbalanced synthetic training data has been illustrated in Figure la. The data have two features, Ho and 21. The perfect classification boundary is a sine curve. We introduce the bias on training data by sampling more points in regions where zo E 1 35] U We expect the reweighting process could make the whole data distribution fairly balanced.
Density- Weights Computed from Our Similarity- based Measure. We then visualize the unnormalized score similarity in Figure 3b. Points with lower error, correspond- ing to higher relative probability density, are shaded in blue and are expected to be downweighted. The visualization clearly distinguishes these points, with two regions of rela- tively high probability density appearing in blue. In contrast, points with truly low probability are represented by higher aggregated error, shown in purple, and can be easily up- weighted. Our similarity measure not only differentiates density clearly but also maintains a consistent and stable numerical range, making it suitable for sample reweighting.
[ 2m]
Exact Likelihood Computation by Probability Flow ODE, As referred in Section 3.3, we can compute the exact log-likelihood through a probability flow ODE in Eq. 10. Following Song et al., we apply the Skilling- Hutchinson trace estimator (Skilling, 1989; Hutchinson, 1990) to esti- mate V fo (x(1), 1) and finally achieves the log- lieeliiel..ood.
Extreme Value Problem Brought by Exact Likelihood. The log- ligelliooo of p(x) represents the estimated probabil- ity density at the feature level, making it useful for guiding the reweighting process. We compute the mean of the es- timated probabilities across all features to determine the final sample-level probability density. In Figure 3a, color
In summary, compared to the exact log- likeliihoood calcula- tion, our method more effectively captures relative probabil- ity densities across the dataset. By computing similarity, any two sample points with significantly different probability densities are clearly distinguished and distinctly colored. In contrast, the exact log-likelihood tends to overemphasize few high-density points, obscuring the distinction between other high- and low-density samples.
0.2
4.7. Ablation Studies
-150
We conduct the sensitivity analysis in this section. As re- ferred in Algorithm 1, our method has three hyperparame- ters. T denotes the number of selected timesteps for com- puting aggregation similarity. 7 controls the strength of reweighting based on SimDiff(-). K represents the number of repeated sampling iterations for computing scores. The
-390
(a) Computed log likelihood (b) Score similarity as proxy of p(x)
Figure 3. Our score similarity is good proxy for density.
8

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Table The experiments deploying tree-based models. Compared JTT, our score-hased weights could consistently improve the worst-case performance with tree based models.
Adult
Default
Phishing
Taxi (NYC)
Method
Worst
Mean
Worst
Mean
Worst
Mcan
Worst
Mcan
CatBoost
50.10 0.85
75.38 0.15
31.79 + 0.38
64.70 0.33
97.18 + 0.11
97.60 + 0.12
66.95 0.17
81.30 + 0.40
CatBoost + JTT
48,94 0.40
31.44 0.15
0.96
64.29 0.07
95.88 0.84
97.70
6.992
81.52 = 0.70
CatBoost + ours
64.63 = 0.29
7
33.43 # 0.29
0.81
97.15 0.20
98.10
67.60 1.03
81.23 = 0.61
LightGBM
50.33 + 0.66
75.79
+
32.00 a 0.38
64,98 # 0.15
96.10
97.20 0.14
68.35 + 1.00
81.67 + 0.32
LightGBM + JTT
49.03 # 1.01
75.26 0.32
31.91
65.03 # 0.36
96.57 + 0.32
98.07 +: 0.14
68.55 1.55
81.88 + 0.54
LightGBM ours
0.002
79.53 0.00
34.94 # 0.66
65.97
96.57 = 0.23
97.92 # 0.08
69.77
82.40 + 0.80
XGBoost
50.76 1.70
0..4
31.90 1.04
63.84
0.39
0.44
6..28
81.65 # 1.20
XGBoost JTT
49.42
74.63 0.06
02.04.0..5
96.33 63.86 = 0.30
0.112
1.42
81.32 + 2.36
XGBoost + ours
64.93 + 0.39
78.93 0.21
34.19 + 0.25
64.83 0.68
96.62 + 0.38
97.98 + 0.30
69.10 1.70
81.47 # 0.50
discussion on K is deferred to Appendix C.2, as it does not functionally impact the score computation.
strategies. To overcome these limitations, we propose using score-based models to capture the latent data distribution. Specifically, we estimate scores at several fixed timesteps and use their similarity to model the relative probability density of each sample. A score-based reweighting strat- egy is then employed train a robust classification model. Our approach requires no prior information during training and ensures that the reweighting process aligns with the original data distribution. Experiments on seven datasets demonstrate that our method effectively balances the origi- nal training data globally and achieves robust performance under distribution shifts.
Table 4, Ablation studies of the number of chosen timesteps T on Adult dataset.
15
Weee
Mias
Wone
Miain
Mase
Acc. (%)
3133 2017
56.79 sous
1402 aeas
56,00 us.s
37.40
The ablation study results of I' on Adult are shown in Table 4. There are two main observations:
The mean accuracy shows minimal varia- tion with changes in T.
Acknowledgements
Wors (er))
Mew hes (e-2)
This work was supported in part by the National Key Research and Development Program of China (2024YFB3312900) National Natural Science Foundation of China (62376243, 62441605), "Pioneer" and "Leading Goose" R&D Program of Zhejiang (2025C02037), Key R&D Program of Zhejiang (2025C01 128) (a Fundamen- tal Research Funds for the Central Universities), and Ant Group. All opinions in this paper those of the authors and do not necessarily reflect the views of the funding agen- cies. We also express our deepest gratitude to Yulin Kang for his contributions and support to this work.
Wore (emh)
Maa (oml))
Wort Art (e-4)
Max Ace (o-4)
Worn Ane (rd))
Man Are (10))
The worst-case accu- racy improves as T increases. This is because T indirectly represents the number of aggregated noisy points in the sampled neighborhood. A larger T leads to a more accu- rate computation of aggregated similarity in Eq. 11. We set the default value of T to 10, balancing fast computa- tion with optimal performance.
20
25
36
35
49
45
50
Figure Ablation study on tem- perature of reweighting scale
In Figure 4, we show the classification results by using dif- ferent r. Our model performs consistently well with all the values. The worst accuracies under different r all surpass the baseline methods in Table 1, which proves that our method does not heavily rely on the values of hyperparameters.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
5. Conclusion
References
In this paper, we tackle the challenge of improving model robustness, measured across all non-causal covariates rather than focusing on a single attribute. We observe that pre- vious methods are limited by their inability to model the original joint data distribution and apply effective balancing
Arjovsky, M., Bottou, L., Gulrajani. I., and Lopez- Paz, D. Invariant risk minimization. arXiv preprint arXiv: 009999299,,... 2019.

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Becker, B. and Kohavi, R. Adult.
UCI
Ma- DOI:
Higgins, I., Matthey, L.., Pal, A., Burgess, C. P., Glorot, X., Botvinick, M. M., Mohamed, S., and Lerchner, A. beta- vae: Learning basic visual concepts with constrained variational framework. ICLR (Poster), 3, 2017.
chine Learning Repository, 1996.
htpp://ddo.or//122444//////
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. In Advances in neural information processing systems, pp. 6371-6583, 2018.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in Neural Information Process- ing Systems, 33:6840-681) 2020.
Chen, T. and Guestrin, C. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd acm sigkdd inter- national conference on knowledge discovery and data mining, pp. 785-794, 2016
Hu, Z., Zhang, F., Chen, L., Kuang, K., Li, Gao, K., Xiao, J., Wang, X., and Zhu, W. Towards better alignment: Training diffusion models with reinforcement learning against sparse rewards. arXiv preprini 1rX1v:2503.11240. 2025.
Choi, S., Jung, S., Yun, H., Kim, T., Kim, S., and Choo, J. Robustnet: Improving domain generalization in urban-scene segmentation via instance selective whiten- ing. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11580-11590, June 2021.
Hutchinson, M. F. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Com- munications in SkaiitiissSimulltti and Computation, 19(2)433 450, 1990.
Ilse, M., Tomczak, J. M., Louizos, C., and Welling, M. Diva: Domain invariant variational autoencoders. In Arbel, T., Ben Ayed, I., de Bruijne, M., Descoteaux, M., Lombaert, H., and Pal, C. (eds.), Proceedings of the Third Conference on Medical Imaging with Deep Learning, volume 121 of Proceedings of Machine Learning Research, pp. 322-348. PMLR, 06-08 Jul 2020. httptt httpp/ppeeeeng..e/////// mlr. prr.prsss///////.o v121/ilse20a. html.
Creager, E.,, Jacobsen, J.-H., and Zemel, R. Environment inference for invariant learning. In International Confer- on Machine Learning, 2021.
Dendorfer, P., Elflein, S., and Leal-Taie Taixé, L. Mg-gan: A multi-generator model preventing out-of-distribution sam- ples in pedestrian trajectory prediction. Proceedings of the IEEF/CVF International Conference on Computer Vision (ICCV), pp. 55158-13167 October 2021.
Kadkhodaie, 7., Guth, F., Simoncelli, E. P., and Mallat, S. Generalization in diffusion models arises from geometry- adaptive harmonic representations. In The Twelfth In- lernational Conference Learning Representations, 2024. URL httpt: /ttpas:///ppeeeeeeeeee net ntt.e/uuuntt/ouru id=ANvmVS2YrG.
Ding, F., Hardt, M., Miller, J., and Schmidt, L. Retiring adult: New datasets for fair machine learning. Advances in neural information processing systems, 34:6478-6490, 2021.
Duchi, J. C. and Namkoong, H. Learning models with uni- form performance via distributionally robust optimization. The Annals of Statistics, 49(3):1378-1444 1406, 2021.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-based based generative models. Advances in neural information processing systems, 35: 26565 26577, 2022.
Ghaznavi, M., Asadollahzadeh, H., Noohdani, F. H., Tabar, S. V., Hasani, H., Alvanagh, T. A., Rohban, M. H., and Baghshah, M. S. Trained models tell us how to make them robust to spurious correlation without group annotation. arXiv preprint arXiv:2410.05344 2024.
Ke, G., Meng, Q., Finley, T., Wang. T., Chen, W., Ma, W., Ye, Q. and Liu, T.-Y. Lightghm: highly efficient gradient boosting decision tree. Advances in neural infor- mation processing systems, 30, 2017.
Grinsztajn, L., Oyallon, E., and Varoquaux, G. Why do tree- based models still outperform deep learning on typical tabular data? Advances neural information processing systems, 35:507-520. 2022.
LaBonte, T. Muthukumar, V., and Kumar, A. Towards last-layer retraining for group robustness with fewer an- notations. Advances in Neural Information Processing Systems, 36: 1552-111115 2023,
Gu, Y., Fang, C., Bühhmann, P., and Fan, J. Causality pursuit from heterogeneous environments via neural adversarial invariance learning. arXiv preprint XXiv:2405.04715. 2024.
Levy, D., Carmon, Y., Duchi, J. C., and Sidford, A. Large- scale methods for distributionally robust optimization. Advances in Neural Information Processing Systems, 33: 8847-8860, 2020.
10

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Lin, Z.,, Tang, Z., Liu, X., Gong, Y., Cheng, Y., Chen, Q. Li, H., Xin, Y., Yang. Z.,, Yang, K., et al. Sigma: Differential rescaling of query, key and value for efficient language models. arXiv preprint arXiv:2501.13629, 2025.
Sakar, C. and Kastro, Y. Online Shoppers Purchasing Inten- tion Dataset. UCI Machine Learning Repository, 2018. DOI: http:://bo..rry/1024333///8888
Shen, X., Bühlmann, P., and Taeb, A. Causality-oriented ro- bustness: exploiting general additive interventions. arXiv preprint arXiv:2307.10299 2023.
Liu, E. Z., Haghgoo, B., Chen, A. S., Raghunathan, A., Koh, P. W., Sagawa, S., Liang, P., and Finn, C. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pp. 6781-6791 PMLR, 2021.
Shen, Z., Cui, P., Zhang, T., and Kunag, K. Stable learning via sample reweighting. In Proceedings of the AAAI Con- ference on Artificial Intelligence, volume 34, pp. 5692- 5699, 2020.
Liu, J., Wang, T., Cui, P., and Namkoong. H. On the need for a language describing distribution shifts: Illustrations tabular datasets. Advances in Neural Information Processing Systems, 36, 2024.
Skilling, J. The eigenvalues of mega-dimensional matrices. Maximum Entropy and Bayesian Methods, pp- 455- 466. Springer, 1989.
Lv, Z., Zhang, W., Chen, Z.., Zhang, S., and Kuang, K. Intelligent model update strategy for sequential recom- mendation. In Proceedings of the ACM Web Conference 2024, pp. 3117-3128, 2024.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Moro, S., Rita, P. and Cortez, P. Bank Market- ing. UCI Machine Learning Repository, 2012. DOE: http:://ooiorr/1002444222/
Song. Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A.,
Ermon, S., and Poole, Score-based generative mod- eling through stochastic differential equations. In In- ternational Conference Learning Representations, httpp:://ppeereeeeee howpnreview.eeoeeoooooe net/ forum?
Nam, J., Cha, H., Ahn, S., Lee, J., and Shin, J. Learning from failure: De- biasing classifier from biased classifier. Advances in Neural Information Processing Systems, 33: 20673-20684, 2020.
2021. URL https:/tp////..... id 1d-PXTIG12RRHS.
Sun, B. and Saenko, K. Deep coral: Correlation alignment for deep domain adaptation. In European conference on computer vision, pp. 443 450. Springer, 2016.
Navas, M. quito and
Taxi routes of mexico city, more, 2018. URL https: wtt.e//atttt////aaaaaaaaaeeaaa
Tang, Z., Lv, Z., Zhang, S., Wu, F., and Kuang, K. Mod- elgpt: Unleashing llm's capabilities for tailored model generation. arXiv preprint arXiv:2402.12400 2024a.
Petzka. H., Kamp, M.. Adilova, L.. Sminchisescu, C., and
Boley, M. Relative flatness and generalization. In Ad- vances in Neural Information Processing Systems, vol- ume 34. Curran Associates, Inc., 2021.
Tang, Z., Lv, Z., Zhang, S., Zhou, Y., Duan, X., Wu, F., and Kuang. K. Aug-kd: Anchor-based mixup generation for out-of-domain knowledge distillation. arXiv preprint arXiv:2403.0700 2024b.
Prokhorenkova, L., Gusev, G., Vorobev, A., Dorogush, A. V., and Gulin, A. Catboost: unbiased boosting with categori- cal features. Advances in neural information processing systems, 31, 2018.
Tong. Y., Yuan, J., Zhang, M., Zhu, D., Zhang, K., Wu, F, and Kuang. K. Quantitatively measuring and contrastively exploring heterogeneity for domain generalization. In Proceedings of the 29th ACM SIGKDD Conference Knwwledge Discovery and Data Mining. KDD '23. ACM, August 2023. doi: 10. 3599481. URL 33994481.
Qiu, S., Potapczynski, A., Izmailov, P., and Wilson, A. G. Simple and Fast Group Robustness by Automatic Fea- ture Reweighting, International Conference Machine Learning (ICML), 2023.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF con- ference on computer vision and paltern recognition, pp. 10684-10695, 2022.
Van der Maaten, L., and Hinton, G. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008
Sagawa". S., Koh", P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks. In International Conference on Learning Representations, 2020.
Vapnik, V. An overview statistical learning theory. IEEE transactions on neural networks, 10(5):988-99) 1999.

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Wu, C., Restrepo, D., Nakayama, L. F., Zago Ribeiro, L., Shuai, Z., Barboza, N. S., Sousa, M. L. V., Fitterman, R. D., Pereira, A. D. A., Regatieri, C. V. S., et al. A portable retina fundus photos dataset for clinical, demo- graphic, and diabetic retinopathy prediction. Scientific Data, 12(1):323, 2025.
Zhang, X., Cui, P., Xu, R., Zhou, L., He, Y., and Shen, Z. Deep stable learning for out-of-distributiin general- ization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 5372- 5382, 2021.
Zhao, C., Liu, M., Wang. W., Chen, W., Wang, F,, Chen, H., Zhang, B., and Shen, C. Moviedreamer: Hierarchi- cal generation for coherent long visual sequence. arXiv preprint arXiv:2407. 16655, 2024.
Xu, Y., Liu, Z., Tegmark, M., and Jaakkola, Poisson flow generative models. arXiv preprint arXiv:2209.11178, 2022.
Xu, Y., Tong, S., and Jaakkola, S. Stable target field for reduced variance score estimation in diffusion models. In The Eleventh International Conference on Learning Representations, 2023, URL https: openreview. net/ Eorum? wwWWIwwTdITT.............
Zhao, C., Liu, M., Zheng. H., Zhu, M., Zhao, Z., Chen, H., He, T., and Shen, C. Diception: A generalist diffu- sion model for visual perceptual tasks. arXiv preprint arXiv:2502. 2025.
Zhu, D., Song, Y., Shen, T., Zhao, Z., Yang, J., Zhang, M.. and Wu, C. Remedy: Recipe merging dynamics in large vision- models. In The Thirteenth International Conference on Learning Representations, 2025.
Yang, J., Chen, Z., Xiao, T., Lin, Y., Zhang, W., and Kuang, K. Leveraging invariant principle for heterophilic graph structure distribution shifts. In Proceedings of the ACM on Web Conference 2025, pp. 1196-1200 1204, 2025a.
Zou, Y., Kawaguchi, K., Liu, Y., Liu, J., Lee, M.-L., and Hsu, W. Towards robust ouitt--distributionn generalization bounds via sharpness. arXiv preprint arXiv:2403.06339. 2024.
Yang, J., Jin, D., Tang, A., Shen, L.., Zhu, D., Chen, Z.,,
Wang, D., Cui, Q. Zhang, Z., Zhou, J., et al. Mix data or merge models? balancing the helpfulness, honesty, and harmlessness of large language model via model merging. arXiv preprint 2rXiv:2502.06876 2025b.
Yang, L., Tong, Y., Gu, X., and Sun, Explain tempo- ral black-box models via functional decomposition. In Forty-first International Conference on Machine Learn- ing. 2024.
Yeh, I-C. Machine Learning Repository, 2016. DOI:
Default of Credit Card Clients.
UCI
https://bb.orr/l022443//5555
Zhang. F., Kuang, K., Chen, L.. Liu, Y., Wu, C., and Xiao, J. Faire contrastive learning with partially anno- tated sensitive attributes. In The Eleventh International Conference on Learning Representations, 2022.
Zhang. F., He, Q., Kuang, K., Liu, J., Chen, L., Wu, C., Xiao, J., and Zhang, H. Distributionally generative augmenta- tion for fair facial attribute classification. In Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pp. 22797-22808, 2024a.
Zhang, F., Shuai, Z., Kuang. K., Wu, F., Zhuang, Y., and Xiao, Unified fair federated learning for digital health- care. Patterns, 5(1), 2024b.
Zhang, H., Zhang, J., Srinivasan, B., Shen, Z., Qin, X., Faloutsos, C., Rangwala, H., and Karypis, G. Mixed- type tabular data synthesis with acore-based diffusion in latent space. In The twelfth International Conference on Learning Representations, 2024c.
12

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
The supplementary materials structured follows:
Appendix provides the pseudocode of our method.
Appendix B.1 details the default network preconditioning, following the practice of EDM (Karras et al., 2022).
Appendix B.2 lists the selected non-causal attributes used for robustness evaluation.
Appendix C.I presents the standard deviations of our experiments.
Appendix C.2 offers a sensitivity analysis of hyperparameter K.
Appendix C.3 conducts the analysis on sample size's influence on our method.
Appendix C.4 investigates how our method balances the imbalanced real-world dataset.
Appendix C.5 provides the visualization results of our score-based weights on real-world dataset.
Appendix D.1 compares our method with stable learning-baa -baed methods.
Appendix D.2 discusses the relation between mean and the worst-case accuracy on synthetic dataset.
Appendix D.3 discusses the relation between mean and the worst-case accuracy on real dataset.
A. The pseudocode of our method
Algorithm The pseudocode of our method in binary classification problem.
Input:
Training dataset = {x, y) {a,, 3/ N-11, fixed noises along with specific timesteps (o(t)) the temperature of the reweighting scale r. the repeated sampling times K
Parameters to be optimized: VAE Φz((), two score-based models {s)(s-)).s e {0, 10, final classification model (.)
Stage One: Training Distribution Modeling
Train a VAE oz = E'n Dec with Eq. 4
Obtain latent representation with trained encoder: (x)
Separate into ²y=0 and Zy=1 according to their labels, train score model sy with 2y separately by Eq. // Stage Two: Probability Density Aligning
2, compute the aggregated similarity with target class model Sim sym9.,) ) and with non-target class model Sim 2;1 $yz through Eq. 11
For each latent
Compute SimDiff( 2;) through Eq. 12 to indicate 2;'s relative probability density
Stage Three: Unbiased Learning on Distribution-balane
Data
Reflect SimDiff into weights w; through Eq.
Train the final unbiased classification model with Eq. 14
Return Enc and for testing
B. Supplementary details of the setting in our experiments
B.1. The choice of network preconditioning
In our method, we follow the practice from EDM (Karras et al., 2022) to train our neural network. The key lies in using preconditioning techniques to make the output of neural network stable, instead of varying with the scale of variance a (t)e. We take the default choice of scaling factors from Karras al.; Zhang al.. The details listed in Table 5
set Juzaa = 0.5 our experiment as TabSyn (Zhang et al., 2024c) did.
13

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Table The choices various scaling factors for denoiser Do(:.) Eq.
Skip scaling Cskip(7) (7)
Raaa/(o² + +.
Output scaling Coul (0)
dduaa o data +2
Input scaling Cin (T)
1/ 0 + data
Noise cond. Cnoise (0)
II((r)
B.2. Details of the selected non-causal attributes for measuring overall robustness
As referred in Section 4.2, measure overall robustness by selecting several non-causal attributes and recording their worst-case prediction results. Here we provide the details of selected attributes.
Adult (Becker & Kohavi, 1996): We select marital status, race, and sex as sensitive attributes.
Bank (Moro et al., 2012): We select age, housing status, marital status, and the last contact duration as sensitive attributes.
Default (Yeh, 2016): We select age, sex, and the amount of the given credit as sensitive attributes.
Shoppers (Sakar & Kastro, 2018): We select the traffic type, the visitor type as returning or new visitor, a Boolean feature indicating whether the date of the visit is weekend sensitive attributes.
Taxi (Navas, 2018): We select the indicator for weekday, the month of picking up. and the direction as sensitive attributes.
US-Wide ACS PUMS Data (Ding et al., 2021): We select race and sex as sensitive attributes.
We demonstrate the Pearson correlation coefficients between the attributes and target vairable in training and test data in Table 6.
We select sensitive attributes for evaluation based on the following criteria:
Weak linear correlation with the target variable: Selected covariates should not exhibit a strong correlation with the target variable. For instance, in the Adult dataset, the attributes marital status, race, and were chosen because their Pearson correlation coefficients with the target variable are relatively smaller compared to attributes like native country and workclass.
Divergent correlation statistics across training and test datasets: Selected attributes should show notable differences in correlation coefficients between the training and test datasets. For example, in the Taxi dataset, the selected attributes exhibit varying correlation coefficients in the training and test datasets, suggesting that these attributes are not direct causes of the target variable in this context.
C. Further experiments
C.1. Details of the standard deviation of our experiments
We randomly conducted each experiment three times using different seeds and computed the mean classification results, as shown in Tables and 2. The standard deviation across these three runs is provided in Tables 7 and 8, respectively.
C.2. The ablation studies of repeated sampling times
In Section 4.7, we conduct sensitivity analysis the number of selected timesteps T and the temperature which controls the reweighting scale. These two hyperparameters directly influence score computation. Additionally, we have another hyperparameter, K, which governs the number of repeated sampling iterations. While K does not functionally affect score computation, it impacts the robustness of our computed values. Therefore, we will analyze the effect K independently in this section to distinguish it from T and т.
14

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Table The Pearson correlation coefficients between our selected sensitive attrihutes (with regular font) and the target variahle in training
and test data, For ease of comparison, we also show some typical attributes (with italic font) not selected for evaluation. We set attributes that either (1) exhibit low correlation with the target variable across both training and test datasets, or (2) demonstrate significant variation in correlation with the target variable between the training and test datasets as sensitive (non-causal) attributes. Then we evaluate the
model bias concerning these attributes.
Pearson Coefficient
Pearson Coefficient
Dataset
Attrihute
Dataset
Attribute
train
test
train
test
marital status
-0.0345
-0.0287
weekday
0.0200
+0.0169
Adult
race
-0.0852
-0.0807
Taxi
month
-0.0004
-0.0243
sex
0.0785
0.0700
direction
0.0180
0.0754
native country
0.3872
0.3915
distance
0.4501
0.4618
work class
-0.2160
-0,2119
hour
0.0413
0.1220
age sex
0.0193
0.0224
age
-0.0886
-0.0130
Default
-0.0396
-0,0430
housing slatus
-0,0690
-0.0607
Bank
given credit
-0.0368
-0.0287
marital status
-0.0884
-0.0757
education
-0.1407
-0.1365
duration
0.0287
0.0263
payment
0.3297
0.2814
job
0,2807
0.2718
loan
-0.1395
-0.1368
traffic type visitor type
-0.0548
-0.0644
Shoppers
-0,0276
-0.0533
race
-0.1127
-0.1312
ACS Income (AZ)
weekend
0.0277
0.0445
sex
-0.1312
0.1205
administrative
0.1422
0.1401
marital status
0,2307
0.2472
administrative duration
-0.1042
-0.1030
age
0.2658
0.2713
ACS Income (MA)
race
-0.1030
-0,1486
ACS Income (MI)
race
-0,0620
-0.0324
sex
0.1435
0.1086
sex
0.1806
0.2030
marital status
0.2925
0.2704
marital status
0,2518
0,2405
age
0.2748
0,2560
age
0.2469
0.2127
The experimental results listed in Table 9. We could find larger K corresponds to longer computation time and more
stable performance. We choose 32 as the default value for K, which achieves good trade-off between sampling time and model performance.
C.3. The analysis sample size's influence on our method
Our method use score-based model to model the original data distribution, serving purpose similar to traditional density estimation methods such as 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.
In this section, we examine the impact of training samples' size on our score-basee based proxy. We create subsets of varying sizes from the original training dataset and use these subsets to train score models. The final classification model is trained on the weights from new 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 Table 10. It demonstrates that the weights generated from our soore-based proxy is insensitive to the sample size. The robustness of our method arises from score model's ability to construct the latent score field. As Kadkhodaie et al. (2024) stated, 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 and R 0.8). However, when 0.2. the number of training samples becomes insufficient to construct an accurate score field, which makes the estimated weights less accurate. However, the performance of our model still outperformed the results of the other baselines using the full dataset in Table 1.
when we used only one-fifth of the data for training,
In a word, 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 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.
15

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Table
The standard deviation of the classification results three runs five datasets.
Adult
Bank
Default
Shoppers
Taxi
Methods
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
ERM
0.47
1.64
0.35
1.85
0.21
3.93
0.73
2.46
0.14
2.61
CVaR-DRO
0.68
2.63
1.52
3.59
0.33
2.44
0.31
2.01
2.00
5.93
-DRO
0.73
1.78
0.81
3.56
0.19
2.12
1.72
3.60
0.99
1.08
KL-DRO
0.14
0.90
0.60
1.75
3.25
5.10
1.08
3.98
8.96
23.68
EIIL
2.50
8.15
2.49
6.91
1.06
7.55
7.04
15.02
0.07
4.43
JTT
1.89
3.65
0.67
1.50
0.90
1.58
0.75
3.35
0.42
2.60
FAM
1.65
1.95
3.89
5.25
0.19
2.07
0.09
2.05
0.75
0.63
SRDO
2.43
6.82
2.55
5.44
8.06
10.94
0.35
2.31
2.07
2.50
Ours
0.28
2.01
0.39
1.56
0.12
1.22
1.44
3.39
0.16
2.02
Table 8 The standard deviation of the classification results across three runs ACS dataset.
AZ.
MA
MI
AZ.
MA
MA -
MI
MI +
AZ
Methods
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
ERM
0.85
1.31
0.46
1.23
0.49
2.66
0.78
1.62
0.32
2.50
0.67
1.82
CVaR-DRO
0.14
3.70
1.27
1.88
0.74
4.97
0.35
2.51
0.14
0.65
0.28
2.55
DRO
0.14
4.60
0.92
1.58
1.38
2.06
0.25
1.15
0.18
0.67
0.60
1.76
KL-DRO
1.23
3.00
0.32
2.41
0.07
0.37
1.20
2.91
0.78
1.71
0.25
0.97
EIIL
1.80
6.65
0.67
5,86
0.49
3.19
1.70
5.52
1.59
5.88
1.51
4.95
JTT
1.06
2.49
0.25
2.70
0.53
1.84
0.78
3.14
0.28
2.33
0.35
0.62
FAM
0.57
2.17
0.03
3.33
2.09
6.29
0.53
2.45
0.11
2.44
1.27
5.24
SRDO
2.22
10.23
1.01
2.49
1.22
5.75
1.83
8.57
0.49
4.26
0.35
0.96
Ours
0.11
0.89
0.14
1.48
1.77
1.61
0.07
0.40
0.35
0.30
0.35
0.35
C.4. The balanced samples after deploying our weights
To intuitively demonstrate how our score-based weights balance the original dataset, divide the samples into groups based the sensitive attribute and target label, then compute the of weights for each group. Without accounting for training challenges caused by group-specific variance, expect the weights achieve balanced distribution. Specifically, the majority group is expected to have lower weighted sum compared to its unweighted sum. Table 11 presents the results for four groups after applying our score-based weights. For nearly all the listed non-causal attributes, the sum of weights for the majority group decreases, while the sum for the minority group increases. This observation confirms that our weights effectively balance the original distributional shift, offering clear explanation of how our method operates on real-word sol-w datasets.
C.5. The t-SNE visualization comparison with our score-based weights
previously provided an intuition for how our score-based weights help balance datasets using a synthetic example in Figure 1. In this section, we present a straightforward visualization on Default dataset demonstrate how our method balances the real-world dataset. To be specific, we use t-SNE (Van der Maaten Hinton, 2008) to reduce the latent representation two dimensions for visualization. Note that in Table 11, the Default dataset predominantly exhibits label y shift. Therefore, we divide the representations into majority group (y 0) and minority group (v 1). In Figure 5a, points from the minority group (y = 1) are shaded in dark green, while the majority group (y = 0) is colored brown. We could observe that minority samples are often situated in the marginal regions clusters, indicating that they have low probability densities. Correspondingly, we visualize these samples' score-basee based weights in Figure 5b. Points with higher weights are represented by warmer colors like purple. We expect that the points from the minority group in Figure 5a be
16

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Table 9, Ablation studies of the repeated sampling iterations K' Adult dataset.
K
Mean (%)
Worst (%)
Running Time (s)
8
73.97
10.50
52.47 10.71
7.35 10.01
16
74.13 +0.67
54.18
1.90
14.28 =1.50
32
74.33
+0.28
54.79
1:2.11
25.32 =0.01
48
75.29
+0.40
56.28
+ 1.44
38.48 2.33
Table 10. Experimental performance under different ratio on Adult dataset. The number of original training and test samples is 32561 and 16281.
R = 0.2
R 0.5
R 0.8
R = 1, original
Mean
Worst
Mean
Worst
Mean
Worst
Mean
Worst
Acc. (%) 73.81 11003
54.24 15.48
74.23
1..25
54.57
1198
74.78
10.62
54.64 1.190
74.33 10.28
54.79 #2.01
assigned higher weights in Figure 5b. Comparing these two figures, we can observe that our method significantly increases the weights for samples in the minority group. Almost all the dark green points in the left figure are shaded closer to purple in the right figure, which validates the effectiveness of our score-based weights.
llihh Weight
60
60
40
20
-20
-20
40
Majority Group
-40-
Minority Group
-60
-60
Low Weight
-50
-25
25
50
-50
-25
0
25
50
(a) Group-spectinn specified latent representation
(b) Latent representation with our weights
Figure 5. The t-SNE visualization the original latent representations and their corresponding weights in Default datasct.
D. Other discussions beyond our method
D.1. The comparison with stable learning
In this section, we provide a comprehensive comparison between our method and stable learning. Stable learning aims to decorrelate features to achieve a uniform and balanced data distribution, which similar to our approach. However, our method offers three sigmificant advantages over stable learning, as detailed below:
Ability to handle potential Y-shift problems: Our method use a similarity difference measure address implicit -shiift discussed in Section 3.3. In contrast, stable learning focuses solely on the decorrelation of covariates without differentiating the information carried by labels y.
Balanced Distribution via Original Distribution Modeling vs. Feature Independence: Our method obtains a balanced
17

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Table 11. The number of samples divided into different groups originally well the weighted sum of these samples after deploying our score- score-based weights. The reweighted samples of different groups are more balanced, which is conducive to subsequent unbiased classification learning.
Default
Shoppers
sensitive
attribute
age
given credit
traffic type
visitor type
weekend
original
weighted
original
weighted
original
weighted
original
weighted
original
weighted
original
weighted
r= y=
9144
6943.19
9834
7390.46
12902
9766.14
4635
2298.18
8221
3951.20
2120
1025.14
==0, y=1
3236
5972.06
4180
7690.34
2586
4751.89
1004
3855.73
383
1413.91
1282
4930.48
1=1. y=1
2733
5061.28
1789
3343.00
3383
6281.44
725
2707.13
1346
5148.95
447
1632.38
a=0.
y=0
11887
9023.47
11197
8576.20
8129
6200.52
4733
2235.96
1147
582.94
7248
3509
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, is important to note that balanced distribution does not equate 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 xo and 3i, with samples distributed such that < x < 1 and xo < . < zig + 1. Suppose that original data distribution imbalanced, e.g., samples with 2g > 0.5 all share a same higher density than those with 20 0.5. Under this circumstance, our method can transform the original data distribution p(zo, x1) into a uniform data distribution U (z0, Ø1) at the sample level easily. This is achievable because the score model captures and recovers the original data distribution p(zo, r.).
In contrast, stable learning, which enforces feature decorrelation achieve independence, cannot produce uniform balanced distribution at the sample level while maintaining independence between zo and æ1. The geometry of the sample space-a parallelogram iijion-dictates that zo and xi cannot be independent while maintaining uniform sample-level distribution, i.e., U0 x1) # p(zo) p(z1). Thus, stable learning fails to balance the data distribution at the sample level while simultancously decorrelating features.
Decorrelating features to achieve feature independence is good idea, but it is insufficient to ensure 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 robust balancing operation based on the estimated distribution.
Enhanced Experimental Outcomes Relative to Stable Learning: To facilitate a quantitative comparison between our method and stable learning. we conducted experiments Tables and Our approach shows a minimum improvement of 5% in the worst group accuracy compared the SRDO baseline across both tables. 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 predictor to estimate probabilities, which involves the generation of synthetic samples. For instance, SRDO (Shen et al., 2020) employs empty vectors to create samples, whereas StableNet (Zhang et al., 2021) uses a Random Fourier Transformation. The quality of these synthetic samples crucially affects the training process the predictor, thereby making the estimated probabilities potentially unreliable. In contrast, the score our method represents the gradient of the likelooodoo of the original distribution p(z). 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. utilizes the score, i.e., the gradient of estimated likeliiodd original data distribution, to perform sample reweighting, which is flexible and easy to conduct. The additional experiments further confirm the effectiveness of our method.
D.2. Discussion about the relation between mean and the worst-case accuracy on synthetic dataset
The trade-off between mean accuracy and worst-group accuracy is a well-documented phenomenon. In many cases, a specific optimization objective may prioritize either higher mean accuracy or higher worst-case accuracy. However, we
18

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
Mean Ac)(%)
y 5000
74.5
y,=1 5000
74.0
73.5
73.0
Ny= 5000
72.5
4000
72.0
71.5
71.0
Ny-0 5000
N= 3000
70.5
47.5
50.0
52.5
55.0
57.5
60.0
Worst Worst Acc(%)
Figure 6. The Pareto curve for the synthetic experiment in Section D.2.
want emphasize that the trade-off is not always stable, even in synthetic data. To illustrate this point, we conducted new synthetic experiment inspired by Zhang al..
We designed binary classification task with explicit -ssifft. Data for each class were generated from two distinct multivariate normal distributions with different means and covariance matrices as listed in Table 12. The sample size for class y = 0 was fixed at 5000, while the number of class y = 1 samples varied incrementally from 3000 to 5000 steps of 100. This variation mimicked the effect reweighting, akin to our score-based balancing approach. Models were trained on these mixtures, each representing different proportion of y samples. Each experiment was repeated 500 times. The results are visualized in Figure 6.
Table 12. Data generation parameters for the synthetic dataset.
Class
Mean
Covariance Matrix
Number of Samples
5
5
= 0
[ 1, 0]
fixed at 5000
5
5
15
5
1
[1, 0]
from 3000 to 5000
5
15
In Figure 6, we observe the following:
Effect of Dataset Balance Worst-Case Accuracy: As the dataset becomes balanced, the worst-case accuracy increases consistently, mirroring the trend observed with our score-based reweighting strategy. This observation suggests that balancing the dataset is crucial for improving worst-case accuracy.
Dynamic Interaction Between Mean and The Worst-Case Accuracy: When the number class y = 1 samples
increases from 3000 4000, both mean accuracy and worst-case accuracy increase simultaneously. with no evident
19

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
trade-off. However, the number increases from 4000 to 5000, mean accuracy drops significantly. This phenomenon highlights that the trade-off between accuracy and worst-case accuracy is not always persistent. Instead, their interaction depends on how the optimization process influences the training trajectories.
D.3. Discussion about the relation between mean and the worst-case accuracy on real dataset
Mean Acc(%)
80.0
75.5
cm (Pure ERM)
79.5
75.0
mm0.5
79.0
m-0.8
I- (Our Methad)
76.5
62.7
a- (Pure ERM)
74.0
a-0.2
76.
2m88
73.
62.6
o-0.8
71.5
Method)
62.
35.5
36.0
37.0
57.5
38.5
56.0
77.0
55
56
58
59
61
73.0
57
58
59
60
61
62
Woed Acc(%)
Worst Accf%)
Wersl Acc(%)
(a) Default dataset
(b) Shoppers dataset
(c) ACS dataset from AZ to MA
Figure 7. The trade-off curve of our method between mean accuracy and worst-group accuracy on three settings which exhibit different kinds of shift.
To better understand how the relationship between average accuracy and worst-group accuracy evolves when optimizing with our score-based weights, we conducted a new experiment and present the Pareto curve in Figure 7.
In fact, by simply combining the loss functions of empirical risk minimization (ERM) and our method, we can achieve balance between both optimization objectives. Specifically, we defined mixed optimization objective as Cmix weeegggeedd + (1 a LERM, where 0 S S 1, and train classification models using Lmix with varying a values. Here, Eweighted represents the loss computed with our boore-based weights, while LERM corresponds to the standard ERM loss. The parameter controls the the influence of our weights in the optimization process. As a increases, the optimization objective aligns more closely with our score-based balancing strategy, while a lower gives greater weight to ERM. By varying a, we can compare model performance and gain insights into how these two optimization objectives interact and influence the performance.
We evaluate the trained models under three scenarios: (1) the Default dataset (Column 6-7 in Table 1), (2) the Shoppers dataset (Column 8-9 in Table 1), and (3) training on data from AZ of the ACS income dataset while testing data from MA (Column 10-11 11l11nn in Table 2).
Figure 7a reveals clear trade-off curve between mean accuracy and worst- group accuracy. Notably, the overall trend of the curves forms a near Pareto frontier, supporting the existence of a trade-off between these two accuracies. Furthermore, compared to ERM's standard optimization objective, our method more effectively improves worst-group accuracy.
In Figure 7b, the curve does not form an exact Pareto frontier. Within a certain range, both worst-case accuracy and mean accuracy exhibit similar trends under the distribution shift present in the Shoppers dataset. This suggests that method can simultaneously enhance both accuracies.
Figure 7c exhibits a curve distinct from Figure 7a. Here, the trade-off between mean and worst-group accuracy is no longer the sole dynamic at play. We attribute this phenomenon to changes in the causal mechanism, specifically Y|h-shift caused by selection bias across different states as highlighted in previous studies. From the perspective of optimization, such shifts violate the assumption independent and identically distributed data, introducing challenges for ERM. Since our method employs a reweighting strategy to balance the dataset, its optimization goal better suited to this setting than ERM to some extent, resulting in improvements both mean and worst-group accuracy. However, when compared to methods explicitly designed for scenarios involving causal mechanism changes, our score-based reweighting falls short in achieving the best mean accuracy.
In summary, the relationship between mean accuracy and worst-grond grosp accuracy can take many forms. The trade-off between these metrics plays a significant role in optimization, but how this trade-off quantitatively evolves is a complex problem. To the best of our knowledge, there are currently no methods the community capable of predicting this trend in advance.
20

---

Latent Score-Based Reweighting for Robust Classification on Imbalanced Tabular Data
However, in cases where trade-off exists, such in the Default dataset, can construct mixed optimization objective combining our loss function and ERM. This allows for control over mean and worst-case accuracy values, demonstrated in Figure 7a, effectively serving as "knob" for balancing these metrics. Ultimately, the optimization process determines how mean and worst-case accuracy interact. Notably, method can outperform another on both metrics if its optimization is better suited to the specific distribution shifts present in the dataset. Our reweighting- bastid optimization objective primarily designed to globally optimize for the worst-group accuracy. It consistently achieves the best worst-case accuracy across nearly all evaluated datasets, Additionally, since our method can address both X-shift and Y-shift through score-based modeling, performs better than some baselines in terms of mean accuracy under specific types of shifts. These factors collectively contribute the improved overall performance observed in Column 12-13 of Table and Column 8-9 and 16-17 of Table 2.
21


Reviews:
[{"review_id": "owKNTiNWIk", "reviewer": "Reviewer_xfS3", "summary": "The paper proposes a latent score-based reweighting framework to improve robustness in machine learning models on tabular data, addressing biases from imbalanced distributions. Unlike existing methods that require prior group labels or focus only on P(Y|X), the approach leverages score-based (diffusion) models to estimate the joint distribution P(X, Y). By using directional similarity of score vectors as a proxy for density, it identifies and upweights underrepresented data regions without relying on unstable raw density estimates. Experiments on tabular datasets under distribution shifts demonstrate improved performance and fairness, making the method broadly applicable in scenarios with unknown biases.\n\n## update after rebuttal\n\nI believe the authors have adequately addressed all of my concerns. I have also reviewed the points raised by the other reviewers and continue to find this work solid, novel, and impactful. There are no remaining concerns—minor or major—from my side, so I have adjusted my score accordingly.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank you for providing helpful comments and positive feedbacks. Below are our responses to your concerns.\n\n> Lack of details about noise preconditioning factor $\\sigma$ and hyper-parameter $\\tau$\n\n**[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. \n\n> Lack of discussion for essential references\n\n**[A2]** We thank the reviewer for highlighting these previous studies. We will incorporate a detailed discussion of these works in the final manuscript.\n\n> Revise the captions in Figure 1\n\n**[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.\n\n---\n\nWe 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.\n\n[1] Karras, Tero, et al. \"Elucidating the design space of diffusion-based generative models.\" *Advances in neural information processing systems* 35 (2022): 26565-26577.\n\n[2] Zhang, Hengrui, et al. \"Mixed-type tabular data synthesis with score-based diffusion in latent space.\" *arXiv preprint arXiv:2310.09656* (2023).", "claims_and_evidence": "Yes, the claims made in the submission are supported by clear and convincing evidence. The paper provides a well-structured empirical evaluation across six diverse datasets, ensuring a comprehensive validation of its method under different types of distribution shifts. The evaluation metrics focus on worst-group accuracy, aligning with the stated goal of improving robustness. The experimental design follows best practices, including multiple runs with different random seeds and comparisons against strong baseline methods such as ERM, DRO variants, JTT, EIIL, and FAM. The results consistently demonstrate the superiority of the proposed method, with at least a 3% improvement in worst-case accuracy across datasets. The inclusion of standard deviations, detailed dataset preprocessing, and hyperparameter choices further strengthens the credibility of the findings. Additionally, the generalization experiments on the ACS Income dataset validate the method’s robustness across different environments, reinforcing its ability to handle distribution shifts effectively. Overall, the empirical evidence strongly supports the paper’s key claims.", "methods_and_evaluation_criteria": "The proposed method is well-suited for unbiased learning, effectively handling covariate and label shifts through score-based diffusion models on latent representations. The probability density proxy via similarity difference offers a practical alternative to exact log-likelihood computation, mitigating extreme weight imbalances. Sample reweighting ensures balanced training without requiring prior knowledge.", "theoretical_claims": "Yes, the problem formulation and overall theoretical claims are correct.", "experimental_designs_or_analyses": "The experimental design appears methodologically sound in leveraging score-based diffusion models for unbiased learning. The use of latent representations via VAEs ensures computational efficiency while capturing meaningful semantics, reducing reliance on raw feature biases. The probability density proxy (SimDiff) is a novel way to approximate sample importance without explicit likelihood estimation, which helps prevent extreme weighting issues. However, it is unclear whether the choice of noise schedule (σ(t)) and temperature parameter (τ) was systematically validated. Sensitivity analysis on these hyperparameters would strengthen the robustness of the method.\nAdditionally, standard benchmarks with known biases (e.g., Waterbirds, CelebA for fairness, UCI datasets for covariate shift) were not included.", "relation_to_broader_scientific_literature": "The key contributions of this paper build upon and extend several important areas in the broader scientific literature, particularly in score-based generative modeling, bias mitigation, and distributional robustness. The foundation of the work relies on score-based diffusion models (Song & Ermon, 2019), which have demonstrated state-of-the-art performance in generative modeling by estimating probability densities through iterative noise perturbation and reconstruction. By integrating latent diffusion techniques (Rombach et al., 2022), the proposed method efficiently models data distributions in a lower-dimensional space, aligning with prior research that has shown latent representations to be more effective for capturing meaningful semantic structures.\n\nIn terms of bias mitigation and robustness, this work shares objectives with reweighting-based debiasing approaches (e.g., importance sampling in fairness-aware learning) but introduces a novel probability density proxy (SimDiff) to achieve unbiased sample weighting. Unlike traditional methods that rely on explicit likelihood estimation—which can be sensitive to extreme values—the similarity-based density estimation method provides a more stable alternative. The paper also builds on findings in fair representation learning (Higgins et al., 2017; Karras et al., 2022), where latent space transformations have been used to decorrelate sensitive attributes from predictions. However, its class-conditional score modeling introduces a new perspective by explicitly accounting for label shifts alongside covariate shifts. This distinction enhances its applicability to fairness-sensitive domains, aligning with recent work on distributionally robust optimization (DRO) and adaptive sample weighting for fairness (Xu et al., 2023).", "essential_references_not_discussed": "Maybe some more recent works on debiasing like SELF (LaBonte et al.) or EVaLS (Ghaznavi et al.) would be nice to mention.", "other_strengths_and_weaknesses": "The paper demonstrates strong originality by combining score-based generative modeling with unbiased learning techniques, presenting a novel way to address distributional shifts without requiring prior assumptions. This integration of latent diffusion models with class-conditional probability estimation is a meaningful contribution, as it removes restrictive assumptions found in traditional bias mitigation methods. The approach is particularly significant for fairness-sensitive applications, as it provides a solution that adapts to both covariate and label shifts.\n\nAdditionally, the visualizations and figures are a notable strength, as they effectively illustrate key concepts such as probability density estimation and reweighting. However, Figure 1’s captions are unclear and difficult to follow, making it harder for readers to interpret the figure’s intended message. Improving clarity in figure descriptions would enhance the overall readability.", "other_comments_or_suggestions": "The captions in Figure 1 are mixed and hard to follow. Revising them for clarity would improve the reader’s understanding of the figure’s significance.", "questions_for_authors": "See previous parts.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5PnciPiiyw", "reviewer": "Reviewer_kxsQ", "summary": "This paper introduces a score-based approach to address data distribution imbalance. First, a variational autoencoder (VAE) is used to transform raw data into a latent representation space, where a diffusion model is applied to learn the joint data distribution. The proposed method estimates relative density using score similarities between noisy data points, mitigating the extreme weight imbalances caused by traditional log-likelihood calculations. A reweighting scheme is then introduced, where each sample is assigned a weight based on its relative density difference, ensuring a more balanced training distribution. Finally, an unbiased classification model is trained using these reweighted samples, and during inference, only the trained encoder and classifier are used for prediction. The proposed method show its effectiveness on several tabular datasets.", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your detailed reviews. Below are our responses to your concerns (**R**) and questions (**A**):\n\n**[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.\n\n**[R2]** We appreciate your suggestions on the undiscussed studies and will incorporate them in our manuscript.\n\n**[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.\n\n**[A2]** Sure. We have conducted additional experiments on three datasets from [3] and [4], as you suggested. The results are listed in [here](https://anonymous.4open.science/r/ICML-rebuttal-9750/new_dataset.md). 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**.\n\n**[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.\n\nHowever, 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.\n\n**[A4]** Nearly all datasets used in our experiments contain categorical features. How we preprocess these features into latent embeddings is detailed in **A1**.\n\n**[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.\n\nRegarding 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.\n\n---\n\nWe 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.\n\n[1] Zhang, Hengrui, et al. \"Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Space.\" *The Twelfth International Conference on Learning Representations*.\n\n[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).\n\n[3] Asuncion, Arthur, and David Newman. \"UCI machine learning repository.\" Nov. 2007.\n\n[4] Gijsbers, Pieter, et al. \"Amlb: an automl benchmark.\" Journal of Machine Learning Research 25.101 (2024).", "claims_and_evidence": "The mathematical formulations and logical progression appear well-structured with no major issues.", "methods_and_evaluation_criteria": "The proposed method appears to be appropriately designed for achieving unbiased learning. However, benchmark tabular datasets are highly diverse, with varying characteristics, and evaluations typically involve a large number of datasets. Since the proposed method has not been evaluated on a large number of tabular datasets, additional validation is needed.", "theoretical_claims": "I did not identify any major issues.", "experimental_designs_or_analyses": "When measuring performance, only three runs were conducted, and it is unclear why the standard deviation was reported separately in a different table. Additionally, no information is provided on dataset statistics, such as the number of categorical features or the number of samples. Furthermore, there is no discussion of time complexity or scalability.", "relation_to_broader_scientific_literature": "If the method is shown to be effective for classification on general tabular datasets, it would have significant potential for broader applications and extensions.", "essential_references_not_discussed": "The paper has omiited literature on tabular classification [1,2,3].\n\n[1] Gorishniy, Yury, et al. \"TabR: Tabular Deep Learning Meets Nearest Neighbors.\" The Twelfth International Conference on Learning Representations.\\\n[2] Holzmüller, David, Léo Grinsztajn, and Ingo Steinwart. \"Better by default: Strong pre-tuned mlps and boosted trees on tabular data.\" Advances in Neural Information Processing Systems 37 (2024): 26577-26658.\\\n[3] Prokhorenkova, Liudmila, et al. \"CatBoost: unbiased boosting with categorical features.\" Advances in neural information processing systems 31 (2018).", "other_strengths_and_weaknesses": "**Strengths**\n\nS1. This paper is well-written.\\\nS2. The motivation is clear.\n\n**Weaknesses**\n\nSee above.", "other_comments_or_suggestions": "No.", "questions_for_authors": "* Does using an autoencoder impose restrictions on the choice of network architecture?\n\n* Could you provide experimental results on more datasets?\n\n* In tabular data classification, there are many datasets where tree-based models outperform neural networks. Do the authors think it is acceptable not to compare the proposed method with tree-based models?\n\n* Do the used datasets not contain categorical features?\n\n* In tabular datasets, where each dataset has vastly different characteristics, I believe that, as shown in [2], comparing performance across various benchmark datasets [4, 5] is necessary to validate the effectiveness of the proposed method. What do you think about this opinion?\n\n[2] Better by Default: Strong Pre-Tuned MLPs and Boosted Trees on Tabular Data\\\n[4] Asuncion, Arthur, and David Newman. \"UCI machine learning repository.\" Nov. 2007,\\\n[5] Gijsbers, Pieter, et al. \"Amlb: an automl benchmark.\" Journal of Machine Learning Research 25.101 (2024): 1-65.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vun30Z92KV", "reviewer": "Reviewer_eedu", "summary": "This paper introduces a latent score-based reweighting framework for improving classification robustness on imbalanced tabular datasets. The approach leverages score-based generative models (diffusion models) to estimate the joint distribution P(X,Y), identifying underrepresented data regions and upweighting samples accordingly. Experiments on six real-world tabular datasets show that the method achieves improved worst-case accuracy under distribution shifts while maintaining competitive mean accuracy compared to baseline robustness-enhancing methods.\n\n## update after rebuttal\n\nI read the author's rebuttal carefully and have updated my evaluation accordingly.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your detailed review and valuable feedback. Below is our concise response:\n\n**[A1]** The explanation is as follows:\n\n- 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.\n- 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. \n\nIn 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.\n\n**[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. \n\n**[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. \n\nIn 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.\n\n**[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](https://anonymous.4open.science/r/ICML-rebuttal-9750/tree_as_base_model.md). 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.\n\n---\n\nWe 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.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "The idea of balancing the data distribution density in feature space makes sense for more robust classification. Conceptually it should help models to focus more on the underrepresented group with lower density in specific feature subspace.", "theoretical_claims": "NA", "experimental_designs_or_analyses": "The experimental designs generally make sense, but I wonder how this method works with tree-based models such as XGBoost and LightGBM, which are known to generally better than neural networks on tabular data.", "relation_to_broader_scientific_literature": "It may serve as a general technique for more robust learning on tabular datasets.", "essential_references_not_discussed": "The references seem appropriate to my best knowledge.", "other_strengths_and_weaknesses": "Strengths:\n1.\tThe paper is well-written and easy to follow.\n2.\tThe proposed framework does not require pre-defined group labels, this can be an advantage in real-world scenarios where the group labels may not be available due to privacy concerns.\n3.\tI like the visualizations that help illustrate the core ideas.\n\nWeaknesses: Please see my questions.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1.\tIt’s not intuitive to me how the method can benefit the learning of the worst group. In my opinion, since the method is upweighting samples in low-density area, then the authors are assuming that the worst group is under-represented group with fewer samples distributed in a low-density feature space? I would appreciate clarification on this.\n2.\tIt’s unclear how the selected datasets contain distribution shifts, please add more detailed description in the dataset section or appendix.\n3.\tAlso, why do shifts matter for robust classification? The proposed method is simply trying to balance the distribution in feature space via reweighting and helps learning in low-density area. It should be beneficial even without train-test shift.\n4.\tFor tabular data, it is known that tree-based ensembles (like XGBoost or LightGBM) are typically better than neural networks and is more commonly used in practical industry scenarios. I also wonder how this technique works well with those classifiers.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces a novel framework for improving robustness in machine learning models trained on imbalanced tabular data. The authors propose using score-based models (diffusion models) to capture the joint data distribution P(X,Y) without requiring prior group labels. Their approach estimates sample density using score-based models. By identifying low-density regions where data is underrepresented, they upweight samples from these regions, aiming to ensure less biased representation during training. The proposed method is tested on several tabular datasets.

Several reviewers highlighted significant strengths of this work. Reviewer xfS3 noted the paper "provides a well-structured empirical evaluation across six diverse datasets" with results that "consistently demonstrate the superiority of the proposed method, with at least a 3% improvement in worst-case accuracy across datasets." The reviewer also highlighted the "novel way to approximate sample importance without explicit likelihood estimation, which helps prevent extreme weighting issues." Reviewer eedu appreciated that the paper is well-written and that the proposed framework does not require pre-defined group labels -- an advantage, since in many applications group labels are not available

Weaknesses were aslo identified by the reviewers. Reviewer kxsQ pointed out that "no information is provided on dataset statistics, such as the number of categorical features or the number of samples" and "there is no discussion of time complexity or scalability." I find these concerns relatively minor and easily addressed in a camera-ready version, particularly since the datasets used are fairly standard.

The main concern raised by kxsQ, however, was about the evaluation, stating "the proposed method has not been evaluated on a large number of tabular datasets." Reviewer eedu went so far as to lower their score, under the suspicion that the result might be cherry picked.

Despite the criticisms by kxsQ, I believe the paper merits acceptance based on its new approach and effective benchmarks -- a sentiment echoed by the other two reviewers. As Reviewer xfS3 noted, the paper does seem to include novel results, and the experimental results consistently show improvements across various datasets. However, I **strongly** encourage the authors to include their code with the final submission -- particularly to assuage the concerns raised by reviewer kxsQ. Failure to do so would significantly weaken the contribution. 

Overall, the paper addresses a significant challenge in machine learning, namely improving robustness over specific subsets of the input space, and provides a practical solution that shows favourable performance on benchmarks.
Decision: Accept (poster)