Paper: BiDST: Dynamic Sparse Training is a Bi-Level Optimization Problem (ICLR-2024)

Reviewer comment (Theory):
The proposed algorithm does not have any convergence guarantee. Different from other bi-level optimization algorithms, the algorithm is based on a rough estimation of the gradient (eliminating the inversion of the second-order derivative) and is facing the discrete issue of the mask (the estimation of the gradient is not at the point of the continuous variable of the mask but a hard discretization of the soft mask). Thus, it is important to have a theoretical guarantee to ensure the performance of the proposed algorithm.

Author response:
We appreciate the reviewer's consideration of convergence guarantees. It's crucial to note that the nature of dynamic sparse training methods, including our proposed BiDST, often involves challenges that make it inherently difficult to provide formal convergence guarantees. Many successful literature has demonstrated that it is not necessary to include such convergence analysis for DST since the performance of DST is already proved to be SOTA in the sparse training research domain. We humbly disagree with the opinion that the lack of convergence guarantee is the major reason for rejection since the uniqueness of BiDST lies in its ability to address the dynamic nature of sparse training by leveraging a rough estimation of the gradient and handling discrete mask issues. On the contrary, we innovatively adopt bi-level optimization with improved efficiency to solve the DST problem, while we also reduce the computation cost in the solution that approximates the discrete mask into continuous variables. These aspects set BiDST apart from conventional approaches.

It's essential to recognize that existing methods in the DST domain also lack strict convergence guarantees due to the intricacies of optimizing sparse structures. Our results, demonstrating superior performance, speak to the effectiveness of BiDST in the context of dynamic sparse training. We believe that the focus should be on the empirical validation and real-world applicability of the algorithm, where BiDST excels, rather than imposing traditional convergence expectations that may not align with the unique challenges of sparse training.

Reviewer comment (Theory):
The upper-level objective is non-natural to me. Since the upper-level problem is also a finite sampling objective, we need some regularization to get good results. Meanwhile, the nature regularization is the l2 norm (i.e., $||\theta^*(\psi)||$). With this regularization, we can formulate the upper-level problem as the same as the lower-level problem. Thus, the problem can be formulated into a single-level problem, which is easier to solve than a bi-level problem. Further, people use bi-level formulation in machine learning usually when upper-level and lower-level objective contains different dataset.

Author response:
We appreciate the reviewer's consideration of the upper-level objective. The upper-level problem is the optimization of the mask parameter based on the trained weights. In our implementation, our code uses one single optimizer, which adopts the standard L2 regularization with the penalty parameter that can be found in Table A.1 in our appendix. However, we do not believe it is necessary to incorporate such regularization terms in the equations. While regularization, including the traditional L2 norm, is a common practice in training neural networks, its omission from both objective formulations is a deliberate choice to maintain clarity and emphasize the bi-level optimization framework. For the lower-level problem, however, since we are optimizing sparse weights, we add the second term as a sparse regularization to stabilize the training without incurring dense computation of all parameters. This regularization is controlled by a coefficient hyper-parameter that can also be found in Table A.1.

We also humbly disagree with the opinion that if we add the L2 regularization to the upper-level problem, the whole problem becomes the single-level optimization problem. Due to the nature of sparse training, there are two objectives that need to be optimized regardless of the method used for such optimization. Those two objectives are weights and masks. Even if we add L2 regularization and use it for mask/weight optimization, the effectiveness cannot be comparable to our design since using L2 to find mask is essentially through the penalization of all mask parameters, but our approach is to analytically learn an optimal mask based on all information flow in training. 

It's also crucial to emphasize that even with the addition of a regularization term in the upper level, the overall optimization, by our intentional design, remains a bi-level optimization problem in nature. We have provided clear illustrations, as shown in Figure 1, demonstrating that while other dynamic sparse training (DST) methods may be inherently bi-level, they lack the optimization techniques that BiDST employs. Our focus is on optimizing both levels simultaneously through our novel bi-level approach, which distinguishes BiDST from other sparse training methods.

Lastly, it is not always true that bi-level optimization is used with different datasets for its upper and lower problem. Please see the literature of [R1]. We acknowledge that some works might use different datasets, but bi-level optimization is a very general solution framework, and can be used in different settings when using proper design. Therefore, we can embrace the possibility that bi-level optimization is a very useful tool to solve real-world problems, and that’s indeed what this paper’s focus is.

[R1] Zhang, Yihua, et al. "Revisiting and advancing fast adversarial training through the lens of bi-level optimization." International Conference on Machine Learning. PMLR, 2022.

Reviewer comment (Reproducibility):
The reason directly optimizing the mask is typically not done is of course that a (hard) binary mask is non-differentiable. Using a soft mask in any form would necessarily be dense, not sparse, training. I think the authors do not focus on explaining what they do with respect to both of these things enough in the paper as it is currently written, as I'm still not confident I understand how this is addressed by BiDST. I believe the authors use the Gumbel-Softmax trick, and while that might explain how they learn a binary mask using gradients, that would mean that BiDST uses dense training as far as I can see. And yet the authors demonstrate sparse training acceleration in their results (and this itself is unclear). I have asked this in the question section in more detail, and a comprehensive answer by the authors is vital in my being able to understand this paper better as a reader and a reviewer. Understanding the fact of if BiDST is in fact doing dense weight/mask updates is key in understanding if it is fair to compare BiDST with existing DST sparse training methods, such as SET, as the main motivation of these methods is to reduce compute during training.

Author response:
A1. According to Section 2.3, we discuss that the BiDST’s key parameters within the computation graph (e.g., weight, mask) are performing sparse computation (i.e., in both forward and backward propagation) to reduce the training cost. The reason is that we only learn partial mask parameters that associate with the active non-zero weights, and the update of mask parameters is also very efficient since it reuses the activation gradients and takes no input data for forward computation. For weights in the neural network, the forward computation is sparse according to equation 4’s lower equation that all terms are masked with the sparse mask. For backward computation, we only update the non-zero weights so that the entire computation is sparse. We also want reviewer to refer to reviewer xhao’s strengths part that acknowledges our method can achieve overall efficiency during sparse training. For the results that demonstrate on-device acceleration, we adopt the compiler-level optimization [R1] that transforms the DNN sparse computation graph with static execution code, thus achieving reduced compact computation. We will add our implementation details in our revision.

Reviewer comment (Theory):
In the explanation below Equation 8 in section 2.2, the authors casually mention that they simply replace the second-order partial derivatives with 0 in their derivation of the weight/mask-update rules for BiDST. Isn't it precisely these second-order partial derivatives that carry much of the information of the relationship between the mask and weight parameters? There is no discussion on the effect of this, or why this is a reasonable thing to do aside from the fact it is expensive to calculate. Note that the first-order partials are retained. After all the fanfare of BiDST doing optimization "properly" compared to existing DST methods in the motivation, this is quite a let-down, and I think the authors should note this earlier/be careful of giving the impression of over-claiming and be much more transparent that they are crudely approximating the bi-level optimization methodology.

Author response:
A2. We acknowledge the reviewer's concern. The decision to replace second-order partial derivatives with 0 is made for computational efficiency reasons. We follow the method used in citation Finn et al. (2017) and Liu et al. (2022a) to achieve this conversion. However, we understand the need for a more detailed explanation. We will revise the paper to include a thorough discussion of the implications and potential limitations of this approximation. Transparency regarding this decision will be emphasized to provide a more accurate portrayal of the optimization methodology used in BiDST.

Reviewer comment (Experiments):
As explained in 3.1, BiDST experiments are presented using fewer training epochs to ensure a "fair" comparison due to the "mask learning computation cost", however it's not detailed what exactly is the fair number of epochs to compare and why, and I didn't see any numbers or details on what the "mask learning computation cost" overhead is w.r.t. existing DST and other sparse training methods.

Author response:
A3. We understand your concern. The reduced training epochs is shown in Table A.1 in appendix. For the mask update, the only computation related to the overall training cost is the Gumbel-softmax, and it becomes the mask mask learning computation cost overhead to the overall training. Since Gumbel-softmax is a very efficient implementation which only needs element-wise computation instead of convolution computation, the overall cost for mask updating is approximately 1% of the total training cost when we train CIFAR on ResNet-32 with 160 epochs. Therefore, we reduce the training epochs by 2 epochs to match the training cost with other baselines that train 160 epochs.

Reviewer comment (Evaluation):
It is not clear how significant some of the ImageNet results are, being within 0.2 percentage points of the baselines in some cases. Because of this it really stands out as suspicious that although the authors stress that ImageNet experiments are all performed three times, we are not given the variance across these three runs to aid us in understanding the significance of the results, as for example done with CIFAR-10/100. The authors say "we omit these because it's less than 0.1% std dev (is this percentage or percentage points?). Why not just show them in the table? I would not ask for multiple imagenet runs normally, but if you've performed them and are stressing that you did, why hold back on showing the variance?

Author response:
A4. The reason that we report std dev for CIFAR and not ImageNet is because the data is meaningful for CIFAR but not quite informative for ImageNet. We understand the importance of providing a comprehensive understanding of the results. We will include the variance across the three runs of the ImageNet experiments in the paper to better convey the robustness and significance of the reported results.

Reviewer comment (Evaluation):
In section 3.3, the authors use IoU to analyze the evolution of the mask during training and compare to existing DST methods, claiming "BiDST achieves better mask development" based on low IoU compared to other DST methods. However, while potentially interesting, the motivation for this analysis is nowhere near convincing enough to make such a bold claim.

Author response:
A5. We appreciate the feedback. We justify the IoU changes for better mask development by providing the final accuracy result. It shows in the Figure 4 that BiDST achieves the best accuracy compared to the baseline methods (RigL, MEST) that uses manually set mask update policy. From the IoU curves, we can see that the purpose for other baselines that use manually set update policy is to ensure a reasonably large search ratio on the mask (e.g., MEST uses different levels of search ratio, while RigL uses a cosine annealing schedule for mask update ratio). However, such heuristic settings cannot always ensure the mask searching process is naturally suitable for the network. Additionally, a similar approach that uses IoU to analyze the accuracy can be found in [R2]. So we believe it is a solid metric. We will provide a more nuanced explanation in section 3.3. The revised section will explicitly justify the use of IoU and clarify the basis for claiming superior mask development by BiDST.

Reviewer comment (Reproducibility):
In section 3.4, when explaining the "sparse training" speedups shown, the authors explain "training acceleration is obtained from compiler optimization that skips the zeroes weights". If this is a compiler optimization, it's static analysis, i.e. a fixed mask. How is this possible for a DST method that is changing masks potentially every iteration, or is the "training engine" actually to speed up inference timings? This explanation is important as speeding up unstructured sparsity on real-world hardware is something that is not easy, and in fact is well worth a standalone publication if it was truly being solved so easily by the authors.

Author response:
A6. We have explained in section 2.4, that the end-to-end on-device training overhead is because of the “recompilation” of the “static machine code”. The reviewer is correct about the execution code for sparse neural network is in the form of static code. However, DST needs to change this static code when the mask is updated (i.e., new computation graph). This recompilation needs time, which becomes the overhead time that slows down the overall training time. The benefit of using compiler code generation is because the execution efficiency is very high, even with the unstructured sparsity. According to the discussion and experiments in MEST, we can see that MEST successfully incorporates different sparsity in the compiler design, which certifies the possibility of using compiler optimization to accelerate BiDST. We use the CSR format to store the non-zero weights and integrate the compact form into the static code to achieve sparse acceleration. However, we must point out that the focus of this paper is in the algorithm-level innovation that not only achieves good DST accuracy with an analytical solution, but also helps to reduce the overhead related to hardware-level implementation.

Reviewer comment (Presentation):
The background is very short, and awkwardly positioned in the paper. It appears much of the relevant background is distributed throughout the method implicitly by citing baselines and methodology, but personally I'm a fan of the traditional consolidated background before a method.

Author response:
A7. We appreciate the reviewer’s opinion. We will re-organized the background section after the introduction section.