提升长文本模型短文本表现,解决训练中性能下降问题。
LongReD: Mitigating Short-Text Degradation of Long-Context Large Language Models via Restoration Distillation
- 通过恢复性蒸馏缩小扩展前后模型的隐藏状态分布差异。
- 在长文本训练中引入短文本蒸馏,保持短文本任务性能不降。
- 适合需要兼顾长短文本能力的通用大模型应用者。
大语言模型通过扩展位置编码和轻量持续预训练实现了更长上下文窗口,但常导致短文本任务性能下降,其原因尚未充分探索。本文识别出两个主因:隐藏状态与注意力分数的分布漂移,以及持续预训练中的灾难性遗忘。为此,提出长上下文预训练恢复性蒸馏(LongReD),通过最小化扩展模型与原始模型间的分布差异来缓解短文本性能退化。除长文本训练外,LongReD还对原始模型在短文本上的选定层隐藏状态进行蒸馏,并引入短到长的蒸馏机制,利用跳过的位置索引对齐短文本输出分布。在常见文本基准测试中,LongReD有效保持了短文本性能,同时在处理长文本方面达到或超过基线水平。代码已开源于 https://github.com/RUCAIBox/LongReD。
原文摘要 · Abstract (English)
Large language models (LLMs) have gained extended context windows through scaling positional encodings and lightweight continual pre-training. However, this often leads to degraded performance on short-text tasks, while the reasons for this degradation remain insufficiently explored. In this work, we identify two primary factors contributing to this issue: distribution drift in hidden states and attention scores, and catastrophic forgetting during continual pre-training. To address these challenges, we propose Long Context Pre-training with Restoration Distillation (LongReD), a novel approach designed to mitigate short-text performance degradation through minimizing the distribution discrepancy between the extended and original models. Besides training on long texts, LongReD distills the hidden state of selected layers from the original model on short texts. Additionally, LongReD also introduces a short-to-long distillation, aligning the output distribution on short texts with that on long texts by leveraging skipped positional indices. Experiments on common text benchmarks demonstrate that LongReD effectively preserves the model's short-text performance while maintaining comparable or even better capacity to handle long texts than baselines. Our code is available at https://github.com/RUCAIBox/LongReD.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。