Mixture Experts with Test-Time Self-Supervised Aggregation for Tabular Imbalanced Regression

Existing techniques for handling imbalanced data mainly focus on imbalanced and long-tailed classification, including rebalancing, logit adjustment, feature enhancement, and ensemble methods. Rebalancing comprises resampling [37, 8, 16, 22, 14, 30] and reweighting [5, 9, 18, 19, 21], both designed to balance classes by the corresponding techniques during the training process. Logit adjustment methods [28, 36, 32, 31, 26] adjust the output logits based on the frequencies of training labels. However, both rebalancing and logit adjustment methods focus on tail classes yet overlook overall class distributions, leading to increased sensitivity to tail classes and consequently, higher model variance [39]. Feature enhancement methods [27, 13] aim to improve the representation of minority classes by integrating features across class distributions, thereby enhancing the model’s ability to recognize underrepresented classes. These methods are tailored for visual data, and thus may not effectively apply to the structured nature of tabular data.

On the other hand, ensemble methods [46, 4, 39] are designed to capture diverse information (e.g., facial characteristics of ages) from imbalanced datasets and effectively aggregate this information through multi-expert approaches. These methods leverage multiple specialized models to address various aspects of the data. Regarding ensemble approaches that are designed to be test-agnostic, LADE [17] utilizes test data distribution as available information to post-adjust model outputs, while SADE [46] introduces a test-time training method for imbalanced classification that adjusts expert weights based on the assumption of discrete target space. However, LADE and SADE were both designed for categorical indices, which cannot be directly adapted to imbalanced regression. On the flip side, our method adopts a region-aware training technique with a regression synthesizer and a self-supervised method to dynamically aggregate models for imbalanced regression. We enhance the representation of few-shot regions through multiple experts and adjust expert contributions through test-time training based on varying test distributions.

Imbalanced regression is less explored compared to classification due to the unclear boundaries of regression targets, which leads to fewer benchmark datasets for imbalanced regression. Existing works on imbalance regression resample data by interpolating features and labels [37] or adding Gaussian noise [3]. These methods of imbalanced regression are directly or indirectly related to SMOTE [6] which was developed for classification tasks that generate synthetic samples for the minority class to balance the dataset. These approaches suffer from a lack of context awareness, and the fact that resampling methods highly rely on statistical properties for synthesizing [34].

[42] curated a deep imbalanced regression (DIR) benchmark and proposed label smoothing using kernel density estimation to smooth and estimate label frequencies for reweighting techniques. [40] improved imbalanced regression by introducing probabilistic smoothing combined with variational inference, which enhances accuracy across underrepresented regions in the image data distribution. For logit adjustment, [33] adjusted the mean squared error (MSE) loss function by incorporating a balancing term, which modifies the standard MSE loss to give more weight to errors from underrepresented (minority) classes or regions in the data distribution. However, these imbalanced regression methods assume balanced test distributions, making them inadequate for addressing test-agnostic evaluation (e.g., imbalanced distributions at test time).

We focus on tabular imbalance regression tasks: consider a dataset $\mathcal{D}=\{(x_{i},y_{i})\}_{i=1}^{N}$, where each $x_{i}\in\mathbb{R}^{d}$ denotes the input features, and $y_{i}\in\mathbb{R}$ represents the associated continuous target. To address the challenge of imbalanced regression, we follow the definition of [42] to define the label space $Y$ as a partitioned set of $B$ non-overlapping intervals or bins, such that $Y=\bigcup_{b=1}^{B}[y_{b-1},y_{b})$. Note that $B$ is a finite set of indices $\{1,\ldots,B\}$ corresponding to these bins, and $\delta_{y}=y_{b+1}-y_{b}$ defines the resolution of interest within the target variable.

Inspired by the successful advancements of utilizing the multi-expert techniques in vision [17, 4, 46, 7] and language [15, 20] applications, where each expert model is specifically skilled in handling certain classes or target regions, we introduced region-aware mixture experts by extending expert models dedicated to specific distributions of imbalanced regression tasks. Specifically, we train our region experts by synthesizing different regions of data based on the distribution information from $D_{T}$.

The process of training target region experts is depicted in the left part of Figure 2. To mitigate the imbalance by generating synthetic samples in underrepresented regions, the training set $D_{T}$ undergoes a synthesizer designed based on SMOGN [3], which includes a minority oversampling technique tailored for tabular regression. The synthesizer automatically defines rare regions for over-sampling within the entire target space $Y$ of $D_{T}$. This results in a dataset, denoted as $D_{S}$, which is more balanced compared to the original data $D_{T}$:

The original data with imbalance distribution may consist of multiple complex distributions where a single model might be biased towards the nearest majority regions [10]. To capture this distributional information and enable multiple models to correspond accordingly, a Gaussian Mixture Model (GMM) is applied to $D_{T}$ to capture the distinct distributions within the original training data. We divide the dataset $D_{T}$ with its labels $Y$ into multiple datasets by fitting GMM on the original dataset and its label. Formally, the GMM is defined as:

where $N$ is the number of Gaussian components, $\pi_{n}$ are the mixture weights, and $\mathcal{N}(Y\mid\mu_{n},\Sigma_{n})$ is the Gaussian distribution with mean $\mu_{n}$ and variance $\sigma_{n}$.

Afterwards, each data point $(x_{i},y_{i})\in D_{T}$ is assigned to a cluster $n$ based on the posterior probability:

Generally, the original dataset $D_{T}$ can be divided into $N$ subsets $D_{T_{n}}$, where each subset contains data points assigned to the $n$-th Gaussian component:

GMM provides a statistical approach to model the complex underlying structure of $D_{T}$, while it fits $D_{T}$ by its label $Y$ into $n$ sets of data, each denoted as $D_{T_{n}}$, where $n$ is decided by the Akaike Information Criterion score [23] for better training and generalization results. With the Gaussian components, we utilize a second synthesizing on $D_{S}$ based on the statistical information of each $D_{T_{n}}$, denoted as $\text{Synthesizer}_{ex}$ as in Figure 2. $\text{Synthesizer}_{ex}$ synthesizes data in the same manner as Synthesizer, except that the over-sampled regions are replaced with statistical information. Specifically, we over-sample on a specified target range of $D_{S}$, and the range is decided by the label mean $\mu_{n}$ and standard deviation $\sigma_{n}$ of each component $D_{T_{n}}$ with an adjustable hyperparameter $\alpha$ multiplied on $\sigma_{n}$:

The outcome of this training step is a collection of datasets after over-sampling a specified target range of $D_{S}$, denoted as $D_{S}^{n}$. The final phase involves the development of multiple expert models, represented as $Expert_{n}$. Each expert is trained on its corresponding training data $D_{S}^{n}$ with the MSE loss, ensuring that every model develops a specialized understanding of a particular target region.

After training region-specific experts, our next objective is to aggregate these expert models without prior knowledge of the test distribution and labels, allowing for effective adaptation to test distributions with varying degrees of imbalance. To that end, our intuitions are that 1) there exists a positive correlation between expertise and prediction stability, i.e., when predicting different views of samples, experts have higher prediction similarity of samples from their favorable regions, and 2) the principle that a model with more expertise in a particular area should assume a more significant role and consequently carry greater weight after the aggregation.

As shown in algorithm 1, we aggregate the trained region experts by minimizing the continuous output gap of different views of a test sample since experts proficient in specific regions should demonstrate enhanced stability in their predictions. Hence, Continuous Prediction Gap Minimization is proposed to ensure that the contribution of experts is effectively adapted to test distributions with varying degrees of imbalance.

For the baseline methods and MATI, we used TabNetRegressor [2] as our base model for fair comparisons. We trained each method for a maximum of 200 epochs, with an early stopping set of 20 epochs. The model size ranged from 3 to 5 for each method, and the results are the average of three different seeds. For some detailed hyperparameters, we set $\alpha$ from 1 to 2 for the synthesizer. To ensure fairness, we fix the number of epochs and the corruption ratio for test-time training in MATI. The number of epochs is set between 20 and 50, depending on the test set size, while the corruption ratio is fixed at 0.1.

Tables 1 and 2 present the performance of MATI and the baselines on the four datasets. We summarize the observations as follows: 1) Superior MATI Performance on Balanced Test Set: Quantitatively, MATI outperforms all baselines on the balanced test set, achieving at least 0.339 and 1.021 lower MAPE in the CHW and KHH house price datasets, respectively. It also surpasses baselines by at least 1.05% and 8.71% in RMSE on the Abalone and Bike Sharing datasets. Despite these baselines focusing on the performance of balanced test distributions, the comparisons with our MATI not only illustrate the need for evaluating various distributions but also reveal the importance of considering different distributions via the corresponding experts. 2) The Advantage of Test-Time Aggregation: We can observe that imbalanced regression baselines (i.e., RRT+LDS) perform well only on balanced evaluation protocols, while MATI excels across different distributions. Vanilla is only superior in many-shot regions which share the same distribution as the training sets, but is deleterious in few-shot regions due to the lack of imbalanced techniques. In contrast, MATI equipped with Test-Time Self-Supervised Expert Aggregation adjusts expert weights for adapting to test distributions. This allows MATI to perform slightly worse than Vanilla on the normal test set, but to achieve optimal results on both the balanced and inverse test sets. This highlights the effectiveness of MATI’s test-time aggregation in handling to varying test distribution.

To assess the proficiency of region experts and validate the test-time aggregation mechanism of MATI, we conducted experiments on the Bike Sharing and CHW datasets, both of which exhibit highly right-skewed distributions. For these datasets, we employed three Gaussian components, sorting experts by their means to define Region 1, Region 2, and Region 3, handled by $Expert_{1}$, $Expert_{2}$, and $Expert_{3}$, respectively. Region-specific test sets were created by sampling based on Gaussian means and standard deviations, enabling performance evaluation of each expert in its corresponding region.Table 3 confirms that each expert excels within its designated region, validating the effectiveness of Gaussian Mixture Synthesizing in modeling region-specific characteristics. Furthermore, MATI’s dynamic weight adjustment during test-time aggregation effectively aligns expert contributions to the test distribution. Table 4 illustrates that $Expert_{1}$, optimized for Region 1, dominates on right-skewed normal test sets; $Expert_{2}$, representing Region 2, is weighted higher on balanced test sets; and $Expert_{3}$, trained for Region 3, is prioritized on left-skewed inverse test sets. These findings demonstrate MATI’s ability to adaptively allocate expert weights, ensuring optimal performance across diverse test scenarios.

To verify that the perturbation method does not generate excessively large or dissimilar samples, thereby affecting the performance, we closely examined the changes in performance with varying perturbation ratios (Algorithm 1, line 5) from 0.1 to 0.7. As illustrated in Figure 4, the performance curves for the four datasets remain relatively flat and stable at lower perturbation ratios. However, when the ratio exceeds 0.4, there is a noticeable and abrupt increase in the curves. This finding indicates that a slight perturbation of the input features allows the perturbation method to generate proper and similar samples for aggregation.