A Mathematics-Inspired Learning-to-Optimize Framework for Decentralized Optimization

Authors: Yutong He, Qiulin Shang, Xinmeng Huang, Jialin Liu, Kun Yuan
Primary area: optimization
Keywords: Learning to Optimize, Decentralized Optimization, Composite Optimization

Abstract:
Most decentralized optimization algorithms are handcrafted. While endowed with strong theoretical guarantees, these algorithms generally target a broad class of problems, thereby not being adaptive or customized to specific problem features. This paper studies data-driven decentralized algorithms trained to exploit problem features to boost convergence. Existing learning-to-optimize methods typically suffer from poor generalization or prohibitively vast search spaces. In addition, they face more challenges in decentralized settings where nodes must reach consensus through neighborhood communications without global information. To resolve these challenges, this paper first derives the necessary conditions that successful decentralized algorithmic rules need to satisfy to achieve both optimality and consensus. Based on these conditions, we propose a novel **M**athematics-**i**nspired **L**earning-to-**o**ptimize framework for **D**ecentralized **o**ptimization (**MiLoDo**). Empirical results demonstrate that MiLoDo-trained algorithms outperform handcrafted algorithms and exhibit strong generalizations. Algorithms learned via MiLoDo in 100 iterations perform robustly when running 100,000 iterations during inferences. Moreover, MiLoDo-trained algorithms on synthetic datasets perform well on problems involving real data, higher dimensions, and different loss functions.

Decision: reject
Reviewer scores: 3, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1729371184198

thread_type:
official_review

<thread_object>
creator: reviewer_ebwg

<content>
summary:
This paper proposes an approach to extend the Learning to Optimize framework to the decentralized setting. The authors compared the trained model with many commonly used decentralized optimization approaches using empirical evaluations.

strengths:
The effort of extending Learning to Optimize framework to the decentralized setting is interesting.

weaknesses:
1. The comparison with existing decentralized optimization results are not fair. It is commonly known that existing decentralized optimization algorithms are sensitive to parameters such as learning rate. How did you select learning rate for these existing decentralized optimization algorithms? If these existing algorithms were to given the best parameters, then their performance cannot be beaten by the proposed algorithm, which uses essentially the same algorithmic structure. 

2. It is well known that exiting decentralized algorithms are sensitive to parameters like learning rate, it seems much easier to learn the learning rate, than to conduct the learning proposed here.

questions:
see weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730296941800

thread_type:
official_review

<thread_object>
creator: reviewer_xkvb

<content>
summary:
This paper addresses the learning to optimize problem in the decentralized setting. The current learning to optimize algorithms suffer from big search space and poor generalization. The paper considers the mathematical conditions requires in the decentralized system to achieve consensus and optimality and proposed MiLoDo. For a fixed network with individually trained agents, MiLoDo demonstrated robust performance with generalization properties for both synthetic and real data.

strengths:
The MiLoDo framework provides insights on the necessary mathematical conditions for ensuring consensus and optimality in decentralized training, which could be of separate interest.

This work bridges the L2O literature with decentralized optimization techniques, the problem formulation is intuitive and the simplification on parameters as well as reductions in parameter search space works well in practice.

I have not checked the technical details of the proof in the paper's appendix, but the analytical results in the main paper seem to be intuitive and aligns with the existing literature in the decentralized optimization literature.

weaknesses:
Perhaps one of the most important aspects of decentralized optimization and training over a network, as noted in papers such as EXTRA and various gradient tracking approaches, is to address the data and/or function heterogeneity across agents. In the numerical experiments, especially for MNIST and CIFAR-10, the data seems to be evenly split across agents. When the loss functions are similar across the decentralized network, the system essentially reduces to a SVRG problem, where the consensus and optimality questions are, in a sense, trivial. This paper would benefit from additional experiments regarding the impact of data heterogeneity in the framework. The actual notion of generalization should also be discussed regarding whether MiLoDo can generalize across various degrees of heterogeneity. I will base my opinion on the acceptance of this paper on how the authors address this question in their rebuttal.

As the authors have mentioned, the current MiLoDo framework requires a fixed network with fully synchronized updates. This assumption is difficult to satisfy for many current applications of decentralized optimization tasks. 

MiLoDo requires a set of parameters to be learned for all connections among agents, which suggests that the worst-case complexity of the problem is O(n^2) with respect to network size. This scaling is undesirable for large-scale applications.

Despite the mathematical inspiration behind MiLoDo, the framework did not directly address the impact of network structure e.g. the connectivity of the decentralized network, with the only assumption that the graph is strongly connected.

Currently each agent in the MiLoDo framework is required to learn individual parameters, which are specific to the agents. A more concise and elegant solution would be training one set of parameters which can be used for all agents in the network. Though this might not satisfy the mathematical conditions noted in the papers, some discussions would be appreciated.

questions:
Please see the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730596875914

thread_type:
official_review

<thread_object>
creator: reviewer_tvre

<content>
summary:
This paper introduces a new approach to learning-to-optimize for distributed optimization. The core contribution is a new parameterization of the optimizer update rule which is motivated by theoretical conditions that any update rule must satisfy for convergence. This is an elegant approach leading to a new method, MiLoDo, that exhibits superior empirical convergence behavior compared to prior work in the literature on small test problems (LASSO, logistic regression, MLP/MNIST, and ResNet/CIFAR).

strengths:
1. The paper is well-written and easy to follow.
2. The contributions are well-situated with respect to the recent literatures on learned optimizers and decentralized optimization.
3. The motivation for the parameterization introduced in equations (12)--(14) as well as the implementation in (22)--(24) is sound and interesting.
4. Thorough experiments with small problems illustrate the promise of this approach.
5. That MiLoDo generalizes to many more iterations than the learned optimizer was trained on is impressive.

weaknesses:
1. The motivation to use decentralized optimizers remains somewhat unclear, and there doesn't appear to be a well-known widely-adopted implementation today. This may change in the future, especially as some companies talk about training models across data centers. However, without that, the overall motivation for this contribution remains limited.
2. The experimental results focus on smaller problems, the largest being a ResNet trained on CIFAR. These are all problems that can be easily solved today in a centralized manner with very modest/accessible compute resources. The empirical results would be much stronger if they demonstrated that the same trends hold at scales where distributed/decentralized training is necessary.
3. Some experimental details are unclear; for example is the same 5-stage MiLoDo training procedure used for all workloads, including MLP and ResNet? How sensitive is MiLoDo training to this procedure? The MiLoDo training setup for MNIST and CIFAR is also less convincing, given that the meta-training is done using (subsets of) the same dataset that evaluation is performed on. It would be much more convincing to see results on neural network training where there is clear separation between training ("optimizees") and the workload used to evaluate the learned optimizer.

questions:
1. For Lines 82--93, are there references that could be cited to support the last three challenges mentioned, especially "weak generalization"?
2. Is there stronger motivation that can be provided for where decentralized (learned) optimization is being applied or will be useful in practice? Similarly, can motivation be provided for decentralized LASSO or logistic regression tasks?
3. Could consistent notation be used throughout the paper? In Sections 2 and 3 the learnable update rule is denoted $\bf{g}$, while in Section 4 $\bf{g}$ is used for subgradients of the regularization term and the learned update rule is $\bf{m}$.
4. I'm confused how the generalization to problems with higher problem dimension works. From Section 4.2 I understand that the update rule is implemented as an LSTM, for which I would have thought the shapes of its weights depended on the problem dimension. Can you please explain?
5. The formulation and derivation of MiLoDo in Section 4.1 is written with $p_i^k$, $p_{i,j,1}^k$ and $p_{i,j,2}^k$ that are memoryless, while in Section 4.2 we see these are implemented with LSTMs, i.e., RNNs that have memory/state across iterations. That this helps or is useful is not surprising since it is known that optimizers with memory/momentum are important for efficient convergence of decentralized optimization (see, e.g., EXTRA, DIGing, SlowMo). How do we reconcile this important difference between Sections 4.1 and 4.2? Is memory helpful/sufficient, but not necessary according to the theory? Do the theoretical claims also apply to this implementation with LSTMs?
6. Any intuition about the loss spike/instability (in $F(x)$) observed in Fig 7 around iteration 500?
7. The experiments illustrate the effectiveness and capabilities of trained MiLoDo optimizers. It would be good to also include some illustration and discussion about the process of training MiLoDo optimizers to provide more intuition and support for the five-stage procedure.
8. I appreciated the inclusion of runtimes and transparency about MiLoDo and learned optimizers requiring more computation per update. Please include more information about the system on which timing experiments were run for results in Table 1. It is impossible to interpret these times without knowing more about the system setup.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693424595

thread_type:
official_review

<thread_object>
creator: reviewer_r2et

<content>
summary:
This paper proposes distributed algorithm called MiLoDo that solve consensus-type problems. Some simulations are run to demonstrate the speed of convergence.

strengths:
- Distributed optimization is an important problem. 
- This paper is fairly easy to read and the methods seem to make sense. 
- A number of simulations are done.

weaknesses:
- It seems that not much in the paper is new? A lot of the tricks are standard in distributed optimization, for example, variable duplicating to create equality constraints are used in ADMM. 
- I'm also not sure what mathematical-inspired means. The methods are follow standard approaches. 
- I would suggest that the paper do a better job in describing what is different between the material in the paper and existing work out there.

questions:
- The update rules in this paper is similar to tuning the step size in iterative algorithms. There are a lot of different methods for tuning these step sizes, it would be useful to compare against some of these.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper considers a learning to optimize (L2O) task for the distributed optimization problems. However, a straight-forward extension of the L2O framework to distributed optimization has difficulties in both the huge size of the search space as well as the lack of a mechanism to ensure consensus. Therefore, instead of searching among the general form first-order algorithm space, the authors propose to limit the space of primal-dual first-order algorithm space. To further simplify the search space, the authors propose to learn how to do diagonal scaling (coordinate-style learning rates tuning like adam) for primal-dual method.  Based on the basic form of primal-dual algorithm, the authors design a learn-to-optimize scheme and train the MiLoDo optimizer. Numerical result has shown the advantage of proposed MiLoDo meta-optimizer. 

However, this work also has some drawbacks.

1. The work is a relatively simple extension from L2O to the distributed setting. The only methodology novelty is properly restricting the search space to some relatively well-known algorithmic space. E.g. primal-dual first-order algorithms with (bounded and positive) diagonal scaling. 

2. This work only inherits the methodology from L2O but has no attempt to resolve the drawbacks of L2O. For example, the L2O meta-optimizer is trying to overfitting the training dynamics of the training problem, while having many limitations in generalizing to different problems. E.g., the optimizer learned to trained ResNet may not work well for CNN and RNN. In fact, as the proposed framework does not consider any dimension-agnostic representations (such as coordinate-wise algorithmic framework), the trained optimizer may not even work for the same problem with different problem dimension. The authors' experiment is a bit cheating in this point as their meta-dataset for LASSO is 20000 dimensional and one may pad the redundant dimensions with zeros in the input, then they can work on different dimension such as 50, 100, etc. But this is definitely not the correct approach. The authors provide additional experiment that samples from the same MNIST dataset with different sampling distribution, but this does not seem enough, and reviewers are not convinced by such experiments. A good meta-optimizer should have the ability to generalize across different datasets. If the meta-optimizer learned on MNIST does not work as well on CIFAR, we cannot say it has good generalization w.r.t. data heterogeneity. Moreover, as the distributed feature of the algorithm introduces additional network topology to the algorithm, this also cause generalization issue due to the change of network topology.   Any small change in the problem setting may require additional retraining of the meta-learner.  All these issues limit the application of the method to practical needs. 

Overall, we think the work is marginally below the acceptance bar and we decide to reject the paper.