[논문 리뷰] Distilled Wasserstein Learning for Word Embedding and Topic Modeling

Hongteng Xu, Wenlin Wang외 2명|arXiv (Cornell University)|2018. 09. 12.

Machine Learning in Healthcare인용 수 33

한 줄 요약

이 논문은 워셔스타인 거리와 유클리드 단어 임베딩을 기반으로 한 비용 함수를 사용하여 단어 임베딩과 토픽 모델을 동시에 학습하는 통합 프레임워크인 Distilled Wasserstein Learning (DWL)을 제안한다. 모델 디스틸리케이션을 통합하여 최적화를 안정화시킴으로써 수렴성과 성능을 향상시켜 임상 입원 기록 분석에서 질병 네트워크 구축, 사망 예측, 절차 추천 등 최첨단 성능을 달성한다.

ABSTRACT

We propose a novel Wasserstein method with a distillation mechanism, yielding joint learning of word embeddings and topics. The proposed method is based on the fact that the Euclidean distance between word embeddings may be employed as the underlying distance in the Wasserstein topic model. The word distributions of topics, their optimal transports to the word distributions of documents, and the embeddings of words are learned in a unified framework. When learning the topic model, we leverage a distilled underlying distance matrix to update the topic distributions and smoothly calculate the corresponding optimal transports. Such a strategy provides the updating of word embeddings with robust guidance, improving the algorithmic convergence. As an application, we focus on patient admission records, in which the proposed method embeds the codes of diseases and procedures and learns the topics of admissions, obtaining superior performance on clinically-meaningful disease network construction, mortality prediction as a function of admission codes, and procedure recommendation.

연구 동기 및 목표

• 단어 임베딩과 토픽 모델링 간의 방법론적 격차를 해소하기 위해 둘을 통합된 프레임워크에서 동시에 학습하는 것.
• 임상 데이터에서 ICD 코드가 의미적으로 순서가 없기 때문에 순차적 단어 순서의 제한을 극복하는 것.
• 디스틸리케이션 메커니즘을 도입하여 공동 학습의 알고리즘 수렴성과 성능을 향상시키는 것.
• 질병/절차 분포 간 최적 운반 이론을 통한 최적의 표현을 통해 환자 입원 기록의 해석 가능하고 임상적으로 의미 있는 표현을 가능하게 하는 것.
• 실제 임상 과제인 사망 예측 및 절차 추천과 같은 분야에서 뛰어난 성능을 입증하는 것.

제안 방법

• 학습된 임베딩의 유클리드 거리로 단어 간 거리를 기반으로 하는 워셔스타인 토픽 모델을 수립한다.
• 문서(입원 기록)를 토픽을 정점으로 하는 워셔스타인 공간 내 가중치 중심으로 표현한다.
• 토픽 분포, 문서 가중치, 단어 임베딩을 동시에 업데이트하기 위해 번갈아 최적화 절차를 사용한다.
• 임베딩 업데이트를 안내하고 수렴성을 향상시키기 위해 안정화된 디스틸리케이션 버전의 거리 행렬을 사용하는 디스틸리케이션 메커니즘을 도입한다.
• 문서 및 토픽 분포 간 효율적인 최적 운반 계산을 위해 싱크본 알고리즘을 활용한다.
• 저자료 환경에서 성능을 향상시키기 위해 사전 학습된 Word2Vec 임베딩을 미세조정하여 단어 임베딩을 초기화한다.

실험 결과

연구 질문

• RQ1 단어 임베딩과 토픽 모델의 공동 학습이 별도 또는 순차적 접근 방식에 비해 임상 텍스트 분석 성능을 향상시키는가?
• RQ2 워셔스타인 거리의 비용으로 단어 임베딩 간 유클리드 거리를 사용할 경우 토픽 모델의 해석 가능성과 품질이 향상되는가?
• RQ3 모델 디스틸리케이션은 워셔스타인 기반 모델의 공동 학습 과정에서 수렴성과 강건성을 크게 향상시키는가?
• RQ4 제안된 방법은 사망 예측 및 절차 추천과 같은 후행 임상 과제에서 얼마나 잘 성능을 발휘하는가?
• RQ5 학습된 임베딩 및 토픽이 질병과 절차 간 임상적으로 의미 있는 관계를 얼마나 잘 반영하는가?

주요 결과

• DWL은 입원 ICD 코드를 사용한 환자 사망 예측에서 최첨단 기법을 능가하며, 더 높은 AUC 스코어를 기록한다.
• 절차 추천 과제에서 뛰어난 성능을 달성하여 임상 워크플로우와의 높은 일치도를 보인다.
• 학습된 단어 임베딩은 KNN 그래프 시각화를 통해 임상적으로 관련된 질병과 절차들이 강하게 군집되어 있음을 보여준다.
• DWL이 학습한 토픽은 임상적으로 해석 가능하며, 각 토픽의 상위 ICD 코드들이 신장 질환, 심혈관계 장애, 신생아 상태와 같은 일관된 임상 상태를 반영한다.
• 사전 학습된 Word2Vec 임베딩에서의 미세조정은 DWL의 성능을 크게 향상시켜 효과적인 전이 학습의 가능성을 시사한다.
• 디스틸리케이션 메커니즘이 학습을 안정화시키고, 특히 고차원적이고 희박한 임상 데이터 환경에서 수렴성을 향상시킨다.