[논문 리뷰] Deep Reinforcement Learning for Multi-class Imbalanced Training

Jenny Yang, Rasheed El-Bouri외 4명|arXiv (Cornell University)|2022. 05. 24.

Imbalanced Data Classification Techniques인용 수 6

한 줄 요약

이 논문은 임상 데이터셋에서 다중 클래스 데이터 불균형 문제를 해결하기 위해 듀얼링 및 더블 DQN 아키텍처를 사용한 딥 강화학습 프레임워크를 제안한다. 또한 맞춤형 보상 함수를 도입하여 소수의 클래스 탐지 성능을 크게 향상시키면서도 모든 클래스에서 균형 잡힌 성능을 유지한다. 이는 최신 기술들보다 뛰어난 성능을 보인다.

ABSTRACT

With the rapid growth of memory and computing power, datasets are becoming increasingly complex and imbalanced. This is especially severe in the context of clinical data, where there may be one rare event for many cases in the majority class. We introduce an imbalanced classification framework, based on reinforcement learning, for training extremely imbalanced data sets, and extend it for use in multi-class settings. We combine dueling and double deep Q-learning architectures, and formulate a custom reward function and episode-training procedure, specifically with the added capability of handling multi-class imbalanced training. Using real-world clinical case studies, we demonstrate that our proposed framework outperforms current state-of-the-art imbalanced learning methods, achieving more fair and balanced classification, while also significantly improving the prediction of minority classes.

연구 동기 및 목표

• 소수의 질병 클래스가 부족하게 나타나는 임상 기계학습 환경에서의 다중 클래스 데이터 불균형 문제를 해결하기 위해.
• 클래스 불균형에 기반해 동적으로 학습을 조정하는 강화학습 기반의 학습 프레임워크를 개발하여 소수 클래스 예측 성능을 향상시키기 위해.
• 기존의 단일 클래스 RL 접근 방식을 실제 임상 데이터를 활용한 다중 클래스 환경으로 확장하기 위해.
• SMOTE, 비용 감안 학습, 앙상블 모델 등을 포함한 최신 기술들과의 비교를 통해 방법의 성능을 평가하기 위해.
• 희귀 임상 진단에 특히 유리한 공정성과 모든 클래스 간 균형 잡힌 성능 향상을 입증하기 위해.

제안 방법

• 학습 안정성과 가치 함수 추정 향상을 위해 듀얼링 및 더블 딥 Q넷(DQN) 아키텍처를 활용한다.
• 소수 클래스의 정확한 분류를 장려하기 위해 클래스 빈도의 역수 비례 보상이 부여되는 맞춤형 밀도 보상 함수를 설계한다.
• 에피소드 기반 학습 절차를 구현하며, 각 에피소드는 학습 에포크에 해당하고, 에이전트(모델)는 누적 보상을 최대화하기 위해 행동(학습 업데이트)을 선택한다.
• 클래스 인식 탐색과 우선순위 경험 재현을 통합하여 불균형 데이터에서의 샘플 효율성과 수렴 속도를 향상시킨다.
• 각 클래스를 에이전트의 행동 공간 내 별도의 결정 목표로 간주함으로써 다중 클래스 분류를 지원한다.
• 클래스 분포 변화에 따라 정책을 조정함으로써 데이터 재샘플링이나 클래스 가중치 조정 없이도 적응형 학습이 가능하다.

실험 결과

연구 질문

• RQ1 기존의 데이터 수준 및 알고리즘 수준의 접근 방식과 비교해 볼 때, 강화학습 프레임워크가 다중 클래스 불균형 임상 데이터셋에서 분류 성능 향상에 효과적으로 기여할 수 있는가?
• RQ2 기본 모델이 실패하는 상황에서, 제안된 RL 방법은 희귀 임상 상태 탐지에 얼마나 잘 작동하는가?
• RQ3 맞춤형 보상 함수가 주로 클래스 성능에 영향을 주지 않으면서도 소수 클래스 학습에 얼마나 효과적인가?
• RQ4 수십 개의 진단 코드를 포함한 다중 클래스 문제에 대해 이 프레임워크는 어떻게 확장되며, 클래스 이질성은 성능에 어떤 영향을 미치는가?
• RQ5 비용 감안 학습이나 SMOTE와 같은 데이터 증강 기법과 비교했을 때, RL 기반 방법이 G-mean, F-score와 같은 균형 잡힌 성능 지표에서 더 뛰어난 성능을 보이는가?

주요 결과

• 제안된 RL 프레임워크는 모든 테스트 세트에서 G-mean 점수를 최고 또는 2위로 기록하여 모든 클래스 간 균형 잡힌 성능이 뛰어나다는 것을 보여준다.
• 이진 COVID-19 예측 작업에서 RL 모델은 평균 민감도 0.703(CI: 0.539–0.785)을 기록하여 기준 신경망(0.236)과 XGBoost(0.288)보다 유의미하게 뛰어난 성능을 보였다.
• 다중 클래스 진단 작업에서 RL은 개별 및 평균 소수 클래스 F-score 모두에서 가장 뛰어난 성능를 기록하여 다양한 임상 조건에서의 강건성을 입증했다.
• 비용 감안 가중치를 추가하면 기준 모델의 민감도는 향상되지만 주류 클래스 정확도는 감소하는 반면, RL은 모든 지표에서 높은 성능를 유지했다.
• SMOTE 증강은 최소한의 성능 향상을 가져왔고, 특히 신경망에서는 노이즈가 많거나 겹치는 합성 샘플로 인해 성능이 떨어지는 경우가 있었다.
• 모델은 급성 신장 이상을 예측하는 데 가장 어려움을 겪었으며, 이는 입력 특징에 혈액 검사나 소변 검사와 같은 핵심 진단 마커가 누락되어 있기 때문일 것이다.