REALM: Retrieval-Augmented Language Model Pre-Training

下载 PDF

REALM: Retrieval-Augmented Language Model Pre-Training - 中文验证版

英文原文卡片:realm_2020.md

状态:已翻译。

元数据

计算设置

论文在系统层面报告设置,但未命名 TPU 代际。它在 64 个 Google Cloud TPU 上预训练 REALM 200K 步,batch size 512,学习率 3e-5,使用 BERT 的默认优化器。用于 MIPS 索引的文档嵌入步骤在 16 个 TPU 上并行化。作者还说明整个模型可在配备 12GB GPU 的单机上运行。按项目规则,确切的 TPU 代际不应超出原文升级;如果放在本地加速器年代地图上,这属于 2018-2020 年 Google Cloud TPU Transformer 时期,但"64 Google Cloud TPUs"是论文声明的设备描述。

检索语料足够大,以至于本身就是一个系统组件。英文 Wikipedia 快照拆分为最多 288 个 BERT wordpiece 的块,产出略超过 1300 万检索候选项。在预训练期间,每个样本检索并对 8 个候选项文档进行边缘化,包含一个空文档;在微调推理期间,系统考虑前 5 个候选项。

瓶颈

瓶颈是梯度学习与大型外部记忆之间的不匹配。REALM 希望有一个神经检索器,其文档分数在 masked LM 预训练期间改进,但对数百万文档的精确边缘化每步不可能实现。最大内积搜索使检索成为次线性的,但它需要预计算的文档嵌入和一个索引。一旦检索器参数改变,那些缓存的文档嵌入就变得过时。

这是一个内存带宽和系统问题,与 NLP 问题同等重要。纯参数模型将事实存储在权重中,使更新昂贵且不透明。retrieval-augmented 模型在语料中暴露知识,但将成本转移到索引构建、嵌入刷新、MIPS 查找和对检索文本的交叉注意力中。

方法适配

REALM 通过将检索搜索与全神经训练解耦来适配硬件范式。检索器用输入嵌入和文档嵌入之间的内积对每个文档评分。因为文档端可以被缓存,MIPS 可以找到近似的 top-k 候选项,而无需在每个批次项上扫描 Wikipedia 的每个块。在检索候选项后,REALM 使用当前参数重新计算这些 top 文档的概率和梯度,减少过时缓存嵌入的损害。

关键工程举措是异步 MIPS 刷新。作者运行一个主训练任务用于梯度更新,以及一个辅助索引构建任务用于嵌入和索引文档。论文说明这大约每 500 个训练步产生一次索引刷新,这种节奏使检索保持对齐,而不拖慢 64 TPU 预训练任务。

模型本身保持足够小以供下游使用。比较列出 REALM 约 330M 参数,论文强调它可在单块 12GB GPU 上运行。这与将所有事实性知识扩展到巨型密集模型相反。大型对象是语料及其索引,而非神经参数张量。

证据

主要 Open-QA 表格报告 REALM 在 NaturalQuestions-Open、WebQuestions 和 CuratedTREC 上超越先前系统 4-16 个绝对百分点。以 Wikipedia 同时作为预训练和知识语料,REALM 在三个基准上报告 39.2、40.2 和 46.8;以 CC-News 为预训练语料、Wikipedia 为知识语料,报告 40.4、40.7 和 42.9。ORQA 比较行在同一 330M 参数量级上为 33.3、36.4 和 30.1。

消融表是最强的计算结构证据。在 NaturalQuestions-Open 上,REALM 达到精确匹配 38.2 和 recall@5 38.5。将检索器重置为基线状态使检索召回降至 13.9,匹配降至 35.3。ORQA 基线为 31.3 匹配和 13.9 recall@5。30 倍过时 MIPS 设置使匹配崩溃到 28.7,召回到 15.1 recall@5。这直接表明索引刷新计划不是实现细节;过时记忆破坏了该方法。

论文还与隐式知识缩放进行了比较。它指出将 T5 从 Base 迁移到 11B 使模型约大 50 倍,获得大约 5 点精度,而 REALM 在约小 30 倍的情况下超越 T5-11B。其主张是某些事实性知识检索比烘焙到密集参数中更廉价。

历史影响

REALM 是 retrieval-augmented 语言模型的一张早期密集检索预训练卡片。它表明外部记忆可以通过语言模型目标本身来学习,而不仅仅是用冻结的 BM25 或启发式检索器附加。历史计算效应是"模型大小"变为更大服务和训练系统的仅一个组件:语料大小、嵌入缓存、MIPS 索引、刷新频率和 top-k 边缘化都是一等计算变量。

它还预示了后来的 RAG 系统,其中事实新鲜度和出处(provenance)在密集权重之外处理。该卡片特别有用,因为它量化了系统失败模式:过时检索消除了收益。

局限

REALM 工程负担重。它需要一个大型文本语料、密集文档嵌入、近似搜索索引、辅助索引构建任务,以及刷新频率的调优。论文未命名确切的 TPU 代际,硬件核算不如后来的前沿模型报告完整。检索还增加了推理时延迟和失败模式:如果答案不在 top 候选项中,读取器无法恢复它;如果语料过时或有噪声,模型可能检索到看似合理但错误的支持。该方法在事实性知识可以用可检索文本块表示时最强。

链接

  • 计算范式:history/compute_regimes/inference_time_compute_post_training/README.md
  • 来源 PDF 和抽取文本见上方元数据。
  • 队列状态:read_complete
  • 方法索引:rag
  • Ledger 更新:compute bottlenecks