LongReD: Mitigating Short-Text Degradation of Long-Context Large Language Models via Restoration Distillation

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Dong, Zican, Li, Junyi, Jiang, Jinhao, Xu, Mingyu, Zhao, Wayne Xin, Wang, Bingning, Chen, Weipeng
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866913862822920192
author Dong, Zican
Li, Junyi
Jiang, Jinhao
Xu, Mingyu
Zhao, Wayne Xin
Wang, Bingning
Chen, Weipeng
author_facet Dong, Zican
Li, Junyi
Jiang, Jinhao
Xu, Mingyu
Zhao, Wayne Xin
Wang, Bingning
Chen, Weipeng
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2502_07365
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle LongReD: Mitigating Short-Text Degradation of Long-Context Large Language Models via Restoration Distillation
Dong, Zican
Li, Junyi
Jiang, Jinhao
Xu, Mingyu
Zhao, Wayne Xin
Wang, Bingning
Chen, Weipeng
Computation and Language
Machine Learning
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.
title LongReD: Mitigating Short-Text Degradation of Long-Context Large Language Models via Restoration Distillation
topic Computation and Language
Machine Learning
url https://arxiv.org/abs/2502.07365