SkyLadder: Better and Faster Pretraining via Context Window Scheduling

  • 作者: Tongyao Zhu, Qian Liu, Haonan Wang, Shiqi Chen, Xiangming Gu, Tianyu Pang, Min-Yen Kan
  • 机构: National University of Singapore, Sea AI Lab, City University of Hong Kong

A1 主要贡献

本文研究的核心问题是，在大型语言模型（LLM）预训练中，不断扩大的上下文窗口是否真的能带来性能提升。普遍观点认为，使用更长的上下文窗口进行预训练可以减少文档截断并保持连贯性，从而获得与短上下文模型相当甚至更好的性能。然而，作者通过严格控制的实验挑战了这一观点。

研究目标与发现：
- 核心发现：在固定的 token 预算下，使用较短上下文窗口预训练的模型在标准基准测试中的平均性能始终优于使用长上下文的模型。这一性能差距即使在使用先进的文档打包策略后也依然存在。
- 面临的权衡：尽管短上下文训练效果更好，但为了让模型具备处理长序列的能力，仍然需要让模型接触长序列。这带来了长上下文能力与预训练效果之间的权衡。

创新点（SkyLadder）：
为解决上述权衡问题，本文提出了 SkyLadder，一种简单而有效的上下文窗口调度策略。
- 核心思想：在预训练过程中逐步扩大上下文窗口的大小，从一个极小的窗口（如8个 token）开始，线性增加到目标的长上下文窗口（如32,768个 token）。
- 主要优势：
1. 性能提升：与基线方法相比，SkyLadder 在通用基准测试中取得高达 3.7% 的一致性增益，同时在长上下文任务上达到或超过基线水平。
2. 效率提升：训练速度最多可提升 22%。
3. 优化注意力模式：SkyLadder 展现出更集中、更有效的注意力模式。

实验验证：
作者通过在 1000 亿个 token 上预训练 1B 参数模型（上下文最长 32K）和 3B 参数模型（上下文 8K）来验证 SkyLadder 的有效性。实验结果表明，该方法在短上下文和长上下文评估任务中均优于朴素的长上下文预训练基线。

结论与建议：
作者认为上下文窗口的长度是预训练中的一个重要维度，应该在训练过程中进行调度。他们推荐一种渐进式方法，从8个 token 的小上下文开始，根据训练步数线性增加。对于给定的目标上下文窗口（例如32K），建议将大约60%的总训练 token 分配给这个扩展阶段，以获得比基线更强的下游性能。


图1：左图：近年来LLM的预训练上下文窗口不断增长。右图：不同预训练上下文窗口大小（颜色编码）的1B参数模型在九个下游任务上的平均性能（百分比）。增加上下文窗口会降低整体性能。

相关工作

上下文窗口调度。早期工作在BERT和GPT2等较小模型中探索了逐渐增加上下文窗口的方法，以提高训练稳定性和效率【索引35，Pre-training a BERT with curriculum learning by increasing block-size of input text，RANLP 2021】【索引28，The stability-efficiency dilemma: investigating sequence length warmup for training gpt models，NeurIPS 2022】【索引21，Growlength: Accelerating LLMs pretraining by progressively growing training length，arXiv 2023】。值得注意的是，Li等人【索引28，The stability-efficiency dilemma: investigating sequence length warmup for training gpt models，NeurIPS 2022】提出了长度预热以实现更稳定的训练，但未显示明确的性能增益，而Jin等人【索引21，Growlength: Accelerating LLMs pretraining by progressively growing training length，arXiv 2023】则专注于400M模型的训练加速。本文扩展了这些发现，首次证明上下文窗口调度在更大规模（高达3B参数）上显著提升了效率和性能。Pouransari等人【索引38，Dataset Decomposition: Faster llm training with variable sequence length curriculum，NeurIPS 2024】提出了一种并行方法，按长度分割训练文档，但Fu等人【索引10，Data engineering for scaling language models to 128k context，ICML 2024】警告说，这种分割可能引入领域偏差，因为较长的文本通常集中在特定领域（如书籍）。最近在长上下文窗口的持续预训练方面的发展【索引37，Yarn: Efficient context window extension of large language models，ICLR 2024】【索引55，When precision meets position: Bfloat16 breaks down rope in long-context training，arXiv 2024】【索引12，How to train long-context language models (effectively)，arXiv 2025】，也可以被看作是采用不同策略的上下文窗口调度（如图2所示）。本文的工作首次展示了上下文窗口调度在有效性和效率两方面的优势，为其在标准和长上下文基准测试中的好处提供了经验证据。


图2：训练期间上下文窗口调度的示意图比较。

长上下文语言模型。长上下文语言模型因其能够捕捉大文本窗口中的长距离依赖关系而备受关注。大多数现有方法遵循持续预训练的范式【索引10，Data engineering for scaling language models to 128k context，ICML 2024】【索引57，Effective long-context scaling of foundation models，NAACL 2023】，即通过专门的微调或额外训练将预训练的骨干模型扩展到更长的上下文。一些工作提出干预位置嵌入以适应更长的序列【索引1，Why does the effective context length of llms fall short?，arXiv 2024】【索引31，NTK-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degration，Reddit 2023】【索引37，Yarn: Efficient context window extension of large language models，ICLR 2024】【索引4，Extending context window of large language models via positional interpolation，arXiv 2023】【索引22，LLM maybe longlm: Selfextend LLM context window without tuning，ICML 2024】，而另一些则在更长序列的语料库上进行扩展预训练【索引12，How to train long-context language models (effectively)，arXiv 2025】【索引55，When precision meets position: Bfloat16 breaks down rope in long-context training，arXiv 2024】【索引32，A controlled study on long context extension and generalization in llms，arXiv 2024】【索引63，Longskywork: A training recipe for efficiently extending context length in large language models，arXiv 2024】。本文的方法与以往不同，我们从头开始训练原生长上下文模型，而不是在后训练阶段修改预训练模型。与采用恒定调度策略的朴素长上下文预训练基线相比，我们的方法在多个长上下文任务上取得了显著的增益，强调了从头开始训练的好处。这些发现表明，我们的方法可能是未来研究构建更长上下文窗口语言模型的一个有前途的方向。


图3：预训练数据准备工作流程示意图，突出了几个关键决策。主要考虑因素包括数据打包方法、使用的注意力掩码类型（因果掩码或文档内掩码），以及确定合适的上下文窗口长度L。

A3 上下文窗口如何影响预训练

研究问题。为了公平可比地研究上下文窗口如何影响预训练，本文在固定的总词元（token）数量下，从头开始预训练了上下文窗口从512到16,384个词元不等的语言模型，并通过困惑度（perplexity）和下游任务基准进行评估。研究分析了上下文窗口大小如何影响模型性能，并探讨了数据打包和掩码（masking）策略与窗口大小的相互作用。

3.1 打包、掩码和上下文窗口

模型架构与数据处理。大多数现代LLM基于一个仅解码器（decoder-only）的Transformer架构【索引54，Attention is all you need，NeurIPS 2017】，具有一个固定的上下文窗口大小，记为L。相比之下，预训练语料库$D = \{d_1, d_2, d_3, \dots, d_n\}$由长度与L不同的各种文档组成。因此，预训练前的一个关键步骤是将这些文档打包成长度为L的序列。形式上，一个打包后的序列$C_i$被构造为$C_i = \text{Trunc}(d_{i,1}) \oplus d_{i,2} \oplus \dots \oplus d_{i,n-1} \oplus \text{Trunc}(d_{i,n})$，其中$\oplus$表示拼接，$\text{Trunc}(\cdot)$表示对文档进行截断以确保$len(C_i) = L$。遵循先前的工作【索引44，In-context pretraining: Language modeling beyond document boundaries，ICLR 2024】【索引64，Analysing the impact of sequence composition on language model pre-training，ACL 2024】，序列$C_i$中的文档边界通过序列结束（[EOS]）标记明确标出。

注意力机制与掩码。序列打包后，输入被送入Transformer层进行下一个词元的预测训练。这些层的一个关键组成部分是注意力机制，可以表示为$A_{i,j} = q_i^\top k_j$，然后是$\text{Attn}(X) = \text{Softmax}(A + M)$。在仅解码器模型中，会应用一个掩码M来引入约束。一种常见的方法是使用因果掩码（causal mask），它通过屏蔽掉（设置为$-\infty$）对应于未来位置的注意力分数来确保每个位置只能关注前面的词元：对于$j > i$，$M_{ij} = -\infty$，否则$M_{ij} = 0$。最近提出的一种掩码方案，称为文档内掩码（intra-doc mask）【索引64，Analysing the impact of sequence composition on language model pre-training，ACL 2024】【索引13，The llama 3 herd of models，arXiv 2024】，施加了一个约束，即只有当词元属于同一文档时才能相互关注。假设每个文档d的起始索引为$s_d$，结束索引为$e_d$，则该掩码可以表示为：当存在某个文档d使得$s_d \le i, j \le e_d$且$j \le i$时，$M_{ij}^{\text{intra}} = 0$，否则$M_{ij}^{\text{intra}} = -\infty$。模型使用标准的交叉熵损失在长度为L的打包序列上进行训练。预训练数据处理的工作流程如图3所示。

3.2 上下文窗口大小的初步研究

实验设计。根据第1节的讨论，本研究首先通过一个受控实验来探究上下文窗口大小对模型性能的影响。具体来说，研究者们预训练了具有不同上下文窗口大小的语言模型，同时保持所有其他实验设置不变。这使得能够纯粹地分析上下文窗口对模型性能的影响，旨在理解更长的上下文窗口是否天生就会导致更好或更差的模型性能。

关键变量。上下文窗口大小决定了每个打包序列中包含的词元数量。然而，如前所述，还有几个额外因素会影响上下文窗口内的内容：（1）打包方法决定了哪些文档构成了上下文窗口，不同的打包策略可以显著改变词元序列的组成；（2）掩码方法决定了在同一上下文窗口内是否启用跨文档注意力。掩码的选择影响了来自不同文档的信息在训练过程中的交互方式。

打包与掩码策略。为了研究打包的影响，本文采用了两种策略：随机打包（random packing）和语义打包（semantic packing）。对于随机打包，文档被随机拼接，没有特定的顺序。对于语义打包，受Shi等人【索引44，In-context pretraining: Language modeling beyond document boundaries，ICLR 2024】的启发，本文从语料库中检索并拼接语义相关的文档，旨在将它们保持在同一个上下文窗口内。在实验了密集检索器【索引20，Unsupervised dense information retrieval with contrastive learning，arXiv 2021】和词汇检索器BM25之后，发现BM25性能更强，并选择其作为研究重点。对于掩码，基线方法是因果掩码，其中每个词元可以关注同一上下文窗口内所有在它之前的词元，无论文档边界如何。相反，最近的研究【索引64，Analysing the impact of sequence composition on language model pre-training，ACL 2024】【索引9，Fewer truncations improve language modeling，ICML 2024】表明，禁用跨文档注意力，从而启用文档内注意力，可以提高性能。为便于后续讨论，本文将随机打包与因果掩码表示为Random，BM25打包与因果掩码表示为BM25，随机打包与文档内掩码表示为IntraDoc。

训练设置。本文使用TinyLlama代码库【索引61，Tinyllama: An open-source small language model，arXiv 2024】从头开始预训练模型，并研究了120M、360M和1B参数的模型。考虑到语义打包中检索相关的巨大计算成本，本文从SlimPajama数据集【索引46，SlimPajama: A 627B token cleaned and deduplicated version of RedPajama，HuggingFace 2023】的CommonCrawl (CC)子集中随机选择了约300亿个词元作为预训练语料库。所有模型都训练了多达1000亿个词元（约3.3个周期）。为了确保实验之间的一致性，本文严格控制了所有其他设置，对所有上下文窗口都保持了相同的批量大小和学习率调度。所有模型都采用了旋转位置编码（RoPE）【索引47，Roformer: Enhanced transformer with rotary position embedding，Neurocomputing 2024】来编码位置信息。附录A.3和A.4提供了进一步的模型架构细节和训练设置。

评估方法。对于所有模型规模，本文使用原始数据集验证文档上的困惑度（PPL）作为关键指标，这与既定实践【索引10，Data engineering for scaling language models to 128k context，ICML 2024】【索引24，Scaling laws for neural language models，arXiv 2020】【索引17，Training compute-optimal large language models，NeurIPS 2022】一致。需要注意的是，在比较不同上下文窗口的模型时（例如，一个2K上下文模型和一个8K上下文模型），必须确保评估序列能够适应较短模型的上下文窗口，以保持公平比较。本文还评估了1B模型在下游标准基准测试上的表现：HellaSwag【索引60，Hellaswag: Can a machine really finish your sentence?，ACL 2019】，ARC-Easy和ARC-Challenge【索引6，Think you have solved question answering? try arc, the ai2 reasoning challenge，arXiv 2018】，Winogrande【索引42，Winogrande: An adversarial winograd schema challenge at scale，Communications of the ACM 2021】，CommonsenseQA【索引48，CommonsenseQA: A question answering challenge targeting commonsense knowledge，NAACL 2019】，OpenBookQA【索引34，Can a suit of armor conduct electricity? a new dataset for open book question answering，EMNLP 2018】，PIQA【索引2，PIQA: Reasoning about physical commonsense in natural language，AAAI 2020】，Social-QA【索引43，Social IQa: Commonsense reasoning about social interactions，EMNLP-IJCNLP 2019】，和MMLU【索引16，Measuring massive multitask language understanding，ICLR 2021】。本文采用OLMES套件【索引15，OLMES: A standard for language model evaluations，NAACL 2025】进行评估，因为它已被证明在使用精心策划的5-shot演示【索引12，How to train long-context language models (effectively)，arXiv 2025】时能提供可靠且稳定的结果。

3.3 实验结果

核心发现。图1展示了使用Random设置和1B参数模型得到的主要实验结果。结果表明，上下文窗口大小显著影响LLM的性能，较短的上下文通常能带来更好的性能。为了进一步探究导致这一观察结果的因素，本文进行了全面的分析，以检查可能影响结论的潜在变量。图4展示了研究结果，并得出了四个关键发现：

四个关键发现。

  1. 在较短上下文上训练的优势在不同模型规模上是一致的。
  2. 这种优势与所采用的打包和掩码方法无关。
  3. 它也与位置编码的使用无关。
  4. 最佳的打包和掩码策略是IntraDoc，其性能优于其他策略，可能是因为它在预训练期间引入了更多的短上下文。

发现（1）和（2）的论证。如图4所示，无论是在(a)中的模型规模，还是在(b)中的打包和掩码方法，较短的预训练上下文窗口通常会导致在基准测试中获得更高的平均性能。基准测试的发现与验证PPL的趋势一致，即较短的上下文窗口总是产生较低的PPL。

发现（3）的论证。当使用较短的上下文窗口时，可能会假设模型更频繁地学习到较近位置的位置编码模式，从而在标准基准测试上表现更好。为了验证这一假设，本文遵循先前的工作【索引25，The impact of positional encoding on length generalization in transformers，NeurIPS 2023】，通过在预训练期间完全排除RoPE来系统地进行消融实验。在图4(c)中，即使在没有位置编码的情况下，使用短上下文窗口训练的模型仍然优于其长上下文的对应模型。这表明短上下文的优势与位置编码无关。

发现（4）的论证。从图4(b)中可以观察到，与Random和BM25相比，IntraDoc在所有上下文窗口大小上都取得了最佳的验证PPL，并且在标准基准测试中也表现出持续更高的性能（参见附录A.7.1）。这就引出了一个问题：为什么IntraDoc表现出色？作者将这一优势归因于IntraDoc的上下文窗口大小分布，它隐式地增加了短上下文的普遍性。如图4(d)所示，尽管序列长度为8K，但实际上只有不到1%的上下文窗口达到了这个限制。虽然先前的工作将IntraDoc的成功归因于减少了上下文噪声【索引64，Analysing the impact of sequence composition on language model pre-training，ACL 2024】，但本文确定了一个互补的因素——平均上下文窗口大小的减小——是其强大性能的关键因素。也就是说，作者假设IntraDoc的有效性也可能与短上下文窗口密切相关。


图4：不同因素在不同上下文窗口大小上的消融研究。请注意，验证PPL是在验证文档上使用512个词元的滑动窗口大小获得的。(a)中的打包策略是随机的，(b)和(c)中的模型大小分别是1B和120M。请注意，(d)中的上下文窗口指进行下一个词元预测时可用的前文词元数量（计算细节见A.6节）。

A2 方法细节：SkyLadder上下文窗口调度

本节介绍SkyLadder方法，该方法在预训练期间逐步扩大上下文窗口。

4.1 方法

核心思想。受学习率调度的启发，我们探索了在预训练期间从短到长动态调度上下文窗口是否能带来性能提升。该方法可以通过对一个长的、打包好的序列应用多个局部的“迷你”因果掩码来实现。我们在图5中说明了这种掩码策略。


图5：使用Random和IntraDoc的SkyLadder示意图。该示例显示了一个由两个文档组成的打包序列（长度为L）。对于SkyLadder，上下文窗口w从一个较小的值开始，并在训练过程中动态调整，最终收敛到Random或IntraDoc的掩码模式。

形式化定义。我们定义一个局部窗口长度$w$。相关的掩码$M^w$定义如下：当$\lfloor i \rfloor_w \le j \le i$时，$M_{ij} = 0$，否则$M_{ij} = -\infty$。其中$\lfloor i \rfloor_w$计算小于或等于$i$的$w$的最大倍数，有效地为位置$i$的查询词元定义了一个块级注意力掩码。我们通过每个训练步骤$t$按一个常数因子线性向上调整窗口大小：$w(t) = \min(w_e, w_s + \lfloor\alpha t\rfloor)$，其中$w_e$和$w_s$分别表示结束和开始的上下文窗口大小。这里，$\alpha$表示扩展速率，$t$对应于训练步骤。随着训练的进行，当动态上下文窗口大小$w(t)$最终达到期望的（长）上下文窗口大小$L = w_e$时，它将保持在该值固定。此时，注意力掩码等同于一个完整的因果掩码。值得注意的是，该方法通过掩码来修改有效的上下文窗口，这与序列的打包方式无关。因此，这个掩码$M^w$可以与保持文档间注意力边界的$M^{\text{Intra}}$集成；它可以与大多数打包和掩码策略无缝结合。

4.2 实验设置

训练配置。我们遵循3.2节中的相同设置，在100B个词元上预训练具有8K上下文的语言模型。我们默认设置$w_s = 32$和$\alpha = 1/8$，这意味着模型大约需要64K步（约64B个词元）才能达到最终期望的上下文窗口L=8192。所有基线模型和SkyLadder模型均使用Flash Attention 2【索引7，FlashAttention-2: Faster attention with better parallelism and work partitioning，ICLR 2024】实现（伪代码见A.5）。为了公平比较，我们固定了所有其他超参数，如学习率调度、批量大小等。由于资源限制，我们没有进行广泛的超参数搜索来获得$w(t)$、$\alpha$和$w_s$的最佳组合。在我们的消融研究中，我们表明只要这些超参数在合理范围内，它们对性能的影响可以忽略不计。

评估配置。我们使用3.2节中提到的相同套件进行标准基准测试。为了在8K长度内评估长上下文问答的性能，我们利用了多文档问答（MDQA）基准【索引30，Lost in the middle: How language models use long contexts，TACL 2024】中的30个文档设置。这是一个被广泛采用的基准，对于1B规模的模型被证明是可靠的【索引38，Dataset Decomposition: Faster llm training with variable sequence length curriculum，NeurIPS 2024】【索引64，Analysing the impact of sequence composition on language model pre-training，ACL 2024】，平均长度约为6K个词元。我们还选择了RULER【索引18，RULER: What’s the real context size of your long-context language models?，CoLM 2024】中的合成任务，如Yen等人【索引59，Helmet: How to evaluate long-context language models effectively and thoroughly，ICLR 2025】所定义的。我们选择的任务设置会填满模型的目标上下文窗口L。

4.3 实验结果

主要发现。表1和表2展示了主要结果，突显了SkyLadder在标准基准、阅读理解任务和长上下文基准上取得的显著改进。例如，与Random基线相比，集成SkyLadder在MMLU（+2.5%）、ARC-E（+7.4%）和HellaSwag（+4%）等标准任务上带来了显著的性能提升。这表明使用SkyLadder的模型在预训练期间更擅长学习常识知识。此外，我们的方法在许多基准上进一步提升了强基线IntraDoc的性能。同时，对于像MDQA这样的真实长上下文基准，SkyLadder的性能与基线相当或超过基线。对于RULER，性能差异可能是由于其合成性质和规模较小【索引55，When precision meets position: Bfloat16 breaks down rope in long-context training，arXiv 2024】引起的波动。更多长上下文评估可以在A.7.3节中找到，证实SkyLadder在长上下文评估上与基线相当或更优。除了Random和IntraDoc，我们还验证了SkyLadder在短任务和长任务上都提升了BM25模型的性能（A.7.4节）。


表1：在100B CC词元上预训练的不同1B模型在标准基准上的性能（准确率%）。*表示相对于基线的统计学显著改进（详见§A.7.3）。


表2：在100B CC词元上预训练的1B模型使用不同方法在阅读理解和长上下文基准上的性能（准确率%）。详细设置见附录A.7.3。

数据质量泛化性。为了解决在短上下文中观察到的益处可能源于CC中高噪声水平的担忧，我们使用FineWeb-Pro数据集【索引65，Programming every example: Lifting pre-training data quality like experts at scale，arXiv 2024】进行了额外的实验，这是一个包含100B词元的精心策划的高质量数据集。如表4所示，改进的数据质量确实带来了显著的性能提升。然而，我们的关键发现保持一致：IntraDoc继续优于Random，而SkyLadder在两个基线上都持续带来了显著的改进。这表明我们的方法可以推广到不同质量的语料库。


表4：在FineWeb-Pro上使用8K上下文窗口预训练的1B模型的性能（任务平均准确率%）。

任务泛化性。我们进一步检验SkyLadder是否能泛化到自然语言任务之外。遵循Ding等人【索引9，Fewer truncations improve language modeling，ICML 2024】的方法，我们使用Starcoder分词器【索引29，Starcoder: may the source be with you!，arXiv 2023】在100B个Python代码上预训练了1B代码模型。我们观察到代码预训练的训练损失（约0.9）低于自然语言（约2.1），这表明代码中的结构使训练更容易。然而，如表3所示，在贪婪解码和采样设置下应用SkyLadder仍然有显著改进，尤其是在目标上下文长度为32K时。这展示了SkyLadder在编码和可能超出自然语言建模的其他推理任务中的潜力。


表3：在100B Python代码数据上预训练的1B模型的性能（%）。我们遵循Huang等人【索引19，Opencoder: The open cookbook for top-tier code large language models，arXiv 2024】的协议在HumanEval【索引3，Evaluating large language models trained on code，arXiv 2021】和BigCodeBench【索引66，BigCodeBench: A large-scale benchmark for multilingual program synthesis and comprehension，arXiv 2024】上进行评估。t是采样温度。SkyLadder显示出持续的改进，尤其对于32K上下文模型。

4.4 可扩展性实验

模型与上下文规模扩展。我们检验了随着模型参数的增加和上下文窗口尺寸的扩展，SkyLadder的改进是否持续存在。我们使用了我们计算资源允许的最大模型和上下文尺寸。

模型大小。我们在Fineweb-Pro数据集上对120M、360M和3B参数这三种模型尺寸进行了实验。表5表明，使用SkyLadder的模型在所有模型尺寸上都一致地取得了更好的标准基准性能。对于长上下文任务，我们的方法对120M模型没有益处，可能是由于它们处理长序列的能力有限。然而，在3B模型上的性能增益是显著的。我们观察到一个积极的扩展趋势：随着模型尺寸的增长，性能提升也随之增加，这表明我们的方法有潜力应用于超出我们当前规模的更大型号的模型。由于需要大量的计算资源，我们将探索更大型号作为未来的工作。


表5：不同尺寸模型的性能（平均准确率%）。

上下文窗口大小。为了检验SkyLadder是否能有效扩展到更长的上下文窗口，我们在100B的FineWebPro词元上训练了具有32K上下文窗口的1B模型。我们将α调整为1/2，以确保最终的上下文窗口在预训练结束前扩展到32K。如表6所示，我们的模型在标准和长基准上都表现出强大的性能。此外，与基线方法（1.8%）相比，SkyLadder在8K和32K模型之间的性能差异（0.9%）大大减少，这缓解了我们早期研究中描述的性能下降问题。值得注意的是，与基线Random方法相比，SkyLadder在早期阶段用逐渐变短的上下文来训练模型。这揭示了一个与直觉相反的见解：即使模型在长上下文上进行评估，用长上下文窗口朴素地训练模型也并非总是最优的。相反，在预训练期间策略性地调度上下文窗口可以产生更好的结果。


表6：在100B FineWeb-Pro词元上使用32K上下文窗口训练的1B模型的性能（%）。

4.5 消融研究

实验设置。我们现在检验SkyLadder调度中超参数的影响。为了控制计算成本，我们采用预训练120M模型，上下文为8K，在100B CC词元上的默认设置。

扩展速率。我们在图6（左）中研究了扩展速率α的影响。我们选择了从最慢（1/12）到最快（1）的不同α。我们的发现表明，对于短上下文，随着扩展速率减慢，性能通常会提高。然而，选择一个过慢的速率（例如1/12）会因为对长上下文的训练不足而对长上下文性能产生负面影响。因此，我们建议将α设置为1/8以取得良好平衡。


图6：具有不同扩展速率α（左）和初始窗口长度ws（右）的模型在512和8K上下文上的验证PPL。

初始上下文窗口。由于最终上下文窗口长度$w_e$固定为L，唯一剩下的超参数是$w_s$。直观上，将$w_s$设置得过大（例如接近L）会给调度留下很小的空间，导致性能次优。在图6（右）中，我们展示了当$w_s$设置得相对较小（例如8）时，可以在短上下文和长上下文中都取得很好的性能。这表明在我们默认设置中仍有进一步改进的潜力。因此，我们建议从一个小的上下文窗口开始，例如8个词元。

调度类型。SkyLadder中默认的调度方法是线性调度。我们评估了不同上下文窗口调度类型（更多细节见附录A.7.4中的表20和图12）：(1) 阶梯式线性（Stepwise Linear）将窗口大小$w(t)$四舍五入到1K的倍数，形成一个阶梯函数；(2) 正弦（Sinusoidal）在早期阶段快速增加，然后减速；(3) 指数（Exponential）开始缓慢但急剧加速；(4) 持续预训练（Continual pretraining）设置，先用4K上下文窗口训练约97B词元，然后在最后3B词元切换到32K上下文。表7显示，线性和正弦调度在长任务上优于指数变体，可能是因为指数调度在开始时有较长的短上下文预训练，未能充分训练长上下文。最后，最常用的持续预训练设置整体表现不佳，表明突然的上下文变化对短和长性能都有害。这些发现表明，上下文窗口调度优于恒定长上下文预训练和持续预训练。


表7：使用不同调度方法训练的32K上下文窗口1B模型的比较。数字为平均准确率（%）。

总结。总的来说，我们得出结论，调度应从一个小的$w_s$开始，并且扩展应该是渐进的。我们将其留给未来的工作来研究更先进的调度并发现最优配置。例如，调度可能需要根据不同的模型大小进行调整。关于与BM25的结合、混合注意力、循环调度以及在计算预算下的调度的更多消融实验可以在附录A.7.4中找到。

4.6 分析与讨论

训练效率。在表8中，我们观察到使用SkyLadder时训练效率有显著提升。对于8K模型，由于计算注意力时上下文窗口减小，SkyLadder将训练时间加快了13%。对于32K上下文窗口，效率增益更为显著：我们的方法节省了22%的训练时间，同时取得了更好的性能。由于注意力计算量的减少，FLOPs的节省比实际时间节省更大。


表8：不同上下文窗口大小L的1B模型的相对训练时间和计算效率比较。FLOPs计算遵循Zhang等人【索引61，Tinyllama: An open-source small language model，arXiv 2024】的方法。更大的上下文窗口带来更多的效率增益。

注意力模式。接下来，我们研究为什么SkyLadder虽然总体上是在短上下文中训练，却能持续优于基线。由于语言模型依赖注意力机制来编码上下文信息，我们研究了注意力模式的变化。具体来说，在预训练期间，我们监控了（i）注意力熵（图7中的实线）的动态，其中较低的熵与更好的下游性能相关【索引62，Attention entropy is a key factor: An analysis of parallel context encoding with full-attention-based pre-trained language models，arXiv 2024】；（ii）注意力沉淀（attention sink）【索引56，Efficient streaming language models with attention sinks，ICLR 2024】，即上下文中的初始词元获得不成比例的高注意力。我们利用Gu等人【索引14，When attention sink emerges in language models: An empirical view，arXiv 2024】中的指标来量化注意力沉淀的幅度。如图7（虚线）所示，与基线Random相比，SkyLadder表现出更低的注意力熵，表明其注意力模式更集中。然而，同时观察到注意力沉淀的出现较慢且幅度较低。这表明SkyLadder的注意力集中在上下文中的关键信息上，而不是初始词元，这解释了其性能的提升。


图7：预训练1B模型（8K上下文）期间注意力沉淀和熵的动态。SkyLadder延迟了注意力沉淀的出现，同时降低了整体熵，表明其注意力模式更有效。

训练稳定性。为了进一步理解SkyLadder性能更好的原因，我们分析了预训练上下文长度对训练动态的影响。我们预训练了不同上下文长度的120M参数模型。我们首先遵循K2【索引50，Kimi k2: Open agentic intelligence，arXiv 2025】的方法，在整个预训练过程中监控最大注意力logits（$S_{max} = \max_{i,j} q_i \cdot k_j$ for all i, j）。大的注意力logit表明一个注意力头功能失常，可能导致数值不稳定。在图8中，我们观察到用16K token的长上下文进行预训练会导致最大注意力logits爆炸，而较短的窗口则导致较低的注意力logits。


图8：不同上下文长度（不同颜色）模型训练期间的最大注意力logits。

损失与梯度行为。接下来，我们通过计算预训练前N=30K步的四个稳定性指标来研究损失和梯度行为，其中$L_t$是训练损失，$G_t$是裁剪前的梯度范数：
- 损失波动性（Loss Volatility）：衡量滑动窗口（w=10）内损失的局部波动，计算公式为$\frac{1}{N} \sum_{t=1}^N \text{Std}(L_{t-w+1}, \dots, L_t)$。值越低表示训练越稳定。
- 损失平滑度（Loss Smoothness）：连续步骤间的平均损失变化，$\frac{1}{N-1} \sum_{t=2}^N |L_t - L_{t-1}|$。值越小意味着收敛越平滑。
- 平均损失比（Mean Loss Ratio）【索引28，The stability-efficiency dilemma: investigating sequence length warmup for training gpt models，NeurIPS 2022】：衡量损失相对于迄今为止最佳损失的暂时增加，$\frac{1}{N-1} \sum_{t=2}^N \frac{L_t}{\min(L_1, \dots, L_{t-1})}$，值越小表示损失尖峰越少。


表9：预训练不同上下文长度的120M模型期间的训练稳定性指标。所有指标都是在前300亿个词元上平均的。↓表示值越小越好。

  • 平均梯度范数（Average Gradient Norm）：$\frac{1}{N} \sum_{t=1}^N \min(G_t, 1)$，值越大表示梯度更新越激进。

在表9中，长上下文模型表现出更高的波动性、不那么平滑的损失曲线、更频繁的向上尖峰和更大的梯度范数，所有这些都表明优化过程不太稳定。相比之下，短上下文模型收敛得更平滑，波动更小，梯度更新也更受控制。总之，这些结果揭示了短上下文预训练在注意行为和优化动态方面都天生更稳定。减少的数值不稳定性和更平滑的收敛可能使得梯度信号更一致，整体收敛更好，从而解释了它们优越的下游性能。

与相关工作比较。我们在表10中将我们的方法与另一种改进预训练的方法进行了比较。如第2节所述，Pouransari等人【索引38，Dataset Decomposition: Faster llm training with variable sequence length curriculum，NeurIPS 2024】提出了数据集分解（Dataset Decomposition, DD），通过将文档分割成不同长度的序列并在预训练期间使用课程学习。然而，这种方法不可避免地会引入领域偏差，因为不同领域中的文档长度是不同的【索引10，Data engineering for scaling language models to 128k context，ICML 2024】。这解释了为什么只有一个从短到长周期的DD未能超越IntraDoc基线。为了缓解这个问题，作者建议在长短数据之间迭代多个周期，这确实显著提高了性能。相比之下，我们的方法通过不根据长度改变数据顺序来避免这种偏差，从而取得了更好的性能。在附录A.7.4中，我们实验了各种循环调度，但没有观察到任何改进。事实上，我们注意到周期之间出现了损失尖峰（图14），这表明可能存在领域转移问题。这进一步支持了我们的方法更安全，因为它不破坏数据的自然顺序和分布。更多与其他相关工作【索引28，The stability-efficiency dilemma: investigating sequence length warmup for training gpt models，NeurIPS 2022】【索引21，Growlength: Accelerating LLMs pretraining by progressively growing training length，arXiv 2023】的讨论见A.8节，其中我们证明了我们的工作提供了新的见解，即在整个训练时间内调度上下文窗口可以提高效率和性能。


表10：SkyLadder与数据集分解（DD）在100B FineWeb-Pro词元上训练的1B模型之间的比较。数字为平均性能百分比。

A4 实验环境

  • 数据集：

    • 预训练：

      • SlimPajama【索引46，SlimPajama: A 627B token cleaned and deduplicated version of RedPajama，HuggingFace 2023】的CommonCrawl (CC) 子集，用于初步研究（约30B tokens）和主要实验（100B tokens）。
      • FineWeb-Pro【索引65，Programming every example: Lifting pre-training data quality like experts at scale，arXiv 2024】：一个高质量的100B tokens数据集，用于验证方法在不同数据质量下的泛化性。
      • Python代码：100B tokens的Python代码数据集，用于验证方法在代码生成任务上的泛化性。

    • 评估：

      • 标准基准：HellaSwag, ARC-Easy, ARC-Challenge, Winogrande, CommonsenseQA, OpenBookQA, PIQA, Social-QA, MMLU。
      • 长上下文基准：MDQA（30文档设置）, RULER（合成任务）, TOEFL, QuALITY。
      • 代码基准：HumanEval, BigCodeBench。
      • 其他：SQuAD, HotpotQA, NaturalQuestions, TriviaQA, RACE-high（阅读理解），以及用于RAG和多示例ICL的多个数据集。

  • 模型架构：

    • 120M, 360M, 1B参数模型基于TinyLlama架构【索引61，Tinyllama: An open-source small language model，arXiv 2024】。
    • 3B参数模型基于Llama3.2架构【索引13，The llama 3 herd of models，arXiv 2024】。
    • 详细配置见附录表11。所有模型均采用RoPE【索引47，Roformer: Enhanced transformer with rotary position embedding，Neurocomputing 2024】进行位置编码。

  • 硬件配置：

    • ≤ 1B模型：在一个内部集群上进行，该集群由配备40GB内存的NVIDIA A100节点组成。例如，一个1B模型在8个A100上训练100B tokens（8K上下文）耗时约200小时。
    • 3B模型：在H100节点上进行。

  • 软件配置：

    • 代码实现：基于开源的TinyLlama预训练框架【索引61，Tinyllama: An open-source small language model，arXiv 2024】。
    • 依赖库：使用了Flash Attention 2【索引7，FlashAttention-2: Faster attention with better parallelism and work partitioning，ICLR 2024】。
    • 优化器与调度：使用AdamW优化器，学习率调度为余弦退火（Cosine schedule）。
    • 评估工具：使用OLMES套件【索引15，OLMES: A standard for language model evaluations，NAACL 2025】进行标准基准评估。

A4 实验结果

  • 核心发现：短上下文预训练更优

    • 在固定的token预算下，使用较短上下文窗口预训练的模型在标准基准上的平均性能一致优于长上下文模型（图1右）。
    • 这一优势在不同模型规模（120M, 360M, 1B）、不同打包/掩码策略（Random, BM25, IntraDoc）下均成立（图4a, 4b）。
    • 该现象与位置编码（RoPE）无关，即使移除RoPE，短上下文模型依然表现更优（图4c）。
    • IntraDoc策略表现最佳，部分原因是它隐式地增加了短上下文的比例（图4d）。

  • SkyLadder方法有效性

    • 标准基准与阅读理解：SkyLadder显著提升了基线模型（Random和IntraDoc）在标准基准（如MMLU, ARC-E, HellaSwag）和阅读理解任务（如HotpotQA, SQuAD）上的性能（表1, 表2）。
    • 长上下文任务：SkyLadder在真实长上下文任务（如MDQA）上表现与基线相当或更好，弥补了短上下文训练可能带来的长程依赖捕捉能力不足的问题（表2, 表13）。
    • 泛化能力：
      • 数据质量：无论是在噪声较大的CommonCrawl数据还是高质量的FineWeb-Pro数据上，SkyLadder都表现出一致的性能提升（表4）。
      • 任务领域：在Python代码生成任务上，SkyLadder同样带来了显著的性能增益，尤其是在32K长上下文设置下（表3）。

  • 可扩展性分析

    • 模型规模：随着模型从120M扩展到3B，SkyLadder带来的性能增益也随之增加，显示出良好的扩展潜力（表5）。
    • 上下文长度：SkyLadder能有效扩展到32K的上下文窗口，并且相比基线方法，它显著缓解了因上下文窗口增长导致的在标准任务上的性能下降问题（表6）。

  • 效率与机制分析

    • 训练效率：SkyLadder通过在训练早期使用较短的上下文，显著提升了训练效率，8K上下文模型提速13%，32K模型提速22%，同时节省了大量FLOPs（表8）。
    • 注意力模式：SkyLadder训练的模型展现出更集中、有效的注意力模式，即更低的注意力熵和更弱的注意力沉淀效应，表明注意力更多地集中在关键信息而非初始token上（图7）。
    • 训练稳定性：短上下文训练本身具有更好的稳定性，表现为更低的注意力logit峰值、更平滑的损失曲线和更小的梯度范数，这可能是其性能更优的根本原因之一（图8, 表9）。

  • 消融研究与方法对比

    • 调度策略：实验证明，从一个小的初始窗口（如8 tokens）开始，并采用平滑、渐进的扩展方式（如线性和正弦调度）效果最好。突变式的调度（如持续预训练中的急剧切换）对性能有害（图6, 表7）。
    • 与其他方法对比：与数据集分解（DD）方法相比，SkyLadder在避免引入领域偏差的同时取得了更好的性能（表10）。

A5 结论

本文通过全面的受控研究，揭示了在预训练中上下文窗口选择的一个重要发现：使用较短的上下文窗口更有利于模型在标准基准测试上的性能。这一发现挑战了当前预训练中追求更长上下文窗口的趋势。

基于此，本文提出了SkyLadder，一种在预训练期间从短到长动态调度上下文窗口的策略。该方法不仅在下游任务性能上取得了显著提升，还大幅提高了计算效率。

核心结论：上下文窗口调度是预训练中一个重要的维度，值得更多关注。

未来工作：计划探索更动态、性能更优的调度策略，使其能够根据模型大小或预训练数据分布进行自适应调整。

A6 附录

A.1 局限性

研究局限。尽管本文进行了广泛的实验来研究上下文窗口对预训练的影响，并展示了SkyLadder的有效性，但我们承认仍存在一些局限性有待解决。首先，我们的实验规模最大到3B模型和32K上下文长度，而最新的大型语言模型通常要大得多，并且能够处理更长的上下文。然而，用长上下文窗口预训练一个大模型需要超出我们预算的巨大计算资源。在我们计算能力范围内，我们已尝试展示SkyLadder在不同语料库、上下文窗口大小、模型大小和下游任务上的泛化性。因此，我们将把SkyLadder应用于更大型号的模型作为未来的工作。其次，我们没有包含理论分析来解释SkyLadder的有效性，因为我们主要关注经验性见解。我们建议未来的工作可以研究上下文窗口与训练计算之间的关系，以获得最优的上下文窗口调度。

A.2 更广泛的影响

社会影响。本工作旨在研究语言模型预训练中上下文窗口选择的影响，并提出一种通过调度上下文窗口来加速预训练的方法。从积极的方面来看，这提高了语言模型预训练的效率，使其更易于获取，并减少了碳足迹。此外，它增强了预训练语言模型的性能，这可能导致在应用中获得更好的下游性能。预训练语言模型可能存在潜在的滥用，这超出了本工作的范围。

A.3 模型架构

模型配置。在表11中，我们列出了所训练模型的架构选择，包括基于TinyLlama架构【索引61，Tinyllama: An open-source small language model，arXiv 2024】的120M、360M和1B模型。3B模型基于Llama3.2架构【索引13，The llama 3 herd of models，arXiv 2024】。


表11：预训练语言模型的模型配置。

A.4 训练配置

训练超参数。我们在表12中包含了训练配置的详细信息。所有模型，无论大小或上下文窗口长度，都使用这同一套超参数进行训练。对于大多数超参数值，我们遵循TinyLlama【索引61，Tinyllama: An open-source small language model，arXiv 2024】项目，因此我们的结果具有高度可复现性。


表12：预训练语言模型的超参数设置。所有预训练模型遵循相同的结构。

A.5 实现

伪代码。我们提供了使用Flash Attention 2【索引7，FlashAttention-2: Faster attention with better parallelism and work partitioning，ICLR 2024】实现SkyLadder的伪代码。唯一的变化是应用大小为w的局部因果掩码，并在IntraDoc场景下将它们与原始文档边界结合。它可以很容易地在计算注意力之前集成到任何模型中。训练流程的其余部分保持不变。

# SkyLadder与Flash Attention 2
# q, k, v: 经过RoPE编码的query, key, value张量
# doc_boundaries: 每个文档的EOS token位置
# is_intradoc: 文档内注意力标志
# training_step: 当前全局步数
# L: 最大上下文窗口长度

# 获取当前窗口大小
w = min(L, get_current_mask_length(training_step))

# 每w个token设置一个断点（如果使用IntraDoc掩码，则在文档边界处也设置断点）
mask_boundaries = np.arange(w, L, w)
if is_intradoc:
    mask_boundaries = np.union1d(mask_boundaries, doc_boundaries)

# 为flash attention计算最大段长度和累积长度
max_seqlen = get_max_seqlen(mask_boundaries, L)
cu_seqlens = get_cu_seqlens(mask_boundaries, L)

attn = flash_attn_varlen_func(
    q, k, v, cu_seqlens, max_seqlen, causal=True
)

A.6 每词元上下文窗口的定义

定义解释。在图4(d)中，我们展示了IntraDoc和Random之间上下文窗口分布的差异。需要澄清的是，上下文窗口大小指的是在进行下一个词元预测时，上下文窗口中可用的前文词元数量。这与(a)和(b)中的上下文长度L（模型的预训练上下文窗口）不同。

形式化定义。考虑一个位于索引$i$的词元和一个注意力掩码矩阵M，其中条目$M_{i,j} = 0$表示词元$i$可以关注词元$j$，否则为$-\infty$。第$i$个词元的上下文窗口大小$C_i$定义为$C_i = \sum_{j=1}^{i-1} \mathbf{1}_{\{M_{i,j} = 0\}}$，其中$\mathbf{1}_{\{\cdot\}}$是指示函数，当$M_{i,j} = 0$时返回1，否则返回0。实质上，$C_i$是第$i$个词元可用的上下文词元数量，而$C_i$在所有预训练词元上的分布如图4(d)所示。

不同策略下的分布。对于Random策略，因果掩码是三角矩阵：第$i$个词元的上下文窗口大小等于$i-1$（即$C_1=0, C_2=1$等）。因此，$C_i$的分布是均匀的。相比之下，IntraDoc通过限制跨文档注意力有效地缩短了上下文长度。

A.7 附加结果

A.7.1 上下文窗口研究

验证困惑度。在图9中，我们绘制了在Random、IntraDoc和BM25设置下，具有不同上下文窗口的模型的验证困惑度。我们观察到一个一致的趋势，即在所有设置下，短上下文模型在较短序列上的评估困惑度较低。


图9：在不同上下文长度下，模型在滑动窗口为512的验证困惑度。

不同策略的性能比较。在图10中，我们绘制了具有不同打包或掩码策略的模型的评估困惑度和下游性能。我们得出结论，总体而言，IntraDoc取得了最佳性能，具有持续较低的PPL和较高的下游准确率。我们认为这部分是由于IntraDoc模型是在较短的上下文窗口上训练的。


图10：左图：不同打包或掩码策略模型的评估困惑度。右图：不同模型在9个任务上的下游性能。

A.7.2 上下文窗口研究的消融实验

RoPE的基底。已有研究表明，RoPE的值可能对模型的长上下文性能有显著影响，且更长的上下文需要更大的基底【索引33，Base of RoPE bounds context length，arXiv 2024】。因此，我们将RoPE基底增加到100,000，根据Men等人【索引33，Base of RoPE bounds context length，arXiv 2024】的研究，这个值足够大。在图11中，我们观察到长上下文模型在长上下文评估上有所改进。然而，短模型和长模型之间的巨大差距依然存在，因此拒绝了RoPE基底是短上下文模型优越性能的关键贡献因素这一假设。


图11：不同上下文窗口和RoPE基底θ下，验证困惑度与训练词元数的关系。评估是在验证文档上使用不同长度（x轴）的滑动窗口进行的。

A.7.3 SkyLadder评估

统计检验。我们在表1中测试了我们的模型和基线之间性能差异的统计显著性。我们使用McNemar检验，因为两个模型是在同一组问题上进行评估的。原始的OLMES套件从每个基准的完整评估套件中抽样1000个样本。相比之下，在进行McNemar检验时，我们在完整的数据集上评估模型，以获得更具统计意义的结果。我们注意到，OpenBookQA只有500个问题，这使得获得统计显著性更加困难。

阅读理解。对于阅读理解，我们评估了以下基准：Hotpot QA (2-shot)【索引58，HotpotQA: A dataset for diverse, explainable multi-hop question answering，EMNLP 2018】，SQuAD (4-shot)【索引41，SQuAD: 100,000+ questions for machine comprehension of text，EMNLP 2016】，NaturalQuestions (NQ) (2-shot)【索引26，Natural questions: A benchmark for question answering research，TACL 2019】，TriviaQA (2-shot)【索引23，TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension，ACL 2017】，以及RACE-high (0-shot)【索引27，RACE: Large-scale ReAding comprehension dataset from examinations，EMNLP 2017】。我们遵循Zhao等人【索引64，Analysing the impact of sequence composition on language model pre-training，ACL 2024】的设置，其中NQ和TriviaQA使用检索到的文档作为上下文。对于RACE，我们使用lm-evaluation-harness【索引11，A framework for few-shot language model evaluation，Zenodo 2024】来比较选项之间的PPL。

长上下文评估。我们对最大的3B模型（8K上下文）提供了额外的长上下文评估，以减轻在小模型上使用合成基准的性能不稳定性。我们首先遵循【索引38，Dataset Decomposition: Faster llm training with variable sequence length curriculum，NeurIPS 2024】在阅读理解基准TOEFL【索引5，Supervised and unsupervised transfer learning for question answering，arXiv 2018】【索引53，Towards machine comprehension of spoken content: Initial toefl listening comprehension test by machine，arXiv 2016】和QuALITY【索引36，QuALITY: Question answering with long input texts, yes!，NAACL 2022】上评估模型准确率。接下来，我们评估模型在检索增强生成（RAG）上的性能，其中模型被提供了许多相关但可能带有噪声的上下文，并需要定位正确的信息。如表13所示，SkyLadder在所有评估的RAG和阅读理解数据集上始终优于基线，突显了其在长上下文中定位正确答案的能力。此外，我们还在文本分类基准上测试了模型的上下文学习能力【索引64，Analysing the impact of sequence composition on language model pre-training，ACL 2024】【索引44，In-context pretraining: Language modeling beyond document boundaries，ICLR 2024】。表14的结果表明，SkyLadder在标签数量多的任务（如DBpedia）上显示出显著增益，同时在二元任务上达到了相当的高性能。


表13：3B模型在检索增强生成（通过精确匹配分数评估）和阅读理解基准（准确率%）的长任务上的性能。


表14：在文本分类基准上的多示例上下文学习性能（准确率）。括号中的数字表示每个任务的标签数。

闭卷问答。我们还评估了模型在不访问任何文档的情况下的闭卷问答（closed-book QA）性能。我们使用Zhao等人【索引64，Analysing the impact of sequence composition on language model pre-training，ACL 2024】的评估协议来衡量精确匹配率。在表15中，我们注意到我们的方法在回答闭卷问题方面相比基线有显著改进。这与我们的模型在包含常识知识的标准基准上表现出改进的结果是一致的。


表15：1B模型（在CC上训练）在闭卷问答任务上的性能（精确匹配率%）。

A.7.4 SkyLadder消融实验

与BM25打包结合。由于SkyLadder仅通过掩码改变上下文长度而不改变底层数据，它与任何先进的数据打包方法（如Shi等人【索引44，In-context pretraining: Language modeling beyond document boundaries，ICLR 2024】，Ding等人【索引9，Fewer truncations improve language modeling，ICML 2024】）是正交的。在表16中，我们将SkyLadder与BM25打包方法结合。我们表明，该模型在短上下文和长上下文评估上都取得了比不使用调度的BM25更好的性能，这也优于Random基线。这揭示了我们的方法可以与更先进的打包技术结合，以进一步提升性能。


表16：不同调度类型的1B模型性能（%）。所有模型都在相同的100B CommonCrawl词元上训练，最终上下文长度为8K。BM25打包与SkyLadder结合后，显著提升了长任务的性能。

与混合注意力结合。我们注意到，近期一个有趣的趋势是在长上下文模型预训练中使用混合注意力结构。例如，Gemma3【索引49，Gemma 3 technical report，arXiv 2025】使用全局和局部注意力层的混合来平衡长上下文模型的效率和性能。我们好奇SkyLadder对此类架构的泛化性，并遵循Gemma3的策略，全局与局部比例为6:1。结果如表18所示。我们观察到SkyLadder在所有评估长度上始终优于基线，验证了其适用性。重要的是，SkyLadder沿时间维度工作，并且可以与不同的注意力变体结合，只要有上下文窗口可以调度。我们还在替代模型结构上验证了SkyLadder的有效性。在表19中，我们遵循Qwen2.5-0.5B结构预训练模型，并同样获得了持续的增益。


表18：Gemma3-like模型在不同评估上下文长度Le下的评估困惑度。所有模型都在CommonCrawl上用100B词元进行训练。


表19：Qwen2.5-0.5B模型在不同评估上下文长度Le下的评估困惑度。所有模型都在CommonCrawl上用100B词元进行训练。

长到短调度。SkyLadder在通常较短的标准基准上优于基线，一个可能的原因是应用掩码后训练数据混合中有更多短上下文数据。为了研究纯数据分布的影响，我们进行了一项消融实验，将原始的短到长调度反转，称之为长到短调度。该调度在变化阶段花费相同数量的词元（64B），然后在L=8K的恒定训练阶段再训练36B词元。在表17中，我们显示长到短调度对模型的短评估和长评估任务的性能都没有帮助。这突出表明上下文窗口需要被调度，而不仅仅是长短上下文的数据混合。


表17：不同调度类型的1B模型性能（%）。所有模型都在相同的100B FineWeb-Pro词元上训练，最终上下文长度为8K。短到长调度始终优于长到短调度。

替代调度类型。我们探索了遵循不同函数的各种短到长调度类型，如4.5节所述。表20显示了作为t的函数的调度细节，图12展示了不同调度类型的示意图。在表7中，我们显示遵循正弦调度的更平滑的增加在长上下文评估中效果最好，同时在标准基准上也取得了强大的性能。


表20：不同上下文窗口调度类型的函数。我们在实验中设置ws = 32和we = 32768。用于取整的r设置为1024。


图12：各种调度类型的示意图。

扩展速率。我们在图13中说明了扩展速率α的影响。由于评估是在8K上下文上进行的，速率较低（且上下文窗口较短）的模型会有较高的损失，因为评估长度是分布外的。然而，最终所有模型的损失在调度达到8K后都会收敛到一个较低的水平。预训练后验证损失的详细数字可以在表21中找到。遵循先前的工作【索引10，Data engineering for scaling language models to 128k context，ICML 2024】【索引17，Training compute-optimal large language models，NeurIPS 2022】【索引24，Scaling laws for neural language models，arXiv 2020】，我们将大于0.01的损失差异视为显著。我们得出结论，设置一个合理的速率1/8可以平衡短上下文和长上下文的损失，这是我们主要实验的默认设置。


图13：不同α影响的示意图。虚线代表每一步的当前上下文窗口w，实线是在8K长度下评估的损失。


表21：不同扩展速率下的验证损失。如果一个单元格与该列最佳值的差异大于0.01，则标记为红色。Le是评估上下文长度。所有模型大小为120M，并在100B个词元上训练。

循环调度。受循环学习率调度【索引45，Cyclical learning rates for training neural networks，WACV 2017】的启发，我们也想知道循环是否对调度有帮助。在图14中，我们展示了两种循环调度。在“跳跃”（Jump）调度中，$w(t)$在达到L后会立即降至$w_s$。另一方面，“渐进”（Gradual）调度意味着在$w_e$和$w_s$之间呈“M”形交替。值得注意的是，在不连续的跳跃调度中，当我们长时间只在短上下文上训练时，我们注意到长上下文困惑度有显著增加。然而，只要w增加回L，性能就会恢复。


图14：具有渐进增加或跳跃的循环调度示意图。虚线代表每一步的上下文长度，实线是在8K长度下评估的损失。c表示循环次数。

在表22中，我们显示这些调度对最终性能没有重大影响。这突出表明该方法没有在数据选择中引入额外的偏差：与现有方法如Pouransari等人【索引38，Dataset Decomposition: Faster llm training with variable sequence length curriculum，NeurIPS 2024】提出的先训练短数据再训练长数据的建议不同，我们不假设对数据有这样的课程。我们认为上下文窗口大小应独立于数据长度，以避免仅在某些数据领域上训练的偏差。


表22：循环调度下的验证损失。Le代表评估上下文长度。所有模型大小为120M，并在100B个词元上训练。

初始窗口长度。我们展示了不同初始窗口长度$w_s$（训练开始时的窗口长度）的影响。在表23中，我们显示最优的起始长度是8个词元。这个趋势在$\alpha = 1/4$和$\alpha = 1/8$两种情况下都是相同的。这表明起始长度应该足够小，与扩展速率无关。这也揭示了之前的研究，如Jin等人【索引21，Growlength: Accelerating LLMs pretraining by progressively growing training length，arXiv 2023】和Pouransari等人【索引38，Dataset Decomposition: Faster llm training with variable sequence length curriculum，NeurIPS 2024】中以256的初始长度开始可能是次优的。


表23：在α=1/4和α=1/8时，使用不同ws训练120M模型100B词元后的最终验证损失。Le代表评估的上下文长度。如果一个单元格的损失与该列最佳值的差异大于0.01，则标记为红色。ws = 8192等于无调度。

计算预算。我们表明，当总词元数量有限时，我们的方法仍然可以提高语言模型的性能。在表24中，我们选择12.5B、25B和50B总词元作为计算预算，并改变扩展速率，使得w在训练过程中的同一点达到L。我们观察到，在不同的词元预算下，性能趋势是最好的：逐渐扩大上下文窗口比快速增加能带来更好的性能。


表24：在不同训练词元预算和扩展速率α下，120M模型的最终验证损失。Le代表用于评估的上下文长度。“% of Token Budget”表示在w(t)增加的扩展阶段花费了多少词元。在所有词元预算下，我们观察到当我们在扩展阶段花费约64%的词元，在稳定阶段花费36%的词元时，性能有一致的改善。

滑动窗口扩展。SkyLadder的一个可能替代方案（默认使用局部因果掩码）是使用滑动窗口注意力，窗口大小为$w(t)$，随训练时间变化。形式上，掩码变为：

$$\begin{aligned} M_{i, j}=\begin{cases}0 & \text { if } i-w \leq j \leq i \\ -\infty & \text { otherwise. }\end{cases} \end{aligned}$$

这样上下文中的每个词元都有一个固定大小为w的前文。当$w(t)$达到L时，该掩码变得与因果掩码等效。我们在表25中比较了两者的性能，并观察到滑动窗口方法在长任务上表现稍好，在标准基准上表现稍差。这可能是因为总体上滑动窗口方法有更多的词元具有更长的前文。在两种情况下，SkyLadder都优于Random基线。我们认为未来的工作可以进一步研究使用因果注意力和滑动窗口注意力的SkyLadder实现之间的差异，例如注意力沉淀的形成【索引14，When attention sink emerges in language models: An empirical view，arXiv 2024】。也可能存在两者的组合：例如，首先使用局部掩码来禁用干扰，随着训练的进行再启用滑动窗口。


表25：不同掩码方案的1B模型性能（%）。所有模型都在相同的100B FineWeb-Pro词元上训练，最终上下文长度为8K。SkyLadder的两种实现都优于基线，滑动窗口方法在长任务上表现出色，但在标准基准上性能略有下降。

A.8 与相关工作的附加比较

与先前工作的关系。我们承认有几项先前的工作发现了类似的从短到长预训练的模式。例如，Li等人【索引28，The stability-efficiency dilemma: investigating sequence length warmup for training gpt models，NeurIPS 2022】发现，在预训练的初始步骤中使用序列长度预热可以提高模型稳定性。然而，他们主要关注训练损失的稳定性，并未在多个评估和更大规模上显示出明确的性能增益。此外，我们证明了调度模型上下文窗口的好处不仅仅局限于预热阶段。在表21的第一行中，仅用8B词元对模型进行预热，与较慢的扩展速率相比，性能是次优的。这验证了上下文窗口应被视为在整个训练过程中需要调度的因素，这也将我们与仅考虑预热阶段的Li等人【索引28，The stability-efficiency dilemma: investigating sequence length warmup for training gpt models，NeurIPS 2022】区分开来。

与另一相关工作的比较。另一个相关工作是Jin等人【索引21，Growlength: Accelerating LLMs pretraining by progressively growing training length，arXiv 2023】，作者使用渐进式序列长度来加速训练。然而，他们的方法在相同词元预算下导致性能下降，而我们的SkyLadder在相同词元数量下既节省了时间又提高了性能。我们怀疑这可能是因为他们使用了次优的调度。此外，他们的研究仅限于观察小模型（最大410M参数）的训练损失，而我们全面展示了在多个语料库、模型大小、上下文大小和各种任务上的性能增益。总的来说，我们系统地进行了关于上下文窗口调度在预训练中影响的受控实验，为解释这些先前的研究提供了见解。

A.9 计算信息

计算资源。我们所有针对≤1B尺寸模型的实验都在一个由NVIDIA A100节点组成的内部集群上进行，每个节点配备40G内存。3B模型的实验则在H100节点上进行。还有一些我们未在论文中包含的初步实验，它们占总计算量的一部分。每个实验的详细计算情况如下：对于上下文窗口的初步研究，用100B词元（8K上下文）预训练一个1B模型，在一个8个A100的节点上大约需要200小时。不同尺寸模型的计算时间相应缩放。例如，绘制图4(a)和(b)总共需要在一个单节点上进行159天的预训练。对于SkyLadder实验，使用各种语料库的基线预训练耗时相同，而SkyLadder根据上下文长度的不同，将训练速度提高了13%至22%。

A.10 数据集统计

数据集详情。在本节中，我们提供了研究中使用的数据集的详细统计信息。这些包括预训练语料库的文档长度分布、评估数据集的特征，以及标准推理基准的输入长度统计。

表26报告了两个预训练语料库CommonCrawl和FineWeb-Pro的文档长度统计。两个分布都呈现强烈的右偏，表明长文档很少见。与FineWeb-Pro相比，CommonCrawl通常包含更长的文档，而FineWeb-Pro则经过了更仔细的清洗和过滤。


表26：预训练语料库的文档长度统计，单位为每个文档的词元数。均值、中位数和标准差描述了集中趋势和变异。P25和P75表示第25和第75百分位数，而偏度和峰度捕捉了分布的不对称性和尾部厚重程度。

表27显示了常见推理和知识基准的输入长度特征，包括ARC、CSQA、HellaSwag、OBQA、PIQA、SocialIQA、Winogrande和MMLU。虽然这些基准由相对较短的上下文组成，但它们仍然是评估模型事实一致性和推理能力的标准。重要的是，一个长上下文模型即使用户提供短查询时也应保持稳定的行为。


表27：常见推理和知识基准的输入长度特征。尽管相对较短，这些任务对于衡量知识和推理一致性至关重要。

最后，表28总结了在阅读理解和长上下文评估中使用的评估数据集的特征。这些包括问答基准，如MDQA、RULER、SQuAD、HotpotQA、NQ、TriviaQA和RACE。这些数据集在输入长度上差异很大，反映了推理深度和上下文复杂性的多样性。


表28：阅读理解和问答评估数据集的长度统计。这些基准捕捉了不同层次的输入复杂性，从简短的事实性问答到多跳推理任务。

总结。总的来说，本工作中使用的​​数据集涵盖了广泛的输入长度和领域，从大规模预训练语料库到短上下文和长上下文评估基准，确保我们的分析既全面又有代表性。

A.11 资产许可

许可信息。本文主要使用了以下公共数据集或代码库：SlimPajama【索引46，SlimPajama: A 627B token cleaned and deduplicated version of RedPajama，HuggingFace 2023】，遵循CommonCrawl基金会的使用条款；FineWeb-Pro【索引65，Programming every example: Lifting pre-training data quality like experts at scale，arXiv 2024】，采用ODC-By 1.0许可证；以及TinyLlama【索引61，Tinyllama: An open-source small language model，arXiv 2024】，采用Apache 2.0许可证。