Natural Perturbations for Black-box Training of Neural Networks by Zeroth-Order Optimization

Thank you for reviewing our paper and evaluating that the contribution is sufficient and the paper is very well written.

If the authors include a discussion on the memory-costs of their algorithm and provide experimental results on (a subsection of) the ZO-Bench datasets with LLM fine-tuning, it would widen the audience of their method

Below, we report and discuss the memory costs of our algorithm not only for already reported neural networks but also for a newly constructed larger-size neural network with one million parameters. However, it is difficult for our current algorithm to perform experiments on the ZO-Bench datasets with LLMs that have much larger number of parameters (perhaps hundreds of millions at the smallest). We feel that we need a different strategy, e.g., a much stronger approximation of the FIM than making it block diagonal, and would like to contribute as future work in another paper.

What I am missing is a study for the memory-usage of this method. ... While this paper explicitly doesn't target memory-reductions, it would have been a great addition to make this method appealing to a wider audience.

The table below shows the memory cost corresponding to the settings of Table 1. The difference between ZO-I and ZO-NP roughly corresponds to the additional memory cost of introducing natural perturbations.

Memory footprint (GB) corresponding to the settings of Table 1

We can reduce the memory cost by decreasing the maximum block size $N_\mathrm{max}$, as the following table shows. Together with Figure 7 in the submitted main text, we understand that ZO-NP trades off the test accuracy and the memory cost overhead.

Memory footprint (GB) for FashionMNIST corresponding to the setting of Figure 7

We also measured the memory footprints of a larger MLP-mixer for CIFAR10, as shown in the table below. The number of mixers has increased from 3 to 12 and the channel width has increased from 32 to 256, making the number of parameters increased from 33,642 to 1,706,762. Again, the difference between ZO-I and ZO-NP roughly corresponds to the additional memory cost of introducing natural perturbations, and we can reduce the memory cost by decreasing the maximum block size $N_\mathrm{max}$.

Memory footprint for CIFAR10 with a larger MLP-mixers

We will include these results in the revised manuscript to clarify the memory cost of the proposed algorithm for a wide variety of neural networks.

In terms of datasets, many papers in this literature study ZO for LLM fine-tuning on datasets mentioned here: https://github.com/ZO-Bench/ZO-LLM It would be great to see Natural Perturbations ZO applied to those.

We will include the references to ZO-LLM and the corresponding paper [1] in the revised manuscript, and state that applying the idea of natural perturbations to ZO-based LLM fine-tuning is a future work.

[1] Zhang, Yihua, et al. "Revisiting zeroth-order optimization for memory-efficient LLM fine-tuning: A benchmark." arXiv preprint arXiv:2402.11592 (2024).

Thanks for the nice suggestions to widen the audience of our method.

Thank you for reviewing our paper and evaluating the idea of designing the sampling distribution as innovative.

Can the authors elaborate on how the method might be adapted or scaled for neural networks with orders of magnitude more parameters (e.g., in the millions), especially regarding the computational cost of Jacobian and FIM estimation?

Yes, stemming from the experimental condition of the last row (CIFAR10) of Table 1, we additionally performed experiments with a larger MLP-mixer. The number of mixers has increased from 3 to 12 and the channel width has increased from 32 to 256, making the number of parameters increased from 33,642 to 1,706,762. The following table shows the computational cost with the million parameters. The difference between ZO-I and ZO-NP roughly corresponds to the cost of Jacobian and FIM estimation. While the elapsed time (reported as seconds/epoch) overhead is less than 3%, the memory footprint cost is significant. We can reduce the memory footprint cost by decreasing the maximum block size $N_\mathrm{max}$ as shown in the last two rows. As a tradeoff, however, this increased the number $B$ of blocks and consequently the elapsed time in seconds per epoch.

Elapsed time in seconds per epoch and memory footprint for CIFAR10 with a larger MLP-mixers

One minor potential concern is the additional query cost for computing the FIM, which is addressed via block partitioning—but its impact on extremely large-scale networks remains to be further validated. ... Although the block-coordinate method mitigates query costs, it remains to be seen how the approach scales to very large-scale networks (e.g., millions of parameters).

According to our current procedure for computing the FIM shown in Algorithm 2, the additional query cost is not affected by the scale of a network (the number $N$ of parameters) as long as the block size $N_b\leq N_\mathrm{max}$ is the same. This is because only one block $b$ is sampled in line 9. The problem for a large-scale network instead is that the number $B$ of blocks is large and it takes many epochs to compute the FIMs for all blocks. In other words, the FIMs of many blocks remain as the initialized identity matrix $\mathbf{I}$ until many epochs have passed.

The following table shows such situations. The second row corresponds to the last row in Table 1, where the number $B$ is not so large. The third row corresponds to the above introduced larger MLP-mixer with million parameters. Since the number $B$ is large in this larger network, some blocks remained as initialized even after 1000 epochs with the update frequency hyperparameter $T_\mathrm{ud}=100$. We feel that we can show the scalability and also some limitations of our current method/algorithm by additionally reporting the results like these two tables in the revised manuscript. Thanks for raising these issues.

The number of epochs (1st row) and the number of blocks (the remaining rows) whose FIM are computed.

Essential References Not Discussed: Wierstra D, Schaul T, Glasmachers T, et al. Natural evolution strategies[J]. The Journal of Machine Learning Research, 2014, 15(1): 949-980.

We will refer to this paper in the revised manuscript and discuss the differences described in the following table. It shows that the proposed method of natural perturbations is superior in the sense that it can be used for a larger problem with a larger number $N$ of parameters than natural evolution strategies due to the size of FIM. Thanks for suggesting the reference.

Differences between Natural perturbations and Natural evolution strategies

Thank you for reviewing our paper and evaluating that the proposed method makes sense, experimental improvements are consistent, and the theoretical analysis provides a solid foundation.

How does the proposed natural perturbations method scale to larger neural network architectures, particularly in terms of computational resources and training time?

Stemming from the experimental condition of the last row (CIFAR10) of Table 1, we additionally performed experiments with a larger MLP-mixer. The number of mixers has increased from 3 to 12 and the channel width has increased from 32 to 256, making the number of parameters increased from 33,642 to 1,706,762. The following table shows how the computational resources in terms of memory footprint and training time (seconds per epoch) have increased with the 50 times larger number of parameters.

Memory footprint and elapsed time in seconds per epoch for CIFAR10 with MLP-mixers

We observe that the memory footprint and the training time have increased by 5.5 and 14.5 times, respectively, which is rather mild considering the 50 times larger number of parameters. More critical is a linearly increasing (50 times) number $B=3334$ of blocks. This is due to the nature of the block coordinate approach with the same $N_\mathrm{max}$, and it takes many epochs to compute the FIMs for all $B=3334$ blocks. See the second table in the reply to Reviewer 73sf for a related experimental result. We think that the proposed natural perturbations method scales to larger networks to some extent but a network with one million parameters might be a limitation. We will include such a discussion in the revised manuscript. Thanks for raising this issue.

-- The experiments primarily focus on small datasets, and it would be beneficial to explore performance on larger datasets and more diverse tasks including NLP.

As reported above, we have tested a larger neural network with one million parameters. However, it is difficult for our current algorithm to perform experiments on larger datasets like with LLMs that have much larger number of parameters (perhaps hundreds of millions at the smallest). We feel that we need a different strategy, e.g., a much stronger approximation of the FIM than making it block diagonal, and would like to contribute as future work in another paper.

-- The paper could benefit from more detailed ablation studies to disentangle the effects of sample number.

We might misunderstood this comment, but if 'sample number' means the number $Q$ of perturbations in Eq. (10), we agree that we should additionally report some results with varying numbers $Q$. So, we have obtained the following results. The query budget for each $Q$ was determined by the number of queries that ZO-I and ZO-co consumed for 100 epochs. We observe that the proposed method ZO-NP consistently outperformed the other two methods, except the extreme cases ($Q=1,2$), where the number of epochs allowed was small since the relative query consumption cost for the black-box Jacobian computation becomes large. We will include such results in the revised manuscript. Thanks for the comment.

Test accuracies of FashionMNIST with varying the number $Q$ of perturbations in (10)

Thank you for reviewing our paper and evaluating that the concept of natural perturbation is well established.

It would be interesting to see if the proposed method outperforms the baseline with the same computation budget, or the proposed method gives an interesting trade-off.

The computation cost can be characterized by the elapsed time and the memory footprint. For the former, we already show the corresponding results in the right plot of Figure 6 and the bottom plots of Figure 12. These show that the proposed method ZO-NP outperformed the baselines with the same elapsed time. For the latter, we newly report the results in the reply to Reviewer 2tFg (please see the first two tables). With these tables and Figure 7, we understand that ZO-NP trades off test accuracy and memory overhead. Thanks for raising this issue. We will include these memory footprint results in the revised manuscript.

Why in Table 1, ZO-I also uses blocking. Can we add the comparison to the most naive baseline i.e. ZO-I without blocking?

The reason to use blocking also for ZO-I is that the performance gets better with an appropriate block size $N_\mathrm{max}$, as Figure 7 shows. And we already have some results without blocking in Figure 11 in Appendix D.2. While the main intention there was to compare with CMA-ES, we can also compare the results with ZO-I.

For ZO-co and ZO-NP, are they using the same Q? ... Since ZO-NP has a higher computation overhead, it seems to make sense to compare it with ZO-co with a larger Q (which have similar computation cost to ZO-NP) for a fair comparison.

We used the same Q for ZO-co and ZO-NP, and ZO-co results in a sparse perturbation (a sparse estimated gradient as a weighted sum of perturbations to be exact) as you pointed out. For more related information, see the discussion on the number of changing parameters in Section 5.3.2. Regarding the computational overhead of ZO-NP, we have already made a fair comparison in terms of elapsed time (the right plot of Figure 6 and the bottom plots of Figure 12).

In figure 6, it looks like it is not fully converged and I am wondering what happens if we increase the epochs till it fully converges. I am wondering if ZO-NP has the advantage just in convergence speed, or also in the ending test accuracy.

The two tables below show the situations with more epochs (three times more). The epochs allowed for ZO-NP, which consumed extra queries for the Jacobian computation, are shown in the parentheses. The training losses continued to decrease, but the speed efficiency was ZO-I < ZO-co < ZO-NP. The test accuracies converged in ZO-co and ZO-NP but not in ZO-I. It seems that ZO-NP entered the range of overfitting.

Training loss

Test accuracy

In line 167 left column, it mentions that a good sampling strategy should maximize the entropy of the distribution. Why is that? Is there theoretical or empirical evidence supporting it? Actually, in the experiment, can you add another baseline with sampling from another distribution with a smaller entropy?

The reason is to make the sampled perturbations explore as widely as possible. This statement is our new way of understanding the sampling in ZO optimization. Section 3.1, especially the paragraph starting from line 213 left column, provides theoretical support as the existing sampling strategy from $\mathcal{N}(\mathbf{0}, \mathbf{I})$ can be derived from this statement. We could perform some experiments where the entropy of a sampling distribution is randomly reduced independent of the FIM. Please let us know in a further reply if you prefer to see the results.

In line 321left column, it mentions that FIM also becomes block diagonal. Why?

This is because the block matrix size is the same for $\Sigma_b$ and $\mathbf{F}_{\theta^{(b)}}$ by (34).

In line 216 right column, it mentions a too large FSD leads to unstable training. Is there evidence to support it?

A good lecture note can be found at https://www.cs.toronto.edu/~rgrosse/courses/csc2541_2022/readings/L03_metrics.pdf and Figure 1 there on the Rosenbrock function is very intuitive.

In line 87 right column, what is z and what does p(z|y) means?

The specific form of p(z|y) depends on the type of task the neural network performs. For example, if the neural network performs a classification task, p(z|y) is a multinomial distribution as explained in Section 4.1.1. And z is a random vector that is marginalized in the FIM computation (28) and specified by a target vector in the loss computation (1).