효율적인 사전 학습 길이 스케일링

Efficient Pretraining Length Scaling

April 21, 2025

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

cs.AI

초록

최근 대규모 언어 모델의 발전은 사후 학습(post-training) 과정에서 길이 스케일링(length scaling)의 효과를 입증했지만, 사전 학습(pre-training)에서의 잠재력은 아직 충분히 탐구되지 않았습니다. 본 논문에서는 사전 학습 중 효율적인 길이 스케일링을 가능하게 하면서도 추론 효율성을 유지하는 새로운 프레임워크인 PHD-Transformer(Parallel Hidden Decoding Transformer)를 제안합니다. PHD-Transformer는 원본 토큰과 숨겨진 디코딩 토큰을 구분하는 혁신적인 KV 캐시 관리 전략을 통해 이를 달성합니다. 장거리 의존성을 위해 원본 토큰의 KV 캐시만 유지하고, 숨겨진 디코딩 토큰은 사용 후 즉시 폐기함으로써, 기존 트랜스포머와 동일한 KV 캐시 크기를 유지하면서도 효과적인 길이 스케일링을 가능하게 합니다. 성능을 더욱 향상시키기 위해 두 가지 최적화된 변형을 도입했습니다: PHD-SWA는 슬라이딩 윈도우 어텐션(sliding window attention)을 사용하여 지역적 의존성을 보존하고, PHD-CSWA는 청크 단위 슬라이딩 윈도우 어텐션(chunk-wise sliding window attention)을 구현하여 사전 채우기 시간의 선형 증가를 제거합니다. 다양한 벤치마크에서의 광범위한 실험을 통해 일관된 성능 향상을 입증했습니다.

English

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 (PHD-Transformer), a novel framework that enables efficient length scaling during pre-training while maintaining inference efficiency. 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: PHD-SWA employs sliding window attention to preserve local dependencies, while PHD-CSWA implements chunk-wise sliding window attention to eliminate linear growth in pre-filling time. Extensive experiments demonstrate consistent improvements across multiple benchmarks.