[論文レビュー] Distilled Wasserstein Learning for Word Embedding and Topic Modeling

Hongteng Xu, Wenlin Wang+2 more|arXiv (Cornell University)|Sep 12, 2018

Machine Learning in Healthcare被引用数 33

ひとこと要約

本稿では、ユークリッド語彙埋め込みを基本的なコストとして Wasserstein 距離を用いることで、語彙埋め込みとトピックモデルを統合的に学習する統一フレームワークである「ドミナント Wasserstein 学習（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コードが意味的に順序付けられていない臨床データにおける語の順序の制限を克服すること。
• 蒸留機構を導入することで、連合学習におけるアルゴリズムの収束性と性能を向上させること。
• 疾患／処置分布間の最適輸送に基づき、入院記録の解釈可能で臨床的に意味のある表現を可能にすること。
• 死亡予測や処置推薦などの実世界の臨床タスクにおいて、優れた性能を示すこと。

提案手法

• 語の間の距離を学習済み埋め込みのユークリッド距離とする Wasserstein トピックモデルを定式化する。
• 文書（入院記録）を、トピックを頂点とする Wasserstein 空間における重み付きバリセントルとして表現する。
• トピック分布、文書の重み、語の埋め込みを交互に更新するための最適化手順を採用する。
• 距離行列の安定化・蒸留されたバージョンを用いて埋め込みの更新をガイドする蒸留機構を導入する。
• 文書とトピック分布間の効率的な最適輸送計算のため、Sinkhorn アルゴリズムを活用する。
• 性能を高めるために、事前学習済み Word2Vec 埋め込みを微調整して語の埋め込みを初期化する。

実験結果

リサーチクエスチョン

• RQ1 語彙埋め込みとトピックモデルの連合学習は、別個または逐次的手法と比較して、臨床テキスト解析の性能を向上させることができるか？
• RQ2 語彙埋め込み間のユークリッド距離を Wasserstein 距離のコストとして用いることで、トピックモデルの解釈可能性と品質が向上するか？
• RQ3 モデル蒸留は、Wasserstein 基盤のモデルにおける連合学習プロセスの収束性と頑健性を顕著に改善できるか？
• RQ4 提案手法は、死亡予測や処置推薦といった下流の臨床タスクでどれほど優れた性能を発揮するか？
• RQ5 学習された埋め込みとトピックは、疾患と処置の間の臨床的に意味のある関係をどれほど的確に反映しているか？

主な発見

• DWL は、入院記録の ICD コードを用いた患者の死亡予測において、最先端の手法を上回り、AUC スコアが向上した。
• 処置推薦タスクにおいても、優れた性能を達成し、臨床ワークフローと良好に整合した結果を得た。
• 学習された語の埋め込みは、KNN グラフの可視化により、臨床的に関連する疾患や処置が強くクラスタリングされていることが示された。
• DWL が学習したトピックは臨床的に解釈可能であり、各トピックの上位 ICD コードが、腎疾患、循環器疾患、新生児疾患などの一貫性のある臨床状態に対応していた。
• 事前学習済み Word2Vec 埋め込みからの微調整が DWL の性能を顕著に向上させたことから、有効なトランスファーラーニングが可能であることが示された。
• 蒸留機構により訓練が安定化し、特に高次元かつスパースな臨床データ環境において収束性が向上した。