Escalonamento Eficiente do Comprimento de Pré-treinamento

Efficient Pretraining Length Scaling

April 21, 2025

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

cs.AI

Resumo

Avanços recentes em modelos de linguagem de grande escala demonstraram a eficácia do escalonamento de comprimento durante o pós-treinamento, mas seu potencial no pré-treinamento permanece pouco explorado. Apresentamos o Parallel Hidden Decoding Transformer (PHD-Transformer), uma estrutura inovadora que permite o escalonamento eficiente de comprimento durante o pré-treinamento, mantendo a eficiência na inferência. O PHD-Transformer alcança isso por meio de uma estratégia inovadora de gerenciamento de cache KV que distingue entre tokens originais e tokens de decodificação oculta. Ao reter apenas o cache KV dos tokens originais para dependências de longo alcance e descartar imediatamente os tokens de decodificação oculta após o uso, nossa abordagem mantém o mesmo tamanho de cache KV que o transformer convencional, permitindo um escalonamento eficaz de comprimento. Para aprimorar ainda mais o desempenho, introduzimos duas variantes otimizadas: o PHD-SWA emprega atenção de janela deslizante para preservar dependências locais, enquanto o PHD-CSWA implementa atenção de janela deslizante em blocos para eliminar o crescimento linear no tempo de pré-preenchimento. Experimentos extensivos demonstram melhorias consistentes em vários 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.