임베딩(Embedding)을 활용한 의료 인공지능 성능 향상

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

이번 글에서는 의료 인공지능 개발 프로젝트의 성능 개선을 위해 사용한 임베딩을 소개합니다. 또, MDwalks EXI에서 선보였던 MIMIC-III 30일 이내 재입원 예측 모델(https://linewalks.com/archives/6142)에 임베딩을 적용해 이를 적용하지 않았던 이전 모델과 성능을 비교해보았습니다.

이 글이 임베딩의 개념과 EHR 데이터에서 임베딩의 필요성을 이해하는데 도움이 되기를 바랍니다.

임베딩(Embedding)이란?

임베딩이란 범주형(Categorical) 자료를 연속형 벡터로 치환하는 것을 말합니다. 예를 들어 A, B, C, D라는 카테고리가 있다고 할 때, 이를 다음과 같이 변환하는 것입니다.

이 변환이 왜 필요할까요? 쉬운 예제로 임베딩을 하지 않았을 때와 임베딩을 하였을 때 어떻게 다른지 알아보겠습니다.

임베딩이 필요한 이유

직장인의 가장 큰 고민인 “오늘 점심에 뭐 먹지?”를 해결해 주는 로봇, 점심봇이 있다고 가정합시다.

점심봇은 어제 먹은 점심 메뉴를 알려주면 오늘 먹을 메뉴를 추천해 주는 기능을 합니다. 그리고 우리는 추천된 메뉴를 반드시 먹어야만 한다고 가정합니다.

우리는 다음과 같이 물어볼 수 있습니다.

“어제 김치찌개 먹었어. 오늘은 뭐 먹을까?”

그럼 점심봇은 다음과 같이 대답합니다.

“오늘은 된장찌개 드세요.”

너무 완벽하죠? 하지만 실제 세상에선 메뉴가 이렇게 단순하지 않습니다. 김치찌개도 들어가는 재료에 따라 종류가 여러가지입니다. 실제 세상의 우리는 다음과 같이 물어봅니다.

“어제 참치김치찌개 먹었어. 오늘은 뭐 먹을까?”

그럼 점심봇은 다음과 같이 대답합니다.

“오늘은 스팸김치찌개 드세요.”

그럼 우리는 열받습니다. 비록 들어간 재료가 달라 다른 메뉴지만 김치찌개를 또 먹어야한다니… 이런 멍청한 점심봇!

지금부터 점심봇이 각 메뉴를 어떻게 인식하고 있는지를 확인해봅시다. (편의를 위해) 점심봇에게 현재 입력되어 있는 메뉴는 다음과 같이 총 6개라고 가정합시다. 그리고 이 메뉴들에는 정수형 ID가 하나씩 붙어 있습니다. ID는 임의로 붙여진 번호로 아무 의미가 없습니다.

위와 같은 데이터를 범주형(Categorical) 자료라고 합니다. ID에는 아무 의미가 없기 때문에 이 데이터를 넣기 위해선 다른 전처리가 필요합니다. 가장 일반적인 방법은 One-Hot Encoding 입니다. 전체 범주 갯수만큼의 배열을 만들고 해당 범주의 index만 값을 1로 주는 방법입니다. 위 6개 메뉴를 One-Hot Encoding하면 다음과 같이 6차원 벡터로 변환됩니다.

위의 변환된 벡터 값을 이용하여 참치김치찌개와 다른 메뉴 사이의 거리를 계산해 봅시다.

참치김치찌개와 다른 모든 메뉴 사이의 거리가 똑같군요! 점심봇이 보기엔 스팸김치찌개나 된장찌개나 참치김치찌개와 다른 정도가 모두 같기 때문에 어떤 메뉴가 서로 비슷한 메뉴인지 알 수가 없습니다. 그래서 아무리 어제 참치김치찌개 먹었다고 해도 스팸김치찌개와 생고기김치찌개를 제외하고 추천해줄 수가 없는거죠!

임베딩을 해봅시다

임베딩을 직접 해보겠습니다. 변환되는 벡터의 차원은 3차원으로 정하고 각각 찌개 강함 정도, 고기 선호도, 건강한 정도를 나타낸다고 합시다. 주관적인 의견을 많이 넣어서 다음과 같이 정의해보았습니다.

그럼 이제 참치김치찌개와 다른 메뉴 사이의 거리를 다시 한 번 계산해 볼까요?

이제 거리 차이가 나는 게 보입니다. 참치김치찌개와 비슷한 스팸김치찌개, 생고기김치찌개가 가장 가깝고 다른 메뉴들은 상대적으로 멀게 되었습니다. 이제 점심봇은 어떤 메뉴가 비슷한 메뉴인지 그나마 알게되었기 때문에 우리가 어제 먹은 메뉴와 너무 비슷한 메뉴는 추천하지 않도록 작동할 것입니다.

하지만 이 임베딩은 문제가 있습니다. 왜냐하면 필자의 머리에서 나온 임베딩이라 주관적인 요소가 많고 사실 다른 메뉴들과 김치찌개들의 거리를 멀게 만드는 게 생각보다 엄청 어려웠습니다. 고작 6개 메뉴인데도 이렇게 어려우면 현실의 수많은 메뉴들을 임베딩하는 건 거의 불가능에 가깝다고 말할 수 있습니다. 마찬가지로 수많은 코드 정보가 포함된 EHR 데이터도 직접 임베딩하는 건 불가능합니다. 그래서 임베딩을 컴퓨터에게 학습시키도록 하는 방법이 필요합니다.

EHR 데이터에서의 임베딩

EHR 데이터에는 수많은 코드 정보가 있습니다. 수많은 질병과 그 병을 치료하기 위한 의약품,  의료행위(수술, 시술 등), 또 그 의료행위에서 사용되는 여러 치료재료 등 모두 합치면 수십만개 이상이 된다고 볼 수 있습니다. 다음 코드 예를 봅시다.

위 세 코드는 세계 보건 기구에서 발표하는 ICD-9(국제질병사인분류 9차 개정판)에 포함된 코드입니다. 코드 설명을 보면 487.0과 V04.81이 인플루엔자에 연관된 내용으로 비슷한 걸로 보이는데 코드 자체만으로는 둘이 비슷한지 전혀 알 수가 없습니다.

이렇게 코드 자체로는 알 수 없는 정보들을 서로 의미있는 것끼리 서로 가깝게 두는 임베딩 과정을 거치면 예측 모델의 성능을 향상 시킬 수 있을 것으로 예상했습니다. 

Skip-Grams 모델

이 글에서 EHR 데이터 임베딩에 사용한 모델은 Skip-Grams 모델입니다. Skip-Grams 는 자연어 처리에서 Word2Vec을 생성하기 위해 연구된 모델로 한 단어를 기준으로 그 주변의 단어들을 예측하는 모델입니다. EHR 데이터에서는 방문에서 사용된 코드를 기준으로 같은 방문에서 사용된 다른 코드들을 예측하는 방법으로 학습하였습니다. 모델의 구조는 다음 그림과 같습니다.

이렇게 학습된 모델에서 Embedding Machine만을 가지고 Target Code를 입력해 Embed Representation을 출력 받아 이용하는 것입니다.

임베딩 결과

다음은 MIMIC-III 데이터의 ICD-9 진단 코드 중 100회 이상 사용된 코드 874개를 Skip-Grams 모델을 이용하여 50 차원의 벡터로 임베딩한 결과를 Tensorboard의 T-SNE를 이용하여 2D로 표현한 그림입니다.

50차원의 벡터가 2차원으로 변환되며 여러 곳에 군집되어 있는 것을 볼 수 있습니다. 오른쪽 아래에 모인 군집을 확대하여 자세히 살펴보겠습니다.

표시되어 있는 진단명들을 보면 accidents, fall, closed fracture 등의 여러번 나타나는 공통된 단어들이 보입니다. 어느 정도 비슷한 진단명들이 모여 있음을 확인할 수 있습니다.

자세한 임베딩 결과는 다음 링크의 Colab Notebook을 이용하세요. → 링크 바로가기

모든 셀을 실행한 뒤, 출력되는 마지막 링크를 클릭하시면 임베딩 결과가 출력되는 Tensorboard 화면을 보실 수 있습니다.

임베딩 결과 활용

MDwalks-EXI에서 선보였던 MIMIC-III 30일 이내 재입원 예측 모델에 임베딩을 적용하여 기존 모델과 비교해 보았습니다. 기존 모델은 코드들을 One-Hot Encoding 하였고, 임베딩 적용 모델은 의약품과 진단 코드만 임베딩을 적용하였습니다. 환자를 기준으로 5-Fold Cross Validation을 적용하여 AUROC의 평균과 분산을 계산하였습니다.

그 결과, 기존 모델의 AUROC는 

0.7026 ± 0.0224

임베딩 적용 모델의 AUROC는

0.7320 ± 0.0193

로 성능이 약 4% 증가하였고, 임베딩이 One-Hot Encoding보다 더 많은 정보를 주고 있음을 알 수 있습니다.

마치며

임베딩에 대한 기본 설명과 EHR 데이터에서 임베딩 적용에 대해서 알아보았습니다. 요약하자면 의미 없는 코드의 나열을 컴퓨터가 더 잘 이해할 수 있도록 하는 작업을 임베딩이라 할 수 있습니다.

의료분야에는 질병과 증상에 따라 처방하는 일종의 가이드라인들이 잘 마련되있습니다. 이런 가이드라인에 따라 EHR 데이터의 각 코드 간의 관계가 의미가 있고, 중요한 정보들을 담고 있습니다. 따라서 EHR 데이터를 사용하는 기계학습 모델의 성능 향상을 위해 임베딩이 매우 중요하게 작용할 것이라 예상합니다.

그러나 지금까지는 EHR 데이터 임베딩은 걸음마 단계라고 볼 수 있습니다. 위에서 소개한 Skip-Grams 모델은 2013년에 발표된 모델로 그 사이에 많은 모델들이 연구 및 발표되었고, 대부분의 임베딩 모델은 자연어 처리 분야의 모델입니다. 의료 분야에서도 코드나 방문을 표현하려는 시도들이 있고 논문들이 나오고 있으나 타 분야에 비하면 많이 부족한 상태입니다.

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


No comments yet.

Leave a comment

Your email address will not be published.