DPLS: Dynamic Partial Label Smoothing Loss for Enhancing Pre-training Stability of Large Language Models

DPLS is a novel loss function that addresses the limitations of the fixed strategy in traditional label smoothing during large language model pre-training by dynamically adjusting the label smoothing strategy, achieving more stable convergence and better generalization performance. This method is built based on the nanoGPT framework and FineWeb-Edu-100B dataset, and has advantages such as plug-and-play, low computational overhead, and strong interpretability, providing a new regularization tool for large language model pre-training.

In large language model pre-training, traditional cross-entropy loss using hard labels (one-hot encoding) tends to cause overfitting, excessively high confidence, and limited generalization ability. As a classic regularization technique, label smoothing converts hard labels to soft labels to alleviate these issues, but the traditional strategy is fixed and cannot adapt to dynamic changes during training.

The key innovations of DPLS lie in dynamic adjustment and partial smoothing features: dynamic adjustment adaptively changes the smoothing intensity based on training steps, model confidence, etc. (high in the early stage, low in the later stage); partial smoothing reduces smoothing for predictions that the model is confident about, while maintaining high smoothing for uncertain ones. The implementation is based on the nanoGPT and TPA codebases, supporting the FineWeb-Edu-100B dataset. Key parameters include dpls_epsilon (smoothing baseline), dpls_top_k (top-k tokens involved in smoothing), and resume_dir (resume from checkpoint and enable DPLS).

The pre-training phase supports torchrun distributed multi-GPU training, provides baseline vs. DPLS comparison scripts, and integrates SwanLab for visualizing training logs; the evaluation phase calculates perplexity (PPL) and entropy, uses EleutherAI's lm-evaluation-harness for downstream task evaluation, and supports multiple standard benchmark tests.

The tech stack includes Python 3.12+, PyTorch 2.8.0, and supports multi-GPU distributed training; hardware recommendations are 8 Pro6000 GPUs for large-scale pre-training, while fewer GPUs can be used for small-scale experiments.

DPLS addresses the core problem of balancing fitting and generalization in large language model pre-training. Its advantages include plug-and-play (can be used as a drop-in replacement for existing processes), low computational overhead, and strong interpretability. For pre-training researchers and engineers, it is a tool worth trying, which may improve stability and final performance.