PaperlayerView exact revision on arXiv

Source-linked AI summary

Balanced MSE for Imbalanced Visual Regression

Jiawei Ren, Mingyuan Zhang, Cunjun Yu, Ziwei Liu

arXiv:2203.16427v1cs.CV

TL;DR

Imbalanced regression involves continuous labels under skewed training distributions, where MSE can be ineffective and rare-label performance suffers. The paper proposes Balanced MSE, which uses the training label distribution and multiple implementations to restore balanced predictions, and reports strong results across synthetic, uni-dimensional, and multi-dimensional benchmarks. The method also addresses scenarios without prior knowledge of the training label distribution, while remaining limited to bias arising from label imbalance.

  • Problem

    Imbalanced regression concerns continuous, potentially high-dimensional labels under skewed training distributions, and existing approaches remain insufficient for this setting.

  • Method

    Balanced MSE revisits MSE statistically and uses the training label distribution to convert balanced conditional predictions into a trainable objective, with implementations including online distribution estimation.

  • Results

    Balanced MSE outperforms existing methods across uni- and multi-dimensional imbalanced regression benchmarks, including age estimation, depth estimation, and IHMR.

  • Takeaways & Limitations

    Balanced MSE provides a general solution for high-dimensional imbalanced regression and supports real-world scenarios that lack prior training-label distribution knowledge.

  • Takeaways & Limitations

    The method addresses bias from imbalanced label distributions, but other dataset biases may remain and can still produce negative social impacts.

Abstract

from arXiv · show

Data imbalance exists ubiquitously in real-world visual regressions, e.g., age estimation and pose estimation, hurting the model's generalizability and fairness. Thus, imbalanced regression gains increasing research attention recently. Compared to imbalanced classification, imbalanced regression focuses on continuous labels, which can be boundless and high-dimensional and hence more challenging. In this work, we identify that the widely used Mean Square Error (MSE) loss function can be ineffective in imbalanced regression. We revisit MSE from a statistical view and propose a novel loss function, Balanced MSE, to accommodate the imbalanced training label distribution. We further design multiple implementations of Balanced MSE to tackle different real-world scenarios, particularly including the one that requires no prior knowledge about the training label distribution. Moreover, to the best of our knowledge, Balanced MSE is the first general solution to high-dimensional imbalanced regression. Extensive experiments on both synthetic and three real-world benchmarks demonstrate the effectiveness of Balanced MSE.

1. Introduction

Imbalanced regression addresses continuous, potentially high-dimensional labels under skewed training distributions, where standard MSE can produce inferior rare-label performance. The paper proposes Balanced MSE, multiple implementations, and benchmarks showing strong performance across uni- and multi-dimensional settings.

  • Continuous labels can be boundless and high-dimensional, making imbalanced regression more challenging than imbalanced classification.
  • Imbalanced regression remains under-explored and lacks an effective approach for handling skewed label distributions.
  • MSE can carry training-label imbalance into predictions, causing inferior performance on rare labels.
  • Balanced MSE uses the training label distribution to statistically restore balanced predictions, with implementations that can estimate this distribution online without prior knowledge.
  • On a synthetic benchmark, Balanced MSE remains consistent as training-label skew increases, whereas reweighting incurs significantly larger prediction error.
  • Balanced MSE achieves strong empirical results on age and depth estimation and on the proposed multi-dimensional IHMR benchmark.

2. Related Works

Logit adjustment methods modify classification mapping functions using an offset proportional to log ptrain(y) to estimate balanced conditional probabilities.

  • Logit adjustment modifies logits in mappings such as Softmax or Sigmoid by an offset proportional to log ptrain(y).
  • These techniques can serve as either train-time loss functions or test-time adjustments.
  • Recent work identifies logit adjustment as providing Bayes-optimal estimation of pbal(y|x).

3. Methodology

The method reframes imbalanced regression as estimating a balanced conditional label distribution rather than the skewed training conditional. Balanced MSE performs this statistical conversion during training and connects to logit adjustment in imbalanced classification.

  • Problem setting: Imbalanced regression targets pbal(y|x) under a uniform test-label distribution, while training uses skewed ptrain(y).The label-conditional distribution p(x|y) is assumed unchanged between training and testing.
  • Revisiting MSE: MSE is equivalent to Gaussian negative log likelihood and therefore learns ptrain(y|x), causing underestimation for rare labels.The mismatch follows because ptrain(y|x) is proportional to pbal(y|x) multiplied by the imbalanced training prior.
  • Balanced MSE: Balanced MSE restores pbal(y|x) by converting it into ptrain(y|x) with the training label prior before computing the training negative log likelihood.At test time, the conversion is skipped and the regressor directly outputs its balanced prediction.
  • Balanced MSE: Balanced MSE combines the standard MSE term with a balancing term, and standard MSE is recovered when the training label distribution is uniform.The balancing term requires an integral over the label space.
  • Connection with imbalanced classification: The statistical conversion underlying Balanced MSE becomes the logit-adjustment form for discrete labels, unifying imbalanced regression and classification.The paper presents these as two instantiations of the same theorem.

4. Experiments

Synthetic and real-world experiments show that Balanced MSE is robust to label-distribution skew and improves performance across uni- and multi-dimensional regression benchmarks.

  • Synthetic benchmark: Balanced MSE achieves a marginal label distribution closest to uniform in two-dimensional regression.
  • Synthetic benchmark: Balanced MSE remains robust across different training-distribution skewness, noise, nonlinear mappings, and multi-dimensional regressions.Reweighting can fail even when the true label distribution is known, whereas Balanced MSE remains robust across these settings.
  • Datasets and metrics: The experiments use age, depth, and a new multi-dimensional IHMR benchmark, extending imbalanced regression evaluation beyond one-dimensional labels.IHMR estimates 3D human meshes from monocular images.
  • Age estimation: Balanced MSE substantially outperforms prior methods on IMDB-WIKI-DIR age estimation, while BMC exceeds SOTAs without preprocessed label-distribution knowledge.The method also trades performance between frequent and rare labels toward balanced estimation.
  • Depth estimation: Both closed-form and numerical Balanced MSE implementations achieve clear improvements over SOTA on NYUD2-DIR depth estimation.Because depth has inter-pixel dependency, the evaluation fixes σnoise at 1 and uses BNI for the numerical option.
  • Human mesh recovery: Balanced MSE beats the HMR baseline by 3.4 on bMPVPE, while PM-Net performs better on tail-5% bMPVPE.PM-Net’s regression initialization is described as complementary to Balanced MSE.

5. Discussion and Conclusion

The paper concludes that Balanced MSE addresses MSE’s ineffectiveness under imbalanced regression and performs well across diverse benchmarks, while acknowledging that label imbalance is only one source of bias.

  • Conclusion: Balanced MSE revisits MSE probabilistically and offers closed-form and numerical implementations for imbalanced regression.
  • Conclusion: Balanced MSE outperforms existing methods on uni- and multi-dimensional imbalanced regression benchmarks.
  • Broader impacts: The method addresses bias from imbalanced labels, but other dataset biases may remain and can still produce negative social impacts.
  • Discussion: When ptrain(y) is uniform, Balanced MSE reduces to the standard MSE loss up to an additive constant.
  • Discussion: MSE is equivalent to negative log likelihood and therefore models the training conditional distribution ptrain(y|x).

A.3. GAI Loss Derivation

The GAI derivation evaluates the Balanced MSE integral analytically by exploiting Gaussian-product identities, then substitutes the resulting expression into the loss.

  • Gaussian integration: The derivation evaluates the Gaussian integral in Balanced MSE and defines S_i as a Gaussian evaluated at the prediction.
  • Closed-form loss: The resulting closed-form Balanced MSE loss combines the negative log Gaussian likelihood with log Σ_i φ_iS_i.
  • Synthetic setup: The synthetic construction samples imbalanced labels, adds Gaussian noise, and generates inputs through an invertible oracle relation.
  • Synthetic setup: The test set uses a uniform label distribution without added noise, enabling evaluation against balanced coverage of the label space.

B.1.2 Training Details

Training details specify optimization settings, benchmark configurations, label-distribution visualizations, and ablations for synthetic and real-world experiments.

  • Synthetic training: The synthetic benchmark uses batch size 256, SGD for one-dimensional linear regression, and Adam for nonlinear and two-dimensional regressions.Training lasts 2K epochs for one-dimensional linear regression and 10K epochs for the other synthetic settings.
  • Synthetic ablation: Balanced MSE remains robust to different sampled noises in the synthetic benchmark.
  • Evaluation: IMDB-WIKI-DIR contains visible test-distribution tails, so overall metrics may inadequately assess children and senior adults.
  • Real-world training: The age-estimation setup uses a ResNet-50 backbone, while depth estimation uses a ResNet-50 encoder-decoder and GMM-based label modeling.
  • Real-world training: The HMR setup uses a pretrained SPIN feature extractor, retrains the linear regressor, and evaluates on 3DPW after training on specified 2D and 3D datasets.
  • Ablations: The experiments include quantitative synthetic results and ablations over noise scale and GMM component count.

C. Experiment on random seeds

Balanced MSE is robust across random seeds, whereas reweighting is highly sensitive and can vary drastically in performance.

  • Reweighting’s performance varies drastically when the random seed changes, while Balanced MSE remains robust to different noise sampling results.The sensitivity may arise because reweighting magnifies rare-label noise when the zero-mean noise assumption does not hold.

D. Quantitative results for the synthetic benchmark

Synthetic experiments evaluate one- and two-dimensional regression under Normal, Exponential, and Multivariate Normal label distributions with varying skewness. Balanced MSE variants outperform Vanilla and Reweighting, especially under high skewness, while practical implementation choices remain comparable.

  • The synthetic benchmarks cover Normal and Exponential one-dimensional regression plus Multivariate Normal two-dimensional regression.
  • GAI and BMC significantly outperform Vanilla and Reweighting, particularly when label-distribution skewness is high.
  • BMC achieves comparable performance to the closed-form GAI implementation.
  • Using a learned noise scale achieves comparable performance to using the true noise scale.

F. Ablations

Ablations examine noise-scale selection and Gaussian-mixture complexity on IMDB-WIKI-DIR, identifying trade-offs and a small component count sufficient for optimal performance.

  • Larger σnoise shifts performance toward tail labels, while jointly optimizing σnoise finds an effective trade-off point.
  • GAI performance becomes optimal when the Gaussian mixture model uses K ≥ 2 components, possibly because the training label distribution is relatively simple.

G. Additional Discussions and Analysis

Additional analysis reports negligible training overhead and no added inference cost, explains Balanced MSE’s Bayes-optimal connection under its Gaussian formulation, and identifies modeling-scope limitations and future extensions.

  • Computational Cost: Balanced MSE adds negligible train-time overhead compared with the overall cost and no additional computational cost during inference.
  • Bayes-optimal Prediction: Under the Gaussian formulation, Balanced MSE’s predicted mean equals the mode and therefore gives the Bayes-optimal prediction for a balanced test set.
  • Additional Discussions and Analysis: Balanced MSE assumes isotropic Gaussian noise; finer-grained correlated-noise modeling is outside its scope and may improve regression performance.
  • Additional Discussions and Analysis: Balanced MSE estimates constant noise and gains from imbalance handling, while extending balanced losses to L1 and Huber remains future work.
Loading 2203.16427v1…