DPaI: Differentiable Pruning at Initialization with Node-Path Balance Principle

Authors: Lichuan Xiang, Quan Nguyen-Tri, Lan-Cuong Nguyen, Hoang Pham, Khoat Than, Long Tran-Thanh, Hongkai Wen
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: Prunning at Initialization, Sparsity, Neural Architecture Search

Abstract:
Pruning at Initialization (PaI) is a technique in neural network optimization characterized by the proactive elimination of weights before the network's training on designated tasks. This innovative strategy potentially reduces the costs for training and inference, significantly advancing computational efficiency. A key factor leading to PaI's effectiveness is that it considers the saliency of weights in an untrained network, and prioritizes the trainability and optimization potential of the pruned subnetworks. Recent methods can effectively prevent the formation of hard-to-optimize networks, e.g. through iterative adjustments at each network layer. However, this way often results in large-scale discrete optimization problems, which could make PaI further challenging. This paper introduces a novel method, called DPaI, that involves a differentiable optimization of the pruning mask. DPaI adopts a dynamic and adaptable pruning process, allowing easier optimization processes and better solutions. More importantly, our differentiable formulation enables readily use of the existing rich body of efficient gradient-based methods for PaI. Our empirical results demonstrate that DPaI significantly outperforms current state-of-the-art PaI methods on various architectures, such as Convolutional Neural Networks and Vision-Transformers. Code is available at https://github.com/QuanNguyen-Tri/DPaI.git

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

Reviews and discussion:
<thread>
creation date:
1730178844760

thread_type:
official_review

<thread_object>
creator: reviewer_v68w

<content>
summary:
Update:
Thanks for the authors' feedback. My concerns have been fully addressed.

----- original review -----
The authors propose a novel differentiable algorithm for deep network pruning at initialization. The key idea of the algorithm is to optimize the number of effective paths and nodes while pruning the network parameters. The non-differentiable issue can be addressed by Straight-Through Estimator. The authors demonstrate the superiority of the proposed method on several popular benchmark datasets.

strengths:
The proposed algorithm presents a novel approach to Node-Path Balancing in network pruning. It distinguishes itself by incorporating this methodology into a differentiable framework. Additionally, the article provides extensive numerical results for various vision tasks and includes detailed ablation studies, offering readers valuable insights into the performance of the proposed algorithm under different conditions. Furthermore, the authors provide some convergence analysis, which can be useful for those interested in understanding the underlying principles of the method.

weaknesses:
* The experiments are mostly for vision models. It would be nice to see some LLM experiments. For exmaple, how does the method perform on Llama-3.2-90B?
*  It is unclear why choosing logarithm in maximizing NPB principle. For example, why not using square root or other functions? Please suggest the theoretical support of this choice, and conduct numerical ablation studies to verify the optimality of this design. A possible approach is to make this loss function learnable, then find an optimized loss function on a set of small datasets, finally generalize to larger datasets.

* Using STE to convert non-differentiable optimization to differentiable one has been proposed in many previous works. Please justify the novelty of this approach when appling STE to this work. Any special adaptation to make STE more efficient / robust in this problem?

questions:
* Why choosing logarithm in NPB principle? Is logarithm loss an optimal design in this problem?
* Is the method applicable to LLM? How it performs on popular LLMs such as Llama-3.2 families?
* Why using STE while there are several possible methods to relax the differentiable issue? For example the popular [0,1] + L1 norm relaxation seems a good choice too.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730257336270

thread_type:
official_review

<thread_object>
creator: reviewer_gcqx

<content>
summary:
(+) The authors propose a novel pruning method referred to as Differentiable Pruning at Initialization (DPaI), which involves differentiable optimization of the pruning mask. DPaI adopts a dynamic and adaptable pruning process, allowing easier optimization processes and better solutions. 

(+) DPaI’s differentiable formulation enables easy use of the gradient-based methods for Pruning at Initialization (PaI). 

(+) The empirical results demonstrate that DPaI significantly outperforms current state-of-the-art PaI methods on various architectures, including Convolutional Neural  Networks (CNN) and Vision-Transformers (ViT).

strengths:
(+) The author analyzes the interpretations of the proposed method (DPaI) and proves its behavior using convergence theory.

Please see the strength of the paper in summarization further.

weaknesses:
(-) It isn't easy to follow how the proposed method (DPaI) finds the best nodes and paths using only an importance score. 

Through extensive ablation studies on the update rule (line 8) stated in Algorithm 1, the authors should show the effectiveness: the importance of nodes and paths seems to depend only on score s_{i,j}. The empirical observations (small neural networks) with the ablation studies help us understand the contributions of the proposed method, DPaI, and the convergence theory. 

(-) What about pre-trained model-based DPaI performances of ViT?

questions:
Please see the above weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730619402826

thread_type:
official_review

<thread_object>
creator: reviewer_cgaa

<content>
summary:
This paper extends Node-path Balancing Principle (NBP) into differentiable formulation by the learnable masks. To solve the gradient flow problem of top-k operation, authors introduce the Straight-Through Estimator (STE) technique during mask updating, while using Erdo ̋s-Rényi Kernel (ERK) technique to determine the layer-wise sparsity ratios.

strengths:
1. The differentiable NPB formulation achieves competitive performance compared to previous Pruning at Initialization (PaI) methods.
2. The proposed method can be easily integrated into the previous gradient-based methods.

weaknesses:
1. This paper employs methods that are well-established in prior works, such as the top-k operation with Straight Through Estimator (STE) [1, 2, 3], and the Erdős-Rényi Kernel (ERK) technique[4]. Although the differentiable masks have not been previously applied in the PaI domain, many methods[1, 2, 3, 5, 6] in pruning during/post training methods have used the differentiable masks for model compression. I suggest the authors provide a detailed discussion regarding the difference.

2. This paper uses the tanh activation function for the mask in the top-k with STE method, particularly when sigmoid and its variants are more commonly employed in previous works[1, 2, 3]. Therefore, what is the rationality of the tanh function. 

3. The models and datasets employed in the paper are predominantly small-scale validations. How does the method proposed in this paper perform on larger-scale models and datasets, such as experiments using Resnet-50 on ImageNet-1k?

4. In Fig.3, why does the proposed pruning method in the paper take more time in the Resnet20 on CIFAR-10 experiments compared to other methods, while it takes less time in the Resnet18 on Tiny-ImageNet and VGG19 on CIFAR-100 experiments? 

5. I suggest including a brief introduction or citation for the abbreviations of previous methods used in this paper, as relying solely on abbreviations from related work can reduce the readability of the paper for readers who are not fully familiar with PaI-related research.

[1] Disentangled Differentiable Network Pruning.
[2] Growing Efficient Deep Networks by Structured Continuous Sparsification.
[3] Learning sparse neural networks through l_0 regularization.
[4] The unreasonable effectiveness of random pruning: Return of the most naive baseline for sparse training.
[5] Filter Pruning for Efficient CNNs via Knowledge-driven Differential Filter Sampler.
[6] Convolutional Networks with Adaptive Inference Graphs.

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

</thread>
<thread>
creation date:
1730693415405

thread_type:
official_review

<thread_object>
creator: reviewer_qu2t

<content>
summary:
This paper introduces DPaI (Differentiable Pruning at Initialization), a novel method for pruning neural networks at initialization using the Node-Path Balancing (NPB) principle. Unlike traditional methods, DPaI applies a differentiable approach to optimize pruning masks, enabling efficient gradient-based optimization of sparse networks without needing training data. By balancing the number of effective nodes and paths, DPaI maintains information flow in sparse networks, achieving superior accuracy at high sparsity levels. Experimental results show that DPaI outperforms state-of-the-art methods across various network architectures and datasets, including CIFAR-10, CIFAR-100, Tiny-ImageNet, and ImageNet-1K. DPaI also achieves lower pruning time and reduced FLOPs, demonstrating its efficiency and practical applicability.

strengths:
1. Differentiable Optimization: DPaI introduces a differentiable approach, making pruning more efficient and compatible with standard neural network training, unlike discrete, layer-wise pruning methods.
2. High Performance at Extreme Sparsity: By balancing effective nodes and paths, DPaI maintains high accuracy even at sparsity levels of 96-99%, outperforming traditional methods.
3. Architecture-Agnostic: DPaI works effectively across different architectures, such as ResNet, VGG, and EfficientNet, making it broadly applicable.
4. Reduced Pruning Time and FLOPs: DPaI achieves low pruning times and significantly reduces FLOPs, enhancing computational efficiency compared to other methods.

weaknesses:
1. Hyperparameter Sensitivity: I assume DPaI’s performance will depend on hyperparameters which may require tuning for different datasets and sparsity levels, as this usually happens to pruning at initialization methods.
2. Implementation Complexity: DPaI involves complex gradient calculations and differentiable mask updates, which may be harder to apply to different task or models.

questions:
I currently hold a positive view of this work, but I’m somewhat lacking the sense of recent progress about pruning at initialization. I believe the most critical aspect is comparing it with other works in the same track, so I may further adjust my review after considering feedback from other reviewers.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732204739247

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Transformer Results

comment:
In response to the reviewer's request for experiments on ViT, we provide the following results:

## ViT-B/16 on ImageNet-1K: 90% Sparsity

| **Method** | **Eff Nodes** | **Log Eff Paths** | **Test Accuracy (%)** |
|------------|---------------|-------------------|-----------------------|
| Synflow    | 47,848        | 149.80609         | 64.33                 |
| DPal       | 40,325        | 172.09947         | 66.27                 |

---

## Experiments with Pre-trained ViT-B/16 Models: 90% Sparsity


| **Method** | **Dataset**       | **ViT** | **Pre-trained ViT** |
|------------|-------------------|---------|---------------------|
| Synflow    | CIFAR-10          | 77.98   | 84.09               |
|            | CIFAR-100         | 49.14   | 55.10               |
|            | Tiny-ImageNet     | 34.94   | 42.28               |
| DPal       | CIFAR-10          | 78.53   | 85.11               |
|            | CIFAR-100         | 51.33   | 58.23               |
|            | Tiny-ImageNet     | 36.23   | 43.31               |

---

## Beyond Vision Models: Experiments with  Transformer (41 million parameters) on IWSLT14 De-En (160K training samples)

Conducting experiments on LLMs is highly resource-intensive, and their evaluation often involves zero-shot in-context learning, which does not align with the goal of PaI—reducing the cost of training models. Therefore, we conducted experiments by training smaller language models from scratch to demonstrate the performance of PaI beyond vision models. Specifically, we trained an encoder-decoder transformer architecture (with 41 million parameters) on a machine translation task using the IWSLT14 De-En dataset (160K training samples). The results, evaluated using BLEU scores with beam search settings of 1 and 5 beams, are presented in the following table:

| **Method**  | **Sparsity** | **BLEU (beam=1)** | **BLEU (beam=5)** |
|-------------|--------------|-------------------|-------------------|
| **Dense**   | 0%           | 34.49            | 35.24            |
| **SynFlow** | 90%          | 29.54            | 30.64            |
| **DPaI**    | 90%          | 30.94            | 31.76            |

---
DPaI outperforms the baseline SynFlow by 1.4 BLEU at beam=1 and 1.12 BLEU at beam=5 while maintaining performance levels close to that of the dense model. This demonstrates the versatility and effectiveness of our proposed method across different model architectures and tasks.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732544813594

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Expressing Gratitude to the Reviewer and Seeking Further Feedback

comment:
Dear Reviewers,

We would like to express our sincere gratitude for your valuable feedback and constructive comments, which have significantly helped us improve our manuscript. We are pleased that you acknowledged the core novelties and strengths of our work, as listed follows:

>**Novel Differentiable Pruning Method (DPaI)**: Introducing a differentiable approach to optimize the pruning mask, making pruning more efficient and compatible with standard neural network training processes.**(Reviewer QU2t, CGAa, gcQx, v68w)**

>**High Performance at Extreme Sparsity**: Maintaining high accuracy even at extreme sparsity levels (96–99%) by balancing effective nodes and paths, outperforming traditional pruning methods.**(Reviewer QU2t, CGAa, gcQx,v68w)**

>**Architecture-Agnostic Applicability**: Effectively applying our method across various neural network architectures—including ResNet, VGG, EfficientNet, CNNs, and Vision Transformers (ViT) Language Models(LMs)—demonstrates broad utility. **(Reviewer QU2t)**

> **Reduced Pruning Time and Computational Load**: Achieving consistently lower pruning times and significantly reducing FLOPs at different initial model sizes, enhancing computational efficiency compared to other pruning techniques.**(Reviewer QU2t)**

> **Extensive Empirical Validation**: Providing comprehensive numerical results and detailed ablation studies on various vision tasks, showing significant performance improvements over state-of-the-art PaI methods.**(Reviewer QU2t, CGAa, gcQx,v68w)**

>**Theoretical Insights and Convergence Analysis**: Offering convergence theory analysis to interpret the method's behaviour, aiding in understanding the underlying principles and ensuring reliable performance.**(Reviewer gcQx,v68w)**

We also thank the reviewers for raising insightful questions and perspectives that have helped us further improve our paper. In response, we have provided additional experiments and clarifications:

>**ViT-B/16 on ImageNet-1K (90% Sparsity)**: We conducted experiments with ViT-B/16 on the ImageNet-1K dataset at 90% sparsity, demonstrating the effectiveness of our method on large-scale models and datasets.

>**Pre-trained ViT-B/16 Models (90% Sparsity)**: We evaluated our method on pre-trained ViT-B/16 models, showcasing its applicability and performance benefits when applied to pre-trained networks.

>**Language Models**: To illustrate the versatility of our approach, we conducted experiments with a Transformer model (41 million parameters) on the IWSLT14 De-En dataset (160K training samples), demonstrating the effectiveness of our method beyond vision tasks.

>To better understand the proposed method's convergence, we conducted empirical observations through ablation studies, examining how the number of effective nodes, paths, and kernels changes during the optimization of each objective. 

In response to the reviewers' requests for clarity, we have also made the following clarifications and will include those texts in our paper to improve the clarity:

>**Clarified Implementation Complexity**: We discussed why our differentiable objective brings simplicity compared with previous PaI methods that require iterative mask schema and large-scale discrete optimization 

>**Rationale Behind the Design of Our Optimization Objective**: We elaborated on the reasons for choosing logarithmic scaling, the Tanh function, and the use of the Straight-Through Estimator (STE) in our optimization objective, providing theoretical justification and empirical evidence.

>**Pruning Efficiency with Increasing Model Size**: We included additional analysis and discussion on how our method scales efficiently with model size, demonstrating its practicality for larger models.

>**Compared with Previous Methods**: We provided a more detailed discussion comparing our method with previous pruning methods during or after training, as well as other PaI works, highlighting the simplicity and adaptability of our approach.

Once again, we sincerely thank the reviewers for their time and effort in reviewing our paper. Your constructive feedback has been invaluable in helping us improve the quality and clarity of our work. We would like to know if you have any further concerns about our response. We would be happy to engage in any follow-up discussion and address any additional comments.

Sincerely,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper presents Differentiable Pruning at Initialization (DPaI), a method leveraging the Node-Path Balancing principle to optimize pruning masks, achieving improvements over state-of-the-art methods across various architectures. While reviewers raised concerns about hyperparameter sensitivity, scalability to larger models, and logarithmic scaling, the authors addressed these issues through experiments, ablation studies, and clarifications. The Area Chair recommends acceptance, with suggestions to further explore DPaI's scalability to larger models and alternative functions/objectives.