The Journey Matters: Average Parameter Count over Pre-training Unifies Sparse and Dense Scaling Laws

TL;DR

The paper tackles the high cost of scaling large language models and investigates sparse pre-training as a compute-efficient alternative. It introduces a unified scaling law using the average parameter count during training, $L(\bar{N},D) = \frac{A}{\bar{N}^\alpha} + \frac{B}{D^\beta} + E$, and validates it theoretically and empirically for both sparse and dense pre-training. Through an 80-schedule sweep and analysis of pruning dynamics, it identifies that starting pruning at 25% and ending at 75% of total compute yields near-optimal final loss, and that hyperparameters from dense pre-training transfer effectively. Practically, sparse pre-training matches final evaluation loss at equivalent compute with substantial reductions in final model size, enabling inference savings and broader compute efficiency.

Abstract

Pruning eliminates unnecessary parameters in neural networks; it offers a promising solution to the growing computational demands of large language models (LLMs). While many focus on post-training pruning, sparse pre-training--which combines pruning and pre-training into a single phase--provides a simpler alternative. In this work, we present the first systematic exploration of optimal sparse pre-training configurations for LLMs through an examination of 80 unique pruning schedules across different sparsity levels and training durations. We find that initiating pruning at 25% of total training compute and concluding at 75% achieves near-optimal final evaluation loss. These findings provide valuable insights for efficient and effective sparse pre-training of LLMs. Furthermore, we propose a new scaling law that modifies the Chinchilla scaling law to use the average parameter count over pre-training. Through empirical and theoretical validation, we demonstrate that this modified scaling law accurately models evaluation loss for both sparsely and densely pre-trained LLMs, unifying scaling laws across pre-training paradigms. Our findings indicate that while sparse pre-training achieves the same final model quality as dense pre-training for equivalent compute budgets, it provides substantial benefits through reduced model size, enabling significant potential computational savings during inference.

Figures (6)

• Figure 1: We show the predictive power of average active parameters by creating two families of models. The first is sparse, starting from a dense model with 138 million prunable parameters in the linear layers and targeting final sparsity levels of 20%, 40%, 60%, and 80%. The second is dense, created by adjusting the hidden dimension to match the average number of active parameters throughout sparse-pre-training for each sparse models. In the left plot, we represent sparse models with dashed lines and dense models with solid lines. Each sparse-dense pair, with matching average active parameters, is shown in the same subfigure. Each pairs of model shares the same total training compute. In the right plot, despite differences in pre-training techniques, sparse and dense models with matching average active parameters (indicated by matching colors) achieve similar final loss.
• Figure 2: Left: Loss vs effective compute for 410M models for various sparsity levels and training durations. Right: Estimated $\alpha$ coefficient in the scaling law for 410M model.
• Figure 3: Predicted eval loss from our fitted scaling law versus the actual achieved final loss.
• Figure 4: Optimal sparsity schedule sweep for 162M-10$\times$ (left) and 162M-20$\times$ (right) models. Each tuple on the x-axis, $(t_d,t_s)$, represents the percentage of training time spent for dense traning ($t_d$), and percentage of time spent gradually pruning ($t_s$).
• Figure 5: Batch size and learning rate sweep for 162M models.