Paper: Flag Aggregator: Scalable Distributed Training under Failures and Augmented Losses using Convex Optimization (ICLR-2024)

Reviewer comment (Theory):
There are no specific examples that can show FA is theoretically better than existing methods, and thus makes the claim not persuasive to some audience. It can be, possibly, FA only performs well in the specific datasets, models, the way of distributing datasets among compute nodes in the experiments presented in this work. Even a naive toy example would help people understand why FA is theoretically better.

Author response:
Please see Figure 3 for a specific example. Here, we extracted gradients of parameters of ResNet-18 at $p=15$ workers, of which $f=2$ produce uniformly random gradients for 1000 iterations. At each iteration, we compute the mean of gradients and use it to calculate the values of workers. We also compute the optimal subspace provided by FA and calculate the values of workers using that at each iteration. The histogram of values computed using the mean gradient is much lower since most values are close to zero whereas when the optimal subspace provided by FA aggregation is used, the value is enhanced. We observed this phenomenon on many datasets, and architectures at various training stages, and we are happy to give more details with exact checkpoints in the Appendix for scientific replication purposes. In summary, we believe that this decreased value while using average gradients is due to the high dimensional phenomenon and can be easily replicated. 

We used the standard data-parallel training setup used in practice. First, we shuffle the data to randomize it, using a seed for consistency. Then, we create partitions proportional to the number of workers, ensuring each chunk is a different part of the dataset. These partitions are lists of indexes, representing how the data is divided. This setup allows distributing different data chunks to multiple workers in a training process, ensuring each worker gets a unique, random subset of the data for independent and identically distributed gradients.

We assure the reviewer that the way of distributing datasets did not have any effect on our experiments. We consider the reproducibility aspects of our experiments seriously, so we will release all the datasets, along with code, model, a blog post, as well as a tutorial to reproduce all the figures in our submission on acceptance.

Reviewer comment (Evaluation):
There are no comparisons between FA and existing aggregators in terms of computation complexity. It can be that, FA is too expensive to use in every iteration.

Author response:
Please refer to the potential limitation in the appendix and the corresponding Figure~18. Please note that although the iteration complexity of FA is higher, the utility of each iteration is higher, which is reflected in the time-to-accuracy comparison. The figure shows two main points: 1) FA converges to a higher accuracy, and more importantly, 2) it reaches to that level in less (wall-clock) time. In this sense, we think that FA's somewhat higher per-iteration complexity is well justified. In the future, we plan on evaluating eigenvalue solvers to further accelerate training time. We have provided a more detailed explanation, along with preliminary evaluation of per-iteration complexity in Section G in the Appendix. In essence, we believe that this gap in per-iteration time can be further reduced using modern fast eigenvalue solvers.

Reviewer comment (Theory):
The intuition behind the development of FA is not clear enough to me, this goes to the first bullet, can the authors provide an example or a simple problem model to illustrate this? And why is FA optimal as claimed in line 83? Can the authors elaborate on this point?

Author response:
FA is optimal in the sense that it seeks to find an approximate solution to the optimization problem formulation given in Equation (5). Our derivation using MLE in Sec 2.2 shows that problem (5) can be used whenever the value of the gradients needs to be maximized or in other words, when we want to maximize the values provided by gradients from workers in a distributed training pipeline. Finally, to see that FA aggregation guarantees convergence, please notice that our update is given by $YY^{\top}G{\bf 1}$ where $G{\bf 1}$ simply denotes the sum of gradients from workers. Now since $YY^{\top}$ is a symmetric positive semidefinite matrix, it is guaranteed to have an eigenbasis. Intuitively, this means that the update rule itself simply rotates the gradients, scales them in specific coordinates, and rotates them back to their original basis. This is our main intuition in proposing subspace based aggregation that is currently not explored in distributed training pipelines used in practice.

Reviewer comment (Writing):
In line 22, the description of quadratic function is not clear to me, can the authors put it simpler?

Author response:
We mean that $A_{g_1,\cdots,g_p}(Y)=\sum_{i=1}^p\\|Y-g_i\\|_2^2$ where $Y\in\mathbb{R}^{n\times 1}$
 (or simply $C=\mathbb{R}^n$) where $n$ is the number of training parameters.

Reviewer comment (Theory):
How will the augmented data and stable diffusion influence the mutual dependence of distributed data, can the authors comment on that?

Author response:
Thank you for the question! When data augmentation is independent of the training data set, then the mutual independence in data obtained by augmentation is preserved since functions of independent variables are independent, so no dependence after augmentation. However, when models such as stable diffusion or latent space or feature based augmentations are used, independence is not preserved anymore. This is because parameters of stable diffusion, encoders used in obtaining latent representations, features are trained using training data set -- information from individual data points in the training data set can influence the parameters of the models.  So we cannot assume independence of augmented data using these models, and dependence increases.

Reviewer comment (Writing):
In line 72, what does it mean for noise to be nonlinear?

Author response:
If $g$ is the true gradient, then noise $\epsilon$ is *linear* if the obtained or observed gradient can be written as $g+\epsilon$. If the obtained gradient cannot be decomposed in this way, we call the noise to be nonlinear. The same condition can be used for training data samples $x$. In our submission, we mention two specific nonlinear noise models in Sec 3.1 near L283-290 viz., Lotka-Volterra and Arnold Cat Map, for imaging datasets.

Reviewer comment (Theory):
In line 73 & 74, how does the discrete hyper parameters hamper convergence of the overall training procedure, can you elaborate on that?

Author response:
For every value of discrete hyperparameter, the training usually needs to be restarted from scratch, and so we may require multiple training runs to choose the optimal value.

Reviewer comment (Experiments):
In line 76, the authors mention 'sparse', does that correspond to the choice of the second dimension of $Y$, i.e., $m$, how is $m$ chosen?

Author response:
Yes. We believe that the reviewer is pointing out this is a discrete parameter in our setting. Thank you, we completely agree with this subtle observation. Fortunately, while using FA, there is an ordering in our setting due to orthogonality constraints as explained in L135-148, so a higher $m$ can be used in practice. In this sense, it is advantageous to use FA compared to discrete parameters that are used in $k$-NN based algorithms. Experimentally, we observed that choosing a larger $m$ in the definition of $Y$ increases the wall clock time since a larger SVD has to be computed at every iteration but *not* the number of iterations for convergence.

Reviewer comment (Writing):
In line 115, it seems that $g$ should $g_i$, and there is no need to use trace on a scalar right?

Author response:
Yes, $g$ should be $g_i$, thank you for pointing. We have fixed the typo. The $\text{tr}(\cdot)$ is used in tractability analysis of our formulation in Equation (5). So we thought that it would be nice to hint to the reader of further developments in our submission. We are happy to implement suggestions from the reviewer to improve the readability of our submission.

Reviewer comment (Writing):
In figure 3, I don't understand what is optimal subspace, do you mean optimality in the sense of formulation (5)?

Author response:
Yes. We have answered this in detail in your previous question asking why is FA optimal. We will include our explanation and any other suggestions from the reviewer.

Reviewer comment (Theory):
Why is beta distribution used in the formulation of FA?

Author response:
In general, Beta distributions can be used to model events which are constrained to take place within an interval defined by a minimum and maximum value or bounded sets. Since the value of a worker, $v_i$ is in the unit interval $[0,1]$ by definition,  assuming values to be sampled from a Beta distribution is valid. Moreover,  parameters $\alpha,\beta$ used in the definition of Beta distribution can be used to indicate different mathematical quantities such as convexity, dispersion of the distribution of values in our formulation (5). We believe that the additional flexibility can enable faster training procedures (in terms of iteration complexity) which may be important in the distributed setups.

Reviewer comment (Experiments):
In the experiment starting from line 300, what attacks are being used?

Author response:
In this experiment, $f$ workers produce uniformly random gradients in $[0, 1]$ so has coordinatewise mean  is $0.5$ and standard deviation is $ \sqrt{\frac{1}{12}}$.

Reviewer comment (Experiments):
The choice of setting the subspace rank $m$ to $(p+1)/2$ in all experiments raises a question. Is this decision rooted in theoretical analysis or other considerations?

Author response:
**Q.** Is the choice of setting the subspace rank $m$ to $(p+1)/2$ in all experiments rooted in theoretical analysis or other considerations?

**Ans.** Great question! We decided to use $m$ to be $(p+1)/2$ since it corresponds to a naive prior that each worker may provide noisy gradient with probability $1/2$. In this case, the expected number of clean gradients will be the integer greater than or equal to  $p/2$.  Our choice also corresponds to a realistic setting in most use cases -- that the majority (i.e., more than half) of workers provide correct gradients.

Reviewer comment (Experiments):
The paper introduces a general framework in Section 2, considering a general norm and suggesting SDP for solving the system. However, the focus shifts to $\ell_1$ regularization later due to its ease of optimization. An experimental comparison of different regularization terms could bridge the apparent gap between Sections 2 and 3.

Author response:
**Q.** The paper introduces a general framework in Section 2, considering a general norm and suggesting SDP for solving the system. However, the focus shifts to regularization later due to its ease of optimization. An experimental comparison of different regularization terms could bridge the apparent gap between Sections 2 and 3.

**Ans.** We use the SDP relaxation of the MLE in Equation (4) to argue that solving FA problem may be tractable since SDPs can be solved efficiently from the theoretical standpoint. We provided the regularized version with two specific regularization functions $\mathcal{R}$ that our tractability analysis using SDP can handle since these two choices of $\mathcal{R}$ are convex functions with respect to $Y$. We provided preliminary experiments with pairwise data dependent regularizer in Sec F.2, Figure 10. We plan to do a detailed study on different regularizers with FA in the near future. We will definitely clarify this in the camera-ready.

Reviewer comment (Experiments):
It would be useful to explore whether the proposed algorithm offers advantages even when $f=0$.

Author response:
**Q.** It would be useful to explore whether the proposed algorithm offers advantages even when $f=0$


**Ans.** We agree, thank you for pointing this out!  Our update rule is given by $YY^{\top}G{\bf 1}$ where $G{\bf 1}$ simply denotes the sum of gradients from workers. Reviewer will be able to see that since $YY^{\top}$ is a symmetric positive semidefinite matrix, it is guaranteed to have an eigenbasis. Intuitively, this means that the update rule itself simply rotates the gradients, scales them in specific coordinates, and rotates them back to their original basis. So, we may think of $YY^{\top}$ as an adaptive preconditioner and can be used whenever the loss landscape is ill-conditioned. In the Appendix, we have added a convergence plot in Figure 11 for the noiseless setting i,e, $f=0$. In this setting, a robust aggregator should be as good as the standard distributed SGD with Mean as the aggregator, which we can clearly observe here.  In Figure 11 we can see that FA, performs better than mean, which is consistent with the findings in [Mankovich et al. 2022]. 

[Mankovich et al. 2022] The Flag Median and FlagIRLS, CVPR 2022.

Reviewer comment (Evaluation):
I am not sure how much overhead the SVD might bring in practice, could you provide some real-world measurement? So far all the empirical results are epoch-wise measuring.

Author response:
**Q.** How much overhead the SVD might bring in practice, could you provide some real-world measurement?

**Ans.** Yes, please refer to the potential limitation in Section G in the appendix. Indeed, there is nontrivial overhead to execute FA in each training iteration. Fortunately, this overhead is compensated by reduction in total number of training iterations in all our experiments, thus providing significantly lesser wall clock time. We have provided more explanations and suggestions for potentially accelerating training time by using approximate eigenvalue calculation. When using exact eigenvalue calculations, the overheads per iteration can be seen in Figure 18.

Reviewer comment (Experiments):
Most of the baseline compared in the experiments seem to be from at least five years ago (2018); I wonder if the authors can compare their approach with latest algorithms? For instance Allouah et al. (2023a;b); Farhadkhani et al. (2022) as mentioned in the related work.

Author response:
**Q.** Most of the baseline compared in the experiments seem to be from at least five years ago (2018). Can the authors compare their approach with latest algorithms? For instance Allouah et al. (2023a;b); Farhadkhani et al. (2022)

**Ans.** Thank you for the suggestion. We compared FA to RESAM [Farhadkhani et al. (2022)] which adapts the concept of gradient momentum to distributed architectures. RESAM involves each honest worker computing and storing the momentums of their stochastic gradients as well as the gradients, and sending the momentum to the server which increases computational complexity per iteration. Similar to [Allouah et al. (2023a;b)], RESAM still relies on resilient aggregation similar to our baselines at the server, so we complement it with the existing baselines including CGE [Gupta et al. (2021)]. Nevertheless, we show the comparison in Figure~13 of the updated submission while ignoring the overheads of RESAM. The figure shows FA's advantage under three different scenarios. In addition, FA is designed as a robust standalone aggregator that needs no extra computation or space for calculating gradient momentums and storing them on the worker side.

[Gupta et al. (2021)] Byzantine Fault-Tolerant Distributed Machine Learning Using Stochastic Gradient Descent (SGD) and Norm-Based Comparative Gradient Elimination (CGE)

[Farhadkhani et al. (2022)] Byzantine machine learning made easy by resilient averaging of momentum

[Allouah et al. (2023a)] Distributed Learning with Curious and Adversarial Machines

[Allouah et al. (2023b)] Fixing by Mixing: A Recipe for Optimal Byzantine ML under Heterogeneity

Reviewer comment (Experiments):
I'll increase my rating if comparison to more recent algorithms is provided.

Author response:
**Q.** Most of the baseline compared in the experiments seem to be from at least five years ago (2018). Can the authors compare their approach with latest algorithms? For instance Allouah et al. (2023a;b); Farhadkhani et al. (2022)

**Ans.** Thank you for the suggestion. We compared FA to RESAM [Farhadkhani et al. (2022)] which adapts the concept of gradient momentum to distributed architectures. RESAM involves each honest worker computing and storing the momentums of their stochastic gradients as well as the gradients, and sending the momentum to the server which increases computational complexity per iteration. Similar to [Allouah et al. (2023a;b)], RESAM still relies on resilient aggregation similar to our baselines at the server, so we complement it with the existing baselines including CGE [Gupta et al. (2021)]. Nevertheless, we show the comparison in Figure~13 of the updated submission while ignoring the overheads of RESAM. The figure shows FA's advantage under three different scenarios. In addition, FA is designed as a robust standalone aggregator that needs no extra computation or space for calculating gradient momentums and storing them on the worker side.

[Gupta et al. (2021)] Byzantine Fault-Tolerant Distributed Machine Learning Using Stochastic Gradient Descent (SGD) and Norm-Based Comparative Gradient Elimination (CGE)

[Farhadkhani et al. (2022)] Byzantine machine learning made easy by resilient averaging of momentum

[Allouah et al. (2023a)] Distributed Learning with Curious and Adversarial Machines

[Allouah et al. (2023b)] Fixing by Mixing: A Recipe for Optimal Byzantine ML under Heterogeneity