PoSE:通过位置跳跃式训练实现大语言模型上下文窗口的高效扩展
计算与语言
2024-02-22 v3 机器学习
摘要
大语言模型(LLMs)以预定义的上下文长度训练,这限制了其在需要长输入场景中的使用。以往使LLMs适应更长长度的尝试通常需要在该目标长度下进行微调(全长度微调),训练代价高昂。为解耦训练长度与目标长度以实现高效的上下文窗口扩展,我们提出位置跳跃式(Positional Skip-wisE, PoSE)训练,利用固定上下文窗口智能地模拟长输入。其做法是先将原始上下文窗口划分为若干块,再设计不同的跳跃偏置项以操控每块的位置索引。这些偏置项及每块的长度对每个训练样本均发生改变,使模型能适应目标长度内的所有位置。实验结果表明,与全长度微调相比,PoSE大幅降低了内存与时间开销,且对性能影响极小。利用该优势,我们已成功使用2k训练上下文窗口将LLaMA模型扩展至128k tokens。此外,我们实证确认PoSE兼容所有基于RoPE的LLM及位置插值策略。值得注意的是,我们的方法潜在支持无限长度,仅受推理时内存使用的限制。随着高效推理的持续进展,我们认为PoSE可进一步将上下文窗口扩展至128k以上。
引用
@article{arxiv.2309.10400,
title = {PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training},
author = {Dawei Zhu and Nan Yang and Liang Wang and Yifan Song and Wenhao Wu and Furu Wei and Sujian Li},
journal= {arXiv preprint arXiv:2309.10400},
year = {2024}
}
备注
ICLR 2024