[논문 리뷰] Learning to Re-weight Examples with Optimal Transport for Imbalanced Classification

Dandan Guo, Zhuo Li외 4명|arXiv (Cornell University)|2022. 08. 05.

Imbalanced Data Classification Techniques인용 수 7

한 줄 요약

이 논문은 클래스 불균형 문제를 해결하기 위해 최적 운반 기반의 재가중 방법을 제안하며, 클래스 재가중을 분포 근사 문제로 공식화하여 불균형 데이터와 균형 잡힌 데이터 분포 간의 최적 운반(OT) 거리를 최소화한다. 이전의 방법들과 달리, 가중치 학습을 분류기에서 분리하여 이미지, 텍스트, 포인트 클라우드 데이터셋에서 상태의 기술(SOTA) 성능을 달성하며, 클래스 불균형과 클래스 수에 대한 강건성도 향상시킨다.

ABSTRACT

Imbalanced data pose challenges for deep learning based classification models. One of the most widely-used approaches for tackling imbalanced data is re-weighting, where training samples are associated with different weights in the loss function. Most of existing re-weighting approaches treat the example weights as the learnable parameter and optimize the weights on the meta set, entailing expensive bilevel optimization. In this paper, we propose a novel re-weighting method based on optimal transport (OT) from a distributional point of view. Specifically, we view the training set as an imbalanced distribution over its samples, which is transported by OT to a balanced distribution obtained from the meta set. The weights of the training samples are the probability mass of the imbalanced distribution and learned by minimizing the OT distance between the two distributions. Compared with existing methods, our proposed one disengages the dependence of the weight learning on the concerned classifier at each iteration. Experiments on image, text and point cloud datasets demonstrate that our proposed re-weighting method has excellent performance, achieving state-of-the-art results in many cases and providing a promising tool for addressing the imbalanced classification issue.

연구 동기 및 목표

• 딥러닝에서 다수 클래스에 의해 모델이 지배당하고 소수 클래스에서 성능이 떨어지는 클래스 불균형 문제를 해결하기 위해.
• 기존 재가중 방법의 한계를 극복하기 위해, 가중치 학습을 분류기의 기울기와 결합함으로써 정확한 가중치 최적화가 어려운 문제를 해결하기 위해.
• 최적 운반을 사용하여 재가중을 분포 근사 문제로 공식화하고, 불균형 학습 데이터를 균형 잡힌 메타 분포와 정렬하기 위해.
• 프로토타입 기반 분포를 사용하여, 많은 클래스를 포함한 데이터셋에 적합한 강건하고 유연하며 메모리 효율적인 접근법을 개발하기 위해.

제안 방법

• 이 방법은 불균형 학습 데이터를 각 샘플이 가중치를 확보한 확률질량으로 가지는 이산 경험 분포 P로 표현한다.
• 균형 잡힌 메타 데이터셋은 균일한 확률질량을 가지는 이산 경험 분포 Q로 모델링되며, 이는 균형 잡힌 목표 분포를 나타낸다.
• P와 Q 사이의 최적 운반(OT) 거리를 최소화함으로써 가중치를 학습하며, 특징과 진짜 레이블을 기반으로 한 비용 함수를 사용한다.
• OT 손실을 직접 가중치 네트워크 학습에 사용함으로써 분류기의 기울기와 가중치 최적화를 분리하여 안정적이고 독립적인 학습이 가능하다.
• 메모리 비용을 줄이기 위해 프로토타입 기반 OT 손실을 도입하여 개별 샘플 대신 클래스 프로토타입으로 Q를 구성한다.
• 학습 가능한 가중치 네트워크를 사용한 엔드 투 엔드 학습과 OT 기반 직접 가중치 할당을 모두 지원하며, 가중치 학습 과정에서 분류기 의존도를 최소화한다.

실험 결과

연구 질문

• RQ1 분류기의 기울기를 활용하지 않고도 최적 운반을 효과적으로 사용하여 불균형 분류에서 샘플 가중치를 학습할 수 있는가?
• RQ2 다양한 데이터 모odal에서 기존 자동 재가중 방법과 비교해 볼 때, OT 기반 재가중 방법은 성능과 강건성 측면에서 어떻게 다른가?
• RQ3 특히 장꼬리(롱테일) 설정에서 클래스 수가 증가할 경우에도 제안된 방법이 높은 성능을 유지할 수 있는가?
• RQ4 프로토타입 기반 OT 변형은 전체 샘플 OT와 비교해 메모리 효율성과 정확도 측면에서 어떻게 다른가?
• RQ5 이 방법은 이미지, 텍스트, 3D 포인트 클라우드와 같은 다양한 데이터 유형으로 일반화되는가?

주요 결과

• SST-2 및 SST-5 텍스트 분류 벤치마크에서 극심한 불균형(1000:100) 조건 하에서 각각 87.08% 및 87.13%의 정확도를 달성하여, Logit Adjustment 및 Hu 등(2023)의 방법을 포함한 모든 베이스라인을 초월했다.
• SST-5에서 500:50의 클래스 불균형 조건 하에서 44.95%의 정확도를 기록하여, 다음으로 우수한 성능을 보인 제약 기반 재가중 방법(44.62%)을 크게 앞서며 고불균형 조건에서도 강건성을 입증했다.
• 이미지 분류에서 여러 데이터셋에서 상태의 기술(SOTA) 성능을 달성하며, 기존 재가중 및 메타러닝 기반 베이스라인 대비 일관된 우수성을 보였다.
• 프로토타입 기반 OT 변형은 전체 샘플 OT와 유사한 성능을 달성하면서도 메모리 사용을 극적으로 줄여 대규모 클래스 데이터셋에 대한 확장성을 확보했다.
• 이 방법은 이미지, 텍스트, 3D 포인트 클라우드 데이터에서 일관된 성능을 보이며, 다양한 데이터 모달 간의 일반화성과 강건성을 입증했다.
• 제거 실험을 통해 분류기에서 가중치 학습을 분리함으로써 장꼬리 상황에서 더 안정적이고 정확한 가중치 최적화가 가능하다는 점을 확인했다.