[논문 리뷰] Sepsis World Model: A MIMIC-based OpenAI Gym "World Model" Simulator for Sepsis Treatment

Amirhossein Kiani, Chris Wang외 1명|arXiv (Cornell University)|2019. 12. 15.

Machine Learning in Healthcare참고 문헌 5인용 수 4

한 줄 요약

이 논문은 변분 오토에인코더(VAE)와 혼합 밀도 네트워크-순환신경망(MDN-RNN)을 사용하여 패혈증 치료에서의 환자 상태 전이를 모델링하고 노이즈를 줄이며 불확실성을 포착하는 MIMIC 기반의 OpenAI Gym 시뮬레이터인 Sepsis World Model을 제안한다. 이 모델은 심층 강화학습 정책을 훈련시켜 임상 의사 결정을 더 잘 모방하게 하며, 특히 젖산과 SOFA 점수와 같은 생리적 변화를 기반으로 한 중간 보상 신호를 사용할 경우에 특히 효과적이다.

ABSTRACT

Sepsis is a life-threatening condition caused by the body's response to an infection. In order to treat patients with sepsis, physicians must control varying dosages of various antibiotics, fluids, and vasopressors based on a large number of variables in an emergency setting. In this project we employ a "world model" methodology to create a simulator that aims to predict the next state of a patient given a current state and treatment action. In doing so, we hope our simulator learns from a latent and less noisy representation of the EHR data. Using historical sepsis patient records from the MIMIC dataset, our method creates an OpenAI Gym simulator that leverages a Variational Auto-Encoder and a Mixture Density Network combined with a RNN (MDN-RNN) to model the trajectory of any sepsis patient in the hospital. To reduce the effects of noise, we sample from a generated distribution of next steps during simulation and have the option of introducing uncertainty into our simulator by controlling the "temperature" variable. It is worth noting that we do not have access to the ground truth for the best policy because we can only evaluate learned policies by real-world experimentation or expert feedback. Instead, we aim to study our simulator model's performance by evaluating the similarity between our environment's rollouts with the real EHR data and assessing its viability for learning a realistic policy for sepsis treatment using Deep Q-Learning.

연구 동기 및 목표

• 심각한 치료 분야에서의 패혈증 환자 궤적을 위한 현실적이고 미분 가능한 시뮬레이터를 개발하기 위해.
• 변분 오토에인코더(VAE)를 사용하여 전자적 건강기록(EHR) 데이터의 노이즈를 줄이고 보정된 잠재 표현을 학습하기 위해.
• 다중 정규분포 혼합 모델을 사용한 MDN-RNN을 통해 다음 상태에 대한 확률 분포를 예측함으로써 환자 상태 전이의 불확실성을 모델링하기 위해.
• 심층 Q-네트워크(DQN)를 사용하여 다양한 보상 구성 방식을 적용했을 때 시뮬레이터가 임상적으로 타당한 치료 정책을 훈련시킬 수 있는지 평가하기 위해.
• 젖산과 SOFA 점수와 같은 중간 임상 지표를 기반으로 한 시간에 따라 변화하는 보상 신호를 통합함으로써 정책의 다양성과 현실성을 향상시키기 위해.

제안 방법

• 모델은 46개의 정규화된 EHR 기능을 30차원의 잠재 공간으로 압축하기 위해 세 개의 밀집 층으로 구성된 VAE를 사용하며, 평균 제곱 오차를 최소화한다.
• VAE에서 생성된 잠재 상태 표현은 액션과 은닉 상태에 조건부인 다중 정규분포 혼합 모델을 사용하여 다음 잠재 상태의 확률 분포를 모델링하는 MDN-RNN에 입력된다.
• 각 시간 단계에서 예측된 혼합 분포에서 샘플링을 통해 확률적 환자 상태 전이를 시뮬레이션하며, 온도 조절을 통해 불확실성 수준을 조절한다.
• 시뮬레이터에는 상태, 종료, 결과 예측 모델이 통합되어 있으며, 전체 시스템은 강화학습 훈련을 위해 OpenAI Gym 환경으로 노출된다.
• 세 가지 보상 구성 방식을 평가한다: (1) 희박한 종료 시점 보상, (2) 벌점이 있는 지연 보상, (3) 젖산과 SOFA 점수 변화를 기반으로 한 밀도 높은 중간 보상.
• 심층 Q-네트워크(DQN) 에이전트를 시뮬레이션 환경에서 훈련시키며, 그로 얻어진 정책을 실제 의사 행동의 행동 분포, 에피소드 길이, 누적 보상 측면에서 비교한다.

실험 결과

연구 질문

• RQ1 VAE와 MDN-RNN 기반의 월드 모델은 노이즈가 많은 EHR 데이터로부터 패혈증 환자 궤적의 잠재 역학을 효과적으로 정제하고 표현할 수 있는가?
• RQ2 확률적 MDN-RNN을 통해 불확실성을 모델링할 경우, 결정론적 RNN 기반 모델 대비 더 현실적이고 다양한 치료 정책 학습이 가능한가?
• RQ3 특히 밀도 높은 중간 보상 전략이 DQN으로 학습된 패혈증 치료 정책의 현실성과 다양성에 어떤 영향을 미치는가?
• RQ4 시뮬레이터가 학습한 정책이 실제 의사 행동에서 관찰된 행동 분포, 에피소드 길이, 보상 궤적을 어느 정도 재현하는가?
• RQ5 이 시뮬레이터가 극단적이거나 단조로운 조치에 과적합되지 않는, 임상적으로 타당하고 시간에 민감한 치료 정책을 훈련시키는 데 기여할 수 있는가?

주요 결과

• 노이즈가 많은 EHR 데이터를 다룰 때 특히, 결정론적 RNN 기반 모델 대비 VAE+MDN-RNN 모델이 더 안정적이고 현실적인 상태 궤적을 생성했다.
• 희박한 종료 시점 보상만을 사용할 경우 정책 붕괴 현상이 발생하여 에이전트가 단일 행동에 집중하고 매우 짧은 에피소드를 생성했다.
• 젖산과 SOFA 점수 변화를 기반으로 한 중간 보상은 행동 분포의 다양성과 현실성을 높였으며, 의사 행동과 유사하게 나타났다.
• 중간 보상을 사용한 모델는 더 균형 잡힌 에피소드 길이와 더 나은 행동 다양성을 확보하여 시간에 따른 정책 학습이 향상됨을 시사했다.
• 개선에도 불구하고 종료 모델은 여전히 과적합되어 보상 구성 전략이 향상되었음에도 불구하고 일관되게 짧은 에피소드 길이를 보였다.
• 시뮬레이터를 통해 DQN 훈련이 가능했으며, 특히 실시간 임상 변화를 반영한 보상 신호를 사용할 경우 기존 기반 모델 대비 정책의 현실성이 높아졌다.