GRAM: Graph-based Attention Model for Healthcare Representation Learning文献阅读记录

Abstract

目前的深度学习模型在医疗领域的应用面临着两个问题：

• 数据不充分：深度学习的模型需要比较大的数据量，在医疗领域的建模有时数据量不太够
• 解释性：表示学习所学得的方法应该能够与医学知识匹配

Introduction

本文首先介绍了一下EHR在医疗领域的应用，不再赘述。
然后提到关于训练深度学习模型时可能面临的数据量不足的问题。举例说明了cerebral degenerations (e.g. Leukodystrophy, Cerebral lipidoses) or developmental disorders (e.g. autistic disorder, Heller’s syndrome)的预测是很困难的，因为训练数据集中对于这些罕见病的学习机会是很少的。

解决上述问题的方法之一是引入医学本体，从而将医学概念的结构与关系都进行编码，从而在缩小search space的同时也不损失信息。
这里提到了几个应该可以用的本体：

• International Classification of Diseases (ICD)
• Clinical Classifications Software (CCS)
• Systematized Nomenclature of Medicine-Clinical Terms (SNOMED-CT)

这些本体结构中比较相近的nodes可能会和相似的患者有更紧密的联系，可以为模型进行knowledge transfer提供可能。
这篇文章中提出的方法叫做GRAM（GRaph-based Attention Model），这个模型可以通过注意力机制将医学本体的知识融合到深度学习模型中。

结果比较由以下几个方面展开：

• 预测性能与RNN在两个sequential diagnoses预测任务和heart failure预测任务中展开，GRAM在那些不常见的疾病预测任务上最多可以获得比RNN高10%的accuracy
• 对表示的结果进行了可视化，GRAM所表示的医学概念中更相似的位置也会更相近

Methodology

Basic Notation

• 整个EHR中的medical codes记录为$c_1,c_2,...,c_{|\mathcal{C}|}$的集合，总共有$|\mathcal{C}|$种medical codes
• 每个患者的clinical records可以看作是多次visits，即$V_1,V_2,...,V_T$，而每次visit中又包含着一系列medical codes。$V_t$可以表示成binary vector，${\mathbf{x}}_t\in \{0,1{\}}^{|\mathcal{C}|}$，其中第$i$个元素代表这个患者的这次记录中是否包含了code $c_i$
• 本体结构 $\mathcal{G}$所展示的是医学概念的层次结构，是以$parent-child$的形式展示关系的，而上面的$\mathcal{C}$形成的是所有的叶子结点。本体结构 $\mathcal{G}$可以以directed acyclic graph(DAG) 的形式进行表示，这个图中的所有结点为$\mathcal{D}=\mathcal{C}+{\mathcal{C}}^{\prime }$，其中的${\mathcal{C}}^{\prime }=\{c_{|\mathcal{C}|+1},c_{|\mathcal{C}|+2},...,c_{|\mathcal{C}|+|{\mathcal{C}}^{\prime }|}\}$，为所有的非叶子结点（也就是叶子结点的ancestors）。像ICD-9和CCS这样的本体结构中的分支都是没有交叉的，而SNOMED-CT中会存在少量少量交叉，但是本文主要只是考虑$parent-child$的关系。

Knowledge DAG and the Attention Mechanism

GRAM的目的是通过本体$\mathcal{G}$中的$parent-child$的关系在数据量不足的情况下学习到稳健的表示。
GRAM的设计中面临着使用本体的信息和数据中的信息的一个权衡，这会影响所学医学概念的特异性，这就是说如果用本体结构的信息更多，那么学到的概念应该是更加范化的，用数据的信息更多（数据量足够的话）那么学到的概念就会比较聚焦。

在上图中，下面几个实线的圈代表叶子结点，上面的虚线代表不同层次的$parent$，其中叶子结点$c_i$的表示可以由它几个$parent$的表示组合得到，然后可以使用几个$parent$的向量排出来的$G$矩阵进行嵌入学习，再结合注意力机制，可以将患者的$v_t$表示为一个向量${\mathbf{x}}_t$

在knowledge DAG中，每一个结点$c_i$都有一个基础的嵌入向量${\mathbf{e}}_i\in {\mathbb{R}}^m$。对应就有${\mathbf{e}}_1,{\mathbf{e}}_1,...,{\mathbf{e}}_{|\mathcal{C}|}$一直到${\mathbf{e}}_{|\mathcal{C}|+1},{\mathbf{e}}_{|\mathcal{C}|+2},...,{\mathbf{e}}_{|\mathcal{C}|+|{\mathcal{C}}^{\prime }|}$分别为叶子结点和非叶子结点的嵌入向量。而一个叶子结点的最终表示由它自己的嵌入向量和它的ancestors的嵌入向量的凸组合来表示：
$${\mathbf{g}}_i=\sum _{j\in \mathcal{A}(i)}{\alpha }_{ij}{\mathbf{e}}_j,\sum _{j\in \mathcal{A}(\mathbf{i})}{\alpha }_{ij}=1,{\alpha }_{ij}\ge 0forj\in \mathcal{A}(i)$$
其中${\mathbf{g}}_i$表示叶子结点的最终表示，$\mathcal{A}(i)$表示$c_i$的所有祖先和自己的集合。${\alpha }_{ij}$是注意力机制的权重值，用来权衡使用第$j$个祖先的嵌入来表示$i$的信息量权重，可以由下式计算得到:
$${\alpha }_{ij}=\frac{\exp (f({\mathbf{e}}_{\mathbf{i}},{\mathbf{e}}_{\mathbf{j}}))}{{\sum }_{k\in \mathcal{A}(i)}\exp f({\mathbf{e}}_{\mathbf{i}},{\mathbf{e}}_{\mathbf{j}})}$$
其中$f({\mathbf{e}}_{\mathbf{i}},{\mathbf{e}}_{\mathbf{j}})$度量的是${\mathbf{e}}_i$与${\mathbf{e}}_k$的匹配度，本文使用了一个单层的MLP来计算:
$$f({\mathbf{e}}_i,{\mathbf{e}}_j)={\mathbf{u}}_a^T\tanh ({\mathbf{W}}_a\left[\begin{array}{c}{\mathbf{e}}_i\\ {\mathbf{e}}_j\end{array}\right]+{\mathbf{b}}_a)$$
其中的${\mathbf{W}}_a\in {\mathbb{R}}^{l\times 2m}$用于对${\mathbf{e}}_{\mathbf{i}}$和${\mathbf{e}}_{\mathbf{j}}$拼接后的向量进行线性组合，$\mathbf{b}\in {\mathbb{R}}^l$是偏置向量，${\mathbf{u}}_a^T$是一个用于生成标量值的向量。式子中的$l$其实就是MLP的中间神经元的个数。

文章在这里特意说明了一下，虽然在这个研究中本体结构基本上是single path形式的（应该是指一个子结点只有一个父母），但是也可以拓展到多path的形式。

End-to-End Training with a Predictive Model

当得到了每一个code的表示${\mathbf{g}}_1,{\mathbf{g}}_2,...,{\mathbf{g}}_{|\mathcal{C}|}$后，可以对它们进行拼接然后得到嵌入矩阵$\mathbf{G}\in {\mathcal{R}}^{m\times |\mathcal{C}|}$，${\mathbf{g}}_i$对应着这个矩阵的第$i$列。
对于每个患者的一个visit $V_t$，可以通过一个多热向量${\mathbf{x}}_t$与$\mathbf{G}$作点乘再通过非线性转换获得一个向量表示${\mathbf{v}}_t$，然后可以预测最终的label ${\mathbf{y}}_t$。本文使用RNN来预测疾病序列。也就是说，给定了前几次的visits $V_1,V_2,...,V_t$以预测下一次的$V_{t+1}$。可以表示为如下形式:
$$\begin{gathered} {\mathbf{v}}_1,{\mathbf{v}}_2,...,{\mathbf{v}}_{\mathbf{t}}=\tanh \mathbf{G}[{\mathbf{x}}_1,{\mathbf{x}}_2,...,{\mathbf{x}}_t], \\ {\mathbf{h}}_1,{\mathbf{h}}_2,...,{\mathbf{h}}_{\mathbf{t}}=RNN(\mathbf{G}[{\mathbf{v}}_1,{\mathbf{v}}_2,...,{\mathbf{v}}_t,{\theta }_r]), \\ {\hat{y}}_t={\hat{x}}_{t+1}=Softmax({\mathbf{W}\mathbf{h}}_t+\mathbf{b}) \end{gathered}$$
其中${\mathbf{x}}_t$是第$t$次visit对应的多热向量，然后可以对应得到表示${\mathbf{v}}_t$，${\mathbf{h}}_t$为对应的隐藏层向量，${\theta }_r,\mathbf{W},\mathbf{b}$为参数。作者说这里最后预测使用的是Softmax而非dimension-wise sigmoid是因为前者有更好的预测效果。

模型使用的损失函数为交叉熵函数:
$$\mathcal{L}({\mathbf{x}}_1,{\mathbf{x}}_2,...,{\mathbf{x}}_T)=-\frac{1}{T-1}\sum _{t=1}^{T-1}({\mathbf{y}}_t^T\log ({\hat{\mathbf{y}}}_t))+(1-{\mathbf{y}}_t^T)\log (1-{\hat{\mathbf{y}}}_t))$$
这里的损失是针对一个患者的，并且是对所有的时刻的损失进行了加和。在实际的操作中，还需要对所有的患者进行取平均的操作。
下面是整个算法的伪代码：

Initializing Basic Embeddings

本文使用的是medical codes的共现关系来学习它们的基础表示，具体的方法用的是GloVe.
在这篇文章中，medical codes及它们的共现矩阵是根据患者的每次visit进行计数的。
这里举了一个例子：

其中$count(c_i,V_t^{\prime })$是$c_i$在$V_t^{\prime }$中的个数。在这个例子中，$c_a$与$c_c$的共线值为$3\times 2$, $c_i$与$c_a$的共线值为$1\times 3$。最终可以得到一个共现矩阵$\mathbf{M}\in {\mathbb{R}}^{|\mathcal{D}|\times |\mathcal{D}|}$

然后可以使用下面文献的方法来学习嵌入表示。

Jeffrey Pennington, Richard Socher, and Christopher D Manning. 2014. Glove: Global Vectors for Word Representation. In EMNLP.

Experiments

Experiment Setup

Prediction tasks and source of data
本文的预测任务是序列疾预测任务，即使用这次以前所有的患者记录以预测下一次患者可能诊断出的所有病。
使用了俩数据集，Sutter Palo Alto Medical Foundation (PAMF) 和MIMIC-III。前者每个患者有比较多的visits，后者只有比较少的visits。
疾病ICD通过CCS降低了类别数量以降低任务难度和加快训练速度，这里疾病编码出现的频率反映着数据的不充足程度。使用的指标是$Accuracy@k$，其含义是给定$V_t$，如果预测的top $k$疾病中有$V_t$中的疾病那么得到一个1，反之为0。此外，本文还整了一个预测心衰（Heart Failure）的预测任务。
本文使用的knowledge Graph 是CCS multi-level diagnoses hierarchy。

最后的结果从预测性能、表示的可视化以及注意力机制的得分等方面展开，不再赘述。