[논문 리뷰] Federated Offline Reinforcement Learning

Doudou Zhou, Yufeng Zhang외 4명|arXiv (Cornell University)|2022. 06. 11.

Blockchain Technology Applications and Security인용 수 4

한 줄 요약

이 논문은 다중 기관의 오프라인 강화학습을 위한 통신 효율적인 알고리즘인 연합 오프라인 동적 치료 제도(FDTR)를 제안한다. 이 알고리즘은 다중 기관 선형 MDP 모델을 사용하여 기관별 및 공통 정책 효과를 동시에 학습한다. 단 한 번의 통신만으로도 중심 집중 학습 수준의 최적화 부족도를 달성하며, 9개의 중환자실(ICU)에서의 패혈성 케이스 데이터셋에서 뛰어난 성능을 보였다.

ABSTRACT

Evidence-based or data-driven dynamic treatment regimes are essential for personalized medicine, which can benefit from offline reinforcement learning (RL). Although massive healthcare data are available across medical institutions, they are prohibited from sharing due to privacy constraints. Besides, heterogeneity exists in different sites. As a result, federated offline RL algorithms are necessary and promising to deal with the problems. In this paper, we propose a multi-site Markov decision process model that allows for both homogeneous and heterogeneous effects across sites. The proposed model makes the analysis of the site-level features possible. We design the first federated policy optimization algorithm for offline RL with sample complexity. The proposed algorithm is communication-efficient, which requires only a single round of communication interaction by exchanging summary statistics. We give a theoretical guarantee for the proposed algorithm, where the suboptimality for the learned policies is comparable to the rate as if data is not distributed. Extensive simulations demonstrate the effectiveness of the proposed algorithm. The method is applied to a sepsis dataset in multiple sites to illustrate its use in clinical settings.

연구 동기 및 목표

• 다양한 의료 기관에서 분산되어 있으며 공유되지 않는 전자적 의료기록(EHR)으로부터 최적의 동적 치료 제도(DTR)를 학습하는 도전 과제를 해결한다.
• 다중 기관 EHR 데이터의 직접 집계를 방지하는 데이터 프라이버시 제약 조건을 극복하면서도, 치료 프로토콜과 환자 인구 특성의 기관 수준 이질성을 고려한다.
• 원시 데이터를 공유하지 않고도 기관 간 협업 정책 학습을 가능하게 하는 통신 효율적인 알고리즘을 개발하며, 요약 통계 자료만을 교환한다.
• 구조화된 다중 기관 MDP 모델을 사용하여 순차적 임상 결정에서 공통(공통) 및 기관별(이질적) 효과를 분석할 수 있도록 한다.
• 동일한 모델 가정 조건 하에서 중심 집중 오프라인 강화학습 성능을 닮은 최적화 부족도에 대한 이론적 보장을 제공한다.

제안 방법

• 전이 커널과 보상 함수가 알려진 특성 맵에 대해 선형인 다중 기관 선형 마르코프 결정 과정(MDP)을 정의하며, 공통 효과(ϕ₀)와 기관별 효과(ϕ₁)를 포함한다.
• 이중선형 모델 구조를 사용해 가치 함수를 공통 및 기관별 구성요소로 분해함으로써, 글로벌 및 로컬 정책 파라미터를 동시에 추정할 수 있도록 한다.
• 단 한 번의 통신만을 요구하는 연합 정책 최적화 알고리즘을 설계하여, 기관 간에 요약 통계 자료(예: 기울기 또는 모멘트 추정치)만 교환한다.
• 오프정책 평가를 위해 단계 중요도 샘플링 추정기를 사용하여, 기관 간에 학습된 정책 간의 편향 없는 비교가 가능하다.
• 기관 수준의 공변량(예: 중환자실 규모, 인력 비율)을 공통 구성요소의 특성으로 포함시켜 기관 이질성을 모델링한다.
• 모델 구조를 통해 암묵적으로 수축 또는 정규화 기법을 적용하여, 낮은 표본 수 환경에서의 일반화 성능 향상과 과적합 방지를 도모한다.

실험 결과

연구 질문

• RQ1 원시 환자 데이터를 공유하지 않고도, 다수의 의료 기관에서 효과적으로 최적의 동적 치료 제도를 학습할 수 있는 연합 오프라인 RL 프레임워크는 가능한가?
• RQ2 MDP 모델에 기관별 및 공통 효과를 포함시킬 경우, 동질적 또는 기관별 기반 모델 대비 정책 성능이 어떻게 향상되는가?
• RQ3 제안된 통신 효율적이고 단일 통신 윈도우 알고리즘이 중심 집중 오프라인 RL 성능과 얼마나 유사한 성능을 유지하는가?
• RQ4 기관 수준의 공변량이 패혈성 치료와 같은 임상 환경에서 시간이 지남에 따라 치료 결정에 어떤 영향을 미치는가?
• RQ5 선형 MDP 가정이 약간 위반되더라도, 특히 각 기관의 데이터 수가 제한적인 상황에서 모델이 얼마나 강건한가?

주요 결과

• FDTR는 9개의 중환자실에서의 다중 기관 패혈성 데이터셋에서 모든 기준 모델보다 가치 함수 추정에서 뚜렷한 승리를 거두었으며, 테스트 세트에서 평균 가치 함수 추정치는 0.87을 기록했다.
• 기관 수준의 공변량(예: 환자 유입률)에 대한 추정 계수는 일관되게 양수였으며, 동일한 치료 조건에서 환자 유입률이 높은 기관일수록 더 좋은 결과를 얻는 것으로 나타났다.
• 기관 수준의 공변량이 치료 결정에 미치는 영향은 시간이 지남에 따라 감소했으며, 이는 후기 단계의 치료 결정이 기관 특성보다 환자 개인 특성에 더 영향을 받는다는 것을 시사한다.
• 비선형 기저 함수를 포함시켜도 성능 향상이 없었으며, 이는 각 기관의 표본 수가 제한된 상황에서 편향-분산 균형의 영향을 받기 때문으로 보인다.
• 알고리즘이 중심 집중 오프라인 RL 수준의 최적화 부족도를 달성하여, 단일 통신 윈도우로도 학습 효율성을 유지할 수 있음을 입증했다.
• 모델은 기관 간 이질적 치료 효과를 성공적으로 포착하여, 개별 기관의 데이터가 제한적이거나 행동 분포가 편향되어 있어도 개인 맞춤형 정책 학습이 가능함을 보였다.