Improve Representation for Imbalanced Regression through Geometric Constraints

TL;DR

This work tackles the problem of learning uniform representations for deep imbalanced regression (DIR), where labels are continuous and ordered. It introduces geometry-inspired losses—enveloping loss to encourage a latent trace to envelop the unit hypersphere $S^{n-1}$ and homogeneity loss to promote smooth, evenly spaced representations along the trace—and integrates them via Surrogate-driven Representation Learning (SRL). SRL constructs a surrogate, centroid-based representation across all label bins and pairs it with a contrastive objective, enabling global geometric regularization; the authors further introduce Imbalanced Operator Learning (IOL) as a practical benchmark. Across real-world regression tasks and IOL, SRL improves accuracy, especially in few-shot regions, and yields gains when combined with existing DIR methods, all with manageable computational cost. This geometry-based approach offers a principled path to robust representations under label imbalance and expands DIR evaluation through IOL.

Abstract

In representation learning, uniformity refers to the uniform feature distribution in the latent space (i.e., unit hypersphere). Previous work has shown that improving uniformity contributes to the learning of under-represented classes. However, most of the previous work focused on classification; the representation space of imbalanced regression remains unexplored. Classification-based methods are not suitable for regression tasks because they cluster features into distinct groups without considering the continuous and ordered nature essential for regression. In a geometric aspect, we uniquely focus on ensuring uniformity in the latent space for imbalanced regression through two key losses: enveloping and homogeneity. The enveloping loss encourages the induced trace to uniformly occupy the surface of a hypersphere, while the homogeneity loss ensures smoothness, with representations evenly spaced at consistent intervals. Our method integrates these geometric principles into the data representations via a Surrogate-driven Representation Learning (SRL) framework. Experiments with real-world regression and operator learning tasks highlight the importance of uniformity in imbalanced regression and validate the efficacy of our geometry-based loss functions.

Figures (9)

• Figure 1: 2D feature space of vanilla baseline and ours from UCI-Airfoil asuncion2007uci. The vanilla feature space lacks uniformity and is dominated by samples from the Many-shot region. In contrast, our approach achieves a more uniform distribution over the feature space, improving the performance, especially in the Medium and Few-shot regions. (For visualization purposes, we curated the dataset to ensure equal partitions across the three regions.)
• Figure 2: t-SNE visualization van2008visualizing of feature comparison. The first row corresponds to the original UCI Airfoil Dataset asuncion2007uci, while the second row corresponds to its curated version, with an additional few-shot region in the middle of the label range. Colored arrows point to the few-shot regions and their corresponding positions in the feature distributions. We evaluate feature distributions using: MSE Loss (Baseline), SRL without uniformity loss (w/o $\mathcal{L}_{\text{env}}$), SRL without homogeneity loss (w/o $\mathcal{L}_{\text{homo}}$), and complete SRL (ours). The baseline leads to feature collapse to many-shot regions and inadequate distinction of few-shot samples. In w/o $\mathcal{L}_{\text{env}}$, features collapse into a trivial shape, not fully utilizing the feature space. In w/o $\mathcal{L}_{\text{homo}}$, features spread out along the trace. Different from the previous ones, our SRL uniformly and smoothly "fills" the feature space.
• Figure 3: 2D schematic overview of two geometric losses. The arrow indicates the improvement of the loss function. Enveloping loss encourages the representations to fill the latent space, and homogeneity loss encourages the smoothness and even distribution of the representations along the trace.
• Figure 4: Overview of Surrogate-driven Representation Learning (SRL). (1) Every mini-batch is encoded to the latent space. Some bins may not be present in the current batch. To address this, (2) it takes centroids corresponding to the missing bins from the previous epoch. These stored centroids are used to "re-fill" the missing bins in the current batch. (3) Average the representations for bins that appear multiple times, creating centroids for these bins. This surrogate, containing a representation for the full label range, allows for the effective application of geometric loss across all bins. (4) Loss calculation based on the surrogate. (5) Update the surrogate in memory to ensure enveloping and homogeneity. The training of the first epoch is driven by MSE loss only.
• Figure 5: SRL performance gain compared to VANILLA across age ranges on AgeDB-DIR. The gray histogram in the background shows the distribution of samples across age groups. SRL substantially improves the performance on the medium-shot and few-shot regions (age $<$ 20 and $>$ 70).

Theorems & Definitions (2)

• Theorem 1
• proof