의료 인공지능 성능 개선을 위한 시도 – Matrix Factorization Embedding

2020.05.27

라인웍스에서는 전자의무기록(Electronic Medical Record,이하 EMR) 데이터를 이용하여 다양한 머신러닝 프로젝트를 진행하고 있습니다. 

이번 글에서는 의료 인공지능 개발 프로젝트의 성능 개선을 위해 사용한 Matrix Factorization 알고리즘을 이용한 임베딩을 소개합니다. 또, 지난 글(https://linewalks.com/archives/6408)에서 소개한 Skip-grams 모델을 이용한 임베딩과 성능비교를 통해, Matrix Factorization 알고리즘의 장단점을 알아봅니다. 

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


Skip-Grams 모델

지난 글에서 임베딩에 사용한 모델은 Skip-Grams 모델입니다. Skip-Grams 모델은 자연어 처리에서 Word2Vec을 생성하기 위해 연구된 모델로 한 단어를 기준으로 그 주변의 단어들을 예측하는 모델입니다.

<Skip-gram 모델 구조 (Source: https://arxiv.org/pdf/1301.3781.pdf)>

예를 들어, 다음과 같은 문장이 있습니다.

Linewalks is hiring great machine learning engineers.

문장 중간의 한 단어인 machine이라는 단어가 입력되면 모델은 그 옆의 great, learning, hiring, engineers 라는 단어를 예측하는 것입니다.

이를 실제로 구현하면, 모델의 입력과 출력은 다음과 같은 구조가 됩니다.

한 단어가 입력되고 그 단어 주변에 있을 거라고 예상되는 단어들에 높은 확률을 부여하는 방식입니다.

Skip-Grams 모델에는 window size라는 중요한 hyperparameter가 있습니다. 입력 단어 양옆으로 어디까지를 “정답 단어”로 볼 것인가를 결정하는 요소입니다. 위의 예에서 앞 부분의 “Linewalks”와 “is”는 정답 단어에 포함이 되지 않았는데 이는 window size가 2로 설정되었기 때문입니다. 이 값을 늘리면 “Linewalks”와 “is”도 정답 단어에 포함되어 학습이 됩니다. 

Skip-Grams 모델은 기본적으로 “가까이 있는 단어끼리는 서로 연관성이 있다”를 전제로 하고 있고, window size에 따라서 “가까이 있는 단어”의 정의가 바뀌기 때문에 모델의 임베딩 성능이 많이 달라질 수 있습니다.

EMR 데이터에서의 Skip-Grams 모델

EMR 데이터로 예를 들어보겠습니다. 

철수가 감기에 걸려 병원에 갔습니다. 첫 날엔 열이 있어 해열제를 처방 받았습니다. 둘째 날엔 코가 점점 막혔고 해열제를 하루 더 먹었습니다. 셋째 날엔 기침이 나고 콧물이 나 기침약과 콧물약을 처방 받았습니다. 넷째 날엔 기침은 그대로인데 다시 열이 많이 나서 기침약과 해열제를 먹고 주사를 맞았습니다.

이 상황을 Skip-Grams 예시와 같이 문장으로 표현하면 다음과 같습니다.

Skip-Grams 모델에 굵게 표시된 “열”을 넣어봅시다. Window size가 2라면 모델은 “콧물약”, “기침”, “기침약”, “해열제”를 가장 가까이 있다고 예측할 것입니다. 주사는 열이 많이 나서 맞은 것인데 제외되게 되었습니다. Window size를 늘리면 해결이 될 것도 같지만 이 문제에 한정 지어서 해결이 되는 것이라 모든 상황에 대한 해결책은 되지 않습니다.

예제와 같이 EMR 데이터에서의 이벤트(진단, 의약품 등)를 나열한 리스트는 자연어 문장과 성격이 다릅니다. 기준 이벤트와 멀리 있는 이벤트가 가까이 있는 이벤트보다 연관성이 적다라는 가정이 성립하지 않습니다. 또한 자연어의 경우 단어의 순서가 바뀔 경우 의미가 바뀌기 때문에 매우 중요하지만, EMR 데이터에서는 그렇지 않습니다. 예제에서 넷째 날에 “기침”과 “열”의 순서나 “기침약”과 “해열제”, “주사”의 순서가 바뀌어도 의미는 전혀 바뀌지 않습니다. 위의 예제의 나타난 이벤트들 대부분이 순서가 바뀌어도 크게 이상하지 않고 동시에 발생했다고 보는게 더 타당한 경우도 있습니다.

EMR 데이터에서 Skip-Grams 모델이 가지는 한계라고 보았고, 다른 알고리즘을 시도해보기로 하였습니다.

Matrix Factorization

Matrix Factorization(행렬 인수분해)은 추천 시스템에서 주로 사용되는 Collaboration Filtering(협업 필터링) 알고리즘 중 하나입니다. 하나의 행렬을 특정한 구조를 가진 다른 행렬의 곱으로 나타내는 수학적 방법을 이용한 알고리즘입니다.

영화 추천 시스템을 예로 들어 설명하겠습니다. 다음과 같이 유저가 영화를 시청하였는지 여부를 행렬로 표현해보았습니다. 9명의 유저(U)와 8개의 영화(M)의 데이터로 만들어진 행렬입니다.

만들어진 행렬의 크기는 9 x 8 이며 이를 U(유저) x M(영화) 으로 표현하겠습니다. 이 행렬에 Matrix Factorization 알고리즘을 적용하면, U x K,  M x K 인 두 개의 행렬 P와 Q로 분해할 수 있습니다. 여기서 행렬 P와 Q를 찾는 일을 Matrix Factorization이 수행합니다. 

실제 알고리즘을 이 행렬에 적용한 결과를 보겠습니다. 임의의 수 K를 2로 설정해 다음과 같은 결과를 얻었습니다.

각각의 유저와 영화가 2차원 배열로 표현된 것을 확인할 수 있습니다. 유저를 표현한 행렬이 P, 영화를 표현한 행렬이 Q가 됩니다. 이 행렬들을 다시 곱하면 다음과 같은 행렬로 복원됩니다.

복원된 행렬에서 하늘색 바탕으로 표시된 부분이 기존에 유저가 시청했던 영화(기존 행렬의 값이 1)입니다. 대부분 1에 가까운 값으로 복원이 된 걸 확인할 수 있습니다.

 그리고 붉은 색 바탕으로 표시된 부분은 기존에 유저가 시청하지 않은 영화(기존 행렬의 값이 0) 중 복원된 값이 0.8 이상인 부분입니다. 복원된 값이 높을 수록 유저가 영화를 좋아할 확률이 높다라고 볼 수 있습니다.

이렇게 기존 행렬을 분해했다가 복원하는 과정을 거치면, 기존 행렬에서 비어있던 값을 채울 수 있습니다. 채워진 값을 이용해 유저에게 새로운 영화를 추천해주는 시스템을 만들 수 있습니다.

EMR 데이터에서의 Matrix Factorization

위의 영화 추천 시스템 예제에서 사용된 개념을 EMR 데이터에서는 다음과 같이 바꾸어 사용할 수 있습니다.

유저→ 환자
영화→ 이벤트(진단, 의약품, 의료행위 등)

이를 한 번 임의의 행렬로 표현해 보겠습니다.

위 표를 보면 영화 추천 시스템의 예제 행렬과 크게 다른 점이 없습니다. 

EMR 데이터에 똑같이 Matrix Factorization을 적용하면 영화 추천 시스템과 유사하게, 환자에게 어떤 약을 추천할 지, 어떤 질병에 걸릴 것 같은지를 예측해볼 수 있습니다.

하지만 이번 글에서는 범위를 좁혀 행렬 Q에 집중했습니다.

예제 행렬 Q가 영화를 K차원의 배열로 표현했듯, EMR 데이터 행렬 Q는 의료 이벤트를 K차원의 배열로 표현합니다. 이렇게 EMR 데이터의 이벤트가 표현된 K차원의 배열이 Skip-Grams 모델을 거쳐 나온 임베딩 결과를 대체할 수 있다고 생각했습니다.

이렇게 Matrix Factorizatioin 모델을 사용하면 얻을 수 있는 몇 가지 장점이 있습니다. 

첫째, Matrix Factorization모델은 Skip-Grams 모델의 한계인 동시 발생 이벤트 처리가 가능합니다. 이 모델은 각 이벤트들의 순서는 전혀 고려하지 않고, 해당기간 동안 발생한 이벤트가 무엇인가에 대한 정보만으로 임베딩을 합니다. 

두번째, 임베딩을 사용하지 않은 예측 모델이 있다면, 그 데이터를 즉시 Matrix Factorization 모델에 입력해 사용할 수 있습니다. 예를 들면 위에 언급한 EMR 데이터 행렬로 입원환자의 퇴원 후 30일 이내 재입원 예측을 진행한 적이 있었는데 이 행렬에 바로 Matrix Factorization 모델을 적용할 수 있습니다. 반면 Skip-grams 모델은 이벤트 발생 순서를 고려한 데이터 전처리가 필요해, EMR 데이터 행렬을 바로 사용할 수 없습니다.

물론 단점도 존재하는데, Matrix Factorization 모델은 이벤트 순서는 전혀 고려하지 않으므로, Skip-grams 모델과 서로 반대되는 한계점이 생깁니다. 질병 A에 걸린 환자가 합병증으로 질병 B가 발생한 경우, Matrix Factorization 모델은 질병 A와 B를 모두 같이 발생한 질병으로 판단해 임베딩할 것입니다.

이런 장단점, 한계를 고려하여 Matrix Factorization 모델을 MDwalks-EXI 재입원 예측 모델에 적용해 기존 방법과 성능의 차이를 실험해봤습니다.

Matrix Factorization 모델 적용 실험 결과

MDwalks-EXI에서 선보였던 MIMIC-III 30일 이내 재입원 예측 모델에 각각 임베딩을 적용하여 30일 이내 재입원 여부를 예측해보았습니다.

임베딩을 사용한 두 가지 모델 모두 임베딩을 사용하지 않은 모델에 비해서 AUROC가 향상된 것을 볼 수 있습니다. 반면, 두 모델 간의 성능은 큰 차이가 없다고 볼 수 있습니다. 위에서 설명했던 각 모델의 한계점이라고 볼 수 있을 것 같습니다.

마치며

Matrix Factorization 모델을 이용한 임베딩 방법을 Skip-Grams 모델과 비교하여 알아보았습니다. 실험 결과 Matrix Factorization 모델도 Skip-Grams 모델만큼의 성능을 낼 수 있다는 것을 알 수 있었습니다. 그러나, 순서가 전혀 고려되지 않는다는 한계점 때문인지 Skip-Grams 모델을 능가하는 성능을 보이진 못했습니다. 

그럼에도 불구하고 Matrix Factorization 모델의 사용성은 Skip-grams 모델보다 오히려 더 좋습니다. One-Hot Encoding을 이용한 모델에서 사용하는 데이터 전처리를 동일하게 사용할 수 있어 코드 재사용이 쉽고, 임베딩 학습 시 메모리 사용량도 데이터의 크기가 (환자 수) x (이벤트 종류 개수)로 고정되기 때문에 환자의 이벤트 Sequence에 따라 데이터의 크기가 달라지는 Skip-Grams 모델보다 예측이 쉽기 때문입니다.

이번 작업을 통해 두 모델의 장단점을 파악할 수 있었고, 예측할 문제에 따라 알고리즘 선택의 중요성을 다시 한 번 확인했습니다. 더 나아가, EMR 데이터에서 두 모델의 한계를 극복할 다른 모델을 연구하여 인공지능 모델 성능 향상을 이룰 수 있도록 노력해야겠습니다.


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

Jungwoo Chae

machine learning engineer

Jungwoo Chae