通过注意力头分析发现,合成数据可有效训练长文本检索能力。
Understanding Synthetic Context Extension via Retrieval Heads
- 用合成数据微调模型,通过注意力头追踪检索能力
- 合成数据训练的模型在真实任务上表现略差但可预测
- 关键注意力头重合度高,是性能提升的核心机制
长上下文大模型在检索增强生成等应用中需求日益增长。为降低长文本预训练成本,近期工作采用合成上下文扩展方法:在后训练阶段使用合成生成的长文本数据微调模型。然而,该方法如何及为何能提升下游长文本任务能力仍不清晰。本文针对需检索与推理的三类长文本任务,考察不同“针”概念(被检索内容)的现实性与“草堆”上下文(背景信息)多样性的影响,从基于LLM生成文档到模板化关系和符号数据集均有覆盖。结果发现,合成数据训练的模型性能低于真实数据,但其差距可通过一类特定的注意力头——检索头(Wu et al., 2024)——进行解释甚至预测。合成数据训练出的检索头与真实数据训练出的检索头高度重合,且头部召回率与下游性能强相关。通过注意力屏蔽和激活修补实验,进一步证实检索头对模型性能必要且具解释性,虽非充分条件。研究揭示了合成数据微调的内在机制,为构建更真实长文本能力训练数据提供了方向。
原文摘要 · Abstract (English)
Long-context LLMs are increasingly in demand for applications such as retrieval-augmented generation. To defray the cost of pretraining LLMs over long contexts, recent work takes an approach of synthetic context extension: fine-tuning LLMs with synthetically generated long-context data in a post-training stage. However, it remains unclear how and why this synthetic context extension imparts abilities for downstream long-context tasks. In this paper, we investigate fine-tuning on synthetic data for three long-context tasks that require retrieval and reasoning. We vary the realism of "needle" concepts to be retrieved and diversity of the surrounding "haystack" context, from using LLMs to construct synthetic documents to using templated relations and creating symbolic datasets. We find that models trained on synthetic data fall short of the real data, but surprisingly, the mismatch can be interpreted and even predicted in terms of a special set of attention heads that are responsible for retrieval over long context, retrieval heads (Wu et al., 2024). The retrieval heads learned on synthetic data have high overlap with retrieval heads learned on real data, and there is a strong correlation between the recall of heads learned and the downstream performance of a model. Furthermore, with attention knockout and activation patching, we mechanistically show that retrieval heads are necessary and explain model performance, although they are not totally sufficient. Our results shed light on how to interpret synthetic data fine-tuning performance and how to approach creating better data for learning real-world capabilities over long contexts.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。