[논문 리뷰] A Conservative Q-Learning approach for handling distribution shift in sepsis treatment strategies

Pramod Kaushik, Sneha Kummetha외 2명|arXiv (Cornell University)|2022. 03. 25.

Sepsis Diagnosis and Treatment인용 수 5

한 줄 요약

이 논문은 분포 이탈을 완화하기 위해 고도로 보수적인 Q-학습(CQL) 기반 접근법을 제안하여, 패혈증 치료에 대한 오프라인 강화학습에서 의사의 임상 실천과의 정렬도를 향상시킨다. 분포를 벗어난 행동의 과대평가를 방지하기 위해 Q-함수를 정규화함으로써, CQL 모델은 특히 중증(고-SOFA) 환자에서 더 임상적으로 정확한 혈관 수축제 및 정맥 주입액 복용량 정책을 학습한다. 이는 표준 DQN보다 의사의 행동을 더 잘 모방하는 데서 유의미한 성능 향상을 보였다.

ABSTRACT

Sepsis is a leading cause of mortality and its treatment is very expensive. Sepsis treatment is also very challenging because there is no consensus on what interventions work best and different patients respond very differently to the same treatment. Deep Reinforcement Learning methods can be used to come up with optimal policies for treatment strategies mirroring physician actions. In the healthcare scenario, the available data is mostly collected offline with no interaction with the environment, which necessitates the use of offline RL techniques. The Offline RL paradigm suffers from action distribution shifts which in turn negatively affects learning an optimal policy for the treatment. In this work, a Conservative-Q Learning (CQL) algorithm is used to mitigate this shift and its corresponding policy reaches closer to the physicians policy than conventional deep Q Learning. The policy learned could help clinicians in Intensive Care Units to make better decisions while treating septic patients and improve survival rate.

연구 동기 및 목표

• 고-SOFA 환자에서 데이터가 희소하고 모델 일반화 능력이 떨어지는 오프라인 강화학습에서의 분포 이탈 문제를 해결하기 위해.
• 특히 중증(고 SOFA) 환자 상태에서 실제 의사의 결정 과정과의 유사도를 높이기 위해 학습된 치료 정책의 정밀도를 향상시키기 위해.
• 희귀하거나 분포를 벗어난 상태(예: 고용량 혈관 수축제 사용)에서 딥 Q-네트워크의 과대평가 편향을 완화하기 위해.
• 임상 실천과 일치하는 보수적이고 신뢰할 수 있는 ICU 의사결정 지원 정책을 개발하여 사망률 위험을 낮추기 위해.

제안 방법

• MIMIC-III v1.4 데이터를 사용하여 패혈증 치료에 적합한 듀얼링 더블-DQN 아키텍처를 변형하였으며, 상태 특징은 매 4시간마다 집계하였다.
• 분포를 벗어난 행동의 과대평가를 방지하기 위해 Q-함수 손실에 보수적 페널티 항목을 적용하였으며, 로그-합-지수(log-sum-exp) 공식을 사용하였다.
• CQL 손실를 다음과 같이 정의하였다: $ L(\theta) = \alpha \mathbb{E}_{s_t \sim D} \left[ \log \sum_a \exp Q_\theta(s_t, a) - \mathbb{E}_{a \sim D} Q_\theta(s_t, a) \right] + L_{\text{DoubleDQN}}(\theta) $, 이는 보수적인 가치 추정을 촉진한다.
• 혈관 수축제 및 정맥 주입액 복용량에 대해 이산 행동 공간을 사용하였으며, 각각 4분위수(0–4)로 나누었으며, 0은 투여 없음을 의미한다.
• 환경 상호작용 없이 고정된 데이터셋에서 학습하여, 역사적 ICU 데이터에 의존하여 최적의 치료 정책을 학습하였다.
• 2D 히스토GRAM을 통해 SOFA 중증도 그룹별 행동 분포를 평가하고, 의사 및 기준 DQN 정책과 비교하였다.

실험 결과

연구 질문

• RQ1 고도로 보수적인 Q-학습은 특히 데이터가 적고 고-SOFA 환자에서의 오프라인 딥 강화학습에서 과대평가 편향을 줄일 수 있는가?
• RQ2 CQL 정책은 다양한 환자 중증도 수준에서 실제 의사의 치료 패턴과 얼마나 잘 일치하는가?
• RQ3 보수적 정규화를 통해 분포 이탈을 줄이면 시뮬레이션된 치료 결과에서 사망률이 낮아지는가?
• RQ4 표준 DQN이 데이터가 희소하여 실패하는 중증 환자(고-SOFA)에 대해 CQL 모델은 일반화 능력을 보일 수 있는가?

주요 결과

• CQL 모델은 의사의 행동을 매우 잘 모방하는 정책을 학습하였으며, 특히 고-SOFA 환자에서만 고용량 혈관 수축제를 투여하는 경향을 보였다. 이는 표준 DQN 모델과는 다름.
• 고-SOFA 상황에서 CQL 모델은 표준 DQN보다 더 높은 혈관 수축제 복용량을 처방하여 임상 실천과 일치하고, 정책의 현실성도 향상시켰다.
• CQL 정책은 처방 복용량과 의사의 복용량 간의 차이가 최소일 때 사망률이 낮게 나타났으며, 특히 중간 및 고-SOFA 그룹에서 두드러졌다.
• CQL 모델은 특히 희귀하거나 중증 환자 상태에서 분포를 벗어난 행동의 Q-값 과대평가를 보완하기 위해 보수적 가치 정규화를 강제함으로써 감소시켰다.
• 2D 행동 히스토GRAM에서 CQL 정책은 모든 SOFA 중증도 수준에서 의사 정책과 매우 유사한 행동 분포를 보이며 뛰어난 행동을 보였다.
• 향후 임상 적용 전 신뢰성 검증을 위해 정량적 오프-폴리시 평가가 필요하다는 점을 이 연구는 규명하였다.