의료 인공지능 성능 개선을 위한 시도 – 온톨로지 임베딩

2020.07.17

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

이번 글에서는 의료 인공지능 개발 프로젝트의 성능 개선을 위해 코드 체계를 이용한 온톨로지 임베딩을 소개합니다. 지난 글(https://blog.linewalks.com/archives/6408)에서 소개한 일반 Skip-grams 모델과의 차이점을 알아보고, 재입원 예측 모델에 적용한 성능을 비교합니다.

이 글은 기술 설명 위주로 작성했습니다. 임베딩에 대한 기초적인 개념은 지난 글을 참고 바랍니다.

질병 이름은 편의를 위해 구글 번역기를 이용해 번역을 같이 표시했습니다. 번역이 정확하지 않을 수 있습니다.


코드 체계

대부분의 코드는 관리의 편의성을 이유로 각자의 체계가 있습니다. 코드를 아무 의미 없이 단순 나열한 것이 아니라, 정해진 기준에 따라 분류되어지고 상하위 개념이 있는 트리나 그래프 구조로 이루어지는 경우가 많습니다.

지난 글에서 세계 보건 기구에서 발표하는 ICD-9(국제질병사인분류 9차 개정판)에 대해 간단히 언급했습니다. ICD-9는 분류에 따라 다르지만 총 5단계로 나누어져 있습니다.

ICD-9 코드 중, 487.0 Influenza with pneumonia (폐렴이 있는 인플루엔자)에 대해서 자세히 알아보겠습니다. 이 코드는 총 4단계의 분류로 이루어져 있습니다. 이 코드의 상위 분류는 다음과 같습니다.

  • 460–519 Diseases of the respiratory system (호흡 체계의 질병)
  • 480–488 Pneumonia and influenza (폐렴 및 인플루엔자)
  • 487 Influenza (인플루엔자)

1단계 분류인 Diseases of the respiratory system는 6개의 2단계 분류로 나뉘어집니다. 그 중 하나인 Pneumonia and influenza 또한 8개의 3단계 분류로 나뉘어지고, Influenza가 그 중 하나 입니다. Influenza는 4단계 분류인 487.0 Influenza with pneumonia 487.1 Influenza with other respiratory manifestations로 나누어집니다.

이 코드 체계를 통해서 487.0 Influenza with pneumonia487.1 Influenza with other respiratory manifestations과 아주 비슷한 유형의 질병이란 것을 알 수 있고, Pneumonia and influenza에 속하는 다른 질병들과도 어느 정도 유사하다는 것을 알 수 있습니다.

EHR 데이터의 코드

EHR 데이터에 입력되어 있는 코드는 어떨까요? MIMIC-III 데이터에서 480–488 Pneumonia and influenza (폐렴 및 인플루엔자)에 속하는 코드들을 검색해보았습니다.

총 36개의 코드가 나왔고 다음은 그 중 일부입니다.

3,4,5 단계의 코드가 혼용되어 사용되고 있습니다. 또한, 데이터만으로는 검색된 36개의 코드가 같은 2단계 분류에 포함된다는 사실을 전혀 알 수 없습니다.

이 코드 체계에 대한 정보를 제공한다면 인공지능이 코드를 더 잘 이해할 수 있지 않을까요?

이에 대한 고민을 한 논문(참고문헌.1)을 하나 찾았습니다. 논문은 온톨로지 트리(Ontology Tree)와 BERT를 이용한 임베딩 방법을 제시하고 있는데, 이 중에서 온톨로지 트리를 이용한 방법만을 적용해보겠습니다.

온톨로지 임베딩

온톨로지 임베딩은 코드 체계의 상위 분류와 하위 분류의 코드들을 함께 활용해서 임베딩을 만드는 방법입니다.

온톨로지 임베딩은 두 번의 과정으로 나누어집니다.

먼저 하위 분류 코드의 임베딩을 집계해 상위 분류 코드의 임베딩을 정의합니다. 예를 들어, 487 Influenza는 하위 분류의 코드인 487.0 Influenza with pneumonia 487.1 Influenza with other respiratory manifestations를 이용해 임베딩합니다.

그 다음, 임베딩된 상위 분류 코드들을 다시 집계해 하위 분류 코드를 임베딩합니다. 예를 들어, 487.0 Influenza with pneumonia은 상위 분류 코드들인 460–519 Diseases of the respiratory system, 480–488 Pneumonia and influenza, 487 Influenza를 이용해 임베딩합니다.

집계시 사용하는 함수는 합계(sum), 평균(average) 등 어떤 것이든  사용할 수 있습니다.

자세한 방법과 수식은 논문(참고문헌.1)을 참조하세요.

Skip-Grams 모델

이전 글에서 소개했던 Skip-Grams 모델의 구조를 다시 보겠습니다.

Skip-Grams 모델은 한 코드를 기준으로 같은 방문에서 함께 사용된 다른 코드들을 예측하는 방법으로 학습하여 임베딩을 만드는 방법입니다. 여기에서 Embedding Machine 부분만 온톨로지 임베딩으로 바꾸어 기존 임베딩과의 성능을 비교해보았습니다.

기존에 사용했던 임베딩 방법은 One-Hot Encoding된 코드에 임베딩 가중치를 이용하여 임베딩 값을 뽑아내는 것입니다. 이 방법의 문제점 중 하나는 사용된 횟수가 적은 코드들로 인해 임베딩이 힘들어지고 성능이 하락한다는 것입니다.

반면, 온톨로지 임베딩은 코드 체계를 추가로 입력 받아 입력된 코드에 대해서 임베딩을 수행합니다. 덕분에 데이터에 사용된 횟수가 적더라도, 상위 분류나 하위 분류를 통해 어느 정도 해당 코드에 대한 정보를 줄 수 있습니다. 이를 통해 기존 모델보다 이해력이 높은 임베딩을 만들 수 있을 것입니다.

실험 결

집계 함수는 합계(sum)와 평균(average)를 사용하였고, 온톨로지 임베딩의 경우, 기존 임베딩에 학습 되는 속도가 느려 Epoch를 더 많이 주었습니다. 100회 이상 사용된 진단 코드만 사용한 것과 모든 진단 코드를 사용한 것도 비교해보았습니다.

Skip-Gram Loss

전체적으로 기존 Skip-Grams에 비해 Log Loss가 낮아졌고, 모든 코드를 사용하였을 때, 더 낮아진 것을 볼 수 있습니다.

이 임베딩을 재입원 예측 모델에 적용하여 성능을 비교해보겠습니다.

진단 코드만 임베딩하였고, 이에 따른 성능 변화를 더 확실히 보기 위해, 재입원 예측 모델에도 진단 정보만 넣어서 성능을 측정했습니다. 환자 레벨에서 분리하여 5-Fold Cross Validation을 수행했습니다.

재입원 예측 모델 성능

One-Hot Encoding보다는 성능이 좋지만, 기존 Skip-Grams와는 큰 차이가 없는 것을 볼 수 있습니다.

코드 체계에 대한 정보를 넣어서 임베딩을 헀지만, 성능엔 큰 변화가 없는 이유가 무엇일까요? 다음과 같은 이유를 생각해볼 수 있습니다.

  1. Skip-Grams 모델은 온톨로지 임베딩을 학습하기에 부족하다.

이번 글에서는 온톨로지 임베딩의 효과만을 확인하기 위하여 기존에 사용했던 Skip-Grams 모델을 그대로 사용했습니다. 하지만 참고한 논문에서는 훨씬 복잡하고 큰 모델인 BERT를 사용해 임베딩을 하고 있습니다.

  1. 진단 코드로는 예측할 수 있는 재입원 여부는 한계가 있다.

재입원 여부는 진단만으로 결정되지 않습니다. 데이터 내외적으로 수많은 요인이 있기 때문에 아무리 진단 코드에 대해서 이해가 완벽하다고 해도 재입원 예측 성능이 오르는 것은 불가능할 수 있습니다.

마치며

코드 체계를 이용한 온톨로지 임베딩 방법을 기존 임베딩 방법과 비교하여 알아보았습니다. 기존 임베딩 방법과 비슷한 성능은 낼 수 있었으나 능가하지는 못했고, 현재 실험 설계의 한계를 알게되었습니다.

EHR 데이터에서 사용되는 코드들에 체계가 있는 것은 널리 알려져 있지만, 이 데이터를 엔지니어가 활용하기에는 정제가 잘 되어 있지 않습니다. 이 글에서 사용한 코드 체계인 ICD-9도 일정한 포맷으로 정리된 페이지가 없어 정제하는데 많은 시간을 사용했습니다. ICD-9는 업데이트가 중단된 예전 코드 체계입니다. 다행히 최근 사용되고 있는 코드 체계는 정제가 잘 되고 있고, 데이터베이스화되어 제공되는 경우도 있습니다. 앞으로는 더 쉽게 코드 체계를 이용하여 인공지능 모델을 발전시킬 수 있을 것이라 생각합니다.


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


참고문헌

  1. Junyuan Shang, Tengfei Ma, Cao Xiao & Jimeng Sun (2019). Pre-training of Graph Augmented Transformers for Medication Recommendation (https://arxiv.org/pdf/1906.00346.pdf)

Jungwoo Chae

machine learning engineer

Jungwoo Chae