AI前沿｜将训练数据按相关性排序可以增强模型效果
当下语言模型被训练成可以在给定文档前缀后预测字词。现有预训练流水线将各种短文本随机拼接来构成输入上下文的方法有个缺点，即先前的文档对预测下一个文档无法提供信息。作者便提出上下文预训练，使得模型能在一个相关文档的序列上训练，因此显性地鼓励模型跨文档间阅读并推理。该方法可以通过简单改变文档顺序来实现，并直接应用到预训练流水线上。但文档排序问题有挑战性，需要在不重复数据的情况下将十亿级别的文档排序以最大化上下文相似性。作者用高效最近邻搜索来找到最相关文档并用一个图遍历算法构建一致的输入上下文。实验展示上下文预训练提供了简单且可规模化的方法来大大增强模型表现，在需要复杂上下文推理的任务上，包括上下文学习（+8%）、阅读理解（+15%）、先前上下文忠诚度（+16%）、长上下文推理（+5%）和检索增强（+9%）。上下文预训练的总览见figure1，对于图中右上角的例子，当预测后一句关于2022年FIFA设置了$40M奖金的句子，前一句是2022年前FIFA没有设置超过过$10M的奖金，而右下角的普通的训练数据可能前面一句话就来自另外一个主题。
为解决大数据按话题排序的问题，作者先大规模查找相关文档以创建文档图，然后通过遍历文档图构建预训练输入上下文。沿着路径，文档被连接成一个序列，随后被分割以形成固定大小的输入上下文（如8192个token）。在大规模查找相关文档这步，作者使用检索模型（contriever模型，其通过对文档中token的最后一个隐藏表示进行平均池化，将文档映射到嵌入中）链接预训练语料库D中的文档。具体来说，对于每个文档di ∈ D，使用稠密检索模型来检索前k个最相似的文档，表示为N(di)。检索模型使用近似最近邻搜索（近似最近邻搜索、乘积量化和倒排文件FAISS索引）来进行任意两个文档之间的高效成对相似性比较，该方法可扩展。在这个过程中作者发现很多近似重复的文档，便进一步利用检索分数从预训练语料库中消除接近重复的文档，该重复数据删除步骤对于实现语言模型的良好性能至关重要。第二步，为了依次沿路经排序，一种已经在RAG预训练中使用的就是kNN，但用kNN时，某些文档经常出现为其他文档的最近邻居，导致不同的输入上下文包含重复文档从而过拟合，因此作者的目标是以每个文档仅包含一次的方式构建一组上下文，这可以转化为最大化旅行商问题（将每个文档表示为图中的一个节点，并使用文档相似度作为边权重。），当然准确地解决大型旅行商问题是NP-hard的，贪心算法可作为有效的近似解决方案。具体见algorithm1，结合figure2，作者的算法首先选择一个具有最小度数的尚未访问的文档作为起始节点(Doc 0)。然后，该算法通过导航到具有最高权重的未访问的相邻文档（Doc 9），将文档节点添加到路径中，逐步扩展当前路径。该过程持续到路径到达所有相邻文档都被访问过的节点，这是因为这个的图不完整且仅包含k个最相邻文档之间的边。在本例中，作者将权重为0边加到图中和随机的未访问的拥有最小度的文档（Doc 1）相连，便可继续上述过程。从最小程度的文档开始的动机是它们最有可能首先访问所有邻居，因此在最终路径中连接到不同的文档。至此，所有文档都按照相似度被排序好了。
为评估上下文预训练的有效性，在CommonCrawl数据集的300B个token上预训练了0.3B到7B个参数的LM。在所有模型规模上，上下文预训练语言模型(ICLM)都展示了强大的语言建模和下游任务性能，优于在同一语料库上使用标准方法预训练的语言模型。效果见figure3等，太多了懒得截图了。