Learning Communication-Efficient Optimizers

Major:

1. The experiments are conducted on toy MLPs and CNNs, which cannot convince me of the usability of the proposed methods. Communication-efficient optimizers are typically deployed in two main contexts: datacenter distributed training and federated learning. In the former, the scale of the models leads to communication lags due to the vast volume of data exchanged. In the latter, although the models might not be that large, optimization is complicated by factors such as data heterogeneity and partial client participation. The experiments in this study do not resonate with the complexities of either setting, even without testing on a small ResNet. When using such toy models, the need for communication-efficient optimizers becomes questionable.

2. Insufficient Insight into the Methodology. Even if we assume that the proposed method performs well in practical settings, this paper does not provide insight into its success. Given that the learned optimizer can perform complex operations and that local updates add to the complexity., it's crucial to delve deeper and find out the key contributing factors.

3. The authors are urged to improve their writing. I struggled to grasp the authors' intentions during reading. For example, the description of the proposed methodology in Section 3 is unclear. The text preceding Algorithm 1 does not reference the 'adafactor feature' $A_t$, yet it unexpectedly appears within the algorithm. This oversight can lead to confusion. The terms 'adafactor feature' and 'learned optimizer state' are also vague, requiring more illustration. The loss summation in eq (1) references the step $t$, but its absence in the summand is perplexing. In the experimental section, the x-axis labels of the figures should more accurately read 'the number of communication rounds' rather than 'training iterations' .

Minor:

1. Missing reference: Lv, Kaifeng, Shunhua Jiang, and Jian Li. "Learning gradient descent: Better generalization and longer horizons." International Conference on Machine Learning. PMLR, 2017.

2. Grammatical mistakes/typos:

   • Section 1, lines 3-4: "the communication overhead associated with distributed SGD can often lead to inefficient use of computing resources and reduced wall clock times (Lin et al., 2018)". reduced->increased
   • Section 3, line 3: "We take H local step of SGD", step->steps

   The authors should go through their paper carefully to correct all such mistakes.

1. It is unclear whether the trained optimizers can generalize to different local update hyperparameters other than the number of inner-loop iterations $H$, e.g. the local learning rate $\gamma$, the local batch size $\mathcal{B}_{loc}$.
2. The optimizer learned by LAgg-A cannot generalize to settings with different number of workers, because the input size $K+38$ is correlated with the number of workers $K$.
3. The train/test problems for the learned optimizers are mainly shallow networks on image classification tasks. The lack of diversity in test problems may not justify well on the generalization properties of learned optimizers.

1. Regarding the chosen architecture in the experiments: the studies employed basic models such as MLP and small-scale CNNs as their neural architectures. I think these models are overly simplistic, raising doubts about the applicability of the proposed method to more complex, widely discussed architectures like ResNet and Transformer-based models.

2. Concerning the generalization capability of the proposed optimizer: while it is asserted in the abstract that the optimizer shows generalization, it often falls short in this aspect. For example, the optimizer diverges in figures 6d and 6f. This raises concerns about the method's practicality since retraining the optimizer for each new problem is time-consuming.

3. The presentation of experimental outcomes is somewhat unclear. Specifically, in figures 1–5, the x-axis represents training iterations, which equals the number of communication iterations. This seems unfair to SGD and ADAM, as algorithms based on local SGD undertake more updates per communication. It would be more appropriate to display the count of local updates on the x-axis. Moreover, in real-world scenarios, communication time often overlaps with computation, making it less time-consuming. Hence, for most neural network structures, local updates remain the primary bottleneck.

4. In figures 2, 5, and 6, the proposed optimizer sometimes exhibits divergence. Incorporating a trust-region check to ensure convergence of the optimizer would be beneficial.

• The practicality of LAgg-A is called into question, particularly in the context of modern neural networks with numerous layers and millions/billions of parameters. The experimental application of LAgg-A is limited to simple 2/3-layer networks with thousands of parameters. Moreover, LAgg-A appears to be an architecture-specific optimizer although it shows some generalization abilities. This restricts its usability compared to widely applicable optimizers like SGD and Adam.

• The novelty of LAgg-A appears to be incremental, as it extends the small_fc_lopt optimizer proposed in [1] for distributed learning. The main innovation seems to involve substituting the first/second gradient momentum in [1] with local updates from multiple clients as input features for the learning optimizer MLP. A more substantial leap in innovation would be beneficial.