[论文解读] kNN-BOX: A Unified Framework for Nearest Neighbor Generation

Wenhao Zhu, Qianfeng Zhao+5 more|arXiv (Cornell University)|Feb 27, 2023

Natural Language Processing Techniques被引用 5

一句话总结

kNN-BOX 是一个统一的开源框架，将最近邻生成分解为三个模块化组件——数据存储库、检索器和组合器——实现了 kNN-MT 及其他序列到序列生成模型的快速原型设计、公平比较和交互式可视化。该框架在机器翻译、文本简化、改写生成和问题生成任务中均实现了显著的性能提升，BLEU 最高提升 6.1 点，SARI 最高提升 2.4 点。

ABSTRACT

Augmenting the base neural model with a token-level symbolic datastore is a novel generation paradigm and has achieved promising results in machine translation (MT). In this paper, we introduce a unified framework kNN-BOX, which enables quick development and interactive analysis for this novel paradigm. kNN-BOX decomposes the datastore-augmentation approach into three modules: datastore, retriever and combiner, thus putting diverse kNN generation methods into a unified way. Currently, kNN-BOX has provided implementation of seven popular kNN-MT variants, covering research from performance enhancement to efficiency optimization. It is easy for users to reproduce these existing works or customize their own models. Besides, users can interact with their kNN generation systems with kNN-BOX to better understand the underlying inference process in a visualized way. In the experiment section, we apply kNN-BOX for machine translation and three other seq2seq generation tasks, namely, text simplification, paraphrase generation and question generation. Experiment results show that augmenting the base neural model with kNN-BOX leads to a large performance improvement in all these tasks. The code and document of kNN-BOX is available at https://github.com/NJUNLP/knn-box.

研究动机与目标

• 解决当前缺乏统一、可扩展且交互式的 kNN 基础序列生成模型平台的问题。
• 实现对目前分散在多个代码库中的多样化 kNN-MT 方法的公平比较与快速复现。
• 通过可视化、交互式分析 kNN 生成中的检索过程，提升模型可解释性。
• 通过解耦数据存储库、检索器和组合器组件，实现模块化开发，促进模型融合。
• 将 kNN 增强生成技术的应用范围从机器翻译扩展到其他序列到序列任务。

提出的方法

• 将 kNN 生成分解为三个解耦模块：数据存储库（用于记忆键值对）、检索器（用于相似性搜索）和组合器（用于输出分布融合）。
• 利用教师强制法，以解码器隐藏状态作为键，目标词元作为值，从平行语料库构建符号化数据存储库。
• 采用 L2 平方距离并结合温度缩放，计算查询表示与存储键之间的相似性得分。
• 对基础模型的输出分布与检索到的 k 个最近邻分布进行加权插值。
• 支持基于网页的交互式可视化，用于检查检索结果、相似性分布及推理动态。
• 提供预实现的变体，涵盖性能增强（如 Robust kNN、PLAC）和效率优化（如 EfficientDatastore）方案。

实验结果

研究问题

• RQ1 统一框架能否简化多样化 kNN 生成方法的开发、比较与部署？
• RQ2 kNN-BOX 在机器翻译之外的多种序列到序列任务中，能在多大程度上提升生成质量？
• RQ3 模块化设计在支持灵活融合先进组件（如轻量级数据存储库与鲁棒组合器）方面有多高效？
• RQ4 kNN-BOX 内部的交互式可视化能否增强基于检索的生成机制的可解释性？
• RQ5 在多样化生成任务中，通过 kNN-BOX 增强基础模型，可实现多大的性能提升？

主要发现

• kNN-BOX 能够可靠复现七个已有的 kNN-MT 模型，其在 OPUS 和 TED 数据集上的结果与原始实现高度一致。
• 通过 kNN-BOX 增强基础模型，问题生成任务中 BLEU 分数最高提升 6.1 点，文本简化任务中 SARI 分数最高提升 2.4 点。
• PLACDatastore 与 RobustCombiner 的组合在仅使用原始数据存储库 55% 大小的情况下，于 OPUS-Law 数据集上达到 63.7 BLEU，展现出极高的效率与性能。
• 该框架支持新组件的无缝集成，实现了模块间状态最先进技术的模块化开发与融合。
• kNN-BOX 中的交互式可视化揭示了有意义的检索模式与分布偏移，有助于提升可解释性与模型调试。
• kNN-BOX 在所有评估任务中均显著提升性能——包括机器翻译、文本简化、改写生成与问题生成——充分证明了其广泛适用性。