Bi-Directional Communication-Efficient Stochastic FL via Remote Source Generation

We thank Reviewer Bbin for their positive evaluation of our work and for their thoughtful comments. We answer their question as follows.

• Q1/W1 (Other remote generation schemes) The choice of MRC is motivated mainly by practicality and ease of exposition. A natural improvement is ordered random coding (ORC) [R1], which reduces the entropy of the indices selected for transmission using the Gumbel-Max trick. This extension can directly be applied in our method with a minor adaption that we refrained from including here for clarity. On a related note, the authors of ORC [R1] further expose an interesting connection to the Possion functional representation mentioned by the reviewer. While such proposals allow exact remote generation, lossless sampling schemes usually incur substantially larger communication costs and sampling complexities, often rendering those methods impractical. We believe that the overhead required for lossless compression compared to the lossy variants used in our methods are insufficiently motivated in a federated learning context, since the variance of the training process likely outweighs the negative effects of the noise of the sampling process conducted for compression. If the reviewer deems it useful, we are happy to add a respective discussion in the paper.

[R1] L. Theis and N. Yosri, “Algorithms for the communication of samples,” in International Conference on Machine Learning, 2022, pp. 21 308–21 328.

Dear all,

Thank you for your clarification, and I appreciate your work. I do not have further questions.

I have also read the comments (especially the weaknesses) from other reviewers. I agree that larger-scale experiments are preferable, and that the availability and practicality of shared randomness might be a weakness (which the authors have addressed in the paper and responded to in the rebuttal). Adding some discussion on privacy and fairness would also be beneficial. From my personal perspective, introducing remote source generation to FL and demonstrating academic-level empirical validation are already meaningful contributions. I also find the writing quality of this paper to be good. I intend to keep my positive rating.

Best regards,

Reviewer Bbin

We thank Reviewer GJUK for their positive evaluation of our work. We address their concerns as follows.

• W1 (Partial client participation) Although we do not consider partial client participation in the paper, our methods are compatible with such extensions. For BiCompFL-PR, partial client participation can readily be implemented on top of the methods we propose, as they do not rely on full prior synchronization among the clients. BiCompFL-GR on the other hand requires that all clients' priors are synchronized at each iteration. To allow partial client participation, the federator is required to transmit the previous global model to the clients that were absent in the previous iteration. This ensures full synchronization of the global model estimate used for communication-efficient uplink communication. Hence, while partial client participation is supported by BiCompFL-PR without additional communication overhead, BiCompFL-GR's compatibility with partial client participation comes with a potentially increased downlink communication cost for clients that were previously passive. If the reviewer deems it helpful, we will be happy to add a discussion on this in the paper.
• W2 (Practical strengths and justification of small number of clients): The number of clients is not an inherent limitation in our methods. While we experimentally examined our methods with at most 50 clients, our methods exhibit favorable scaling properties in the number of clients. The more clients, the better the averaging effects of the noise introduced by the sampling operations for uplink compression, effectively reducing the error made in estimating the clients' average posterior. We selected the number of clients used in the experiments in alignment with common studies in the literature and our available resources.
• W3 (Practical tips for the choice of quantization intervals): In general, the choice of the number of quantization intervals depends on the model architecture and dataset complexities, affecting the number of iterations required until convergence. The simpler the learning task, the fewer quantization intervals are sufficient for convergence. In addition to generic arguments, we particularly want to highlight that the number of quantization levels should be carefully selected respecting the expected variance of the gradients. If gradients exhibit inherently large variance, e.g., through small mini-batch sizes, large variance of the local datasets, or substantially over-parameterized networks, small number of quantization intervals might be particularly efficient, and increasing that number would only negligibly improve the performance. We also want to remark that one-bit quantization was often shown to be trivially effective [R1, R2]. We hope these explanations provide a glimpse on the problem of quantization interval selection.
• Formatting: We thank the reviewer for this comment and will unify the citation style accordingly.

[R1] F. Seide, H. Fu, J. Droppo, G. Li, and D. Yu, “1-bit stochastic gradient descent and its application to data-parallel distributed training of speech DNNs,” in Interspeech, 2014, pp. 1058–1062.

[R2] S. P. Karimireddy, Q. Rebjock, S. Stich, and M. Jaggi, “Error feedback fixes SignSGD and other gradient compression schemes,” in International Conference on Machine Learning, vol. 97, 2019, pp. 3252–3261.

We thank Reviewer oHgr for their comments. We address the reviewer's concerns as follows.

• W1 (Clarity issues): We are unsure which parts are referred to as overly dense. If the reviewer has particular pointers, we are happy to modify the passages accordingly.

• W2/Q3 (Algorithm workflow clarity): We are happy to add such a conceptual diagram to the paper highlighting the differences between the two proposed methods, which lie mostly in the way shared randomness is leveraged.

• W3/Q1 (Scalability validation to model architectures): Our work proposes a novel method within the scope of academic research. Given standard resources, we conducted experiments with complexities that are comparable to all related works in this line of research.Given the numerical and theoretical evidence, we could not identify any particular negative effects as the task complexities scale; on the contrary, the positive effects of substantially reduced communication complexity gain increasing importance. The communication costs are quantified both in theory and experiments for varying model sizes and levels of heterogeneity. We kindly refer the reviewer to the respective sections and various experiments, also in the appendix.

• W4 (Runtime and energy costs): Like other papers in this line of work, we propose a conceptual framework that uses standard models and building blocks, such as sampling from parametric distributions. Many works are particularly concerned with resource-efficient implementations of such generic building blocks, e.g., the sampling operations. To complement our experiments, we theoretically analyze the fundamental sampling complexity of our compression methods. However, implementing and deploying any method on hardware in practice requires additional engineering efforts to make it efficient. Thus, measuring the energy or runtime of our methods on our servers would not yield meaningful comparisons, unless tailoring each method individually to the specific hardware at hand. We, hence, refrain from showing such plots, noting that the results would be orthogonal to the message of our paper. Our experiments, which include 50 clientsm align well with the settings considered in the literature. Importantly, we substantially benefit from increased number of clients on the uplink through positive averaging effects of the noise introduced in the sampling operations. Hence, the averaged posterior estimates observed by the federator improve as the number of clients increases. We thus believe that our ablation studies provide sufficient evidence of the positive scaling properties of our method.

• W5/Q2 (Generalization of theoretical guarantees): We deliberately choose to conduct a tailored analysis for Bernoulli distributions. While it is possible to run a more general analysis for other classes of distributions (as detailed in the sequel), we found that the specific optimization method proposed in the paper is particularly communication-efficient and allows for inherent communication-regularized training, a unique property in this line of work that renders our methods substantially more communication-efficient while preserving state-of-the-art performance.

Nonetheless, our analysis can indeed be extended to other parametric distributions, such as multi-variate Gaussians. The key ingredient is to replace the bias bound in Lemma 2 with the generic results for MRC known from [R1]. Although weaker than our result, their result holds for all classes of distributions. Having such an upper bound on the bias, one can follow our proof strategy to prove the communication costs in Theorem 1, i.e., using the convexity of KL divergence and decomposing the error in parameter estimation into a bias term (where the result from [R1] generically applies) and employing concentration bounds, such as the Hoeffding inequality for subGaussian random variables, to bound the sampling error. The remaining steps follow analogously to our proof. As for the contraction property in Lemma 1 used to establish convergence guarantees for the case of conventional FL, similar adaptations can be made. However, we note that for the composition of stochastic quantization with MRC to turn conventional FL into a stochastic setting, the incentive to use distributions other than Bernoulli might be weaker than for the natively stochastic settings mentioned by the reviewer. We are happy to include a discussion related to this in the paper.

• W6/Q4 (Privacy and fairness implications): Although privacy and fairness are issues widely studied in the literature, this work does not primarily focus on those aspects. For instance, privacy in FL with relative entropy coding was considered in [R2]. On the other hand, as a general statement, compression methods positively affect the clients' privacy in FL, which mainly stems from the fundamental data processing inequality [R3]. In our case, the noise introduced throughout the sampling operations used for compression dilutes the individual contributions of the clients observed by the federator, thus naturally enhancing the clients' privacy and reducing the risk of privacy breaches through, e.g., membership inference attacks. Concerning fairness, we note that adaptive block allocation does not have implications on potential bias w.r.t. particular clients in the aggregation process. The same amount of samples are observed for each client, ensuring balance in the aggregation process. Further testing the particularities of variance with regard to demographic groups is beyond the scope of this work. However, we are happy to add a respective discussion in our paper.

• W7 (Convergence instability under high heterogeneity): We kindly point the reviewer to our extensive experiments in the appendix, where we examine the stability of our methods with high data heterogeneity modeled by a Dirichlet with parameter $\alpha=0.1$, which is considered an extremely heterogeneous regime [R4].

• W8 (Overhead of adaptive block allocation): All our experiments consider the quantitative cost of adaptive block allocation incurred to transmit the block locations. We find that the overhead for transmitting block indices is negligible compared to the positive effects of block size adaptation on the communication cost. As the training progresses, adaptive partitioning successively increases the block sizes, thus drastically reducing the communication costs over the runtime of the algorithm.

[R1] S. Chatterjee and P. Diaconis, “The sample size required in importance sampling,” The Annals of Applied Probability, vol. 28, no. 2, pp. 1099–1135, 2018.

[R2] A. Triastcyn, M. Reisser, and C. Louizos, DP-REC: Private and communication-efficient federated learning, 2022.

[R3] T. Cover and J. A. Thomas, Elements of information theory. Wiley-Interscience, 2006.

[R4] Y. Allouah, S. Farhadkhani, R. Guerraoui, N. Gupta, R. Pinot, and J. Stephan, “Fixing by mixing: A recipe for optimal Byzantine ML under heterogeneity,” in International Conference on Artificial Intelligence and Statistics, 2023, pp. 1232–1300.

We thank Reviewer EDwN for their thorough review and helpful suggestions. We address their concerns as follows.

• W1/Q1 (Empirical validation limited to small vision benchmarks): Our work proposes a novel method within the scope of academic research. Given standard resources, we conducted experiments with complexities that are comparable to all related works in this line of research. Moreover, through our extensive numerical and theoretical results, we did not identify any particular negative effects as the task complexities scale; on the contrary, the positive effects of substantially reduced communication complexities gain increasing importance with larger models. % Nonetheless, we are currently running larger-scale experiments for follow-up works.

• W2 (MRC sampling complexity scales exponentially in the KL-divergence): As the theory suggests, the sampling complexity of MRC scales exponentially in the KL-divergence between prior and posterior. To overcome this limitation, we employ the method of block allocation. Therefore, we fix a per-block target KL-divergence and partition the model parameters into blocks to match the desired target divergence. We further highlight that the KL-divergence is not static and decreases as the training progresses, often leading to negligible divergences when advancing in the training. As a result, higher-dimensional models actually support meeting the KL-conditions, as larger dimensionality increases the degrees of freedom for block partitioning. With regard to the scaling properties, partitioning the models into blocks of equal KL-divergences leads to a linear dependence of the sampling complexity on the model size, alleviating concerns regarding exploding sampling costs.

• W3 (Shared randomness difficult to maintain in real-world FL): We clearly articulate the assumption that common randomness is required, and specifically propose a variant for cases where global randomness among all clients might be difficult to maintain. Nonetheless, the fact that our methods are fully synchronous facilitate easy piggy-backing of reasonably large seeds onto global model updates sent to the clients. This relative cost of this overhead diminishes, and becomes negligible, as the model size increases. In our experiments, the one-time transmission of few-bit seeds is sufficient to achieve state-of-the-art performance with substantially reduced communication costs.

• Q2 (Runtime and energy costs): Similar to other papers in this line of research, we propose a conceptual framework that uses standard models and building blocks, such as sampling from parametric distributions. Many works are particularly concerned with resource-efficient implementations of such generic building blocks, e.g., the sampling operations. To complement our experiments, we theoretically analyze the fundamental sampling complexity of our compression methods. However, naturally, implementing and deploying any method on hardware in practice requires additional engineering efforts to make it efficient. Thus, measuring the energy or runtime of our methods on our servers would not yield meaningful comparisons, unless tailoring each method individually to the specific hardware at hand. We, hence, refrain from showing such plots, noting that such statements would be orthogonal to the message of our paper.

• Q3 (Exploding KL gap for extreme heterogeneity): This is a very good point brought up by the reviewer, which we carefully investigated (although not discussed much in the paper). Firstly, we do conduct our experiments in regimes considered extremely heterogeneous in the literature (Dirichlet distribution with $\alpha=0.1$), e.g., [R1], and did not observe exploding KL gaps in such settings. We initially assumed such an effect would be present and studied whether it can actually be leveraged for improved system design. Assuming the clients' posteriors are moving substantially away from the global model when data is very heterogeneous, communication costs with private randomness would benefit substantially from using the previous posterior of the clients as prior instead of the previous global model. However, we couldn't observe notable positive effects from such individual prior selection. The key reason is the implicit regularization of our optimization methods. Specifically, the mirror descent approach studied in our paper naturally regularizes the posterior optimization w.r.t. the KL-divergence between prior and posterior, avoiding KL blow-ups in the optimization. This positively affects both the variance in local optimization in the clients and the incurred communication costs.

• Q4 (Compatibility with sparsification or structured pruning-based compression): Yes, our methods are entirely compatible with other compression methods that sparsify or prune the models/gradients before transmission. Applied prior to our compression methods, sparsified models can be readily fed into our compression methods without further adaptations. The pruned parameters/distribution are simply not included in the block selection of model parameters, thereby further reducing the incurred communication costs. We are happy to add a respective remark in the paper.

[R1] Y. Allouah, S. Farhadkhani, R. Guerraoui, N. Gupta, R. Pinot, and J. Stephan, “Fixing by mixing: A recipe for optimal Byzantine ml under heterogeneity,” in International Conference on Artificial Intelligence and Statistics, 2023, pp. 1232–1300.