[22′ CVPR] Balanced MSE for Imbalanced Visual Regression

1. 들어가며

이번 포스팅에서는 22년 CVPR에 발표된 Balanced MSE for Imbalanced Visual Regression 논문의 핵심 내용을 정리해보겠습니다. 이 논문은 long tail regression 문제를 다루고 있으며, 기존 MSE loss를 long tail regression에 그대로 적용했을때의 한계를 지적하며, MSE loss를 학습 데이터 분포에 맞게 balancing 하는 방법을 제안하고 있습니다. 이번 글에서는 이번 논믄을 balanced MSE라고 부르겠습니다.

long tail 문제는 꽤나 오래된 역사를 갖고 있습니다. 딥러닝 연구가 활발해지면서 딥러닝이 성능이 굉장히 좋지만 학습 데이터 분포가 치우쳐진 형태일 경우 성능이 저하된다는 현상이 발견되었죠. 이후로 학습 데이터 분포가 치우쳐져있는 long tail 문제를 개선하기 위한 다양한 방법이 연구되고 있습니다.

특히 class imbalance classification은 다양한 연구가 활발히 연구되고 있습니다. 반면 long tail regression은 classification 만큼 활발히 연구되고 있지는 않은데요. 대표적인 원인으로는 classification과 달리 ‘빈번한 정도’를 정의하기 어렵다는 이유를 들 수 있습니다. classification의 경우 클래스의 빈도수만 측정하면 데이터 분포가 얼마나 imbalance 한지를 측정할 수 있고, 따라서 이를 바탕으로 다양한 개선 방법이 제안되고 있죠. 반면 regression은 클래스가 아닌 연속값으로 분포하기에 이러한 접근 방식을 그대로 적용하기 어렵습니다.

그럼에도 long tail regression 문제를 개선하기 위한 여러 접근 방법들이 제안되었는데요. 대표적으로 Reweighting 논문은 KDE로 y 분포를 측정해 빈도수에 반비례하는 가중치를 loss에 더해주는 방법을 제안했습니다. classification의 접근법을 응용한것으로 볼 수 있죠. 또한 DIR에서는 regression 문제에 맞게 feature space에서 인접한 y들끼리 align을 맞추는 방법을 제안했습니다. 이는 기존 classification과 차별화되는 regression만의 특징을 모델링에 멋지게 녹여낸 방법이라고 할 수 있습니다.

그럼에도 long tail regression 문제는 완벽하게 해결되지 않고 있습니다. 특히 MSE loss가 regression에서 어떠한 의미를 갖는지를 근본적으로 파고들어 접근한 방법은 없었죠. Balanced MSE에서는 MSE loss와 regression 문제의 상관관계를 분석하고, long tail에 적용하기 위해 MSE를 개선할 수 있는 방법을 제안합니다.

2. 제안 방법

2-1. MSE를 학습한다는 의미

보통 regression 문제는 MSE loss로 학습하는데요. MSE loss는 확률적인 관점에서, 정규분포를 따르는 데이터의 평균값을 맞추는 문제라고 볼 수 있습니다. 이는 수식을 따라가보면 이해할 수 있는데요. 먼저 MSE loss는 다음과 같이 예측값과 정답값의 제곱값의 평균으로 표현됩니다.

한편 y^을 평균으로, sigma를 표준편차로 하는 정규분포에서의 y 확률밀도함수는 다음과 같이 표현되는데요.

여기에 log를 취하면 다음과 같이 깔끔하게 변환됩니다.

즉 MSE는 정규분포를 이루는 데이터에서의 확률값을 최대화하는 의미라는것을 알 수 있습니다. 그 말은 학습 데이터 분포는 정균분포를 이룬다는 가정 하에, 예측 값의 분포 평균을 데이터 정규 분포 평균과 맞추는 작업인거죠. 그리고 정규분포에서의 확률값을 최대화하기 위해 가장 빈도수가 많은 평균값에 치우쳐 학습한다는것을 알 수 있습니다.

2-2. 기존 방식의 문제점

문제는 long tail regression 문제에서의 전제가 학습 데이터 분포는 skewed 되어 있지만 테스트 분포는 uniform 한다는 것입니다. 이는 공평한 성능 평가를 위해서인데요. 테스트 분포를 학습 데이터처럼 skewed 분포로 구성할 경우, 빈번한 영역의 데이터를 잘 맞추는 모델이 최종 성능 평가에서 좋은 성적을 받게 됩니다. 이는 우리가 원하는 ‘잘 예측하는’ 모델의 모습은 아니죠. 따라서 테스트 데이터는 학습 데이터와 달리 출력값이 uniform 분포를 이루도록 구성해줍니다.

이렇게 학습 데이터 분포와 테스트 데이터 분포가 달라지면서 학습된 모델은 테스트 데이터에서의 성능이 크게 떨어지게 되는데요. 학습 데이터의 평균에 맞춰 학습한 모델이기에 갑자기 테스트 데이터로 넘어오면서 맞춰야 하는 평균이 달라지면서 성능이 급격하게 떨어지는 것이죠.

2-3. 제안 방법 : Balanced MSE

이에 저자들은 balanced MSE라는 MSE loss를 개선한 방법을 제안합니다. 순서대로 따라가 보겠습니다. 먼저 기존 MSE Loss는 위에서 설명한것과 같이 정규분포의 NLL와 동일합니다.

이대로 학습할 경우 데이터 분포가 달라질경우 제대로 동작하지 못한다는 문제가 있죠. 따라서 이를 보정해 줄 정규화항을 추가해줄겁니다.

위 수식이 이번 balanced MSE loss의 핵심이므로 하나씩 뜯어보며 의미를 생각해볼게요. 우선 y’은 학습 데이터에서의 모든 y를 의미합니다. 따라서 첫 번째 적분항은 모델이 예측한 y^을 중심으로 하는 정규분포에서 학습 데이터 y들의 확률값의 합을 의미합니다. 그리고 두 번째 적분항은 학습 데이터에서의 y들의 확률이므로, 쉽게 빈도수라고 이해할 수 있습니다. 즉 모델에게 패널티를 줄건데, 모델이 예측한 정규분포가 학습 데이터 분포와 다를수록 패널티를 주는거죠. 그런데 그냥 다를수록 주는게 아니고, 빈도수에 따라 패널티를 줄겁니다. 따라서 모델은 빈도수가 큰 값을 예측할때마다 패널티를 받는겁니다. 결과적으로는 기존에는 빈도수가 낮은 값은 무시하고 학습했는데, 이제 이 값들을 잘 예측하도록 노력하게 되겠죠.

최종적으로 balanced MSE loss는 다음과 같이 표현됩니다.

하나의 log로 묶으면 다음과 같습니다.

실제 구현을 위해 모든 y에 대해 적분하는 대신 배치 샘플을 전체 y와 동일하다고 가정하면 다음과 같이 표현할 수 있습니다.

이는 다음과 같이 하나의 log로 묶어 정리할 수 있는데요.

위 수식은 softmax와 동일한 수식임을 알 수 있습니다. 따라서 cross entropy loss로 쉽게 구현할 수 있습니다.

2-4. 파이썬 코드 구현

Cross Entropy로 구현한 Balanced MSE loss 파이썬 코드는 다음과 같습니다. 배치에 대해 각각의 MSE를 계산해준 뒤 Cross Entropy를 구하는 방식으로 구현해줄 수 있습니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

class BalancedMSELoss1D(nn.Module):
    def __init__(self, tau=2.0):
        """
        tau: 온도 파라미터 (temperature), 논문에서는 tau = 2 * sigma^2 로 사용됨
        """
        super().__init__()
        self.tau = tau

    def forward(self, y_pred, y_true):
        """
        y_pred: (B, 1) - 모델의 예측값
        y_true: (B, 1) - 실제 정답값 (레이블)
        """

        # B: 배치 크기
        B = y_pred.size(0)

        # (B, 1) → (B,)로 reshape
        y_pred = y_pred.view(B)
        y_true = y_true.view(B)

        # 예측값과 실제값 간의 pairwise L2 거리 제곱 계산 (B x B)
        # y_pred[i] - y_true[j] 를 모두 계산 → (i, j) 쌍으로 구성된 거리 행렬
        diff = y_pred.unsqueeze(1) - y_true.unsqueeze(0)  # (B, 1) - (1, B) → (B, B)
        dist_sq = diff ** 2  # 제곱 (B, B)

        # similarity score 계산: negative distance / tau
        logits = -dist_sq / self.tau  # (B, B), 마치 유사도처럼 해석 가능

        # 정답 인덱스: diagonal 위치가 각 샘플의 올바른 매칭
        labels = torch.arange(B, device=y_pred.device)  # (B,)

        # cross-entropy loss 계산 (softmax over row + NLL)
        # 각 row에서 정답은 diagonal 위치 (i-th row의 i-th column)
        loss = F.cross_entropy(logits, labels)

        return loss

3. 실험 결과

이번에는 논문에서 제안한 balanced MSE 방식이 실제 long tail regression 문제에서 효과가 있는지 다양한 실험 결과를 통해 살펴보겠습니다.

3-1. Synthetic dataset

먼저 1차원 회귀 문제로 합성한 데이터셋을 사용하여 다양한 방식의 성능을 비교해봤는데요. 성능은 다음과 같습니다.

왼쪽 두개 열은 학습 데이터가 normal 분포인 경우, 오른쪽 두개 열은 학습 데이터가 long tail 분포인 경우의 결과를 보여주고 있습니다.

두 번째 열을 살펴보겠습니다. 학습 데이터는 normal 분포이지만 (파란색), 테스트 데이터는 uniform 분포 (분홍색)를 보여주고 있는데요. 일반 MSE로 학습한 모델은 학습 데이터 분포에 가까운 정규 분포 형태의 출력 분포 (노란색) 을 보여주고 있습니다.

반면 balanced MSE (빨간색)은 테스트 분포와 거의 유사한 uniform 분포 출력을 보여주고 있습니다. 네 번째 열을 살펴보겠습니다. 학습 데이터가 long tail 분포를 보이는 상황인데요. 역시 일반 MSE는 테스트 분포와 꽤 차이를 보이는 모습입니다. 반면 balanced MSE는 test 분포를 잘 따라가는 모습을 보입니다.

3-2. Real world dataset

이어서 실제 데이터셋에서의 성능을 살펴보겠습니다. 먼저 Age Estimation 문제입니다.

Balanced MSE는 Med, Few 빈에 대해 큰 효과를 내고 있고, Many 빈에 대해서는 오히려 약간의 성능 하락 모습을 보여주고 있습니다.

위 그림은 나이별 성능 개선폭을 시각화한 플롯인데요. 확실히 데이터가 적은 고령층에서의 성능 개선 폭이 두드러지는 모습입니다. 이러한 모습은 balanced MSE의 의도에 맞게 빈번한 데이터에 패널티를 주면서 trade off를 반영한 결과라고 볼 수 있습니다. 이어서 Depth Estimation 벤치마크 실험 결과입니다.

마찬가지로 Med, Few 빈에서의 성능이 크게 개선되는 모습입니다. 마지막으로 Human Mesh Recovery 벤치마크 실험 결과입니다.

일부 PM-Net에서의 성능이 가장 우수하기도 하지만, 대체로 balanced MSE를 적용했을때 rare pose에서 성능이 개선되는 모습을 볼 수 있습니다.

4. 마치며

이번 글에서는 long tail regression 문제를 다룬 balanced MSE 논문의 주요 내용을 정리해봤습니다. 지금까지의 long tail regression 문제는 class imbalance classification 아이디어를 regression 형태로 변형하거나 regression 문제의 특성을 모델링에 반영하는 방법들이었는데요. balanced MSE에서는 regression에서의 MSE loss의 의미를 수학적으로 분석하고, long tail 문제에 맞게 loss 수식을 변형하는 제안을 했습니다.

5. 참고자료

• balanced MSE paper
• focal loss 논문 리뷰
• reweighting 논문 리뷰
• DIR 논문 리뷰

Ffightingseok