[22′ ICML] RankSim: Ranking Similarity Regularization for Deep Imbalanced Regression

1. 들어가며

이번 포스팅에서는 2022년 ICML에 발표된 RankSim: Ranking Similarity Regularization for Deep Imbalanced Regression 논문의 핵심 내용을 정리해보겠습니다. 이번 논문을 RankSim이라고 부르겠습니다!

RankSim은 Long Tail Regression 문제를 다루고 있는데요. 이 전에도 long tail regression을 다루는 다양한 연구들이 있었습니다. 그 중에서도 RankSim과 핵심 아이디어를 공유하고 있는 방법은 DIR입니다. Regression은 Classification 문제와는 차별화되는 특징들이 있습니다. 그중에서도 대표적으로 예측값이 연속적이라는 특증이 있는데요. 그 말은 feature space에서도 인접한 feature들끼리는 가깝고, 먼 label feature들 끼리는 멀어야 한다는 것이죠. 이러한 아이디어를 구현한 방식이 DIR 이었습니다. 하지만 DIR은 모든 feature들간의 관계를 모델링하지는 못했고, 인접 feature 들끼리의 관계에만 초점을 맞췄다는 한계가 존재했습니다.

이에 RankSim에서는 DIR에서의 아이디어를 더욱 발전시켜 global한 관점에서 모든 feature들끼리의 관계를 align 하는 방법을 제안합니다.

2. 제안 방법

핵심 아이디어는 아주 간단합니다. 미니 배치에서 각 label을 기준으로 정렬을 할 수 있잖아요. 말하자면 이 순서가 미니 배치에서의 랭킹이 될텐데요. 이들 label에서의 랭킹이 feature space에서의 랭킹과 동일해지도록 모델이 학습하자는 겁니다.

위 그림을 보면 입력으로 받은 사람 이미지의 나이를 예측하는 태스크에서의 RankSim loss 계산 방법을 표현하고 있는데요. 위의 파란색은 label space에서의 상호간 랭킹을, 아래 주황색은 feature space에서의 상호간 랭킹을 표현하고 있습니다. 이 둘이 같아지는 방향으로 학습을 유도하는게 RankSim입니다!

그러자면 몇 가지 유의할 점이 있는데요. 먼저 배치 구성입니다. 원래 미니 배치는 train loader에서 랜덤으로 구성하잖아요. 하지만 이렇게 구성할 경우 label끼리 겹치게 되고, 랭킹이 꼬이기 때문에, 미니 배치를 구성할때 각 label이 한 번씩만 등장하도록 구성해줘야 합니다.

또한 랭킹을 매기는 sorting 알고리즘은 기본적으로 미분이 불가능합니다. 이러면 gradient 계산이 안되어서 모델 학습이 불가능한데요. 이 문제는 Combinatorial Solver 기반 근사화를 통해 해결합니다.

파이썬 코드로 보면 RankSim loss를 쉽게 이해할 수 있습니다.

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

# 유사도 함수: cosine similarity
def cosine_sim(a, b):
    a_norm = F.normalize(a, dim=-1)
    b_norm = F.normalize(b, dim=-1)
    return torch.matmul(a_norm, b_norm.T)  # (B, B)

# 정렬 순위 벡터 계산 (ranking)
def rank_tensor(tensor):
    sorted_indices = torch.argsort(tensor, dim=-1, descending=True)
    ranks = torch.argsort(sorted_indices, dim=-1)
    return ranks.float()

# RankSim Loss
class RankSimLoss(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, features, labels):
        """
        features: (B, D) - feature vectors from backbone
        labels: (B,) - scalar regression targets
        """
        with torch.no_grad():
            # label space similarity: negative absolute distance
            label_sim = -torch.abs(labels.unsqueeze(1) - labels.unsqueeze(0))  # (B, B)
            rank_label = rank_tensor(label_sim)  # (B, B)

        # feature space similarity: cosine
        feat_sim = cosine_sim(features, features)  # (B, B)
        rank_feat = rank_tensor(feat_sim)

        # loss: MSE between rank vectors
        loss = F.mse_loss(rank_feat, rank_label)
        return loss

# 간단한 회귀 모델 (MLP)
class ToyRegressor(nn.Module):
    def __init__(self, input_dim=1, hidden_dim=32):
        super().__init__()
        self.feature = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        self.regressor = nn.Linear(hidden_dim, 1)

    def forward(self, x):
        z = self.feature(x)
        y_pred = self.regressor(z).squeeze(-1)
        return y_pred, z  # prediction, features

# 예시 학습 루프
if __name__ == "__main__":
    torch.manual_seed(0)
    N = 64
    x = torch.linspace(-3, 3, N).unsqueeze(1)
    y = x.squeeze()**3 + 0.1 * torch.randn(N)  # 비선형 함수 + 잡음

    model = ToyRegressor()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion_mse = nn.MSELoss()
    criterion_ranksim = RankSimLoss()

    for epoch in range(2000):
        model.train()
        optimizer.zero_grad()
        y_pred, z = model(x)

        loss_mse = criterion_mse(y_pred, y)
        loss_ranksim = criterion_ranksim(z, y)
        loss = loss_mse + 0.5 * loss_ranksim  # γ=0.5

        loss.backward()
        optimizer.step()

        if epoch % 200 == 0:
            print(f"[Epoch {epoch}] MSE: {loss_mse.item():.4f}, RankSim: {loss_ranksim.item():.4f}")

핵심 클래스인 RankSimLoss를 볼게요. 먼저 label ranking 을 계산하기 위해 미니 배치 내 각 label 간의 차이를 계산해줍니다. 그리고 torch.argsort 함수를 사용해 label ranking을 계산해줍니다. 마찬가지로 feature ranking도 계산해줄건데요. 먼저 feature를 추출해주고, cosine sim을 기반으로 동일하게 ranking을 계산해줍니다. 최종 loss는 전체 배치에 대해 (B,B) 사이즈에서의 MSE loss로 구현해줍니다. 그 말은 각 미니 배치에서 모든 데이터간의 랭킹 관계를 Ground Truth와 같아지도록 학습하겠다는거죠.

3. 실험 결과

이렇게 제안한 RankSim 성능은 얼마나 좋은지, 다양한 실험 결과들을 살펴보겠습니다.

3-1. benchmark experiments

먼저 기존 연구들과 동일하게 다양한 벤치마크에서 실험을 진행했는데요. IMDB 데이터셋에서의 성능 비교 실험 결과는 다음과 같습니다.

다음은 AgeDB 데이터셋 실험 결과입니다.

이어서 STS-B 데이터셋 실험 결과입니다.

기존 방법들과 병렬로 적용할 수 있다는점이 큰 장점이고, 모든 구간에 대해서 유의미하게 성능이 개선되는 모습입니다.

3-2. Ablation

RankSim loss는 기존 loss들과 혼합해서 사용할 수 있는데요. 어떤 loss랑 가장 궁합이 잘 맞을까요?

MSE랑 가장 좋은 궁합을 보여주네요.

Feature 랭킹은 다양한 similarity 방법으로 측정할 수 있는데요.

대체로 Cosine Sim을 사용할때 가장 효과적인 모습을 보여줍니다.

배치 샘플링 전략 또한 성능에 큰 영향을 줄텐데요.

배치 내에 중복 라벨을 방지해주면 성능 개선폭이 더 큰 모습을 볼 수 있습니다.

3-3. Qualitative Analysis

애초에 RankSim에서 중요하게 생각했던건 label space 상에서의 랭킹과 feature space 상에서의 랭킹이 유사해야한다는 것인데요. 과연 학습된 feature들이 진짜 label 순서에 맞게 분포하고 있을까요?

위 그림은 다양한 학습 방법에 따른 feature 분포를 비교해주고 있습니다. RankSim은 다른 방법들과 비교했을때 확실히 feature 공간상에서 순차적으로 잘 배열되어 있는 모습을 볼 수 있습니다.

다양한 Anchor Age를 기준으로 feature들간의 cosine similarity를 측정해봐도, 일관되게 align 되어 있는 모습을 확인할 수 있습니다.

3-4. zero shot

RankSim의 재미있는 장점 중 하나는 zero shot 성능이 향상된다는 점인데요. 생각해보면 RankSim 방식으로 학습하면 학습 데이터가 없어도 좋은 정보를 줄 수 있습니다. 예를 들어 31세 데이터가 없었다고 생각해볼게요. 하지만 우리는 29세, 30세, 32세, 33세 데이터를 받으며 이 들이 순차적으로 존재해야 한다는 내용을 학습합니다. 그럼 아무래도 31세 데이터를 본적은 없지만 대략적으로 30세와 32세 사이에 존재하도록 모델이 학습되지 않겠어요?

위 그림은 의도적으로 일부 데이터를 누락시켜 모델을 학습한 뒤 성능 평가를 진행한 결과를 보여주고 있는데요. 학습 데이터가 존재하지 않는 영역에서 큰 폭의 성능 개선을 보여주고 있습니다.

4. 마치며

Long Tail Regression 문제를 다룬 RankSim 논문의 핵심 내용들을 정리해봤습니다. RankSim은 기존 연구들을 충실하게 이어간다고 볼 수 있는데요. 기존 연구인 DIR에서의 feature space 상에서 데이터들끼리의 align이 맞아야 한다는 아이디어를 더욱 개선한 연구라고 볼 수 있습니다.

5. 참고자료

• RankSim paper
• DIR 논문 리뷰

Ffightingseok