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

Authors: Tian Jin, Ahmed Imtiaz Humayun, Utku Evci, Suvinay Subramanian, Amir Yazdanbakhsh, Dan Alistarh, Gintare Karolina Dziugaite
Primary area: foundation or frontier models, including LLMs
Keywords: pruning, sparsity, large language model, pretraining

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.

Decision: accept (poster)
Reviewer scores: 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730695062542

thread_type:
official_review

<thread_object>
creator: reviewer_kf8y

<content>
summary:
The paper studies the scaling law for pretraining sparse language models. It argues that sparse pretraining reduces both pretraining and inference costs since the final model can be smaller and requires less computation during inference. The authors modify the Chinchilla scaling law for dense models to use the average number of active parameters (receiving gradients) in training. The modified scaling law unifies both sparse and dense models.  They also conduct extensive studies to search for optimal sparse pretraining configurations for LLMs. The key findings include pruning starting at 25% and ending at 75% obtains near optimal final eval loss, optimal learning rates, and batch sizes for sparse training match dense models under the same compute budget. The paper further provides both theoretical justifications and empirical validation for the scaling law and presents a simple recipe for sparsely pretrained models.

strengths:
- The paper studies an interesting research problem, the scaling law for pretraining sparse models. It develops a modified scaling law that unifies both dense and sparse models. 
- The analysis and empirical validation are solid and convincing. The findings are supported by both theoretical justifications and experiments over 80 sparse pretraining configurations. The optimal settings for pretraining sparse models provide valuable practical lessons.
- The paper itself is well-written, clearly structured, and easy to follow.

weaknesses:
- One major concern is that the scaling law for dense models seems to have less impact, especially given that generally more data and compute resources are better to obtain better and more capable LLMs, LLaMA 2 and LLaMA 3 papers provide such evidence. How would unifying the scaling law for both dense and sparse models be useful under such a context?
- The paper mainly focuses on LLaMA style model architecture for sparse pretraining, any discussion on other sparse architectures like the mixture of experts (MoE)? Otherwise, it may be a bit overclamied. 
- The other concern is that previous scaling laws at least study model parameters with 1B scale, but this paper only discusses models under 500M parameters, making the claim of scaling law for LLM weaker. 
- While the paper acknowledges the limitation of evaluation loss, it would be very useful to report the performance numbers of downstream tasks, especially given the size of models is small. No compute limitation prohibits such evaluation.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730701169235

thread_type:
official_review

<thread_object>
creator: reviewer_tyda

<content>
summary:
This paper investigates sparse pre-training as an alternative to post-training pruning for large language models (LLMs) and proposes a modified scaling law based on the average active parameter count. Through extensive experimentation with over 80 configurations, the authors claim that sparse pre-training can yield models with similar performance to dense pre-training, offering a more efficient approach by beginning with dense training and progressively sparsifying. However, due to limitations in hardware and software, the paper lacks direct evidence of computational savings, which is a primary motivation for sparse pre-training.

strengths:
+ Modifying the Chinchilla scaling law to consider average active parameters is a novel approach, bridging dense and sparse pre-training frameworks effectively
+ The study evaluates a wide array of configurations and sparsity schedules, providing a thorough analysis of optimal sparse pre-training practices

weaknesses:
- Despite the theoretical focus on efficiency, the paper lacks a demonstration of actual computational savings due to the current limitations in sparse matrix support in hardware/software, which could weaken the case for real-world applicability
- The study uses evaluation loss as the sole metric, without investigating the effects on real-world downstream tasks. This limits the ability to gauge the model’s practical effectiveness or generalization capabilities
- The framework’s reliance on finely tuned, phase-specific compute allocations (dense, pruning, recovery) introduces implementation complexity, which could be challenging to replicate or scale, particularly in resource-constrained environments

questions:
See weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730712985409

thread_type:
official_review

<thread_object>
creator: reviewer_pohq

<content>
summary:
This paper conducts an overwhelming analysis of the effect of sparse pretraining on large language models (LLMs). Specifically, it proposes a new scaling law that is modified from the Chinchilla scaling law using a novel concept of the average number of active parameters during training, i.e., the averaged number of parameters that receive gradients in LLM pretraining among all training steps. The experimental results show that the proposed new scaling law sufficiently fits the evaluation loss prediction of LLM pretraining, while leveraging sparse pretraining could match the sparse, pruned model's evaluation loss to a larger, dense LLM.

strengths:
- This paper extends the Chinchilla scaling law and adapts it to the scenario of sparse pretraining, where the modified scaling law uses the concept of averaged activate parameter number to sufficiently model the evaluation loss of LLMs.
- When design a pretraining pruning schedule using the scaling law proposed in this paper, one can pretrain an LLM with fewer parameters yet matching the evaluation performance of a dense model.
- The theoretical and empirical analysis in this paper demonstrates strong proof of the scaling law proposed in this paper.

weaknesses:
- The scale of models being used in this paper is limited, as the authors have addressed in the limitation section. Larger model experiments are definitely useful and could bring wider impact to this paper.
- The only evaluation metric being used in this paper is the pretrained model's perplexity (or pretraining evaluation loss) without downstream task evaluations. As models could forget knowledge in pretraining because of pruning, a finer-grained analysis with task-specific evaluations could be beneficial.

questions:
- What is the exact pruning method being used in the iterative pruning phase? A detailed description of the pruning criteria could be helpful for the audience to understand (e.g., structured/unstructured pruning and how is the pruning decision being made)
- $\bar{N}$ and $D$ are independent variables shown in equation (2), however, there seems no analysis in this paper showing that the effects of $\bar{N}$ and $D$ bring to the model's evaluation loss are independent.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733288675892

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Summary

comment:
As we reach the rebuttal deadline, we sincerely thank all reviewers for making such a productive and constructive discussion period possible!

We are happy to see positive recognition of our work. Reviewer poHq described our analysis as *"an overwhelming analysis of the effect of sparse pretraining on large language models (LLMs)."* Similarly, Reviewer Kf8Y highlighted that *"the analysis and empirical validation are solid and convincing. The findings are supported by both theoretical justifications and experiments over 80 sparse pretraining configurations."*

During the rebuttal period, we provided additional experimental validations to strengthen our claims:

- We extended our findings to models exceeding 1 billion parameters, demonstrating generalizability [[link](https://openreview.net/forum?id=ud8FtE1N4N&noteId=LfNH4MThhM)].
- We complemented our perplexity-based evaluations with downstream task results, covering benchmarks like PIQA, ARC, Lambada, and Winogrande [[link](https://openreview.net/forum?id=ud8FtE1N4N&noteId=LfNH4MThhM)].
- We conducted detailed analysis into how sparse scaling laws can guide practical LLM training, emphasizing their implications for optimizing both training and deployment configurations and resource consumption [[link](https://openreview.net/forum?id=ud8FtE1N4N&noteId=v7jkCLPM8v)].

Thank you again for your time and thoughtful input, which have contributed to refining and improving this work.

Authors of 9129.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a new scaling law that modifies the Chinchilla scaling law to use the average number of active parameters during training, rather than just the total parameter count. The authors claim that this modified scaling law can effectively model the evaluation loss for both sparsely and densely pre-trained LLMs.

Overall, the paper makes a solid contribution to the understanding of sparse pre-training for large language models. The novel scaling law, extensive empirical analysis, and practical insights are valuable. While the paper has some limitations in model scale and evaluation metrics, the reviewers generally find the work to be sound, well-presented, and a good contribution to the field. Therefore, I recommend accepting it.