PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Pretraining Length Scaling

Bohong Wu, Shen Yan, Sijun Zhang, Jianqiao Lu, Yutao Zeng, Ya Wang, Xun Zhou

arXiv:2504.14992v2cs.CL

TL;DR

Pre-training length scaling remains underexplored, and naive token repetition increases KV-cache size and inference costs. PHD distinguishes original from hidden decoding tokens, retaining only original-token KV for long-range dependencies; its SWA and CSWA variants add local attention and chunk-wise dependencies. The paper reports consistent improvements across benchmarks, with the PHD-CSWA-3-16-32 configuration averaging 2.0% higher accuracy and 0.034 lower training loss.

  • Problem

    Pre-training length scaling is underexplored, while naive token repetition raises KV-cache size, pre-filling time, and decoding latency.

    Source: §1 · p. 2

  • Method

    PHD repeats input tokens, retains original-token KV for long-range dependencies, and discards hidden-token KV after use; SWA and CSWA add windowed and chunk-wise attention variants.

    Source: §1 · p. 2

  • Results

    Across benchmarks, PHD-CSWA-3-16-32 introduces an average of 2.0% accuracy improvement and 0.034 decrease in training loss.

    Source: §3.2 · p. 6, §5 · p. 10

  • Takeaways & Limitations

    The experiments support length scaling as an effective and resource-efficient approach during inference.

    Source: §5 · p. 10

Abstract

from arXiv · show

Recent advances in large language models have demonstrated the effectiveness of length scaling during post-training, yet its potential in pre-training remains underexplored. We present the Parallel Hidden Decoding Transformer (\textit{PHD}-Transformer), a novel framework that enables efficient length scaling during pre-training while maintaining inference efficiency. \textit{PHD}-Transformer achieves this through an innovative KV cache management strategy that distinguishes between original tokens and hidden decoding tokens. By retaining only the KV cache of original tokens for long-range dependencies while immediately discarding hidden decoding tokens after use, our approach maintains the same KV cache size as the vanilla transformer while enabling effective length scaling. To further enhance performance, we introduce two optimized variants: \textit{PHD-SWA} employs sliding window attention to preserve local dependencies, while \textit{PHD-CSWA} implements chunk-wise sliding window attention to eliminate linear growth in pre-filling time. Extensive experiments demonstrate consistent improvements across multiple benchmarks.

1 Introduction

The introduction identifies pre-training length scaling as promising but underexplored, while naive token repetition increases KV-cache size and inference costs. PHD retains original-token KV for long-range dependencies and discards hidden-token KV after use; SWA and CSWA variants preserve local information while addressing memory and pre-filling costs.

1 Introduction

  • Earlier pre-training approaches used inserted plain text or latent representations, while this work explores length scaling through repeated input tokens without middle-layer post-processing.The introduction says prior approaches achieved reasoning-benchmark gains at the cost of increased KV-cache size and decoding latency.Source: §1 · pp. 1–2
  • Naive token repetition yields loss and performance scaling, but linearly increases KV-cache size, pre-filling time, and decoding latency.The introduction notes that pre-filling can grow superlinearly with repetition.Source: §1 · p. 2, Figure 1 · p. 2
  • PHD retains original-token KV for long-range dependencies and discards hidden-decoding-token KV after use, maintaining vanilla-Transformer cache size while enabling length scaling.The repeated input tokens are used in parallel, and only the final copy contributes to next-token prediction loss.Source: §1 · p. 2, Figure 2 · p. 3
  • PHD-SWA preserves local hidden-token dependencies with a sliding-window cache, adding only O(1) KV memory but retaining sequential dependencies that increase pre-filling time.The introduction presents this variant as a way to recover performance benefits from hidden-token KV.Source: §1 · p. 2
  • PHD-CSWA restricts sequential dependencies to chunks, so only the final chunk's pre-filling time increases linearly.This addresses the linear pre-filling growth associated with PHD-SWA.Source: §1 · p. 2
  • The PHD series is reported to improve accuracy across multiple benchmarks while maintaining acceptable computational overhead in pre-filling and decoding.The introduction frames efficient pre-training length scaling as the work's contribution.Source: §1 · pp. 2–3

2 Approach

The PHD-Transformer series supports length scaling through separate handling of original and hidden decoding tokens. PHD-SWA adds local attention, while PHD-CSWA confines it to chunks to reduce pre-filling overhead.

2.1 Notations

  • PHD repeats input tokens K times and names the repeated copies after the first as hidden decoding tokens.The section defines a training sample as t tokens with hidden representations of dimension d and an attention mask M_mn between tokens.Source: §2.1 · p. 3

2.2 PHD

  • PHD retains original-token KV cache for global dependencies and discards hidden-token cache after parallel decoding, keeping cache size unchanged from the vanilla Transformer.Pre-filling time also remains constant with scaling factor K; inference requires K times the FLOPs, processed in parallel.Source: §2.2 · pp. 3–4

2.3 Kernel Design

  • Rearranging original and hidden decoding tokens groups masked attention positions into a continuous block, reducing attention-computation complexity to O(K).This layout makes the sparse attention pattern more efficient to optimize.Source: §2.3 · p. 4, Figure 3 · p. 4

2.4 PHD-SWA and PHD-CSWA

  • PHD-SWA attends globally to original tokens and locally to the W most recent hidden decoding tokens, improving performance with O(1) additional KV cache memory.The hidden-token cache has sequential dependencies, causing K-times pre-filling overhead.Source: §2.4 · p. 4
  • PHD-CSWA confines sliding-window attention to independent chunks, limiting extra pre-filling overhead to K repetitions within the final chunk.This avoids the across-sequence overhead caused by sequential hidden-token dependencies while preserving local attention patterns.Source: §2.4 · pp. 4–5

3 Experiments

Experiments report consistent benchmark improvements for PHD-CSWA, evaluate window and chunk choices, and show scaling gains with limited decoding-latency overhead. PHD-CSWA also keeps pre-filling overhead small as sequence length increases.

3 Experiments

  • The experiments use OLMo2 as their codebase, and the main-results evaluation presents PHD-CSWA as the most practical and effective strategy.The authors describe its performance improvement as steady with acceptable overhead.Source: §3 · p. 5, §3.2 · p. 5

3.2 Main Results

  • PHD-CSWA-2-16-32 improves average accuracy by 1.5% and decreases training loss by 0.025; PHD-CSWA-3-16-32 improves average accuracy by 2.0% and decreases training loss by 0.034.The results cover the evaluated benchmarks for 1.2B-parameter dense models.Source: Table 1 · p. 6, §3.2 · p. 6

3.3.1 Chunk-wise Sliding Window Attention

  • Increasing the window from W = 1 to W = 4 significantly reduces training and validation loss, while W = 16 yields only marginal further improvement.The authors conclude that a small local KV cache for hidden decoding tokens can approach optimal performance while remaining computationally efficient and hardware-friendly.Source: §3.3.1 · pp. 6–7, Figure 6 · p. 7
  • Adding chunks produces negligible differences in training and validation loss, and larger chunk sizes yield progressively better results.The authors select C = 32 as their balance between computational efficiency and model performance for subsequent experiments.Source: Figure 7 · p. 7, §3.3.1 · p. 7

3.4 Decoding Token Scaling

  • PHD-SWA improves performance as the scaling factor rises; at K = 5, loss decreases by nearly 0.06 and average accuracy improves by 1.8% across benchmarks.The reported scaling experiments use a fixed window size.Source: §3.4 · p. 7, Figure 8 · p. 7, Table 2 · p. 8

3.5 Pre-filling Speed and Decoding Speed

  • At increasing sequence lengths, PHD-CSWA incurs only a small pre-filling increase over vanilla, while PHD-SWA incurs superlinear pre-filling overhead.Increasing K to 256 adds no more than 20% decoding latency for either variant.Source: Figure 9 · p. 8, §3.5 · p. 8

4 Related Works

The related work situates PHD among sparse-attention, KV-cache, and latent-thinking approaches. It distinguishes PHD by applying tailored attention patterns to hidden decoding tokens while retaining full attention for original tokens.

4.1 Sparse Attention Mechanisms

  • PHD-SWA applies sliding-window patterns only to hidden decoding tokens, retaining full attention for original tokens; data-dependent sparse methods instead estimate patterns at runtime.The paper describes such estimation as substantial overhead that can limit effectiveness for long-context scenarios.Source: §4.1 · p. 9
  • PHD integrates efficient attention patterns during training, targeting hidden decoding tokens rather than general block sparsity as Native Sparse Attention does.The paper states that this design maintains the original transformer’s KV-cache footprint while improving performance.Source: §4.1 · p. 9

4.2 KV Cache Optimization

  • Unlike post-hoc KV-cache compression, PHD shares cache across repeated tokens to maintain the original transformer’s cache size while increasing computational depth.PHD is also presented as complementary to KV-cache allocation and decoding-efficiency methods, and orthogonal to attention-kernel optimizations.Source: §4.2 · p. 9

4.3 Latent Thinking Transformers

  • Earlier hidden-token approaches use inserted text or continuous signals, but the paper identifies increased KV-cache size as limiting their applicability during inference.PHD also uses continuous signals, with KV-cache management intended to keep the increase acceptable.Source: §4.3 · p. 10
  • Recurrent latent-thinking methods scale decoding computation by reusing transformer blocks or layers, but the paper says their lack of parallel generation limits inference efficiency.PHD’s approach instead focuses on computational depth without proportional increases in KV-cache size.Source: §4.3 · p. 10

5 Conclusion

The paper presents PHD as a pre-training length-scaling framework that repeats input tokens while retaining only original-token KV cache. Its variants preserve local dependencies or reduce pre-filling growth, and experiments report improvements across benchmarks with resource-efficient inference.

5 Conclusion

  • PHD-Transformer repeats input tokens while retaining only original-token KV cache, achieving performance gains without increasing KV-cache size.The conclusion frames pre-training length scaling as an efficient paradigm for enhancing transformer models.Source: §5 · p. 10
  • PHD-SWA uses sliding-window attention to preserve local dependencies, while PHD-CSWA uses chunk-wise sliding-window attention to eliminate linear pre-filling latency growth.Source: §5 · p. 10
  • Experiments demonstrate consistent improvements across multiple benchmarks, supporting effective length scaling with resource-efficient inference.Source: §5 · p. 10
Loading 2504.14992v2…