Beyond Homoscedasticity: Decoupled Uncertainty Optimization for Deep Imbalanced Regression

1 Introduction

• Long-tailed DIR biases optimization toward the abundant, easy-to-fit “Many” labels, reducing reliability on scarce “Few” samples despite the potential importance of rare cases [40].The paper notes that tail categories can have substantial collective influence and individual rare cases may matter in application-specific settings.Source: §1 · p. 1
• In processed AAV2-DIR data, the fraction labeled active drops from approximately 10% to 0.3% as mutation count reaches ≥6, approaching zero beyond 21 mutations.This scarcity makes reliable prediction in sparse regions important for identifying unusual viable variants.Source: §1 · p. 2
• MSE assumes uniform uncertainty, but predictive uncertainty can vary with image quality and label density: at age 80, a high-quality sparse-tail sample has σ≈10, versus σ≈2 for common-age samples.At age 35, clear images yield ages 33-37 (σ≈2), while blurry images yield ages 30-40 (σ≈5).Source: §1 · p. 2
• Homoscedastic methods can underfit uncertain tail samples and entangle cross-label semantics; jointly optimizing mean and variance in NLL can impair mean fitting and cause DIR optimization inertia.Prior uncertainty-aware DIR approaches use probabilistic smoothing or multi-expert aggregation, adding modeling or inference complexity.Source: §1 · p. 2
• DUO models targets as conditional Gaussian distributions, decouples mean-variance optimization with gradient detachment, and uses Bhattacharyya overlap to guide contrastive pairs; it is evaluated across multimedia and biological DIR benchmarks.The paper reports consistent gains on imbalance-aware few-shot metrics.Source: §1 · p. 2

2.1 Imbalanced Classification

• Imbalanced classification methods use reweighting/resampling[8, 14, 18] or boundary and prior calibration[39, 41], but their discrete-label assumptions limit transfer to continuous-label regression.Examples include Class-Balanced Loss[10], Focal Loss[22], LDAM[7], decoupled training[20], Logit Adjustment[23], and Balanced Softmax[28].Source: §2.1 · p. 2

2.2 Imbalanced Regression

• In imbalanced regression, continuous labels and absent class boundaries yield skewed, sparse, noisy tail regions, while most methods intervene at only one level.This leaves the coupling between sample scarcity and predictive uncertainty unmodeled.Source: §2.2 · p. 2
• Input methods SMOTER[34] and SMOGN[4] resample rare-label regions but miss tail feature collapse and prediction bias; output methods DenseWeight[32], BMSE[29], and Dist Loss[25] correct bias without guiding upstream representations.LDS-based reweighting[38] smooths empirical label densities before assigning sample weights.Source: §2.2 · p. 3
• Feature methods LDS/FDS[38], RankSim[13], and ConR[21] preserve continuous-label feature geometry, but do not incorporate predictive uncertainty into feature learning.LDS and FDS share statistics across local label neighborhoods; RankSim aligns relative label-feature ordering, and ConR uses contrastive learning to penalize label-feature inconsistency.Source: §2.2 · p. 3
• VIR[37] quantifies long-tailed uncertainty and UVOTE[17] uses it to select ensemble experts, but neither uses uncertainty to construct representation-space relations.Their uncertainty modeling remains primarily at the output or predictive-distribution level.Source: §2.2 · p. 3
• DUO explicitly models sample-level uncertainty and uses it as unified guidance for optimization and representation learning to improve tail-sample modeling in long-tailed regression.It addresses the gap in jointly modeling sample scarcity, predictive uncertainty, and feature geometry in continuous-label spaces.Source: §2.2 · p. 3

3.1 Problem Formulation

• Tail samples have greater uncertainty and noise than head samples, so assuming constant variance prevents adaptation to sample difficulty and limits tail generalization.Conventional MSE models only the conditional mean and implicitly assumes homoscedasticity.Source: §3.1 · p. 3

3.2 Probabilistic Regression Framework

• DUO models p(y|x) as a Gaussian with a predicted mean and instance-level variance, explicitly representing each sample’s ambiguity.The mean is the distribution’s mode, while variance depends on the input rather than a class-level prior.Source: §3.2 · p. 3

3.3 Network Architecture

• A shared feature backbone feeds parallel mean and uncertainty heads, parameterizing the Gaussian predictive distribution with negligible computational overhead.A positive parameterization ensures a valid standard deviation.Source: §3.3 · pp. 3–4

3.4.1 Gradient Coupling in NLL Loss.

• NLL couples the mean gradient to inverse variance, which can drive variance toward squared residuals and weaken updates for high-error tail samples.At this equilibrium, the mean-gradient magnitude approaches 1/|y−μ|, so the model can reduce loss by increasing variance instead of improving the mean; DUO uses stop-gradient to decouple their optimization.Source: §3.4.1 · p. 4, §3.4.2 · p. 4

3.4.2 Decoupled Objective.

• DUO’s decoupled mean loss uses label-bin scarcity and stop-gradient uncertainty to amplify errors without sending gradients into the variance estimate.The variance loss fits uncertainty to a fixed mean residual while balancing residual fitting against a penalty on unnecessarily large variance; its unique global minimizer makes predicted variance an implicit residual proxy.Source: §3.4.2 · pp. 4–5

3.5.1 Pair Assignment via Distribution Overlap.

• Samples are modeled as Gaussian proxies centered on ground-truth labels with predicted variance; BC overlap supplies adaptive semantic tolerance, letting high-uncertainty tail samples interact with broader semantic neighborhoods despite large label distances.This overcomes limitations of rigid distance-based thresholds.Source: §3.5.1 · p. 5

3.5.2 Distance-Aware Contrastive Objective.

• The distance-aware contrastive objective upweights tail samples and repels feature-similar negatives in proportion to their inverse distribution overlap.Alignment begins only after a warm-up that optimizes the mean and variance losses, allowing uncertainty estimates to stabilize first.Source: §3.5.2 · pp. 5–6

3.6 Overall Objective Function

• The overall objective jointly optimizes mean prediction, variance estimation, and feature alignment, with λ balancing alignment against the other terms.Source: §3.6 · p. 6

4.1 Evaluation protocol and metrics.

• Evaluation covers AgeDB-DIR, IMDB-WIKI-DIR, NYUD2-DIR, and AAV2-DIR, with All, Many, Median, and Few partitions; age and protein tasks use MAE, bMAE, and GM, while depth uses RMSE and δ1.The datasets span age estimation, indoor depth prediction, and AAV2 protein-variant activity prediction.Source: §4 · p. 6, §4.1 · p. 6

4.1.1 Main results for age estimation.

• DUO’s few-shot MAE is slightly higher than LDS on IMDB-WIKI-DIR (22.940 vs. 22.755) and Balanced MSE on AgeDB-DIR (9.917 vs. 9.690), while tail gains outweigh head compromises.AgeDB Many-shot MAE is 6.980 vs. Vanilla at 6.743, while Few-shot MAE is 9.917 vs. 13.381.Source: §4.1.1 · p. 6

4.1.4 Feature visualizations.

• DUO’s t-SNE features are more continuous and less entangled than the compared methods, and its KNN-MAE is lowest in compound-tail and median-shot regions.Vanilla shows central collapse, FDS substantial local overlap, and RankSim an over-compressed line-like manifold; t-SNE evidence is qualitative.Source: Figure 4 · p. 6, §4.1.4 · p. 6

4.1.1 Main results for age estimation.

• DUO achieves the best few-shot bMAE and GM on AgeDB-DIR and IMDB-WIKI-DIR, best NYUD2-DIR bMAE (1.696), and best AAV2-DIR few-shot MAE, bMAE, and GM (4.106, 4.329, and 3.777).The results span scalar age regression, dense depth prediction, and biological sequence activity prediction.Source: Table 1 · pp. 6–7, §4.1.1 · p. 6, §4.1.2 · p. 6

4.2.1 Component-wise ablation of DUO.

• The decoupled objective improves few-shot performance, and full DUO gains further from contrastive alignment; capacity alone is insufficient, while removing weighting or replacing BC degrades few-shot metrics.Controlled variants use matched dual-head capacity.Source: §4.2.1 · p. 7, Table 3 · p. 8

4.2.2 Analysis of the hyper-parameter.

• A moderate contrastive weight λ improves few-shot performance, whereas an excessively large λ harms performance across regions; λ = 1.0 is used thereafter.Small λ provides insufficient structural constraints; excessive λ disrupts the primary regression objective and can over-separate features.Source: Figure 5 · p. 7, §4.2.2 · pp. 7–8

4.1.5 Pairwise comparison.

• Adding DUO to AAV2 baselines improves MAE and bMAE in every pair and most GM values, though RankSim’s GM slightly worsens.This indicates useful but metric-dependent complementarity.Source: §4.1.5 · pp. 6–7, Table 4 · p. 8

5 Conclusion

• DUO models each target as a conditional Gaussian distribution, computationally decouples mean and variance optimization, and refines feature learning through distribution-guided contrastive alignment.Source: §5 · p. 8
• Across AgeDB-DIR, IMDB-WIKI-DIR, NYUD2-DIR, and AAV2-DIR, DUO shows its strongest gains in scarce regions, particularly on imbalance-aware metrics, with modest trade-offs on some overall or many-shot results.The findings support sample-wise uncertainty as an effective signal for robust long-tailed regression.Source: §5 · p. 8

A Proof of Maximum Entropy Optimality for Gaussian Distribution

• Among continuous distributions with fixed mean μ and variance σ², the unique maximum-entropy distribution is Gaussian N(μ, σ²) [4].Source: A Proof of Maximum Entropy Optimality… · p. 10
• The proof uses constrained entropy maximization, deriving an extremum condition and applying the mean and variance constraints to establish the Gaussian form.Source: A Proof of Maximum Entropy Optimality… · pp. 10–11

B.1 Problem Formulation and Notation

• The regression model outputs an input-dependent mean μ(x) and positive variance σ²(x), with residual e(x) defined as y − μ(x).Source: §B.1 · p. 11
• Its single-sample objective is the heteroscedastic Gaussian negative log-likelihood (NLL), analyzed using residual e(x) and variance v = σ²(x).Source: §B.1 · p. 11

B.2 Gradient Computation

• Proposition B.1 frames the first-order partial derivatives of the NLL with respect to mean μ and variance v = σ².Source: §B.2 · p. 11
• A large predicted variance directly suppresses the mean update because its gradient is scaled by 1/σ²(x).Source: §B.2 · p. 12

B.3.1 Gradient Descent (GD)

• When e^2 > v_t, gradient descent strictly increases variance toward the unique stable equilibrium v*=e^2 and converges linearly.Source: §B.3 · p. 12
• Near v*=e^2, the Hessian is positive and bounded, the gradient is Lipschitz, and gradient descent converges linearly when 0 < η < 2e^4 with factor ρ = |1 − η/e^4| < 1.Source: §B.3 · p. 12
• Under the approximate equilibrium σ^2(x) ≈ e(x)^2, hard samples with e^2 ≫ σ^2 rapidly inflate predicted variance to the squared residual’s magnitude.Source: §B.3 · p. 12

Theorem B.3 (Quadratic Convergence under Newton’s Method)

• Newton’s method converges quadratically to v*=e^2 for initial points sufficiently close to the equilibrium.Source: §B.3 · p. 13

B.4 Proof of the Vanishing Mean Update Signal

• At equilibrium, σ² = e² + o(e²) as |e| grows, and the mean update signal for hard samples vanishes.The proof concludes that the mean update signal vanishes for hard samples.Source: §B.4 · p. 13
• To reduce hard-sample NLL, the optimizer inflates predicted variance, suppressing the mean’s learning signal toward zero.As a result, the model becomes almost unable to improve mean predictions for challenging samples.Source: §B.4 · p. 13
• NLL optimization drives predicted variance toward the squared residual, so variance absorbs hard-sample errors during training.The passage states this behavior for both Gradient Descent and Newton’s method.Source: §B.4 · p. 13
• The proof identifies a unique global minimum of L_Var for nonzero residuals and examines local optimization in linear and logarithmic parameterizations.The supplied passage states the proof objective but does not provide further results about these parameterizations.Source: §B.4 · p. 14

C.1 Problem Formulation

• During variance optimization, stop-gradient treats the mean prediction as constant to decouple the two optimization processes.The formulation also specifies β > 0 as a positive scaling coefficient.Source: §C.1 · p. 14
• For a fixed sample with nonzero residual, the formulation sets E to the squared stop-gradient residual and S to predicted variance, simplifying the objective to a function of S.Both E and S are positive.Source: §C.1 · p. 14

C.2.1 Global Unique Minimum

• J(S) has its unique global minimum at S*=E; its derivative changes from negative to positive there, and global convexity is neither claimed nor required.For a fixed mean prediction, L_Var and L_NLL share the same stationary variance.Source: §C.2 · p. 14

C.2.2 Convergence Rate in Linear Space

• When variance inflates to σ²=cE with c≫1, L_NLL suppresses the mean gradient by factor c, whereas L_Var’s mean-gradient magnitude is invariant to variance inflation.L_Var updates the mean head through an independent loss term.Source: §C.2 · p. 14
• Near E, Newton’s method on J(S) converges quadratically; for gradient descent, β rescales local curvature and must be considered jointly with step size.Source: §C.2 · pp. 14–15

C.2.3 Convergence under χ²-Divergence

• Locally, gradient descent on J(S) induces a contraction in χ²-divergence space.Source: §C.2 · p. 15

• The appendix presents this derivation as mathematical verification of the closed-form formulation for adaptive semantic learning.Source: D Closed-form Derivation of the… · p. 15
• The derivation substitutes Gaussian density functions into the overlap integral and completes the square for the exponential terms.Source: D Closed-form Derivation of the… · p. 15

E Asymptotic Analysis of Repulsive Gradients Based on Reciprocal Distribution Overlap

• Under topological hallucination, the relative repulsive gradient grows squared-exponentially with label distance when predicted variances are bounded.The analysis assumes large semantic label distance despite near-identical feature-space clustering and compares the target sample’s gradient with that of an ordinary hard negative.Source: E Asymptotic Analysis of Repulsive… · p. 16
• The reciprocal-overlap weighting penalizes erroneous feature-manifold connections, while high predicted variance suppresses gradient amplification to maintain training stability.For samples with large semantic gaps that are erroneously clustered, the squared-exponential repulsive gradient corrects the distorted manifold; its growth is constrained by predicted variances.Source: E Asymptotic Analysis of Repulsive… · p. 17

F Experiment Details

• DUO is evaluated against representative DIR baselines using splits that follow Yang et al..Training uses inverse empirical-frequency bin weights multiplied by detached instance-level uncertainty in the mean loss; squared error is the training objective, with MAE, bMAE, and GM used for evaluation.Source: F Experiment Details · p. 17

F.1 Age estimation

• Age estimation uses a ResNet-50 encoder with separate fully connected heads for μ and σ, trained for 120 epochs.DUO uses β=1.0, λ_align=1.0, τ=0.07, and δ=0.5 on these benchmarks.Source: §F.1 · p. 17

F.2 Depth estimation

• For NYUD2-DIR depth estimation, μ and σ are predicted pixel-wise, and contrastive sample proxies use spatial means over valid pixels.Depth maps are output at 152x114 resolution, with pixel-wise losses computed over the valid-pixel mask and averaged.Source: §F.2 · p. 18

F.3 Protein estimation

• AAV2-DIR contains 42,329 protein-variant sequences, embedded using mean-pooled ProtBERT residue representations and modeled with a two-layer Transformer encoder.The DUO branch outputs μ and σ; training is configured for up to 200 epochs with inverse-frequency reweighting.Source: §F.3 · p. 18

Controlled component ablation

• Under matched dual-head capacity, pure dual-head NLL does not beat single-head Vanilla, while removing decoupling, dynamic weighting, or BC consistently worsens few-shot metrics.These controls indicate that gains do not come from extra parameters and that each of the three DUO components contributes to tail performance.Source: §F.4 · p. 18, Table S1 · p. 19

Uncertainty and overlap controls

• Replacing learned uncertainty with instantaneous absolute residual removes the variance head and harms most tail metrics; rigid interval-IoU overlap also underperforms the full model.The learned predictor provides a smoother difficulty estimate than a batch-wise residual.Source: §F.4 · p. 18, Table S2 · p. 19

Feature-space smoothness

• On AgeDB-DIR, DUO has the lowest KNN-MAE (k=5) in compound-tail and median-shot regions, supporting the local-continuity interpretation of the visualization.KNN-MAE measures average label discrepancy between each sample and its nearest neighbors in feature space.Source: §F.4 · pp. 18–19, Table S3 · p. 19

Warm-up sensitivity

• AgeDB-DIR results show a stable U-shaped trend around the selected warm-up epoch T_w=15, rather than sensitivity to a single unstable setting.Source: §F.4 · p. 19, Table S4 · p. 19