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

## Abstract

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, the vast search space of communicating choices and final goal to reach the global solution via limited neighboring communication cast more challenges in decentralized settings. 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.

## 1 Introduction

With the ever-growing scale of data and model sizes in modern machine learning and optimization, there is an increasing demand for efficient distributed algorithms that can harness the power of multiple computing nodes. Traditional centralized approaches that rely on global communication and synchronization face significant communication overhead and latency bottlenecks. This challenge has given rise to decentralized learning, an emerging area that promises to alleviate these issues.

In decentralized learning, computing resources like CPUs/GPUs (known as nodes) are connected via a network topology and only communicate with their immediate neighbors, averaging model parameters locally. This neighbor-based averaging eliminates the need for global synchronization, drastically reducing communication costs compared to centralized methods. Moreover, decentralized algorithms exhibit inherent robustness, maintaining convergence despite node or link failures, as long as the network remains connected.

Motivations for data-driven decentralized algorithms. Most existing decentralized algorithms are handcrafted, driven by optimization theories and expert knowledge. Notable examples include primal algorithms such as DGD (Nedic & Ozdaglar, ; Yuan et al., ) and Diffusion (Lopes & Sayed, ; Chen & Sayed, ), dual algorithms like dual averaging (Duchi et al., ), and primal-dual algorithms such as decentralized ADMM (Shi et al., ), EXTRA (Shi et al., ), Exact-Diffusion (Yuan et al., ) (also known as NIDS (Li et al., )), and Gradient-Tracking (Nedic et al., ; Xu et al., ; Di Lorenzo & Scutari, ). These handcrafted decentralized algorithms are designed to address a wide range of optimization problems, making them versatile and broadly applicable. Furthermore, their convergence guarantees are valid in worst-case scenarios, ensuring strong reliability. However, due to their emphasis on theoretical guarantees and broad applicability, handcrafted algorithms often fail to leverage problem-specific features in data and thus exhibit sub-optimal performance in practice.
This motivates us to explore data-driven decentralized algorithms that exploit problem-specific features to improve performance.

Learning to Optimize (L2O). Our main idea draws inspiration from the L2O paradigm (Gregor & LeCun, ; Andrychowicz et al., ; Bengio et al., ; Monga et al., ; Chen et al., ) that utilizes machine learning techniques to develop optimization algorithms (also known as “optimizers”). Specifically, L2O employs a data-driven procedure where an optimizer is trained by its performance on a set of representative example problems (which we call “optimizees”). Through this training process,
the learned optimizer becomes tailored and adaptive to the structures of problems similar to those in the training set, potentially outperforming general-purpose, handcrafted algorithms.

Two mainstreams in L2O are algorithm unrolling (Gregor & LeCun, ; Monga et al., ) and the generic L2O (Andrychowicz et al., ). Algorithm unrolling conceptualizes each iteration of a certain hand-crafted optimization algorithm as a layer in a neural network, inducing a feed-forward network. In contrast, the generic L2O does not impose any prior mathematical knowledge on the optimizer to be learned. Instead, it crudely parameterizes the optimizer with a recurrent neural network and learns it through end-to-end training.

Challenges in applying L2O to decentralized optimization. While algorithm unrolling and generic L2O have demonstrated strong empirical successes (Andrychowicz et al., ; Lv et al., ; Wichrowska et al., ; Wu et al., ; Metz et al., ; Chen et al., ; Micaelli & Storkey, ; Metz et al., ; Liu et al., ; Gregor & LeCun, ; Moreau & Bruna, ; Chen et al., ; Liu & Chen, ; Ito et al., ; Yang et al., ; Zhang & Ghanem, ; Adler & Öktem, ; Solomon et al., ), their direct application to decentralized settings poses several challenges.

- •
Memory bottleneck.
Algorithm unrolling requires storing a neural network with as many layers as optimization iterations, which easily exhausts available memory, especially as the number of unrolled iterations increases. This situation becomes even more challenging in decentralized optimization, where researchers are usually constrained to test their algorithms on simple target problems (Nedic et al., ; Shi et al., ; Xu et al., ; Shi et al., ; Yuan et al., ) within dozens of nodes (Chen & Sayed, ; Wang et al., ; Shi et al., ; Xu et al., ; Nedic et al., ), and we shall maintain much more memory than traditional decentralized algorithms during the training stage.
- •
Vast search space.
While more memory-efficient, the generic L2O faces a significant challenge: how to parameterize an optimizer properly.
The parameter space of the generic L2O is vast, rendering its training highly ineffective.
This challenge is exacerbated in decentralized optimization due to the extra need to learn inter-node interaction (e.g., when and with whom to communicate, and what information to exchange), thereby further expanding the parameter space.
- •
Consensus constraint. In decentralized optimization, nodes must achieve consensus through local communication with immediate neighbors. This consensus constraint incurs significant complexity to L2O, as nodes, with trained optimizers, must adapt their behaviors to ensure convergence towards a common solution despite lacking global communication.
- •
Weak generalization. L2O often struggles to generalize to out-of-distribution tasks. Without theoretical guidance, it is challenging for L2O to handle the more sophisticated loss landscapes encountered in unseen problems. This challenge naturally carries over in decentralized L2O.

Contributions. To address the aforementioned challenges, this paper proposes a novel Mathematics-inspired Learning-to-optimize framework for Decentralized optimization (MiLoDo). MiLoDo adopts the generic L2O strategy to circumvent memory bottlenecks. However, instead of learning an optimizer directly from an unconstrained parameter space, we introduce mathematical structures inherent in decentralized optimization to guide MiLoDo’s learning process. This significantly narrows the parameter space, enforces asymptotic consensus among nodes, and ensures generalization across out-of-domain tasks. Our contributions can be summarized as follows:

- •
We derive fundamental mathematical conditions that learning frameworks for decentralized optimization that converges quickly to the exact solution should satisfy. These conditions will serve as guiding principles for training decentralized optimizers that can achieve consensus and optimality.
- •
Building upon these conditions, we derive a math-inspired neural network structure for MiLoDo. Utilizing this structure, we demonstrate that for MiLoDo-trained optimizers, any fixed point attains consensus across nodes and achieves the solution to the target decentralized problem.
- •
We develop effective training strategies for MiLoDo, which are critical to ensuring the fast and robust convergence of learned optimizers. We conduct extensive experiments to validate the strong generalization and superior convergence of MiLoDo-trained optimizers.

Experimental results. Our experimental results demonstrate that MiLoDo-trained optimizers exhibit strong generalization to out-of-distribution tasks. Specifically, they can adapt to tasks with varying data distributions, problem types, and feature dimensions. For instance, in the high-dimensional LASSO problem illustrated in Fig. [1](#S1.F1), while trained to operate for $100$ iterations when solving prob-

Figure: Figure 1: Numerical comparison between decentralized algorithms in solving LASSO problems of dimension 30,000. Experimental details are deferred to Sec. [5](#S5).
Refer to caption: /html/2410.01700/assets/x1.png

lems in the training dataset, the MiLoDo-trained optimizer
performs exceptionally well for way more iterations (e.g., 200,000 iterations) when solving unseen problems in the test dataset.
Furthermore, our observations indicate that even when trained on LASSO problems of dimension 300, these optimizers can proficiently solve the LASSO problem with dimension 30,000, as shown in Fig. [1](#S1.F1). More impressively, MiLoDo-trained optimizer achieves about $1.5\times\sim 2\times$ speedup in convergence compared to state-of-the-art handcrafted decentralized algorithms. These phenomena justify the necessity to incorporate mathematical structures into L2O for decentralized optimization.

Related work on decentralized L2O.
Previous studies have designed various L2O algorithms for decentralized optimization, most of which are based on algorithm unrolling. Kishida et al. () and Ogawa & Ishii () use algorithm unrolling to learn decentralized algorithms for the consensus problem. Noah & Shlezinger () unrolls D-ADMM, while Wang et al. () unrolls prox-DGD and PG-EXTRA. Hadou et al. () proposed an unrolled algorithm called U-DGD. However, none of these learned optimizers can operate for more than 100 iterations due to the explosive memory cost caused by algorithm unrolling.
Additionally, a recent work Zhu & Lu () employs a reinforcement learning agent to control local update rules through a coordinator linked with all computing nodes, which is not fully decentralized. More related works on decentralized optimization and learning-to-optimize are discussed in Appendix [A](#A1).

## 2 Preliminaries

Generic L2O.
Let’s begin by addressing a fundamental inquiry: Given a set of optimization problems $\mathcal{F}$, how can we learn an optimizer from this dataset? Consider a parameterized optimizer seeking to minimize $\min_{\bm{x}}f(\bm{x})$, represented as:

$$ $\displaystyle\bm{x}^{k+1}=\bm{x}^{k}+\bm{g}^{k}(\nabla f(\bm{x}^{k});\bm{\theta}),\quad k=0,1,\cdots,K-1$ (1) $$

where $\bm{g}^{k}$ is a learnable update rule typically implemented as a deep neural network(^1^11A common approach is using Recurrent Neural Network (RNN) $\bm{m}$: $(\bm{g}^{k},\bm{h}^{k})=\bm{m}(\nabla f(\bm{x}^{k}),\bm{h}^{k-1};\bm{\varphi})$, where $\bm{h}^{k}$ is the hidden state at iteration $k$, and $\bm{\varphi}$ represents the learnable parameters in the RNN model $\bm{m}$.) parameterized by $\bm{\theta}$. To determine $\bm{\theta}$, we evaluate and refine the performance of the optimizer ([1](#S2.E1)) over the initial $K$ steps on the dataset $\mathcal{F}$. Specifically, this entails minimizing a loss function $\mathbb{E}_{f\in\mathcal{F}}[\frac{1}{K}\sum_{k=1}^{K}f(\bm{x}^{k})]$. This loss minimization process is termed training an optimizer, with the employed dataset $\mathcal{F}$ referred to the training set. Upon determining $\bm{\theta}$, the learned update rule $\bm{g}^{k}(\cdot,\bm{\theta})$ will map the gradient $\nabla f(\bm{x}^{k})$ to a desirable descent direction per iteration. Compared to standard gradient-based algorithms, such an optimizer is tailored to $\mathcal{F}$ and ideally exhibits faster convergence on unseen problems similar to those in the training set. In this context, the generalization of a trained optimizer involves two aspects: generalizing to iterations beyond $K$ and generalizing to diverse problems $f\not\in\mathcal{F}$.

Decentralized optimization. In this paper, we aim to learn a decentralized optimizer. To formally define decentralized optimization, we first introduce several definitions used throughout this paper.

###### Definition 1 (Decentralized network topology) :

We use $\mathcal{G}=(\mathcal{V},\mathcal{E})$ to denote the undirected network topology in a decentralized system, where $\mathcal{V}=\{1,2,\cdots,n\}$ represents the set of all $n$ nodes and $\mathcal{E}=\{\{i,j\}\mid i\neq j,\mbox{ node }i\mbox{ can communicate with node }j\}$ denotes the set of all edges.

Throughout this paper, we assume the decentralized network is strongly connected, meaning there is always a path connecting any pair of nodes in the network topology.

###### Definition 2 (Families of objective functions) :

We define the following function families:

$$ $\displaystyle\mathcal{F}(\mathbb{R}^{d})=$ $\displaystyle\ \left\{f:\mathbb{R}^{d}\rightarrow\mathbb{R}\mid f\mbox{ is closed, proper and convex}\right\},$ $\displaystyle\mathcal{F}_{L}(\mathbb{R}^{d})=$ $\displaystyle\ \left\{f:\mathbb{R}^{d}\rightarrow\mathbb{R}\mid f\mbox{ is convex, differentiable and }L\mbox{-smooth}\right\}.$ $$

This paper targets to solve the following problem over a network of $n$ collaborative computing nodes:

$$ $\displaystyle\bm{x}^{\star}=\arg\min_{\bm{x}\in\mathbb{R}^{d}}\left\{\frac{1}{n}\sum_{i=1}^{n}f_{i}(\bm{x})+r(\bm{x})\right\}$ (2) $$

Here, the local cost function $f_{i}(\bm{x})\in\mathcal{F}_{L}(\mathbb{R}^{d})$ is privately maintained by node $i$, and $r(\bm{x})\in\mathcal{F}(\mathbb{R}^{d})$ is a regularization term shared across the entire network. We assume each node $i$ can locally evaluate its own gradient $\nabla f_{i}(\bm{x})$ and must communicate to access information from other nodes. Additionally, communication is confined to the neighborhood defined by the underlying network topology; if two nodes are not direct neighbors, they cannot exchange messages.

A naive approach to extend generic L2O to decentralized optimization.
A straightforward approach to extend generic L2O ([1](#S2.E1)) to the decentralized setting is as follows:

$$ $\displaystyle\bm{x}_{i}^{k+1}=\bm{x}_{i}^{k}+\bm{g}_{i}^{k}\Big{(}\{\bm{x}_{j}^{k},\nabla f_{j}(\bm{x}_{j}^{k})\}_{j\in\mathcal{N}(i)\cup\{i\}};\bm{\theta}_{i}\Big{)},\quad\forall i\in\mathcal{V},$ (3) $$

where $\bm{x}_{i}^{k}$ represents the local
variable
maintained by node $i$ at iteration $k$, $\bm{g}_{i}$ denotes a learnable update rule with parameter $\bm{\theta}_{i}$ retained by node $i$, and notation $\mathcal{N}(i)$ signifies the set of immediate neighbors of node $i$.
While general, the naive update rule $\bm{g}_{i}$ encounters two significant challenges: (I) its vast search space, as finding an effective $\bm{g}_{i}$ requires exploring all possible combinations of the iterative variables, gradients, and neighbors’ information. This complexity makes it difficult to identify a suitable rule, especially when training samples are limited. (II) The update rule in ([3](#S2.E3)) lacks a mechanism to ensure that all nodes reach consensus and converge to the common solution of problem ([2](#S2.E2)), i.e., $\bm{x}_{1}^{\star}=\cdots=\bm{x}_{n}^{\star}=\bm{x}^{\star}$, where $\bm{x}_{i}^{\star}$ represents the limit of sequence of $\{\bm{x}_{i}^{k}\}_{k=1}^{\infty}$.
These two limitations inspire us to introduce mathematical structures into $\bm{g}_{i}$ to narrow the search space and enforce asymptotic consensus among nodes.

###### Definition 1 (Decentralized network topology) :

###### Definition 2 (Families of objective functions) :

## 3 Mathematics-inspired Update Rules for Decentralized Optimization

This section establishes the mathematical principles underlying decentralized optimization and utilizes them to motivate the learning-to-optimize update rules for decentralized optimization. In the subsequent subsections, we will first determine the base update rules that decentralized optimizers should follow, and then specify the concrete structure for each base rule.

### 3.1 Base update rules

The base rule serves as a fundamental mechanism to update optimization variables. While it delineates the necessary inputs, it does not specify a particular structure that the rule must follow. Examples include $\bm{g}(\cdot,\bm{\theta})$ in generic L2O ([1](#S2.E1)) and $\bm{g}_{i}(\cdot,\bm{\theta}_{i})$ in naive decentralized L2O ([3](#S2.E3)). This subsection aims to identify improved base update rules
that resolve the aforementioned issues.

Decentralized optimization interpreted as constrained optimization. One limitation in the naive update rule in ([3](#S2.E3)) is that it cannot explicitly enforce variable consensus during updates. To address this limitation, we reformulate the unconstrained problem ([2](#S2.E2)) as the following constrained optimization:

$$ $\displaystyle\min_{\bm{X}\in\mathbb{R}^{n\times d}}$ $\displaystyle\quad\frac{1}{n}\sum_{i=1}^{n}f_{i}(\bm{x}_{i})+r(\bm{x}_{i}),\quad\mathrm{s.t.}\quad\bm{x}_{i}=\bm{x}_{j},\,\forall\,\{i,j\}\in\mathcal{E}.$ (4) $$

Here, the optimization variable $\bm{X}=[\bm{x}_{1}^{\top},\bm{x}_{2}^{\top},\cdots,\bm{x}_{n}^{\top}]^{\top}\in\mathbb{R}^{n\times d}$ stacks the local variables across all nodes. The consensus constraints in ([4](#S3.E4)) are imposed according to the structure of the underlying network topology. Since the network is strongly connected, we have $\bm{x}_{1}=\bm{x}_{2}=\cdots=\bm{x}_{n}$.

Primal-dual algorithm and its implication. The Lagrangian function of ([4](#S3.E4)) is given as follows

$$ $\displaystyle\mathcal{L}(\bm{X},\{\bm{\nu}_{i,j}\}_{\{i,j\}\in\mathcal{E}})=$ $\displaystyle\frac{1}{n}\sum_{i=1}^{n}\left(f_{i}(\bm{x}_{i})+r(\bm{x}_{i})\right)+\sum_{\{i,j\}\in\mathcal{E}}\langle\bm{\nu}_{i,j},\bm{x}_{i}-\bm{x}_{j}\rangle,$ (5) $$

with $\bm{\nu}_{i,j}$ the dual variable of constraint $\bm{x}_{i}=\bm{x}_{j}$. The primal-dual algorithm to solve ([4](#S3.E4)) is given by:

$$ $\displaystyle\bm{x}_{i}^{k+1}=$ $\displaystyle\ \bm{x}_{i}^{k}-\frac{\gamma}{n}\left(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{g}_{i}^{k+1}+\bm{y}_{i}^{k}\right),\quad\bm{g}_{i}^{k+1}\in\partial r(\bm{x}_{i}^{k+1}),$ (6) $\displaystyle\bm{y}_{i}^{k+1}=$ $\displaystyle\ \bm{y}_{i}^{k}+2\gamma\sum_{j\in\mathcal{N}(i)}(\bm{x}_{i}^{k+1}-\bm{x}_{j}^{k+1}),$ (7) $$

where $\bm{y}_{i}:=n\sum_{j\in\mathcal{N}(i)}(\bm{\nu}_{i,j}-\bm{\nu}_{j,i})$
explicitly satisfies $\sum_{i=1}^{n}\bm{y}_{i}=0$. Updates ([6](#S3.E6)) and ([7](#S3.E7)) imply optimality and consensus in the optimization process. To see it, let $\bm{x}_{i}^{\star}$ and $\bm{y}_{i}^{\star}$ denotes the fixed points that updates ([6](#S3.E6)) and ([7](#S3.E7)) converge to for any $i\in\{1,\cdots,n\}$. It follows that

- •
Update ([7](#S3.E7)) implies consensus. With $\bm{x}^{k}_{i}\to\bm{x}_{i}^{\star},\bm{y}^{k}_{i}\to\bm{y}_{i}^{\star}$, we have
$\displaystyle\bm{x}_{i}^{\star}=({1}/{|\mathcal{N}(i)|)}\sum_{j\in\mathcal{N}(i)}\bm{x}_{j}^{\star},\quad\forall i\in\{1,\cdots,n\}\quad\Longrightarrow\quad\bm{x}_{1}^{\star}=\cdots=\bm{x}_{n}^{\star}.$
(8)
- •
Update ([6](#S3.E6)) implies optimality. With the consensus property established in ([8](#S3.E8)), we introduce $\bm{x}^{\star}:=\bm{x}_{i}^{\star}$. Since $\bm{x}^{k}_{i}\to\bm{x}^{\star},\bm{y}^{k}_{i}\to\bm{y}_{i}^{\star}$ and $\sum_{i=1}^{n}\bm{y}^{\star}_{i}=0$, we have
$\displaystyle\nabla f_{i}(\bm{x}^{\star})+\partial r(\bm{x}^{\star})+\bm{y}_{i}^{\star}\ni 0\quad\Longrightarrow\quad\frac{1}{n}\sum_{i=1}^{n}\nabla f_{i}(\bm{x}^{\star})+\partial r(\bm{x}^{\star})\ni 0,$
(9)
which implies that the consensual fixed point $\bm{x}^{\star}$ is the solution to problem ([2](#S2.E2)).

Mathematics-inspired base update rules. To learn better update rules than the handcrafted primal-dual updates ([6](#S3.E6)) and ([7](#S3.E7)), we propose the following base parameterized update rules

$$ $\displaystyle\bm{x}_{i}^{k+1}=$ $\displaystyle\ \bm{x}_{i}^{k}-\bm{m}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k}),\bm{g}_{i}^{k+1},\bm{y}_{i}^{k};\bm{\theta}_{i,1}),\quad\bm{g}_{i}^{k+1}\in\partial r(\bm{x}_{i}^{k+1}),$ (10) $\displaystyle\bm{y}_{i}^{k+1}=$ $\displaystyle\ \bm{y}_{i}^{k}+\bm{s}_{i}^{k}(\{\bm{x}_{i}^{k+1}-\bm{x}_{j}^{k+1}\}_{j\in\mathcal{N}(i)};\bm{\theta}_{i,2}),$ (11) $$

where $\bm{m}_{i}$ and $\bm{s}_{i}$ are primal and dual update rules, maintained by node $i$ and parameterized by $\bm{\theta}_{i,1}$ and $\bm{\theta}_{i,2}$, respectively. We expect the learned updates ([10](#S3.E10)) and ([11](#S3.E11)) to enforce optimality and consensus when $\bm{m}_{i}$ and $\bm{s}_{i}$ satisfy certain conditions (see Sec. [3.2](#S3.SS2) for details). The formats of the inputs to the base update rules in ([10](#S3.E10)) and ([11](#S3.E11)) are inspired by ([6](#S3.E6)) and ([7](#S3.E7)).

However, the above update rules can be further improved. Note that, ([10](#S3.E10)) does not utilize the communicated information when updating the primal variables $\bm{x}_{i}$.
Intuitively, it is more efficient to use neighbors’ information, i.e., $\{\bm{x}_{j},\nabla f_{i}(\bm{x}_{j}),\bm{g}_{j}\}_{j\in\mathcal{N}(i)}$ to update both the primal and dual variables. Therefore, we propose

$$ $\displaystyle\bm{z}_{i}^{k+1}=\$ $\displaystyle\bm{x}_{i}^{k}-\bm{m}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k}),\bm{g}_{i}^{k+1},\bm{y}_{i}^{k};\bm{\theta}_{i,1}),\quad\bm{g}_{i}^{k+1}\in\partial r(\bm{z}_{i}^{k+1}),$ (12) $\displaystyle\bm{y}_{i}^{k+1}=\$ $\displaystyle\bm{y}_{i}^{k}+\bm{s}_{i}^{k}(\{\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}\}_{j\in\mathcal{N}(i)};\bm{\theta}_{i,2}),$ (13) $\displaystyle\bm{x}_{i}^{k+1}=\$ $\displaystyle\bm{z}_{i}^{k+1}-\bm{u}_{i}^{k}(\{\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}\}_{j\in\mathcal{N}(i)};\bm{\theta}_{i,3}).$ (14) $$

where $\bm{z}_{i}$ is a local auxiliary variable to estimate $\bm{x}_{i}$ after one local update ([12](#S3.E12)) within node $i$, and $\bm{u}_{i}$ is the newly introduced update rule to update $\bm{x}_{i}$ with neighbor’s information. Base update rules $\bm{m}_{i},\bm{s}_{i}$ and $\bm{u}_{i}$ serve as foundations to our MiLoDo framework.

### 3.2 Structured update rules

To ensure the sufficient capacity of update rules $\bm{m}_{i},\bm{s}_{i}$ and $\bm{u}_{i}$ in practice, one should opt for neural networks to parameterize them. Inspired by the universal approximation theorem, which states that neural networks can approximate any continuous functions, it follows that searching the parameter space of a neural network model is similar to searching the entire continuous function space. Therefore, in this subsection, we suppose $\bm{m}_{i},\bm{s}_{i}$ and $\bm{u}_{i}$ are picked from the following space without specific parameterization.

###### Definition 3 (Family of learnable functions) :

Given a domain $\mathcal{Z}$, we let $\mathrm{J}\bm{f(z)}$ denote the Jacobian matrix of the map $\bm{f}:\mathcal{Z}\rightarrow\mathbb{R}^{d}$ and $\|\cdot\|_{F}$ denote the Frobenius norm. We define

$$ $\displaystyle\mathcal{D}_{C}(\mathcal{Z})=\left\{\bm{f}:\mathcal{Z}\rightarrow\mathbb{R}^{d}\mid\bm{f}\mbox{ is differentiable, }\|\mathrm{J}\bm{f(z)}\|_{F}\leq C,\forall\bm{z}\in\mathcal{Z}\right\}$ $$

as the family of learnable functions.

Specifically, we assume $\bm{m}_{i}\in\mathcal{D}_{C}(\mathbb{R}^{3d})$, and $\bm{s}_{i},\bm{u}_{i}\in\mathcal{D}_{C}(\mathbb{R}^{|\mathcal{N}(i)|d})$. This ensures that our results do not depend on particular parameterizations but rather reflect general principles.

Mathematical conditions that a good update rule should satisfy.
One may naturally ask: do all the mappings in $\mathcal{D}_{C}(\mathcal{Z})$ serve as effective rules within the framework of ([12](#S3.E12))–([14](#S3.E14))? If not, can we identify the subset of $\mathcal{D}_{C}(\mathcal{Z})$ containing desirable update rules by considering fundamental conditions that these rules must fulfill?

Now we examine the mathematical conditions that base update rules ([12](#S3.E12))–([14](#S3.E14)) need to satisfy in order to guarantee both consensus and optimality, i.e., $\bm{x}_{1}^{\star}=\cdots=\bm{x}_{n}^{\star}=\bm{x}^{\star}$.
Inspired by the primal-dual algorithm discussed in Section [3.1](#S3.SS1), we refer to the base update rules in ([12](#S3.E12))–([14](#S3.E14)) as good ones if they satisfy the following two conditions.

###### Condition 1 ( Fixed point) :

For any $\bm{x}^{\star}\in\arg\min_{\bm{x}\in\mathbb{R}^{d}}f(\bm{x})+r(\bm{x})$, $\bm{g}_{i}^{\star}\in\partial r(\bm{x}^{\star})$ and $\bm{y}_{i}^{\star}=-\nabla f_{i}(\bm{x}^{\star})-\bm{g}_{i}^{\star}$, it holds for any $i\in\mathcal{V}$ that

$$ $\displaystyle\lim_{k\rightarrow\infty}\bm{m}_{i}^{k}(\nabla f_{i}(\bm{x}^{\star}),\bm{g}_{i}^{\star},\bm{y}_{i}^{\star})=\lim_{k\rightarrow\infty}\bm{s}_{i}^{k}(\{\bm{0}_{d}\}_{j\in\mathcal{N}(i)})=\lim_{k\rightarrow\infty}\bm{u}_{i}^{k}(\{\bm{0}_{d}\}_{j\in\mathcal{N}(i)})=\bm{0}_{d}.$ (15) $$

Remark. Here $\bm{x}^{\star}$ is an optimal primal solution, and $\bm{y}_{i}^{\star}$ is the dual solution retained at node $i$, according to ([9](#S3.E9)).
Condition [1](#Thmcondition1) is derived from a fundamental requirement for a good update rule: if $(\bm{x}_{i}^{k},\bm{y}_{i}^{k},\bm{z}_{i}^{k})$ stay at an optimal solution $(\bm{x}^{\star},\bm{y}_{i}^{\star},\bm{x}^{\star})$, the next iterate $(\bm{x}_{i}^{k+1},\bm{y}_{i}^{k+1},\bm{z}_{i}^{k+1})$ should be fixed. By substituting $\bm{x}_{i}^{k}=\bm{z}_{i}^{k}=\bm{x}^{\star},\bm{y}_{i}^{k}=\bm{y}_{i}^{\star}\in-\nabla f_{i}(\bm{x}^{\star})-\partial r(\bm{x}^{\star})$ and $\bm{x}_{i}^{k+1}=\bm{z}_{i}^{k+1}=\bm{x}^{\star},\bm{y}_{i}^{k+1}=\bm{y}_{i}^{\star}$ into ([12](#S3.E12))-([14](#S3.E14)), we will obtain $\bm{m}_{i}^{k}(\nabla f_{i}(\bm{x}^{\star}),\bm{g}_{i}^{\star},\bm{y}_{i}^{\star})=\bm{0}_{d}$, $\bm{s}_{i}^{k}(\{\bm{0}_{d}\}_{j\in\mathcal{N}(i)})=\bm{0}_{d}$, and $\bm{u}_{i}^{k}(\{\bm{0}_{d}\}_{j\in\mathcal{N}(i)})=\bm{0}_{d}$, and Condition [1](#Thmcondition1) reflects these conditions.

###### Condition 2 ( Global convergence) :

For any sequences
generated by the base update rules in ([12](#S3.E12))-([14](#S3.E14)), there exists $\bm{x}^{\star}\in\arg\min_{\bm{x}\in\mathbb{R}^{d}}f(\bm{x})+r(\bm{x})$, $\bm{y}_{i}^{\star}\in-\nabla f_{i}(\bm{x}^{\star})-\partial r(\bm{x}^{\star})$ such that

$$ $\displaystyle\lim_{k\rightarrow\infty}\bm{x}_{i}^{k}=\lim_{k\rightarrow\infty}\bm{z}_{i}^{k}=\bm{x}^{\star},\quad\lim_{k\rightarrow\infty}\bm{y}_{i}^{k}=\bm{y}_{i}^{\star},\quad\forall i\in\mathcal{V}.$ (16) $$

Remark. With Condition [2](#Thmcondition2), any fixed points of ([12](#S3.E12))-([14](#S3.E14)) will be the optimal primal and dual solution to problem ([2](#S2.E2)). This condition enforces both consensus and optimality for update rules in ([12](#S3.E12))-([14](#S3.E14)).

Deriving mathematical structures for base update rules. The following theorem derives the mathematical structures that the base update rules in ([12](#S3.E12))–([14](#S3.E14)) should possess to satisfy the necessary mathematical conditions mentioned above:

###### Theorem 1 ( Mathematics-inspired structure) :

Given $f_{i}\in\mathcal{F}_{L}(\mathbb{R}^{d})$, $r\in\mathcal{F}(\mathbb{R}^{d})$ and base update rules $\{\bm{m}_{i}^{k},\bm{s}_{i}^{k},\bm{u}_{i}^{k}\}_{k=0}^{\infty}$ with $\bm{m}_{i}^{k}\in\mathcal{D}_{C}(\mathbb{R}^{3d})$, $\bm{s}_{i}^{k},\bm{u}_{i}^{k}\in\mathcal{D}_{C}(\mathbb{R}^{|\mathcal{N}(i)|d})$, if Conditions [1](#Thmcondition1) and [2](#Thmcondition2) hold, there exist $\bm{P}_{i}^{k},\bm{P}_{i,j,1}^{k},\bm{P}_{i,j,2}^{k}\in\mathbb{R}^{d\times d}$ and $\bm{b}_{i,1}^{k},\bm{b}_{i,2}^{k},\bm{b}_{i,3}^{k}\in\mathbb{R}^{d}$ satisfying

$$ $\displaystyle\vspace{1mm}\bm{m}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k}),\bm{g}_{i}^{k+1},\bm{y}_{i}^{k})$ $\displaystyle=\bm{P}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{g}_{i}^{k+1}+\bm{y}_{i}^{k})+\bm{b}_{i,1}^{k},$ (17) $\displaystyle\bm{s}_{i}^{k}(\{\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}\}_{j\in\mathcal{N}(i)})$ $\displaystyle=\sum_{j\in\mathcal{N}(i)}\bm{P}_{i,j,1}^{k}(\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1})+\bm{b}_{i,2}^{k},$ (18) $\displaystyle\bm{u}_{i}^{k}(\{\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}\}_{j\in\mathcal{N}(i)})$ $\displaystyle=\sum_{j\in\mathcal{N}(i)}\bm{P}_{i,j,2}^{k}(\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1})+\bm{b}_{i,3}^{k},$ (19) $$

with $\bm{P}_{i}^{k},\bm{P}_{i,j,1}^{k},\bm{P}_{i,j,2}^{k}$ uniformly upper bounded and $\bm{b}_{i,1}^{k},\bm{b}_{i,2}^{k},\bm{b}_{i,3}^{k}\rightarrow\bm{0}_{d}$ as $k\rightarrow\infty$. If we further assume $\bm{P}_{i}^{k}$ to be positive definite, base update rule ([12](#S3.E12)) can be uniquely determined through

$$ $\displaystyle\bm{z}_{i}^{k+1}=\mathrm{prox}_{r,\bm{P}_{i}^{k}}\left(\bm{x}_{i}^{k}-\bm{P}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{y}_{i}^{k})-\bm{b}_{i,1}^{k}\right),$ (20) $$

where notation $\mathrm{prox}_{\phi,\bm{M}}(\bm{x}):=\arg\min_{\bm{y}}\phi(\bm{y})+\frac{1}{2}\|\bm{y}-\bm{x}\|_{\bm{M}^{-1}}^{2}$ and $\|\bm{x}\|_{\bm{M}}:=\sqrt{\bm{x}^{\top}\bm{M}\bm{x}}$ for positive definite $\bm{M}$.

Remark. Theorem [1](#Thmtheorem1) illustrates that the update rules $\bm{m}_{i},\bm{s}_{i}$, and $\bm{u}_{i}$ are not completely free under Conditions [1](#Thmcondition1) and [2](#Thmcondition2). It suggests mathematically-inspired structures for the base update rules, as shown in ([17](#S3.E17))–([19](#S3.E19)), where $\bm{P}_{i}^{k},\bm{P}_{i,j,1}^{k},\bm{P}_{i,j,2}^{k}$ can be regarded as preconditioners, while $\bm{b}_{i,1}^{k},\bm{b}_{i,2}^{k},\bm{b}_{i,3}^{k}$ are bias terms. We name ([17](#S3.E17))–([19](#S3.E19)) as structured update rules. As shown in Sec. [5](#S5), compared to the base update rules in ([12](#S3.E12))–([14](#S3.E14)), these structured update rules benefit from a significantly more condensed parameter space and ensure consensus and optimality upon convergence.

###### Definition 3 (Family of learnable functions) :

###### Condition 1 ( Fixed point) :

###### Condition 2 ( Global convergence) :

###### Theorem 1 ( Mathematics-inspired structure) :

## 4 MiLoDo: An Efficient Math-inspired L2O Framework

Inspired by the structured update rules derived in ([17](#S3.E17))–([19](#S3.E19)), this section develops a practical L2O framework that can be used to learn effective decentralized optimizers for solving problem ([2](#S2.E2)).

### 4.1 Making structured update rules efficient to learn

To ensure computational efficiency of the update rules in ([17](#S3.E17))–([19](#S3.E19)), we specify those $\bm{P}$ matrices as diagonal ones. Inspired by Theorem [1](#Thmtheorem1), which indicates that the bias terms $\bm{b}_{i,1}^{k},\bm{b}_{i,2}^{k},\bm{b}_{i,3}^{k}$ vanish asymptotically, we eliminate these terms. Specifically, we set:

$$ $\displaystyle\bm{P}_{i}^{k}=\mathrm{Diag}(\bm{p}_{i}^{k}),\,\bm{P}_{i,j,1}^{k}=\mathrm{Diag}(\bm{p}_{i,j,1}^{k}),\,\bm{P}_{i,j,2}^{k}=\mathrm{Diag}(\bm{p}_{i,j,2}^{k}),\,\bm{b}_{i,1}^{k}=\bm{b}_{i,2}^{k}=\bm{b}_{i,3}^{k}=\bm{0}_{d},$ (21) $$

where $\bm{p}_{i}^{k},\bm{p}_{i,j,1}^{k},\bm{p}_{i,j,2}^{k}\in\mathbb{R}_{+}^{d}$, and $\bm{p}_{i,j,1}^{k}=\bm{p}_{j,i,1}^{k},\,\forall\{i,j\}\in\mathcal{E}$. With ([21](#S4.E21)), the structured update rules in ([17](#S3.E17))–([19](#S3.E19)) can be further simplified as:

$$ $\displaystyle\bm{z}_{i}^{k+1}=$ $\displaystyle\ \mathrm{prox}_{r,\mathrm{Diag}(\bm{p}_{i}^{k})}\left(\bm{x}_{i}^{k}-\bm{p}_{i}^{k}\odot(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{y}_{i}^{k})\right),$ (22) $\displaystyle\bm{y}_{i}^{k+1}=$ $\displaystyle\ \bm{y}_{i}^{k}+\sum_{j\in\mathcal{N}(i)}\bm{p}_{i,j,1}^{k}\odot(\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}),$ (23) $\displaystyle\bm{x}_{i}^{k+1}=$ $\displaystyle\ \bm{z}_{i}^{k+1}-\sum_{j\in\mathcal{N}(i)}\bm{p}_{i,j,2}^{k}\odot(\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}).$ (24) $$

Here, $\odot$ denotes element-wise production. We name ([22](#S4.E22))–([24](#S4.E24)) as MiLoDo update rules.

Remark. The above MiLoDo update rules cover state-of-the-art handcrafted decentralized algorithms. If $r\equiv 0$ and we let $\bm{p}_{i}^{k}=\gamma\cdot\bm{1}_{d}$, $\bm{p}_{i,j,1}^{k}=\left(w_{ij}/(2\gamma)\right)\cdot\bm{1}_{d}$, $\bm{p}_{i,j,2}^{k}=(w_{ij}/2)\cdot\bm{1}_{d}$, where $\bm{1}_{d}=[1,1,\cdots,1]^{\top}\in\mathbb{R}^{d}$, MiLoDo update rules reduce to Exact-Diffusion (Yuan et al., ) with symmetric doubly-stochastic gossip matrix $W=(w_{ij})_{n\times n}$ and learning rate $\gamma$. However, MiLoDo update rules are more general than Exact-Diffusion due to the learnable preconditioner $\bm{p}_{i}$ and mixing weight $\bm{p}_{i,j}$.

The following theorem provides theoretical guarantees for MiLoDo update rules, demonstrating that their fixed points are the primal and dual optimal solutions to problem ([2](#S2.E2)).
To our knowledge, no existing decentralized L2O algorithms could guarantee that their fixed points are optimal solutions.

###### Theorem 2 ( Exact convergence) :

Assume $\mathcal{G}=(\mathcal{V},\mathcal{E})$ is strongly connected, $\{f_{i}\}_{i\in\mathcal{V}}\subset\mathcal{F}_{L}(\mathbb{R}^{d})$, $r\in\mathcal{F}(\mathbb{R}^{d})$ and there exists $0<m<M<\infty$ such that $[\bm{p}_{i}^{k}]_{l}\geq m$, $m\leq[\bm{p}_{i,j,1}^{k}]_{l}\leq M$, $|[\bm{p}_{i,j,2}^{k}]_{l}|\leq M$ for all $k\geq 0$ and $1\leq l\leq d$. Here $[\bm{x}]_{l}$ denotes the $l$-th coordinate of the vector $\bm{x}$.
If a sequence
generated by ([22](#S4.E22))-([24](#S4.E24)) with initialization $\bm{y}_{i}^{0}=\bm{0}_{d}$ converges to $\{\bm{x}_{i}^{\star},\bm{y}_{i}^{\star},\bm{z}_{i}^{\star}\}_{i=1}^{n}$, then this limit must be the primal and dual optimal solutions to problem ([2](#S2.E2)). In other words, there exists $\bm{x}^{\star}\in\arg\min_{\bm{x}\in\mathbb{R}^{d}}f(\bm{x})+r(\bm{x})$ such that $\bm{x}_{i}^{\star}=\bm{z}_{i}^{\star}=\bm{x}^{\star}$ and $\bm{y}_{i}^{\star}\in-\nabla f_{i}(\bm{x}^{\star})-\partial r(\bm{x}^{\star})$ holds.

###### Theorem 2 ( Exact convergence) :

### 4.2 LSTM Parameterization for MiLoDo update rules

This subsection discusses how to learn $\bm{p}_{i}^{k}$, $\bm{p}_{i,j,1}^{k}$ and $\bm{p}_{i,j,2}^{k}$ in MiLoDo update rules ([22](#S4.E22))–([24](#S4.E24)). To this end, we parameterize $\bm{p}_{i}^{k}$, $\bm{p}_{i,j,1}^{k}$ and $\bm{p}_{i,j,2}^{k}$ through three local coordinate-wise LSTM neural networks $\phi_{M,i}$, $\phi_{S,i}$, $\phi_{U,i}$. Each network is constructed with a single LSTM cell, followed by a 2-layer MLP and an output activation layer. Specifically,

$$ $\displaystyle\bm{p}_{i}^{k},\bm{h}_{M,i}^{k+1}=$ $\displaystyle\ \phi_{M,i}(\nabla f(\bm{x}_{i}^{k}),\bm{y}_{i}^{k},\bm{h}_{M,i}^{k};\bm{\theta}_{M,i}),$ (25) $\displaystyle\{\tilde{\bm{p}}_{i,j,1}^{k}\}_{j\in\mathcal{N}(i)},\bm{h}_{S,i}^{k+1}=$ $\displaystyle\ \phi_{S,i}(\{\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}\}_{j\in\mathcal{N}(i)},\bm{h}_{S,i}^{k};\bm{\theta}_{S,i}),$ (26) $\displaystyle\{\bm{p}_{i,j,2}^{k}\}_{j\in\mathcal{N}(i)},\bm{h}_{U,i}^{k+1}=$ $\displaystyle\ \phi_{U,i}(\{\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}\}_{j\in\mathcal{N}(i)},\bm{h}_{U,i}^{k};\bm{\theta}_{U,i}),$ (27) $$

where $\bm{h}_{M,i}^{k},\bm{h}_{S,i}^{k},$ $\bm{h}_{U,i}^{k}$ are hidden states in the LSTM modules with random-initialization, $\bm{\theta}_{M,i},\bm{\theta}_{S,i},$ $\bm{\theta}_{U,i}$ are learnable parameters in $\phi_{M,i},\phi_{S,i},\phi_{U,i}$, respectively. To achieve $\bm{p}_{i,j,1}^{k}=\bm{p}_{j,i,1}^{k}$, we compute

$$ $\displaystyle\bm{p}_{i,j,1}^{k}=\bm{p}_{j,i,1}^{k}=(\tilde{\bm{p}}_{i,j,1}^{k}+\tilde{\bm{p}}_{j,i,1}^{k})/2.$ (28) $$

Combining the structured update rules with LSTM parameterization, we obtain the complete architecture of MiLoDo framework, as illustrated in Algorithm [1](#algorithm1).

Figure: Algorithm 1 MiLoDo Framework

### 4.3 Training MiLoDo framework

To determine $\bm{\Theta}=\{\bm{\theta}_{M,i},\bm{\theta}_{S,i},\bm{\theta}_{U,i}\}_{i=1}^{n}$ in ([25](#S4.E25)) – ([27](#S4.E27)), we evaluate and refine the performance of the optimizer over the initial $K$ steps on
a batch of training optimizees $\mathcal{F}_{B}$, i.e.,

$$ $\displaystyle\min_{\bm{\Theta}}\mathcal{L}_{K}(\bm{\Theta},\mathcal{F}_{B}):=\frac{1}{|\mathcal{F}_{B}|}\sum_{f\in\mathcal{F}_{B}}\left[\frac{1}{K}\sum_{k=1}^{K}f(\bar{\bm{x}}^{k})\right].$ (29) $$

The variable $\bar{\bm{x}}^{k}=\frac{1}{n}\sum_{i=1}^{n}\bm{x}_{i}^{k}$ is the average of local variables $\bm{x}_{i}$’s. Typically, we set $K=100$ and train the model by truncated Back Propagation Through Time (BPTT) with a truncation length of $K_{T}=20$, following a common setup in previous L2O approaches
(Chen et al., ; Lv et al., ; Wichrowska et al., ; Metz et al., ; Cao et al., ; Chen et al., ; ).
More specifically, we divide the $K$ iterations into $K/K_{T}$ segments of length $K_{T}$ and train the optimizer on them separately. Such a training strategy is denoted by $(K_{T},K)=(20,100)$ throughout this paper.
More training techniques such as initialization and multi-stage training are in Sec. [E.1](#A5.SS1).

## 5 Experimental results

This section presents numerical experiments to validate the strong generalization capability of the MiLoDo-trained optimizer to out-of-distribution tasks. Additionally, we compare it with state-of-the-art handcrafted optimizers such as Prox-DGD, PG-EXTRA, Prox-ATC, Prox-ED, DAPG (Ye et al., ), ODAPG (Ye & Chang, ), as well as the learned optimizer LPG-EXTRA (Wang et al., ).
Note that LPG-EXTRA, an algorithm unrolling method, is confined to solving unseen problems in the test dataset for a maximum of 100 iterations due to the memory bottleneck imposed by its unrolling structure.
Conversely, all handcrafted and MiLoDo-trained optimizers can be tested over much longer horizons, typically in the order of $10^{5}$.

Experimental setup.
In our experiments, we use a special initialization and a multi-stage training strategy discussed in Sec. [E.1](#A5.SS1). Specifically, we train MiLoDo in five stages with $(K_{T},K)=$ $(5,10)$, $(10,20)$, $(20,40)$, $(40,80)$, and $(20,100)$, using Adam with learning rates of
$5\times 10^{-4}$,$1\times 10^{-4}$,$5\times 10^{-5}$,$1\times 10^{-5}$, and $1\times 10^{-5}$, for 20, 10, 10, 10, and 5 epochs, respectively. Throughout all stages, the Adam optimizer is configured with momentum parameters $(\beta_{1},\beta_{2})=(0.9,0.999)$ and the batch size is set at 32. More data collection/generation and training details can be found in Appendix [E](#A5).

Target problems. Our target problems include LASSO, logistic regression, MLP and ResNet.
In all experiments, we use the shape $(n,d,N,\lambda)$ to represent different characteristics of the optimizees, where $n$ represents the number of nodes in the decentralized network, $d$ represents the feature dimension, $N$ represents the number of data samples held by each worker, and $\lambda$ represents the $\ell_{1}$ regularization coefficient. Without further clarification, we consider a ring topology for the network.

Training sets. MiLoDo optimizers in this section are trained on two different training sets: specialized and meta training set. Specifically, the specialized training set consists of 512 synthetic LASSO$(10,300,10,0.1)$ instances, while the meta training set consists of 1280 synthetic instances with various sizes, including 64 LASSO$(10,500,N,0.1)$ for each $N\in\{5,10,15,\cdots,100\}$.

Figure: Figure 2: MiLoDo-optimizer trained on synthetic LASSO$(10,300,10,0.1)$ and tested on unseen LASSO$(10,300,10,0.1)$ instances.
Refer to caption: /html/2410.01700/assets/x2.png

Figure: Figure 4: MiLoDo-optimizer trained on meta training set and tested on LASSO$(10,200,10,0.05)$ with real dataset BSDS500(Martin et al., ).
Refer to caption: /html/2410.01700/assets/x6.png

Figure: Figure 6: MiLoDo-optimizer trained on MLP$(10,\newline 13002,1000,0)$ with MNIST dataset and tested on MLP$(10,13002,5000,0)$.
Refer to caption: /html/2410.01700/assets/x10.png

Generalization to longer testing iterations. MiLoDo optimizer, trained to operate for a small number of iterations with training problem sets, performs well for significantly more iterations when solving unseen problem sets. As illustrated in Fig. [3](#S5.F3), MiLoDo trained on LASSO$(10,300,10,0.1)$ with $(K_{T},K)=(20,100)$, performs robustly for up to 100,000 testing iterations on unseen LASSO$(10,300,10,0.1)$ instances. In contrast, the other learned optimizer, LPG-EXTRA, can only be applied for 100 iterations limited by its memory bottleneck. Moreover, compared with handcrafted optimizers, MiLoDo achieves a $1.7\times$ speedup in convergence and more than a $2\times$ speedup in consensus.

Generalization to higher problem dimensions. MiLoDo optimizer trained with low-dimensional problems can be generalized to solve problems with much higher dimensions.
As illustrated in Fig. [3](#S5.F3), MiLoDo trained on LASSO$(10,300,10,0.1)$ with a problem dimension 300 performs consistently well on LASSO$(10,30000,1000,0.1)$ instances with a much higher dimension of 30,000.

Generalization to real data distributions. MiLoDo optimizer trained with meta training dataset (synthetic LASSO) can be generalized to real data distributions. As illustrated in Fig. [5](#S5.F5), MiLoDo trained on the meta training set performs consistently well on LASSO$(10,200,10,0.05)$ constructed with real dataset BSDS500 (Martin et al., ), achieving more than a $2.5\times$ speedup in both convergence and consensus rate.

Generalization to different problem types. MiLoDo optimizer trained with meta-training set can generalize to different problem types. As depicted in Fig. [5](#S5.F5), MiLoDo trained on the meta-training set, which consists solely of LASSO problems, converges precisely to the global solutions of unseen logistic regression problems with varying feature dimensions $d\in\{50,200,500,1000\}$.

Efficacy in neural network training scenarios. The efficacy of MiLoDo extends to the realm of neural network training, a domain characterized by high computational complexity and strong non-convexity. As shown in Fig. [7](#S5.F7), MiLoDo consistently achieves a $2\times$ speedup in training MLP on the MNIST (Deng, ) dataset, compared to other baseline methods. MiLoDo also achieves a $2\times$ speedup in training ResNet on the CIFAR-10 (Krizhevsky, ) dataset, as illustrated in Fig. [7](#S5.F7). This performance underscores MiLoDo’s ability to efficiently navigate neural networks’ complex loss landscapes, significantly enhancing distributed deep learning.

Scalability to more complex topologies and larger networks. MiLoDo optimizer consistently performs well on complex and large-scale networks, showcasing its superior scalability. As shown in Fig.[9](#S5.F9), MiLoDo consistently enhances efficiency on more complex network topologies and larger networks, achieving a $1.5\times$ speedup on an exponential graph topology, and a $3\times$ speedup on a 50-node network. Further tests confirm its effectiveness across various topologies and a 100-node network. Detailed results are in Appendix [E.4](#A5.SS4).

Influence of math-inspired structures. The mathematics-inspired structures are crucial for the success of MiLoDo. To see this, we directly parameterize the base update rules ([12](#S3.E12))-([14](#S3.E14)) targeting a smooth problem LASSO$(10,10,5,0)$. Detailed experimental setups are deferred to Appendix [E.5](#A5.SS5). As illustrated in Fig. [9](#S5.F9), directly parameterizing the base update rules fails to learn good optimizers.

Figure: Figure 8: MiLoDo-optimizer trained on different network settings: (left) an exponential graph topology of 10 nodes, (right) a network of 50 nodes.
Refer to caption: /html/2410.01700/assets/x14.png

Runtime studies.
When compared to handcrafted optimizers like Prox-ED, MiLoDo faces a higher per-iteration computational cost due to additional neural network calculations for $\{\bm{p}_{i}^{k},\bm{p}_{i,j,1}^{k},\bm{p}_{i,j,2}^{k}\}$. To assess whether its benefits outweigh these costs, we compre its running time with Prox-ED. As illustrated in Table [1](#S5.T1), MiLoDo exhibits only a slight increase in computational cost, about 18.4%, while achieving a significant convergence speedup, resulting in a $1.6\times\sim 2.1\times$ speedup in total.

**Table 1: Runtime Comparison on LASSO$(10,30000,1000,0.1)$.**
|  | Stopping condition: Gap $<10^{-7}$ | Stopping condition: Gap $<10^{-15}$ |  |  |  |
| --- | --- | --- | --- | --- | --- |
|  | Time/Iters | Iters | Total Time | Iters | Total Time |
| MiLoDo | $5.91$ ms | 2.45e+04 | 144.80 s | 1.62e+05 | 957.42 s |
| Prox-ED | $4.99$ ms | 6.22e+04 | 310.38 s | 3.07e+05 | 1531.93 s |

More experimental results. More testing results on MiLoDo optimizers trained with meta training set and Logistic regression optimizees are in Appendix [E.4](#A5.SS4). We also conduct ablation studies on the mixing matrices, see Appendix [E.5](#A5.SS5).

## 6 Conclusions and Limitations

We propose MiLoDo, a mathematics-inspired L2O framework for decentralized optimization. With its mathematics-inspired structure, the MiLoDo-trained optimizer can generalize to tasks with varying data distributions, problem types, and feature dimensions. Moreover, MiLoDo-trained optimizer outperforms handcrafted optimizers in convergence rate. The primary limitation of MiLoDo, which is shared by most state-of-the-art decentralized L2O works, is that it cannot generalize to different network sizes and topologies, which we will address in future work.

## Appendix A More related work

Learning to optimize. The concept of L2O dates back to the 1990s (Thrun & Pratt, ). Different L2O approaches exist: Plug-and-Play (PnP) (Venkatakrishnan et al., ) approximates expensive functions in traditional algorithms; algorithm unrolling (Gregor & LeCun, ; Moreau & Bruna, ; Chen et al., ; Liu & Chen, ; Ito et al., ; Yang et al., ; Zhang & Ghanem, ; Adler & Öktem, ; Solomon et al., ) models the entire procedure as a neural network, effective for domains like image/signal processing; generic L2O (Andrychowicz et al., ; Lv et al., ; Wichrowska et al., ; Wu et al., ; Metz et al., ; Chen et al., ; Shen et al., ; Harrison et al., ; Micaelli & Storkey, ; Metz et al., ; ; ; Jain et al., ; Liu et al., ), which is more related to this paper, parameterizes update rules using current states, enabling flexibility across applications. Some other studies also atempts to learn machine learning models to accelerate the discrete problems solving(Bengio et al., ).

Decentralized optimization.
Decentralized optimization has been extensively studied, dating back to early algorithms like decentralized gradient descent (DGD) (Nedic & Ozdaglar, ; Yuan et al., ), Diffusion (Lopes & Sayed, ; Chen & Sayed, ), and dual averaging (Duchi et al., ) from the signal processing and control communities. These were followed by primal-dual methods such as ADMM variants (Shi et al., ), explicit bias-correction techniques (Shi et al., ; Yuan et al., ; Li et al., ), Gradient-Tracking (Xu et al., ; Di Lorenzo & Scutari, ; Nedic et al., ). More recently, decentralized stochastic gradient descent (DSGD) (Lian et al., ) has gained significant attentions in deep learning. For non-smooth optimization problems, effective algorithms like PG-EXTRA (Shi et al., ), PG-Exact-Diffusion (Yuan et al., ), and PG-Gradient-Tracking (Alghunaim et al., ) utilize the proximal gradient method to solve them. A unified decentralized framework is developed in Alghunaim et al. () and Xu et al. () to unify various decentralized algorithms. All these algorithms are driven by expert knowledge.

## Appendix B Missing proofs

### B.1 Preliminaries

###### Lemma 1 ( Liu et al. ( 2023 ) , Lemma 1) :

For any operator $\mathbf{o}\in\mathcal{D}_{C}(\mathbb{R}^{m\times n})$ and any $\mathbf{x}_{1}$, $\mathbf{y}_{1}$, $\cdots$ $\mathbf{x}_{m}$, $\mathbf{y}_{m}\in\mathbb{R}^{n}$, there exists matrices $\mathbf{J}_{1}$, $\mathbf{J}_{2}$, $\cdots$, $\mathbf{J}_{m}\in\mathbb{R}^{n\times n}$ such that

$$ $\displaystyle\mathbf{o}(\mathbf{x}_{1},\cdots,\mathbf{x}_{m})-\mathbf{o}(\mathbf{y}_{1},\cdots,\mathbf{y}_{m})=\sum_{j=1}^{m}\mathbf{J}_{j}(\mathbf{x}_{j}-\mathbf{y}_{j}),$ $$

and $\|\mathbf{J}_{1}\|_{F}\leq\sqrt{n}C$, $\cdots$, $\|\mathbf{J}_{m}\|_{F}\leq\sqrt{n}C$.

This lemma is an extension of the mean value theorem.

###### Lemma 1 ( Liu et al. ( 2023 ) , Lemma 1) :

### B.2 Proof of Theorem 1

###### Proof.

By Lemma [1](#Thmlemma1), there exists $\bm{Q}_{i,1}^{k},\bm{Q}_{i,2}^{k},\bm{Q}_{i,3}^{k}\in\mathbb{R}^{d\times d}$ such that $\|\bm{Q}_{i,1}^{k}\|_{F}\leq\sqrt{d}C$, $\|\bm{Q}_{i,2}^{k}\|_{F}\leq\sqrt{d}C$, $\|\bm{Q}_{i,3}^{k}\|_{F}\leq\sqrt{d}C$ and

$$ $\displaystyle\bm{m}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k}),\bm{g}_{i}^{k+1},\bm{y}_{i}^{k})=$ $\displaystyle\bm{m}_{i}^{k}(\nabla f_{i}(\bm{x}^{\star}),-\nabla f_{i}(\bm{x}^{\star})-\bm{y}_{i}^{\star},\bm{y}_{i}^{\star})+\bm{Q}_{i,1}^{k}(\nabla f_{i}(\bm{x}_{i}^{k})-\nabla f_{i}(\bm{x}^{\star}))$ $\displaystyle+\bm{Q}_{i,2}^{k}(\bm{g}_{i}^{k+1}+\nabla f_{i}(\bm{x}^{\star})+\bm{y}_{i}^{\star})+\bm{Q}_{i,3}^{k}(\bm{y}_{i}^{k}-\bm{y}_{i}^{\star})$ $\displaystyle=$ $\displaystyle\bm{Q}_{i,2}^{k}(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{g}_{i}^{k+1}+\bm{y}_{i}^{k})+(\bm{Q}_{i,1}^{k}-\bm{Q}_{i,2}^{k})(\nabla f_{i}(\bm{x}_{i}^{k})-\nabla f_{i}(\bm{x}^{\star}))$ $\displaystyle+(\bm{Q}_{i,3}^{k}-\bm{Q}_{i,2}^{k})(\bm{y}_{i}^{k}-\bm{y}_{i}^{\star})+\bm{m}_{i}^{k}(\nabla f_{i}(\bm{x}^{\star}),-\nabla f_{i}(\bm{x}^{\star})-\bm{y}_{i}^{\star},\bm{y}_{i}^{\star}).$ $$

Letting

$$ $\bm{P}_{i}^{k}=\bm{Q}_{i,2}^{k}$ $$

and

$$ $\displaystyle\bm{b}_{i,1}^{k}=$ $\displaystyle(\bm{Q}_{i,1}^{k}-\bm{Q}_{i,2}^{k})(\nabla f_{i}(\bm{x}_{i}^{k})-\nabla f_{i}(\bm{x}^{\star}))+(\bm{Q}_{i,3}^{k}-\bm{Q}_{i,2}^{k})(\bm{y}_{i}^{k}-\bm{y}_{i}^{\star})$ $\displaystyle+\bm{m}_{i}^{k}(\nabla f_{i}(\bm{x}^{\star}),-\nabla f_{i}(\bm{x}^{\star})-\bm{y}_{i}^{\star},\bm{y}_{i}^{\star}),$ $$

the above equation can be reorganized into ([17](#S3.E17)). Addtionally, we have $\|\bm{P}_{i}^{k}\|_{F}\leq\sqrt{d}C$ and $\bm{b}_{i,1}^{k}\rightarrow\bm{0}_{d}$ thanks to Condition [1](#Thmcondition1). Similarly, there exists $\bm{Q}_{i,j,1}^{k},\bm{Q}_{i,j,2}^{k}\in\mathbb{R}^{d\times d}$ such that $\|\bm{Q}_{i,j,1}^{k}\|_{F}\leq\sqrt{d}C$, $\|\bm{Q}_{i,j,2}^{k}\|_{F}\leq\sqrt{d}C$ and

$$ $\displaystyle\bm{s}_{i}^{k}(\{\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}\}_{j\in\mathcal{N}(i)})=$ $\displaystyle\bm{s}_{i}^{k}(\{\bm{0}_{d}\}_{j\in\mathcal{N}(i)})+\sum_{j\in\mathcal{N}(i)}\bm{Q}_{i,j,1}^{k}(\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}),$ $\displaystyle\bm{u}_{i}^{k}(\{\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}\}_{j\in\mathcal{N}(i)})=$ $\displaystyle\bm{u}_{i}^{k}(\{\bm{0}_{d}\}_{j\in\mathcal{N}(i)})+\sum_{j\in\mathcal{N}(i)}\bm{Q}_{i,j,2}^{k}(\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1}).$ $$

Thus, it is sufficient to obtain ([18](#S3.E18)) and ([19](#S3.E19)) by letting $\bm{P}_{i,j,1}^{k}=\bm{Q}_{i,j,1}^{k}$, $\bm{P}_{i,j,2}^{k}=\bm{Q}_{i,j,2}^{k}$, $\bm{b}_{i,2}^{k}=\bm{s}_{i}^{k}(\{\bm{0}_{d}\}_{j\in\mathcal{N}(i)})$, and $\bm{b}_{i,3}^{k}=\bm{u}_{i}^{k}(\{\bm{0}_{d}\}_{j\in\mathcal{N}(i)})$.

Now we rewrite ([12](#S3.E12)) as

$$ $\displaystyle\bm{z}_{i}^{k+1}=\bm{x}_{i}^{k}-\bm{P}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{g}_{i}^{k+1}+\bm{y}_{i}^{k})-\bm{b}_{i,1}^{k},\quad\bm{g}_{i}^{k+1}\in\partial r(\bm{z}_{i}^{k+1}).$ (30) $$

If we further assume $\bm{P}_{i}^{k}$ to be symmetric positive definite, ([30](#A2.E30)) implies

$$ $\displaystyle\bm{0}_{d}\in\partial r(\bm{z}_{i}^{k+1})+(\bm{P}_{i}^{k})^{-1}\left(\bm{z}_{i}^{k+1}-\bm{x}_{i}^{k}+\bm{b}_{i,1}^{k}+\bm{P}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{y}_{i}^{k})\right).$ $$

Consequently, $\bm{z}_{i}^{k+1}$ coincides with the unique solution to the following strongly-convex optimization problem:

$$ $\displaystyle\min_{\bm{x}\in\mathbb{R}^{d}}\quad r(\bm{x})+\frac{1}{2}\|\bm{x}-(\bm{x}_{i}^{k}-\bm{P}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{y}_{i}^{k})-\bm{b}_{i,1}^{k})\|_{(\bm{P}_{i}^{k})^{-1}}^{2},$ $$

i.e.,

$$ $\displaystyle\bm{z}_{i}^{k+1}=\mathrm{prox}_{r,\bm{P}_{i}^{k}}(\bm{x}_{i}^{k}-\bm{P}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{y}_{i}^{k})-\bm{b}_{i,1}^{k}),$ $$

which finishes the proof.
∎

###### Proof.

### B.3 Proof of Theorem 2

###### Proof.

Let $\mathcal{I}_{l}=\{i\in\mathcal{V}\mid[\bm{z}_{i}^{\star}]_{l}=\max_{j\in\mathcal{V}}[\bm{z}_{j}^{\star}]_{l}\}$ be the set of indices $i$’s with the largest $[\bm{z}_{i}^{\star}]_{l}$’s. We first prove the following statement:

$$ $\displaystyle i\in\mathcal{I}_{l}\Rightarrow\mathcal{N}(i)\subset\mathcal{I}_{l}.$ (31) $$

Suppose there exists $i\in\mathcal{I}_{l}$ and $j\in\mathcal{N}(i)\backslash\mathcal{I}_{l}$, it holds that $[\bm{z}_{i}^{\star}]_{l}>[\bm{z}_{j}^{\star}]_{l}$. Define $\delta=[\bm{z}_{i}^{\star}]_{l}-[\bm{z}_{j}^{\star}]_{l}>0$ and choose an $\epsilon\in\left(0,\frac{m\delta}{2(nM+1)}\right)$. By convergence property there exists $K\geq 0$ such that for any $k\geq K$,

$$ $\displaystyle|[\bm{x}_{i}^{k}]_{l}-[\bm{x}_{i}^{\star}]_{l}|\leq\epsilon,\quad|[\bm{y}_{i}^{k}]_{l}-[\bm{y}_{i}^{\star}]_{l}|\leq\epsilon,\quad|[\bm{z}_{i}^{k}]_{l}-[\bm{z}_{i}^{\star}]_{l}|\leq\epsilon,\quad\forall i\in\mathcal{V},1\leq l\leq d.$ $$

By iteration step ([23](#S4.E23)) we have

$$ $\displaystyle[\bm{y}_{i}^{k+1}]_{l}=[\bm{y}_{i}^{k}]_{l}+\sum_{j\in\mathcal{N}(i)}[\bm{p}_{i,j,1}^{k}]_{l}([\bm{z}_{i}^{k+1}]_{l}-[\bm{z}_{j}^{k+1}]_{l}).$ $$

Fix $k\geq K$ and let $\mathcal{N}_{1}=\{j\in\mathcal{N}(i)|[\bm{z}_{i}^{k+1}]_{l}\geq[\bm{z}_{j}^{k+1}]_{l}\}$, $\mathcal{N}_{2}=\{j\in\mathcal{N}(i)|[\bm{z}_{i}^{k+1}]_{l}<[\bm{z}_{j}^{k+1}]_{l}\}$,

we have

$$ $\displaystyle m(\delta-2\epsilon)\leq$ $\displaystyle[\bm{p}_{i,j,1}^{k}]_{l}([\bm{z}_{i}^{k+1}]_{l}-[\bm{z}_{j}^{k+1}]_{l})\leq\sum_{\tau\in\mathcal{N}_{1}}[\bm{p}_{i,\tau,1}^{k}]_{l}([\bm{z}_{i}^{k+1}]_{l}-[\bm{z}_{\tau}^{k+1}]_{l})$ $\displaystyle=$ $\displaystyle[\bm{y}_{i}^{k+1}]_{l}-[\bm{y}_{i}^{k}]_{l}-\sum_{\tau\in\mathcal{N}_{2}}[\bm{p}_{i,\tau,1}^{k}]_{l}([\bm{z}_{i}^{k+1}]_{l}-[\bm{z}_{\tau}^{k+1}]_{l})$ $\displaystyle\leq$ $\displaystyle 2\epsilon+M(n-1)\cdot(2\epsilon),$ $$

which implies

$$ $\displaystyle m\delta\leq 2(nM+1)\epsilon,$ $$

a contradiction. Consequently, ([31](#A2.E31)) holds and thus together with the strongly connectivity we obtain $\mathcal{I}_{l}=\mathcal{V}$, i.e., $[\bm{z}_{1}^{\star}]_{l}=[\bm{z}_{2}^{\star}]_{l}=\cdots=[\bm{z}_{n}^{\star}]_{l}$. By arbitrariness of $l$, we conclude that there exists $\bm{x}^{\star}\in\mathbb{R}^{d}$ such that $\bm{z}_{i}^{\star}=\bm{x}^{\star}$ for any $i\in\mathcal{V}$. By iteration step ([24](#S4.E24)) and $|[\bm{p}_{i,j,2}^{k}]_{l}|\leq M$ we have

$$ $\displaystyle\bm{x}_{i}^{\star}=\lim_{k\rightarrow\infty}\bm{x}_{i}^{k+1}=\lim_{k\rightarrow\infty}\left(\bm{z}_{i}^{k+1}-\sum_{j\in\mathcal{N}(i)}\bm{p}_{i,j,2}^{k}\odot(\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1})\right)=\bm{x}^{\star},\quad\forall i\in\mathcal{V}.$ $$

By iteration step ([22](#S4.E22)), we have

$$ $\displaystyle\bm{z}_{i}^{k+1}=$ $\displaystyle\mathrm{prox}_{r,\mathrm{Diag}(\bm{p}_{i}^{k})}\left(\bm{x}_{i}^{k}-\bm{p}_{i}^{k}\odot(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{y}_{i}^{k})\right)$ $\displaystyle=$ $\displaystyle\mathop{\arg\min}_{\bm{x}\in\mathbb{R}^{d}}r(\bm{x})+\frac{1}{2}\|\bm{x}-\bm{x}_{i}^{k}+\bm{p}_{i}^{k}\odot(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{y}_{i}^{k})\|_{(\mathrm{Diag}(\bm{p}_{i}))^{-1}}^{2},$ $$

which is equivalent to

$$ $\displaystyle 0\in\partial r(\bm{z}_{i}^{k+1})+(\mathrm{Diag}(\bm{p}_{i}))^{-1}\left(\bm{z}_{i}^{k+1}-\bm{x}_{i}^{k}+\bm{p}_{i}^{k}\odot(\nabla f_{i}(\bm{x}_{i}^{k})+\bm{y}_{i}^{k})\right)$ $\displaystyle\Leftrightarrow$ $\displaystyle-\nabla f_{i}(\bm{x}_{i}^{k})-\bm{y}_{i}^{k}-(\mathrm{Diag}(\bm{p}_{i}))^{-1}(\bm{z}_{i}^{k+1}-\bm{x}_{i}^{k})\in\partial r(\bm{z}_{i}^{k+1}).$ $$

Denote $\bm{g}_{i}^{k+1}=-\nabla f_{i}(\bm{x}_{i}^{k})-\bm{y}_{i}^{k}-(\mathrm{Diag}(\bm{p}_{i}))^{-1}(\bm{z}_{i}^{k+1}-\bm{x}_{i}^{k})\in\partial r(\bm{z}_{i}^{k+1})$, we have

$$ $\displaystyle r(\bm{x})\geq r(\bm{z}_{i}^{k+1})+\langle\bm{g}_{i}^{k+1},\bm{x}-\bm{z}_{i}^{k+1}\rangle,\quad\forall\bm{x}\in\mathbb{R}^{d}.$ (32) $$

Since $r\in\mathcal{F}(\mathbb{R}^{d})$ inherits lower semi-continuity, and $0<[\bm{p}_{i}^{k}]_{l}^{-1}\leq 1/m$, ([32](#A2.E32)) implies

$$ $\displaystyle r(\bm{x})\geq$ $\displaystyle\mathop{\lim\inf}_{k\rightarrow\infty}r(\bm{z}_{i}^{k+1})+\lim_{k\rightarrow\infty}\langle\bm{g}_{i}^{k+1},\bm{x}-\bm{z}_{i}^{k+1}\rangle$ $\displaystyle\geq$ $\displaystyle r(\bm{x}^{\star})+\langle-\nabla f_{i}(\bm{x}^{\star})-\bm{y}_{i}^{\star},\bm{x}-\bm{x}^{\star}\rangle,\quad\forall\bm{x}\in\mathbb{R}^{d}.$ $$

As a result, $\bm{g}_{i}^{\star}:=-\nabla f_{i}(\bm{x}^{\star})-\bm{y}_{i}^{\star}\in\partial r(\bm{x}^{\star})$. The last thing is to show $\bm{x}^{\star}\in\arg\min_{\bm{x}\in\mathbb{R}^{d}}f(\bm{x})+r(\bm{x})$. Adding ([23](#S4.E23)) for all $i\in\mathcal{V}$, we have

$$ $\displaystyle\sum_{i\in\mathcal{V}}\bm{y}_{i}^{k+1}=$ $\displaystyle\sum_{i\in\mathcal{V}}\left(\bm{y}_{i}^{k}+\sum_{j\in\mathcal{N}(i)}\bm{p}_{i,j,1}^{k}\odot(\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1})\right)$ $\displaystyle=$ $\displaystyle\sum_{i\in\mathcal{V}}\bm{y}_{i}^{k}+\sum_{\{i,j\}\in\mathcal{E}}(\bm{p}_{i,j,1}^{k}-\bm{p}_{j,i,1}^{k})\odot(\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1})$ $\displaystyle=$ $\displaystyle\sum_{i\in\mathcal{V}}\bm{y}_{i}^{k},\quad\forall k\geq 0,$ (33) $$

where the last equality uses $\bm{p}_{i,j,1}^{k}=\bm{p}_{j,i,1}^{k}$. By initialization $\bm{y}_{i}^{0}=\bm{0}_{d}$, ([33](#A2.E33)) implies

$$ $\displaystyle\sum_{i\in\mathcal{V}}\bm{y}_{i}^{\star}=\lim_{k\rightarrow\infty}\sum_{i\in\mathcal{V}}\bm{y}_{i}^{k}=\lim_{k\rightarrow\infty}\bm{0}_{d}=\bm{0}_{d},$ $$

thus

$$ $\displaystyle\partial r(\bm{x}^{\star})\ni\frac{1}{n}\sum_{i=1}^{n}\bm{g}_{i}^{\star}=\frac{1}{n}\sum_{i=1}^{n}-\nabla f_{i}(\bm{x}^{\star})-\bm{y}_{i}^{\star}=-\nabla f(\bm{x}^{\star}),$ $$

which is exactly $\bm{x}^{\star}\in\arg\min_{\bm{x}\in\mathbb{R}^{d}}f(\bm{x})+r(\bm{x})$.
∎

###### Proof.

## Appendix C Illustration of MiLoDo framework

To better understanding the two components, i.e., the MiLoDo update rules ([22](#S4.E22))-([24](#S4.E24)) and the LSTM parameterization ([25](#S4.E25))-([27](#S4.E27)) and how they make up the whole MiLoDo optimizer, we illustrate the interaction beween them in Fig. [10](#A3.F10).

Figure: Figure 10: The interaction between MiLoDo update rules ([22](#S4.E22))-([24](#S4.E24)) and the LSTM parameterization ([25](#S4.E25))-([27](#S4.E27)).
Refer to caption: /html/2410.01700/assets/structure.png

## Appendix D Robust implementation of decentralized algorithms

Common implementation. Traditional decentralized algorithms like Prox-ED, PG-EXTRA, Prox-ATC use a doubly-stochastic gossip matrix $W$ for information aggregation. A common implementation of the aggregation step $\tilde{\bm{X}}=W\bm{X}$ is to compute

$$ $\displaystyle\tilde{\bm{x}}_{i}=\sum_{j\in\mathcal{N}(i)}w_{ij}\bm{x}_{j}$ (34) $$

on each node $i$. However, ([34](#A4.E34)) is not a robust implementation. As illustrated in Fig. [11](#A4.F11), using ([34](#A4.E34)) and the same hyperparameter settings, Prox-ED with FP32 fails to converge to the desired precision while that with FP64 succeeds.

Robust implementation. When represented with FP32, the elements in matrix $W$ tend to have bigger noise, which largely violates the row-stochastic property. Continually applying such an inexact estimation of $W$ is the major reason behind the failure of the common implementation. This motivates us to consider the following equivalent implementation:

$$ $\displaystyle\tilde{\bm{x}}_{i}=\bm{x}_{i}-\sum_{j\in\mathcal{N}(i)}w_{ij}(\bm{x}_{i}-\bm{x}_{j}).$ (35) $$

Implementation ([35](#A4.E35)) is more robust as it maintains the row-stochastic property of the gossip matrix no matter how much noise is added to $W$ by the low presentation precision. As illustrated in Fig. [11](#A4.F11), Prox-ED with robust implementation successfully converges to the desired precision under the same hyperparameter settings.

Figure: Figure 11: Testing results on synthetic LASSO$(10,50,10,0)$ of Prox-ED with different implementation under varying precision settings. ’Common’ and ’Robust’ represent common implementation ([34](#A4.E34)) and robust implementation ([35](#A4.E35)), respectively.
Refer to caption: /html/2410.01700/assets/x18.png

Robustness of MiLoDo. It’s worth noting that, the structured update rules of MiLoDo does not depend on the doubly-stochastic matrix $W$. The utilization of term $\bm{z}_{i}^{k}-\bm{z}_{j}^{k}$ when aggregating neighboring information in update rules of MiLoDo is also similar to the robust implementation ([35](#A4.E35)). In practice, we observe that applying FP32 in our experiments does not affect the exact convergence of MiLoDo-trained optimizers.

## Appendix E Experimental specifications

### E.1 Training strategies

Initialization strategies.
We consider two initialization strategies for MiLoDo training: random and special initialization. In random initialization, all learnable parameters $\{\bm{\theta}_{M,i},\bm{\theta}_{S,i},\bm{\theta}_{U,i}\}_{i\in\mathcal{V}}$ are randomly initialized using PyTorch defaults. In special initialization, parameters are initialized to mimic traditional decentralized algorithms by setting the weights of the final affine layers to zero and biases to desired output values. Specifically, given the gossip matrix $W=(w_{ij})_{n\times n}$ and learning rate $\gamma$ utilized in Exact-Diffusion, biases in the final affine layers for $\bm{p}_{i}^{k}$, $\tilde{\bm{p}}_{i,j,1}^{k}$, $\bm{p}_{i,j,2}^{k}$ are initialized as $\gamma$, $\ln(w_{ij}/(2\gamma))$, $w_{ij}/2$, respectively. Applying $\ln(\cdot)$ accounts for Exponential activation.

We would like to remark that each of the two initialization strategies has its pros and cons. With random initialization, the objective function value is likely to blow up quickly, leading to excessively large gradients or meaningless values (e.g., inf/NaNs).
With special initialization, it may be too close to local minima, such that MiLoDo might not gain enough advantage over handcrafted algorithms.

Multi-stage training. As discussed above, random initialization of the MiLoDo optimizer results in numerical instability during training. To address this issue, we initially teach the model to optimize within a few iterations by using a short training length such as $(K_{T},K)=(5,10)$. As the model starts exhibiting desired behaviors, such as $\sum_{k=1}^{5}f(\bar{\bm{x}}^{k})>\sum_{k=6}^{10}f(\bar{\bm{x}}^{k})$, we progressively increase the training length. This iterative process is repeated across several stages until reaching a training length of $(K_{T},K)=(20,100)$. Empirically, employing multi-stage training also enhances performance for special initialization. With multi-stage training, both initialization approaches yield MiLoDo-trained optimizers with comparable performance, prompting our focus on special initialization due to its reduced warmup stages.

### E.2 Target problems

LASSO regression. Decentralized LASSO regression problem with shape $(n,d,N,\lambda)$ is defined as:

$$ $\displaystyle\min_{\bm{x}\in\mathbb{R}^{d}}\quad\frac{1}{2n}\sum_{i=1}^{n}\|\bm{A}_{i}\bm{x}-\bm{b}_{i}\|_{2}^{2}+\lambda\|\bm{x}\|_{1},$ $$

where $\bm{A}_{i}\in\mathbb{R}^{N\times d}$ and $\bm{b}_{i}\in\mathbb{R}^{N}$ are kept on nodes $i$ out of a total of $n$ nodes.
To generate LASSO optimizees with shape $(n,d,N,\lambda)$, we first sample $\bm{A}\in\mathbb{R}^{nN\times d}$ and a vector $\bm{x}^{\star}\in\mathbb{R}^{d}$ from normal distribution. Then, we pick 75% of $\bm{x}^{\star}$’s entries with the smallest magnitude and reset them to zero. Afterwards, we generate $\bm{b}=\bm{A}\bm{x}^{\star}+\epsilon\bm{z}$, where $\epsilon=0.1$ is the noise scale and $\bm{z}\in\mathbb{R}^{nN}$ is sampled from standard Gaussian. Finally, we distribute $\bm{A}$ and $\bm{b}$ evenly to each node so that each $\bm{A}_{i}\in\mathbb{R}^{N\times d}$ and $\bm{b}_{i}\in\mathbb{R}^{N}$.

Logistic regression. Decentralized logistic regression problem with $\ell_{1}$-regularization and shape $(n,d,N,\lambda)$ is defined as:

$$ $\displaystyle\min_{\bm{x}\in\mathbb{R}^{d}}\quad\frac{1}{n}\sum_{i=1}^{n}\left[\frac{1}{N}\sum_{j=1}^{N}b_{ij}\ln\left(1+\exp(-\bm{a}_{ij}^{\top}\bm{x})\right)+(1-b_{ij})\ln\left(1+\exp(\bm{a}_{ij}^{\top}\bm{x})\right)\right]+\lambda\|\bm{x}\|_{1},$ $$

where $\bm{A}_{i}=(\bm{a}_{i1}^{\top},\cdots\bm{a}_{iN}^{\top})^{\top}\in\mathbb{R}^{N\times d}$ and $\bm{b}_{i}=(b_{i1},\cdots,b_{iN})^{\top}\in\{0,1\}^{N}$. To generate synthetic logistic regression optimizees with shape $(n,d,N,\lambda)$, we first sample $\bm{A}\in\mathbb{R}^{nN\times d}$ and $\bm{x}^{\star}\in\mathbb{R}^{d}$ from normal distribution. Then we pick 75% of $\bm{x}^{\star}$’s entries with the smallest magnitude and reset them to zero. Afterwards, we generate $\bm{b}=(b_{1},\cdots,b_{nN})^{\top}$ by $b_{i}=\mathbbm{1}_{\{\bm{a}_{i}^{\top}\bm{x}^{\star}\geq 0\}}$. Finally, we distribute $\bm{A}$ and $\bm{b}$ evenly to each node so that each $\bm{A}_{i}\in\mathbb{R}^{N\times d}$ and $\bm{b}_{i}\in\{0,1\}^{N}$.

MLP training. We consider a decentralized MLP training problem using MNIST dataset. The model structure is illustrated as in Fig. [13](#A5.F13). The total number of trainable parameters in the MLP is 13002. The optimizees are constructed by randomly selecting data from MNIST’s training dataset for all nodes.

ResNet training. We consider a decentralized ResNet training problem using CIFAR-10 dataset. The model structure is illustrated as in Fig. [13](#A5.F13). The total number of trainable parameters in the ResNet model is 78042. The optimizees are constructed by randomly selecting data from CIFAR-10’s training dataset for all nodes.

Figure: Figure 12: MLP model structure.
Refer to caption: /html/2410.01700/assets/mlp-3.png

### E.3 Implementation details

Model structure. We use the same model structure throughout our experiments. Specifically, $\phi_{M,i}$ has input dimension 2 and output dimension 1 with ReLU activation, $\phi_{S,i}$ has input dimension $|\mathcal{N}(i)|$ and output dimension $|\mathcal{N}(i)|$ with Exponential activation. $\phi_{U,i}$ has input dimension $|\mathcal{N}(i)|$ and output dimension $|\mathcal{N}(i)|$ with ReLU activation. We use ReLU activation in the middle of the 2-layer MLP. The hidden/output dimensions of the LSTM cells, input/hidden/output dimensions of the MLP are all set to 20.

Training details. In our experiments, we employ special initialization and a multi-stage training strategy. As described in Sec. [5](#S5), we continually train MiLoDo in five stages with training lengths $(K_{T},K)=(5,10)$, $(10,20)$, $(20,40)$, $(40,80)$ and $(20,100)$ by Adam with learning rate 5e-04, 1e-04, 5e-05, 1e-05, 1e-05, for 20, 10, 10, 10, 5 epochs, respectively. Throughout all stages, the Adam optimizer is configured with momentum parameters $(\beta_{1},\beta_{2})=(0.9,0.999)$ and the batch size is fixed to 32.

LASSO with real data. To generate LASSO$(10,200,10,0.05)$ from BSDS500(Martin et al., ) dataset, we first extract a $10\times 10$ patch from testing images and flatten to vector $\bm{b}\in\mathbb{R}^{100}$. We normalize $\bm{b}$ by subtracting the mean. Afterwards, we conduct K-SVD(Aharon et al., ) to obtain $\bm{A}\in\mathbb{R}^{100\times 200}$. Finally, we distribute $\bm{A}$ and $\bm{b}$ evenly to each node so that each $\bm{A}_{i}\in\mathbb{R}^{10\times 200}$ and $\bm{b}_{i}\in\mathbb{R}^{10}$. We generate a total of 1000 instances as the testing set in the experiments.

Construction of meta training set. As illustrated in Sec. [5](#S5), the meta training set consists of synthetic LASSO problems with 20 different shapes: $(10,500,N,0.1)$ where $N\in\{5,10,15,\cdots,100\}$. We generate 64 distinct problem instances for each shape, hence 1280 instances in total.

Evaluation metric. We evaluate solution $\bm{X}=[\bm{x}_{1}^{\top},\bm{x}_{2}^{\top},\cdots,\bm{x}_{n}^{\top}]^{\top}$ of decentralized problem ([4](#S3.E4)) via loss $f(\bar{\bm{x}})+r(\bar{\bm{x}})$ and consensus error $\frac{1}{n}\sum_{i=1}^{n}\|\bm{x}_{i}-\bar{\bm{x}}\|_{2}$, where $\bar{\bm{x}}=\frac{1}{n}\sum_{i=1}^{n}\bm{x}_{i}$. All testing curves display averaged performance on 512 instances, except for problems with over 10,000 dimensions which are highly time-consuming to test. For those high-dimensional problems, we display testing performance on a single instance chosen randomly, as results on other instances are quite similar.

Implementation of baseline algorithms. Following Appendix [D](#A4), we have the following robust implementation for the considered baselines, where the learning rate $\gamma$ is manually tuned optimal for each experiment, and we use $W=(w_{ij})_{n\times n}$ with $w_{ij}=1/3\cdot\mathbbm{1}_{\{i=j,\mbox{ or }\{i,j\}\in\mathcal{E}\}}$ as the doubly-stochastic goissp matrix for the ring topology.

- •
Prox-DGD. Initialized with $\bm{x}_{i}^{0}=\bm{0}_{d}$, Prox-DGD uses the following update rules:
$\displaystyle\bm{z}_{i}^{k+1}=$
$\displaystyle\bm{x}_{i}^{k}-\gamma\nabla f_{i}(\bm{x}_{i}^{k}),$
$\displaystyle\bm{x}_{i}^{k+1}=$
$\displaystyle\mathrm{prox}_{\gamma r}\left(\bm{z}_{i}^{k+1}-\sum_{j\in\mathcal{N}(i)}w_{ij}(\bm{z}_{i}^{k+1}-\bm{z}_{j}^{k+1})\right).$
- •
Prox-ATC. Initialized with $\bm{x}_{i}^{0}=\tilde{\bm{y}}_{i}^{0}=\bm{z}_{i}^{0}$, Prox-ATC uses the following update rules:
$\displaystyle\bm{z}_{i}^{k+1}=$
$\displaystyle\bm{x}_{i}^{k}-\gamma\nabla f_{i}(\bm{x}_{i}^{k}),$
$\displaystyle\tilde{\bm{z}}_{i}^{k+1}=$
$\displaystyle\tilde{\bm{y}}_{i}^{k}-\bm{z}_{i}^{k+1}+\bm{z}_{i}^{k},$
$\displaystyle\bm{y}_{i}^{k+1}=$
$\displaystyle 2\tilde{\bm{y}}_{i}^{k}-\tilde{\bm{z}}_{i}^{k+1}+\sum_{j\in\mathcal{N}(i)}w_{ij}(\tilde{\bm{z}}_{i}^{k+1}-\tilde{\bm{z}}_{j}^{k+1}),$
$\displaystyle\tilde{\bm{y}}_{i}^{k+1}=$
$\displaystyle\bm{y}_{i}^{k+1}-\sum_{j\in\mathcal{N}(i)}w_{ij}(\bm{y}_{i}^{k+1}-\bm{y}_{j}^{k+1}),$
$\displaystyle\bm{x}_{i}^{k+1}=$
$\displaystyle\mathrm{prox}_{\gamma r}(\tilde{\bm{y}}_{i}^{k+1}).$
- •
PG-EXTRA. Initialized with $\bm{x}_{i}^{0}=\bm{0}_{d}$, PG-EXTRA uses the following update rules:
$\displaystyle\bm{z}_{i}^{k+1}=$
$\displaystyle\bm{x}_{i}^{k}-\sum_{j\in\mathcal{N}(i)}w_{ij}(\bm{x}_{i}^{k}-\bm{x}_{j}^{k})-\gamma\nabla f_{i}(\bm{x}_{i}^{k}),$
$\displaystyle\tilde{\bm{z}}_{i}^{k+1}=$
$\displaystyle\begin{cases}\bm{z}_{i}^{k+1},&\mbox{if }k=0,\\
\bm{z}_{i}^{k+1}+\tilde{\bm{z}}_{i}^{k}-\bm{x}_{i}^{k-1}+\frac{1}{2}\sum_{j\in\mathcal{N}(i)}w_{ij}(\bm{x}_{i}^{k-1}-\bm{x}_{j}^{k-1})+\gamma\nabla f_{i}(\bm{x}_{i}^{k-1}),&\mbox{if }k>0,\end{cases}$
$\displaystyle\bm{x}_{i}^{k+1}=$
$\displaystyle\mathrm{prox}_{\gamma r}(\tilde{\bm{z}}_{i}^{k+1}).$
- •
Prox-ED. Initialized with $\bm{x}_{i}^{0}=\tilde{\bm{y}}_{i}^{0}=\bm{z}_{i}^{0}=\bm{0}_{d}$, Prox-ED uses the following update rules:
$\displaystyle\bm{z}_{i}^{k+1}=$
$\displaystyle\bm{x}_{i}^{k}-\gamma\nabla f_{i}(\bm{x}_{i}^{k}),$
$\displaystyle\bm{y}_{i}^{k+1}=$
$\displaystyle\tilde{\bm{y}}_{i}^{k}+\bm{z}_{i}^{k+1}-\bm{z}_{i}^{k},$
$\displaystyle\tilde{\bm{y}}_{i}^{k+1}=$
$\displaystyle\bm{y}_{i}^{k+1}-\frac{1}{2}\sum_{j\in\mathcal{N}(i)}w_{ij}(\bm{y}_{i}^{k+1}-\bm{y}_{j}^{k+1}),$
$\displaystyle\bm{x}_{i}^{k+1}=$
$\displaystyle\mathrm{prox}_{\gamma r}(\tilde{\bm{y}}_{i}^{k+1}).$

Computational resources. We conduct all the experiments within a single NVIDIA A100 GPU server with a GPU memory of 80G.

### E.4 Additional results

Training on logistic regression. Fig. [14](#A5.F14) displays the in-distribution testing results of MiLoDo optimizer trained on a specialized dataset including 512 synthetic Logistic$(10,50,100,0.1)$ optimizees. Fig. [15](#A5.F15) displays the testing results of MiLoDo optimizer trained on a specialized dataset including 512 real data Logistic$(10,14,100,0.1)$ optimizees using Census Income (Kohavi, ) dataset.

More testing results of MiLoDo optimizer trained on the meta training set.
As a supplement to the results in Sec. [5](#S5), Fig. [16](#A5.F16), we further tests MiLoDo trained on the meta training set on synthetic LASSO$(10,20000,1000,0.1)$. While trained on non-smooth optimizees only, the MiLoDo-trained optimizer is consistently fast in solving smooth optimization problems such as linear regression, as illustrated in Fig. [17](#A5.F17).

Testing results of MiLoDo optimizer trained on more complex topologies. Beyond the findings presented in Sec.[5](#S5) and Fig.[9](#S5.F9) (left), further tests were conducted on commonly used topologies. Fig[18](#A5.F18) demonstrates that MiLoDo optimizer exhibits consistent performance, achieving a 2 to 3 times acceleration, which highlights its scalability and robustness across various topologies.

Testing results of MiLoDo optimizer trained on a larger network. Extending the analyses discussed in Sec.[5](#S5) and illustrated in Fig.[9](#S5.F9) (right), additional experiments were carried out on networks with 100 nodes. As depicted in Fig.[20](#A5.F20), MiLoDo optimizer maintained a high level of effectiveness, delivering a $2\times$ to $3\times$ speedup. This not only confirms the optimizer’s efficiency but also highlights its scalability and robustness in larger networks.

Figure: Figure 14: MiLoDo optimizer trained on synthetic Logistic$(10,50,100,0.1)$ and tested on unseen Logistic$(10,50,100,0.1)$ instances.
Refer to caption: /html/2410.01700/assets/x19.png

Figure: Figure 15: MiLoDo optimizer trained on Logistic$(10,14,100,0.1)$ with Census Income (Kohavi, ) dataset and tested on Logistic$(10,14,100,0.1)$ with unseen data in Census Income dataset.
Refer to caption: /html/2410.01700/assets/x21.png

Figure: Figure 16: MiLoDo trained on meta learning set and tested on LASSO$(10,20000,1000,0.1)$.
Refer to caption: /html/2410.01700/assets/x23.png

Figure: Figure 17: MiLoDo optimizer trained on meta training set and tested on linear regression problems as LASSO$(10,15000,1000,0)$.
Refer to caption: /html/2410.01700/assets/x25.png

Figure: (a)
Refer to caption: /html/2410.01700/assets/topo-grid.png

Figure: (a)
Refer to caption: /html/2410.01700/assets/x31.png

Figure: Figure 20: MiLoDo optimizer trained on a large network with 100 nodes, and tested on LASSO(10, 300, 10, 0.1).
Refer to caption: /html/2410.01700/assets/x33.png

### E.5 Ablation studies

Ablation studies on the mixing matrices for baseline algorithms. The adaptive preconditioners and mixing weights are critical to MiLoDo’s performance gain. To better address MiLoDo’s advantages, we conduct ablation experiments on the mixing matrices used in the baseline methods. Fig. [19](#A5.F19) demonstrates that the performance of using strategically designed and fixed weights (1/3 in our experiments) are almost the same , which provides a stronger validation of MiLoDo ’s advantages.

Ablation on base update rules ([12](#S3.E12))-([14](#S3.E14)). We specify detailed experimental setups for directly learning optimizers from the base update rules ([12](#S3.E12))-([14](#S3.E14)), as discussed in Sec. [5](#S5). For simplicity, we use $r\equiv 0$ so that the implicit rule in ([12](#S3.E12)) can be explicitly modeled as

$$ $\displaystyle\bm{z}_{i}^{k+1}=\bm{x}_{i}^{k}-\bm{m}_{i}^{k}(\nabla f_{i}(\bm{x}_{i}^{k}),\bm{y}_{i}^{k};\bm{\theta}_{i,1}).$ $$

Without coordinate-wise structures, the scale of the neural network has to be correlated with the optimizees’ dimension. Consequently, we fix the problem dimension $d=10$ and use LASSO$(10,10,5,0)$ as the training and testing optimizees. We parameterize each of the base update rules with a LSTM model consists of a single LSTM cell and a 2-layer MLP with ReLU activation. The input sizes of the LSTM models are $20$, $10|\mathcal{N}(i)|$, $10|\mathcal{N}(i)|$ for $\bm{m}_{i}$, $\bm{s}_{i}$, $\bm{u}_{i}$, respectively. The output sizes are 10 according to the problem dimension. All hidden dimensions in the LSTM cells and MLPs are set to 100. We use random initialization and multi-stage training strategy similar to MiLoDo to train the parameterized base update rules.

### E.6 Hyperparameter settings

We specify the manually-tuned optimal learning rates of baseline algorithms for all the experiments in Table [2](#A5.T2).

**Table 2: Optimal learning rates of baseline algorithms chosen in different experiments.**
| Experiment | Prox-ED | PG-EXTRA | Prox-ATC | Prox-DGD | DAPG | ODAPG |
| --- | --- | --- | --- | --- | --- | --- |
| LASSO$(10,300,10,0.1)$ | 0.03 | 0.02 | 0.025 | 0.04 | 0.01 | 0.02 |
| LASSO$(10,30000,1000,0.1)$ | 0.03 | 0.02 | 0.025 | 0.04 | 0.01 | 0.02 |
| LASSO$(10,200,10,0.1)$ | 0.05 | 0.04 | 0.045 | 0.05 | 0.02 | 0.03 |
| LASSO$(10,20000,1000,0.1)$ | 0.05 | 0.04 | 0.045 | 0.05 | / | / |
| LASSO$(10,15000,1000,0.0)$ | 0.08 | 0.05 | 0.085 | 0.09 | / | / |
| Logistic $(10,50,100,0.1)$ | 1.0 | 0.8 | 0.4 | 1.0 | / | / |
| Logistic $(10,14,100,0.1)$ | 1.9 | 1.7 | 1.8 | 2.0 | / | / |
| MLP$(10,13002,5000,0)$ | 0.09 | 0.06 | 0.06 | 0.05 | 0.03 | 0.055 |
| ResNet$(5,78042,5000,0)$ | 0.1 | 0.07 | 0.08 | 0.05 | 0.05 | 0.07 |

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
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.
soundness: 2
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.
summary: This paper proposes distributed algorithm called MiLoDo that solve consensus-type problems. Some simulations are run to demonstrate the speed of convergence.
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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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.
soundness: 3
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.
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).
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.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: Please see the weaknesses.
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: see weakness.
soundness: 2
strengths: The effort of extending Learning to Optimize framework to the decentralized setting is interesting.
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.
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.

Meta-review: 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.

Decision: Reject