ICLR2022 | GREASE LM: 图推理增强QA上的LM

本文是Christopher D. Manning和Jure Leskovec合作的一篇文章，NLPer和GNNer应该对这两个名字不陌生，一个讲了CS224N，一个讲了CS224W。

本文提出一种用图推理增强QA的LM的架构：GREASE LM。GREASE LM通过对LM和GNN进行多层次深度交互，有效捕捉导GNN的结构信息与LM的语义约束，从而提升了模型在QA任务上的性能，和处理复杂问题的能力。

研究背景

回答文本相关的问题需要从陈述的上下文(context)和它背后的知识(knowledge)来推理。

QA任务的基本范式是pre-trained LM。尽管在普通的benchmark上表现良好，但是当用于测试的example的分布与精调的数据集不相同时，模型的表现会变的挣扎。因为这种学习方式希望通过简单(偶尔是错误)的模式匹配直接走捷径获得答案，虽然能从LM中捕捉到context的情景约束和细微差别，但不能很好的表达概念之间的潜在关系。我们希望能用context提供的显性信息和内在的隐性知识来进行健壮的、结构化的推理。

Knowledge Graph(KG)由描述实体关系的三元组构成，被认为蕴含有大量的知识，在QA之类的推理任务上作用显著，因此用KG给LM扩充知识的结构化表示成为热门问题。然而，要将KG的推理优势扩展到一般的QA(问题和答案用自然语言表示，不容易转换成严格的逻辑查询)，需要根据QA example提供的信息和约束找到KG中正确的知识集成。

早期LM+KG的工作通常将两种模态在浅层并且以非交互的方式融合，比如各自独立编码最后预测时再融合，或者用一个模态去加强另一个模态的输入。一般分为三种：

1. 双塔结构，没有交互;
2. KG支撑LM，比如用KG的编码来增强QA example的文本表示;
3. 第三种是LM支撑KG，用文本表示(LM的最后一层)去增强从example提取的KG。

这三种结构信息流动最多有一个方向，这样两种模态信息交互的能力有限。如果想更好的模拟结构化的情景推理，需要实现双向的信息交互，让KG获得LM的细微语义约束，KG和LM进行深层次的融合。最近的一些工作探讨了这两种模态更深层次的集成。有的通过在结构化的KG数据上训练LM来将隐层知识编码进LM，然后用LM去生成针对QA的小KG。但是这样KG在转化成文本之后结构信息就丢掉了。QA-GNN用消息传递来联合更新LM和GNN，但是他们将LM池化作为整个架构的文本组件，限制了文本表示的更新能力。还有一些工作在预训练的时候将KG融入到LM中，但是模态的交互方向更多是将知识送到语言中。 Sun的工作与GREASE LM相似，但是它交互的bottleneck需要高精度的实体链接；并且LM和KG的参数共享限制了表达能力。

模型架构

符号定义:
MCQA(multiple choice question answering)的数据集包括上下文段落 $c$ ，问题 $q$ ，答案候选集 $A$ ，都以文本的形式表示。本工作中，还假设有额外的知识图谱 $G$ 提供背景知识。

给定QA example $(c,q,A)$ 和KG $G$ ，判断正确的答案 $a\in A$ 。

我们将自然语言中的token序列表示为 { w 1 , … , w T } \{w_1,\dots,w_T \} { w1​,…,wT​} ，其中 $T$ 表示token数量，token $w_t$ 在 $l$ -th layer的表示为 $h_t^{(l)}$ 。KG的点集表示为 { e 1 , … , e J } \{e_1,\dots,e_J\} { e1​,…,eJ​} ，其中 $J$ 表示节点的数量。节点 $e_j$ 在 $l$ -th layer的表示为 $e_j^{(l)}$ 。

输入表示：
先将 $c,q,a$ 和分离符并起来作为模型输入 $[c;q;a]$ ,转换成token序列 { w 1 , … , w T } \{w_1,\dots,w_T \} { w1​,…,wT​} ;然后用输入序列去检索(retrieval)出 $G$ 的子图 $G_{sub}$ , $G_{sub}$ 提供跟QA example相关的知识. $G_{sub}$ 的点集表示为 { e 1 , … , e J } \{e_1,\dots,e_J\} { e1​,…,eJ​} .

KG Retrieval：
首先根据文本中的实体从 $G$ 中链接出一个初始点集 $V_{linked}$ 。然后将 $V_{linked}$ 中任意点对之间的2-hop路径(长度为2,也就是中间只有一个点，也就是桥点)的桥点加进去形成 $V_{retrieved}$ 。然后再对 $V_{retrieved}$ 里的点计算相关分数(relevance score) : 将node name并在QA example后面，通过LM得到node name的output score，作为relavance score。我们取 $V_{retrieved}$ 中分数最高的200个点为 $V_{sub}$ ，剩下的都扔掉。最后，将所有链接两个 $V_{sub}$ 中的点的边加进去形成 $G_{sub}$ 。另外， $G_{sub}$ 里的每个点都做一个标记，标记这个点对应的实体是从哪里来的，来自上下文 $c$ / 询问 $q$ / 答案 $a$ / 这些点的相邻节点。本文之后的KG都是表示 $G_{sub}$ .

图1. GREASELM模型架构图

GREASE LM 整体架构有两个堆叠组件:

1. 单模态的LM层*N:获得输入token的初始表示
2. 交叉模态的GREASELM层*M:将LM层的文本表示与KG的图表示融合在一起

Language Pre-Encoding

将 { w 1 , … , w T } \{w_1,\dots,w_T \} { w1​,…,wT​} 的token、段、位置嵌入求和作为 $l=0$ 时的表示 { h i n t ( 0 ) , h 1 ( 0 ) , … , h T ( 0 ) } \{h_{int}^{(0)},h_1^{(0)},\dots,h_T^{(0)}\} { hint(0)​,h1(0)​,…,hT