CogLTX:应用BERT处理长文本

论文标题:CogLTX: Applying BERT to Long Texts
论文链接://arxiv.org/abs/2008.02496
论文来源:NeurIPS 2020

一、概述

BERT由于其随文本长度二次增长的内存占用和时间消耗,在处理长文本时显得力不从心。通常BERT最大支持输入序列长度为512,这对于标准的benchmark比如SQuAD和GLUE数据集是足够的,但对于更加通常的情况下,比如更加复杂的数据集或者真实世界的文本数据,512的序列长度是不够用的。

  1. 相关工作

目前BERT处理长文本的方法有截断法、Pooling法和压缩法。本文提出的CogLTX(Cognize Long TeXts)属于压缩法的一种。下面简单介绍下这三种处理方法:

  • 截断法

截断法的处理方式其实就是暴力截断,分为head截断、tail截断和head+tail截断。head截断就是从头开始保留限制的token数,tail截断就是从末尾往前截断,head+tail截断是开头和结尾各保留一部分。

  • Pooling法

将长文本分成多个segment,拆分可以使用暴力截断的方法,也可以使用断句或者划窗的方法。每一个segment都通过BERT进行encoding,然后对得到的[CLS]的embedding进行max-pooing或者mean-pooling,亦或将max-pooing、mean-pooling进行拼接。如果考虑性能,只能使用一个Pooling的话,就使用max-pooing,因为捕获的特征很稀疏,max-pooling会保留突出的特征,mean-pooling会将特征打平。

这种方法有明显缺点,首先需要将多个segment进行encoding,文本越长,速度越慢。另外这种拆分文本的方式也牺牲了长距离token之间进行attention的可能性,举例来说,如下图,这是HotpotQA中的一个例子,解答问题的两个关键句子之间长度相差超过512,因此他们不会出现在任何一个segment里,因此没法做attention:

example

  • 压缩法

压缩法是将长文本按照句子分为多个segment,然后使用规则或者单独训练一个模型来剔除一些无意义的segment。

使用滑窗一类的方法通常会将每个segment的结果做aggregate,通常这一类方法会使用max-pooing或者mean-pooling,或者接一个额外的MLP或者LSTM,但是这样牺牲了长程注意力的机会并且需要O(512^2\cdot L/512)=O(512L)的空间复杂度,这样的复杂度在batch size为1,token总数为2500,BERT版本为large的情况下对于RTX 2080ti的GPU仍然太大。并且这种方法只能优化分类问题,对于其他任务比如span extraction,有L个BERT输出,需要O(L^2)的空间来做aggregate。

  1. CogLTX

BERT之所以难以处理长文本的问题根源在于其O(L^2)的时间和空间复杂度,因此另外一个思路是简化transformer的结构,但是目前为止这一部分的成果很少能够应用于BERT。

CogLTX的灵感来自working memory,这是人类用来推理和决策的信息储存系统。实验表明working memory每次只能保留5~9个item或者word,但是它却有从长文本中进行推理的能力。 Working memory中的central executive,其功能就像有限容量的注意力系统,职责是协调综合信息。研究表明working memory中的内容会随着时间衰减,除非通过rehearsal来保持,也就是注意和刷新头脑中的信息,然后通过retrieval competition从长时记忆中更新忽略的信息用来推理和决策。

CogLTX的基本理念是通过拼接关键句子来进行推理。CogLTX的关键步骤是将文本拆分成多个blcok,然后识别关键的文本block,在CogLTX中叫做MemRecall,这是CogLTX的关键步骤。CogLTX中有另一个BERT模型,叫做judge,用来给block的相关性进行打分,并且它和原来的BERT(叫做reasoner)是jointly train的。CogLTX能够通过一些干预(intervention)来将特定于任务的标签(task-oriented label)转换成相关性标注(relevance annotation),以此来训练judge。

二、方法

  1. 方法论

CogLTX的基本假设是:对于大多数NLP任务来说,一些文本中关键的句子存储了完成任务所需要的充分且必要的信息,具体地,对于长文本x,其中存在由某些句子组成的短文本z,满足reasoner(x^+)\approx reasoner(z^+),其中x^+z^+是下图所示的BERT(reasoner)的输入:

CogLTX for different tasks

我们将长文本x分解成多个blockx=\left [ x_{0}\; \cdots x_{T-1}\right ],当BERT的长度限制为L=512时,每个block的最大长度限制为B=63。关键短文本z由部分x中的block组成,也就是z=\left [ x_{z_{0}}\; \cdots x_{z_{n-1}}\right ],满足len(z^+)\leq L并且z_{0}< \cdots < z_{n-1},接下来我们用z_{i}表示x_{z_{i}}z中的所有block都会自动排序,以保持x中的原始相对顺序。

CogLTX中两个重要的要素是MemRecall和两个jointly train的BERT,其中MemRecall是利用judge模型来检索关键块然后输入到reasoner中来完成任务的算法。

  1. MemRecall

做QA任务时MemRecall的整个流程如下:

MemRecall

  • 输入

前面图中(a)Span Extraction Tasks、(b)Sequence-level Tasks和(c)Token-wise Tasks三种不同类型的任务有不同的特定设置,在任务(a)、(c)中,问题Q和子句x[i]作为query来检索相关的block,但是在任务(b)中没有query,并且相关性只被训练数据隐式地定义。

z^+是被MemRecall维护的关键短文本,MemRecall除了接受x外还接受一个额外的初始z^+,在任务(a)、(c)中,query成为初始的z^+,judge在z^+的辅助下学习预测特定于任务的相关性。

  • 模型

MemRecall只使用一个模型,这是一个为每个token的相关性进行打分的BERT。假设z^{+}=\left [ [CLS]\; Q\; [SEP]\; z_{0}\; [SEP]\; \cdots\; z_{n-1}\right ],则有:

judge(z^{+})=sigmoid(MLP(BERT(z^{+})))\in (0,1)^{len(z^{+})}

一个blockz_{i}的得分是这个block内所有token得分的均值,记作judge(z^{+})[z_{i}]

  • 流程

MemRecall首先进行retrieval competition,每个blockx_i都会被打一个粗相关性得分judge(z^{+}\; [SEP]\; x_i)[x_i]。得分最高的几个“winner” block被插入到z中直到len(z^+)<L

接下来的rehearsal-decay过程会给每个z_i分配一个精相关性得分judge(z^{+})[z_{i}],得分最高的被保留在z^+中,得到一个new\; z^+。进行精相关性打分的原因是在没有block之间交互和比较的情况下粗相关性得分不够精确。

然后使用这个new\; z^+重复进行retrieval competition和rehearsal-decay,整个过程可以重复多次。通过这个迭代的过程可以实现multi-step reasoning。需要注意的是如果被z^+中的新block的更多信息证明相关性不够,在上一步中保留在new\; z^+中的block也可能decay,这是之前的multi-step reasoning方法所忽略的。

  1. 训练
  • judge的监督学习

通常span extraction tasks会把答案block标记为relevant,即使是multi-hop的数据集比如HotpotQA,通常也会标注支持的句子。在这些情况下,judge通常使用监督学习的方式来训练:

loss_{judge}(z)=CrossEntropy(judge(z^{+}),relv\_label(z^{+}))\\
relv\_label(z^{+})=[\underset{for\; query}{\underbrace{1,1,\cdots ,1}}\; \; \underset{z_0\; is\; irrelevant}{\underbrace{0,0,\cdots ,0}}\; \; \underset{z_1\; is\; relevant}{\underbrace{1,1,\cdots ,1}}\; \; \cdots ]\in [0,1]^{len(z^{+})}

这里训练的样本z是从x中采样出的多个连续blockz_{rand}(对应retrieval competition的数据分布)或者所有相关和随机选择的不相关blockz_{relv}(对应 rehearsal的数据分布)。

  • reasoner的监督学习

理想情况下,训练时reasoner的输入应该由MemRecall来生成,但是并不能保证所有的相关block都能被检索到。以QA任务为例,如果答案的blcok没有被检索到,reasoner就无法通过检索到的block进行训练,因此解决方案为做一个近似,将所有相关blcok和retrieval competition中的“winner” block输入到reasoner中进行训练。

  • judge的无监督学习

大多数的任务不会提供相关性的label。对于这种情况我们使用干预的手段来推断相关性标签:通过从z中剔除某个block来看它是否是不可或缺的。

假设z种包含所有的相关block,则根据我们的假设则有:

loss_{reasoner}(z_{-z_{i}})-loss_{reasoner}(z)>t,\forall z_{i}\in z,\; \; (necessity)\\
loss_{reasoner}([zx_{i}])-loss_{reasoner}(z)\approx 0,\forall x_{i}\notin z,\; \; (sufficiency)

z_{-z_{i}}是从z中移除z_{i}的结果并且t是一个阈值。每次训练reasoner后,我们会剔除每个z中的z_i,然后根据loss的增加调整它们的相关性标签。如果loss的增加是不显著的,则表明这个block是不相关的,它可能在下一个epoch中不会再赢得retrieval competition,因为它将在下一个epoch中被标记为irrelevant来训练judge。在实际中,阈值t会被划分为t_{up}t_{down},保留一个buffer zone来避免标签的频繁切换。

下图展示了20News数据集上无监督学习的一个例子:

example

  • 总结

训练的流程图总结如下:

训练

三、实验

  1. Reading comprehension

Reading comprehension

  1. Multi-hop question answering

Multi-hop question answering

  1. Text classification

Text classification

  1. Multi-label classification

Multi-label classification

  1. 内存和时间消耗

内存和时间消耗