Memorizing Transformers
Yuhuai Wu Markus N. Rabe DeLesley Hutchins Christian Szegedy
Abstract
Language models typically need to be trained or finetuned in order to acquire new knowledge, which involves updating their weights. We instead envision language models that can simply read and memorize new data at inference time, thus acquiring new knowledge immediately. In this work, we extend language models with the ability to memorize the internal representations of past inputs. We demonstrate that an approximate $k$ NN lookup into a non-differentiable memory of recent (key, value) pairs improves language modeling across various benchmarks and tasks, including generic webtext (C4), math papers (arXiv), books (PG-19), code (Github), as well as formal theorems (Isabelle). We show that the performance steadily improves when we increase the size of memory up to 262K tokens. On benchmarks including code and mathematics, we find that the model is capable of making use of newly defined functions and theorems during test time.
中文速览
训练或微调语言模型来获取新知识需要反复更新权重、耗费大量算力，能不能让模型在推理时直接"读进去"新知识？这项工作为Transformer语言模型引入了一个基于近似k近邻（k-nearest-neighbor, kNN）检索的外部记忆模块：模型在处理长文档时，会把历史上下文的键值对（key-value pairs）存入一个不可微的外部记忆库，推理时用近似kNN快速查找最相关的历史表示并融合进当前注意力计算，记忆容量最高可扩展至26万个token。在代码、数学论文、书籍、网页文章和形式化定理等五类长文本任务上，模型困惑度随记忆容量的增大持续下降，性能提升甚至超过将模型规模扩大5倍所带来的收益，且模型确实学会了在测试时去检索并利用此前未见过的函数定义和数学引理。这一结果说明，借助高效的近似检索机制，无需改动模型权重就能让语言模型实现"即时学习"，为突破上下文长度瓶颈提供了一条简洁而有效的路径。
原文 arXiv:2203.08913;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2203.08913v1