Federated Automatic Latent Variable Selection in Multi-output Gaussian Processes

The MGP was originally known as cokriging in the geostatistics field [13], where it gained significant attention as an effective tool for modeling the relatedness between regionalized variables. Beyond geostatistics, the ability to handle multivariate data has allowed MGPs to achieve notable success in broader fields, including healthcare (e.g., [14]), transportation (e.g., [15]), system prognostics (e.g., [16]), and ecology (e.g., [17]). This broad applicability has spurred research into designing MGPs to address specific challenges in different scenarios. For instance, Moreno-Munoz et al. [18] developed a heterogeneous MGP, designed for modeling outputs with heterogeneous data types, such as categorical and time series outputs. Wang et al. [19] developed an MGP that can handle input domain inconsistency and negative knowledge transfer [20] occurring due to a lack of shared information among some outputs. Chung et al. [21] proposed a weakly-supervised MGP to address cases where output membership is unknown for some observations. Additionally, Soleimani et al. [22] introduced an approach that combines an MGP and a survival analysis model for the joint modeling of longitudinal and time-to-event data.

While there are many success stories, a prominent challenge in MGPs is their instability with datasets containing numerous observations. This issue stems from standard GPs that suffer $\mathcal{O}(N^{3})$ time complexity when modeling $N$ observations due to the inversion of an $N\times N$ covariance matrix. This complexity is extended to $\mathcal{O}(M^{3}N^{3})$ for MGPs when modeling $M$ outputs, each with $N$ observations. To tackle this issue, several approaches have been proposed. For example, inspired by sparse approximation of GPs [23], Alvarez and Lawrence [7] introduced a sparse approximation for the shared latent processes using pseudo-inputs, which reduces the complexity of inverting the approximated covariance matrix. This approach was later extended by Alvarez et al. [24], who proposed a variational inference (VI)-based framework to approximate the posterior of the sparse latent processes, addressing the limitations of assuming smooth shared latent functions. Nguyen et al. [25] developed a scalable MGP model built upon the LMC with sparse approximation, amenable to stochastic VI [26] that facilitates the use of mini-batches in optimization. Recently, Bruinsma et al. [27] showed that the orthogonal instantaneous linear mixing models can achieve scalable inference and learning even without sparse approximation, offering an alternative approach to address scalability issues.

As demonstrated by the studies mentioned above and others (e.g., [28, 29, 30]), modeling each output as a linear combination of shared latent functions is highly popular in constructing an MGP. However, the literature on determining the appropriate number of latent functions for specific data is quite scarce. Notably, Titsias and Lazaro-Gredilla [31] introduced a method for selecting latent functions using spike-and-slab priors. However, this method is only applicable within a centralized regime and scales poorly with the number of observations. While our model also automatically selects latent functions, its unique advantages lie in its compatibility with federated systems and scalability through stochastic optimization.

With recent revolutionary advances in technologies that miniaturize chips with computing power, units at the edge, which used to merely create or collect data, increasingly possess significant computing and data storage capabilities. FL, since its first introduction by McMahan et al. [32], has garnered explosive interest in recent years as a collaborative framework to train models using the local computing resources of multiple units while keeping their data stored locally. This approach effectively addresses various challenges in FL such as data/system heterogeneity [33], client drifts [34, 35], fairness [36], and more. Despite these advancements, the majority of FL literature has focused predominantly on deep learning models [37].

Recently, several studies have aimed to extend FL methods beyond deep learning. A line of these efforts focused on GPs in federated settings. Achituve et al. [38] studied personalized FL by collaboratively estimating a GP with deep kernels across units and employing personalized GP classifiers for each unit. Yu et al. [39] investigated deep kernel learning in federated settings but avoided kernel sharing across units by working directly in the feature space. Yue and Kontar [40] explored a method that directly estimates GP hyperparameters via FL and studied its theoretical properties. Furthermore, Zhang et al. [41] examined GP regression under federated settings in the presence of units subject to Byzantine attacks, while Guo et al. [42] investigated a framework that learns a sparse GP model using FL algorithms. Note that these studies primarily focused on learning a single-output GP under federated settings.

In contexts where units in a federated regime have separate but correlated GPs, our study relates to recent work by Chung and Kontar [43]. Both studies establish an FL framework for MGPs where each output corresponds to a different unit. However, our model offers distinctive advantages: it can infer the appropriate number of latent variables needed to extract common latent functions across units, and thus, enhance prediction accuracy for each unit as well as enable efficient prediction for a new unit by using latent functions selectively instead of using all of them.

In contrast to independent modeling for individual units, the LMC is advantageous in analyzing connected systems that involve multiple units by accounting for dependencies across units. However, several challenges are encountered when it is directly deployed to connected systems in practice.

Determining $L$. Identifying non-trivial commonalities across units demands a sufficient number of latent functions to help construct a rich correlation structure capable of estimating complex dependencies between outputs. Yet, having more latent functions can significantly increase the risk of overfitting and computational burden. Thus, determining an adequate $L$ becomes crucial. A naive brute-force search for $L$ is substantially inefficient, given the computational complexity $\mathcal{O}(N^{3}M^{3})$ of MGPs.

Centralized model building and learning. The LMC model construction and inference implicitly assume a centralized process for managing all data and model parameters. That is, model building and training are done at a central location (e.g., a central server on a cloud) that amasses data from all units. To see this, building the covariance matrix $\mathbf{C}^{\sf LMC}_{\mathbf{f},\mathbf{f}}$ needs to calculate $\sum_{l=1}^{L}b_{m,m^{\prime}}^{l}k_{l}(\mathbf{x}_{m,n},\mathbf{x}_{m^{\prime},n^{\prime}};{\boldsymbol{\xi}}_{l})$ in (3), which requires sharing observations $\mathbf{x}_{m,n},\mathbf{x}_{m^{\prime},n^{\prime}}$ as well as the model parameter $b_{m,m^{\prime}}^{l}$. Furthermore, (3) shows that the derivation of the predictive distribution needs a posterior distribution $p(\mathbf{f}|\mathbf{X},\mathbf{y})=p(\mathbf{f}|\mathbfcal{D})$ conditioning on data across all units $\mathbfcal{D}$, again implying the need for a centralized process. Unfortunately, such centralized frameworks have critical drawbacks attributed to (i) the communication inefficiency induced by transmitting units’ raw data to the central server, (ii) the compromised privacy due to data sharing, and (iii) the need for massive computational power at the central cloud, often overwhelmed in a large IoT-enabled connected system that involves a lot of units.

One important characteristic of the LMC is that it features a hierarchical structure, where $f_{m}$ can be viewed as generating from the shared latent functions $\{u_{m}\}_{l=1}^{L}$ along with unit-specific parameters $\{w_{m,l}\}_{l=1}^{L}$. This implies that $\{f_{m}\}_{m=1}^{M}$ are conditionally independent of each other when the entire length of $\{u_{l}\}_{l=1}^{L}$ is given. Such hierarchical structure aligns with the hub-and-spoke structure of IoT-enabled connected systems, where multiple units at the edge are connected to the common central server. Given this structural accordance, we build an LMC-based hierarchical Bayes model that mirrors the same hub-and-spoke structure with connected systems, while incorporating a prior distribution that imposes sparsity on the coefficient parameters $\{w_{m,l}\}$. The prior distribution encourages some coefficients to be shrunk to zero so that only a set of latent functions needed to capture between-unit dependencies remains active. Interestingly, our hierarchical modeling allows for getting around the direct calculation of (3). This, in turn, facilitates to build an FA-based inferential algorithm (discussed in Section 4.2) where units collaboratively train the hierarchical model while sharing only needed information with the central server and securing unit-specific information that may compromise privacy if disclosed.

Our hierarchical modeling starts with placing a prior, which imposes sparsity on the coefficients $\mathbf{w}=[\mathbf{w}_{m}^{\top}]^{\top}_{m=1,\cdots,M}$ with $\mathbf{w}_{m}=[w_{m,l}]_{l=1,\cdots,L}^{\top}$ in (3). Specifically, we place a spike-and-slab prior, written as

where identical priors are independently placed on $w_{m,l}$. The spike-and-slab prior is characterized by a mixture of the Dirac delta $\delta_{0}(w_{m,l})$ centered at zero and a Gaussian distribution with a mean of zero and a variance of $\sigma^{2}_{\mathbf{w}}$, with a mixing probability $\pi$. Intuitively, the prior indicates that $w_{m,l}$ follows a Gaussian distribution centered at zero with probability $\pi$, or is simply zero with probability $1-\pi$. In the presence of $\delta_{0}(w_{m,l})$, the prior puts a substantial mass on zero. This results in encouraging some $w_{m,l}$ to be zero in model inference, thereby eliminating the influence of $u_{l}$ in characterizing $f_{m}$. Consequently, a set of only necessary $u_{l}$ in describing relatedness across output functions $\{f_{m}\}_{m=1}^{M}$ can be identified.

Now, let us tackle the issue of the direct calculation of covariances in (3). We want to estimate the covariance without the need to share inputs across units. Our idea is based on the sparse MGP method [7]. Specifically, it introduces a set of inducing or auxiliary variables, denoted by $\mathbf{Z}_{l}=[\mathbf{z}_{l,q}]^{\top}_{q=1,...,Q}\in\mathbb{R}^{Q\times d}$, at which the latent functions are evaluated, i.e., $\tilde{\mathbf{u}}_{l}=[u(\mathbf{z}_{l,q})]^{\top}_{q=1,...,Q}$. Here $\{f_{m}(\cdot)\}_{m=1}^{M}$ are assumed to be conditionally independent given $\tilde{\mathbf{u}}_{l}$, in place of the entire information of $\{u_{l}(\cdot)\}_{l=1}^{L}$. The rationale is that $\tilde{\mathbf{u}}_{l}$ may characterize $u_{l}(\cdot)$ sufficiently well. This method indeed originates from the effort to reduce the computational complexity of MGPs by choosing $Q\ll N$ [23]. Interestingly, it in turn eliminates the direct calculation of covariance in the LMC. To see this, we present

with

where (7) represents the latent independent GPs for $\tilde{\mathbf{u}}_{l}$ evaluated at $\mathbf{Z}_{l}$ and (8) implies the conditional independence. It is important to note that covariance matrices directly calculating covariances between units (e.g., $\mathbf{C}_{\mathbf{f}_{m},\mathbf{f}_{m^{\prime}}}$ or $\mathbf{C}_{\mathbf{u}_{m,l},\mathbf{u}_{m^{\prime},l}}$ for $m\neq m^{\prime}$) no longer appear in (8) and (7). As such, we can calculate (8) and (7) without sharing data across units.

Finally, the probability distribution of outputs can be represented as

given regression modeling with the Gaussian noise in (1).

The complete form of our hierarchical model is given by (6), (7), (8) and (9). Figure 1 illustrates the hierarchical structure of our model. One should note that the structure aligns with the natural hierarchy of IoT-enabled connected systems. Based on this structure, our proposed model can estimate the relatedness across units while situating their data $\mathbfcal{D}_{m}$ and unit-specific parameters $\mathbf{w}_{m}$ in the storage of respective units at the edge. Instead, only needed information ($\tilde{\mathbf{u}}_{l}$ and $\mathbf{Z}_{l}$) is shared through the central server.

Unfortunately, despite the hierarchical model construction, the inference and prediction are still troublesome because (i) predictions using the proposed model involve a posterior distribution given all data from units, which requires the collection of local datasets to a central location and (ii) the posterior distribution is intractable. These challenges can be clearly seen in the posterior distribution of our model, represented as

Here we can see that the posterior $p(\mathbf{f},\tilde{\mathbf{u}},\mathbf{w}|\mathbf{y},\mathbf{X},\mathbf{Z})$ assumes all data $\mathbfcal{D}\equiv(\mathbf{X},\mathbf{y})$ is given, and further, its closed form is intractable due to the likelihood

of which integration is intractable due to the presence of the spike-and-slab prior $p(\mathbf{w})$.

In the next section, we discuss a posterior inference framework for the proposed hierarchical model, which tackles the challenges mentioned above.

This section discusses our proposed inference framework. The core idea of our approach is to approximate the original posterior using an estimated surrogate distribution without centralized computation. Consequently, we can use the surrogate distribution in place of the original posterior which causes an issue due to the need to collect all data. We first discuss the surrogate distribution and an optimization problem to be solved for its inference (Section 4.2.1). Then we delve into an FL-based algorithm for the optimization problem, which facilitates the use of local computing power and shares only needed information while preserving privacy (Section 4.2.2).

Based on our framework, we develop an efficient approach for incorporating new units into the system after model estimation is completed. A straightforward method to address this is by retraining an MGP with an updated dataset that includes both existing and new units. However, this approach requires all units to participate in the learning process each time a new unit is added, resulting in significant inefficiency. In contrast, our approach utilizes only the computing power of the new unit for its integration. The key assumption is that if the new unit’s data shares the common latent pattern characterized by the existing units, it is sufficient to use the shared parameters estimated from the existing units and only infer the parameters specific to the new unit. Furthermore, this approach benefits from our automatic latent variable selection, enhancing efficiency by using only the selected latent functions instead of all latent functions.

Suppose we have estimated parameters $\hat{\boldsymbol{\Theta}}$ and a set of selected latent functions $l\in\mathcal{S}\subset\{1,...,L\}$ inferred from existing units. Now, consider a new unit $p\notin\{1,...,M\}$. As useful latent functions are already selected, we can build an LMC that also includes the new unit without the spike-and-slab prior. The ELBO for this model can be written as $\mathcal{L}^{p}_{\sf ELBO}({\boldsymbol{\psi}}_{p};\hat{\boldsymbol{\Theta}})=\sum_{m\in\{1,...,M,p\}}\mathbb{E}_{q(\mathbf{f}_{m})}[\log p(\mathbf{y}_{m}|\mathbf{f}_{m})]-\sum_{l\in\mathcal{S}}{\sf KL}(q(\tilde{\mathbf{u}}_{l})\|p(\tilde{\mathbf{u}}_{l}))$ where now we have $q(\mathbf{f}_{m})=\int p(\mathbf{f}_{m}|\tilde{\mathbf{u}})q(\tilde{\mathbf{u}})d\tilde{\mathbf{u}}$, with $p(\mathbf{f}_{m}|\tilde{\mathbf{u}})$ and $q(\tilde{\mathbf{u}})$ defined similarly as (8) and (12), respectively, but $\mathbf{w}$ is treated as a parameter rather than a random variable and $\tilde{\mathbf{u}}_{l}$ is used for $l\in S$ instead of all $l\in\{1,...,L\}$.

If the commonality across the existing and new units is similar to the commonality among the existing units, we can continue using $\hat{\boldsymbol{\Theta}}=\{\hat{\boldsymbol{\theta}},\{\hat{\boldsymbol{\psi}}_{m}\}_{m=1}^{M}\}$ when learning the new unit $p$. Thus, it would be sufficient to maximize $\mathcal{L}^{p}_{\sf ELBO}({\boldsymbol{\psi}}_{p};\hat{\boldsymbol{\Theta}})$ with respect to ${\boldsymbol{\psi}}_{p}$ only, while keeping $\hat{\boldsymbol{\Theta}}$ fixed. Removing the terms in $\mathcal{L}^{p}_{\sf ELBO}$ that are independent of ${\boldsymbol{\psi}}_{p}$, the optimization problem simplifies to

with $\hat{\boldsymbol{\theta}}$ being fixed. Notably, the objective function in (17) does not involve any parameters or data from other units $m\in\{1,...,M\}$. It is solely related to the new unit $p$ and the estimated global parameters $\hat{\boldsymbol{\theta}}$ that remain fixed throughout optimization. Therefore, solving (17) can be locally performed using the computing power of unit $p$, without the involvement of other units.

Once model parameters and variational distributions are estimated, each unit can make predictions for a new input. Let $\mathbf{y}_{m,*}$ and $\mathbf{f}_{m,*}$ represent observations and function values at $N_{*}$ new inputs $\mathbf{X}_{m,*}\in\mathbb{R}^{N_{*}\times d}$ for unit $m$, respectively. The predictive distribution can be written as

where $\mathbf{A}^{*}_{m,l}$ and $\mathbf{K}^{*}_{m,l}$ are calculated similarly to $\mathbf{A}_{m,l}$ and $\mathbf{K}_{m,l}$ but with $\mathbf{X}_{m,*}$; and we use the hat notation for the estimated parameters. To address the intractable integral in (18), we resort to the Monte Carlo method [45]. This leads to the following approximation

where $w_{m,l}^{(s)},\alpha_{l}^{(s)}\overset{\text{i.i.d.}}{\sim}q(\tilde{w}_{m,l},\alpha_{l})$ in (13) with the estimated parameters $\hat{\mu}_{w_{m,l}}$,$\hat{\sigma}^{2}_{w_{m,l}}$ and $\hat{\gamma}_{l}$, for $s=1,...,S$.