Meta-learning Optimizers for Communication-Efficient Learning
Authors: Charles-Étienne Joseph
Benjamin Thérien
Abhinav Moudgil
Boris Knyazev
Eugene Belilovsky
Venue: Accepted by TMLR ()

Abstract: Communication-efficient variants of SGD, specifically local SGD, have received a great deal of interest in recent years. These approaches compute multiple gradient steps locally on each worker, before averaging model parameters, helping relieve the critical communication bottleneck in distributed deep learning training. Although many variants of these approaches have been proposed, they can sometimes lag behind state-of-the-art adaptive optimizers for deep learning. In this work, we investigate if the recent progress in the emerging area of learned optimizers can potentially close this gap in homogeneous data and homogeneous device settings while remaining communication-efficient. Specifically, we meta-learn how to perform global updates given an update from local SGD iterations. Our results demonstrate that learned optimizers can substantially outperform local SGD and its sophisticated variants while maintaining their communication efficiency. Our learned optimizers can even generalize to unseen and much larger datasets and architectures, including ImageNet and ViTs, and to unseen modalities such as language modeling. We therefore show the potential of learned optimizers for improving communication-efficient distributed learning.

Published in Transactions on Machine Learning Research (03/2025)
Meta-learning Optimizers for Communication-Efiiienl Learning
Charles-Etiees Joseph'
charles-etienne, joseph @rmila, quebee benjamin. therien @mmila. quebec bolsinavmmoudgee Errrela quebee kneuzeeb @rrilla. quebec eugene. Beeene..bellossky@omoooode
@omila.
Benjamin Thérien"
Abhinav Moudgil
Boris Knyazev
Eugene Belilovsky
Deportment of Computer Science and Operation Research, Université de Montréal, Montréal, Canada
Deportment of Computer Science and Software Engineering, Concordia University, Montréal, Canada
Mila, Montréal, Canada Samsung SAIT AI Lab, Montrénl, Canada +
Reviewed on OpenReview: https: // openreview. net/ net/ forum? id- wRb f9ANARS
Abstract
Communication-effin -fffccienn variants of SGD, specifically local SGD, have received a great deal of interest recent years. These approaches compute multiple gradient steps locally on each worker, before averaging model parameters, helping relieve the critical communication bot- tleneck in distributed deep learning training. Although many variants of these approaches have been proposed, they sometimes lag behind state-of-the-art adaptive optimizers for deep learning. In this work, we investigate the recent progress in the emerging area of learned optimizers can potentially close this gap in homogeneous data and homogeneons device settings while remaining communication-effieen. Specifically, we meta-learn how to perform global updates given an update from local SGD iterations. Our results demon- strate that learned optimizers can substantially outperform local SGD and its sophisticated variants while maintaining their communication efficiency. Our learned optimizers can even generalize to unseen and much larger datascts and architectures, including ImageNet and ViTs, and to unseen modalities such as language modeling. We therefore show the potential of learned optimizers for improving communication distributed learning.
1 Introduction
Rapidly training large-scale deep learning models is problem of continued interest the community. It re- quires great deal of distributed computing resources that often challenging to efficiently utilize. In many distributed learning settings, the communication overhead associated with distributed SGD can lead to inef- ficient use of computing resources and increased wall clock times (Lin et al. 2018). This reliance on frequent communication is especially impractical for training large models over heterogeneous hardware (Yuan et al., 2022). Moreover, it can increase the cost and complexity designing data centers and other infrastructure to support the heavy communication constraints.
The primary communication overhead of distributed SGD comes from the synchronization of gradients computed by different workers. A recently popular direction alleviate this overhead is local SGD (Stich) 2019), where each worker computes multiple (H) gradient steps independently before aggregating the weights
'Equal contribution

---

Published in Transactions on Machine Learning Research (03/2025)
updated params
W I+1
Giobal learned optimizer
44+
eccumulators state
updated state
A
A
H
local update steps
Worker
Worker
Worker K
Model
Model
Model
1 W
current params
Figure 1: In local SGD, workers take H local update steps (i.e., without communicating gradients) of SGD before communicatigg local parameter deltas (A*). This effectively reduces the number of communccation steps by a factor H. Instead of averaging deltas at communication steps, we meta-train a global learned optimizer aggregate the deltas into more effective update.
(or deltas A of their local models (fig. 1). This reduces the communication costs. There is an algorithmic connection here to federated averaging McMahan et al. (2016), an analogous algorithm in the federated learning setting. The only difference bet the two are assumptions about the workers computations capabilities, the data, and worker participation. We consider the setting homogeneous high-resource workers, homogencous data split among the workers, and full worker participation. We, therefore, refer Local SGD (Stich 2019) the main prior work upon which we build they also consider this setting.
Local SGD, however, has number of challenges limiting its practical use. As the number of local steps H increases the local models may diverge from each other leading to degradation of performance (Wang et al. 2019 Local SGD also introduces a complex dynamic between the local and global updates, which can for example lead to complex interactions between hyperparametens such as global and local learning rates [Reddi et al. 2020).
Learned optimization through meta-learning has been an increasingly important topic of research inter- est (Andrychowicz et al. 2016). Advances have been made in scalable architectures (Wichrowska et al. 2017 Metz et al.] (002aa), meta-learning strategies (Vicol et al. 2021 and the diversity and scale of meta- learning tasks (Metz et al. 2022b). Notably, Metz et al. (2022a) analyzed different learned optimizers in a large-scale study and introduced a highly efficient and simple per-parameter MLP optimizer and strong gradient- basod features. However, these recent learned optimizors have not been studied in a communication- efficient distributed selting.
In this work, we propose learned optimization as an approach to alleviate the challenges of communication- efficient distributed learning. Specifically, follow the setup of local SGD Stich (2019) with homogeneous devices and homogeneous data split among them and demonstrate that our global learned optimizers (fig. I meta-trained for this setting can outperform Local SGD and SlowMo (Wang et al. 2019 as well as data- parallel Adam and SGD. Our main contributions are:
We demonstrate, for the first time, that learned optimizers can be used to improve local SGD for communication-eficionn distributed learning, outperforming strong basclines and maintaining benefits even for high number local steps.
We propose and evaluate two architectures for the learned optimization of local SGD, worker- aware optimizer (LAgg-A) and a depending the use-case.
optimizer (LOpt-A), from which one can choose
2

---

Published in Transactions on Machine Learning Research (03/2025)
We demonstrate that our learned optimizers, even when meta-learned on single few architecture and dataset combinations, can generalize to new and much larger datasets and architectures, includ- ing ImageNet, ResNets, Vision Transformers (ViIs), and new modalities such as language modeling, obtaining competitive results in commuiiationn-fficeet distributed settings.
2 Related Work
2.1
Local SGD and Communication--fficeeee DL
Local SGD has been analyzed in number of works (Stich 2019 Lin et al. 2018 which demonstrated that it both theoretically and empirically can lead to communication savings. It has also been shown that local SGD, particularly when combined with phases regular SGD, can lead to better generalization (Lin et al. 2018 depending on the task scale (Ortiz et al. 2021).
Wang al. 2019 introduced SlowMo using global or server-side momentum and showed that can ac- number of decentralized and asynchronous stochastic algorithms. A closely related algorithm has been proposed and extensively used in federated learning for communication effi- ciency (McMahan et al. 2017: Li et. al. 2019). Work in this field has largely focused on addressing the heterogeneity of data across workers or clients (Karimireddy et al.; 2020 Mishchenko et al., 2022). Most of the aforementioned algorithms struggle to handle large numbers of local steps due to client drift, even in the non-federated setting. To address this, Douillard et al. (2024) proposes to use AdamW as the client side optimizer and Nesterov momentum to aggregate the deltas, allowing scaling to 500 local steps. These advancements are generally achieved by hand-designed algorithmic enhancements, whereas our approach relics more flexible and potentially more powerful learnable mechanisms that may generalize these and more complex algorithms.
celerate local SGD well
Another approach to communication-eficiin learning is to compress the gradients or parameters. Two pop- ular strategics in this setting are sparsification (Stich et al. 2018 Shi et al., 2019) and quantization (Alistarh et al. 2017) of the gradient. These stralegies have also been combined by Wang el al. (2023). This line of work thus orthogonal but complementary to our proposal. Communication efficiency has also been studied
in the decentralized setting (Nabli & Oyallon 2022 Nabli et al. 2023 Lian et al.) 2018). Our work focuses the centralized training setting but the methods also be extended to decentralized training.
2.2 Learning to Optimize (L20)
The idea of learning learn and meta-learning has a long history (Schmidhuber 1992 Thrun & Pratt 2012). Many early works in this area focused on learning to efficiently acquire general knowledge or inductive bias. Hochreiter et al. (2001) proposed to use meta-learning in direct combination with gradient-based optimization to learn separate network, which can be seen as a learned optimizer, which performs updates another network. Andrychowicz et al. (2016) extended these ideas to a more scalable LSTM-based per-parameter architecture and demonstrated that the learned optimizer can generalize to new problems.
A large number of follow up works have improved I.20 methods (Wichrowska et al., 2017 Metx et al. 2019 Chen et al.] 22200; Metz et al., 2020 Harrison al. 2022: Lv et al. (017)) (sce Chen et al. (2022): Amos (2022) for surveys). These methods introduced different types of hierarchy into the learnable optimizer while simplifying its architecture in favor of stronger predefined features to improve its efficiency (Metz et al, 2022a). However, compared to our work, these have not considered distributed setting, where learnable optimizers may significantly improve local SGD which is challenging to combine with adaptive optimizers.
Ji et al. (2019) proposed to learn the aggregation gradients from workers in a distributed learning frame- work with a recurrent network. However, the focus was on improving non-local SGD while our work focuses on the communicttin efficiency in settings where each worker returns a message computed from multiple update steps. Furthermore, our approach is shown to generalize to new architectures and datasets.
3

---

Published in Transactions on Machine Learning Research (03/2025)
3
Background
3.1
Local SGD
We consider distributed training setup with K workers. In local SGD (Stich 2019 round 1 consists of H local SGD steps performed independently all K workers. Each local step h uses minibatch of size Bise After completing the local updates, global update is computed by averaging the local weight deltas and subtracting the average from previous weights, as shown on line 9 of algorithm
each communication
Algorithm 1: Learned optimizers vs Local SGD. Steps used both algorithms are not colored.
Input:
T
Number of communication steps
K
Number of workers
H
Number of local steps
Local learning rate
Wo.
Initial weights
D
Dataset
C
Loss function
F
Learned optimizer
Uo
Initial accumulators state
1:
for t 0 to - do
2:
for k = to K - in parallel do
3:
for 1= to - do
X (k) + GET (OTT_MIIIATCCC()
4:
(-) GA++1 + ((h)) t.
V W (x,, (()),, w((h)
5:
1.
end for
A
W
(k) 1.0 HH
(k)
w(()
// Difference in weights after H local steps
8:
end for
A + K Σ A
(k)
9:
10:
A,, U++1 + ADA 2.,, U, A
// Compute Ada features and update state
LAgg-A A
411.1
W4L10 + F6 AA, A,
(0, ...... -...
11:
LOpt-A
(41.2
WHL00
Fo (A,, A.)
// Compute global update
Local SGD
Stich
2019
W+L00 Weo A.
12: end for
3.2 Learned Optimizer Input: Ada Features
In the learning-to-optinngg and optimization literature, it is common to maintain accumulators of gradient (or in our case A,) statistics. For instance, Adam (Kingma & Ba 2017 maintains per-pane cor-parameter accumulators for the first and second moment of the gradient and Adafactor (Shazeer & Stern 2018) maintains column- wise and row-wise sums of these moments. Although it rarely donc for hand-designed optimizers, is possible to leverage more information by maintaining multiple (> 2) accunulators of these statistics different time scales (e.g. multiple momentums with different coefficients.. In fact, recent work (Metz et al, 2022a) develops set of foatures based on this idea for training learned optimizers. Drawing inspiration from existing adaptive optimizers that often maintain accumulators similar to we henceforth refer to these features as " Ada features" in the absence of prior name.
We will now describe how we adapt Ada features the local SGD setting. Ada features maintain three different per-paree paramaraetter momentum accumulators (maa) and one variance accumulator (v)). In addition, they also maintain six accumulators of the column-wise (ct,s) and row-wise (re,)) mean of the gradient. The
4

---

Published in Transactions on Machine Learning Research (03/2025)
accumulator update is given as follows:
my,i = 8.14-1. + (1 -
i (1,2,3)
Dt = 81% + (1 - B4)
Tt, = 8,7-1111 + (1 -3)) ROW ROW_REEAN NOW_MEAN(A?),
iE (5,6,
Ct2 8,4-1,1 + (1) COL_MEAN(4
i (5,6,7).
U, :- [m,,1 m4,2, M4,3, Ut, 11,5, Tr,6, Tt,7, Cr,5, Cr,6, c.7l
(1)
Here, we slightly abuse notation and define U, be the entire accumulator state for all parameters in the optimizee Coolmmnwwise and wow- features are repeated for notational convenience). Note that Metz et al. (2022a) use V. instead of A, since they train learned optimizers in fully-centralized setting. We propose using 4e, the average local weight delta, instead of the gradient in our local SGD setting. After updating these accumulators at each communication step, one must compute additional learned optimizer input features to concatenated with the accumulator values (U) and an 11-dimensional timestep embedding (T). The Icarned optimizer input-fcature matrix A, created as follows:
T. = [tanh for x E {1, 10, 30, 100, 300, 1000, 3000, 10000, 30000, 100000)];
1
1
m fo
ML3
R2 VTL.5 VTL. VTL, Ver.5 Vas VeuT
H. (1,17,,50,,5 14,57,6645 7,,77, Am ArrE,CC, Am177))
A: = W.O.G 0 T H.O.R R.
Where @ denotes matrix concatenation across the feature dimension, T. are time embeddings, Rs are recip- rocal features, H, are adafactor normalized features, and W, are the current parameters of the optimizee. In algorithm 1 A, is computed as follows:
A,, U1+1 ADA W. A.).
(2)
Note that the complete matrix of input Ada features MMtzz et al. 2022a) is given by A. 3 Ar, however, we leave out A in the equations above to make explicit the differences of our proposed optimizers with respect to A. We provide further details about these features in section A of the appondix.
4 Methodology
Our method builds upon local SGD. After H local steps, we employ per-parameter learned optimizer F based on (Metz et al. 2022a to compute the updated centralized weights (algorithm D By computing the centralized update using a small MLP, Fş: that is much more expressive than hand-designed updates, our method can be seen as a generalization of existing update methods such as taking the average iterate (Stich 2019 or computing server-side momentum updates (Wang et al. 2019).
4.1
Learned Optimizer Training and Architectures
We consider the meta-learning framework with learned optimizer Fg parameterized by & used to optimize model with parameters W. In the meta-learning formulation, is obtained by solving the following optimization problem:
T-IK-1 TKH ΣΣΣ
((X.Y: w/ h=0
min E(p,C.Wo--T E(XYn>
(3)
For simplicity we remove the subscripts inside the sum term, but note that the exact value of the summands does depend and h. Here, T is distribution optimization tasks defined tuples of dataset D.
5

---

Published in Transactions on Machine Learning Research (03/2025)
objective function C, and initial weights Wo associated with a particular neural architecture, ø represents the parameters of the learned optimizer, H is the number of local steps, K is the number of workers, and T is the number of communication steps which we write as fixed quantity for simplicity. In practice, during meta-training, is varied according to a truncation schedule (Metz et al. 2022a]. We provide extended
details of the mela-training process for our global learned optimizers in section B of the appendix.
In our experiments, Fo is a two hidden layer 32 hidden dimension MLP with ReLU activations mapping the imput Ada features for each parameter, p, in the optimizee to two-dimensional vector, mpl. At step t, the learned optimizer update for all p is given as follows:
AAp C [Atp)) = m.
Pt Pt-1 X A (2mo.y)
(4)
Where Ap are the ada features computed from statistics of p and \ and X2 are constants set to 0.001. We drop the timestep subscript for some values of to avoid making the notation cumbersome. We propose two variants of our global learned optimizers, LAgg-A A and LOpt-A (algorithm D LAgg-A takes advantage of individual deltas from all the workers and so can learn better optimizers when the number of workers is known and fixed beforehand. LOpt-A operates on the averaged delta, thus it is more versatile as can be applied the setting with an arbitrary number of workers, however, it can be less powerful than LAgg-A in certain cases show empirically.
4.1.1 Worker-aware Optimizer (LAgg-A)
Our first learned optimizer lakes advantage of pre-aggregated information from each worker. Specifically, it takes as input A (1) ...,, (Ath) along with the Ada features computed from A, (the average of A ..... A We refer to it as a learned aggregator (LAgg-A) as it learns to aggregate the workers' weight updates. With its access to pre-aggregated information, LAgg-A can learn complex interactions between workers potentially making more powerful weight updates. However, it requires fixing the number of workers K before training, which in our experience not an essential problem because oftentimes the distributed training assumes some standard fixed budget of workers.
4.1.2 Worker-invariant Optimizer (LOpt-A)
Our second proposed learned optimizer directly takes A., the average of the updates from workers, as an input feature along with the Ada features computed from it. This process analogous to existing learned optimization proposed by Metz et al. (2022a) where the role of the gradient is replaced with A,. The advantage of LOpt-A versus LAgg-A A is that it has the same number of parameters (løl) regardless of the number of workers K. This can be usefull when the same learned optimizer is applied for settings with variable K. However, this approach is less powerful as it cannot take advantage of individual deltas from all the workers.
4.2 Practical Considerations and LOpt-A and LAgg-A A Overhead
As discussed by Reddi et al. (2020) the class of local algo can be described with a server-side optimizer and worker-side optimizer. For example, SlowMo (Wang et al.) 2019 can be interpreted as adding momentum to the server optimization. Our design of algorithm
such that the learned optimizer lives entirely on the server-side, making its use more practical and scalable than in non-commmiicatonnccciioii settings.
Specifically, standard learned optimizers have overhead of memory and compute. The memory must store state information and intermediate activations of the learned optimizer. In the case of our learned optimizer, this overhead (Metz et al. 2022a) is only incurred at the aggregation stage and can therefore live entirely on the server if one is available. Similarly, while the computations cost of the forward pass of learned optimizers provides a substantial overhead compared to simple add and multiply operations of SGD and Adam, the case of our global learned optimizers this cost becomes small with respect to the large amount of data processed on workers during local updates. For instance, table [ we show that the overhead of our optimizers relative to Local SGD shrinks the number of local steps grows. The figure reports total timings

---

Published in Transactions on Machine Learning Research (03/2025)
for the forward and backward pass and optimizer step of Local SGD, LOpt-A, and LAgg-A for training ResNet50 (a) and ViT-Base (b) on 224 and 128 sized ImageNet, respectively. All timings are measured when training across K=8 A6000 GPUs. We further expand on this point in section of the appendix.
fwbw + Opt. Step Time (s)
fwbw + Opt. Step Time (s) Local SGD LOpt-A LAgg-A A
Local Steps
Local SGD LOpt-A LAgg-A
Local Steps
H
0.58
0.64
0.68
8
0.57
0.65
0.71
overhead
1.00
1.11
1.17
overhead
1.00
1.14
1.25
H-16
2.05
2.13
2.16
H-16
2.11
2.19
2.25
overhead
1.00
1.04
1.05
overhead
1.00
1.04
1.07
H-32
3.93
4.02
4.04
H-32
4.19
4.25
4.32
overhead
1.00
1.02
1.03
overhead
1.00
1.01
1.03
H-64
7.71
7.83
7.86
H-64
8.51
8.60
8.66
overhead
1.00
1.02
1.02
overhead
1.00
1.01
1.02
(a) ViT-Base ImageNet-128
(b) ResNet50 ImageNet-222
Table 1: The overhead of our optimizers relative to Local SGD shrinks to ZERO as the number of local steps grows. We report total timings for the forward and backward pass and optimizer step of Local SGD, LOpt-A, and LAgg-A for training ResNet50 (a) and ViT-Base (b) on 224 and 128 sized ImageNet, respectively. All timings are measured when training across K=8 A6000 GPUs. All reported values are the median time seconds of the final 40 training steps across different training restarts. We provide more details in section C of the appendix.
5 Experiments
Learned optimizers are a relatively recent area and the experiments are usually run small-scale datasets due to the challenges of meta-training and applying learned optimizers (Metz et al. 2022a).. However, local
SGD and its variants are typically studied a large-scale distributed sctup (Wang et al. 2019). There- fore, compared to the previous learned optimizers literature, we not only perform sinall-scale experiments but also experiment with larger and stronger architectures such as ViI',, including larger datasets such ImageNet (Russakovsky et al. 2015) and more modalities such as language modeling (LM1B (Chelba et al. 2013 All our experiments in the main manuscript assume distributed setting with homogeneous devices and homogeneous data split among them. Following the convention in learned optimization [Metz et al. 2022a 2019); Harrison et al. (2022), we mainly report training loss for simplicity comparing among differ- ent optimizers. However, we do demonstrate that models trained by our global learned optimizers compare favorably on held-out data to models trained with other optimizers (see Figures and G].
5.1 Experimental Details
In the following two sections, we detail the training and evaluation tasks (optimizees) and the optimizers that we compare. We note that our experiments use standard datasets and evaluation protocols learned optimization (Metz al. (202aa). Our method currently implemented in simulation. We meta-train and evaluate using NVIDIA A100. For the presented results, each curve is an average over 10 trials with different seeds. Shaded rogions represent standard error from the mean.
5.1.1 Datasets
We use the Fashion MNIST (FMNIST) dataset (iiao et al. 2017) (10 classes) with 28 images. We also use the CIFAR-10 datasct (Krizhevsky al.] (009) (10 classes) with 32 x images. Finally, we scale our setup the ImageNet dataset (Russakovsky al. 2015) (1000 classes) with downsampled 32 32 and 64 x 64 images. We designate the dataset ImageNet when the larger images are used. For the language modeling task, we use LM1B (Chelba et al. 2013)

---

Published in Transactions on Machine Learning Research (03/2025)
'Final Loss/Acc.
7.00
0.06
SSD B B-Bmc
Test Test Accurary
(K=1 H=1)
Train Loss
6,50
0.04
SGD
ADAM
0.03
Local SGD
0.02
SlowMo
5.75
0.01
LOpt-A
0.00
LAgg-A
200
400
600
800
1000
200
400
600
800
1000
Communication Steps
Communication Steps
(a) MLP ImageNet (In Distribution)
(b) MLP ImageNet (In Distribution)
0.15
Test Accuracy
Train Loss
0.10
0.05
0.00
0
200
400
800
800
1000
200
400
600
800
1000
Communication Steps
Communication Steps
(c) ResNet50, ImageNet* (OoD)
(d) ResNet50, ImageNet" (OoD)
7.0
0.12
6.5
Test Accuracy
0.10
Train Loss
0.08
6.0
0.06
5.5
0.04
5.0
4.5
0.00
200
800
1000
200
Communication Steps
Communication Steps
(e) ViT, ImageNet* (OoD)
(1) ViT, ImageNet* (OoD)
7.0
7.0
6.5
6.5
Train Loss
6.0
Test Loss
6.0
5.5
5.5
5.0
5.0
4.5
4.5
4.0
200
400
600
800
1000
200
400
600
800
1000
Communication Steps
Communication Steps
(g) Decoder-only, LMIB (OoD)
(h) Decoder-only, LMIB (OoD)
Figure 2: Meta- Gettaaaennrlliitiiionn to ResNee55, VIT, and a decoder-only language model. We report train loss and test performance for each task. Note that the amount of data seen for ImageNet training, K- H 1000 Bloe; corresponds to roughly 3.2 epochs, so the trained models do not reach convergence. In Figures ba,,), Our LAgg-A and LOpt-A optimizers, meta-trained on a 3-layer MLP (0.5M params) 32 x 32 TmageNet. (IN32) task for 1000 steps, outperform extensively tuned baselines on the in-distribution task. Moreover, these optimizers generalize to IN64 on ResNet50 (c,d) (50x larger) and VIT (e.f) (10x larger). Finally, we also show (g.h) that the optimizers are useful for training a decoder-only transformer language model (38x larger). The orange SGD baseline is trained for K H 1000 optimization steps with batch size Bioc 128. The dotted orange line corresponds to the final performance of this baseline.
8

---

Published in Transactions on Machine Learning Research (03/2025)
5.1.2 Neural Architectures
As for neural network architectures that our learned optimizers are going to optimize, we use multilayer perceptron (MLP) of two different sizes, both with ReLU activations. The first has two layers of 128 hidden nodes each and we refer to it as 2-Layer MLP. The second has three hidden layers of 128 hidden nodes each and we refer to it as 3-Layer MLP. We also use convolutional neural network (CNN) of 3 layers with ReLU activations. All 3 layers have convolution kernels of size 3 x 3 and use "same" padding. The first layer has 32 units and stride 2, while the two other layers have 64 units and stride 1. We refer to this architecture as CNN. We also use standard architectures such as ResNet55 (He et al. 2016), ViT equivalent in size to DeiTT tiny (Touvron et al. 2021). ) and for the language task a decoder-only transformer with hidden size 192, heads, and 12 layers.
5.1.3 Meta-training LOpt-A and LAgg-A
To meta-train our learned optimizers we estimate gradients using Persistent Evolution Strategies (PES) Vii col et al. 2021) and take gradient descent steps using AdamW and a linear warmup plus cosine decay schedule. Each gradient estimated from a batch of 8 task each unrolled to a specific number of steps T. T varies from 100 to 1000 during training according to log-uniform truncation schedule. In our exper- iments, gradients are estimated with respect to the optimizee's 's training loss, except for the curves in fig, [ whose gradients were estimated with respect to the optimizee's validation loss. During meta-training, the learning rate is warmed up for 100 steps to maximum learning rate before being decayed (following a cosine decay schedule) to 1/3 of the maximum value. All the mela-training details are provided in appendix B
Since our learned optimizers are meta-trained for 1000 steps, any evaluations beyond this training horizon is duttof-distribution. Therefore, we will evaluate our optimizers within their meta-training horizon of 1000 steps, which may not lead to convergence on all lasks, especially those requiring longer training (e.g.- TmageNet).
5.1.4 Non-local SGD Baselines
We follow the setup of SlowMo (Wang et al. 2019) and provide a comparison to non-local algorithms. To do s0, we train models using SGD Robbins 1951 and Adam (Kingma. & Ba 2017) for a number of steps equivalent to the total number of communication rounds used for the local methods. At each step, these baselines compute updates using the same elfective batch size K x H x Bioe the local optimizers they are compared to. We also include an SGD baseline that uses K H x 1000 steps and batch size Bioe- = 128 in select experiments (figures 2 and 10 The hyperparameters of these optimizers and details of their grid
search are provided in appendix D
5.1.5 Local SGD-based Baselines
Since our method focuses on improving server-side optimization and other client improving methods are orthogonal to our work, we provide two sufficient commuication-effcciontt distributed basclines: local SGD Sttch 2019) and SlowMo (Wang et al., 2019). An extensive hyper-parameter search is conducted for each baseline in every configuration. We detail the search process and report the best hyperparameters
in appendix
D For each task, we use a local batch size Bid of 128.
5.2 Evaluating LAgg-A and LOpt-A In-distribution
In this section, we evaluate our proposed optimizers on FMNIST, CIFAR-10, and ImageNet using H = iterations and K = 8 workers. Following the evaluation protocol of |Mtz et al. (2022a), in each case, we meta-train on task (dataset and architecture pair) and perform evaluation distribution evaluations test the generalization of the optimizer to a new initialization of the model and new ordering of the data. Results reported in table 2 show that our learned optimizers, when evaluated in-distribution, enjoy faster convergence than local SGD and consistently outperform SlowMo. Figure 2
new seed. That is, in
I Note that unless otherwise stated our experiments all the tasks in batch correspond to the same dataset and architecture, but different initial weights (see section section 4.1 for details).
9

---

Published in Transactions on Machine Learning Research (03/2025)
Final Loss
0.8
0.8
0.8
SGD limllee
Train Loss
Train Less
Loss
[K=1 H=])
SCDD
Local SGD
0.2
0.2
0.2
L0pt-A
LAgy-A
200
400
600
800
1000
200
400
600
850
1000
200
400
600
800
1000
Communication Steps
Communication Steps
Comm Commatiiaiion Steps
(a) =
(b) H
(c) H 16
Figure 3: LAgg-A A outperforms all optimizers for H E 84, , 16) local steps. All training curves are for FMNIST 2-Layer MLP. The orange SGD baseline is trained for K 1000 optimization steps with a batch size Bioc 128. The dotted orange line corresponds the final performance of this baseline.
(a.b) presents the training and accuracy curves for ImageNet 3-Tayer MLP. Figure 9 in appendix shows the training curves for FMNIST 2-Layer MLP (fig. (a)) and CIFAR-10 CNN (fig, 9b.. We observe that LAgg-A and LOpt-A consistently converge faster than all other baselines from the start of training. Note that Slow Mo is well-tuned and represents very competitive approach in the class of methods that perform local updates Waan et al. 2019)
Table 2: Speedup with respect to local SGD; reported as the ratio of the number of communicttions required by local SGD to the number of communications required by the other optimizer to achieve local SGD's minimum training loss (higher is better). In-distribution denotes that LOpt-A/LAgg-A are trained the same task the evaluation lask. For mela-generalizatioo,, LOpt-A/LAgg-A are trained ImageNet 3-layer MLP. A hyphen (-) indicates that the local SGD's minimum loss value was not achieved in the training run (1000 communication steps). When taking averages, hyphens are ignored.
Optimizer
In-distribution
Meta-genoralization
AVG
FMNIST
CIFAR1O MLP
ImageNet MLP
ImageNet ResNet50
ImageNet ViT
LMIB
MLP
Transformer
Local SGD
1.00
1.00
1.00
1.00
1.00
1.00
1.00
Adam
+.42
1.69 0.16
2.00 + 0.15
2.10 + 0.06
2.18 0.04
3.56 0.09
2.44
SlowMo
4.10 0.38
4.81
2.48
1.99 0.08
1.26 0.02
2.93
1.Opt-A
8.47 + 0.72
9.62 + 0.36
4.26 + 0.41
2.56 + 0.12
2.31 0.12
6.76 0.43
5.66
LAgg-A A
9.80 = 0.54
7.69 + 0.70
4.63 0.26
2.48 + 0.09
1.88 + 0.08
6.02 0.56
5.42
5.3
Effect of Local Iterations (H)
We now analyze our learned optimizers' capability to scale to larger number of local iterations (H). Specifically, vary H (4,8, 16) and meta-train learned optimizers the FMNIST 2-Layer MLP task for each case (note that generalization to different H is possible we show in fig, performance of corresponding tuned baselines with the equivalent batch size (fig. 3 We also show the communication efficiency compared to local SGD and SlowMo in table We observe that even for relatively high И (Lin et al. 2018) there is an improvement over the strong communication efficient baselines. As expected, table illustrates higher yields more rapid convergence on a per communication step basis (due more samples being processed). We also observe that LAgg-A begins show a substantial advantage compared to LOpt-A at this higher H value. We believe that using information from all the (k) allows for LAgg-A to learn a non-trivial aggregation scheme (compared to averaging), meaning it outperforms I.Opt-A when the local models drift as H gets higher.
7 We report the
In their recent work, Douillard al. (2024) proposes DiLoCo, a new algorithm for the low-commminatiion datacenter training of LLMss. They AdamW as the client side optimizer and Nesterov momentum to
10

---

Published in Transactions on Machine Learning Research (03/2025)
7.0
1.2
SGD
ADAM
1.0
Local SGD
Test Loss
SlowMo
LOpt-A
Test Loss
0.8
LOpt- A
6.2
0.6
LÄgg-A
6.0
0.4
0.2
5.4
0.0
200
400
600
800
1000
200
400
600
800
1000
Communication Steps
Communication Steps
(a) 3-yayer ML. ImageNet
(b) 2-Tayer ML.P FMNIST
Figure 4: Directly targeting validation loss during meta-training obtains strong performance the test set. Hand-designed optimizers were hyper-parameter-tundd to the validation set, while LAgg-A and A.OptAA were meta-trained to optimize validation loss on their respective tasks. We observe that learned optimizers trained to optimize validation loss during meta-training generalize seamlessly to the test set in our commmicatiion-fffci setting.
aggregate worker deltas, allowing scaling to H 2000 local steps. However, they show that the step-to-loss improvement from using more local steps diminishes for H 100. While experiments show that LAgg-A and 1.Opt-A can improve the efficiency of local SGD for up to H 16 steps, increasing H further will eventually become limited due to using SGD as the client-side optimizer. However, LAgg-A and LOpt-A are orthogonal to the choice of client-side optimizer and could casily combined with AdamW on the client side to enable training with very large H values at increased efficiency compared to DiL.Co.
Table Communication rounds until achieving 0.2 loss value for different optimizers at different H values (lower better).
Optimizer
H=4
H=8
H=16
Local SGD
721
625
SlowMo
311
182
121
LOpt-A
119
121
89
LAgg-A
122
81
55
5.4 Outer Loop Generalization
Following conventions in the learned optimization literature (Metz et al. 2022b a our focus in this work has been demonstrating the efficient convergence of the learned optimizer. Thus in our experiments, the outer loop of the meta-learning problem (see equation in section 4.1) evaluates the training data. In this section, we demonstrate that we can also obtain strong performance on the validation data using our learned optimizer. Figure 4 the test loss of learned optimizers meta-trained using the validation loss objective and baselines tuned using validation loss on 3-Layer MLP ImageNet (fig. 4ah and 2-Layer MLP FMNIST (fig. 4bb. We observe similar trends to training loss plots (figs. 3 and 10 In fig, 4a observe that both LAgg-A and LOpt-A converge significantly faster and obtain lower final test loss than the basclines. In fig. 4b LAgg-A A and LOpt-A converge faster than other baselines reaching a test loss around ileration 200 that baselines only reach after 600 iterations of training. While both plots show similar relative trends between our learned optimizers and the basclines, we note that they represent distinct scenarios: in fig. 4a the model is far from convergence, while in fig. 4a the model converges and is close to overfitting. Our learned optimizers handle both situations gracefully.
11

---

Published in Transactions on Machine Learning Research (03/2025)
1.50
0.
Sondle Lagpvied
1.00
Traits
0.30
3.25
1000
1000
220
400
1000
Steps
Communication Stege
Communication Stopa
(a) FMNNIT, 2-Layer MLP
(b) FMNIST, 3-Layer MLP
(c) FMNIST, CNN
0.75
Lo05
5.50
3.2
Local 5GD
Loot 5C0
1.0
Logeo4
03
1.50
220
405
000
300
1003
300
430
6200
032
1000
200
400
6b2
100
1000
Communication Stops
Communicatton Steps
Communication Stops
(d) CIFAR-10, 2-Layer MLP
(e) CIFAR-10, CNN
(I) ImageNet, CNN
Figure 5: Meta Meta-geneealliztioon to new datasets and new architectures. All optimizers were meta- trained and hyper-parameter tuned for task 5a Meta-generaliin getereeeeraliztionnn is evaluated in three progressively more difficult settings: new architectures same dataset (bb 5 y new dataset same architecture (5d), and different architectures on the same dataset, but experience difficulties optimizing the same architecture on a new dataset. However, the improvements of performance from LAgg-A-f to LAgg-A-cf in plot (5f) shows that these issues can be miligated by scaling training tasks. Finally, both learned optimizers evaluated generalize outside of the training data distribution and architecture in plots (5e) and 5f
new dataset and new architecture (5c| 5fp. Learned optimizers achieve strong generalization
SlowMe
LOpt-A
0.40
LDpt.A
Test Accuracy
0.35
Test Loss
Train Loss
5
0.30
5-
0.25
0.20
0.10
SlowMo LOpt-A
0.05
2
0.00
500
1000
1500
1000
1500
2000
500
Communication Steps
Communication Steps
Communication Steps
(a) Test Loss, 15s
(b) Train Loss, ResNet-155 152
(c) Top-1 1 Accuracy, ResNet- 152
Figure Training ResNet-152 on Imagenet64. We meta-train new LOpt-A optimizer for 2000
communication steps and compare to well-tuned Slowmo baseline.
5.5
Meta-gemeralization
The results are reported in figs. and In fig. 5 we evaluate generalization in three progressively more difficult settings: new architectures same dataset (figs. 5b and 5ch. new dataset same architecture
(fig. 5d], and new datasot and new architecture (figs. 5c and 5fß. In fig. 7 we evaluate the capability of our learned optimizers trained at one H value to generalize to another. Both these figures report results from optimizers meta-trained and hyperparameter tuned on FMNIST and are therefore of smaller scale. In contrast, fig. 2 reports larger scale experiments showing the generalization performance of optimizers meta-trained and hyperparameter tuned on ImageNet to much larger tasks, such as ResNet50, ViT, and a decoder-only transformer which especially relevant in deep learning.
12

---

Published in Transactions on Machine Learning Research (03/2025)
0.5
SlowMo
LAgg-A H=4
0.4
LAgg-A A H=16
Train Loss
0.3
0.2
0.1
0.0
200
600
800
1000
Communication Steps
Figure 7: LAgg-A trained H generalizes to H 4. We observe that LAgg-A H=16 trained H 16, K improves upon a strong SlowMo baseline = 4, K 8.
5.5.1 Meta-trained FMNIST and CIFAR-10
fig. LAgg-A-f and LOpt-A-f are trained on the FMNIST, 2-Layer MLP task, while LAgg-A-ci is trained on a two-dataset lask using FMNIST and CIFAR-10 with the same 2-Layer MLP. All baseline models use hyperparameters tuned on the FMNIST 2-Layer MLP task. Every model is trained using 8 and H 4 with the exception of LAgg-A H=16 (trained using K and H 16).
5.5.2 Generalization to Unseen Architectures
We observe that our learned optimizers can generalize to unseen architectures (figs. 5b and 5cp. In particular, LAgg-A- A-f trained on 2-Layer MT.P tasks can perform well on a CNN and an MLP of different depth, showing gencralization in our communicttion-effciin sctting. Performance in the case of the CNN is particularly strong without having seen this task during training.
5.5.3 Generalization to Unseen Datasets
We observe that LAgg-A meta-trained on FMNIST 2-Layer MLP struggles to optimize the same architecture on CIFAR-10 (fig, 5d) and ImageNet (fig. 12 appendix E.4p. We note, however, that including an additional task (CIFAR-10, MLP) during meta-learning can significantly improve performance. Specifically, we observe that this learned optimizer (TAgg-A-cf) is able to generalize to both of its in-distribution tasks (CIFAR-10 and FMNIST MLP) as well improve performance generalization can be achieved by scaling the training tasks in our communication--lliieee setting has been demonstrated for standard optimization settings in the learned optimization literature (Metz et al. 2022b
ImageNet MLP. This suggests that stronger meta-
5.5.4 Generalization to Unseen Datasets and Architectures
Interestingly, observe (figs. 5e and 5f] that both learned optimizers, LAgg-A-f and LAgg-A-cf achieve strong generalization when varying both the dataset (CIFAR-10 and ImageNet) and the architecture (CNN). This is perplexing when contrasted with the poor generalization observed in the previous paragraph when training the same MLP architecture on CIFAR-10 and ImageNet. We hypothesize that these difficulties arise from changes in MLP dimensions required to accommodate CIFAR-10 and ImageNet 3-channel images as compared to FMNIST's singlo-channel images. As for the strong performance when optimizing the CNN, we believe this due the architecture's inductive biases for image processing, making it relatively casier to optimize.
13

---

Published in Transactions on Machine Learning Research (03/2025)
5.5.5 Scaling up: Meta-trained ImageNet
We now consider larger-scale meta-training task along with an array of target modern architectures and lasks. Figure reports mela-ggnneralizationn results to ResNet50, VIT model, and Decoder-only LM. Lagg- A and I.Opt-A were meta-trained on the 3-layer MLP ImageNet task, while the baselines were extensively iyperparanneter-uuned for this task.
Figure 2 (c,d) shows meta- generalization results for ResNet50 trained on ImageNet" (64 x 64 images). We observe strong generalization of I.OptA,, outperforming all baselines, while LAgg-A performs well at the beginning, but encounters some instability later on in training. The generalization of LOpt-A particularly notable ResNet50 has 50x more parameters than the architecture seen during meta-traint veeta-training and the imput images are two times larger.
Figure 2 (e. shows meta-generalization results for a ViT model of the same size as DeïT tiny (Touvron et al. 2021 Both Lagg-A and 1.0pt-A show strong generalization, but L.pt-A performs best again. Interestingly, for this task, both global learned optimizers outperform the communication-effiieese baselines by a large margin.
Figure 2 hg shows meta-generalization to decoder-only transformer for causal language modeling on LMIB. We observe notably strong generalization performance from both LAgg-A and LOpt-A, improving on all the basclines by a large margin. As shown in table 2 both optimizers reach the minimal loss value achieved by local SGD in over 6 times fewer communication steps.
These results establish the existence of highly promising meta-generrllizatonn capabilities Metli et al. 2022a by for learned optimizers in communication-offcien settings that appear to improve with scaling the meta- training task. Moreover, they demoustrate that such optimizers can also generalize to different values of H, suggosting that it is possible to obtaín learned optimizers that general in H and in tasks by scaling training compute and task variely while using higher H values.
5.5.6 Training ResNet- 152 on ImageNet
To demonstrate that our method can smoothly train large models on ImageNet, meta-train learned optimizer on four 3-layer MLP ImageNet-32 classification tasks of width w E {64, 128, 512} for 2000 steps in a = 8 and = 8 setting. We compare exclusively to a tuned SlowMo baseline since this has been the strongest performing bascline throughout the paper. For both optimizers, we use gradient clipping to a norm of and weight decay of le - for the local steps. Our results in Figure B demonstrate that LAgg-A outperforms SlowMo and reaches 40% top-1 accuracy after 2000 communication steps (se 14 epochs). While do not test more steps this exceeds meta-training widow, expect that learned optimizers meta-trained for longer should achieve analogous strong performance.
6 Limitations
Despite our method's strong performance in variety of settings, it still has some limitations. In some cases meta-genere gettaagnneralizttion is limited (fig. 5d]. This problem is generally observed in previous L20 literature and some recent methods (e-g. STAR (Harrison et al. 2022 Thérien et al. (2024)) can complement our method
to further improve gettaagenealiiztionn..
This work is first step towards the use of L20 for commmiiatiion-efiiien learning and our focus in this work is learning the global step of local learning schemes. We believe that that combining these methods with other communictiion-cffciinn tochniques such gradient sparsification (appendix E..5 may prove effective and have results going in this direction, but we leave more detailed study of this to future work.
Finally, our method does not have local learnable components and relies on vanilla SGD steps performed locally. This was a design choice to allow for a more simple and scalable investigation. We leave room in future research to address these limitations.
14

---

Published in Transactions on Machine Learning Research (03/2025)
7 Conclusion
We have demonstrated the utility of learned optimization for improving communication-elliceen distributed training homogeneous data and device settings. Specifically, we proposed two learned optimizer archi- tectures for this setting LAgg-A and LOpt-A. Our results illustrate that these optimizers can effectively be applied in caamunication-cffciinn distributed settings. We highlight their generalization capabilities unseen architect and datasets. These findings establish learned optimization improving commuiiction--fficies distributed training algorithms for deep learning while scaling to diverse architectures, datascts, and H values. They also hold promise not only in the current context but also in decentralized and federated learning scenarios.
promising direction for
Broader Impact Statement
This paper proposes learned algorithmic enhancements for communication efficient optimizers. Although there are many potential future societal consequences of our work, the application of methods proposed herein does not directly pose societal consequences. That being said, in the long term, data-driven optimization algorithrs do offer the possibility of creating inscrutable optimization algorithms which could be tailored to user's specification and may not always be aligned with human values.
Acknowledgements
We acknowledge support from Mila-Samsung Research Grant, FRQNT New Scholar [E.B.], the Canada Excellence Research Chairs Program in Autonomous AI, and the FRQNT Doctoral (B2X) scholarship [B.T., A.M.J. We also acknowledge resources provided by Compute Canada and Calcul Québec.
References
Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. Qsgd: Commuiication-effiienn sgd via gradient quantization and encoding. Advances in neural information processing systems, 30, 2017.
Brandon Amos. Tutorial amortized optimization for learning optimize over contimuous domains. arXiv e-prints, pp. arXiv 22X2,,220,, 2022.
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
Debraj Basu, Deepesh Data, Can Karakus, and Suhas Diggavi. Qsparse-local-sgd: Distributod sgd with quantization, sparsification, and local computations, 2019.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, and Phillipp Koehn. One billion word benchmark for mcasuring progress in statistical language modeling. CoRR, abs/1312.3005, 2013. URL
Tianlong Chen, Weiyi Zhang, Zhou Jingyang, Shiyu Chang, Sijia Liu, Lisa Amini, and Zhangyang Wang. Training stronger baselines for learning optimize. Advances Neural Information Processing Systems, 33:7332 7343, 2020.
Tianlong Chen, Xiaohan Chen, Wuyang Chen, Zhangyang Wang, Howard Heaton, Jialin Liu, and Wotao Yin. Learning to optimize: A primer and benchmark. The Journal of Machine Learning Research, 23 (1):8562-863) 8620, 2022.
Arthur Douillard, Qianli Feng, Andrei A. Rusu, Rameswar Chhaparia, Yordan Donchev, Adhiguna Kuncoro, Marc' Aurclio Ranzato, Arthur Szlam, and Jonathan Shen. Diloco: Distributed low-communicttion train- ing of language models. arXiv preprint arXiv:231100800 2024. URL https::/arra www.riii..gg/ba...e org/abs/2311,
08105
15

---

Published in Transactions on Machine Learning Research (03/2025)
James Harrison, Luke Metz, and Jascha Sohl-Dickstein. A closer look at learned optimization: Stability, robustness, and inductive biases. Advances in Neural Information Processing Systems, 35:3758-3777 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 TEFE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pp- 770-778. IEEE Computer Society, 2016. URL https://doi. Strg//0..109/CVPP. 2016.90 90)
Sepp Hochreiter, A Steven Younger, and Peter R Conwell. Learning to learn using gradient descent. In Arlificial Neural Nelworks Dedworks-ICANN ICANN 2001: International Conference Vienna, Austria, August 21-25, 2001 Proceedings 11, pp. 87-94. Springer, 2001.
Jinlong Ji, Xuhui Chen, Qianlong Wang, Lixing Yu, and Pan Li. Learning to learn gradient aggregation by gradient descent. In IJCAI, pp. 2614-2620, 2019.
Sai Pranceth Karimiroddy, Satyen Kale, Mchryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International conference on machine learning, pp. 5132-52 5143. PMLR, 2020.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Alex Krizhevsky et al. Learning multiple layers of features from tiny images, 2009.
Qinbin Li, Yiqun Diao, Quan Chen, and Bingsheng Hc. Federated learning on non-iid data silos: An experimental study, 2021.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang, On the convergence of fedavg non-iid data, arXiv preprint arXim:1907.021 2019.
Xiangru Lian, Wei Zhang, Ce Zhang, and Ji Liu. Asynchronous decentralized parallel stochastic gradient descent. In International Conference on Machine Learning, pp. 0043-3052. PMI.R, 2018.
Tao Lin, Sebastian U. Stich, and Martin Jaggi. Don't use large abs/1808.07211, 07211, 2018. URL ptttp://arxin http//aaww..oooarrraaaa/aaa///ee 0olbss/888..077 07211
use local SGD. CoRR,
Kaifeng Lv, Shunhua Jiang, and Jian Li. Learning gradient descent: Better generalization and longer horizons. CoRR,
2017. URL http:///aaxi...rg/aaa////t org/abs/1703. 03633
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. caamunication-cffcieee learning of deep networks from decentralized data. In Artificial intelligence and slatistics, pp. 1273- 1282. PMLR, 2017.
H. Brendan McMahan, Eider Moore, Daniel Ramage, and Blaise Agüera y Arcas. Federated learning of deep networks using model averaging. CoRR, abs/1602.00629) 2016. URL http:://txixit// wrg//bs/162//// http:////aaa...o//a/////1/0 05629
Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Under- standing and correcting pathologics in the training of learned optimizers. In International Conference on Machine Learning, pp. 4556-4565. PMLR, 2019.
Luke Melz, Niru Maheswaranathan. C Daniel Freeman, Ben Poole, and Jascha Sohl-Dickstein. Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselyes. arXiv preprint arXiw:2009.112: 11243, 2020.
Luke Metz, C. Daniel Freeman, James Harrison, Niru Maheswaranathan, and Jascha Sohl-Dickstein. Prac- tical tradeoffss between memory, compute, and performance in learned optimizers, 2022a.
Luke Metz, James Harrison, Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, et al. Velo: Training versatile learned optimizers by scaling up. Xixiv preprint ar Niw:2211.09760 2022b.
16

---

Published in Transactions on Machine Learning Research (03/2025)
Konstantin Mishchenko, Grigory Malinovsky, Scbastian Stich, and Peter Richtárik. Proxskip: Yes! local gradient steps provably lead to communication acceleration! finally! In International Conference Machine Learning, pp. 15750 15769. PMLR, 2022.
Adel Nabli and Edouard Oyallon. Dadao: Decoupled accelerated decentralized asynchronous optimization for varying gossips. arXiv preprinl arxiv:2208. 00779, 2022.
Adel Nabli, Eugene Belilovsky, and Edouard Oyallon. A²CID:: Accelerating asynchronous communication in decentralized deep learning. arXiv preprint arXim:2306. 08289, 2023.
Jose Javier Gonzalez Ortiz, Jonathan Frankle, Mike Rabbat, Ari Morcos, and Nicolas Ballas. Trade-offs of local sgd at scale: An empirical study. arviv preprint arXiv:2110. 0tt::2100088. 0..0000813,,. 2021.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint orXiv:2008.00295 2020.
Herbert E. Robbins. A stochastic approximation method. Annols of Mathematical Statistics, 22:400-405 407, 1951. URL /ttts://pppp www.nticschollea.e........ wtt.s/antccccooooooooooooopaaaaa 16945044
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 2015.
Jürgen Schmidhuber. Learning to control weitht memories: An alternative to dynamic recurrent net- works. Neural Computation, 4(1):131-139, 1992.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, 1CML 2018, Stockholmsmüssann, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 4603 4611. PMLR, 2018. /tt:ppptp// mlr. press/v80/ shazeer 88mhtmll
Shaohuai Shi, Xiaowen Chu, Ka Chun Cheung, and Simon See. Understanding top-k sparsification in distributed deep learning. CoRR, abs/1911.08777 108772, 2019. URL
Sebastian U. Stich. Local sgd converges fast and communicates little, 2019.
Sebastian U. Stich, Jean-Baptiste Cordonnier, and Martin Jaggi. Sparsified SGD with memory. CoRR, abs/1809.07599, 2018. URL http:://arxv.orr//////1880 17599
Benjamin Thérien, Charles-Etienne Joseph, Boris Knyazev, Edouard Oyallon, Irina Rish, and Eu- gene Belilovsky, plo: Compute-efficient neta-generallzation of learned optimizers. arXiv preprint arXiv:2406.001 2024.
Scbastian Thrun and Lorien Pratt. Learning to learn. Springer Science & Business Media, 2012.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-cfficient image transformers distillation through attention. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 10347 10357. PMLR, 2021. URL http://pproceedings mlr wwwppeaw.....................eee wtwppppwte/eeepeeppppppppppp httl//11199too tttp//oovoon2aaa html.
Paul VIcol, Luke Metz, and Jascha Sohl-Dickstein. Unbiased gradient estimation in unrolled computation
graphs with persistent evolution strategies. In Marina Meila and Tong Zhang (eds.), Proceedings of the 98th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event,
volume 139 of Proceedings of Machine Learning Research, pp. 18553-10566 PMLR, 2021. URL http: //proceedings.mlr..r mlr wtt...rree//11////cccl//l
17

---

Published in Transactions on Machine Learning Research (03/2025)
Jianyu Wang, Vinayak Tantia, Nicolas Ballas, and Michael G. Rabbat. Slowmo: Improving communication- efficient distributed SGD with slow momentum. CoRR, abs/1910.006433, 2019. URL http:://axxi..orrrr abs/ abs/1910.00643
Jue Wang, Yucheng Lu, Binhang Yuan, Beidi Chen, Percy Liang, Christopher De Sa, Christopher Re, and Ce Zhang. CocktailSGD: Finc-tuning foundation models over 500Mbps networks. In Andreas Krause, Emna Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett. (eds.); Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of mlr.
Machine Learning Research, pp. 36058-36074 PMLR, 23-29 Jul 2023. URL press/v202/wange wang23t. html
Olga Wichrowska, Niru Maheswaranathan, Matthew W Hoffman,, Sergio Gomez Colmenarejo, Misha Denil, Nando Freitas, and Jascha Sohl-Dickstein. Learned optimizers that scale and generalize. In International conference on machine learning, pp. 3751-3760. PMLR, 2017.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.0711 2017.
Binhang Yuan, Yongjun He, Jared Davis, Tianyi Zhang, Tri Dao, Beidi Chen, Percy Liang, Christopher Re, and Ce Zhang. Decentralized training of foundation models in heterogeneous environments. Advances in Neural Information Processing Systems, 35:25464-25477, 2022.
18

---

Published in Transactions on Machine Learning Research (03/2025)
A Learned Optimizers Architecture and Features
Both our proposed global learned optimizers, LOpt-A and LAgg-A, are 2-hidden-layer MLPs with 32 hidden nodes per layer and ReLU activation functions. In addition to the input Ada features introduced in sec- tion 3.2 L, has K other imput features: (a for each of the K workers. Therefore, LAgg-A A has a total of 38 + K input features. LOpt-A, has single additional input feature: 4,, the average of A, "..,..,,!"" A, but the timestep features are normalized to have second moment of across the tensor. With all of this in mind we can compute the number of LOpt-A has a total of |o| = 2402 meta- pattteaars,meess, while LAgg-A for values K E {8, 16, 32} respectively have |0) E (2626, 2882, 3394) meta-paramecers.
A...,,!!)
All
& in the MLP for each of our learned optimizers.
Table 4: Ada features used with our global learned optimizers. All the coefficients, fi, are learnable parameters adjusted during meta-optimization.
Type
#
Description
Equation
3 Momentum accumulators with coefficients
me, 3 -,))
B,, E (1,2,3).
Second moment accumulator with coefficients B..
2 310-11 + (1-33))A3
Accumulators
3 Adafactor row accumulator with coefficients
Tt, 3 + (1 B:) row row_mean(4 (A?)
B, E 7,.
Adafactor accumulator with coefficients 3, E (5, 6.
3
o,, B. 1, + (1- Bt) col col_meard mea_menn A?)
3 Momentum values normalized by the square
miii V0
root of the second moment for i E (5, 6, 7}. The reciprocal square root of the second mo- ment value.
1
Accumulator Features
The reciprocal square root of the Adafactor ac- cumulators.
6
Vru, OR o2,
3
Adafactor delta features for i E (5,6,7)
A. X rt,s X Ca.i
Adafactor momentum features for i E {55, 1), (6,2), (7,3))
3
myy x rt,s X Cz.
Time Features for æ E (1,3, 30, 100,300, 1000, 3000, 10", 3 10" 10").
Time Features
tanh (4)
1
Parameter value.
We
Parameters
1 Average Delta across workers. (LOpt-A only)
A ΣK
===1 A
K
Per-worker Deltas. (LAgg-A only)
A 11.., A (6)
B Meta-training Process
Algorithm 2 describes the meta-training of our global learned optimizers. We adapt the algorithms from Vicol et al. (2021 (algorithm 1) with minimal changes to notation for reader convenience. Note that use PES to estimate gradients for improved meta-training efficiency. However, other algorithms than PES can be used to estimate the global Icarned optimizer's gradients. On lines 4 and 18 we sample the inner problem length log uniformly between Imin and Imax, ensuring that the start of the inner-problem is seen more frequently during meta-training. The states s(t) contain information (e.g., the accumulator state u, the iteration count, the onpimizec's weights W, etc.) that must be propagated from one truncated unroll to the next. g()) track perturbations across unrolls, tracks the total inner problem length, and âPES is the PES gradient estimate. For simplicity, our meta-training algorithms only shows meta-training for single task (Wa, D, L) where
19

---

Published in Transactions on Machine Learning Research (03/2025)
Algorithm 2: Meta-training our Global Learned Optimizers with PES. Note that this algorithm has been adapted from algorithm 1 of (Vicol ct al. 2021 with minimal changes to the notation.
Input:
Optimizer parameters Pt, gradients
Input:
initial stale
T.
truncation length for partial unrolls
local workers
H
local steps
N
number of particles standard deviation of perturbations
a
local learning rate outer steps minimum inner steps maximum inner steps
O
Imin
Iax
training objective
meta loss
Dataset
0
Learned Optimizer Parameters
unroll(.) corresponds to Algorithm with T T.
Initialize < NeuralNetworklntt))
3: Initialize (()) + so for E (1,.... N}
4: Initialize g(() - 0 for E N1
5: Initialize ( +- LOC UNIFORM mmax) for E [I,...,N]
Initialize 0
for I,...,C O do 0
ÔPES
f
for 1, N do
(1) draw from N(O, o21), if is odd
10:
(1-1))
if is even
8(1), L... unroll( (0) , (0) , C, T,, H,
11:
12:
e(i) + -00 (e(
13:
ÔPES OPES + {(,) L
14:
+ c+ T, -
15:
end for
16:
ÔPES f No CLAMW(GPPSS)
1 N. ÔPES
17:
8 8
18:
for L,..., do
19:
if C; 2 (()) then
20:
(11) + 80
((0) + 0
(0) + LOG__LU_I_S MOG__UNIFOMM/nnn Imax)
C; +
21:
end if
22: end for
23: end for
only the initialization of W, is varied across resets of (0) but not D or C. The multitask setting is analogous and simply requires maintaining additional (s(6), (()) ke e) for each task.
As stated scction 4.1 our meta-learning objective, Lr., the average loss over T, iterations. This optimization problem usually requires long unrolls of the compute graph. In our study, we meta-train our global optimizers using I Tmin 100 and Tmax 1000. The only exceptions to this the optimizers trained
20

---

Published in Transactions on Machine Learning Research (03/2025)
8
LOpt-A fwhw Time
LOgt-A fwbo Time
L Opt, Time
8
Step
LOpt- fwbbw + Opt. Step Time
LoghtA febw Ogt Step Time
Local SGD twbw Time
Local SCD Neww Time
SCD Stup
SCD-
Time
Local SGD twbw Opt. Step Time
Local SCD Arbw Opt. Stop Time
Time (8)
LAggi fwbw Time
LAgg-A febw Time
LAgg-A- Upt.
Time
LAggA Opt. Step
LÁgga - fwbw: Opt. Step Time
4
LAgg-A febw Opt. Step Time
2
16
32
64
16
32
64
H Local Steps
H Local Steps
(a) ResNet50 ImageNet-221
(b) VIT ImageNet-128
Figure 8: The overhead of our optimizers relative to Local SGD shrinks as the number of local steps grows. We report timings for the forward and backward pass, optimizer step time, and total step time local SGD, LOpt-A, and LAgg-A for training ResNet50 (a) and ViT-Base (b) on 224 and 128 sized ImageNet across K=S GPUs. All reported values are the median timings of the final 40 training steps across different training restarts.
in fig. for which we found using a truncation schedule leads overfitting later in optimizee training. Therelore, we opted lo train these opimizers with Iiin 1000 and Imax 1000 and the optimizer section 5.5.6 which used Imin 100 and 4max 2000. For most of the learned optimizers in our study, meta-trained for O =5 000 outer steps. The only exceptions are the learned optimizers used in section 5.3 that were meta-trained for O =10 000 outer steps. During meta-training, we used AdamW as our optimizer with a warmup cosine decay schedule. The learning rate starts at 3e-10 and warms up linearly to the peak value 3e-3 after the first 100 iterations. then decays to the final value of le-3 until the end of meta-traini meta-nraining.
C Practical Considerations
Without any optimizations beyond jax. jit (e.g. no quantization of the LÔ or custom kernel), we compare the step time of LOpt-A and LAgg-A to Local SGD. Specifically, report timings for the forward and backward pass, optimizer step time, and lotal step time of the optimizers when training ResNet50 and ViT-Base on 224 and 128 sized ImageNet across K=8 GPUs. In Table ! of the main text, we report the per-stop overhead of our learned optimizers relative to Local SGD. We observe that the overhoad vanishes as the number local steps grows because the cost of local steps dominates. Therefore, learned optimizers are quite appealing for low-communicttion optimization. Tn addition, figure reports the explicit timings as the number of local steps increased.
D Baselines
For every configuration in which used the baseline optimizers, namely the architecture, the dataset and exhaustive hyperparameter sweep the following values. For
the different values of K and H, we SGD and Adam, we scarched over the learning rate a E {1, 5c-1, lc-1, 5e-2, 1c-2, 5c-3, 1e-3, 50-4, le-4, 5e-5, le-5). For local SGD, we searched over the local learning rate E {(1 .5. For SlowMo, we varied the local learning rate 7 E (1,0.5.0.3, 0.1}, the slow learning rate a E (1/r, 5e-1/7, 1e-1/7, 5c-2/7, 1e-2/r, 5c-3 vo-3/yy 1c-3/r, 5c oe-4/7, 1c 1c-4/%, 5c-5/y, 1c-5/y} and the momentum
0.99, 0.95,
0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.6, 0.55, 0.5 }. The best hyperparameters for each configuration are regrouped in table
21

---

Published in Transactions on Machine Learning Research (03/2025)
Table 5: Best hyperparameters for baselines. When it is not specified, the tuning targets training loss.
Configuration
Objective
SGD (a)
Adam (a)
Local SGD (?)
SlowMo /))
FMNIST, 2-Layer MLP, K = 8, H =
Validation loss
0.1
0.001
0.5
0.3 / 0.01 / 0.8
FMNIST, 2- 2-Layer MLP, K 8, H
Train loss
0.1
0.01
0.3
0.1 / 0.95
FMNIST, 2-Layer MLP, K H
Train loss
0.005
0.3
0.1 / 0.95
FMNIST, 2-Layer MLP,
16
Train loss
0.1 / /
FMNIST, 2-Layer MLP, K 16, H
Train loss
0.1
0.005
0.5
0.1 1/ 0.95
FMNIST, 2-Layer MLP,, K = 32, H =4
Train lass
0.1
0.005
0.5
0.3 / 1.66 0.9
CIFAR-IO, CNN, K
Train lass
0.01
2
ImageNet, L-yayer MLP,
Train loss
0.001
0.3
0.1 0.85
Table 6: Best hyperparameters for k 1, H = 1, B Bioc SGD baseline. When it is not specified, the tuning targets training loss.
Configuration
Objective
Steps (K*H*1000)
SGD (a)
FMNIST, 2-Layer MLP, K = 1, H
Train loss
32000
0.05
FMNIST, Layer MLP, K 1,
Train loss
64000
0.05
FMNIST, Layer MIP, K 1, 1,H=1
Train loss
128000
0.05
ImageNet, L-Layer MLP, K 1, 11111 Train loss
32000
0.05
E Extended results
E.1
Evaluating LAgg-A and LOpt-A in-distribution
As mentioned in section 5.2 present the in-distribution training curves for FMNIST 2-Layer MLP and CIFAR-10 CNN in fig.
E.2
Effect of the Number of Workers (K)
In fig. 10 we evaluate the performance of our method as the number of workers (K) increases. Similarly vary K {8, 16, 32). For each different value of K, meta-train our learned optimizers the FMNIST 2-Layer MLP task. We observe that our learned optimizers workers, reaching lower loss in fewer ilerations than all baselines by significant. margin in each case. While LAgg-A performs better, it needs to be retrained for each K. In contrast, LOpt-A does not have to be retrained. Therefore, each optimizer needs to be carefully chosen depending on the use-case.
section 5.3
gracefully handle more
E.3 Ablating Ada Features
Our learned optimizers leverage powerful per-parameter optimization features proposed in Metz al.) (0022)). Here we investigate how important these are to the performance of the optimizers. Specifically, we consider directly feeding the A, only or each A,
(k) to the learned optimization MLP network without adding
any of the Ada features described in appendix (excluding the -A). We observe that a large improvement in convergence and training stability obtained note that the performance of LOpt and LAgg alone still experiences improved convergence early in training with respect to local SGD. These baselines have no momentum calculations and the optimizer is an MLP (as opposed to recurrent model) thus there is no way to maintain history information (unlike SlowMo's momenlum). It is therefore notable that LAgg can achieve similar, albeit slower, convergence to SlowMo during the first 600 iterations. However, LAgg does seem to cause training instability from iteration 800 onwards. Interestingly, the models trained with Ada features do not suffer from such instabilities, despite being trained with the same schedule as LAgg, further demonstrating their benefit.
We denote these baselines as LOpt and L.Agg, respectively
by using Ada features in both cases (fig. 11 However,
22

---

Published in Transactions on Machine Learning Research (03/2025)
1.2
SGD
2.5
Local SGD
Train Loss
Train Less
LÖptA
1.5
0.6
LAgg-A
1.0
0.4
0.0
400
600
1000
Communication Steps
Communication Steps
(a) FMNIST
(b) CIFAR-10
Figure 9: Learned optimizers enable communicaioon-eflcieet learning. Our LOPt-A and LAgg-A outperform strong commuicctionn-eficiont basclincs such as SlowMo and local SGD. They also outperform well-tuned standard optimization strategies at equivalent effective batch sizes.
1.0
1.0
SGD D-B..
0.8
SGD
ADMM
Train Loss
018
0.8
Local
0.6
Trenn
0.0
0.0
0.0
Communikation tion Stepa
Communication Steps
Communication Steps
(a) K=8
K=16
(c) K=32
Figure 10: LAgg-A outperforms all optimizers for K E (8, 16, workers. All training curves reported for the 28 x 28 FMNST dataset. The top plots training curves for small CNN, while the bottom row plots training curves for an MLP. All experiments use H 4. The orange SGD baseline is trained for K H 1000 optimization steps with batch size Blue = 128. The dotted orange line corresponds to the final performance of this basclinc.
E.4 Meta-generalization
As mentioned in section 5.5.1 we present the meta-generallzztion results for the ImageNet 2-Layer MLP task here in fig. 12
E.5 Learned Optimization with Compressed Updates
As mentioned in section section 2.1 gradient or parameters compression techniques, such as gradient spar- sification (Stich et al. 2018 Shi et al. 2019) are orthogonal approaches to reducing communication cost in distributed deep learning. We show that our method works in conjunction with gradient sparsification and compare with baselines that are also applying sparsification. Specifically, we meta-train learned optimizers while implementing top-k sparsification for the deltas, using top-k values {1, 0.1, 0.01) (fraction of the deltas that are communicated each step). Hyperparameters for the baselines can be found in table 7
Figure 13 presents the performance achieved by learned optimizers versus the allotted communication budget (in log, bits). We observe that our learned optimizer achieves better training loss while communicat- ing less. For example, LAgg-0.01 achieves a lower training loss than SlowMo-1, while having a much lower communication budget. Figure| [ shows the train loss achieved by of our learned optimizer during training for different top-k values. We can see that for each value, our learned optimizers achieve a lower training loss than the basclines. Finally, fig. presents the effect of different top-k value on both our learned optimizers, LOpt-A and L.gg-A. Both optimizers robust top-k values down to 0.01.
23

---

Published in Transactions on Machine Learning Research (03/2025)
0.9
Local SGD
0.8
SlowMa
0.7
LOpt
Train Loss
0.6
LOpt-A
LAgg
0.5
LAgg-A
0.4
0.3
0.2
Kaw
0.1
0.0
200
600
800
1000
Communication Steps
Figure 11: Effect of Ada features on optimizer performance. Each learned optimizer trained and tested on FMNIST L-Layer MLP at H and K 8.
7.0
Train Loss
6.6
6.4
Local SCD
6.2
LOpt-Atf
LAgg-Avf
LAgg-A-cf A-f
200
400
600
800
1000
Communication Steps
(a) ImageNet, 2-Layer MLP
Figure 12: Meta-generallzatton to new datasets and new architectures. All optimizers meta- trained and hyper-parameter tuned for task 5a The plot above shows generalization to ImageNet using the same 2-layer MLP architecture as when meta-training, We observe that our learned optimizers exclusively training on 2-layer MLP FMNIST fail to generalize to ImageNet,, but that LAgg-A improves substantially when it was also meta-trained for CIFAR-10. This suggests that meta-generalization can be improved in our communictioon-fficeett setting by simply adding more tasks.
All in all, we observe that our proposed learned optimization framework for distributed learning can similarly provide improvements for sparsification demonstrating that as in the non-learned setting, combining different local learning compression techniques, like sparsification and quantization [Bauu et al. 2019 can further improve communication efficiency.
E.6 Learning Optimizers for Federated Learning
Motivated by privacy and collaboration, foderated learning (McMahan al.] 2016 allows a model to be collaboratively trained by any number of devices without centrally storing or even exchanging client data. Relative to centralized training, this setting has many difficulties. First, distributed and federated stochastic gradient descent algorithms, such FedSGD (McMahan et al.] 2016), have high communication cost re- sulting from the synchronization gradients across all workers. To alleviate this, federated Icarning algorithms such as FedAvg [McMahan el. al.) (016) allow clients to take multiple local steps with their local models before averaging the weights and updating the global model. This results in reduced communication cost and faster convergence per communication round. However, FodAvg can have convergence issucs in part due the local (per client) models diverging from each other (Karimireddy et al. (020). phenomenon known client drift. Variants of FedAvg, such Fed Adagradd Fed and FedAdam (Reddi et al. 2020).
24

---

Published in Transactions on Machine Learning Research (03/2025)
1.0
0.8
LAgg A-1
LAgg-A-0.1
LAgg-A-0001
0.6
LOpt
Trinnnn
L.ptsA-0.1
L
SlowMo-1
SlowMo.d.1
SlowMo-q.01
Local SCD-0.1
Local SGD-0.01
1.25 x
2
Log2 Number Communicated
Figure 13: Performance of top-k learned optimizers versus the communication budget. We show train loss versus the logg number of bits communicated. Full lines represent top-1, dashed lines show top-0.1 and dotted lines are top-0.01.
1.0
Sowtth
LApeA
0.4
0.4
0.4
0.3
0.0
Communicatoo Stepe
Commmniicttio Staps
Genemancaation Strie
(a) Top-1
(b) Top-0.1
(c) Top-0.01
Figure 14: Performance top-k learned optimizers with different top-K values. Our learned optimizer enjoy better performance than baselines for each value of top-k.
other techniques like SCAFFOLD (Karimireddy et al. 2020), which is directly targeting client drift, can improve the convergence speed and final accuracy. Another well-known problem within the FL setting is label-based data heterogeneity, also known statistical heterogeneity, which can also affect client drift. Given the similarities of federated learning with local SGD, LOpt-A and LAgg-A can also be meta-trained in a FL setting.
E.6.1 Methodology
We employ a per-parameter learned optimizer defined by the function Fg to compute the global update. We designed a client invariant optimizer operating on A, corresponding to LOpt-A, that we name FedLOpt. We use this learned optimizer to replace the averaging step in the FedAvg, as described in algorithm 3 The notable difference with algorithm [ is that we sample a subset K of clients, each with their own data (cither homogencously or hetcrogencously distributed depending on the experiment), which expected for federated learning setting.
E.6.2 Dirichlet Partitioning
To simulate label-basod hetcrogencity among the client, we use Dirichlet partitioning (Li ct al. 2021 This process is controlled by a hyperparameter that controls the amount of data heterogeneity. A high value of will induce homogeneity while low value will create heterogeneity, demonstrated in fig. 16
25

---

Published in Transactions on Machine Learning Research (03/2025)
7hp-1
0.5
0.5
0.6
Vhatiing
0.4
(a) LOPt-A
(b) LAgg-A
Figure 15: Effect of the top-k value on the performance of learned optimizers.
Table 7: Baselines best hyperparameters for different top-k values.
Optimizer
Top-K Value
Y
B
So
0.1
0.95
SlowMo
0.1
0.3
1.66
0.8
SlowMo
0.01
0.1
5
0.3
Local SGD
1
0.3
Local SGD
0.3
Local SGD
0.01
0.5
E.6.3 Experiments
For our experiments, clients sampled uniformly random without replacement each round. New clients are selected each round. We use a total of 400 clients throughout the experiments with participation rate of 0.025 for a total of K = 10 clients each round. Each client gets the same nunber of samples. We compute the number of local steps И as described in algorithm 3 with H = EB where E is the number of local epochs we want, B the total number of samples on each client and Bioe is the local batch size used to compute one local step. For all experiments, we fix all experiments, we used T = 1000 rounds. We used 3 different initializations, each with different random seed, and report the average value and the shaded region in plots corresponds to standard deviation. We use the data from all the test split of the dataset to compute the accuracy. We chose this metric as it is more common in federated learning literature.
Biss
1, unless specified otherwise, and Bioc 20. For
We compare our learned optimizer to the following FL baselines: FedAvg (McMahan et al, 2016 FedAda- grad, YogY and FedAdam (Reddi et al., 2020), each combining FedAvg with a different form of adaptive optimization techniques. We fine-tune our baselines for each task with a grid search over the local learning rate 7). and global learning rate My: using the validation loss metric. We fix 0.001, 81 0.9 and 3o 0.99 for all adaptive federated optimizers. We also compare our method with SCAFFOLD (Reddi et al. 2020). that differs from other FL optimizers by directly aiming at reducing the effect of client drift by introducing control variates in local updates. As was done in Reddi et al. (2020). we fix the global learning rate to and sweep over different values of local learning rate. The values of the different byperparamoters can be found lable
E.6.4 Results
We use two tasks, namely FMNIST with MIP and CIFAR-10 with CNN. For both task, we evaluate two settings: one where the data homogeneously distributed
100) and one where data is heterogeneous
(a 0.1) among clients. Results reported in figs. 17 and 18
26

---

Published in Transactions on Machine Learning Research (03/2025)
Algorithm 3: FedT.Opt and FedAvg. Shared steps are not colored.
Data: Number of training rounds T; Number of workers K; Number of local steps H; Local learning rate 7 Initial weights Wo,0; Loss function C; Learned optimizer Fo; Initial accumtlators state uo for E (0,1,....
1) do
Sample subset K of K clients for k E K in parallel do for EE 10.1..... 1 do
4
( ((() GET (I))
W(() (8) wikk) WW (x% y wi(h)
1,++1 t.h
Difference in weights after H local steps: A w (k) w.())
1H
Averaging: A. & Σ
Compute Ada features ($3.2 and update state:
AG, u441 < ADA(W 1.0, uz 4.)
10
Compute global update:
FedLOpt ($4.1.2): WOLOO + Fg (1)
FedAvg McMahan et al. 2016
W4110 + Wa A.
Nuatbor of Sanglees
Number of Samples
(a) 100 (homogeneous)
(b) 0.1 (heterogeneous)
Figure 16: Dirichlet partitioning of clients. Each class label is represented by different color. Figure 16a
shows homogencous data among clients. Figure 16b shows label-basod data heterogencity among clients.
Our learned optimizer achieves better test accuracy in the homogeneous and heterogeneous settings than all well-tuned basclines in our study. This is the case for both FMNIST MLP and CIFAR-10.
27

---

Published in Transactions on Machine Learning Research (03/2025)
Table 8: Best hyperparameters for federated baselines
Optimizer
Task
Ct
7
FedAdam
FMNIST MLP
100
0.1
0.001
FedAdam
FMNIST MLP
0.1
0.1
0.001
FedAdam
CIFAR-10 CNN
100
0.01
FedAdam
CIFAR-10 CNN
0.1
0.1
0.01
YedYogi
FMNIST MLP
100
0.1
0.001
YedYogi
PMNIST MLP
0.1
0.1
0.001
Yogi
CIFAR-10 CNN
100
0.1
0.01
Yod
CIFAR-10 CNN
0.1
0.01
FodAdagrad
FMNIST MLP
100
0.1
0.01
FedAdagrad
FMNIST MLP
0.1
0.1
0.01
FedAdagrad
CTFAR-10 CNN
100
0.1
0.1
FedAdagrad
CIFAR-10 CNN
0.1
0.1
FedAvg
PMNIST MLP
100
0.1
FedAvg
PMNIST MLP
0.1
0.1
FedAvg
CIFAR-10 CNN
100
FedAvg
CIEAR-10 CNN
0.1
0.1
SCAFFOLD
FMNIST MLP
100
0.01
SCAFFOLD
FMNIST MLP
0.1
0.01
SCAFFOLD
CTFAR-10 CNN
100
0.05
SCAFFOLD
CIFAR-10 CNN
0.1
0.05
1
0.00
Acoureyy
Test Accuracy
0.86
0.86
0.81
FedAvg
0.84
FodAdam
FedYogi
0.82
Fed Andgradd
0.82
FedLOpt
Scaffald
0.80
0.80
800
1000
400
800
1000
Treining Rourds
Training Rounds
(a)
100 (homogeneous)
(b) 0.1 (heterogeneous)
Figure 17: Evaluation on FMNIST MLP. Both learned optimizers were meta-trained on the same task
they are evaluated on.
28

---

Published in Transactions on Machine Learning Research (03/2025)
0.8
0.8
0.7
0.7
0.6
0.6
Test Acouraoy
0.5
Test Anuazeyy
0.5
0.4
FodArg
FedAdam
PodYogi
FedAdagred
FedLOpt
Scaffold
0.0
200
400
800
800
1000
200
400
600
800
1000
Training Rounds
Training Rounds
(a) 100 (bomogeneous)
(b) (heterogeneous)
Figure 18: Evaluation CIFAR-10 CNN. Both learned optimizers were meta-trained the task
they are evaluated on.
29


Reviews:
[{"review_id": "Kx5lcqhaLW", "reviewer": "Reviewer_3zHo", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I enjoyed this paper. The authors do a good job of motivating the method, describing its particulars, and benchmarking performance. The combination of meta-learning and local SGD is intuitive. I especially liked the clarity of the experimental sections. \n\nHowever, there are some issues which I would suggest having addressed before publication:\n* For completeness it will be helpful to include the full pseudocode for the meta-training algorithm (along with details such as hyperparameters) in the main paper. I can probably reconstruct some of it by reading Appendix A but am not certain. \n* For completeness it will be helpful to include the full description of the Ada features in the main paper.\n* $u_t$ --- the accumulator variable? -- is used in Algorithm 1 but not defined in the main paper (as far as I can tell).  \n\nSome other minor comments:\n* I didn't quite understand the phrase \"shared steps are not colored\" in the algorithm description. Do you mean \"shared steps are colored\"?\n* It may be helpful to reflect more in detail what exactly the learned optimizer looks like. (Interpreting MLPs can be hard, but for example -- is it very different from just averaging the gradients (like local SGD does)? Are any particular features it is picking up on? Do inductive biases matter? etc.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a new approach for distributed learning that is a variant of local SGD of Stich (2019). In local SGD, workers take several local gradient updates, communicate the cumulative parameter diffs (deltas) to the server, who then averages the deltas to update the weights. Instead of vanilla averaging, the paper instead proposes using a learnable function that maps deltas (and a list of auxiliary \"Ada\" scalar features that are derived from the deltas) to updates. This learnable function can be a shallow MLP, and can be meta-learned by optimizing over a distribution of tasks and datasets. \n\nThe proposed method is validated via a range of experimental results and is shown to improve upon local SGD. In particular, the optimizer shows generalization abilities to larger architectures, bigger datasets, and even new modalities (such as language model training).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Please move the description of Ada features to the main paper, preferably by fleshing out Section 3.2 in greater detail.  \n* Please address (and attempt to fix) the questions listed in the weaknesses above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Qf5se6kii5", "reviewer": "Reviewer_QMtu", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1.The topic studied is meaningful for improving the model training efficiency.\n\n2.Extensive experiments are conducted to verify the effectiveness of the proposed methods.\n\n\nWeaknesses:\n\n1.Please highlight the contributions and the novelty of the paper. It seems that this work is a direct and simple combination of local SGD and meta learning. There is no special design for the two optimizers, LAgg-A and LOpt-A. \n\n2.Could you please compare the total training time of different optimizers in the experiments.\n\n3.Besides the losses during training, the model performance (e.g., accuracy) is expected to be compared. In some cases, the optimizer that achieves fast convergence rate can result in degraded performance.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper investigates the potential of learned optimizers, which are meta-learned to perform global updates based on local SGD iterations. These optimizers aim to close the performance gap between local SGD and state-of-the-art adaptive optimizers while maintaining communication efficiency. This work proposes two architectures for learned optimizers: a worker-aware optimizer (LAgg-A) and a worker-invariant optimizer (LOpt-A). These optimizers are meta-learned to aggregate local updates into effective global updates. It is demonstrated that the learned optimizers can generalize to unseen datasets, architectures, and modalities, such as language modeling, when meta-learned on a single or few combinations of architectures and datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Major revision", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "UxFORfSekN", "reviewer": "Reviewer_y11Q", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Novel application of meta-learning to communication-efficient distributed optimization.\n2. Comprehensive empirical evaluation across multiple datasets and architectures.\n3. Demonstration of generalization capabilities to unseen tasks and larger models.\n4. Integration with existing communication-efficient techniques (e.g., gradient sparsification).\n5. Improvement over strong baselines like local SGD and SlowMo in terms of convergence speed and communication efficiency.\n\nWeaknesses:\n\n1. Limited theoretical analysis:\n  The paper predominantly focuses on empirical results without providing theoretical foundations for the proposed learned optimizers, even in simple settings such as smooth, strongly convex functions [1,2,3,4,5]. While Section 4.1 presents the meta-learning framework and optimization problem, it lacks crucial theoretical guarantees on convergence or optimality. \n\n2. Insufficient analysis of heterogeneous data scenarios:\n  The paper assumes uniformly distributed data across workers, overlooking a critical aspect of real-world distributed learning applications. In federated learning settings, data heterogeneity is a fundamental challenge [6,7], with clients typically having non-IID (Independent and Identically Distributed) data distributions. The paper doesn't investigate how LAgg-A or LOpt-A perform under such heterogeneous conditions, such as when workers have significantly different data distributions or varying amounts of data. This omission is particularly significant given the extensive literature demonstrating the impact of data heterogeneity on distributed optimization [8,9,10] and the various methods developed to address these challenges.\n\n3. LAgg-A's limited applicability in dynamic worker settings:\n  As described in Section 4.1.1, LAgg-A's architecture depends on pre-aggregated information from each worker, making its input size proportional to the number of workers K. While the authors dismiss this limitation, stating \"oftentimes the distributed training assumes some standard fixed budget of workers,\" this design choice severely restricts the optimizer's applicability in practical scenarios. Particularly, LAgg-A cannot handle partial participation - a fundamental aspect of cross-device federated learning where device availability varies over time [11,12]. This limitation is especially problematic given that partial participation is the default training paradigm in modern federated learning systems [13].\n\nReferences:\n\n[1] Nesterov, Y. \"Introductory Lectures on Convex Optimization: A Basic Course.\" *Springer*, 2004.\n\n[2] Bottou, L., et al. \"Optimization Methods for Large-Scale Machine Learning.\" *SIAM Review*, 2018.\n\n[3] Stich, S.U. \"Local SGD Converges Fast and Communicates Little.\" *ICLR*, 2019.\n\n[4] Koloskova, A., et al. \"A Unified Theory of Decentralized SGD with Changing Topology and Local Updates.\" *ICML*, 2020.\n\n[5] Woodworth, B., et al. \"Local SGD With a Communication Compression Guide.\" *NeurIPS*, 2021.\n\n[6] McMahan, H.B., et al. \"Communication-Efficient Learning of Deep Networks from Decentralized Data.\" *AISTATS*, 2017.\n\n[7] Li, T., et al. \"Federated Optimization in Heterogeneous Networks.\" *MLSys*, 2020.\n\n[8] Karimireddy, S.P., et al. \"SCAFFOLD: Stochastic Controlled Averaging for Federated Learning.\" *ICML*, 2020.\n\n[9] Li, X., et al. \"On the Convergence of FedAvg on Non-IID Data.\" *ICLR*, 2020.\n\n[10] Hsu, T.M.H., et al. \"Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification.\" *NeurIPS*, 2019.\n\n[11] Kairouz, P., et al. \"Advances and Open Problems in Federated Learning.\" *Foundations and Trends in Machine Learning*, 2021.\n\n[12] Yang, Q., et al. \"Federated Learning: Challenges, Methods, and Future Directions.\" *IEEE Signal Processing Magazine*, 2019.\n\n[13] Bonawitz, K., et al. \"Towards Federated Learning at Scale: System Design.\" *MLSys*, 2019.\n4. Writing quality and mathematical rigor have significant shortcomings:\n1) The paper fails to clearly justify its criticism of \"hand-designed algorithmic enhancements\". For example, in the context of:\n```\nWork in this field has largely focused on addressing the heterogeneity of data across workers or clients (Karimireddy et al., 2020; Mishchenko et al., 2022). These advancements are generally achieved by hand-designed algorithmic enhancements, whereas our approach relies on more\nflexible and potentially more powerful learnable mechanisms that may generalize these and more complex algorithms.\n```\nThe authors don't provide evidence that their approach improves upon established methods like SCAFFOLD or ProxSkip either theoretically or empirically.\n\n2) The paper uses unconventional terminology without proper definition. Terms like \"gradient feature\" and \"handcrafted features\" are not standard in the field and remain undefined throughout the text.\n\n3) Mathematical notation lacks clarity, particularly in Algorithm 1. For instance, A_t is introduced in line 8 but is not defined in either Section 3.2 or Appendix A, nor are examples provided for classic algorithms like ADAM or AdaGrad.\n\n4) The definition of $F_\\phi$ in Section 4 is ambiguous:\n```\nOur method builds upon local SGD. After $H$ local steps, we employ a per-parameter learned optimizer $F_\\phi$ based on (Metz et al., 2022a) to compute the updated centralized weights (algorithm 1). By computing the centralized update using an expressive neural net $F_\\phi$, our method can be seen as a generalization of existing update methods such as taking the average iterate (Stich, 2019) or computing server-side momentum updates (Wang et al., 2019).\n```\nThe paper alternates between describing $F_\\phi$ as a \"learned optimizer\" and an \"expressive neural net\" without providing a precise mathematical definition.\n\n5) The paper contains imprecise mathematical language, such as referring to \"optimizing a model with parameters w\" rather than optimizing an objective function.\n\n6) The meta-learning objective function:\n\n$$\n\\min \\phi \\mathbb{E}_{\\left(\\mathcal{D}, w_0\\right) \\sim \\mathcal{T}} \\mathbb{E}_{(X, Y) \\sim \\mathcal{D}}\\left(\\frac{1}{T K} \\sum_{t=0}^{T-1} \\sum_{k=0}^{K-1} \\mathcal{L}\\left(X, Y ; F_\\phi(\\cdot)\\right)\\right)\n$$\n\nlacks clarity regarding how the loss depends on indices $t$ and $k$, with no rigorous formulation provided even in the appendix.\n\n7) References lack proper formatting, with inconsistent capitalization and missing journal information. Examples include:\n```\nQsgd: Communication-efficient sgd via gradient quantization and encoding. \nLocal sgd converges fast and communicates little\nQsparse-local-sgd: Distributed sgd with quantization, sparsification, and local computations\n```\n\n5. Experimental evaluation has several limitations:\n\n1) Scale of experiments is insufficient for modern deep learning:\n  While Section 5.5.5 includes experiments with ResNet50 (25M parameters), ViT (5M parameters), and a decoder-only transformer (19M parameters), these architectures are relatively small compared to state-of-the-art models with hundreds of millions or billions of parameters. The paper's primary focus on 2-layer/3-layer MLPs and small CNNs (as evidenced in Table 1) raises concerns about the scalability and practical applicability of the proposed methods to contemporary deep learning models used in research and industry.\n\n2) Absence of standard accuracy metrics:\n  The paper deviates from conventional practice in image classification literature by only reporting training loss values. Standard metrics such as Top1 or Top5 accuracy on both training and test sets, which are crucial for assessing practical utility, are notably absent from the evaluation.\n\n3) Limited comparison with relevant baselines:\n  While the paper demonstrates improvements over Local SGD in communication-efficient settings, it omits comparison with numerous recent methods that have already shown significant advances over Local SGD/FedAvg [14,15,16,17,18,19,20]. This omission makes it difficult to assess the true practical value of the proposed methods relative to the current state-of-the-art.\n\n4) Visualization and presentation issues:\n  Several figures suffer from presentation issues that hinder proper evaluation. Figure 3 lacks logarithmic scaling, making it difficult to assess performance at small loss values. In Figure 2, the similar order of magnitude of losses does not make the experimental contribution seem promising.. These visualization choices limit the paper's ability to effectively communicate the advantages of the proposed algorithms.\n\nReferences\n\n[14] Karimireddy, Sai Praneeth, et al. \"Breaking the Centralized Barrier for Cross-Device Federated Learning.\" *Advances in Neural Information Processing Systems*, vol. 34, 2021, pp. 28663–28676.\n\n[15] Li, Tian, et al. \"Federated Optimization in Heterogeneous Networks.\" *Proceedings of Machine Learning and Systems*, vol. 2, 2020, pp. 429–450.\n\n[16] Reddi, Sashank, et al. \"Adaptive Federated Optimization.\" *arXiv preprint*, arXiv:2003.00295, 2020.\n\n[17] Mishchenko, Konstantin, et al. \"Proxskip: Yes! Local Gradient Steps Provably Lead to Communication Acceleration! Finally!\" *International Conference on Machine Learning*, 2022, pp. 15750–15769.\n\n[18] Yuan, Xiao-Tong, and Ping Li. \"Sharper Analysis for Minibatch Stochastic Proximal Point Methods: Stability, Smoothness, and Deviation.\" *Journal of Machine Learning Research*, vol. 24, no. 270, 2023, pp. 1–52.\n\n[19] Yuan, Xiaotong, and Ping Li. \"On Convergence of FedProx: Local Dissimilarity Invariant Bounds, Non-Smoothness and Beyond.\" *Advances in Neural Information Processing Systems*, vol. 35, 2022, pp. 10752–10765.\n\n[20] Karimireddy, Sai Praneeth, et al. \"Scaffold: Stochastic Controlled Averaging for Federated Learning.\" *International Conference on Machine Learning*, 2020, pp. 5132–5143.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a novel approach to communication-efficient distributed learning by applying meta-learning to optimize the aggregation of local updates in a distributed setting. The main contributions are:\n1. Two learned optimizer architectures, LAgg-A and LOpt-A, designed to improve upon local SGD and its variants.\n2. Demonstration of the learned optimizers' effectiveness across various tasks, including generalization to unseen architectures and datasets.\n3. Empirical analysis of the optimizers' performance with different numbers of local steps (H) and workers (K).\n4. Integration of the learned optimizers with gradient sparsification techniques.\n5. Empirical evidence showing improved convergence rates and communication efficiency compared to standard optimizers and communication-efficient baselines.\n\nThe authors evaluate their method on various datasets (FMNIST, CIFAR-10, ImageNet) and architectures (MLPs, CNNs, ResNet50, ViT, and a decoder-only transformer), demonstrating both in-distribution performance and generalization capabilities.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "1. Mathematical rigor and clarity:\n   - Clearly define all mathematical notation, especially A_t and F_φ\n   - Provide precise formulation of the meta-learning objective, explaining dependencies on t and k\n   - Fix reference formatting and ensure consistent mathematical terminology\n\n2. Experimental evaluation:\n   - Include experiments with larger models (>100M parameters)\n   - Report standard accuracy metrics (Top1/Top5) for image classification tasks\n   - Add comparisons with recent federated learning methods [14-20]\n   - Use logarithmic scale for loss plots where appropriate\n   - Conduct ablation studies to analyze the impact of different components\n\n3. Practical considerations:\n   - Analyze computational overhead, especially for large models\n   - Test robustness to non-IID data distributions\n   - Investigate behavior with partial worker participation\n\n4. Implementation details:\n   - Release source code", "broader_impact_concerns": "1. Generalization and Robustness: The paper shows some generalization capabilities, but it's crucial to ensure that these learned optimizers are robust across a wide range of scenarios to prevent unexpected behavior in critical applications.\n2. Privacy Considerations: In distributed learning scenarios, especially federated learning, privacy is a key concern. The paper does not address how the proposed methods might interact with privacy-preserving techniques like differential privacy.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Overall, making decision for the paper was hard due to several reasons:
- absence of strong theoretical guarantees while performing empirical analysis similar to prior works and comparing only training loss.
- comparison with the strong prior baseline on SlowMo, while usage of specific settings for SGD baseline.
- small-scale experiments due to limited resources and expensive / unknown better optimization of meta-learning while large scale experiments are needed for empirical strong support.

Conceptually, I believe that we need to improve validation of results from the prior works to improve applicability of empirical / theoretical finding of optimization methods by comparing validation loss / accuracy, using large models and datasets, considering different variants of the SGD baseline training.

At the same time having TMLR guidelines, I believe that
- justification for small scale experiments authors provided are acceptable, so this should not affect the final decision if provided results are correct and useful for the community (e.g. paper definitely can serve as an initial step for the following works to properly select how to perform scaled version of analysis).
- comparison with prior works on the communication efficient optimizers is valid, thus the SGD baseline itself did not diminish authors’ contributions and the main story they tell.
- while strong theoretical guarantees are intractable, well established empirical analysis serves as a strong piece of work, as TMLR welcomes both empirical and theoretical results.

As some reviewers and I (AE) conceptually believe that empirical results miss the validation loss in the light of absent theoretical guarantees, I understand the authors' point that prior works and optimization theory itself mainly consider the train loss leaving generalization aside. The paper itself brings interesting results and insights and for the cases where the validation performance was provided it looks even stronger from a practical point. Having all contributions in the paper and strong effort on ablations even those which are orthogonal to the work central topic (e.g. federated learning), I believe community will benefit from the work as a starting point to continue unifying local SGD and federated learning research clearly considering different regimes of data / model / capacity of device / number of devices and showing differences in optimization between these regimes.

However, I request authors to include the following changes in the final revision (which I believe minor and can be done):
- Please include test accuracies in the main paper results to boost practicality of the proposed optimizers and drive usage of them in practice where people are interested in generalization. Based on the results from the paper, I believe all your test accuracies should be in line with the training loss.
- DiLOCO paper was published at ICML 2024 workshop, so it is not fair to say that it is concurrent work (as the first version appeared earlier). I agree that comparing with DiLOCO is not needed, as their proposal is orthogonal to meta learning proposed by authors + it was tested only on LM task, not vision. However, I encourage authors to include some general discussion in the related works comparing with DiLOCO and e.g. discussing how many local steps is possible to do in both works.
- Please add another SGD baseline to quantify exactly performance people can expect if they want to switch from regular training to a communication efficient one.
	- While I understand that authors follow the prior works on that baselines, I believe for practical applications we need to have the baseline standard in the community w/o communication efficient optimization.
	- While I got the point of authors that they consider SGD baseline with the same wall-clock time, having complexity of optimization with very large batches, we do not train them in practice (or do some tricks to be able to do it, e.g. progressive batch scaling, etc.) and thus training longer but with smaller batch (ending up with ~total same compute, but longer wall-clock) gives better results. Also for the setting authors consider (matching the training time and total communication volume of local SGD) SGD is trained with H*B_local having H steps of grad accumulation. This means that setting of SGD with B_loc batch but KxH steps will be slow only by communication cost which happens every step instead of once per H steps as in local SGD. There is always the trade off between compute, wall clock and performance as authors mentioned. However, proposing a communication efficient method one could expect to see the same performance as before but now with reduced wall-clock. But absence of that SGD baseline (training KxH steps with B_loc) cannot help to answer the question if one gets the same performance as before while having speedup in wall clock due to reduced communication by H times.
Decision: Accept (poster)