[论文解读] 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学习（DWL），通过使用欧氏距离的词嵌入作为基础代价，联合学习词嵌入与主题模型。通过引入模型蒸馏以稳定优化过程，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性能，表明迁移学习有效。
• 蒸馏机制稳定了训练过程并提升了收敛性，尤其在高维稀疏临床数据设置下表现突出。