Summary
This work proposes BAFFLE, a method for zero-order optimization in the federated setting. The main motivation behind zero-order optimization is its memory efficiency, as no activations need to be stored for a backward pass, it is compatible with hardware designed for inference only and, furthermore, its memory efficiency allows for training within a trusted execution environment (TEE). The recipe is quite simple, in that the server communicates a random seed and a model to the clients and then each of the clients does $2K$ forward passes with weights $w \pm \delta_k$ and their local data. The loss differences are then communicated and aggregated at the server where, together with the perturbations $\delta_k$, an approximation to the model gradient can be computed. The authors provide guidelines that allow such a scheme to be successful and then show, for moderately large values of $K$, that BAFFLE can reach performance comparable to vanilla backpropagation.
Strengths
The paper overall attempts to tackle an interesting problem that is useful in practice, especially given the success of large models. It is well written and clear in the description of the method and motivation. As for other strengths:
- The memory efficiency of BAFFLE is especially useful in the context of training large models on resource constrained devices. This memory efficiency also allows for training within a TEE.
- The zero-order nature might allow BAFFLE to be compatible with inference only models and accelerators.
- The authors also propose useful guidelines for BAFFLE, which I believe can be useful for zero-order optimization in general.
- The authors provide a theoretical convergence analysis of BAFFLE.
- Experiments cover a decent variety of settings, from different settings of data non-i.i.d.-ness to computation / memory efficiency and robustness
Weaknesses
I believe this work has relatively limited novelty, especially given prior work. More specifically, the authors discuss both forward finite differences and central finite differences but in the experiments they mention using forward finite differences, due to it being better for the specific amount of sample budget they consider. As far as I see, forward finite differences is essentially the exact same thing as Evolution Strategies (e.g., see [1]) which has already been considered for training neural network type of models in parallel (again, see [1]).
Furthermore, some of the claims of the method could use more details. For example, the authors mention that BAFFLE can apply to models optimised for inference, e.g., models that have been quantised and pruned. However, adding Gaussian noise to the weights can break both the quantization and (unstructured) pruning, so it is unclear whether these models would be able to run on their dedicated accelerators (which sometimes require weights to, e.g., be properly qunatized in a uniform grid). The authors also argue that secure aggregation (SAG) is compatible with BAFFLE, however, as far as I know, SAG requires the inputs to be quantised and this is not discussed in the main text.
Finally, if memory constraint is an issue, I believe an important baseline is missing and that is gradient checkpointing for BP, which can also reduce memory costs at the expense of more compute (which is the tradeoff BAFFLE also has).
[1] Evolution Strategies as a Scalable Alternative to Reinforcement Learning, Salimans et al., 2017
Questions
I have the following questions and suggestions
- Do the authors quantise the loss differences in order to allow for SAG? If yes, what is the scheme? How does it affect the convergence behaviour?
- The convergence based on a finite number of forward passes depends on a constant $C_0$. How large is that constant in practice?
- The harms of batch normalisation and its replacement by group normalisation, that the authors argue for as a guideline, is not new and has been discussed in various other works in federated learning, e.g, [2]. I would suggest that the authors properly attribute these prior works.
- Are the experiments shown with full client participation? Partial client participation is more common in cross-device settings (which is usually the case where resource constraints appear), so it would be interesting to see results with that, especially given that it increases gradient variance.
- I would suggest to the authors to elaborate on any differences with respect to [1].
- I would suggest to the authors to elaborate about the feasibility of evaluating quantized / pruned models with Gaussian noise on their weights on dedicated hardware accelerators.
[2] The non-IID data quagmire of decentralized machine learning, Hsieh et al., 2019
Rating
3: reject, not good enough
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.