CDM AI Challenge: Forecasting AKI 1위 후기

2021.12.21

라인웍스는 직원 3명으로 구성된 티미룸이라는 팀으로 서울대학교병원에서 주최하는 CDM AI Challenge: Forecasting AKI에 참가하였습니다. 10개의 팀을 선별한 뒤, 1주일간 진행되는 대회였습니다.

티미룸 팀은 이 대회에서 1위라는 성적을 거두었습니다. CDM 형태의 정형화된 EHR(Electronic Health Records) 데이터를 다루는 대회였기에, 라인웍스가 그동안 연구하고 공개해왔던 기술들을 적용해볼 수 있는 기회였습니다. 1주일이라는 짧은 시간 동안 어떤 접근 방법을 통해 1위를 달성할 수 있었는지 공유해 드리려고 합니다.

접근 방법

최종적으로 결정된 데이터 처리부터 학습까지의 과정은 다음과 같습니다.

Baseline 모델 설정

먼저 Baseline이 될 모델을 선택했습니다. 라인웍스는 EHR 데이터를 이용한 모델들에 대한 경험이 많았기에 이를 선택하는 것은 어렵지 않았습니다. 기존 연구들에서 대체로 좋은 성능과 빠른 학습 속도를 보였던 LightGBM을 Baseline 모델로 선택했고, 기존의 연구에서 뛰어난 성능을 보였던 하이퍼패러미터를 그대로 가지고 왔습니다.

실험 결과, 해당 모델이 어느 정도 준수한 성능을 보이는 걸로 나타났고, 이를 이용해 처리된 데이터들을 추가한 뒤, 성능을 보며 데이터 처리에 대한 검증을 진행했습니다.

Feature Extraction

임상적 의미

의료 데이터는 비슷한 의미를 가진 데이터들이 더 세분화되어 여러 개로 나뉘어 있는 경우가 있습니다. 아주 쉬운 예를 들면, 감기라는 개념이 목감기, 코감기, 열감기 등으로 나뉘어 있는 것입니다. 이렇게 나뉘어 있는 데이터들이 중요한 의미를 갖는 경우(목감기냐, 코감기냐가 중요한 경우)도 있지만, 해당 개념들을 모두 묶어서 하나로 보는 것이 더 중요한 경우(감기에 걸렸냐, 안 걸렸냐가 중요한 경우)도 있습니다. 이러한 의미를 파악하여, 나뉜 개념들을 묶어서 입력했을 때 예측력이 좋아지기도 합니다.

이번 대회의 문제는, 수술 후, 48시간 이내, 급성 신손상(Acute kidney injury, 이하 AKI), 중증 급성 신손상(Sepsis-induced Acute Kidney Injury, 이하 SAKI) 발생을 예측하는 것이었습니다. 이와 관련된 논문을 조사하여 급성 신손상 발생에 영향을 주는 요소들을 찾고, 참고할만한 Risk Score들을 선정했습니다.

데이터에서 찾는 의미

EHR 데이터는 환자의 상태에 대한 의료진들의 행동이 담겨 있는 데이터입니다. 이 때문에 데이터 안에도 임상적인 의미가 담겨져 있고, 이를 찾아내 이용할 수 있습니다.

예측 시점 이전에 받은 진단, 처방받은 의약품, 검사받은 항목, 받은 의료 행위 등 환자에게 발생한 의료 이벤트들이 의미가 있을 것으로 생각했습니다. Baseline 모델 실험을 통해, 예측 시점 이전 90일 이내 발생한 이벤트들을 사용하기로 했습니다. 90일 이내 발생한 이벤트더라도 데이터 내에 일정 수 이하로 발생해 모델에 충분한 정보를 줄 수 없는 이벤트는 제외했습니다.

이벤트 분류에 따라 실험을 통해, 단순히 이벤트 수행 여부만을 사용할지, 이벤트 수행 횟수를 사용할지를 결정했습니다. 검사처럼 수치로 표현이 되는 분류의 경우는 해당 기간 발생한 검사 수치의 통계치를 입력으로 활용했습니다.

모델 선택

1주일이라는 짧은 시간이었기 때문에, 경험적인 부분을 살려서 익숙한 모델들을 활용하기로 했습니다. Baseline 모델로 활용했던 LightGBM과 간단한 Neural Network 모델을 학습해 두 모델을 앙상블 하기로 결정했습니다. 두 개의 모델의 알고리즘이 완전 다른 형태이기 때문에, 데이터에서 서로 다른 의미를 뽑아낼 수 있을 것으로 생각했습니다.

Neural Network 모델은 LightGBM에 비해 Noise에 약하고 데이터 전처리 방법에 따른 성능 차이가 큽니다. 이를 보완하기 위해 다른 머신러닝 모델을 이용해 추가적인 처리를 해주었습니다.

Neural Network를 위한 처리

Feature Elimination

라인웍스 블로그에 올라왔던 글(LightGBM 알고리즘을 활용한 예측모델 구축과 성능 비교)에서도 보여줬던 방법입니다. 학습된 LightGBM의 Feature Importances를 확인하여 모델에게 주는 영향이 없거나 미미한 Feature를 제거해주는 방법입니다. 이를 이용해 Neural Network의 성능이 올랐던 경험이 있어, 이번 모델에도 적용해보았습니다.

Matrix Factorization Embedding

이 방법 또한, 블로그에 올라왔던 글(의료 인공지능 성능 개선을 위한 시도 – Matrix Factorization Embedding)에서 보여줬던 방법입니다. Embedding을 적용하는 것은 데이터 안에서 의미를 파악해 모델에 더 상세한 의미를 주는 방법입니다. Embedding을 만들기 위한 여러 알고리즘이 있지만, Matrix Factorization을 선택한 이유는 Embedding을 만들기 위한 추가 전처리가 필요하지 않고, 예측 모델에 입력되는 데이터를 그대로 이용할 수 있기 때문입니다. Matrix Factorization을 선택하여 구현을 위한 시간을 아낄 수 있었고, 성능 향상 효과도 볼 수 있었습니다.

모델 학습

AKI와 SAKI. 두 개의 예측 모델

이 대회에서 예측해야 할 것은 AKI와 SAKI, 두 가지였습니다. 경험상, 이벤트 간의 연관성이 확실하지 않다면, 대부분 여러 이벤트를 동시에 예측하는 모델은 성능이 좋지 않았습니다. AKI와 SAKI가 어떤 연관이 있는지, 둘의 주요 영향 factor가 무엇인지에 대한 임상적인 정보가 없었기 때문에, 두 개의 모델을 따로 학습하기로 결정했습니다.

두 개의 모델에 입력되는 데이터와 모델의 하이퍼패러미터는 시간 문제와 구현의 편의상 동일하게 가져가기로 했습니다.

아쉬웠던 점

1주일이라는 짧은 시간 동안 진행되었기 때문에, 포기한 것들이 있었습니다.

Time-Series 활용

수술 후, 48시간 이내라는 짧은 시간의 이벤트를 예측하는 문제였기 때문에, 수술 중 일어난 일들이 큰 영향을 끼쳤을 가능성이 높다고 생각했습니다. Recurrent Neural Network(이하 RNN) 모델을 이용해서 짧은 시간 동안 일어난 일들의 순서를 고려하여 입력할 수 있었으면 더 좋은 성능이 나오지 않았을까 생각이 듭니다.

하지만 RNN 모델은 데이터를 처리하는 과정이 상대적으로 복잡하고 학습 시간도 오래 걸리기 때문에 초반 접근 방법 설계에서 배제했습니다.

수술 중, Vital Sign(활력 징후) 데이터 활용

수술 중의 Vital Sign 테이블이 있었기에 이를 이용하여 모델에 입력해보았고, 교차 검증 과정에서 성능이 상승하는 것도 확인했습니다. 하지만, 테스트 데이터를 로드하는 과정에서 pandas를 이용하는 경우, CSV 포맷과 관련된 에러가 발생했습니다. CSV를 직접 읽어오는 코드를 구현할 수도 있었지만, 테스트 데이터에는 직접 접근할 수가 없었기 때문에 시간이 오래 걸리는 일이 될 거라고 생각해서, 바로 포기했습니다.

AKI, SAKI 다른 데이터 처리

AKI와 SAKI의 예측 모델 학습 시, 편의상 같은 데이터와 하이퍼패러미터를 사용하기로 했다고 위에서 언급했습니다. 하지만, 실제 데이터를 추가하며 학습해볼 때, 특정 데이터들에서 두 가지 모델에 다른 결과를 보이는 경우(데이터를 추가하면 AKI 예측의 성능은 상승하지만, SAKI 예측의 성능은 하락한다거나 등등)가 발생했습니다. AKI와 SAKI를 각각 예측하기 위한 최적의 데이터는 다르다고 인지하였지만, 이를 위한 실험 시간과 개별 데이터 로더를 구현하는 구현 시간, 각각 데이터를 불러옴으로써 발생하는 학습 시간 증가 등을 고려하여 성능의 하락을 감수하기로 했습니다.

마치며

의료 데이터를 이용하는 AI 대회들이 점점 많아지고 있습니다. 하지만 아직 대부분 이미지나 영상, 신호, 음성 등 비정형 데이터에 집중되어 있습니다.

정형화된 EHR 데이터는 환자의 상태를 보여주는 동시에 의료진이 치료한 기록도 포함하는 데이터입니다. 때문에, 급성 또는 만성으로 진행되는 상태를 잘 예측하고 분석하는 AI를 만드는데 좋은 재료로 사용될 데이터입니다.  그 가능성을 다양한 질환에서 확인할 수 있도록, 앞으로도 EHR 데이터를 이용하는 대회가 점점 더 많아지길 바랍니다.

라인웍스에서는 EHR 데이터를 포함한 의료 분야에서 사용되는 다양한 데이터를 이용한 머신러닝 연구를 계속 진행하고 있습니다. 저희와 함께 연구를 진행할 머신러닝 엔지니어를 모시고 있습니다. 자세한 내용은 채용 공고(https://linewalks.com/recruit)를 확인해주세요.

Jungwoo Chae

Engineer

Jungwoo Chae