[24′ CVPR] Deep Imbalanced Regression via Hierarchical Classification Adjustment

1. 들어가며

이번 포스팅에서는 24년 CVPR에 발표된 Deep Imbalanced Regression via Hierarchical Classification Adjustment 논문의 핵심 내용들을 정리해보겠습니다. 이번 논문은 Hierarchical Classification Adjustment의 앞글자를 따서 HCA라고 부르겠습니다.

HCA는 long tail regression 문제를 다루고 있습니다. 기존에도 long tail 문제를 다루는 다양한 연구들이 있었는데요. 주로 빈도수에 따라 loss 가중치를 조절하거나, 기존 MSE loss를 long tail 분포에 맞게 바꿔주는 방법등을 사용했습니다. 한편 비슷한 문제인 class imbalance classification 또한 활발하게 연구되고 있는데요. Regression 문제와 달리 Classification은 class별로 빈도수를 명확하게 계산할 수 있어서 보다 수월하고, 따라서 더 활발히 연구되고 있습니다.

HCA에서는 기존 long tail regression 방법들과는 다르게 classification 접근법을 그대로 적용하는 시도를 보여줍니다. 이름에서도 볼 수 있듯이 다양한 사이즈의 classification model을 결합한 Hierarchical Classification 방식을 제안합니다. 또한 이렇게 다양한 Classification 모델을 보정하는 방법인 Adjustment도 제안해주고 있습니다.

2. 제안 방법

2-1. Hierarchical Classifier

HCA의 핵심 아이디어는 아주 직관적입니다. long tail regression 문제를 regression이 아닌, classification 문제로 풀자는거에요. 예를 들어 1~100까지 예측해야 하는 문제라면 10개씩 10덩어리로 나누어 classification으로 풀자는거죠. 그럼 3을 예측해야 했다면 1번 클래스를, 15를 예측해야 했다면 2번 클래스를 예측해야 하는 문제로 바뀌는거죠.

이렇게 regression 문제를 classification으로 변환해서 풀면 기존 class imbalance classification의 다양한 방법론들을 그대로 적용할 수 있다는 장점이 있습니다. 하지만 치명적인 단점은 정확도가 떨어진다는건데요. 정확한 값을 예측하지 못하고, 대략적인 ‘덩어리’ 만 예측할 수 있기 때문이죠. 이러한 문제를 보완하기 위해 하나의 단일 classifier가 아닌, 여러 사이즈의 classifier를 종합하는 방식을 사용합니다. 예를 들어 다시 1~100을 예측해야 하는 문제라면, 50개씩 2덩어리, 20개씩 5덩어리, 10개씩 10덩어리, 5개씩 20덩어리를 나누어 예측하는 각각의 classifier를 학습해주는거에요. 그리고 나서 최종 정답은 이들 다양한 classifier들의 예측값을 종합하여 보정한뒤 판단하는거죠.

위 그림은 HCA 제안 방법의 핵심 아이디어를 보여주고 있습니다. 가장 위쪽은 coarse classifier만으로 예측한 상황을 보여주고 있는데요. 정답과 22만큼 차이나는 모습입니다. 이제 아래쪽으로 갈수록 fine grained classifier로부터의 예측값을 사용하여 정답을 보정해주는데요, 정답과의 차이가 11, 4로 점점 줄어드는 모습을 보여주고 있습니다.

모델 구성 방법은 지금까지 설명한 핵심 아이디어를 그대로 모델로 옮겨주는 과정입니다.

위 그림은 HCA 모델의 동작 방식을 보여주고 있습니다. 총 3개의 classifier들을 종합한 케이스를 보여주고 있는데요. 이때 T는 Transition Matrix로, 각 classifier들의 클래스간의 상대위치를 보정해주는 역할을 해줍니다. 예를 들어 다시 1~100을 예측해야 하는 문제에서 10개씩 10덩어리, 50개씩 2덩어리를 예측하는 2개의 classifier가 있다고 가정해볼게요. 그럼 2덩어리 classifier의 첫 번째 덩어리는 10덩어리 classifier의 5개 클래스에 해당합니다. 이를 mapping 해주는 matrix입니다.

가장 오른쪽에서는 여러 classifier들의 예측값을 최종 예측값으로 환산하는 과정을 보여주고 있는데요. HCA에서는 총 두가지 방법을 제안하고 있습니다. 첫 번째는 Addition 방식입니다. 각 예측값을 Transition Matrix로 최종 클래스로 변환한 뒤 모두 더해주는 방식이에요. 두 번째 방식은 Multiplication 방식으로, 마찬가지로 모든 예측값을 최종 클래스로 변환해준 뒤 곱해서 최종 예측값을 계산해주는 방식입니다. Addition 방식을 파이썬 코드로 보면 다음과 같습니다.

class HCAAdjuster(nn.Module):
    def __init__(self, class_bins):
        super().__init__()
        self.class_bins = class_bins
        self.H = len(class_bins)
        self.transition_matrices = self.build_transition_matrices()

    def build_transition_matrices(self):
        matrices = []
        for h in range(self.H - 1):
            coarse, fine = self.class_bins[h], self.class_bins[-1]
            mat = torch.zeros(coarse, fine)
            size = fine // coarse
            for i in range(coarse):
                mat[i, i * size:(i + 1) * size] = 1
            matrices.append(mat)
        return matrices  # 각 h에 대해 [C_h x C_H]

    def forward(self, probs):
        p_H = probs[-1]  # finest
        adjusted = p_H.clone()
        for h in range(self.H - 1):
            T = self.transition_matrices[h].to(probs[h].device)
            projected = torch.matmul(probs[h], T)  # [B, C_H]
            adjusted += projected
        return adjusted

2-2. Distillation Model

이렇게까지 구성하면 HCA의 핵심 아이디어는 완성이 되는데요. 저자들은 여기서 그치지 않고 Distillation Model을 제안합니다. 여기까지만 구성된 모델의 성능은 생각보다 크게 우수하지 않은데요. 크게 두가지 문제점이 있습니다. 첫 번째는 여러개의 classifier로 구성되어 있으므로 inference 단계에서 속도가 느리다는 점이구요. 두 번째는 계층간 classifier들의 예측값이 일관되지 않을 수 있다는겁니다. 첫 번째 classifier는 첫 번째 덩어리라고 예측했는데, 마지막 classifier는 마지막 덩어리라고 예측하면 곤란하겠죠?

이러한 상황을 보완하기 위해 저자들은 Distillation Model을 제안합니다. 기존 Hierarchical Classifier가 예측한 logit을 따라하도록 distillation을 학습하는거죠. 이렇게 distillation 방식으로 구성했을때 예상되는 장점은 다음과 같은데요. 우선 첫 번째로 당연히 하나의 classifier만 사용하게 되므로 inference 속도가 빨라집니다. 두 번째로 일반화 성능의 향상 효과를 기대할 수 있는데요. 기존 classifier는 hard label을 학습하는 반면, distillation model은 hierarchical classifier logit을 학습하기 때문에 soft label을 학습하게 됩니다. 따라서 여기서 오는 일반화 성능의 향상 효과를 기대할 수 있습니다. 세 번째로 가장 결정적인 문제인 classifier들간의 예측값이 불일치하는 상황을 개선할 수 있습니다. Distillation model에서는 이 문제를 해결하기 위해 Range-Preserving Alignment 기능을 추가하는데요. 핵심 아이디어는 coarse classifier가 예측한 클래스 범위 내에서만 Distilled Classifier가 따라하도록 강제해주는겁니다. 즉 fine classifier의 예측이 coarse classifier 예측 범위를 넘어갈 경우 무시하겠다는거죠. coarse classifier 정확도가 가장 높을테니, 이 친구의 예측값을 기반으로 학습해주겠다는 아이디어입니다.

def range_preserving_align(pT, T_mat):
    p_aligned = torch.max(T_mat.unsqueeze(0) * pT.unsqueeze(1), dim=2)[0]
    p_aligned = p_aligned / (p_aligned.sum(dim=1, keepdim=True) + 1e-8)
    return p_aligned

위 코드에서 T_mat은 transistion matrix를, pT는 distilled model의 출력값을 의미합니다. 첫 번째 줄에서 T_mat과 pT의 matrix multiplication을 통해 distilled feature 값을 coarse range 안으로 정렬해주는 모습을 볼 수 있습니다.

최종 distilled model은 KL divergence loss로 학습해줄거고, 따라서 최종 코드는 다음과 같이 구성됩니다.

def hca_distill_loss(pT, ph_list, T_mats):
    total_loss = 0
    for ph, T in zip(ph_list, T_mats):
        pT_aligned = range_preserving_align(pT, T.to(pT.device))
        total_loss += F.kl_div(pT_aligned.log(), ph, reduction='batchmean')
    return total_loss

3. 실험 결과

이번에는 주요 실험 결과를 살펴보겠습니다.

3-1. Comparison with SOTA on Regression Tasks

먼저 IMDB WIKI 데이터셋입니다.

기존 방법인 balanced MSE는 tail 성능은 좋아지지만 head 성능이 크게 손상되는 모습을 보여줍니다. 반면 HCA는 head 성능을 유지하면서 tail 성능을 개선하는 모습을 볼 수 있습니다.

이어서 SHTech 데이터셋입니다.

기존 방법들과 단순 classification 방법과 비교했을때 성능이 개선되는 모습입니다.

4. 마치며

이번 글에서는 classification 방식으로 long tail regression을 푸는 HCA 방법을 살펴봤습니다. classification 방식으로 regression 문제를 접근했을때 명확한 장단점이 있는데요. HCA에서는 장점은 살리면서 단점을 보완하기 위해 Hierarchical 방식, adjustment 방법등을 제안했습니다.

5. 참고 자료

• HCA paper
• Balanced MSE 논문 리뷰
• Reweighting 논문 리뷰

Ffightingseok