DJAM: distributed Jacobi asynchronous method for learning personal models
Inês Almeida João Xavier


Upper limits on the presence of central massive black holes in two ultra-compact dwarf galaxies in Centaurus A
    Jay Strader
    December 30, 2023
==============================================================================================================





Processing data collected by a network of agents often boils down to solving an optimization problem. The distributed nature of these problems calls for methods that are, themselves, distributed. While most collaborative learning problems require agents to reach a common (or consensus) model, there are situations in which the consensus solution may not be optimal. For instance, agents may want to reach a compromise between agreeing with their neighbors and minimizing a personal loss function.
We present DJAM, a Jacobi-like distributed algorithm for learning personalized models. This method is implementation-friendly: it has no hyperparameters that need tuning, it is asynchronous, and its updates only require single-neighbor interactions. We prove that DJAM converges with probability one to the solution, provided that the personal loss functions are strongly convex and have Lipschitz gradient. We then give evidence that DJAM is on par with state-of-the-art methods: our method reaches a solution with error similar to the error of a carefully tuned ADMM in about the same number of single-neighbor interactions.







§ LEARNING PERSONAL MODELS
 
Consider  n agents, each with a personal loss function:  f_i:^p→,  θ_i ↦ f_i( θ_i ), for agent i = 1, …, n. For example, f_i(θ_i) could be the loss of a model parameterized by θ_i on agent i's personal dataset.  The agents are the nodes of an undirected, connected network.


Each agent aims to find a model that minimizes both the mismatch with its neighbors' models and its personal loss.  More specifically, agents aim to solve

    min_θ_1,…,θ_n1/2∑_i<j^n W_ijθ_i-θ_j^2
    +∑_i=1^n f_i(θ_i),

where W = ( W_ij) ∈^n× n is a symmetric matrix that mirrors the topology of the network: W_ij ≥ 0 if agents i and j are connected in the network; W_ij=0 otherwise. 
The weight W_ij controls the degree of agreement we want between agents i and j: a larger W_ij enforces more similarity between the corresponding agents' models.

Problem (<ref>) can model a number of applications, including peer-network recomender systems, distributed (linear) classification <cit.>, opinion propagation, and field estimation <cit.>; the latter is discussed in Section <ref>. The weights W_ij may be chosen, for instance, based on the spatial distance between pairs of agents, or according to the similarity of their personal datasets.


 

Closest related works. Optimization problem (<ref>) has been addressed in <cit.>. For convex loss functions f_i that are quadratic, the authors suggest a distributed algorithm, which we refer to as Model Propagation Algorithm (MPA). In each round of MPA, an agent wakes up at random,  interacts with one of its neighbours, and both go back to sleep; the pattern repeats for the following rounds. MPA is an algorithm that is easy to implement because it is asynchronous (each agent has its own clock to wake up), has no parameter to tune,  and involves only single-neighbour interactions (the agent that wakes up does not need to coordinate message-passing with several neighbours). The authors in <cit.> prove that MPA converges to the solution of (<ref>) in expectation (mean-value), for quadratic loss functions; for these functions, the iterations of the method we propose coincide with those of MPA.
For more general loss functions, those authors suggest a different algorithm, based on ADMM, which needs parameter tuning to reach optimal performance. This ADMM-based algorithm for collaborative learning (CL-ADMM), will be compared with our algorithm in Section <ref>.

Problem (<ref>), with the same kind of asynchronous single-neighbour interactions, can also be tackled by the algorithm proposed in <cit.>. In the language of <cit.>, this corresponds to having agents deviate from the “rational" decision at each round (the rational decision would require each agent to interact will all its neighbors). For such “irrational" decisions, the authors show that, with probability one, the iterations of their algorithm will visit infinitely often a neighborhood of the solution of (<ref>), although the iterations may continually escape that neighborhood.
Finally, a recent follow-up on <cit.> is <cit.>, where a block coordinate descent method with broadcast communications is used to solve problem (<ref>).




Contributions. We show that a simple Jacobi-like distributed algorithm, which we call DJAM, can solve (<ref>) with the same kind of asynchronous single-neighbor interactions. DJAM, which can also be seen as a randomized block-coordinate method, has no parameters that need tuning. For continuously differentiable personal loss functions that are strongly convex and have Lipschitz gradient, that is, such that, for all i,

    (∇ f_i(x) - ∇ f_i(y))^T (x-y) ≥ m_ix-y^2

for some m_i > 0 and all x,y, and

    ∇ f_i(x)-∇ f_i(y)≤ M_ix-y

for some M_i ≥ 0 and all x, y, we show that DJAM converges to the solution of (<ref>) with probability one. The values of m_i and M_i are used for proving convergence but need not be known when implementing DJAM.

DJAM improves on MPA not only because it applies to a larger class of functions than quadratics, but also because it converges in a stronger sense: as the proof of Theorem <ref> ahead shows, the DJAM iterations are uniformly bounded; thus, the convergence in expectation in <cit.> follows by the dominated convergence theorem from our convergence with probability one. Our result only applies to a (somewhat) more restricted class of functions than the one of <cit.>, but our convergence mode is stronger than the one of <cit.>.
Also, unlike in <cit.>, our method does not require knowing the values of M_i upon implementation.






Other related work. Although <cit.> are the closest works that we are aware of, many other distributed algorithms solve variations of problem (<ref>). We now mention some representative work. 

A number of distributed algorithms allow agents to solve an underlying optimization problem by reaching consensus on the solution. They use techniques ranging from distributed (sub)gradient descent <cit.>, <cit.> to more elaborate techniques such as EXTRA <cit.>, distributed ADMM <cit.>, dual averaging <cit.>, and distributed Augmented Lagrangean (AL) <cit.>.  Some algorithms aim at more specific optimization tasks such as distributed lasso regression <cit.>, distributed SVMs <cit.>, and distributed RFVL networks <cit.>.
All of these methods aim at reaching consensus solutions—all agents converge to the same value. Conversely, in problem (<ref>), agents want to find different (personalized) values.


The related problem of network lasso is dealt with in <cit.>; however, the cost in <cit.> puts a strong emphasis on neighbouring models being exactly equal, whereas in our case we want them to be similar, but not necessarily equal. The methods proposed in <cit.> and <cit.> can tackle more general problems, but both require that agents communicate with all their neighbors before updating, while our method needs only communications between two agents at a time.
Problem (<ref>) is also referred to as multitask problem; this problem is solved in <cit.> for a more restricted class of personal losses than ours.








§ DJAM
 
A naive Jacobi-like approach to solve (<ref>) would work as follows: at each round t, one agent i, picked at random, would update its model according to

    θ_i(t+1) = min_θ_i1/2∑_k∈𝒩_i W_ikθ_i-θ_k(t)^2
    + f_i(θ_i),

where 𝒩_i is the set of neighbors of agent i.
This naive approach, however, has a major drawback: it requires that agent i communicates with all its neighbors—to receive their up-to-date models θ_k(t)—before updating its own model. Coordinating such message-passing, at each round, is cumbersome. A lighter scheme, involving only a single pair of agents at a time, is simpler to implement in practice, and requires fewer communications, at the expense of slowing down convergence. 

The key idea, which we borrow from <cit.>, is to have each agent i keep its own model Θ_i^i as well as (often outdated) versions of its neighbors' models, Θ_i^k for k∈𝒩_i. The versions of each pair of neighbors are updated whenever they communicate with each other. 
More specifically, at each round t, agent i wakes up and chooses a neighbor j∈𝒩_i to communicate with. They begin by exchanging information on their models, meaning that Θ_i^j(t+1)=Θ_j^j(t) and Θ_j^i(t+1) = Θ_i^i(t). All other variables remain unchanged. Afterwards, both agents update their own model via

    Θ_l^l(t+1) = min_θ_l1/2∑_k∈𝒩_l W_lkθ_l - Θ_l^k(t+1)^2 + f_l(θ_l)

for l∈{i,j}.

For the purpose of analyzing DJAM, we merge these two steps into a single one. Since the personal model Θ_i^i can be created at any time at agent i via (<ref>), it need not be stored. This means that, at round t of DJAM, two neighboring agents i and j will compute and share their own models with each other:

    Θ_i^j(t+1) = min_θ_j1/2∑_k∈𝒩_j W_jkθ_j - Θ_j^k(t)^2 + f_j(θ_j),

and similarly for Θ_j^i. Mind that the right-hand side of (<ref>) is computed by agent j and sent to agent i, who stores the result in the variable on the left-hand side of (<ref>).[Updates (<ref>) and (<ref>) are equivalent apart from a minor technicality: The values of Θ_i^j are equal for all t, while those of the Θ_i^i may be (finitely) delayed from one implementation to the other. This detail does not affect the validity of our results.]









§ PROOF OF CONVERGENCE FOR DJAM
 
We now prove that DJAM, the algorithm with updates given by (<ref>), converges with probability one to the solution of (<ref>). We omit some laborious (but otherwise painless) technical steps that would make the notation and proofs too lengthy.


Let ℰ be the set of edges of the network that links the agents. The network need not be fully connected: each agent is connected only to a subset of the remaining agents. We assume that at each round (A1) one edge of ℰ is chosen at random, independently of previous choices; and (A2) each edge in ℰ has a fixed, positive probability of being chosen. It is easy to verify that, under assumptions (A1) and (A2), each edge in ℰ is chosen infinitely often with probability one.



























The number of times a given edge (i,j) is chosen between rounds s and t (with s ≤ t) is a random variable defined as S_(i,j)(s,t) := ∑_τ = s^t Y_(i,j)(τ),
where Y_(i,j)(τ)=1 if edge (i,j) is chosen at round τ, and zero otherwise. 
We now define a useful family of stopping times ( T_m )_m ≥ 0. We let T_0 := 0 and, for m ≥ 0,

    T_m+1min{ t | S_(i,j)(T_m + 1,t)≥ 1, ∀ (i,j)∈ℰ}.

In words, T_m+1 is the first round after T_m  by which all edges have been chosen at least once. Assumptions (A1) and (A2) imply that any T_m is finite for any m. Addtionally, T_m→∞ as m→∞ with probability one.













We first state an important consequence of assumptions (<ref>) and (<ref>) on each personal loss function f_j.









 
Let w_j ∑_k∈𝒩_j W_jk, and take the function
F_j(x) = f_j(x) + 1/2w_jx^2. Note that ∇ F_j is a bijective map (with inverse map (∇ F_j)^-1) because, from standard convex theory, ∇ f_j is. Then, for any a and b,

    (∇ F_j)^-1(a) - (∇ F_j)^-1(b)≤(m_j + w_j)^-1a - b.




















Choose x = (∇ F_j)^-1(a) and y =(∇ F_j)^-1(b). Clearly, a - b = (∇ f_j(x) - ∇ f_j(y)) + w_j (x-y). Multiplying both sides of this equality by (x-y)^T yields (x-y)^T (a - b) ≥ (m_j + w_j) x-y^2, where the inequality is due to the strong convexity of f_j, property (<ref>). By the Cauchy-Schwartz inequality, x-ya-b≥(x-y)^⊤ (a-b),
and, thus, x-ya-b≥ (m_j + w_j) x-y^2. Inserting the definitions of x and y  yields a - b ≥ (m_j + w_j )   (∇ F_j)^-1(a) -  (∇ F_j)^-1(b).



We now give our main convergence result.


Let Θ^*=(Θ_1^*,…,Θ_n^*) be the  solution of (<ref>). Let Θ_i^j(t+1), i=1,…,n, j∈𝒩_i, be updated via (<ref>) whenever edge (i,j) is chosen at round t, and similarly for Θ_j^i(t+1). 
Then,  for any pair of agents (i,j), Θ_i^j(t)→Θ_j^* as t →∞, with probability one.




Let F_j, and w_j be defined as in Lemma <ref>. Suppose edge (i,j) is chosen at time t.  
It can be verified from (<ref>) and from the first order condition for optimality that
Θ_i^j(t + 1) = (∇ F_j)^-1(∑_k∈𝒩_jW_jkΘ_j^k(t)), where F_j is as defined in Lemma <ref>;
similarly, we have  that Θ_j^* = (∇ F_j)^-1(∑_k∈𝒩_jW_jkΘ_k^*) for each component of the solution.

Lemma <ref> allows us to find that

    Θ_i^j(t+1) - Θ_j^*
       ≤    (m_j + w_j)^-1∑ W_jk (Θ_j^k(t) - Θ_k^*)
       ≤    (m_j + w_j)^-1∑ W_jkΘ_j^k(t) - Θ_k^* 
       ≤    (m_j + w_j)^-1 w_j max_k Θ_j^k(t) - Θ_k^* 
       ≤    (m_j + w_j)^-1 w_j V(t),

where V(t) := max_l,kΘ_l^k(t) - Θ_k^* is the maximum error at round t between the agents' estimates and the solution.






If edge (i,j) is chosen at round t, we have, by the derivation above, that 
    Θ_i^j(t+1) - Θ_j^*≤  V(t).
 If that edge is not chosen, then Θ_i^j(t+1) = Θ_i^j(t) and, by definition of V(t), Θ_i^j(t+1) - Θ_j^*≤ V(t). We conclude that (<ref>) holds for any pair (i,j) and, so, V(t+1) ≤ V(t). Since V(t) ≥ 0, the limit (which is a random variable) V lim_t→∞ V(t)
is thus always well defined. The goal of the proof is to show that V = 0 with probability one.



Recall that T_m+1 denotes the first round after T_m by which all edges were selected at least once. It should be clear to the reader that the remainder of the proof holds almost surely, since T_m is finite for all m with probability one. Suppose edge (i,j) was selected at round T_m + s. Then

    Θ_i^j(T_m+s) - Θ_j^*≤ (m_j + w_j)^-1 w_j V(T_m),

cf. inequality (<ref>). Since, by definition of T_m and T_m+1, all edges (i,j) in the graph were selected at least once between T_m and T_m+1, inequality (<ref>) holds for all (i,j)∈ℰ when T_m + s = T_m+1. In other words,

    Θ_i^j(T_m+1) - Θ_j^*≤  (m_j + w_j)^-1 w_j  V(T_m)

for all edges (i,j)∈ℰ. It follows, by the definition of V(T_m+1), that

    V(T_m+1)≤β V(T_m) ,
 
where β  max_i { (m_i + w_i)^-1 w_i}∈ [0,1).

Let us take the limit m→∞ in (<ref>). We know that T_m→∞, and it follows from (<ref>) that V≤β V. Thus, owing to  0 ≤β < 1, we must have V=0, which implies Θ_i^j(t)→Θ_j^* as t →∞, for any (i,j).


Inequality (<ref>) implies that V(t) ≤ V(0). The iterations ( Θ_k^l(t) )_t ≥ 0 are, thus, uniformly bounded. We conclude (by the dominated convergence theorem) that our result implies the convergence in expectation result in <cit.>.










§ FIELD ESTIMATION EXAMPLE
 

Setup. Following <cit.>, we consider a field estimation setup that leads to a problem of the form (<ref>). The n agents are spread in a region and wish to profile a certain quantity, say, temperature, over the region: agent i cares only about the value of the quantity at its location, θ_i. Assume that the true values of the temperatures, θ = (θ_1,…, θ_n), are drawn from a prior distribution: a normal distribution with known mean and covariance Σ; as in <cit.>, we assume that the off-diagonal elements of Σ^-1 match the sparsity of the network, that is, (Σ^-1)_ij>0 if and only if (i,j)∈ℰ. 
Agent i measures y_i = θ_i + ν_i, where ν_i models identically distributed sensor noise (for simplicity), which is independent across agents. 




MAP estimation. A maximum 
a posteriori (MAP) approach seeks the θ = ( θ_1, …, θ_n ) that maximizes ∑_i=1^n log(y_i |θ_i) + log(θ_1,…,θ_n); or, equivalently, the θ that minimizes

    1/2(
    ∑_i∼ jσ_ij(θ_i - θ_j)^2 +
    ∑_i=1^n σ_iiθ_i^2
    )
    +
    ∑_i=1^n ϕ(y_i-θ_i),

where σ_ij := (Σ^-1)_ij and ϕ depends on the distribution of the noise ν_i. We let ϕ be a Huber penalty function to handle outliers <cit.>.
Finally, defining the personal loss functions as f_i(θ_i) ϕ(y_i-θ_i) + 1/2σ_iiθ_i^2 puts (<ref>) in the form (<ref>). Also, assumptions (<ref>) and (<ref>) hold.




Results: comparing DJAM with CL-ADMM.
Since the algorithm MPA from <cit.> applies only to quadratic functions, we use the ADMM-based algorithm CL-ADMM from <cit.> to compare with DJAM. Note that both CL-ADMM and DJAM converge to the solution with probability one. The algorithm CL-ADMM, however, being based on ADMM, has a parameter to tune—the parameter in the quadratic penalization part of the augmented Lagrangian function. This parameter, which we refer to as ρ, is known to affect noticeably the convergence speed of ADMM.



The results for a field estimation instance are shown in Figure <ref>. It shows, across rounds t,  the relative error between an agent's private model Θ_i^i(t) and the solution component Θ_i^*: Θ_i^i(t) - Θ_i^*  / Θ_i^*. The relative error was averaged over agents and over 100 Monte Carlo trials where, in each Monte Carlo run, we choose a different set of edges along time.

Figure <ref> confirms that the speed of convergence of CL-ADMM varies with the parameter ρ noticeably. 
In fact, we verified in other simulations (omitted due to lack of space) that the optimal ρ varied significantly with the number of agents, with the range of values for σ_ij, and with the noise distribution—we found the optimal ρ for those simulations by careful hand-tuning. In contrast, DJAM is on par with the best ρ in Figure <ref>, and needs no parameter tuning. 










ieeetr

