Towards Foundation Models for Knowledge Graph Reasoning

目录

• 概
• 符号说明
• ULTRA (a method for Unified, Learnable, and TRAnsferable KG representations)
  • Relation Graph Construction
  • Conditional Relation Representations
• 代码

Galkin M., Yuan X., Mostafa H., Tang J. and Zhu Z. Towards foundation models for knowledge graph reasoning. ICLR, 2024.

概

如何从不同领域的图中学习的有价值的可迁移的信息是通往 graph foundation model 的根本. 本文在多关系建模问题 (主要围绕知识图谱推理) 上给出了一个极为简单却有效的方案.

符号说明

• $\mathcal{V}$, nodes;
• $\mathcal{R}$, edge types;
• $\mathcal{E}$, edges, $e \in \mathcal{E}$ 表述为 triplet $(h, r, t)$.
• $\mathcal{G} = (\mathcal{V}, \mathcal{R}, \mathcal{E})$, knowledge graph;

ULTRA (a method for Unified, Learnable, and TRAnsferable KG representations)

• 只是图谱的推理问题是回答如下的问题:
  $$(h, r, ?): \text{何尾实体与头实体 h 存在关系 r}, \\ (?, r, t): \text{何头实体与尾实体 t 存在关系 r}. \\ $$

• Transductive: 如上图所示, 我们可以得出一个比较直接的结论. 一步步地,

  1. Michael Jackson 创作了 (authored) Thriller, 属于 disco 流派 (genre);
  2. 他的合作者 Quincy Jones 属于 disco 流派;
  3. 则我们可以推断出 Michael Jackson 的流派也是 disco.
  4. 倘若模型学到了这种结构关系, 那么我们很容易也能得出:
     $$\text{Beatles} \mathop{\longrightarrow} \limits^{authored} \text{Let It Be} \mathop{\longrightarrow} \limits^{genre} \text{rock} \wedge \text{Beatles} \mathop{\longrightarrow} \limits^{collab} \text{George Martin} \mathop{\longrightarrow} \limits^{genre} \text{rock}$$
     可以推断出: $\text{Beatles} \mathop{\longrightarrow} \limits^{genre} \text{rock}$.

• 但是学到这种关系的模型往往无法直接运用到一些别的领域上, 经过关系的模式是非常相似的. 比如:

• 几乎是一样的推理规则, 但是由于这些 entities 或者 edge types 未曾见过, 导致模型无法直接应用.

• 本文提出的 ULTRA 来一般性的学习这些关系间的关系, 然后介绍如何推广到 inductive 的环境中.

Relation Graph Construction

• ULTRA 将知识图谱 $\mathcal{G} = (\mathcal{V}, \mathcal{R}, \mathcal{E})$ 转换为新的图 $\mathcal{G}_r = (\mathcal{V}', \mathcal{R}', \mathcal{E}')$:
  • $\mathcal{V}' = \mathcal{R}$;
  • $\mathcal{R}'$ 包括:
    • tail-to-head (t2h): $(r_1, t2h, r_2)$ 说明 $\mathcal{G}$ 中存在 $a \mathop{\longrightarrow} \limits^{r_1} b \mathop{\longrightarrow} \limits^{r_2} c$;
    • head-to-head (h2h): $(r_1, h2h, r_2)$ 说明 $\mathcal{G}$ 中存在 $a \mathop{\longleftarrow} \limits^{r_1} b \mathop{\longrightarrow} \limits^{r_2} c$;
    • head-to-tail (h2t): $(r_1, h2t, r_2)$ 说明 $\mathcal{G}$ 中存在 $a \mathop{\longleftarrow} \limits^{r_1} b \mathop{\longleftarrow} \limits^{r_2} c$;
    • tail-to-tail (t2t): $(r_1, t2t, r_2)$ 说明 $\mathcal{G}$ 中存在 $a \mathop{\longrightarrow} \limits^{r_1} b \mathop{\longleftarrow} \limits^{r_2} c$;
  • $\mathcal{E}'$ 的确定根据如上的关系.

• 上图给了一个实际的例子.

• 大部分的基于 GNN 的推理都离不开邻接矩阵, 我们来说明如何得到 $\mathcal{G}'$ 的邻接矩阵. 首先, 我们定义 $\mathcal{G}$ 的邻接矩阵为 $\bm{A} \in \mathbb{R}^{|\mathcal{V}| \times |\mathcal{R}| \times |\mathcal{V}|}$. 每一条边 $e$, 实际上我们都可以拆成:

  $$(h, r), (r, t)$$

  于是, 我们可以构建得到:

  $$\bm{E}_h \in \mathbb{R}^{|\mathcal{V}| \times |\mathcal{R}|}, \bm{E}_t \in \mathbb{R}^{|\mathcal{V}| \times |\mathcal{R}|},$$

  满足:

  $$[\bm{E}_h]_{ar} \not = 0, \quad \text{if } \exist (a, r, ?) \in \mathcal{E}, \\ [\bm{E}_t]_{br} \not = 0, \quad \text{if } \exist (?, r, b) \in \mathcal{E}.$$

• 则

  $$\bm{A}_{h2h} = \text{spmm}(\bm{E}_h^T, \bm{E}_h) \in \mathbb{R}^{|\mathcal{R}| \times |\mathcal{R}|}, \\ \bm{A}_{t2t} = \text{spmm}(\bm{E}_t^T, \bm{E}_t) \in \mathbb{R}^{|\mathcal{R}| \times |\mathcal{R}|}, \\ \bm{A}_{h2t} = \text{spmm}(\bm{E}_h^T, \bm{E}_t) \in \mathbb{R}^{|\mathcal{R}| \times |\mathcal{R}|}, \\ \bm{A}_{t2h} = \text{spmm}(\bm{E}_t^T, \bm{E}_h) \in \mathbb{R}^{|\mathcal{R}| \times |\mathcal{R}|}, \\ \bm{A}_r = [\bm{A}_{h2h}, \bm{A}_{t2t}, \bm{A}_{h2t}, \bm{A}_{t2h}] \in \mathbb{R}^{|\mathcal{R}| \times |\mathcal{R}| \times 4}.$$

Conditional Relation Representations

• 现在, 我们来将怎么依靠 $\mathcal{G}_r$ 训练和推理. 主要是依靠一些 inductive learning 方法, 本文基于 NBFNet:

  $$\tag{1} \bm{h}_{v|u}^{0} = \mathbb{I}_{u=v} \odot \mathbf{1}^d, \\ \bm{h}_{v|u}^{t+1} = \text{Update} \bigg( \bm{h}_{v|u}^{t}, \text{Aggregate} \big( \text{Message}(\bm{h}_{w|u}^t, \bm{r}) | w \in \mathcal{N}_r(u), r \in \mathcal{R}' \big) \bigg).$$

  注意到, 上面的过程中不涉及特定的 node features, 而是仅为 node $u$ 初始化为 $\mathbf{1}^d$ 其余为 $\bm{0}$, 故而可以推广到任意图. 其它的细节请看 NBFNet.

• 假设我们通过 (1) 在多个 relation graphs $\mathcal{G}_r$ (涉及不同的领域) 训练好, 然后我们想要在一个新的图 $\mathcal{G}$ 上进行推理怎么办?

  1. 将这个新的图 $\mathcal{G}$ 转换为 relation graph, 然后用之前的模型在其上推理, 得到表示 $\bm{R}_q \in \mathbb{R}^{|\mathcal{R}| \times d}$. 注意, 这个表示是 query $q$ 依赖的.
  2. 作者额外再利用一个 NBFNet 进行 link prediction:
     $$\bm{h}_{v|u}^{0} = \mathbb{I}_{u=v} \odot \bm{R}_q[q], \\ \bm{h}_{v|u}^{t+1} = \text{Update} \bigg( \bm{h}_{v|u}^t, \text{Aggregate} \big( \text{Message}( \bm{h}_{w|u}^t, g^{t+1} (\bm{r})) | w \in \mathcal{N}_r(v), r \in \mathcal{R}. \big) \bigg)$$

• 需要注意, 上面两个阶段所涉及的参数都可以在一些数据集上预训练, 然后直接应用在下游的一些任务上, 从而实现 zero-shot 的推理. 也可以在下游任务进行一些微调. 具体的训练方式, 作者采取 BCE 进行训练.

代码

[official]