Focused Transformer: Contrastive Training for Context Scaling
Szymon Tworkowski1,313{}^{1,3}\thanks{Equal contribution \textsuperscript{\textdagger}Work done % while at Google Research.}start_FLOATSUPERSCRIPT 1 , 3 end_FLOATSUPERSCRIPT、Konrad Staniszewski1,3*13{}^{1,3*}start_FLOATSUPERSCRIPT 1 , 3 * end_FLOATSUPERSCRIPT、Mikołaj Pacek1,3*13{}^{1,3*}start_FLOATSUPERSCRIPT 1 , 3 * end_FLOATSUPERSCRIPT、Yuhuai Wu6666†、Henryk Michalewski、Henryk Michalewski3,434{}^{3,4}start_FLOATSUPERSCRIPT 3 , 4 end_FLOATSUPERSCRIPT、Piotr Miłoś1,2,5125{}^{1,2,5}start_FLOATSUPERSCRIPT 1 , 2 , 5 end_FLOATSUPERSCRIPT、Henryk Michalewski、11{}^{1}start_FLOATSUPERSCRIPT 1 end_FLOATSUPERSCRIPTIDEAS NCBR 22{}^{2}start_FLOATSUPERSCRIPT 2 end_FLOATSUPERSCRIPTInstitute of Mathematics, Polish Academy of Sciences 33{}^{3}start_FLOATSUPERSCRIPT 3 end_FLOATSUPERSCRIPTUniversity of Warsaw 44{}^{4}start_FLOATSUPERSCRIPT 4 end_FLOATSUPERSCRIPTGoogle DeepMind 55{}^{5}start_FLOATSUPERSCRIPT 5 end_FLOATSUPERSCRIPTdeepsense.ai 66{}^{6}start_FLOATSUPERSCRIPT 6 end_FLOATSUPERSCRIPTxAI Equal contribution †Work done while at Google Research.
Abstract
Large language models have an exceptional capability to incorporate new information in a contextual manner. However, the full potential of such an approach is often restrained due to a limitation in the effective context length. One solution to this issue is to endow an attention layer with access to an additional context, which comprises of (key, value) pairs. Yet, as the number of documents increases, the proportion of relevant keys to irrelevant ones decreases, leading the model to focus more on the irrelevant keys. We identify a significant challenge, dubbed the distraction issue, where keys linked to different semantic values might overlap, making them hard to distinguish. To tackle this problem, we introduce the Focused Transformer (FoT), a technique that employs a training process inspired by contrastive learning. This novel approach enhances the structure of the (key, value) space, enabling an extension of the context length. Our method allows for fine-tuning pre-existing, large-scale models to lengthen their effective context. This is demonstrated by our fine-tuning of $3B$ and $7B$ OpenLLaMA checkpoints. The resulting models, which we name LongLLaMA111We release the check
中文速览
大语言模型在处理超长上下文时会遭遇"分心问题"——当文档数量增多，无关文档的键（key）向量会与相关文档的键向量高度重叠，导致注意力机制难以准确聚焦到真正有用的信息上。为此，研究者提出了"聚焦变换器"（Focused Transformer，FoT），核心创新在于一种受对比学习启发的"跨批次训练"方法：在训练时，主动将相关文档（正样本）与无关文档（负样本）的键值对同时送入部分注意力层，迫使模型学会区分语义不同的键，从而优化键值空间的结构。基于此方法对3B和7B参数的OpenLLaMA模型进行轻量级微调，得到的LongLLaMA模型在长上下文任务上表现显著提升，能够处理高达256k词元的超长上下文（而训练时最长仅用8k词元）。这项工作的意义在于，它无需修改模型架构、无需从头训练，就能低成本地将现有大模型的有效上下文长度扩展到理论上无限，为大规模知识库检索与推理提供了实用路径。
原文 arXiv:2307.03170;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2307.03170v2