NeurIPS 2020 | 基于图表示学习的缺失特征填充和标签预测模型

论文标题｜ Handling Missing Data with Graph Representation Learning
论文来源｜ NeurIPS 2020
论文链接｜ https://arxiv.org/abs/2010.16418
源码链接｜ https://github.com/maxiaoba/GRAPE

TL;DR

针对数据特征缺失的 feature imputation 和 label prediction 问题，论文提出了一种基于图的处理框架 GRAPE。GRAPE 基于图表示学习的方法来处理缺失特征填充和标签预测问题，主要思路是将观测样本点和特征类型作为二分图中的两种类型节点，其观测特征值作为边，按二分图建模即可：将特征填充问题转为边级别的预测任务，将标签预测问题转为节点级别的预测任务。实验部分在九个基准数据集中验证了 GRAPE 在缺失值填充任务中 MAE 可提高 20% 且在标签预测任务中提高 10%。

Problem Definition

给定 $n$ 个数据点的 $m$ 维特征 $\mathbf{D} \in \mathbb{R}^{n \times m}$ ，缺失值矩阵表示为 $\mathbf{M} \in {\{0, 1\}}^{n \times m}$ ， $M_{ij}=1$ 表示对应数据点 $\mathbf{D}_{ij}$ 的特征值不是缺失的。令 $\mathbf{Y}\in \mathbb{R}^n$ 表示下游任务的标签且 $\mathbf{V}\in\{0,1\}^n$ 表示标签是否已知。现需解决两个问题：

feature imputation：在 $\mathbf{M}_{ij}=0$ 时预测缺失的特征值 $\mathbf{D}_{ij}$ ;
label prediction：在测试集中 $\mathbf{V}_i=0$ 时预测标签 $\mathbf{Y}_i$ ;

Algorithm/Model

论文中提出 GRAPE 框架如下图所示，主要思想是将缺失值问题转为二分图中节点和边级别的预测问题，然后利用 GNN 模型来解决图级别的任务。👍🏻

所以说大佬们的厉害之处就在于可以将本不相关的问题转化为自身领域相关的问题，然后利用自身擅长的方法来解决而且 thoughts 还显得那么 make sense 🤟🏻 不得不膜拜啊！ 🤙🏻

问题转化

给定特征矩阵 $\mathbf{D}$ 及其掩码矩阵 $\mathbf{M}$ 表示为无向二分图 $\mathcal{G}=(\mathcal{V},\mathcal{E})$ ，其中节点 $\mathcal{V}=\mathcal{V}_D\cup \mathcal{V}_F$ $\mathcal{V}_D$ 表示数据类型的顶点集合， $\mathcal{V}_F$ 表示特征类型的顶点集合； $\mathcal{E}$ 表示两种类型节点间存在的边集合 $\mathcal{E}=\left\{\left(u_{i}, v_{j}, \mathbf{e}_{u_{i} v_{j}}\right) \mid u_{i} \in \mathcal{V}_{D}, v_{j} \in \mathcal{V}_{F}, \mathbf{M}_{i j}=1\right\}$ ，其中 $\mathbf{e}_{u_{i} v_{j}}$ 表示对应特征值 $\mathbf{e}_{u_{i} v_{j}}=\mathbf{D}_{ij}$ ，如果为类别特征那么转为 one-hot 编码。 $\mathbf{e}_{u_{i} v_{j}}$ 简化表示 $e_{ij}$ 表示特征矩阵 $\mathbf{D}$ 间关系， $e_{uv}$ 表示图 $\mathcal{G}$ 节点间关系。

此处 edge encoding 可能存在问题：如果数据中既存在连续特征又存在离散特征，那么会导致边向量的维度不同，如何处理？🤔

对应两个任务如下：

edge-level prediction
对于 $\forall \mathbf{M}_{i j}=0$ 预测边值 $\hat{\mathbf{D}}_{i j}=\hat{\mathbf{e}}_{i j}=f_{i j}(\mathcal{G})$ ，优化目标为最小化 $\hat{\mathbf{D}}_{i j}$ 和 $\mathbf{D}_{i j}$ 距离，损失函数 MSE/Cross entropy；
node-level prediction
对于 $\forall \mathbf{V}_{i}=0$ 预测节点标签 $\hat{\mathbf{Y}}_{i}=g_{i}(\mathcal{G})$ ，优化目标为最小化 $\hat{\mathbf{Y}}_{i}$ 和 $\mathbf{Y}_{i}$ 距离；

选择模型

论文中选用的是经典的 GraphSAGE 模型，详细介绍可以参考另一篇文章：NIPS 2017 | GraphSAGE：大规模图上的归纳表示学习模型

为了使 GraphSAGE 适应二分图中的缺失值预测问题，作者对其进行改进优化。

GRAPE GNN

给定二分图 $\mathcal{G}$ ，首先将 edge embeddings 加入 GraphSAGE 中。对于 $l$ 层 GNN，其输入为源节点node embedding $\mathbf{h}_v^{(l-1)}$ 加上边 edge embedding $\mathbf{e}_{uv}^{l-1}$ 连接

$\mathbf{n}_{v}^{(l)}=\operatorname{AGG}_{l}\left(\sigma\left(\mathbf{P}^{(l)} \cdot \operatorname{CONCAT}\left(\mathbf{h}_{v}^{(l-1)}, \mathbf{e}_{u v}^{(l-1)}\right) \mid \forall u \in \mathcal{N}\left(v, \mathcal{E}_{d r o p}\right)\right)\right)$

其中 $\operatorname{AGG}_l$ 表示聚合函数， $\sigma$ 为激活函数， $\mathbf{P}^{(l)}$ 为训练权重， $\mathcal{N}$ 为邻居节点。然后 node embedding $\mathbf{h}_v^{(l)}$ 更新为

$\mathbf{h}_{v}^{(l)}=\sigma\left(\mathbf{Q}^{(l)} \cdot \operatorname{CONCAT}\left(\mathbf{h}_{v}^{(l-1)}, \mathbf{n}_{v}^{(l)}\right)\right)$

其中 $\mathbf{Q}^{(l)}$ 为可训练权重。然后 edge embedding $\mathbf{e}_{uv}^{l}$ 更新为

$\mathbf{e}_{u v}^{(l)}=\sigma\left(\mathbf{W}^{(l)} \cdot \operatorname{CONCAT}\left(\mathbf{e}_{u v}^{(l-1)}, \mathbf{h}_{u}^{(l)}, \mathbf{h}_{v}^{(l)}\right)\right)$

其中 $\mathbf{W}^{(l)}$ 为可训练权重。

对应的 edge-level 预测为

$\hat{\mathbf{D}}_{u v}=\mathbf{O}_{\text {edge }}\left(\operatorname{CONCAT}\left(\mathbf{h}_{u}^{(L)}, \mathbf{h}_{v}^{(L)}\right)\right)$

对应的 node-level 预测为

$\hat{\mathbf{Y}}_{u}=\mathbf{O}_{\text {node }}\left(\hat{\mathbf{D}}_{u} .\right)$

其中 $\mathbf{O}_{\text {node }}$ 和 $\mathbf{O}_{\text {edge }}$ 都是前馈神经网络 MLP。

node augmented

从上面的构图来看 $\mathcal{V}_D,\mathcal{V}_F$ 都是没有特征的，会使模型难以区分不同类型的节点，所以论文中对两种类型的节点扩充 $m=|\mathcal{V}_F|$ 维特征：

$\operatorname{INIT}(v)=\left\{\begin{array}{ll} \mathbf{1} & v \in \mathcal{V}_{D} \\ \text { ONEHOT } & v \in \mathcal{V}_{F} \end{array}\right.$

edge dropout

为了防止过拟合，在训练时对二分图 $\mathcal{G}=(\mathcal{V} ; \mathcal{E})$ 中的边以概率 $r_{drop}$ 进行随机 dropout:

$\text { DROPEDGE }\left(\mathcal{E}, r_{d r o p}\right)=\left\{\left(u_{i}, v_{j},{ }_{i j}\right) \mid\left(u_{i}, v_{j}, \mathbf{e}_{i j}\right) \in \mathcal{E}, \mathbf{M}_{d r o p, i j}>r_{d r o p}\right\}$