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

Is the method applicable to LLM? How it performs on popular LLMs such as Llama-3.2 families? … For example, how does the method perform on Llama-3.2-90B?..

Thank you for raising this important question about the applicability of our DPaI method to Large Language Models (LLMs) such as the Llama-3.2 family, including models like Llama-3.2-90B. Our DPaI method is designed as a Prune-at-Initialization (PaI) technique, focusing on identifying the best subnetwork or substructure trained from scratch. The primary goal is to save computational resources for training while achieving performance comparable to the full model. This approach contrasts with post-training pruning methods, which aim to reduce the size of a pre-trained model by removing parameters while preserving the learned information. At present, we have not applied our DPaI method directly to LLMs like Llama-3.2-90B. Our experiments have focused on models and datasets where full training is computationally feasible, allowing us to evaluate and validate our approach thoroughly. Extending our method to LLMs would require significant computational resources and careful adaptation to handle the specific characteristics of those models. From the technique report released by Llama-3.2 author, training Llama-3.2-90B requires computational resources as follows:

Stage 1 pretraining: 885K H100 hours; Stage 2 annealing: 885K H100 hours; SFT: 3072 H100 hours; RLHF: 2048 H100 hours

We recognise the importance of evaluating our method on LLMs. Training LLMs is often highly costly, and PaI can potentially significantly reduce the cost of training those models (as PaI reduces model size before training). As discussed in the paper, our approach offers great efficiency concerning model size compared to other PaI methods. The differentiable mechanism we employ optimises sparse masks for the entire network simultaneously, which becomes particularly advantageous as model size increases. This suggests that our method could provide more significant benefits over existing state-of-the-art methods when applied PaI to larger models like LLMs. Our preliminary results on pruning the different sizes of the model are shown in Figure. 3 in the paper indicates that our method scales efficiently with model size, making it a promising candidate for larger models in the future.

…It is unclear why choosing logarithm in maximising NPB principle.. Why choose logarithm in NPB principle? Is logarithm loss an optimal design in this problem?

The main purpose of using the logarithmic scale for the number of effective nodes and effective paths in our differentiable objective is to address the imbalance between these objectives, as discussed in Section 3.2, lines 177–183. The number of effective paths is typically far greater than the number of effective nodes or connections. For example, in ResNet18 under 68.38% sparsity, the number of effective paths can reach up to $10^{65}$ (see Figure 2, ablation study on hyperparameters), while the number of nodes is only around 5000. Mathematically, the logarithmic function is an optimal choice for handling such large numbers, as $\log x$ grows the slowest as $x$ approaches infinity.

We kindly point out that the theoretical justification for employing the logarithmic scale is embedded in the computation of the derivative for each objective under the logarithmic scale. The derivative with respect to the path objective in equation (2) is proportional to the number of effective paths that may contain $s_{i,j}^{(l)}$, which itself is a very large number similar to the total number of effective paths. However, with the logarithm scale, the derivative is now divided by the total number of effective paths, making it relatively smaller.

We thank the reviewer for acknowledging the strengths of our work in methods and performance and providing more feedback for our paper's clarity. We hope the answer below solves all the clarity issues. If you have any further questions, please sincerely welcome the reviewer involved in further discussion,

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

Thank you for your questions and suggestions for further clarification regarding the proposed DPaI method. DPaI updates the importance scores iteratively in a differentiable manner, employing STE to estimate the gradients of the importance scores in relation to the discrete optimization problem. The mechanism by which DPaI updates the importance scores to maximise the number of effective nodes and paths can be observed in the derivatives associated with each objective.

• Path Objective (Equation 2): The derivative with respect to the importance score $s_{i,j}^{(l)}$ of a connection shows that its score increases based on the number of incoming paths $P(v_i^{(l-1)})$ and outgoing paths $\frac{\delta \mathcal{R}_P}{\delta P(v_j^{(l)})}$ it can connect. Specifically, $P(v_i^{(l-1)}) \cdot \frac{\delta \mathcal{R}_P}{\delta P(v_j^{(l)})}$ represents the number of effective paths the connection can contributes to creating. This encourages the selection of connections that can form the highest number of effective paths.

• Node Objective (Equation 6): The derivative with respect to the node objective indicates that the importance score of a connection increases if it belongs to an ineffective node. Additionally, the score is influenced by how many effective paths the connection can contribute to creating. This dynamic promotes the selection of parameters from diverse nodes, encouraging more nodes to become effective. These mechanisms demonstrate that the optimised importance scores effectively result in a subnetwork that maximizes the number of effective nodes and paths. Furthermore, through the convergence analysis in Section 3.3, we establish that each step of the iterative update can result in an incremental increase in the number of effective nodes and paths.

To better understand the convergence of the proposed method, we conducted empirical observations through ablation studies, examining how the number of effective nodes, paths, and kernels changes during the optimization of each objective. The following table shows the number of effective nodes, paths (in logarithmic scale), and kernels in the subnetwork obtained by optimising each objective independently.

• The maximum number of effective paths (in ln scale) is 69.22, achieved by optimising the path objective alone. (visualization during training https://postimg.cc/K4QZwFkg)
• The maximum number of effective nodes is 570, achieved by optimising the node objective alone. (visualization during training https://postimg.cc/21tNmSrV)
• The maximum number of effective kernels is 860, achieved by optimising the kernel objective alone. (visualization during training https://postimg.cc/BXXGzFfd)
• Optimising the overall Node-Path Balancing (NPB) objective must strike a balance among these three objectives. (visualization during training https://postimg.cc/fVw6Y0T3)

We have also included detailed visualisation illustrating how the number of effective nodes, paths, and kernels evolves during each training step in the revised version of the manuscript. These figures provide a clear view of the convergence behaviour of our proposed method. In summary, the results show that our objective effectively converges within approximately 3000 steps (or even fewer in some cases). This provides empirical evidence supporting our convergence theory discussed in Section 3.3.

For any confusion, the DPaI objective does not seek to find the best effective nodes and paths but rather to find the subnetwork that contains the highest number of effective nodes and paths.

We sincerely thank the reviewer for the detailed and constructive feedback. Your insights are invaluable, allowing us to refine further and clarify our work. We believe that addressing your comments will significantly enhance the quality and understanding of our paper. We hope these clarifications address your concerns. Thank you again for your thoughtful feedback, and we look forward to your further comments.

I suggest the authors provide a detailed discussion regarding the difference:

We appreciate the reviewer’s observation and suggestion regarding our methods. While it is true that the top-k operation with Straight Through Estimator (STE) and the Erdős-Rényi Kernel (ERK) technique are well-established and have been adopted in prior works [1, 2, 3, 4], our contribution lies in their application and adaptation to the Pruning at Initialization (PaI) domain, a context distinct from the pruning methods used during or after training [1, 2, 3, 5, 6].

Unlike traditional pruning techniques that optimise sparse masks after or during training [1, 2, 3, 5, 6] with task-specific loss function (most are MSE in image classification) with regularisation term, our proposed method (DPaI) introduces differentiable masks, specifically for PaI, allowing the identification of effective subnetworks before training begins. This novel application integrates network topology considerations, such as the Node-Path Balance (NPB) principle, into initialization-stage optimisation. By leveraging the differentiable Node-Path Balancing (d-NPB) framework, our method balances effective nodes and paths to achieve superior trainability and performance by directly optimising the differentiable-NPB objectives, which allows us to perform task-agnostic pruning. Furthermore, while techniques like [1, 2, 3, 5, 6] focus on pruning trained models or dynamically adjusting during training, DPaI eliminates the need for iterative pruning and retraining cycles, significantly reducing computational overhead. For ERK methods, we directly adopted those methods from [4], as are standard techniques for NPB and PHEW; our method is designed based on those findings.

To address the reviewer’s concern, we will provide a more detailed discussion in the paper, contrasting our approach with the methods above. Specifically, we will highlight how DPaI’s use of differentiable masks and continuous optimisation uniquely benefits the PaI domain, overcoming challenges such as layer collapse, large-scale discrete optimisation, and suboptimal network topology prevalent in prior pruning approaches.

what is the rationality of the tanh function.

The $\text{tanh}(\cdot)$ function naturally aligns with our objective since it only counts the number of effective nodes. Given that $N(\cdot) \geq 0$ represents the number of effective paths passing through a node/channel or kernel/connection, the $\text{tanh}(\cdot)$ function outputs $\text{tanh}(N(\cdot)) = 1$ for effective nodes and $\text{tanh}(N(\cdot)) = 0$ for ineffective nodes.

Alternatively, any activation functions $f(x)$ with similar characteristics can be employed. For instance, the sigmoid function $\sigma(x) = \frac{1}{1+e^{-x}}$ can be adapted for this purpose using $f(x) = 2\sigma(x) - 1$. This modification with sigmoid function yields derivatives computed as $f'(x) = 2\sigma(x)(1-\sigma(x))$, which exhibits behaviour consistent with the analysis presented in Section 3.3.

We thank the reviewer for the positive feedback and support of our work. We hope to have answered all of your questions satisfactorily below. Please let us know if you see any further issues in the paper that must be clarified or addressed.

Hyperparameter Sensitivity:

Thank you for raising the concern about hyperparameter sensitivity. We acknowledge that we performed a hyperparameter search for our method, which helped us improve performance over random hyperparameter settings on average. To address this concern, we included a detailed ablation study on different hyperparameters in Figure 2 of our paper. The results show that results are comparable across various hyperparameter configurations and consistently outperform state-of-the-art. This suggests that our method is robust to hyperparameter choices and does not rely on specific settings to perform well.

Implementation Complexity: ..which may be harder to apply to different task or models...

Thank you for raising the concern about implementation complexity and the applicability of our method to different tasks or models. Most previous Prune-at-Initialization (PaI) methods involve either iterative sparse mask selection, as seen in methods like SynFlow, or rely on large-scale discrete optimization techniques applied to each layer individually, such as in NPB. These approaches can be complex and may pose challenges when adapting to various tasks or models. Specifically, the NPB method, which relies on a discrete optimizer, faces several drawbacks related to the intractability and complexity of discrete optimization. First, it divides the overall objective into layer-wise objectives, leading to suboptimal solutions due to the need for inter-layer connectivity. Second, as the number of parameters in the base model increases, the complexity of the discrete optimisation problem grows significantly. To reduce this complexity and maintain pruning times comparable to other baselines, NPB partitions the parameters of each layer into multiple chunks, limiting each chunk to no more than 16,384 parameters, and then solves the discrete optimization for each chunk. This approach further compromises the solution quality, making the method impractical for large-scale models with billions of parameters. In contrast, the implementation of our proposed method is well-suited for integration into neural network training, making it easily applicable to various architectures such as ResNet, EfficientNet, and Transformer. Specifically, our method leverages the forward pass of the neural network to compute the number of effective paths and the backward pass to calculate the derivative with respect to the path objective. These forward and backward operations are efficiently implemented in widely used deep learning frameworks such as PyTorch and TensorFlow. Following this, the derivative with respect to the node objective is computed using the previously calculated path objective derivative, as outlined in Equation (6). DPaI can prune networks independently of the task data or the specific task objective function as a data-agnostic pruning-at-initialisation method.