Sparse Training: Do All Tokens Matter for Long Sequence Generalization?

Kaitao Song, Zihao Li, Cen LU, Xu Tan, Dongsheng Li

Keywords: Large Language Models, Long Sequence, Length Extrapolation, Efficiency

Abstract: Large language models (LLMs) have demonstrated remarkable progress in generating high-quality natural language through performing extensive pre-training over Transformer architectures. However, the quadratic complexity of transformers in sequence computation greatly limits its capability in efficiently modeling long sequences. In this paper, we introduce \method, a simple training technique to optimize the complexity of Transformer models in long-sequence training. Specifically, in \method, the input sequences of the Transformer network are segmented into two distinct components: {the \textit{memory} part and the \textit{target} part.} The target part adheres to the standard next-token prediction for modeling continuous sequences, while the memory part, sampled from longer sequences, serves as the conditional context for the prediction of the target part. To build the memory part, we apply a sparse sampling policy that decays with the distance from the target part, to obtain tokens and preserve their positions. Without any architectural modifications, our method can extend existing Transformer-based LLMs to capture long-range dependencies within a fixed window size during the training. Experimental results on multiple datasets also demonstrate the effectiveness and efficiency of \textsc{Sparse Training} to mitigate the complexity of the Transformer network in building long-sequence dependency.

Primary Area: applications to computer vision, audio, language, and other modalities

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 4609