ADMM-based Networked Stochastic Variational Inference
Hamza Anwar Quanyan Zhu



*
Anwar et al.: ADMM-based Networked Stochastic Variational Inference





ADMM-based Networked
Stochastic Variational Inference
    Hamza Anwar, Member, IEEE,
        and Quanyan Zhu, Member, IEEE
This work is partially supported by the grants EFMA-1441140 and SES-1541164 from National Science Foundation. H. Anwar and Q. Zhu are with the Department of Electrical and Computer Engineering,
New York University, New York, NY 10003 USA (e-mail: ha1082@nyu.edu, qz494@nyu.edu).
Manuscript received , 2017; revised December 30, 2023.
    December 30, 2023
==================================================================================================================================================================================================================================================================================================================================================================================================================


Owing to the recent advances in “Big Data” modeling and prediction tasks, variational Bayesian estimation has gained popularity due to their ability to provide exact solutions to approximate posteriors. One key technique for approximate inference is stochastic variational inference (SVI) <cit.>. SVI poses variational inference as a stochastic optimization problem and solves it iteratively using noisy gradient estimates. It aims to handle massive data for predictive and classification tasks by applying complex Bayesian models that have observed as well as latent variables. This paper aims to decentralize it allowing parallel computation, secure learning and robustness benefits. We use Alternating Direction Method of Multipliers in a top-down setting to develop a distributed SVI algorithm such that independent learners running inference algorithms only require sharing the estimated model parameters instead of their private datasets. Our work extends the distributed SVI-ADMM algorithm that we first propose, to an ADMM-based networked SVI algorithm in which not only are the learners working distributively but they share information according to rules of a graph by which they form a network. This kind of work lies under the umbrella of `deep learning over networks' and we verify our algorithm for a topic-modeling problem for corpus of Wikipedia articles. We illustrate the results on latent Dirichlet allocation (LDA) topic model in large document classification, compare performance with the centralized algorithm, and use numerical experiments to corroborate the analytical results.


variational inference, transfer learning, stochastic optimization, method of multipliers, inference over networks



§ INTRODUCTION

The explosive influx of data and information for modern day technological systems has opened doors to revolutionary possibilities. One of the most vital uses of this data is in modeling, visualizing, and analyzing large data sets through probabilistic tools. Statistical machine learning is at the core of numerous such applications in what is becoming known as Internet of Things (IoT). Such iterative learning mechanisms help in better control performance for many cyber-physical-systems in which estimation of parameters and system-identification is required. Probabilistic graphical modeling is one key research area that has helped in data analysis in inference and prediction tasks, <cit.>. These models visually express assumptions about data and its hidden structure. Posterior inference algorithms have been proven to exploit such models in explaining this hidden structure while being adaptive, robust, parallelizable, and scalable.

Variational inference, from late 90s, is a method that transforms complex inference problems into high dimensional optimization problems. In contrast to Monte-Carlo sampling methods (that simply aim to find exact answer to an approximate problem), the variational Bayesian approach solves for optimal solution under constraints to the right inference problem, <cit.>. On the same lines, stochastic variational inference (SVI) was developed recently that extends variational inference to be solved using stochastic optimization under certain assumptions, <cit.>. SVI works iteratively in gradient ascent fashion using noisy gradient estimates. It provides approximate model posteriors with only a few passes through a large data collection, making it highly scalable. We propose ADMM-based Networked SVI – a distributed stochastic variational inference technique that builds upon standard SVI, retaining most of its benefits, based on the highly parallelizable alternating direction method of multipliers (ADMM) where the agents are connected as in a graph with nodes and edges.



 §.§ Related Work
Numerous extensions have been proposed for the SVI framework in its application to more model classes (by <cit.> and <cit.>), different underlying processes (by <cit.> and <cit.>), and structural exploitations (by <cit.>) making it faster and widely deployable. A variety of works focus on making variational methods distributed to enhance parallelizability. The work on distributed Bayesian nonparametric models by <cit.> is commendable in making variational inference updates distributed, asynchronous, and `streaming' (online) and they've shown it to outperform standard SVI. However, their work is only specific to the Dirichlet process mixture and lacks in generalizability to the class of probabilistic models that SVI can deal with. Similar to <cit.>, another work D-MFVI by <cit.> uses ADMM for decentralizing, like us, but lacks in being extendable to online updates, fast convergence rate, and other desirable properties of SVI. Distributed VBA (<cit.>) also uses ADMM however their approach lacks in scalability to large data without demanding adequate computational resources.

None of these works use stochastic optimization methods to speed up inference and hence they are fundamentally different from standard SVI itself. One recent work `Extended-SVI' by <cit.> retains the benefits of SVI while making it distributed and asynchronous. They employ a rather simple algorithmic change to SVI however their work remains unexplored in terms of depth because they particularly focus on Gaussian mixture models and do not provide how it is extendable to all other probabilistic models for which SVI in general works.

In contrast to all related works highlighted, our approach extends the general SVI framework to a networked stochastic optimization consensus problem. We tackle the issue of generalizability to all graphical models by providing a general solution and make use of the stochastic gradient updates that make it fast. We run ADMM updates along with stochastic gradient ascent for variational objective to reach consensus among a number of distributed learners. SVI itself being a non-convex stochastic optimization problem makes the distributed problem trickier. Our work aims to show that independent learners that use SVI for similar applications, can collaborate by exchanging their results (not the data itself) to benefit from each other improving overall accuracy of results. This approach makes our work unique and applicable to wider distributed large-scale inference problems. Not just that, but this kind of an approach poses a game problem with multiple agents interested in performing their own inference tasks, and simultaneously benefiting from each other through reinforcement learning and cooperation.



§ ADMM-BASED DISTRIBUTED SVI

Building upon the recent work on SVI by Hoffman et al., <cit.>, we consider the SVI problem for a network of learners.

The N observations are x = x_1:N; the vector of global hidden variables is β; the N local hidden variables are z = z_1:N, each of which is a collection of J variables z_n = z_n,1:J; the vector of fixed parameters is α. (Note we can easily allow α to partly govern any of the random variables, such as fixed parts of the conditional distribution of observations. To keep notation simple, we assume that they only govern the global hidden variables.)


 §.§ Optimization problem


    min_λ_k      ∑_k=1^K g_k(λ_k)
    subject to      λ_k - ζ = 0,   k= 1,⋯, K
    
       λ_k ∈Γ_k

where each λ_k is an m-sized vector and Γ_k indicates the feasible set for the variables λ_k (typically Γ_k = ℝ^m_+) and,

    g_k(λ_k) := -𝔼_ϕ(λ_k)[η_g(x,z)] ^⊤∇_λ_k a_g(λ_k) + λ_k^⊤∇_λ_k a_g(λ_k)
     - a_g(λ_k) + const.

which is the standard SVI problem objective function for a single learner. The above optimization problem gives us a solution for K learners when they form a consensus. Using an Augmented Lagrangian approach, as in ADMM, we solve this problem in a distributed iterative fashion for multiple learners.


 §.§ ADMM-based solution

Augmented Lagrangian with a quadratic penalty is used to arrive at the ADMM update iterations. The Lagrange multipliers are denoted by y_k∈Γ_k. Minimization updates for each processor/agent are given as:

    λ_k^t+1 =    min_λ_k(g_k(λ_k) + (λ_k - ζ^t)^⊤ y_k^t + c/2||λ_k-ζ^t||_2^2 ),
    
    ζ^t+1 :=    1/K∑_k=1^K ( λ_k^t+1 + (1/c)y_k^t ),
    
    y_k^t+1 :=     y_k^t + c (λ_k^t+1-ζ^t+1).

where ζ is called the central collector, and c is the quadratic penalty parameter in the augmented Lagrangian which is given as:

    L_c({λ_k},{y_k},ζ) = ∑_k=1^K g_k(λ_k) + (λ_k - ζ) ^⊤ y_k+ c/2||λ_k-ζ||_2^2.


Here, we note that the λ-minimization update which is actually a solution to λ_k^t+1 = min_λ_k∈Γ_k L_c({λ_k},{y_k^t},ζ^t), requires solving a constrained non-convex optimization problem. We solve this in a gradient descent fashion in of itself, as the standard SVI problem was also solved, but the original solution requires inversion of a Hessian matrix. For that we take into account one-step earlier value of λ_k^t — for details about derivation and Hessian inversion approximation used, see Appendix. Thereby, our proposed iterative ADMM methodology runs along with a gradient-descent iterative update of variables which is completely summarized in algorithm <ref>.





















 §.§ Experimental results

New set of experiments for the distributed problem was performed with multiple learners. Here, we show the results with 4 learners. Figure <ref>, shows the convergence properties of our distributed learners, for a metric of the estimated model's fitness, known as the `held-out perplexity'. This same metric has been used by Hoffman et al. <cit.> to show convergence of the algorithm. A comparison of centralized versus distributed two-player SVI algorithms is depicted in Figure <ref>. We conclude that all the learners not only converge to higher precision in estimates (evident from the Figure <ref>), but also achieve accuracy of estimates (evident from Table <ref>), while simultaneously maintaining consensus.



From Table <ref>, we see that the highly probable words for a given topic learned by any of the four learners all direct to a similar kind of subject. For example, in Topic#98, the learners understood it to represent the names of months | even though for this topic, the distribution of word occurrences is different for the four players but it is evident that they all point to the same abstract class of words. Similarly, in other topics as well we see similarity in estimates. Sometimes, we even see that the descending order of the words is exactly the same, e.g., in Topic#72 and Topic#38 all learners have the same ordering of words. Thus, the table shows that despite the fact that all learners use their own independently fetched datasets from Wikipedia articles, the consensus between results is achieved among all the learners, due to the central collection constraint.




§ ADMM-BASED NETWORKED SVI

Now, after conclusive results about distributed `fully-conected' SVI algorithm, we move on to a network of nodes having independent learners residing at each node. The only difference in problem formulation, as we will see is in the equality constraints. We use the network formulation as given in <cit.>.

The network is modeled by an undirected graph denoted by 𝒢(𝒦, ℰ) with 𝒦:={1,…,K} representing the set of nodes, and ℰ representing the set of links between nodes. Node k ∈𝒦 only communicates with his neighboring nodes B_k ⊆𝒦. Note that without loss of generality, graph 𝒢 is assumed to be connected. The network can contain cycles. An example of such a network is is shown in Figure <ref>.

At every node k ∈𝒦, a set of observations 𝒟_k :={x_kn:n = 1, …, N_k} of size N_k is available, where x_kn denotes the n-th observation for the k-th node. Though not explicitly expressed, each x_kn can be a collection of multiple random variables. The vector of global hidden variables for node k is β_k; its N_k local hidden variables are z_k = z_k,1:N, each of which is a collection of J variables z_kn = z_kn,1:J; the vector of fixed parameters is α_k.
 



With the graph formulation given above, we pose the distributed SVI problem for a network of learners given as:

    min_{λ_k}      ∑_k=1^K g_k(λ_k)
    s.t.      λ_k = λ_l, ∀ k∈𝒦, l∈ B_k

with λ_1,⋯,λ_K as variables. Here, g_k:Γ_k →ℝ_+ is a non-linear function of λ_k re-written here:

    g_k(λ_k) := -𝔼_ϕ(λ_k)[η_g(x,z)] ^⊤∇_λ_k a_g(λ_k) + λ_k^⊤∇_λ_k a_g(λ_k)
     - a_g(λ_k) + const.

Optimization problem (<ref>) is equivalent to the following,

    min_{λ_k}      ∑_k=1^K g_k(λ_k)
    s.t.      λ_k = ζ_kl,   ζ_kl = λ_l,   ∀ k∈𝒦, l∈ B_k

where ζ_kl are redundant variables that will facilitate the decoupling of variable λ_k at node k from its neighboring nodes l∈ B_k. This problem will be solved using its dual. We denote the Lagrange multipliers by y_kl1 (y_kl2) for the constraints λ_k = ζ_kl (ζ_kl = λ_l). We observe that for each k we have 2|B_k| equality constraints. The augmented Lagrangian with a quadratic penalty is:

    L_c({λ_k},{ζ_kl},{y_klj}) = ∑_k=1^K ( g_k(λ_k) + ∑_l∈ B_k( y_kl1^⊤ (λ_k - ζ_kl)
     + y_kl2^⊤ (ζ_kl - λ_l) + c/2(||λ_k-ζ_kl||^2 + ||ζ_kl- λ_l||^2 ))),

The augmented Lagrangian can be iteratively minimized with respect to each variable by keeping others constant, which gives us a set of minimization updates for each variable summarized in the following Proposition.


 §.§ Proposition 1

The distributed iterations solving (<ref>) are as follows:

    λ_k^t+1   =  min_λ_k{ g_k(λ_k) + λ_k^⊤∑_l∈ B_k(  y_kl1^t  - y_lk2^t )
     + c/2∑_l∈ B_k( ||λ_k-ζ_kl^t||^2 + ||ζ_lk^t-λ_k||^2) .
    ζ_kl^t+1   = min_ζ_kl{
    - ζ_kl^⊤y_kl1^t + ζ_kl^⊤y_kl2^t + c/2(||λ_k^t+1-ζ_kl||^2)
    
     + c/2(||ζ_kl-λ_l^t+1||^2)
    .
    
    y_kl1^t+1   = y_kl1^t + c (λ_k^t+1-ζ_kl^t+1),  ∀ k∈𝒦, l∈ B_k
    
    y_kl2^t+1   = y_kl2^t + c (ζ_kl^t+1 -λ_l^t+1),  ∀ k∈𝒦, l∈ B_k

and correspond to the standard ADMM solver discussed in <cit.>.




  §.§.§ Proof
 The first task is to cast the problem (<ref>) into standard ADMM problem form in <cit.>. The network description adopted here is similar to the one used in <cit.> and thus we use it to establish equivalence with standard ADMM <cit.>. Thereby, the remaining form of the minimization updates is directly derived from the augmented Lagrangian given in (<ref>). The λ-minimization update (<ref>) is derived by eliminating the terms that do not affect the minimization in augmented Lagrangian:

    λ_k^t+1   = min_λ_k L_c(λ_k,{ζ_kl^t},{y_kl1^t},{y_kl2^t}),
       =  min_λ_k{
    g_k(λ_k) + ∑_l∈ B_k(  λ_k^⊤ y_kl1^t+ c/2||λ_k-ζ_kl^t||^2)
    
    +∑_s∈ B_k( - λ_k ^⊤ y^t_sk2+ c/2||ζ_sk^t-λ_k||^2)
    .,

which upon merging the two summations reduces to (<ref>). Similarly, the ζ-minimization update (<ref>) comes directly from,

    ζ_kl^t+1 = min_ζ_kl L_c({λ_k^t+1},ζ_kl,{y_kl1^t},{y_kl2^t}).

Equations (<ref>)–(<ref>) are the dual variable updates (cf. <cit.>).
▪

Next we reduce the iteration equations to a simpler form. Here, we observe that the ζ_kl update has the following unique solution (by putting the derivative equal to zero and solving),

    ζ_kl^t+1 = 1/2c(y_kl1^t - y_kl2^t) + 1/2(λ_k^t+1+λ_l^t+1)


Putting (<ref>) in (<ref>)–(<ref>) gives,

    y_kl1^t+1 =    1/2(y_kl1^t+y_kl2^t) + c/2 (λ_k^t+1-λ_l^t+1),
    y_kl2^t+1 =    1/2(y_kl1^t+y_kl2^t) + c/2 (λ_k^t+1-λ_l^t+1).

Now, we assume that both the Lagrange multipliers are identically initialized at every node k, as zero y_kl1^0 = y_kl2^0 = 0_m × 1  ∀ k∈𝒦, l∈ B_k. This ensures that y_kl1^1 = y_kl2^1, and y_kl1^2 = y_kl2^2, and son on. We see that only one of the two multipliers per node needs to be updated at each time step. Furthermore, (<ref>) simplifies to,

    ζ_kl^t+1 = 1/2(λ_k^t+1+λ_l^t+1).

Finally the ADMM iterations ∀ k∈𝒦 simplify, summarized in the following Proposition.



 §.§ Proposition 2

Selecting y_k^0 := y_kl1^0 = y_kl2^0 = 0_m × 1 as initialization ∀ k∈𝒦, l∈ B_k, the iterations (<ref>)–(<ref>) reduce to the following,

    λ_k^t+1   =min_λ_kg_k(λ_k) +  λ_k ^⊤ y_k^t + c∑_l∈ B_k||λ_k-1/2(λ_k^t+λ_l^t)|| ^2
    
    
    y_k^t+1   = y_k^t + c∑_l∈ B_k (λ_k^t+1-λ_l^t+1),  ∀ k∈𝒦



  §.§.§ Proof
 Substituting (<ref>) into the objective (<ref>) gives the following,

    min_{λ_k} L_c({λ_k},{λ_k^t},{ζ_kl^t},{y_kl1^t},{y_kl2^t})=
    ∑_k=1^K ( g_k(λ_k) + λ_k^⊤∑_l∈ B_k(  y_kl1^t  - y_lk2^t ) + .
    . c/2∑_l∈ B_k( ||λ_k-1/2(λ_k^t+λ_l^t)||^2 + ||1/2(λ_l^t+λ_k^t)-λ_k||^2)).

Note that {λ_k} is the set of all variables of optimization and {λ_k^t} denote constants known from previous iteration. All-zero initialization of the Lagrange multipliers implies that y_kl1^t = -y_lk1^t ∀ t [cf. from (<ref>)–(<ref>)], and so the first double sum in (<ref>) can be rewritten as:

    ∑_k=1^K∑_l∈ B_kλ_k^⊤(  y_kl1^t  - y_lk2^t ) = 2∑_k=1^K λ_k^⊤∑_l∈ B_ky_kl1^t.

The other two double sums in (<ref>) can be simplified to give,

    c/2∑_l∈ B_k( ||λ_k-1/2(λ_k^t+λ_l^t)||^2 + ||1/2(λ_l^t+λ_k^t)-λ_k||^2)
    
    =c∑_l∈ B_k( ||λ_k-1/2(λ_k^t+λ_l^t)||^2).

By defining y_k^t := 2∑_l∈ B_k y_kl1^t, and substituting (<ref>) and (<ref>) into (<ref>), gives the final form of the augmented Lagrangian which completes the proof:

    min_{λ_k} L_c({λ_k},{λ_k^t},{y_k^t}) = ∑_k=1^K g_k(λ_k) +  ∑_k=1^Kλ_k ^⊤ y_k^t  
     + c∑_k=1^K∑_l∈ B_k||λ_k-1/2(λ_k^t+λ_l^t)|| ^2.

▪



 §.§ Network solution

Now, we present a solution to the ADMM minimization update (<ref>), which is a non-convex optimization problem, similar to the corresponding update in distributed SVI in section <ref>. 
We make use of stochastic gradient descent like standard SVI algorithm for minimization of augmented Lagrangian (cf. <cit.>). It is known that the natural gradient of g_k is given as,
    ∇̂_λ_k g_k(λ^t_k) = λ_k^t - λ̂_k.
The solution is presented in algorithm <ref>.



 §.§ Experimental results

The issue of cross-matching the topics between two different players was solved using a correlation metric. Like discussed earlier, SVI relies on random initialization of the global parameters. So, each player initializes with different global parameters, and as they encounter observations, they update the global parameters. Since, in our setting each player has its own independent dataset, so the trajectory of converging to `true' topics is different for every player. That is why, if two completely independent learners are fed the same data, they converge to similar estimates but with different trajectories (i.e. topic 0 for player A may truly represent the contents of topic 43 for player B, and so on). Thus, in order to match the right topics, a correlation metric was needed. We used the Pearson correlation coefficient for this.

Result for an experiment that used a line-type graph is shown in Fig. <ref>. The network is shown in Fig. <ref>. In this experiment nodes 0 and 1 were provided with exactly same set of data (limited to a fixed 800 documents offline available that were fed repeatedly). The nodes 2, 3, and 4 were provided with an online data i.e. independent and new data points at each iteration. The purpose of this experiment was to see how the connected nodes corroborate and improve estimation accuracy (of node 1) in contrast to accuracy of the independently running learners (i.e. node 0). The perplexity metric trajectory shown in Fig. <ref>, supports our claim that node interaction through ADMM updates certainly benefits. We achieve better accuracy in the estimate of node 1 as compared to that of node 0 because the learning at nodes 2-4 affects that of node 1 due to the consensus constraint.





§ DISCUSSION AND CONCLUSION

We have presented distributed ADMM-based SVI and an extension of it over a network of learners in a grpah – an algorithm that solves separable stochastic optimization problems and merges their results to achieve optimal consensus solution. Applications of distributed learning agent systems are common in IoT framework especially when different learning systems do not want to share data among each other but still agree on partial collaboration and transfer learning. A well-trod example of latent Dirichlet allocation for probabilistic topic models is implemented to show comparative results for the centralized, distributed (thoroughly fully-connected) and networked settings. One key observation in the results of our networked SVI algorithm is that strongly connected networks exhibit substantial transfer learning benefits. This is highlighted in the comparative experiments of Fig. <ref> and Fig. <ref>. Another observation is that accuracy of estimation improves over time, as more and more data is analyzed. In a nutshell, results show that through collaboration without having to share private data, two or more independent model posterior learners for SVI can improve their learning capabilities. Due to the use of stochastic optimization, this algorithm is considerably fast, scalable, and accurate. Moreover, its distributed learning methodology enhances security and robustness aspects that underpin modern deep learning goals. For future, we intend to apply this to cyber-physical dynamical systems along with inspecting guarantees on the convergence properties of this algorithm.
















IEEEtran

