LightGBM 알고리즘을 활용한 예측모델 구축과 성능 비교

라인웍스에서는 Electronic Health Record (이하 EHR) 데이터를 이용하여 다양한 머신러닝 프로젝트를 진행하고 있습니다.

이번 글에서는 의료 인공지능 개발 프로젝트의 성능 개선을 위해 사용한 LightGBM 알고리즘과 수많은 Feature를 줄이기 위한 Feature Selection 방법을 소개합니다. 추가로, MDwalks EXI에서 선보였던 Deep Learning을 사용한 MIMIC-III 30일 이내 재입원 예측 모델(https://linewalks.com/archives/6142)과 LightGBM 사용 모델의 성능을 비교해보겠습니다.

*이번 글은 라인웍스 여름 인턴 프로그램에 참여한 이서호님의 아이디어와 실험 내용을 바탕으로 작성되었습니다.


1. 들어가기

MIMIC-III

모델 구축에 사용한 데이터는 MIMIC-III 데이터베이스입니다. EHR 데이터 중 거의 유일하게 공개된 다년간의 중환자입원 기록으로, 46,520명 환자들의 58,976건의 입원 정보를 포함하고 있습니다. MDwalks EXI로 구축한 30일 재입원 예측 모델(https://linewalks.com/archives/6142)도 동일한 데이터를 사용했습니다.

*MIMIC-III 데이터베이스에 대한 상세 정보는 공식페이지(https://mimic.physionet.org/about/mimic/)를 참조하세요.
*MIMIC-III 데이터베이스를 세팅하는 방법은 라인웍스 블로그 글(https://linewalks.com/archives/6198)을 참조하세요.

LightGBM

LightGBM은 Microsoft에서 발표하고 제공하는 오픈소스 그래디언트 부스팅 결정트리 프레임워크입니다. 기존의 그래디언트 부스팅 결정 트리 프레임워크의 경우 고차원 변수가 포함된 큰 데이터에서 정확도와 효율성이 떨어지는 경향이 있는데, 이런 단점을 Gradient-based One-Side Sampling(GOSS)과 Exclusive Feature Bundling(EFB)라는 두 가지 기술을 이용하여 보완한 프레임워크입니다.

*자세한 내용은 아래 링크를 참고하세요.
*관련 논문 (https://papers.nips.cc/paper/6907-lightgbm-a-highly-efficient-gradient-boosting-decision-tree.pdf)
* LightGBM Documentation (https://lightgbm.readthedocs.io/en/latest/)


2. 데이터 전처리

MIMIC-III 데이터에는 수십개의 테이블이 있습니다. 실제 모델에 사용하기엔 너무 많은 양이라 의료 전문가들의 조언을 얻어 가장 중요하다고 생각하는 일부 테이블을 우선으로 선정하여 사용하였습니다.

사용한 데이터는 다음과 같습니다.

  • 환자 방문 정보
    • 방문 유형, 이전 방문 횟수, 입원 기간, 중환자실 이용 기록 등
  • 환자 기본 정보
    • 나이, 성별 등
  • 진단
    • 해당 방문에서 환자가 진단받은 내역
    • 874개 진단 코드 (전체 데이터에서 100회 이상 사용된 코드)
  • 의료행위
    • 해당 방문에서 수행된 의료행위 내역
    • 258개 의료행위 코드 (전체 데이터에서 100회 이상 사용된 코드)
  • 의약품
    • 해당 방문에서 처방받은 의약품 내역
    • 366개 의약품 코드 (전체 데이터에서 1,000회 이상 사용된 코드)
  • Lab 검사
    • 해당 방문에서 검사한 검사 내역
    • 154개 Lab 검사 코드 (전체 데이터에서 1,000회 이상 사용된 코드)
  • 각종 이벤트 (투여, 배출, 차트 등등)
    • 2,956개 이벤트 코드

코드와 범주형 데이터는 One-Hot Encoding 하였고, 수치형 데이터는 Normalization (표준화)하지 않고 사용했습니다. 모델에 입력된 최종 Feature는 총 4,903개입니다.

46,520명의 환자를 5 Fold로 나누어 Cross Validation을 수행하였고, 전체 방문 데이터는 총 58,976건으로 각 Fold별로 11,794 / 11,860 / 11,818 / 11,687 / 11,817건으로 나누었습니다.


3. 실험 결과

위에서 처리한 데이터를 이용하여 LightGBM 모델을 학습시켰습니다. 그리고 모델간의 성능 비교를 위해 MDwalks EXI에서 사용했던 “Neural Network(Deep Learning) 모델”과 LightGBM과 같은 Tree형 모델이지만 다른 Ensemble 모델을 사용하는 “Random Forest 모델”을 추가로 학습하였습니다. 그리고 각 모델별로 30일 이내 재입원 여부를 예측해보았습니다.

예측력 평가에 사용한 Metric은 AUROC입니다. 모델별 결과는 다음과 같습니다.

LightGBM이 모든 Fold에서 Neural Network와 Random Forest보다 좋은 성능을 보여주고 있습니다.


4. Feature Selection을 통한 성능 개선

 Feature Importance

Neural Network와 달리 Tree형 모델에서는 모델이 어떤 Feature를 중요하게 보는지에 대한 정보(Feature Importance)를 쉽게 얻어낼 수 있습니다. LightGBM도 이를 쉽게 추출하도록 구현이 되어있습니다.

위에서 학습된 LightGBM모델의 Feature Importance을 추출해 상위 10개만 그래프로 그려보았습니다.

 Importance가 가장 높은 Feature는 “ICU(중환자실) 입원 기간(Length of Stay in ICU)”입니다. 비슷한 Feature인 “전체 입원 기간(Length of Stays)”도 10위 안에 포함되어 있습니다. 나이(Age)와 사망 여부(Dead)가 2위와 3위를 차지하고 있고, 이전 365일간 방문횟수(Visit count of last 365 day), 이전 180일간 방문횟수(Visit count of last 180 day)도 포함되어 있습니다. 즉, 이전에도 병원에 방문을 많이 했다면 앞으로도 방문할 가능성이 높다고 유추해볼 수 있습니다. drg_mortality와 drg_severity는 청구 비용을 나타내는 Feature인데, 많은 비용이 청구되었을 수록 재입원 가능성이 높을 것이란 판단도 가능할 것 같습니다. 특정 진단 코드와 차트 검사 코드도 포함되어 있습니다. 사람이 이해할 수 있는 방향으로 모델이 잘 학습되었다고 볼 수 있겠습니다.

이번엔 Feature Importance 값 자체를 보겠습니다. 전체 값들의 간단한 통계 수치는 다음과 같습니다.

  • 전체 수 : 4,903
  • 평균 : 1.2033
  • 분산 : 94.4078
  • 최소값 : 0
  • 최대값 : 347.4
  • 중앙값 : 0

중앙값이 0이라는 점이 눈에 띕니다. Feature 중 50% 이상의 Importance가 0이라는 얘기입니다. 확인을 해보았더니 857개를 제외하고 나머지 82%에 해당하는 4,046개 Feature의 Importance가 모두 0이었습니다. 모델이 이렇게 많은 Feature를 중요하지 않다고 판단했다면, 굳이 모든 Feature들을 사용할 필요가 없어보입니다. Feature 일부를 제거하여 학습하면 계산량과 메모리 사용량을 줄이고 더 빠르게 학습 결과를 도출할 수 있을 것으로 예상됩니다.

Feature Selection

Feature Selection은 성능에 더 많은 영향을 주는 Feature를 우선적으로 선택하여 학습하는 방법입니다. 상위 몇 개 혹은 몇 %를 선택하는 방법, 반복적으로 하위 Feature를 제거해 나가는 방법 등 다양한 방법이 존재합니다.

우선 LightGBM모델에서 Feature Importance가 0인 값만 제거하여 실험해보았습니다. 성능에는 큰 영향을 주지 않았지만, 전체 중 82%의 무의미한 Feature의 값이 제거되었기 때문에 LightGBM이 학습해야 하는 Feature의 수가 줄어들면서 학습 속도가 3배 이상 향상했습니다. 

이번에는 학습된 LightGBM의 Feature Importance를 이용하여 선택된 Feature를 Neural Network 모델에도 적용하여 실험해보았습니다. Neural Network는 Hidden Layer 때문에 Feature Importance를 쉽게 추출할 수 없으므로 다른 모델의 Feature 정보를 이용하면 성능을 더 올릴 수 있지 않을까하는 생각이었습니다.

그 결과,
기존 Neural Network 모델의 AUROC는

0.7242 ± 0.0092

Feature Selection이 적용된 Neural Network 모델의 AUROC는

0.7363 ± 0.0181

로 성능이 1% 증가하였습니다. LightGBM의 성능(0.7444 ± 0.0123)에는 미치지 못하지만, 4배 이상의 학습 속도 향상은 물론 성능까지 개선되는 효과를 볼 수 있었습니다.


5. 마치며

이번 글에서는 재입원 예측에 LightGBM 알고리즘을 적용하고 Neural Network 모델과 성능을 비교해보았습니다. 추가로 Tree 모델의 장점이라 할 수 있는 Feature Importance를 이용해 유의미한 Feature를 추출하고 Neural Network에 적용해 보았습니다. 

LightGBM 적용을 통해 Gradient Boosting 모델의 학습 속도와 성능의 향상을 확인할 수 있었고, 간단한 Neural Network 모델에 비해 LightGBM 모델이 재입원 예측에선 더 좋은 성능을 보임을 확인했습니다. 또한, Gradient Boosting 모델의 Feature Importance를 이용하여 Neural Network 등 다른 모델들의 예측 성능까지 향상시킬 수 있음을 알게 되었습니다.


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


No comments yet.

Leave a comment

Your email address will not be published.