Escalado Eficiente de la Duración del Pretrenamiento

Efficient Pretraining Length Scaling

April 21, 2025

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

cs.AI

Resumen

Los avances recientes en los modelos de lenguaje de gran escala han demostrado la efectividad del escalado de longitud durante el post-entrenamiento, aunque su potencial en el pre-entrenamiento sigue siendo poco explorado. Presentamos el Transformer de Decodificación Oculta Paralela (PHD-Transformer), un marco novedoso que permite un escalado de longitud eficiente durante el pre-entrenamiento mientras mantiene la eficiencia en la inferencia. El PHD-Transformer logra esto mediante una estrategia innovadora de gestión de la caché KV que distingue entre tokens originales y tokens de decodificación oculta. Al retener únicamente la caché KV de los tokens originales para dependencias de largo alcance y descartar inmediatamente los tokens de decodificación oculta después de su uso, nuestro enfoque mantiene el mismo tamaño de caché KV que el transformer convencional, permitiendo un escalado de longitud efectivo. Para mejorar aún más el rendimiento, introducimos dos variantes optimizadas: PHD-SWA emplea atención de ventana deslizante para preservar dependencias locales, mientras que PHD-CSWA implementa atención de ventana deslizante por fragmentos para eliminar el crecimiento lineal en el tiempo de pre-llenado. Experimentos exhaustivos demuestran mejoras consistentes en múltiples benchmarks.

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.