Source-linked AI summary

Beyond Homoscedasticity: Decoupled Uncertainty Optimization for Deep Imbalanced Regression

Juncheng Zhou, Jiaxi Lu, Weijing Zeng, Zhong Li, Hao Qi, Jingsong Cui

arXiv:2609.04995v1cs.LG

TL;DR

Deep imbalanced regression involves scarce tail samples and instance-level uncertainty, while existing approaches often rely on deterministic mappings or coupled heteroscedastic optimization. DUO models conditional Gaussian targets, decouples mean–variance optimization, and uses distribution-guided contrastive learning; it reports strongest imbalance-aware gains across visual and biological benchmarks, with modest trade-offs on some overall or many-shot results.

  • Problem

    Deep imbalanced regression is biased toward abundant labels, while deterministic methods overlook instance-level heteroscedasticity and coupled NLL can weaken learning signals for hard tail samples.

  • Method

    DUO models targets with conditional Gaussian distributions, decouples mean and variance optimization, and refines features through distribution-guided contrastive learning.

  • Results

    DUO shows strongest gains on imbalance-aware few-shot metrics across AgeDB-DIR, IMDB-WIKI-DIR, NYUD2-DIR, and AAV2-DIR, with best few-shot results including AAV2-DIR MAE 4.106, bMAE 4.329, and GM 3.777.

  • Takeaways & Limitations

    The findings support sample-wise uncertainty as an effective signal for robust long-tailed regression beyond visual tasks.

  • Takeaways & Limitations

    Prior imbalanced-classification methods rely on discrete-label assumptions, while DUO’s probabilistic formulation assumes instance-level uncertainty can be represented through σ^2(x).

Abstract

from arXiv · show

Deep Imbalanced Regression (DIR) is pervasive in continuous prediction tasks across diverse modalities, such as age estimation, depth prediction, and protein mutation activity prediction, where label-scarce tail samples often carry higher practical value. However, most existing methods still learn deterministic point mappings under mean squared error or its simple variants, implicitly assuming a uniform uncertainty level across all samples and thereby overlooking the instance-wise heteroscedasticity that is widespread in long-tailed data. We further point out that even heteroscedastic negative log-likelihood suffers from a gradient coupling issue, which, under DIR scenarios, weakens the learning signal of hard tail samples and leads to optimization inertia as well as tail underfitting. To address this, we propose DUO, an uncertainty-aware long-tailed regression framework. Specifically, the proposed method models the regression target as a conditional Gaussian distribution to explicitly characterize instance-level predictive uncertainty, and transforms uncertainty into a dynamic enhancement signal for tail samples through decoupled mean-variance optimization. Furthermore, we design a distribution-guided contrastive learning mechanism that adaptively constructs positive and negative pairs based on the overlap between sample distributions, thereby alleviating feature looseness and cross-label semantic entanglement. Across visual and biological DIR benchmarks, DUO achieves the best few-shot bMAE and GM on IMDB-WIKI-DIR, AgeDB-DIR, and AAV2-DIR while remaining competitive on few-shot MAE.

1 Introduction

Deep imbalanced regression is dominated by scarce, uncertain tail samples, while common losses assume uniform uncertainty and can weaken learning for difficult cases. DUO addresses this by modeling instance-level uncertainty, decoupling mean-variance optimization, and using distribution-guided contrastive learning.

  • Motivation: Long-tailed continuous prediction tasks overfit abundant, easy-to-fit samples, producing systematically less reliable predictions for scarce tail samples.Applications include facial age estimation, depth estimation, and protein mutation activity prediction.
  • Motivation: MSE-based regression implicitly assumes global homoscedasticity, assigning every sample the same prediction uncertainty.This assumption limits the loss’s ability to reflect sample-wise uncertainty.
  • Motivation: Predictive uncertainty varies with image quality and label density, including high uncertainty for sparse-tail samples even when their inputs are high quality.The examples contrast sharp age distributions with blurry inputs and sparse age regions.
  • Limitations of existing methods: Heteroscedastic NLL can couple mean and variance optimization, weakening updates for high-error tail samples and encouraging optimization inertia.Prior uncertainty-aware methods address the problem with probabilistic smoothing or multi-expert aggregation, which add modeling or inference complexity.
  • DUO: DUO models targets as conditional Gaussian distributions and uses decoupled mean-variance optimization plus distribution-guided contrastive learning.The contrastive module uses distribution overlap to construct pairs and address tail feature looseness and semantic entanglement.
  • Contributions: DUO formulates long-tailed regression probabilistically, analyzes NLL gradient coupling, and validates consistent gains on imbalance-aware few-shot metrics.The framework is evaluated across multimedia and biological DIR benchmarks.

2 Related Work

Existing imbalanced-regression methods address data, feature, or output imbalance separately, while classification methods rely on discrete labels that do not transfer directly to continuous label spaces. DUO instead uses sample-level uncertainty as a unified guide for optimization and representation learning.

  • Imbalanced classification: Classification imbalance methods rely on discrete classes, priors, or decision boundaries, limiting direct transfer to continuous imbalanced regression.Continuous regression couples sample scarcity, predictive uncertainty, and feature geometry without explicit class boundaries.
  • Imbalanced regression: Imbalanced regression involves skewed continuous labels with sparse, noisy tails, making many classification methods inapplicable.Its challenges exceed class-frequency imbalance because continuous labels lack explicit class boundaries.
  • Existing regression methods: Input-level methods rebalance data but do not directly address tail feature collapse or prediction bias during optimization.SMOTER and SMOGN oversample rare regions and undersample head regions.
  • Existing regression methods: Feature-level methods preserve local or global label-feature structure, but they primarily target representation geometry rather than explicitly modeling uncertainty.Examples include LDS, FDS, and RankSim.
  • Existing regression methods: Output-level methods correct prediction bias through density weighting, balanced losses, or distribution penalties but do not use uncertainty to guide representation learning.Examples include DenseWeight, LDS-based reweighting, Balanced MSE, and Dist Loss.
  • Uncertainty-aware regression: Recent uncertainty-aware methods quantify uncertainty mainly at the output level, leaving uncertainty-guided feature learning under-explored.VIR uses probabilistic neighborhood smoothing, while UVOTE selects among ensemble experts.
  • DUO: DUO explicitly models the coupling among scarcity, uncertainty, and feature geometry, using uncertainty to guide both optimization and representation learning.This unified guidance is intended to improve tail-sample modeling in long-tailed regression.

3 Method

DUO reformulates regression as conditional Gaussian modeling with instance-specific variance, then decouples mean and variance optimization to avoid NLL’s weak tail gradients. It also uses uncertainty-aware distribution overlap to construct contrastive relationships for tail representation learning.

  • Probabilistic regression: Conventional MSE learns only the conditional mean and implicitly assumes constant variance, although tail samples generally have greater uncertainty and noise.Ignoring this variance limits dynamic adaptation to sample difficulty and tail generalization.
  • Probabilistic regression: DUO reformulates regression as modeling the conditional distribution p(y|x), using a Gaussian to represent continuous-label probability density.The Gaussian formulation models the target with mean and variance while retaining the regression prediction as the distribution mode.
  • Probabilistic regression: Instance-level variance captures ambiguity arising from each input, such as occlusion or noise, rather than imposing a class-level uncertainty prior.This provides a sample-specific measure of predictive uncertainty.
  • Network architecture: A shared encoder with mean and uncertainty heads parameterizes p(y|x) as a Gaussian with predicted mean and variance.The architecture extends standard backbones with a dual-branch prediction structure.
  • Gradient coupling: Standard NLL scales the mean gradient by inverse variance, so large-error samples can receive vanishing updates while variance increases instead of mean accuracy improving.This coupling is especially harmful for tail samples, which typically have larger residuals.
  • Decoupled objective: DUO uses stop-gradient to computationally decouple mean and variance optimization and amplify uncertainty-associated mean gradients rather than suppress them.The detached uncertainty acts as a forward weight without reintroducing backward optimization coupling.
  • Decoupled objective: The variance objective fits residuals while regularizing predicted variance, with its stationary point matching the squared prediction residual.The resulting predicted uncertainty serves as an implicit proxy for empirical prediction error.
  • Distribution-guided contrastive learning: Distribution-guided contrastive learning uses predicted-variance Gaussian proxies and distribution overlap to form adaptive semantic pairs instead of rigid label-distance pairs.The approach allows hard samples to interact with a broader set of semantic neighbors during feature-manifold alignment.

4 Experiment

DUO is evaluated across visual and biological deep imbalanced regression benchmarks using shot-based partitions and imbalance-aware metrics. It improves tail-region performance, with ablations and visualizations supporting contributions from decoupling, uncertainty weighting, and contrastive alignment.

  • Evaluation setup: DUO is evaluated on AgeDB-DIR, IMDB-WIKI-DIR, NYUD2-DIR, and AAV2-DIR across All, Many, Median, and Few shot partitions.Age, depth, and protein mutation activity prediction are represented.
  • Main results: DUO obtains the best few-shot bMAE and GM on both AgeDB-DIR and IMDB-WIKI-DIR, while few-shot MAE remains competitive.Its few-shot MAE is slightly higher than LDS on IMDB-WIKI-DIR and Balanced MSE on AgeDB-DIR.
  • Main results: 1.696: DUO achieves the best few-shot bMAE on NYUD2-DIR, extending its tail-region performance to dense depth prediction.The NYUD2 comparison concerns structured regression with high-dimensional depth maps.
  • Main results: 4.106, 4.329, and 3.777: DUO achieves the best few-shot MAE, bMAE, and GM on AAV2-DIR, respectively.The result extends few-shot advantages beyond visual regression to biological sequence activity prediction.
  • Feature visualization: DUO produces a more continuous and less entangled feature layout than Vanilla, FDS, and RankSim, while KNN-MAE supports improved local smoothness.The t-SNE comparison is qualitative, whereas KNN-MAE provides quantitative support in compound-tail and median-shot regions.
  • Ablation studies: Removing weighting or replacing BC degrades few-shot metrics, showing that capacity alone is insufficient and the objective components contribute to tail performance.The decoupled objective improves few-shot performance, with further gains from contrastive alignment.
  • Ablation studies: Moderate contrastive weighting improves few-shot performance, whereas excessively large λ harms overall performance by disrupting the primary regression objective.Small λ provides insufficient contrastive constraints and yields limited few-shot gains.

5 Conclusion

The paper presents DUO as an uncertainty-aware framework for deep imbalanced regression that decouples mean and variance optimization and uses distribution-guided contrastive alignment. Across four benchmarks, it reports strongest gains in scarce regions, especially on imbalance-aware metrics, with modest trade-offs in some overall or many-shot results.

  • Conclusion: DUO models targets as conditional Gaussian distributions, decouples mean and variance optimization, and refines features through distribution-guided contrastive alignment.These are the framework’s central design elements.
  • Conclusion: Across AgeDB-DIR, IMDB-WIKI-DIR, NYUD2-DIR, and AAV2-DIR, DUO shows its strongest gains in scarce regions, particularly on imbalance-aware metrics.The conclusion also reports modest trade-offs on some overall or many-shot results.
  • Conclusion: The findings support sample-wise uncertainty as an effective signal for robust long-tailed regression.This conclusion is stated within the evaluated deep imbalanced regression setting.

Supplementary Material for “Beyond Homoscedasticity: Decoupled Uncertainty Optimization for Deep Imbalanced Regression”

The supplementary material identifies the paper and its authors from Wuhan University and Tianjin University. Juncheng Zhou and Jiaxi Lu are marked as equal contributors.

  • Authors and affiliations: Juncheng Zhou, Jiaxi Lu, Weijing Zeng, Zhong Li, Hao Qi, and Jingsong Cui are listed as authors from Wuhan University or Tianjin University.The affiliations span cyber science and engineering, mathematics and statistics, and synthetic biology and biomanufacturing.

A Proof of Maximum Entropy Optimality for Gaussian Distribution

This supplementary section justifies the Gaussian target model through maximum-entropy reasoning and analyzes gradient behavior in heteroscedastic Gaussian regression. It establishes Gaussian optimality under fixed mean and variance and shows that variance can suppress mean updates through NLL coupling.

  • Maximum-entropy optimality: Among continuous distributions with fixed mean μ and variance σ^2, the Gaussian uniquely maximizes differential entropy.The theorem frames the Gaussian as the minimum-inductive-bias probabilistic form under these moment constraints.
  • Maximum-entropy optimality: The proof constructs a constrained Lagrangian and derives the Gaussian form by applying the normalization, mean, and variance constraints.The derivation introduces the exponential quadratic density and requires a negative quadratic coefficient for integrability.
  • Maximum-entropy optimality: The derivation rigorously establishes the optimal probabilistic form for instance-level variance modeling from an information-theoretic perspective.The result is presented as the conclusion of the maximum-entropy proof.
  • NLL gradient analysis: In heteroscedastic Gaussian regression, the model outputs a mean μ(x) and positive variance σ^2(x), with NLL analyzed using residual e(x)=y−μ(x) and variance variable v=σ^2(x).The supplementary analysis studies partial derivatives with respect to the mean and variance separately.
  • NLL gradient analysis: A large σ^2(x) directly suppresses the mean update because the mean gradient is scaled by 1/σ^2(x).This is the gradient-coupling mechanism analyzed for the heteroscedastic NLL.

B.3 Convergence of Standard Iterative Methods to a Unique Optimum

Standard iterative optimization drives the predicted variance toward the squared residual for hard samples. This variance inflation suppresses the mean-learning signal, leaving challenging samples underfit.

  • When the squared residual exceeds the current variance, the variance strictly increases toward the unique stable equilibrium v*=e^2.
  • For large residuals, the equilibrium variance satisfies σ^2=e^2+o(e^2), so the mean update signal vanishes asymptotically.
  • Variance inflation reduces the mean gradient for hard samples and can make their mean predictions nearly stop improving.The optimizer lowers NLL by inflating uncertainty, simultaneously suppressing the mean update signal.
  • NLL optimization drives the predicted variance toward the squared residual for hard samples under both gradient descent and Newton’s method.Gradient descent converges linearly, whereas Newton’s method converges quadratically near the equilibrium.
  • The variance objective uses stop-gradient on the mean prediction to decouple variance optimization from mean optimization.The formulation includes a positive scaling coefficient β.

C.2 Convergence Analysis in Linear Space

The variance-only objective has the same stationary variance as NLL for a fixed mean, but its gradient dynamics avoid suppressing the mean update when variance inflates.

  • The variance objective has a unique global minimum at S*=E, although its second derivative changes sign and global convexity is not claimed.The objective decreases up to E and increases thereafter.
  • NLL suppresses the mean gradient by a factor c when σ^2=cE with c≫1, whereas the decoupled variance formulation keeps the mean update independent.
  • At S=E, the variance objective’s Hessian is β times the NLL Hessian, so β must be considered jointly with the optimization step size.
  • Newton’s method achieves quadratic convergence on the variance objective when initialized sufficiently close to E.

D Closed-form Derivation of the Bhattacharyya Coefficient for Gaussian Distributions

The appendix derives a closed-form Bhattacharyya Coefficient for Gaussian distributions by substituting Gaussian densities into the overlap integral and completing the square.

  • The derivation obtains the Gaussian Bhattacharyya Coefficient by evaluating the overlap integral after completing the square in the exponential terms.
  • The resulting closed-form formulation provides the mathematical basis for the adaptive semantic mechanism.

E Asymptotic Analysis of Repulsive Gradients Based on Reciprocal Distribution Overlap

The reciprocal-overlap weighting mechanism amplifies repulsion between semantically distant samples that are erroneously clustered, while uncertainty limits amplification for noisy samples. The section also describes benchmark setup and implementation across visual and protein DIR tasks.

  • E Asymptotic Analysis of Repulsive Gradients Based on Reciprocal Distribution Overlap: The repulsive gradient grows squared-exponentially with label distance under bounded predicted variances and cosine similarities.
  • E Asymptotic Analysis of Repulsive Gradients Based on Reciprocal Distribution Overlap: When semantically distant samples are erroneously clustered, the squared-exponential repulsive gradient penalizes and corrects the distorted feature manifold.The mechanism targets topological hallucination, where large label distance coexists with high feature similarity.
  • E Asymptotic Analysis of Repulsive Gradients Based on Reciprocal Distribution Overlap: High-uncertainty samples receive naturally suppressed gradient amplification because predicted variances appear in the denominator.This suppression is intended to prevent gradient explosion and maintain training stability.
  • Implementation and Evaluation Setup: DUO combines inverse-frequency scarcity weighting with detached instance-level uncertainty in the mean objective.
  • Implementation and Evaluation Setup: The evaluation includes visual DIR benchmarks and AAV2-DIR, using task-specific encoders while retaining Gaussian mean and variance outputs.The depth setup computes distribution proxies from spatially averaged valid pixels, while the protein branch outputs μ and σ.

F.4 Additional Analyses from the Rebuttal

Additional analyses show that DUO’s tail-performance gains depend on decoupled optimization, uncertainty weighting, and distribution-guided contrastive assignment rather than capacity alone. The method also improves feature-space smoothness and remains stable around the selected warm-up setting.

  • Controlled component ablation: Matched-capacity controls show that pure dual-head NLL does not improve over single-head Vanilla, so gains do not arise from extra parameters.Disabling decoupling, dynamic weighting, or BC assignment consistently degrades few-shot metrics.
  • Uncertainty and overlap controls: Replacing learned uncertainty with instantaneous absolute residuals degrades most tail metrics, indicating that the learned predictor supplies a smoother difficulty estimate.The comparison removes the variance head and contrasts learned uncertainty with batch-wise residuals.
  • Uncertainty and overlap controls: Replacing BC with interval IoU reduces Gaussian distributions to rigid intervals, testing whether distributional overlap modeling contributes beyond interval-based matching.
  • Feature-space smoothness: DUO achieves the lowest KNN-MAE in both compound-tail and median-shot regions on AgeDB-DIR, supporting local continuity in feature space.KNN-MAE uses k=5 nearest neighbors and measures average label discrepancy.
  • Warm-up sensitivity: Warm-up sensitivity follows a stable U-shaped trend around the selected T_w=15 on AgeDB-DIR rather than depending on one unstable setting.
Loading 2609.04995v1…