DIMIX: DIminishing MIXing for Sloppy Agents

This paper is motivated by stochastic learning problems in which the goal is to solve

where $\ell:\mathbb{R}^{d}\times\mathbb{R}^{p}\rightarrow\mathbb{R}$ is a loss function, $\mathbf{x}\in\mathbb{R}^{1\times d}=\mathbb{R}^{d}$ is the decision/optimization row vector, and $\xi$ is a random vector taking values in $\mathbb{R}^{p}$ that is drawn from an unknown underlying distribution $\mathcal{P}$. One of the key practical considerations that renders (1) as a challenging task is that the underlying distribution $\mathcal{P}$ is often unknown. Instead, we have access to $N$ independent realizations of $\xi$ and focus on solving the corresponding ERM problem which is given by

where $f(\mathbf{x})$ is the empirical risk with respect to the data points $\mathcal{D}=\{\xi_{1},\ldots,\xi_{N}\}$. We assume that $\ell(\cdot,\cdot)$ is a non-convex loss function, which potentially results in a non-convex function $f(\cdot)$.

In distributed optimization, we have a network consisting of $n$ computing nodes (agents, or workers), where each node $i$ observes a non-overlapping subset of ${m_{i}=r_{i}N}$ data points, denoted by $\mathcal{D}_{i}=\{\xi^{i}_{1},\ldots,\xi^{i}_{m_{i}}\}$, where ${\mathcal{D}=\mathcal{D}_{1}\cup\cdots\cup\mathcal{D}_{n}}$. Here, $r_{i}$ represents the fraction of the data that is processed at node $i\in[n]$. Note that the vector $\mathbf{r}=(r_{1},\dots,r_{n})$ is a strictly positive stochastic vector, i.e., $r_{i}>0$ and $\sum_{i=1}^{n}r_{i}=1$. Thus, the ERM problem in (2) can be written as the minimization of the weighted average of local empirical risk functions $f_{i}$ for all nodes $i\in[n]$ in the network, i.e.,

where $f_{i}(\mathbf{x}):=\frac{1}{m_{i}}\sum_{\xi\in\mathcal{D}_{i}}\ell(\mathbf{x},\xi)=\frac{1}{m_{i}}\sum_{j=1}^{m_{i}}\ell(\mathbf{x},\xi_{j}^{i})$. We can rewrite the ERM problem in (3) as a distributed consensus optimization problem, given by

Consider an $n\geq 2$ agents that are connected through a time-varying network. We represent this network at time $t\geq 1$ by the directed graph $\mathcal{G}(t)=([n],\mathcal{E}(t))$, where the vertex set $[n]$ represents the set of agents and the edge set ${\mathcal{E}(t)\subseteq\{(i,j):i,j\in[n]\}}$ represents the set of links at time $t$. At each discrete time $t\geq 1$, agent $i$ can only send information to its (out-) neighbors in $\mathcal{E}(t)$, i.e., all $j\in[n]$ with $(i,j)\in\mathcal{E}(t)$.

To discuss our algorithm (DIMIX) for solving (4) distributively, let us first discuss its general structure and the required information at each node for its execution. In this algorithm, at each iteration $t\geq 1$, agent $i\in[n]$ updates its estimate $\mathbf{x}_{i}(t)\in\mathbb{R}^{d}$ of an optimizer of (3). To this end, it utilizes the gradient information of its own local cost function $f_{i}(\mathbf{x})$ as well as a noisy/lossy average of its current neighbors estimates, denoted by ${\hat{\mathbf{x}}_{i}(t):=\sum_{j=1}^{n}W_{ij}(t)\mathbf{x}_{j}(t)+\mathbf{e}_{i}(t)}$. Here, $W(t)$ is a row-stochastic matrix that is consistent with the underlying connectivity network $\mathcal{G}(t)$ (i.e., $W_{ij}(t)>0$ only if $(j,i)\in\mathcal{E}(t)$) and $\mathbf{e}_{i}(t)\in\mathbb{R}^{d}$ is a random noise vector. Later, in Section 2.4 we discuss several noisy and lossy information sharing architectures (quantization and noisy communication) that fit in this broad information structure. Now we are ready to discuss the DIMIX algorithm. In this algorithm, using the information available to agent $i$ at time $t$, agent $i$ updates its current estimate by computing a diminishing weighted average of its own state and the noisy average of its neighbors’ estimates, and moves along its local gradient. More formally, the update rule at node $i\in[n]$ is given by

where $\alpha(t)=\frac{\alpha_{0}}{(t+\tau)^{\nu}}$ and $\beta(t)=\frac{\beta_{0}}{(t+\tau)^{\mu}}$ for some $\mu,\nu\in(0,1)$ are the diminishing step-sizes of the algorithm, and $\tau\geq 0$ is an arbitrary shift, that is introduced to accelerate the finite-time performance of the algorithm. The description of DIMIX is summarized in Algorithm 1. For notational simplicity, let

Since $\hat{\mathbf{x}}_{i}(x)=\sum_{j=1}^{n}W_{ij}(t)\mathbf{x}_{j}(t)+\mathbf{e}_{i}(t)$, we can rewrite the update rule in (5) in the matrix format

Let us discuss some important aspects of the above update rule. Note that both $\alpha(t)$ and $\beta(t)$ are diminishing step-sizes. If $\beta(t)=\beta_{0}<1$ and $\alpha(t)=\alpha_{0}<1$ are both constants, then the dynamics in (7) reduces to the algorithm proposed in [26] for both convex and non-convex cost functions. Alternatively, if $\beta(t)=\beta_{0}<1$ is a constant sequence and $E(t)=\mathbf{0}$ for all $t\geq 1$, (7) would be reduced to the averaging-based distributed optimization with diminishing steps-sizes (for gradients), which is introduced and studied in [18] for local convex cost functions $f_{i}(\mathbf{x})$. The newly introduced time-scale/step-size $\beta(t)$ suppresses the incoming noise $\mathbf{e}_{i}(t)$ from the neighboring agents. However, $\beta(t)$ also suppresses the incoming signal level $\sum_{j=1}^{n}W_{ij}(t)\mathbf{x}_{j}(t)$ at each node $i$. This casts a major technical challenge for establishing convergence-to-consensus guarantees for the algorithm over time-varying networks.

On the other hand, the diminishing step-size for the gradient update is $\hat{\alpha}(t)=\alpha(t)\beta(t)$. We chose to represent our algorithm in this way to ensure that the local mixing (consensus) scheme is operated on a faster time-scale than the gradient descent.

In order to provide performance guarantees for DIMIX, we need to assume certain regularity conditions on (i) the statistics of the (neighbors’ averaging) noise process $\{E(t)\}$, (ii) the mixing properties of the weight sequence $\{W(t)\}$, and (iii) the loss function $\ell(\cdot,\cdot)$.

First, we discuss our main assumption on the noise sequence $\{E(t)\}$.

Note that the natural filtration of the random process $\{X(t)\}$ is one choice for $\{\mathcal{F}_{t}\}$. In this case, (8) reduces to $\mathbb{E}\left[\mathbf{e}_{i}(t)\!\mid\!\!X(1),\ldots,X(t)\right]\!=\!0$ and $\mathbb{E}\left[\|\mathbf{e}_{i}(t)\|^{2}\!\mid\!X(1),\ldots,X(t)\right]\!\leq\gamma.$

Next, we discuss the main assumption on the network connectivity, or more precisely, the mixing of information over the time-varying network.

Finally for the loss function $\ell(\cdot,\cdot)$, we make the following assumption.

Note that Assumption 3-b implies a homogeneous sampling, i.e., each agent draws i.i.d. sample points from a data batch. In a related work [35], a stronger assumption has been considered which allows for heterogeneous data samples.

Here, we characterize the convergence rates of our algorithm for the $K$-smooth non-convex loss functions. More precisely, we establish a rate for the temporal average of the expected norm of the gradients for various choice of the time-scale parameters $\nu,\mu$.

The noisy information in (5) is very general, and captures several models of imperfect data used in practice and/or theoretical studies. A rather general architecture that leads to such noisy/lossy estimates is demonstrated in Figure 1: Once the estimate $\mathbf{x}_{j}(t)$ of node $j$ at iteration $t$ is evaluated, node $j$ may apply an operation (such as compression/sparsification or quantization) on its own model to generate $\tilde{\mathbf{x}}_{j}(t)$. This data will be sent over a potentially noisy communication channel, and a neighbor $i$ will receive a corrupted version of $\tilde{\mathbf{x}}_{j}(t)$, say $\tilde{\mathbf{y}}_{i,j}(t)$ from every neighbor node $j$. Upon receiving such channel outputs from all its neighbors, node $i$ computes their weighted average $\mathbf{y}_{i}(t)=\sum_{j=1}^{n}W_{ij}(t)\tilde{\mathbf{y}}_{i,j}(t)$. This can be decoded to the approximate average model $\hat{\mathbf{x}}_{i}(t)$. In the following we describe three popular frameworks, in which each node $i$ can only use an imperfect neighbors’ average $\hat{\mathbf{x}}_{i}(t)$ to update its estimate. It is worth emphasizing that these are just some examples that lie under the general model in (5).

Note that in both examples above $\mathbf{e}_{i}(t)$ and $\mathbf{e}_{j}(t)$ might be correlated, especially when nodes $i$ and $j$ have common neighbor(s). However, this does not violate the conditions of Assumption 3.

We use regularized logistic regression to validation our algorithm on the benchmark dataset MNIST. $\triangleright$ Data and Experimental Setup. In this experiments we train a regularized logistic regression to classify the MNIST dataset over a fixed network. First, we generate a random (connected) undirected Erdös-Renyi graph with the edge probability $p_{c}=0.3$, which is fixed throughout this experiment. We also generate a doubly stochastic weight matrix ${A:=I-(d_{\max}+1)^{-1}L_{G}}$ where $L_{G}$ is the Laplacian matrix and $d_{\max}$ is the maximum degree of the graph. Finally, we use the time-invariant weight sequence $W=I+c\hat{\mathbf{r}}_{\min}\text{diag}^{-1}(\mathbf{r})(A-I)$, where ${\hat{\mathbf{r}}_{\min}=\min_{i\in[n]}r_{i}/(1-A_{ii})}$, $c=0.95$, and $\text{diag}(\mathbf{r})$ is a diagonal matrix with $r_{i}$ as its $i$th diagonal element. It can be verified that $W$ satisfies the conditions of Assumption 2, i.e., it is row-stochastic and satisfies $\mathbf{r}^{T}W=\mathbf{r}^{T}$. We implemented our algorithm with quantizer parameter $s=3$ and tuned the step-size parameters $(\alpha_{0},\nu^{\star})=(0.005,1/6)$, $(\beta_{0},\mu^{\star})=(0.6,1/2)$, with $\tau=2000$. For the fixed step-sizes, we used the termination time $T=7500$, which results in $\alpha(t)=0.001$ and $\beta(t)=0.01$ for all $t\geq 1$.

Here, we provide simulation results to demonstrate the performance of DIMIX on time-varying networks. For this, we conduct simulations on non-convex (CNN) and convex (linear regression) problems. For the sake of comparison, we apply our algorithm on a corresponding fixed network.