分享自:

通过记忆进行泛化:最近邻语言模型

期刊:ICLR

本文属于类型a,即一项原创性的单一研究报告。以下是关于该研究的学术报告。


研究报告:通过记忆实现泛化——最近邻语言模型

1. 研究作者、机构与发表信息

该研究题为《Generalization through Memorization: Nearest Neighbor Language Models》,由斯坦福大学(Stanford University)的 Urvashi Khandelwal 和 Dan Jurafsky,以及 Facebook AI Research 的 Omer Levy、Luke Zettlemoyer 和 Mike Lewis 共同完成。这项工作在 Urvashi Khandelwal 于 Facebook AI Research 实习期间完成,并发表于 2020 年的国际学习表征会议(ICLR 2020)。

2. 学术背景与研究动机

本研究属于自然语言处理(Natural Language Processing, NLP)领域,核心关注点是语言模型(Language Model, LM)。传统的神经语言模型旨在解决两个子问题:其一,将句子前缀映射为固定长度的向量表征(Representation);其二,利用这些表征来预测文本中的下一个词。该研究团队提出了一个全新的假设:表征学习问题可能比预测问题更为简单。例如,任何说英语的人都知道“Dickens is the author of”和“Dickens wrote”这两个上下文,即使他们不确定后续具体是什么词,也能推断其后续词的概率分布应高度相似。

基于此假设,研究旨在探索是否可以利用现有语言模型强大的表征学习能力,通过一种简单的最近邻(k-Nearest Neighbors, kNN)机制来显式地记忆训练数据中的稀有模式,从而替代模型参数中隐式、低效的记忆,最终提升语言模型的整体性能,尤其是在处理长尾分布中的罕见模式(如事实知识、专有名词)时的预测能力。

3. 研究方法与详细工作流程

该研究提出了一种名为 kNN-LM 的方法,该方法无需任何额外训练,即可通过线性插值的方式,将一个预训练神经语言模型与一个k最近邻模型相结合,其工作流程包含三个主要阶段:数据存储库的构建、推理时的检索与分布计算,以及最终的分布插值。

阶段一:数据存储库(Datastore)的构建 首先,研究者使用一个预训练好的、参数固定的语言模型,在其训练数据上进行一次单向的前向传播。预训练模型本身采用基于 Transformer 的纯解码器架构,完全遵循了 Baevski 和 Auli 在 2019 年提出的模型配置与优化方法。该模型拥有16层,每层包含16个自注意力头(Self-Attention Heads),隐藏状态维度为1024,前馈网络层维度为4096,可训练参数总计约2.47亿。对于语言模型将要处理的每一个训练样本,即上下文(Context)ci 及其对应的目标词(Target)wi,研究团队提取模型中间层的一个固定长度向量作为上下文的表征。具体来说,他们使用 Transformer 模型最后一层中,输入到前馈网络(Feedforward Network)之前、且经过层归一化(Layer Normalization)之后的1024维向量作为键(Key)ki = f(ci),而值(Value)vi 就是目标词 wi 本身。由此,整个训练集(Datastore) (K, V) 便由所有这样的键值对(Key-Value Pair)组成,其规模可达数亿乃至数十亿对。

阶段二:推理时的最近邻检索与分布计算 在测试阶段,给定一个新的上下文 x,预训练模型首先会生成其对应的上下文表征 f(x) 以及初始的词表概率分布 pLM(y|x)。接着,系统使用 f(x) 作为查询向量(Query),在数据存储库中进行k最近邻检索。为了在海量高维向量中进行快速检索,研究者采用了开源的 FAISS 库。为了提高效率,他们首先从存储库中随机采样100万个键来学习4096个聚类中心点(Cluster Centroids),并在检索时只搜索最近的32个聚类中心内的向量,同时将键向量量化为64字节以节省内存。检索的距离度量函数 d(·,·) 采用的是平方欧氏距离(Squared L2 Distance),这使得相似度函数实际上等价于一个径向基函数核(Radial Basis Function kernel, RBF Kernel)。检索到 k 个最近邻后,模型会根据这些邻居与查询向量的距离,通过一个温度参数为1的 Softmax 函数计算出一个关于这些邻居目标词的分布 pkNN(y|x)。该分布的计算公式为:p(kNN|y) ∝ Σ 1(y = vi) exp(-d(ki, f(x))),其含义是,对检索到的所有邻居中相同词汇的概率质量(Probability Mass)进行聚合。

阶段三:分布插值 研究的最后一步是将上述最近邻分布与原始语言模型的分布进行线性插值。最终预测的词概率分布由公式 p(y|x) = λ pkNN(y|x) + (1 - λ) pLM(y|x) 决定。其中,插值系数 λ 是一个在验证集上调优的超参数,用于平衡来自语言模型本身的知识和来自数据存储库的显式记忆。

4. 主要实验结果与分析

研究团队在多个规模、多领域的英文语料库上进行了全面的实证评估,取得了一系列显著结果。

核心性能提升(WikiText-103数据集) 在标准基准 WikiText-103 数据集上,kNN-LM 将基线模型(Base LM)18.65 的测试集困惑度(Perplexity, PPL)显著降低至 16.12,创造了当时的最先进水平。当进一步结合一种正交技术——“连续缓存”(Continuous Cache)时,困惑度进一步下降至 15.79,相较于基线模型实现了 2.86 点的巨大提升。在 Books 语料库上,该方法同样将测试集困惑度从 11.89 降低到 10.89,证明了其在百科类文本之外的泛化能力。

高效扩展性:以检索替代训练 在探索 kNN-LM 的可扩展性时,一个极具启发性的实验是:研究者首先仅在包含 1 亿词元(Token)的子集 Wiki-100M 上训练语言模型,然后利用该模型构建一个包含全部 30 亿词元的 Wiki-3B 数据存储库。结果显示,这个组合模型的测试集困惑度达到了 13.73,不仅大幅优于仅在 1 亿词元上训练的基线模型(19.59),甚至显著优于在全部 30 亿词元上直接训练的语言模型(15.17)。这表明,通过 kNN-LM 从海量数据中直接检索邻居,其效果可以超越在同等数据上进行直接训练。研究者还通过变化数据存储库的大小发现,模型性能随着存储库的增大而单调提升,且插值系数 λ 也随之提高,意味着模型越来越依赖检索到的记忆。

出色的领域自适应能力 在领域自适应(Domain Adaptation)任务上,一个在 Wiki-3B 上训练的语言模型在 Books 测试集上的困惑度高达 34.84。但只需在不改变模型参数的情况下,为其添加一个基于 Books 训练集构建的数据存储库,困惑度便急剧下降至 20.47。这说明 kNN-LM 提供了一个极其轻量且高效的领域自适应方案,即一个通用模型配合特定领域的检索数据库,就能胜任多个领域的任务。

超参数分析与机制解释 研究对几个关键超参数进行了分析。在对 Transformer 不同中间层输出作为键函数 f(·) 的比较中发现,最后一层前馈网络的输入效果最佳,这表明 Transformer 可能将表征学习的能力更多地编码在了自注意力层。k值(邻居数量)的增大能单调改善性能,即使 k=8 也能取得最先进水平。通过与 n-gram 模型的对比实验,发现 n-gram 的插值仅带来 0.2 的困惑度提升,远不及 kNN-LM 的 1.9,这凸显了使用神经网络学习的上下文表征来衡量相似性的关键作用。最后,通过关闭 Dropout 强制让模型过拟合来“记忆”所有训练数据,发现这种隐式记忆虽能使训练损失降为零,但验证集困惑度反而急剧上升至 28.59,且将其与原始模型插值也几乎没有提升。这反证了 kNN-LM 的优越性:Transformer 虽然有足够的容量(Capacity)记忆全部数据,但强行使其隐式记忆会破坏其表征的泛化能力,而 kNN-LM 的显式记忆则能在不干扰模型泛化能力的前提下,有效地利用训练数据。

5. 结论与价值

该研究成功引入了 kNN-LM 框架,证明了通过最近邻检索机制实现的显式记忆,能够显著超越仅依赖模型参数进行隐式记忆的标准语言模型。其科学价值在于,它有力地支持了“学习文本序列间的相似性比预测下一个词更简单”这一核心假设,为理解神经语言模型内部工作机制——即表征学习与预测任务的分离——提供了新的视角和实证依据。在应用价值方面,kNN-LM 开辟了一条提升语言模型性能的新路径:无需耗费巨量算力训练越来越大的模型,而是可以用较小的模型学习优质的上下文表征,再通过构建超大规模的检索数据库来实现高效扩展。此外,其灵活的领域自适应能力使得一个预训练模型能够便捷地服务于多个不同领域,极具实用潜力。

6. 研究亮点

本研究的亮点在于其方法的新颖性、结果的颠覆性以及零额外训练的便捷性。首先,kNN-LM 方法本身是全新的,它将预测问题分解为“表征”和“检索”两步,开创性地使用了预训练模型的中间层表征作为相似性检索的依据。其次,其“以检索替代训练”的实验结果极具启发性,颠覆了“更多数据必须通过更大规模训练来利用”的传统观念,指出了在固定模型参数下,通过扩展外部记忆库来高效利用海量数据的全新范式。最后,该方法作为对现有模型的增强,完全无需任何额外的训练步骤,能够即插即用地应用于任何能够生成固定大小上下文表征的神经语言模型之上,这种非侵入式的特性使其具有极高的通用性和推广价值。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com