[论文解读] Balanced MSE for Imbalanced Visual Regression

Jiawei Ren, Mingyuan Zhang+2 more|arXiv (Cornell University)|Mar 30, 2022

Imbalanced Data Classification Techniques被引用 6

一句话总结

本文提出 Balanced MSE，一种基于统计原理的损失函数，通过利用训练标签先验来纠正视觉回归中的标签分布不平衡问题，实现平衡预测。该方法在高偏度情况下显著优于标准 MSE 和重加权方法，是首个计算开销极低的高维不平衡回归通用解决方案。

ABSTRACT

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.

研究动机与目标

• 解决视觉回归任务（如年龄估计和姿态估计）中的数据不平衡问题，其中罕见标签在训练数据中代表性不足。
• 克服标准均方误差（MSE）的局限性，后者会将标签不平衡传播至预测结果，导致罕见类别性能下降。
• 开发一种适用于高维不平衡回归的通用解决方案，即使在缺乏训练标签分布先验知识的情况下亦可应用。
• 确保对标签分布偏度和随机种子变化的鲁棒性，以应对现有基于重加权方法的缺陷。

提出的方法

• 从统计角度重新诠释 MSE，将其视为在高斯预测分布下最小化负对数似然（NLL）的优化过程。
• 提出 Balanced MSE 作为对 MSE 的校正，利用训练标签分布将目标预测分布从平衡测试分布转换为与训练分布兼容的形式。
• 通过两种变体实现 Balanced MSE：一种为闭式解的高斯平均实现（GAI），另一种为可微分的蒙特卡洛近似（BMC），以提升可扩展性。
• 引入可学习的各向同性噪声尺度参数以建模不确定性，该参数与模型权重联合优化。
• 采用包含 K 个分量的高斯混合模型（GMM）来估计训练标签分布，从而灵活建模复杂的标签偏度。
• 将该方法应用于合成基准和真实世界任务，包括年龄估计、深度估计以及 3D 人体网格恢复（HMR）。

实验结果

研究问题

• RQ1 在视觉回归中，标准 MSE 在严重标签分布偏度下表现如何？为何其在罕见标签上失效？
• RQ2 能否设计一种基于统计原理的损失函数，在无需事先知晓标签分布的情况下纠正标签不平衡？
• RQ3 与基于重加权的方法相比，Balanced MSE 在标签分布偏度和随机种子变化下的鲁棒性如何？
• RQ4 Balanced MSE 是否能泛化至高维回归任务（如具有多维连续标签的 3D 人体网格恢复）？

主要发现

• 在合成 1D 和 2D 不平衡回归基准上，Balanced MSE 的平均绝对误差（MAE）显著低于标准 MSE 和重加权方法，尤其在高偏度条件下表现更优。
• 在 IMDB-WIKI-DIR 基准上，Balanced MSE 在所有标签组别（多、中、少、全部）中均将 MAE 降低 1.5–2.5 个点，对罕见（少）标签的改进最大。
• 该方法对随机种子变化具有鲁棒性，而重加权方法在不同随机初始化下表现出较高的性能方差。
• Balanced MSE 的 GAI 与 BMC 实现性能相当，其中 BMC 更具可扩展性，适用于复杂标签分布。
• 噪声尺度参数的联合优化相比固定噪声尺度，能实现对头部与尾部标签性能更优的权衡。
• Balanced MSE 的计算开销可忽略不计——在 IMDB-WIKI-DIR 上每轮仅增加约 0.3 秒，适用于实际部署。