# Federated Learning for Non-factorizable Models using Deep Generative Prior
Approximations

Conor<EMAIL_ADDRESS>1,2, Joshua J. Bon 3, Elizaveta
Semenova 4,
Antonietta Mira 5,6, Kerrie Mengersen 1,2
1 Centre for Data Science, Queensland University of Technology
2 School of Mathematical Sciences, Queensland University of Technology
3 CEREMADE, Université Paris-Dauphine
4 Imperial College London
5 Euler Institute, Università della Svizzera italiana (USI)
6 Insubria University

###### Abstract

Federated learning (FL) allows for collaborative model training across
decentralized clients while preserving privacy by avoiding data sharing.
However, current FL methods assume conditional independence between client
models, limiting the use of priors that capture dependence, such as Gaussian
processes (GPs). We introduce the _S_ tructured _I_ ndependence via deep _G_
enerative _M_ odel _A_ pproximation (SIGMA) prior which enables FL for non-
factorizable models across clients, expanding the applicability of FL to
fields such as spatial statistics, epidemiology, environmental science, and
other domains where modeling dependencies is crucial. The SIGMA prior is a
pre-trained deep generative model that approximates the desired prior and
induces a specified conditional independence structure in the latent
variables, creating an approximate model suitable for FL settings. We
demonstrate the SIGMA prior’s effectiveness on synthetic data and showcase its
utility in a real-world example of FL for spatial data, using a conditional
autoregressive prior to model spatial dependence across Australia. Our work
enables new FL applications in domains where modeling dependent data is
essential for accurate predictions and decision-making.

##  1 Introduction

Federated learning (FL) enables collaborative model estimation across multiple
clients without needing centralized data aggregation (McMahan et al.,, 2017;
Kairouz et al.,, 2021). FL addresses critical data privacy, security, and
ownership concerns by enabling model training on decentralized data. However,
a fundamental limitation of current FL algorithms is the assumption of strict
conditional independence between the data from different clients (Kotelevskii
et al.,, 2022; Hassan et al.,, 2023). This assumption implies that the local
parameters of each client are independent given the global parameters, but
this is often not true in real-world settings, such as spatio-temporal
modeling or graph-structured data, where observations across clients exhibit
complex dependencies. As a result, applying FL to these domains requires
additional, suboptimal assumptions that limit the effectiveness and accuracy
of the learned models.

To address this challenge, we introduce _S tructured Independence via deep
Generative Model Approximation_ (SIGMA) priors, a novel approach to enabling
scalable FL algorithms for models with dependent data. The key idea is to
train a hierarchical variational autoencoder (VAE) to approximate the prior
distribution over client-specific parameters. The SIGMA prior learns a latent
representation that captures dependencies between clients while providing a
conditionally-independent structure, enabling the use of existing efficient FL
algorithms.

We consider probabilistic models of the form

$\displaystyle\bm{\phi}$ $\displaystyle\sim p(\bm{\phi}),$ (1)
$\displaystyle\bm{\theta}_{j}|\bm{\theta}_{-j},\bm{\phi}$ $\displaystyle\sim
p(\bm{\theta}_{j}|\bm{\theta}_{{-j}},\bm{\phi}),\hskip
8.5359ptj\in\\{1,\ldots,J\\},$ (2)
$\displaystyle\bm{y}_{j}|\bm{\theta}_{j},\bm{\phi}$
$\displaystyle\stackrel{{\scriptstyle\rm
ind}}{{\sim}}p(\bm{y}_{j}|\bm{\theta}_{j},\bm{\phi}),\quad
j\in\\{1,\ldots,J\\},$ (3)

where $\bm{\phi}$ are global parameters shared across all $J$ clients,
$\bm{\theta}_{j}$ are local parameters specific to client $j$,
$\bm{\theta}_{{-j}}$ denotes local parameters from all clients except $j$, and
$\bm{y}_{j}$ is the data available to client $j$. Existing FL algorithms, such
as structured federated variational inference (SFVI, Hassan et al.,, 2023) and
FedSOUL (Kotelevskii et al.,, 2022), restrict (2) with a local conditional
independence property, that is,

$\displaystyle
p(\bm{\theta}_{j}|\bm{\theta}_{-j},\bm{\phi})=p(\bm{\theta}_{j}|\bm{\phi}),\quad
j\in\\{1,\ldots,J\\}.$ (4)

However, prior densities that do not satisfy (4) are common in statistics and
machine learning. For example, Markov random field (MRF, Besag et al.,, 1991;
Rue et al.,, 2009) and Gaussian process (GP) priors (Bernardo et al.,, 1998;
Williams and Rasmussen,, 2006) are a common choice to capture spatial and
temporal dependence between observations in applications such as image
inpainting (ruv,, 2014), disease mapping (Lee,, 2011; Riebler et al.,, 2016),
and graph neural networks (Hasanzadeh et al.,, 2020).

Figure 1 illustrates the challenge of applying such priors in a federated
setting. The left panel shows the standard dependency structure for a
conditional autoregressive (CAR) prior in a non-FL setting. In the FL setting
(right panel), the bold borders represent boundaries between clients, and the
red grids indicate values that cannot be evaluated when inferring the value at
the dark-green grids due to the lack of access to data from other clients.

Figure 1: The _standard setting_ on the left shows the dependency structure
for a _conditional autoregressive_ (CAR) prior in a non-FL setting. On the
right, the bold borders represent boundaries for different clients, and the
red grids represent values we cannot evaluate when trying to infer, in an FL
setting, the value at the dark-green grids.

Our work aims to address this challenge by developing a novel approach that
can capture the dependencies among local parameters while still enabling
efficient and privacy-preserving learning in the FL setting. The SIGMA prior
enables the use of computationally scalable FL algorithms for models with
assumed dependencies between observations captured by the prior structure.

The remainder of the paper is structured as follows. Section 2 provides
background on the model class of interest, SFVI, and VAEs. Section 3
introduces the SIGMA prior and our proposed methodology. Section 4 presents
the results of our experiments on synthetic and real-world applications.
Finally, Section 5 concludes the paper with a discussion of the limitations
and future directions for this work.

##  2 Preliminaries

This section lays the foundation for our work by providing an overview of the
techniques used in our approach, namely structured federated variational
inference (SFVI, Hassan et al.,, 2023), (VAEs, Kingma and Welling,, 2013;
Rezende et al.,, 2014), and PriorCVAE (Semenova et al.,, 2023).

We denote the $p\times p$ identity matrix with $I_{p}$. For
$\bm{z}\in\mathbb{R}^{p}$ we say
$\bm{z}\sim\mathcal{N}_{p}(\bm{\mu},\bm{\sigma})$ for $p$-vectors $\bm{\mu}$
and $\bm{\sigma}$ if $z_{i}$ is normally distributed with mean and standard
deviation ($\mu_{i}$, $\sigma_{i}$) and independent of the other elements
$\bm{z}_{-i}$ for each $i\in\\{1,\ldots,p\\}$. If
$\bm{z}\sim\mathcal{N}_{p}(\bm{\mu},\Sigma)$ and $\Sigma$ is a $p\times p$
matrix, then $\bm{z}$ has a multivariate normal distribution with mean
$\bm{\mu}$ and covariance matrix $\Sigma$.

###  2.1 Structured Federated Variational Inference

Variational inference (VI, Blei et al.,, 2017; Zhang et al.,, 2018) is an
algorithmic technique that approximates the posterior distribution
$p(\bm{\theta}|\bm{y})$, which is often intractable for analytical or direct
computation due to the complexity of the model. The key idea of VI is to
introduce a parameterized distribution
$q_{\bm{\lambda}}(\bm{\theta})\in\mathcal{Q}$ called the _variational
approximation_ , that aims to approximate the posterior distribution. The
variational parameters $\bm{\lambda}$ are optimized to find the best
approximation from a fixed _variational family_ of distributions,
$\mathcal{Q}$.

The objective in the VI framework is to maximize the _evidence lower bound_
(ELBO), defined as:

$\displaystyle\mathcal{L}(\bm{\lambda}):=\mathbb{E}_{\bm{\theta}\sim
q_{\bm{\lambda}}}[\log p(\bm{\theta},\bm{y})-\log
q_{\bm{\lambda}}(\bm{\theta})],$ (5)

which acts as a surrogate objective function, providing a lower bound to the
log marginal likelihood. Maximizing the ELBO is equivalent to minimizing the
Kullback-Leibler (KL) divergence between the variational approximation
$q_{\bm{\lambda}}(\bm{\theta})$ and the posterior distribution
$p(\bm{\theta}|\bm{y})$. When the variational approximation matches the true
posterior, the ELBO equals the log marginal likelihood.

Gradient-based optimization algorithms (Ruder,, 2017) are commonly used to
maximize the ELBO, relying on unbiased gradient estimators. A significant
advancement in stochastic VI (Hoffman et al.,, 2013) is the
_reparameterization gradient estimator_ (Titsias and Lázaro-Gredilla,, 2014;
Kingma and Welling,, 2013; Rezende et al.,, 2014), which allows for an
efficient and low-variance estimate of the ELBO. The key idea is to express a
sample of $\bm{\theta}$ drawn from $q_{\bm{\lambda}}(\bm{\theta})$ as a
deterministic function of $\bm{\lambda}$ by transforming random noise
$\bm{\epsilon}$: $\bm{\theta}=f_{\bm{\lambda}}(\bm{\epsilon})$. The noise
variable $\bm{\epsilon}$ is drawn from a distribution $q_{\bm{\epsilon}}$,
which is independent of $\bm{\lambda}$. This reparameterization enables the
computation of low-variance gradients, which is crucial for efficient
optimization in VI. A resulting single-sample unbiased Monte Carlo estimator
of the desired gradient vector is the _sticking-the-landing_ (STL) estimator
(Roeder et al.,, 2017),

$\displaystyle\widehat{\nabla}_{\bm{\lambda}}\mathcal{L}:={\frac{\partial
f_{\bm{\lambda}}(\bm{\epsilon})}{\partial\bm{\lambda}}}^{\top}\nabla_{\bm{\theta}}[\log
p(\bm{\theta},\bm{y})-\log q_{\bm{\lambda}}(\bm{\theta})].$ (6)

Structured federated variational inference (SFVI, Hassan et al.,, 2023)
extends VI to the FL setting, focusing on hierarchical Bayesian models of the
form,

$\displaystyle
p(\bm{y},\bm{\phi},\bm{\theta})=p(\bm{\phi})\prod_{j=1}^{J}p(\bm{\theta}_{j}|\bm{\phi})p(\bm{y}_{j}|\bm{\theta}_{j},\bm{\phi}).$

where the local parameters $\bm{\theta}_{j}$ only appear in the likelihood
component for the $j$-th client, and the sets of local parameters are assumed
to be conditionally independent given the global parameters. To enable
efficient and privacy-preserving learning in the federated setting, SFVI
employs a factorized structured variational approximation (Hoffman and Blei,,
2015; Ranganath et al.,, 2016; Tan et al.,, 2020) of the form:

$\displaystyle
q(\bm{\phi},\bm{\theta})=q_{\bm{\varphi}_{G}}(\bm{\phi})\prod_{j=1}^{J}q_{\bm{\varphi}_{L_{j}}}(\bm{\theta}_{j}|\bm{\phi}),$

where $q_{\bm{\varphi}_{G}}(\bm{\phi})$ is a Gaussian distribution
parameterized by $\bm{\varphi}_{G}$, and each
$q_{\bm{\varphi}_{L_{j}}}(\bm{\theta}_{j}|\bm{\phi})$ is a Gaussian
distribution parameterized by $\bm{\varphi}_{L_{j}}$, potentially including
dependencies on $\bm{\phi}$. This factorization allows for local optimization
and privacy-preserving updates, as the local variational parameters
$\bm{\varphi}_{L_{j}}$ can be updated independently by each client without
sharing them with the server. SFVI leverages the factorized structure of the
variational approximation to decompose the ELBO into a sum of local ELBOs,
each of which can be optimized locally by a client. Two algorithms, SFVI and
SFVI-Avg, are derived for combinations of models and variational
approximations that share this form. Both algorithms enable parallel and
private updates of the local variational parameters $\bm{\varphi}_{L_{j}}$ for
each client, requiring only the communication of gradient information
concerning the global variational parameters $\bm{\varphi}_{G}$ to the server
at each iteration.

In the numerical examples presented in this work, we focus on the SFVI
algorithm, as it returns the exact estimates that fit the equivalent
variational approximation on the full dataset in the non-FL setting. While
SFVI provides a principled and efficient approach to learning hierarchical
Bayesian models in federated settings, it assumes that the local parameters
are conditionally independent given the global parameters. This assumption may
limit the applicability of SFVI to models with dependent local parameters,
which are common in real-world applications. In Section 3, we introduce the
SIGMA prior, a novel approach to relaxing this assumption and enabling the
application of SFVI to a broader class of models with dependent local
parameters. By capturing the dependencies among local parameters while still
allowing for efficient and privacy-preserving learning, the SIGMA prior has
the potential to significantly expand the scope and impact of FL in various
domains.

###  2.2 Variational Autoencoders and PriorCVAE

####  2.2.1 Variational Autoencoders

_Variational autoencoders_ (VAE, Kingma and Welling,, 2013; Rezende et al.,,
2014) are a class of latent variable models that use variational approximation
techniques to facilitate estimation of the latent space $\bm{z}$. A VAE
consists of a generative model,
$\tilde{p}_{\bm{\psi}}(\bm{\theta},\bm{z})=\tilde{p}_{\bm{\psi}}(\bm{\theta}|\bm{z})\tilde{p}(\bm{z})$,
coupled with $\tilde{q}_{\bm{\varphi}}(\bm{z}|\bm{\theta})$, an amortised
variational approximation of $\tilde{p}_{\bm{\psi}}(\bm{z}|\bm{\theta})$, the
posterior distribution implied by the generative model. The use of the
variational approximation facilitates learning the components of the VAE. We
will refer to the generative model of the VAE simply as the _generative model_
and use $\tilde{p}$ to distinguish this model from the true model that it is
approximating.

The task of a VAE is to approximate some distribution $p(\bm{\theta})$ using a
latent variable representation
$\tilde{p}_{\bm{\psi}}(\bm{\theta})=\int\tilde{p}_{\bm{\psi}}(\bm{\theta}|\bm{z})\tilde{p}(\bm{z})\text{d}z$
given by the chosen generative model. As in VI (see Section 2.1), the
objective function employed by VAEs is the appropriate average ELBO,

$\mathcal{L}(\bm{\psi},\bm{\varphi})=\mathbb{E}_{\bm{\theta}\sim
p}\mathbb{E}_{\bm{z}\sim\tilde{q}_{\bm{\varphi}}(\bm{z}|\bm{\theta})}[\log\tilde{p}_{\bm{\psi}}(\bm{\theta},\bm{z})-\log\tilde{q}_{\bm{\varphi}}(\bm{z}|\bm{\theta})],$
(7)

where the first average is taken with respect to the true data generating
process $\bm{\theta}\sim p(\bm{\theta})$. Note that replacing the
approximation $\tilde{q}_{\bm{\varphi}}(\bm{z}|\bm{\theta})$ with the correct
$\tilde{p}_{\bm{\psi}}(\bm{z}|\bm{\theta})$ simplifies (7) to
$\mathcal{L}(\bm{\psi},\bm{\varphi})=\mathbb{E}_{\bm{\theta}\sim
p}[\log\tilde{p}_{\bm{\psi}}(\bm{\theta})]$ which is the full objective
function, as opposed to the ELBO. As such, we can view the variational
posterior, $\tilde{q}_{\bm{\varphi}}(\bm{z}|\bm{\theta})$, as an approximation
to avoid the difficulties in evaluating
$\tilde{p}_{\bm{\psi}}(\bm{z}|\bm{\theta})$.

Deep latent Gaussian models (DLGMs, Rezende et al.,, 2014) are a popular type
of generative model commonly used within a VAE. The generative model for a
single-layer DLGM consists of a latent variable
$\bm{z}\sim\mathcal{N}_{p}(\bm{0},I_{p})$ and conditional variable
$\bm{\theta}|\bm{z}\sim\mathcal{N}_{d}(\bm{\mu}_{\bm{\psi}}(\bm{z}),\bm{\sigma}_{\bm{\psi}}(\bm{z}))$
where
$\mathcal{D}_{\bm{\psi}}=(\bm{\mu}_{\bm{\psi}},\log\bm{\sigma}_{\bm{\psi}})$
are obtained via the joint output of a neural network termed the _decoder_ ,
parameterized by $\bm{\psi}$. The variational approximation
$\tilde{q}_{\bm{\varphi}}(\bm{z}|\bm{\theta})$ is constructed from an
_encoder_ ,
$\mathcal{E}_{\bm{\varphi}}=(\bm{\mu}_{\bm{\varphi}},\log\bm{\sigma}_{\bm{\varphi}})$
such that
$\bm{z}|\bm{\theta}\sim\mathcal{N}_{p}(\bm{\mu}_{\bm{\varphi}}(\bm{\theta}),\bm{\sigma}_{\bm{\varphi}}(\bm{\theta}))$.
The combination of using a decoder in the generative model and using an
encoder to parameterize the variational approximation results in a
probabilistic analog to the autoencoder (Tschannen et al.,, 2018).

VAEs were originally designed to construct a generative model for finite
observed data, say $\\{x_{i}\\}_{i=1}^{n}$. In this case, $p(\bm{\theta})$ is
an empirical distribution,
$p(\bm{\theta})=\sum_{i=1}^{n}\delta_{x_{i}}(\bm{\theta})$, and the outer
expectation of (7) becomes an empirical average over the data
$\\{x_{i}\\}_{i=1}^{n}$. Estimation then typically involves stochastic
gradient descent (SGD, Ruder,, 2017) to handle the inner expectation of (7).
Mini-batching of the data (Hoffman et al.,, 2013) is also frequently used and
provides an unbiased estimate of the outer expectation of (7). VAEs can also
be used when $p(\bm{\theta})$ is a non-empirical distribution that can be
sampled from (e.g. continuous). In this case, SGD is used to handle both
expectations. We present this approach for the case of prior approximation
next.

####  2.2.2 PriorVAE and PriorCVAE

PriorVAE (Semenova et al.,, 2022) adopts a VAE approach to learn an
approximate prior distribution, which can be used in place of the desired
prior. Such an approximation aims to reduce the dimensionality or difficulty
of an inferential task when the prior is the computational bottleneck because,
for example, it is costly to evaluate. PriorVAE uses a decoder
$\mathcal{D}_{\bm{\psi}}=(\bm{\mu}_{\bm{\psi}},\log\gamma)$ with fixed
variance $\gamma^{2}$ and optimizes (7) using SGD with draws $\bm{\theta}\sim
p(\bm{\theta})$ from the desired prior distribution.

The approximate prior is constructed from the VAE by taking
$\bm{\theta}=\bm{\mu}_{\bm{\psi}}(\bm{z})$ with $\bm{z}\sim
N_{p}(\bm{0},I_{p})$, or more generally $\bm{z}\sim\tilde{p}(\bm{z})$, where
$\bm{\mu}_{\bm{\psi}}(\bm{z})$ is the mean output of the decoder. Using the
approximate prior changes the posterior distribution from
$\pi(\bm{\theta}|\bm{x})\propto p(\bm{x}|\bm{\theta})p(\bm{\theta})$ over
$\bm{\theta}$ to the approximate posterior $\tilde{\pi}(\bm{z}|\bm{x})\propto
p(\bm{x}|\bm{\mu}_{\bm{\psi}}(\bm{z}))\tilde{p}(\bm{z})$ over $\bm{z}$. The
$\bm{\theta}$ implied by the approximate posterior is recovered using
$\bm{\theta}=\bm{\mu}_{\bm{\psi}}(\bm{z})$. One can choose the dimension of
the latent variables $p\ll d$, ensuring that the effective dimension of the
approximate posterior is much less than the original posterior. Large
computational benefits during inference were observed by Semenova et al.,
(2022), who motivate their work from this perspective.

PriorCVAE (Semenova et al.,, 2023) extends PriorVAE to approximate conditional
priors of the form
$p(\bm{\theta},\bm{\phi})=p(\bm{\theta}|\bm{\phi})p(\bm{\phi})$ where
$\bm{\phi}$ is a random variable that forms a hierarchy in the prior
structure. Rather than approximating $p(\bm{\theta},\bm{\phi})$ entirely,
PriorCVAE only approximates the conditional prior $p(\bm{\theta}|\bm{\phi})$,
thus retaining $p(\bm{\phi})$ in both the original and approximate prior. With
this extension, the ELBO (7) becomes

$\mathcal{L}(\bm{\psi},\bm{\varphi})=\mathbb{E}_{(\bm{\theta},\bm{\phi})\sim
p}\mathbb{E}_{\bm{z}\sim\tilde{q}_{\bm{\varphi}}(\bm{z}|\bm{\theta},\bm{\phi})}[\log\tilde{p}_{\bm{\psi}}(\bm{\theta},\bm{z}|\bm{\phi})-\log\tilde{q}_{\bm{\varphi}}(\bm{z}|\bm{\theta},\bm{\phi})].$
(8)

For the remainder of the paper, we will only consider the conditional variant
of VAEs. To illustrate the PriorCVAE approximation we present a toy Gaussian
Process regression example in Table 1.

True model | PriorCVAE model
---|---
$\sigma\sim p(\sigma)$ | $\sigma\sim p(\sigma)$
$\phi\sim p(\phi)$ | $\phi\sim p(\phi)$
| $\bm{z}\sim\mathcal{N}_{p}(\bm{0},\mathrm{I}_{p})$
$\bm{\theta}|\phi\sim\mathcal{N}(\bm{0},\mathrm{k}_{\bm{x}}(\phi))$ | $\bm{\theta}=\bm{\mu}_{\bm{\psi}}(\bm{z},\phi)$
$\bm{y}\sim\mathcal{N}(\bm{\theta},\sigma)$ | $\bm{y}\sim\mathcal{N}(\bm{\theta},\sigma)$
Table 1: True and PriorCVAE model comparison for a Gaussian Process regression
example.

In this example, $\mathrm{k}_{\bm{x}}(\phi)$ is some kernel function evaluated
at points $\bm{x}$ for a given value of the length scale $\phi$. The VAE is
trained using prior draws from the true model. In this instance, the model
approximation provided by PriorCVAE replaces the GP prior with a deterministic
reconstruction of the GP given $\bm{z}$ and $\phi$, thus reducing the
dimensionality of the inference task to be proportional to the dimensionality
of $\bm{z}$ rather than $\bm{\theta}$. The latter of which scales with the
number of observations in the case of a GP regression.

##  3 Structured Independence via Generative Model Approximations

In this section, we introduce the _S tructured Independence via Generative
Model Approximation_ (SIGMA) prior. A SIGMA prior approximates the latent
structure of a probabilistic model using a VAE with a generative model that
satisfies desired conditional independence properties. We motivate the use of
SIGMA priors FL scenarios where the prior structure of the desired model does
not factorize over clients. We learn a SIGMA prior with the desired
conditional independence and factorization properties, which replaces the
original prior. The resulting approximate probabilistic model with the SIGMA
prior is then amenable to use with existing FL algorithms, such as SFVI
(Section 2.1).

###  3.1 SIGMA Prior

The SIGMA prior is a novel approach to approximate the hierarchical Bayesian
model described by equations (1)-(2). It introduces $J+1$ latent variables
$\bm{z}_{G},\bm{z}_{L}$, where
$\bm{z}_{L}=(\bm{z}_{L_{1}},\ldots,\bm{z}_{L_{j}})$, and $J$ is the number of
clients, to capture both the shared and client-specific variations in the
local parameters $\bm{\theta}_{j}$. The global latent variable $\bm{z}_{G}$
induces dependence across clients, while the local latent variables
$\bm{z}_{L_{j}}$ drive the client-specific variations. We choose the
hierarchical structure to enable the desired conditional independence
properties across clients, facilitating the use of computationally efficient
FL algorithms. The form of the generative model is

$\begin{split}\bm{\phi}&\sim p(\bm{\phi}),\\\
\bm{z}_{G}&\sim\tilde{p}(\bm{z}_{G}),\\\
\bm{z}_{L_{j}}&\sim\tilde{p}(\bm{z}_{L_{j}}),\quad j\in\\{1,\ldots,J\\},\\\
\bm{\theta}_{j}&\sim\mathcal{N}_{d_{j}}(\bm{\mu}_{\bm{\psi}_{L_{j}},\bm{\psi}_{G}}(\bm{z}_{L_{j}},\bm{\phi}),\gamma
I_{d_{j}}),\quad j\in\\{1,\ldots,J\\},\end{split}$ (9)

where $\bm{\phi}\in\mathbb{R}^{d}$, $\bm{z}_{G}\in\mathbb{R}^{p}$,
$\bm{z}_{L_{j}}\in\mathbb{R}^{p_{j}},\bm{\theta}_{j}\in\mathbb{R}^{d_{j}}$ for
$j\in\\{1,\ldots,J\\}$, and $\gamma>0$. The $J$ functions
$\bm{\mu}_{\bm{\psi}_{L_{j}},\bm{\psi}_{G}}:\mathbb{R}^{p_{j}+p+d}\rightarrow\mathbb{R}^{d_{j}}$
are neural networks, with parameters $\bm{\psi}_{L_{j}}$ and $\bm{\psi}_{G}$
respectively, used to produce the decoder for the VAE.

In this work, we employ a hierarchical structure for the decoders, where a
_global decoder_
$\bm{g}_{\bm{\psi}_{G}}:\mathbb{R}^{p+d}\rightarrow\mathbb{R}^{k}$ is used as
input to each _local decoder_
$\bm{f}_{\bm{\psi}_{L_{j}}}:\mathbb{R}^{p_{j}+d+k}\rightarrow\mathbb{R}^{d_{j}}$
with the following form,

$\bm{\mu}_{\bm{\psi}_{L_{j}},\bm{\psi}_{G}}(\bm{z}_{L_{j}},\bm{\phi})=\bm{f}_{\bm{\psi}_{L_{j}}}(\bm{z}_{L_{j}},\bm{\phi},\bm{g}_{\bm{\psi}_{G}}(\bm{z}_{G},\bm{\phi})).$
(10)

This hierarchical structure aids high-dimensional learning by allowing each
local decoder to focus on learning the local deviations from a shared global
representation captured by the global decoder $\bm{g}_{\bm{\psi}_{G}}$. The
scale parameter $\gamma$ completes the overall decoder specification, denoted
as
$\mathcal{D}_{\bm{\psi}}=(\bm{f}_{\bm{\psi}_{L_{1}}},\ldots,\bm{f}_{\bm{\psi}_{L_{J}}},\bm{g}_{\bm{\psi}_{G}},\log\gamma)$.
Figure 2 illustrates the generative model used to create the SIGMA prior,
approximating the prior density (2).

Figure 2: Graphical description of the generative model fit to input prior
draws $\bm{\theta}$ to create the SIGMA prior. The output is
$\hat{\bm{\theta}}=(\hat{\bm{\theta}}_{1}^{\top},\ldots,\hat{\bm{\theta}}_{J}^{\top})^{\top}$,
where each of the $J$ blocks of parameters $\hat{\bm{\theta}}_{j}$ is
conditionally independent given the global latent variable $\bm{z}_{G}$ and
the deterministic output $\bm{h}$ of a neural network $f_{\bm{\psi}_{G}}$, the
global decoder, that takes $\bm{z}_{G}$ and $\phi$ as input. The SIGMA prior
is trained as a hierarchical variational autoencoder with a specific
structure, such that the local components of the encoders and decoders
factorize over clients.

To learn the generative model, we approximate the posterior of the latent
variables using a variational family
$\tilde{q}_{\bm{\varphi}}(\bm{z}_{L},\bm{z}_{G}|\bm{\theta},\bm{\phi})$ with
the following structure,

$\begin{split}(\bm{z}_{G}|\bm{\phi},\bm{\theta})&\sim\mathcal{N}_{p}(\bm{\mu}_{\bm{\varphi}_{G}}(\bm{\theta},\bm{\phi}),\bm{\sigma}_{\bm{\varphi}_{G}}(\bm{\theta},\bm{\phi}))\\\
(\bm{z}_{L_{j}}|\bm{z}_{G},\bm{\phi})&\sim\mathcal{N}_{p_{j}}(\bm{\mu}_{\bm{\varphi}_{L_{j}}}(\bm{z}_{G},\bm{\phi}),\bm{\sigma}_{\bm{\varphi}_{L_{j}}}(\bm{z}_{G},\bm{\phi})),\quad
j\in\\{1,\ldots,J\\},\end{split}$ (11)

which factorizes as
$\tilde{q}_{\bm{\varphi}}(\bm{z}_{L},\bm{z}_{G}|\bm{\theta},\bm{\phi})=\tilde{q}_{\bm{\varphi}_{G}}(\bm{z}_{G}|\bm{\theta},\bm{\phi})\prod_{j=1}^{J}\tilde{q}_{\bm{\varphi}_{L_{j}}}(\bm{z}_{L_{j}}|\bm{z}_{G},\bm{\phi})$.
We chose the factorization to mirror the conditional independence properties
of the SIGMA prior, enabling efficient learning and inference. The _global
encoder_ is parameterized by
$\bm{\mu}_{\bm{\varphi}_{G}}:\mathbb{R}^{d_{1}+\cdots+d_{J}+d}\rightarrow\mathbb{R}^{p}$
and
$\bm{\sigma}_{\bm{\varphi}_{G}}:\mathbb{R}^{d_{1}+\cdots+d_{J}+d}\rightarrow\mathbb{R}^{p}$,
which are jointly output by a neural network. Similarly, the _local encoders_
are parameterized by
$\bm{\mu}_{\bm{\varphi}_{L_{j}}}:\mathbb{R}^{p+d}\rightarrow\mathbb{R}^{p}$
and
$\bm{\sigma}_{\bm{\varphi}_{L_{j}}}:\mathbb{R}^{p+d}\rightarrow\mathbb{R}^{p}$
for each $j\in\\{1,\ldots,J\\}$. The complete encoder structure is denoted as
$\mathcal{E}_{\bm{\varphi}}=(\bm{\mu}_{\bm{\varphi}_{L_{1}}},\ldots,\bm{\mu}_{\bm{\varphi}_{L_{J}}},\bm{\sigma}_{\bm{\varphi}_{L_{1}}},\ldots,\bm{\sigma}_{\bm{\varphi}_{L_{J}}},\bm{\mu}_{\bm{\varphi}_{G}},\bm{\sigma}_{\bm{\varphi}_{G}})$.
We provide the architecture details of the encoders and decoders used in the
numerical experiments in Section A.

| Generative | Variational | Prior
---|---|---|---
| model | posterior | approximation
PriorVAE | $\begin{aligned} \bm{z}&\sim\mathcal{N}_{p}(\bm{0},I_{p})\\\ (\bm{\theta}|\bm{z})&\sim\mathcal{N}_{d}(\bm{\mu}_{\bm{\psi}}(\bm{z}),\gamma I_{d})\end{aligned}$ | $\begin{aligned} (\bm{z}&|\bm{\theta})\sim\\\ &\mathcal{N}_{p}(\bm{\mu}_{\bm{\varphi}}(\bm{\theta}),\bm{\sigma}_{\bm{\varphi}}(\bm{\theta}))\end{aligned}$ | $\begin{aligned} \bm{z}&\sim\mathcal{N}_{p}(\bm{0},I_{p})\\\ \bm{\theta}&=\bm{\mu}_{\bm{\psi}}(\bm{z})\end{aligned}$
PriorCVAE | $\begin{aligned} \bm{\phi}&\sim p(\bm{\phi})\\\ \bm{z}&\sim\mathcal{N}_{p}(\bm{0},I_{p})\\\ (\bm{\theta}|\bm{z},\bm{\phi})&\sim\mathcal{N}_{d}(\bm{\mu}_{\bm{\psi}}(\bm{z},\bm{\phi}),\gamma I_{d})\end{aligned}$ | $\begin{aligned} (\bm{z}&|\bm{\theta},\bm{\phi})\sim\\\ &\mathcal{N}_{p}(\bm{\mu}_{\bm{\varphi}}(\bm{\theta},\bm{\phi}),\bm{\sigma}_{\bm{\varphi}}(\bm{\theta},\bm{\phi}))\end{aligned}$ | $\begin{aligned} \bm{\phi}&\sim p(\bm{\phi})\\\ \bm{z}&\sim\mathcal{N}_{p}(\bm{0},I_{p})\\\ \bm{\theta}&=\bm{\mu}_{\bm{\psi}}(\bm{z},\bm{\phi})\end{aligned}$
SIGMA Prior | Eq. (9) | Eq. (11) | Eq. (12)
Table 2: Comparison of Prior(C)VAE with SIGMA Prior: the columns show the
generative model, the variational posterior approximation, and the
approximation to the prior used for parameter inference.

###  3.2 Approximation of the prior

We construct the SIGMA prior approximation from the estimated generative model
and variational posterior of the VAE. In particular, we use

$\begin{split}\bm{\phi}&\sim p(\bm{\phi}),\\\
\bm{z}_{G}&\sim\tilde{p}(\bm{z}_{G}),\\\
\bm{z}_{L_{j}}&\sim\mathcal{N}_{p_{j}}(\bm{\mu}_{\bm{\varphi}_{L_{j}}}(\bm{z}_{G},\bm{\phi}),\bm{\sigma}_{\bm{\varphi}_{L_{j}}}(\bm{z}_{G},\bm{\phi})),\quad
j\in\\{1,\ldots,J\\},\\\
\bm{\theta}_{j}&\sim\mathcal{N}_{d_{j}}(\bm{\mu}_{\bm{\psi}_{L_{j}},\bm{\psi}_{G}}(\bm{z}_{L_{j}},\bm{z}_{G},\bm{\phi}),\gamma
I_{d_{j}}),\quad j\in\\{1,\ldots,J\\},\end{split}$ (12)

which is identical to the generative model (9), except that the distribution
of each local latent variable $\bm{z}_{L_{j}}$ is replaced by its variational
approximation. Using the variational approximation of $\bm{z}_{L_{j}}$ as the
prior is a significant advantage of the SIGMA prior compared to methods like
PriorCVAE, as it allows for incorporating the learned variational
approximation into the approximate prior. The hierarchical structure of the
SIGMA prior enables this replacement, as the variational approximations for
the local latent variables are conditioned only on the global latent variable
$\bm{z}_{G}$ and the global parameters $\bm{\phi}$, which are already part of
the approximate prior specification.

The SIGMA prior offers significant benefits by incorporating learned
variational approximations for the local latent variables, which represent the
optimal density of the latent space for reconstructing draws of $\bm{\theta}$.
These local latent variables capture important information about the structure
and dependencies within $\bm{\theta}$, enabling the SIGMA prior to effectively
model and generate realistic samples, setting it apart from approaches like
PriorVAE and PriorCVAE.

Model 1 shows pseudocode for a one-dimensional GP regression. Model 2 shows
pseudocode for an _approximation_ to Model 1 using a SIGMA prior, where the
encoders and decoders are trained using prior draws from the length scale
parameter $\phi$ and the random variables $\bm{\theta}$, denoted in Model 1 as
phi and theta respectively. The pseudocode closely resembles how models are
defined in modern Python probabilistic programming languages such as Pyro
(Bingham et al.,, 2019), Numpyro (Phan et al.,, 2019), and PyMC (Abril-Pla et
al.,, 2023). The function call deterministic is employed to define variables
whose values are computed directly from other variables without any inherent
randomness. The function call sample is used for drawing random variables from
probability distributions, and the context plate represents a mechanism for
indicating conditional independence across different subsets of random
variables and data.

⬇

1def model(y, x, k):

2 # lengthscale prior

3 phi = sample(Uniform(0.01, 1))

4 # observation noise prior

5 sigma = sample(HalfNormal(0, 1))

6 # form covariance matrix

7 cov = deterministic(k(x, x, phi))

8 # GP prior

9 theta = sample(MultivariateNormal(0, cov))

10 # likelihood

11 sample(Normal(theta, sigma), observed=y)

Model 1: Pseudocode for a 1D GP regression model

⬇

1def approximate_model(ys, mus, encoders, decoders):

2 # lengthscale prior

3 phi = sample(Uniform(0.01, 1))

4 # observation noise prior

5 sigma = sample(HalfNormal(0, 1))

6 # prior for global latents

7 z_G = sample(Normal(0, 1))

8 # variables inside plate are conditionally independent across clients

9 with plate("clients", len(mus)):

10 # prior for local latents

11 mu[j], logvar[j] = deterministic(encoders[j](z_G, phi))

12 z_L[J] = sample(Normal(mu[j], exp(0.5*logvar[j])))

13 # estimate likelihood mean

14 theta[j] = deterministic(decoders[j](z_G, z_L[j], phi))

15 # likelihood

16 sample(Normal(theta[j], sigma), observed=y[j])

Model 2: Pseudocode for the SIGMA prior approximation of a 1D GP regression
model

The SIGMA prior approximation enables the inference of the latent variables
$\bm{z}_{G},\bm{z}_{L_{1}},\ldots,\bm{z}_{L_{J}}$ instead of the original
local parameters $\bm{\theta}_{1},\ldots,\bm{\theta}_{J}$. Notably, the SIGMA
prior ensures the required conditional independence:
$\hat{\bm{\theta}}_{j}\perp\\!\\!\\!\perp\hat{\bm{\theta}}_{k}\mid\bm{z}_{G},\bm{\phi}$
for all pairs of clients $j\neq k$. This property is useful in FL settings, as
it allows each client to perform local updates independently without the need
for sharing raw data, thereby facilitating efficient and privacy-preserving
learning. Furthermore, this property enables the use of existing FL algorithms
such as SFVI (Hassan et al.,, 2023) or FedSOUL (Kotelevskii et al.,, 2022),
which rely on the assumption of conditional independence among local
parameters given the global parameters.

##  4 Numerical Examples

This section presents two numerical experiments that showcase the
effectiveness of the SIGMA prior in handling highly dependent priors with
increasing dimensionality and numbers of clients. These experiments
demonstrate the scalability and adaptability of our approach in various
settings, ranging from simple one-dimensional examples to complex real-world
spatial structures.

The first experiment establishes a baseline for the SIGMA prior’s performance
using a one-dimensional GP prior with $p=100$ dimensions and $J=3$ clients. We
provide qualitative assessments through visualizations, including the
empirical covariance of the SIGMA prior compared to the covariance of the GP
prior, reconstructed draws from the SIGMA prior, and comparison of estimates
when using the SIGMA prior to construct an approximate model compared to the
true underlying model that contains the GP prior. The second experiment
demonstrates the SIGMA prior’s ability to handle complex, real-world spatial
structures by considering a prior structure with dimension $p=2227$ and $J=8$
clients, motivated by the Australian geography. This experiment showcases the
scalability of our approach to more significant and realistic problems.

Throughout these experiments, we employ stochastic variational inference
implemented in the Pyro library (Bingham et al.,, 2019) to train the VAE, with
encoders and decoders constructed using PyTorch (Paszke et al.,, 2019). For
posterior estimation in the approximate model, we use the _No-U-Turn sampler_
(NUTS, Hoffman et al.,, 2014) algorithm in the first example, and mean-field
variational inference in the second example. Both algorithms are implemented
in the BlackJAX (Cabezas et al.,, 2024) library.

To evaluate the performance of the SIGMA prior, we use a combination of
qualitative and quantitative metrics tailored to each experiment. Experiment 1
focuses on visual assessments. In Experiment 2, we compare inference estimates
from the approximate SIGMA prior model with MCMC estimates of the true model.
Additionally, we propose a novel correction by viewing the approximate model
created with the SIGMA prior as an auxiliary-variable model, emphasizing how
the SIGMA prior enables similar parameter estimates using an FL algorithm
compared to MCMC and VI results on the actual model.

The following subsections (Sections 4.1 and 4.2) present each experiment in
detail, discussing the specific settings, results, and insights gained.
Appendix A describes the architecture details of the VAE common to the two
experiments.

###  4.1 One-dimensional Gaussian Process Regression

In our first example, we demonstrate the effectiveness of the SIGMA prior in
approximating a one-dimensional GP regression. We generate training data by
sampling from a GP with a radial basis function (RBF) kernel at 100
equidistant points $\bm{x}_{\text{train}}$ in the interval $[0,1]$. The data
is divided among three clients based on the input values: client 1 receives
data associated with $\bm{x}\in[0,0.34)$, client 2 receives data associated
with $\bm{x}\in[0.34,0.67)$, and client 3 receives data associated with
$\bm{x}\in[0.67,1.00]$.

We generate $10,000$ draws from a GP prior with an RBF kernel with the
lengthscale parameter varying uniformly between $0.2$ and $1$ and with
variance fixed at $1$. We train the SIGMA prior architecture for $20,000$
iterations using the Adam optimizer (Kingma and Ba,, 2014) with a learning
rate of $10^{-3}$ in Pyro (Bingham et al.,, 2019). The batch size in each
iteration is $32$, and the dimension of the hidden layers used in the neural
net architectures is $16$. The dimension of the global latent space
$\bm{z}_{G}$ and each of the three local latent spaces
$\bm{z}_{L_{1}},\bm{z}_{L_{2}},\bm{z}_{L_{3}}$ were chosen to be $5$
respectively, and the standard deviation of the decoder is set to
$\gamma=0.5$.

We train the SIGMA prior and compare its empirical covariance with the actual
RBF kernel. Figure 3 visualizes this comparison, demonstrating that the
learned SIGMA prior closely approximates the actual covariance structure for
varying values of the length scale parameter.

Figure 3: Comparison of the empirical covariance of the draws from the RBF
kernel (top row) and the learned SIGMA prior (bottom row) for
$\phi=(0.2,0.5,0.8)$. The SIGMA prior closely approximates the true covariance
structure, capturing client dependencies.

To further illustrate the SIGMA prior’s ability to capture dependencies across
clients, we generate $25$ prior draws from the fitted SIGMA prior (Figure 4).
The red-vertical-dashed lines denote the boundaries between clients, and the
smooth transitions across these boundaries demonstrate that the SIGMA prior
has learned the underlying correlations in the data.

Figure 4: $25$ draws from the fitted SIGMA prior. The red-vertical-dashed
lines denote boundaries between clients, and the colors of each line denote
the length scale of the Gaussian process kernel that the particular line is
approximating.

We now use the SIGMA approximation in an inferential model. We generate 15
random locations $\bm{x}$ in the interval $[0,1]$ and observations
$\bm{y}=\sin(\bm{x})+0.2\cdot{\bm{\epsilon}},$ where
$\bm{\epsilon}\sim\mathcal{N}(0,1)$. Figure 5 shows the mean estimate (solid
blue line) and 90% credible interval (shaded region) for the mean function
using the SIGMA approximation, with the posterior estimated using the NUTS
algorithm.

Figure 5: Mean estimate (solid blue line) and 90% credible interval (shaded
region) for the mean function using the SIGMA approximation. The posterior was
estimated using the NUTS algorithm. The red dashed lines separate the
different clients.

This example showcases the effectiveness of the SIGMA prior in approximating a
one-dimensional GP regression while capturing dependencies across clients in a
FL setting. The learned SIGMA prior closely approximates the covariance
structure, generates prior draws that smoothly transition across client
boundaries, and provides a reasonable estimate of the posterior mean.

###  4.2 Fitting Bayesian spatial models over the Australian geography

In this experiment, we aim to demonstrate the effectiveness of the SIGMA prior
in approximating the spatial structure of the Australian geography, using a
proper conditional autoregressive (PCAR) prior (see Appendix B) to model
dependencies at the Statistical Area Level 2 (SA2). This experiment is
motivated by the Australian Cancer Atlas (ACA, Duncan et al.,, 2019), which
produces estimates of cancer diagnoses, excess deaths, and risk factors across
the Australian geography at the SA2 level. The ACA comprises over one hundred
individual spatial and spatio-temporal models, each employing a form of a CAR
prior to smooth estimates and to account for spatial dependence in the data.
Integrating new data sources into the ACA poses challenges, as it requires
establishing new data agreements across the Australian geography. The SIGMA
prior offers a potential solution by enabling the approximation of complex
spatial structures while facilitating efficient and privacy-preserving
learning in a federated setting.

In this experiment, we consider $J=8$ clients representing Australia’s eight
states and territories. The parameter vector $\bm{\theta}\in\mathbb{R}^{2221}$
is assigned a PCAR prior distribution,
$\bm{\theta}|\phi\sim\mathcal{N}(\bm{0},\mathrm{\Sigma}^{-1})$, where
$\mathrm{\Sigma}\in\mathbb{R}^{2221\times 2221}$ is the precision matrix
defined as $\mathrm{\Sigma}=\mathrm{D}-\phi\mathrm{W}$, where $\mathrm{W}$ is
the adjacency matrix, with $\mathrm{W}_{ij}=1$ if observations $i$ and $j$ are
neighbors and $\mathrm{W}_{ij}=0$ otherwise. The diagonal matrix $\mathrm{D}$
captures the number of neighbors for each observation, with $\mathrm{D}_{ii}$
representing the number of neighbors for observation $i$. The unknown
parameter $\phi\in(0,1)$ is interpreted as a measure of the spatial smoothing
between neighbors, with higher values indicating stronger spatial dependence.

We generate $10,000$ draws from the PCAR prior over the Australian geography,
with $\phi$ varying uniformly between $0$ and $1$. We train the SIGMA prior
architecture for $400,000$ iterations using the Adam optimizer (Kingma and
Ba,, 2014) with a learning rate of $1e^{-4}$ in Pyro (Bingham et al.,, 2019).
The batch size in each iteration is $128$, and the dimension of the hidden
layers used in the neural net architectures is $2024$ for the decoders and the
global encoder and $1048$ for the local encoders. The dimension of the global
latent space $\bm{z}_{G}$ is $100$ and the dimension for each of the local
latent spaces is $50$. The standard deviation of the decoder is set to
$\gamma=0.03$.

Due to the SIGMA prior being a model approximation and assessing the
performance of an approximate algorithm in VI, we chose to use a synthetic
dataset for this example where the response is a draw from the PCAR prior over
the Australian geography. The synthetic dataset was generated by sampling from
the PCAR prior. The precision matrix $\mathrm{\Sigma}$ was constructed using
the adjacency matrix $\mathrm{W}$ derived from the SA2-level spatial structure
of Australia, ensuring that the synthetic data accurately reflects the complex
spatial dependencies present in the real-world geography. By assessing the
approximations of the SIGMA prior and VI against this synthetic dataset, we
can effectively evaluate their performance in capturing the true data-
generating process without the confounding effects of model misspecification
or data quality issues.

We employ (NUTS, Hoffman et al.,, 2014) and mean-field variational inference
(MFVI) algorithms, both implemented in BlackJAX (Cabezas et al.,, 2024), to
estimate the posterior distribution of the model parameters. For NUTS, we run
the sampler for $5,000$ iterations using the default hyperparameters and
summarize the results using the last $2,500$ draws of the chain, discarding
the first $2,500$ iterations as warmup. For MFVI, we run the algorithm for
$10,000$ iterations, utilizing the Adam optimizer (Kingma and Ba,, 2014) with
a learning rate of $1e^{-3}$. We clip the gradients if their norm exceeds
$1e^{-3}$ to ensure stability.

####  4.2.1 True underlying model

To assess the effectiveness of the SIGMA prior in approximating complex
spatial dependencies, we define a true underlying model that serves as the
target for our approximation. The true model for this example is

$\displaystyle\text{logit }\phi$ $\displaystyle\sim\mathcal{N}(0,2),$ (13)
$\displaystyle\bm{\theta}|\phi$
$\displaystyle\sim\mathcal{N}\big{(}\bm{0},(\mathrm{D}-\phi\mathrm{W})^{-1}\big{)},$
(14) $\displaystyle\bm{y}|\bm{\theta}$
$\displaystyle\sim\mathcal{N}(\bm{\theta},0.5).$ (15)

The spatial smoothing parameter $\phi$ is transformed using the logit function
to ensure it remains within the valid range of (0, 1). The prior distribution
for $\text{logit }\phi$ is a Gaussian distribution with mean 0 and variance 2,
allowing for a wide range of possible values for $\phi$.

The spatial random effects $\bm{\theta}$ are assigned a PCAR prior,
conditional on $\phi$. The precision matrix $(\mathrm{D}-\phi\mathrm{W})^{-1}$
captures the spatial dependence structure, where $\mathrm{D}$ and $\mathrm{W}$
are the diagonal matrix of neighbor counts and the adjacency matrix,
respectively, as described in the previous section. The observed data $\bm{y}$
are modeled using a normal likelihood function, conditional on the spatial
random effects $\bm{\theta}$. The standard deviation of the likelihood
function is set to $0.5$, representing the assumed measurement error or
uncertainty in the observations.

####  4.2.2 Deterministic SIGMA prior model

After training the SIGMA prior on prior draws from the true PCAR model across
the Australian geography, we obtain the following approximation, which we
refer to as the _deterministic SIGMA prior model_ ,

$\displaystyle\text{logit }\phi$ $\displaystyle\sim\mathcal{N}(0,2),$ (16)
$\displaystyle\bm{z}_{G}$ $\displaystyle\sim\mathcal{N}(0,1),$ (17)
$\displaystyle\bm{z}_{L_{j}}|\bm{z}_{G},\phi$
$\displaystyle\sim\mathcal{N}(\bm{\mu}_{\bm{\varphi}_{L_{j}}}(\bm{z}_{G},\phi),\bm{\sigma}_{\varphi_{L_{j}}}(\bm{z}_{G},\phi)),\quad
j={1,\ldots,J},$ (18) $\displaystyle\hat{\bm{\theta}}_{j}$
$\displaystyle=\bm{\mu}_{\bm{\psi}_{L_{j}},\bm{\psi}_{G}}(\bm{z}_{L_{j}},\bm{z}_{G},\phi),\quad
j={1,\ldots,J},$ (19) $\displaystyle\bm{y}_{j}|\hat{\bm{\theta}}_{j}$
$\displaystyle\sim\mathcal{N}(\hat{\bm{\theta}}_{j},0.5).$ (20)

From here on, we refer to this form of the approximate model as the “
_deterministic SIGMA prior model_ ”, as reconstructed estimates of
$\bm{\theta}$ are deterministic given the other unknown parameters. In this
approximation, the global latent variable $\bm{z}_{G}$ captures the shared
information across all regions, while the local latent variables
$\bm{z}_{L_{j}}$ capture the client-specific information. The deterministic
functions $\bm{\mu}_{\bm{\varphi}_{L_{j}}}$ and
$\bm{\sigma}_{\varphi_{L_{j}}}$ represent the encoder networks that map the
global latent variable and the spatial smoothing parameter to the mean and
standard deviation of the local latent variables, respectively. The decoder
network $\bm{\mu}_{\bm{\psi}_{L_{j}},\bm{\psi}_{G}}$ maps the local and global
latent variables, along with the spatial smoothing parameter, to the
reconstructed estimates of the spatial random effects $\hat{\bm{\theta}}_{j}$.

Figure 6: Plots compare the difference between the observed data and mean
estimates for $\bm{\theta}$, under each combination of model and algorithm.
The top row shows the true model using VI, while the bottom row shows the true
model using MCMC. The left-hand column shows the deterministic SIGMA prior
model, and the right-hand column shows the auxiliary variable SIGMA prior
model.

The plots on the left-hand side of Figure 6 show the difference between
$\hat{\bm{\theta}}$ and $\bm{y}$ for the deterministic SIGMA prior model
fitted using MFVI. Compared to the same estimate from the true model using
MFVI (top) and the true model using NUTS (bottom), we observe that the
estimates of $\bm{\theta}$ from the deterministic SIGMA prior model overfit to
the data. This overfitting behavior can be attributed to lower variance in
prior draws of $\bm{\theta}$ in the deterministic SIGMA prior model relative
to the true model. Figure 7 (left) further highlights this issue by showing
that the difference in mean estimates of $\bm{\theta}$ between the
deterministic SIGMA prior model and MCMC estimates from the true model has
much greater variance compared to the difference between the MCMC and VI
estimates under the true model.

####  4.2.3 Auxiliary variable SIGMA prior model

Here, we propose an extension to the inferential model one would fit when
using a model approximation with a SIGMA prior. The motivation of the
extension is to reduce the tendencies of the approximate model to overfit to
the observed data and to correct the narrow credible intervals. We note that
these overfitting properties were observed and left as outstanding problems in
both the PriorVAE and PriorCVAE papers, with overfitting appearing to a lesser
extent in the PriorCVAE paper.

Our proposed extension is motivated by first considering the same generative
scenario as PriorVAE, where we have random variables $\bm{\theta}$ that we
want to approximate with an alternative generative model. In the works of
PriorVAE, PriorCVAE, and this paper, at inference, we have treated
$\bm{\theta}$ as deterministic given the added latent variables.
Alternatively, we can consider treating $\bm{\theta}$ as a random variable,
with mean given by the decoder and variance $\bm{\tau}^{2}$. From this lens,
the added latent variable can be viewed as an auxiliary variable, and we would
end up conducting inference over both $\bm{\theta}$ and the latent variables.
If we choose the generative model

$\displaystyle\bm{z}$ $\displaystyle\sim\mathcal{N}_{P}(0,1),$
$\displaystyle\bm{\theta}|\bm{z}$
$\displaystyle\sim\mathcal{N}_{D}\big{(}\bm{\mu}_{\bm{\theta}}+\mathrm{W}_{\bm{\theta}}^{\top}\bm{z},\bm{\tau}^{2}\big{)},$

with unknown parameters $\bm{\mu}_{\bm{\theta}},\mathrm{W}_{\bm{\theta}}$, and
$\bm{\tau}$, we are fitting probabilistic PCA (PPCA, Tipping and Bishop,,
1999) to prior draws from $p(\bm{\theta})$, where the maximum likelihood
estimate for $\bm{\tau}^{2}$ is the average of the $D-P$ smallest eigenvalues
of the empirical covariance matrix of draws from $p(\bm{\theta})$. If viewing
from this perspective, then fitting a PPCA model to prior draws of
$\bm{\theta}$ and, at inference, jointly sampling from
$p(\bm{\theta},\bm{z})=p(\bm{\theta})p(\bm{\theta}|\bm{z})$ can be viewed as
an approximate auxiliary variable method. If the dimension of $\bm{z}$ is
chosen such that $P=D$, then $\bm{\tau}^{2}\rightarrow 0$, and we essentially
fit a deterministic location-scale transform to draws of $\bm{\theta}$.

As the decoder we use is not an affine transformation as in PPCA but a neural
network, we do not get analytical estimates for the variance $\bm{\tau}^{2}$.
However, we offer a post-hoc correction. Under the true model, we can sample
from the prior distributions to generate draws of $\bm{\theta}$, which we
denote as $\tilde{\bm{\theta}}_{\text{True}}$. We do the same for the
deterministic SIGMA prior model and call these draws
$\tilde{\bm{\theta}}_{\text{SIGMA}}$. We assume that $\bm{\theta}$ is Gaussian
distributed, centred on the decoder estimates, with unknown variance
$\bm{\tau}^{2}$. We attempt to match the amount of prior variance using the
SIGMA prior and the true model. Given our formulation, the variance is
additive, i.e.,
$\text{Var}(\tilde{\bm{\theta}}_{\text{True}})=\text{Var}(\tilde{\bm{\theta}}_{\text{SIGMA}})+\bm{\tau}^{2}$.
This expression enables us to estimate $\bm{\tau}^{2}$ after fitting the VAE
but before inference by taking prior draws from the two models and calculating
$\bm{\tau}^{2}=\text{Var}(\tilde{\bm{\theta}}_{\text{True}})-\text{Var}(\tilde{\bm{\theta}}_{\text{SIGMA}})$.
As the estimator for $\bm{\tau}^{2}$ might not be positive, we clip the
estimates to be between the range $0.01$ and $100$. For this example, $99.8\%$
of the elements of $\bm{\tau}^{2}$ did not require clipping.

We call this model the _auxiliary variable SIGMA prior model_ and, for this
example, define it as

$\displaystyle\text{logit }\phi$ $\displaystyle\sim\mathcal{N}(0,2),$
$\displaystyle\bm{z}_{G}$ $\displaystyle\sim\mathcal{N}(0,1),$
$\displaystyle\bm{z}_{L_{j}}|\bm{z}_{G},\phi$
$\displaystyle\sim\mathcal{N}(\bm{\mu}_{\bm{\varphi}_{L_{j}}}(\bm{z}_{G},\phi),\bm{\sigma}_{\varphi_{L_{j}}}(\bm{z}_{G},\phi)),\quad
j\in\\{1,\ldots,J\\},$ $\displaystyle\hat{\bm{\theta}}$
$\displaystyle=\bm{\mu}_{\psi_{L_{j}},\psi_{G}}(\bm{z}_{G},\bm{z}_{L_{j}},\phi),\quad
j\in\\{1,\ldots,J\\},$
$\displaystyle\bm{\theta}|\bm{z}_{G},\bm{z}_{L_{j}},\phi$
$\displaystyle\sim\mathcal{N}(\hat{\bm{\theta}},\bm{\tau}),\quad
j\in\\{1,\ldots,J\\},$ $\displaystyle\bm{y}_{j}|\bm{\theta}$
$\displaystyle\sim\mathcal{N}(\bm{\theta},0.5),\quad j\in\\{1,\ldots,J\\}.$

The addition of the conditional density of
$\bm{\theta}|\bm{z}_{G},\bm{z}_{L_{j}},\phi$ does not introduce any additional
dependencies between clients, as all the dependencies are already captured by
estimating $\hat{\bm{\theta}}$ through the learned SIGMA prior. Thus,
including this density is straightforward from an FL perspective. While
treating $\bm{\theta}$ as a random variable increases the dimensionality of
the model, the computational cost remains manageable when using a variational
inference algorithm with a mean-field assumption for the variational family of
$\bm{\theta}$.

Figure 6 (right-hand side) illustrates the difference between the mean
estimate of $\bm{\theta}$ and the observed data $\bm{y}$ for the auxiliary
SIGMA prior model fitted using mean-field variational inference (MFVI).
Comparing these estimates to those obtained from the true model using MFVI
(top) and NUTS (bottom) reveals that the auxiliary variable SIGMA prior model
achieves similar levels of smoothing with respect to the data. Notably, the
auxiliary variable approach significantly reduces the overfitting observed in
the estimates from the deterministic SIGMA prior model. Moreover, Figure 7
(right-hand side) demonstrates that the difference between the mean estimates
of $\bm{\theta}$ under the auxiliary variable SIGMA prior and the MCMC
estimates under the true model is substantially smaller than the difference
between the mean estimates of $\bm{\theta}$ under the deterministic SIGMA
prior and the MCMC estimates. This finding suggests that the auxiliary
variable SIGMA prior model approximates the true posterior distribution more
closely than the deterministic SIGMA prior model.

Figure 7: Comparison of the mean estimates of $\bm{\theta}$ obtained from
different models and inference methods. The vertical axis represents the
difference between the mean estimates and the MCMC estimates of $\bm{\theta}$
from the true model. Both plots include the difference between the mean
estimates of $\bm{\theta}$ using MFVI with the true model as a baseline. The
left plot displays the difference between the mean estimates of $\bm{\theta}$
using the deterministic SIGMA prior model and the MCMC estimates from the true
model. The right plot shows the difference between the mean estimates of
$\bm{\theta}$ using the auxiliary variable SIGMA prior model and the MCMC
estimates from the true model.

##  5 Conclusion

In this paper, we introduced the _S tructured Independence via deep Generative
Model Approximation_ (SIGMA) prior, a novel approach to approximating the
latent structure of probabilistic models while enforcing desired conditional
independence properties. The SIGMA prior utilizes a hierarchical variational
autoencoder (VAE) to capture dependencies in the original prior specification
and induce conditional independence among subsets of variables. The key
components of the hierarchical VAE include a global latent variable that
captures shared information across all clients, local latent variables that
capture client-specific information, and encoder and decoder networks that map
between the latent space and the original parameter space. By carefully
designing the architecture and training objective of the VAE, the SIGMA prior
enables the approximation of complex prior distributions while maintaining the
desired conditional independence structure.

The development of the SIGMA prior is motivated by the challenges of
horizontal FL, where data is distributed across multiple clients, and direct
access to data from other clients is prohibited. Existing FL algorithms, such
as SFVI (Hassan et al.,, 2023), assume independence between clients, limiting
their applicability to scenarios where dependencies are present. The SIGMA
prior addresses this limitation by approximating the prior distribution over
client-specific parameters and inducing conditional independence between
clients, enabling the use of existing FL algorithms while capturing
dependencies between clients. This advancement is significant, as it allows
for accurate inference and prediction in real-world applications where
dependencies between clients are expected. We evaluated the performance of the
SIGMA prior on a synthetic and real-world application, including a one-
dimensional GP prior and a CAR prior for modeling spatial dependencies in
cancer estimates across the Australian geography. Our experiments showed that
the SIGMA prior effectively captures the dependencies in the original prior
distributions and provides a close approximation to estimates from the actual
model in the non-FL setting. These results highlight the potential of the
SIGMA prior as a general-purpose tool for enabling FL with dependent data. The
successful application of the SIGMA prior to the CAR prior across the
Australian geography demonstrates the real-world impact of our approach.
Spatial dependencies are prevalent in many fields, such as epidemiology,
environmental sciences, and social sciences, where data is often collected at
different spatial resolutions or administrative units. The SIGMA prior
provides a principled way to model these dependencies in an FL setting,
enabling collaboration between organizations or data custodians without
compromising data privacy.

Despite the promising results, there are several limitations and avenues for
future research. One avenue for potential extensions is to explore alternative
generative models for the SIGMA prior, such as normalizing flows (Rezende and
Mohamed,, 2015), generative adversarial networks (Goodfellow et al.,, 2014),
or denoising diffusion models (Song and Ermon,, 2019; Ho et al.,, 2020; Kingma
et al.,, 2021). These models have shown promising results in capturing complex
distributions and could improve the expressiveness and flexibility of the
SIGMA prior. Additionally, the auxiliary variable SIGMA prior model introduced
in this work offers a promising direction for future research. The auxiliary
variable approach can reduce overfitting and improve the calibration of
uncertainty estimates by treating the local parameters as random variables and
introducing an additional conditional distribution. Further investigation into
the theoretical properties and practical implications of the auxiliary
variable SIGMA prior model could lead to more robust and reliable model
approximations.

In conclusion, the SIGMA prior is a powerful and flexible tool for enabling FL
with dependent data. By approximating the prior distribution and inducing
conditional independence, the SIGMA prior allows using existing FL algorithms
while capturing dependencies between clients. Our work opens up new
possibilities for collaborative learning and privacy-preserving inference in
various applications, from spatial modeling in healthcare and beyond. As FL
continues to evolve, ideas surrounding the SIGMA prior could play an essential
role in developing more expressive and accurate models that can leverage the
full potential of decentralized data.

## References

  * ruv, (2014) (2014).  Context-aware patch-based image inpainting using Markov random field modelling, author=Ruzic, Tijana and Pizurica, Aleksandra.  IEEE transactions on image processing, 24(1):444–456.
  * Abril-Pla et al., (2023) Abril-Pla, O., Andreani, V., Carroll, C., Dong, L., Fonnesbeck, C. J., Kochurov, M., Kumar, R., Lao, J., Luhmann, C. C., Martin, O. A., et al. (2023).  Pymc: a modern, and comprehensive probabilistic programming framework in python.  PeerJ Computer Science, 9:e1516.
  * Bernardo et al., (1998) Bernardo, J., Berger, J., Dawid, A., Smith, A., et al. (1998).  Regression and classification using gaussian process priors.  Bayesian statistics, 6:475.
  * Besag et al., (1991) Besag, J., York, J., and Mollié, A. (1991).  Bayesian image restoration, with two applications in spatial statistics.  Annals of the institute of statistical mathematics, 43:1–20.
  * Bingham et al., (2019) Bingham, E., Chen, J. P., Jankowiak, M., Obermeyer, F., Pradhan, N., Karaletsos, T., Singh, R., Szerlip, P., Horsfall, P., and Goodman, N. D. (2019).  Pyro: Deep universal probabilistic programming.  The Journal of Machine Learning Research, 20(1):973–978.
  * Blei et al., (2017) Blei, D. M., Kucukelbir, A., and McAuliffe, J. D. (2017).  Variational inference: A review for statisticians.  Journal of the American statistical Association, 112(518):859–877.
  * Cabezas et al., (2024) Cabezas, A., Corenflos, A., Lao, J., Louf, R., Carnec, A., Chaudhari, K., Cohn-Gordon, R., Coullon, J., Deng, W., Duffield, S., Durán-Martín, G., Elantkowski, M., Foreman-Mackey, D., Gregori, M., Iguaran, C., Kumar, R., Lysy, M., Murphy, K., Orduz, J. C., Patel, K., Wang, X., and Zinkov, R. (2024).  Blackjax: Composable bayesian inference in jax.
  * Duncan et al., (2019) Duncan, E. W., Cramb, S. M., Aitken, J. F., Mengersen, K. L., and Baade, P. D. (2019).  Development of the australian cancer atlas: spatial modelling, visualisation, and reporting of estimates.  International journal of health geographics, 18:1–12.
  * Goodfellow et al., (2014) Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. (2014).  Generative adversarial nets.  Advances in neural information processing systems, 27.
  * Hasanzadeh et al., (2020) Hasanzadeh, A., Hajiramezanali, E., Boluki, S., Zhou, M., Duffield, N., Narayanan, K., and Qian, X. (2020).  Bayesian graph neural networks with adaptive connection sampling.  In International conference on machine learning, pages 4094–4104. PMLR.
  * Hassan et al., (2023) Hassan, C., Salomone, R., and Mengersen, K. (2023).  Federated variational inference methods for structured latent variable models.  arXiv preprint arXiv:2302.03314.
  * Ho et al., (2020) Ho, J., Jain, A., and Abbeel, P. (2020).  Denoising diffusion probabilistic models.  Advances in neural information processing systems, 33:6840–6851.
  * Hoffman and Blei, (2015) Hoffman, M. and Blei, D. (2015).  Stochastic structured variational inference.  In Artificial Intelligence and Statistics, pages 361–369. PMLR.
  * Hoffman et al., (2013) Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J. (2013).  Stochastic variational inference.  Journal of Machine Learning Research.
  * Hoffman et al., (2014) Hoffman, M. D., Gelman, A., et al. (2014).  The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo.  J. Mach. Learn. Res., 15(1):1593–1623.
  * Kairouz et al., (2021) Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., Bonawitz, K., Charles, Z., Cormode, G., Cummings, R., et al. (2021).  Advances and open problems in federated learning.  Foundations and Trends® in Machine Learning, 14(1–2):1–210.
  * Kingma et al., (2021) Kingma, D., Salimans, T., Poole, B., and Ho, J. (2021).  Variational diffusion models.  Advances in neural information processing systems, 34:21696–21707.
  * Kingma and Ba, (2014) Kingma, D. P. and Ba, J. (2014).  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980.
  * Kingma and Welling, (2013) Kingma, D. P. and Welling, M. (2013).  Auto-encoding variational bayes.  arXiv preprint arXiv:1312.6114.
  * Kotelevskii et al., (2022) Kotelevskii, N., Vono, M., Durmus, A., and Moulines, E. (2022).  Fedpop: A bayesian approach for personalised federated learning.  Advances in Neural Information Processing Systems, 35:8687–8701.
  * Lee, (2011) Lee, D. (2011).  A comparison of conditional autoregressive models used in bayesian disease mapping.  Spatial and spatio-temporal epidemiology, 2(2):79–89.
  * McMahan et al., (2017) McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. (2017).  Communication-efficient learning of deep networks from decentralized data.  In Artificial intelligence and statistics, pages 1273–1282. PMLR.
  * Nair and Hinton, (2010) Nair, V. and Hinton, G. E. (2010).  Rectified linear units improve restricted boltzmann machines.  In Proceedings of the 27th international conference on machine learning (ICML-10), pages 807–814.
  * Paszke et al., (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. (2019).  Pytorch: An imperative style, high-performance deep learning library.  Advances in neural information processing systems, 32.
  * Phan et al., (2019) Phan, D., Pradhan, N., and Jankowiak, M. (2019).  Composable effects for flexible and accelerated probabilistic programming in NumPyro.  arXiv preprint arXiv:1912.11554.
  * Ranganath et al., (2016) Ranganath, R., Tran, D., and Blei, D. (2016).  Hierarchical variational models.  In International conference on machine learning, pages 324–333. PMLR.
  * Rezende and Mohamed, (2015) Rezende, D. and Mohamed, S. (2015).  Variational inference with normalizing flows.  In International conference on machine learning, pages 1530–1538. PMLR.
  * Rezende et al., (2014) Rezende, D. J., Mohamed, S., and Wierstra, D. (2014).  Stochastic backpropagation and approximate inference in deep generative models.  In International Conference on Machine Learning, pages 1278–1286. PMLR.
  * Riebler et al., (2016) Riebler, A., Sørbye, S. H., Simpson, D., and Rue, H. (2016).  An intuitive Bayesian spatial model for disease mapping that accounts for scaling.  Statistical methods in medical research, 25(4):1145–1165.
  * Roeder et al., (2017) Roeder, G., Wu, Y., and Duvenaud, D. K. (2017).  Sticking the landing: Simple, lower-variance gradient estimators for variational inference.  Advances in Neural Information Processing Systems, 30.
  * Ruder, (2017) Ruder, S. (2017).  An overview of gradient descent optimization algorithms.
  * Rue et al., (2009) Rue, H., Martino, S., and Chopin, N. (2009).  Approximate Bayesian inference for latent Gaussian models by using integrated nested Laplace approximations.  Journal of the Royal Statistical Society Series B: Statistical Methodology, 71(2):319–392.
  * Semenova et al., (2023) Semenova, E., Cairney-Leeming, M., and Flaxman, S. (2023).  PriorCVAE: scalable MCMC parameter inference with Bayesian deep generative modelling.  arXiv preprint arXiv:2304.04307.
  * Semenova et al., (2022) Semenova, E., Xu, Y., Howes, A., Rashid, T., Bhatt, S., Mishra, S., and Flaxman, S. (2022).  PriorVAE: encoding spatial priors with variational autoencoders for small-area estimation.  Journal of the Royal Society Interface, 19(191):20220094.
  * Song and Ermon, (2019) Song, Y. and Ermon, S. (2019).  Generative modeling by estimating gradients of the data distribution.  Advances in neural information processing systems, 32.
  * Tan et al., (2020) Tan, L. S., Bhaskaran, A., and Nott, D. J. (2020).  Conditionally structured variational Gaussian approximation with importance weights.  Statistics and Computing, 30:1255–1272.
  * Tipping and Bishop, (1999) Tipping, M. E. and Bishop, C. M. (1999).  Probabilistic principal component analysis.  Journal of the Royal Statistical Society Series B: Statistical Methodology, 61(3):611–622.
  * Titsias and Lázaro-Gredilla, (2014) Titsias, M. and Lázaro-Gredilla, M. (2014).  Doubly stochastic variational Bayes for non-conjugate inference.  In International conference on machine learning, pages 1971–1979. PMLR.
  * Tschannen et al., (2018) Tschannen, M., Bachem, O., and Lucic, M. (2018).  Recent advances in autoencoder-based representation learning.  arXiv preprint arXiv:1812.05069.
  * Williams and Rasmussen, (2006) Williams, C. K. and Rasmussen, C. E. (2006).  Gaussian processes for machine learning, volume 2.  MIT press Cambridge, MA.
  * Zhang et al., (2018) Zhang, C., Bütepage, J., Kjellström, H., and Mandt, S. (2018).  Advances in variational inference.  IEEE transactions on pattern analysis and machine intelligence, 41(8):2008–2026.

##  Appendix A Architecture details for the SIGMA prior

The global encoder generates the variational parameters
$\bm{\eta}_{G}:=(\bm{\mu}_{G}^{\top},\bm{\tau}_{G}^{\top})^{\top}$ for the
posterior approximation of $\bm{z}_{G}$ by taking input
$\bm{g}^{G}=(\bm{y}^{\top},\bm{\phi}^{\top})^{\top}$ and applying the
following functions:

$\displaystyle\bm{f}^{G}$
$\displaystyle=\mathrm{W}_{2}^{G}(\sigma(\mathrm{W}_{1}^{G}\bm{g}^{G}+\bm{b}_{1}^{G}))+\bm{b}_{2}^{G},$
$\displaystyle\bm{\mu}_{G}$
$\displaystyle=\mathrm{W}_{\bm{\mu}}^{G}\bm{f}^{G}+\bm{b}_{\bm{\mu}}^{G},$
$\displaystyle\log\bm{\tau}_{G}^{2}$
$\displaystyle=\mathrm{W}_{\bm{\tau}}^{G}\bm{f}^{G}+\bm{b}_{\bm{\tau}}^{G},$

where
$\mathrm{W}_{1}^{G}\in\mathbb{R}^{h_{G}\times(n+1)},\bm{b}_{1}^{G}\in\mathbb{R}^{h_{G}},\mathrm{W}_{2}^{G}\in\mathbb{R}^{h_{G}\times
h_{G}}$, $\bm{b}_{2}^{G}\in\mathbb{R}^{h_{G}}$,
$\mathrm{W}_{\bm{\mu}}^{G}\in\mathbb{R}^{n_{G}\times
h_{G}},\bm{b}_{\bm{\mu}}^{G}\in\mathbb{R}^{n_{G}},\mathrm{W}_{\bm{\tau}}^{G}\in\mathbb{R}^{n_{G}\times
h_{G}},\bm{b}_{\bm{\tau}}^{G}\in\mathbb{R}^{n_{G}}$, and $\sigma(\cdot)$
represents the ReLU activation function (Nair and Hinton,, 2010) applied
elementwise. For each of the $J$ clients, the local encoder generates the
variational parameters
$\bm{\eta}_{L_{j}}:=(\bm{\mu}_{L_{j}}^{\top},\bm{\tau}_{L_{j}}^{\top})^{\top}$
for the posterior approximation of $\bm{z}_{L_{j}}$ using:

$\displaystyle\bm{g}^{L_{j}}$
$\displaystyle=(\bm{z}_{G}^{\top},\bm{\phi}^{\top})^{\top},$
$\displaystyle\bm{f}^{L_{j}}$
$\displaystyle=\mathrm{W}_{2}^{L_{j}}(\sigma(\mathrm{W}_{1}^{L_{j}}\bm{g}^{L_{j}}+\bm{b}_{1}^{L_{j}}))+\bm{b}_{2}^{L_{j}},\quad
j\in{1,\ldots,J},$ $\displaystyle\bm{\mu}_{L_{j}}$
$\displaystyle=\mathrm{W}_{\bm{\mu}}^{L_{j}}\bm{f}^{L_{j}}+\bm{b}_{\bm{\mu}}^{L_{j}},\quad
j\in{1,\ldots,J},$ $\displaystyle\log\bm{\tau}_{L_{j}}^{2}$
$\displaystyle=\mathrm{W}_{\bm{\tau}}^{L_{j}}\bm{f}^{L_{j}}+\bm{b}_{\bm{\tau}}^{L_{j}},\quad
j\in{1,\ldots,J},$

where
$\mathrm{W}_{1}^{L_{j}}\in\mathbb{R}^{h_{L}\times(n_{G}+1)},\bm{b}_{1}^{L_{j}}\in\mathbb{R}^{h_{L}},\mathrm{W}_{2}^{L_{j}}\in\mathbb{R}^{h_{L}\times
h_{L}}$, $\bm{b}_{2}^{L_{j}}\in\mathbb{R}^{h_{L}}$,
$\mathrm{W}_{\bm{\mu}}^{L_{j}}\in\mathbb{R}^{n_{L_{j}}\times
h_{L}},\bm{b}_{\bm{\mu}}^{L_{j}}\in\mathbb{R}^{n_{L_{j}}},\mathrm{W}_{\bm{\tau}}^{L_{j}}\in\mathbb{R}^{n_{L_{j}}\times
h_{L}},\bm{b}_{\bm{\tau}}^{L_{j}}\in\mathbb{R}^{n_{L_{j}}}$. To generate a
reconstructed mean vector for the latent field, the decoder takes the form:

$\displaystyle\bm{g}^{D}$
$\displaystyle=(\bm{z}_{G}^{\top},\bm{\phi}^{\top})^{\top},$
$\displaystyle\bm{h}$
$\displaystyle=\mathrm{W}_{2}^{D}(\sigma(\mathrm{W}_{1}^{D}\bm{g}^{D}+\bm{b}_{1}^{D}))+\bm{b}_{2}^{D},$
$\displaystyle\bm{f}_{L_{j}}^{D}$
$\displaystyle=(\bm{h}^{\top},\bm{\theta}_{L_{j}}^{\top},\bm{\phi}^{\top})^{\top},\quad
j\in{1,\ldots,J},$ $\displaystyle\bm{\mu}_{\bm{\theta}_{L_{j}}}$
$\displaystyle=\mathrm{W}^{L_{j}}D\bm{f}{L_{j}}^{D}+\bm{b}^{L_{j}}D,\quad
j\in{1,\ldots,J},$ $\displaystyle\hat{\bm{\theta}}{L_{j}}$
$\displaystyle\sim\mathcal{N}(\bm{\mu}_{\bm{\theta}_{L_{j}}},\gamma\mathrm{I}),\quad
j\in{1,\ldots,J},$

where
$\mathrm{W}_{1}^{D}\in\mathbb{R}^{h_{D}\times(n_{G}+1)},\bm{b}_{1}^{D}\in\mathbb{R}^{h_{D}},\mathrm{W}_{2}^{D}\in\mathbb{R}^{h_{D}\times
h_{D}},\bm{b}_{2}^{D}\in\mathbb{R}^{h_{D}}$,
$\mathrm{W}^{L_{j}}D\in\mathbb{R}^{n{L_{j}}\times(h_{D}+n_{L_{j}}+1)},\bm{b}^{L_{j}}D\in\mathbb{R}^{n{L_{j}}}$,
and $\gamma>0$ is a fixed scale hyperparameter.

##  Appendix B Details for the proper CAR prior

The CAR prior captures the spatial dependence between neighboring regions by
conditioning the value of a spatial random effect at a particular location on
the values of its neighbors. In the proper CAR prior, the spatial random
effects $\bm{\phi}=(\phi_{1},\ldots,\phi_{n})^{\top}$ are assumed to follow a
multivariate normal distribution with a covariance matrix that depends on the
neighborhood structure of the regions. The joint distribution of the proper
CAR prior is given by:

$\bm{\phi}\sim\mathcal{N}(\bm{0},\sigma^{2}(\mathrm{D}-\alpha\mathrm{W})^{-1}),$

where $\mathrm{D}=\operatorname{diag}(m_{1},\ldots,m_{n})$ is a diagonal
matrix with the number of neighbors for each region, $\mathrm{W}$ is the
adjacency matrix of the neighborhood structure, $\alpha\in(0,1)$ is a spatial
dependence parameter, and $\sigma^{2}$ is a variance parameter. The
conditional distributions for each spatial random effect given its neighbors
are:

$\phi_{i}\mid\bm{\phi}_{-i}\sim\mathcal{N}\left(\alpha\frac{1}{m_{i}}\sum_{j\sim
i}\phi_{j},\frac{\sigma^{2}}{m_{i}}\right),$

where $\bm{\phi}_{-i}$ denotes all elements of $\bm{\phi}$ except $\phi_{i}$,
$m_{i}$ is the number of neighbors of region $i$, and $j\sim i$ denotes that
regions $j$ and $i$ are neighbors. The conditional mean of each random effect
is a weighted average of its neighbors, with the weight determined by the
spatial dependence parameter $\alpha$. The conditional variance is inversely
proportional to the number of neighbors. The proper CAR prior is proper
because $\alpha$ is strictly between 0 and 1. When $\alpha=1$, the proper CAR
prior reduces to the intrinsic CAR (ICAR) prior, which is improper because the
precision matrix $(\mathrm{D}-\mathrm{W})$ is singular.


# Using gravitational wave early warning to pre-point neutron star mergers
with Imaging Atmospheric Cherenkov telescopes

Jacopo Tissino  Gran Sasso Science Institute (GSSI), I-67100 L’Aquila, Italy
Laboratori Nazionali del Gran Sasso (INFN), I-67100 Assergi, Italy

###### Abstract

The LIGO-Virgo-KAGRA (LVK) collaboration has recently made it possible for
early warning alerts to be sent out, potentially before the end of the
gravitational wave (GW) emission from a neutron star binary. If we get such
alerts in this (the fourth) or the next observing run they may arrive up to
tens of seconds before the merger, which is comparable to the slewing times of
the Large Size Telescopes (designed to observe very high energy gamma rays):
it would be therefore possible to point to the source right before it starts
emitting an electromagnetic signal. This new mode of observation would allow
us to detect the TeV component of prompt emission, which is currently poorly
constrained and understood. There are many technical challenges to overcome
before this can be realized: improving the synergy between gravitational
observatories and telescopes, reducing operational latencies and, from the
gravitational wave side, providing more information, such as real-time updates
on early warning candidates and the probability distribution of the
inclination angle. Although we may need to wait a few years—in the worst case
scenario, until the next generation of GW detectors is built—before the first
detection of this kind is made, implementing these improvements is a
necessity.

Neutron star binaries are a prime candidate for multimessenger astronomy: they
emit gravitational waves (GWs) during their inspiral, followed by multi-
wavelength electromagnetic emission. The first coincident detection of a
neutron star binary in gravitational waves with a short gamma ray burst, GRB
170817A, [1, 2] on the 17th of August 2017, was made by Fermi-GBM and
INTEGRAL. Information about the localization of the source from the
gravitational waves would only become available hours later, while the delay
between merger and burst was less than 2 seconds: the discovery was made
serendipitously, these wide-field instruments fortunately happened to be
oriented in such a way that the emission was detectable.

While we hope to get more of this kind of detections, another higher-effort
avenue has the potential to yield a better result: a pointed follow-up for
energies from the keV to the TeV, as is done in the optical and other bands.
Achieving this would make it possible to use the sensitivity of instruments
like Imaging Atmospheric Cherenkov Telescopes (IACTs) or the localization
capabilities of X-ray facilities like Swift-XRT. Unlike lower-energy bands,
however, the detectable part of high-energy emission is often short-lived and
quickly decaying: detecting it requires fast communication with the
observatories. Let us then discuss the current low-latency infrastructure.

The fourth observing run of the LVK collaboration, dubbed O4, started in
spring 2023 and will continue until the end of 2025. Since the start of O4,
the possibility of sending Early Warning (EW) alerts has been implemented.[3]
The in-band duration of the gravitational chirp from a neutron star binary is
typically quite long, on the scale of one to a few minutes, meaning that it
can be analyzed and localized before it ends.[4] This task is performed
employing truncated template banks, which end at various times, ranging from
around 60 to around 10 seconds before merger.[4]

As of now (the end of May 2024), only five early warning alerts have been
issued, but all have been retracted,111Information about this is public and
can be obtained by querying the Gravitational wave Candidate Event Database
(GraceDB: gracedb.ligo.org/search/) with the keyword EARLY_WARNING. as none of
them were followed by a full-bandwidth alert; for comparison, the total number
of retracted alerts was 14. The decision to retract a lone early warning alert
is much easier than the typical considerations regarding retractions: if it
was not followed by a full-bandwidth alert, it is surely false, with no need
to consider other factors.

The goal of early warning alerts is to communicate information about an event
fast enough that actions can be taken soon, ideally before the merger. The
technical challenge this poses is great: after the gravitational waves from a
binary merger reach the interferometers, several things need to happen. (i)
The strain data needs to be calibrated and distributed between computing
centers, (ii) analyzed by search pipelines which will give an alert, (iii)
which is then handled by the event orchestrator gwcelery, (iv) input into the
database GraceDB, and (v) communicated to the larger community through a
General Circular Network (GCN) or an event broker like SCiMMA. After this, a
telescope needs to (vi) ingest the information and (vii) slew in order to
detect the event.

Prior to the start of O4, a study measured the typical latencies in steps (ii)
through (iv) for a set of injections, produced using realistic noise from the
previous observing run (O3): the median for these was around 30 seconds after
merger for all alerts, and around 3 seconds _before_ merger for early warning
alerts; the tail of the distribution, on the other hand, reached down to
around a minute before merger.[3] Step (i) was not evaluated, and it typically
takes around 10 seconds, which already shifts the median as well as the bulk
of the latency distribution for early warning alerts after the merger.

When the alert is received by a telescope optimized for it, the response can
be remarkably fast: for example, the Gamma Ray Burst (GRB), GRB 160821B was
detected by the wide-field Burst Alert Telescope (BAT) on board the Swift
satellite, and the communication about it was received by the Major
Atmospheric Cherenkov Telescopes (MAGIC) within 13 seconds (analogous to items
(v) and (vi) in the previous nomenclature), with the observation starting only
11 seconds later (vii), for a total of 24 seconds of latency.[6]

The general picture is that, for the current observing run, having the
response of a telescope start before the merger is not impossible, but it
would require an extremely loud event occurring under almost perfect
circumstances. Hence, the first study[7] proposing a very-high-energy (VHE,
E$>60\,$GeV) $\gamma$-ray observation performed in response to gravitational
early warning was in the context of a future instrument: the Einstein
Telescope, planned for the 2030s with an extremely good sensitivity in the
low-frequency band, which will make it especially suited for early warning.

This study highlighted the great scientific return that would be achieved with
such an observation: the early and variable part of the emission from short
GRBs, called _prompt_ , is believed to originate through internal dissipation
within the ultrarelativistic jet, which then decelerates in the surrounding
medium, giving rise to the much better understood _afterglow_ phase. What is
the mechanism of this deceleration? What is the composition of the jet? These
questions may be answered at least in part if we had information about the VHE
component of the emission, which we expect to be generated through synchrotron
self-Compton (SSC).[7, 11]

Figure 1:  Evolution of the skymap area with time. Data from the injections
performed by Sachdev et al.[4]

Let us discuss another issue: the early localization of gravitational
transients, which is typically performed during step (iii) using BAYESTAR.[5]
Sky localization areas for gravitational events tend to be broad, which is
especially true for early warning alerts: figure 1 shows the evolution of the
90% localization area with time for various events with early warning,
assuming design sensitivity for the LIGO and Virgo detectors. Very few events
are expected to be localized within less than 100 square degrees earlier than
10 seconds before the merger.

With this in mind, optimizing the observation strategy is crucial. Banerjee et
al.[7] identified two main modes of observation for the VHE counterpart of a
GW signal, which is expected to be extremely bright, potentially requiring
mere tens of seconds of exposure for a significant detection in the early
stages: (a) one-shot observation, for events with a localization area small
enough to be significantly covered by the field of view, and (b) mosaic
tiling, which is currently employed to search for late counterparts (afterglow
emission).

Only strategy (a) can allow us to detect prompt emission in the VHE band. With
the sensitivity of next generation detectors we may achieve small localization
areas well in advance of the merger, but as discussed earlier, this is not
feasible in the next observing run (O5). However, we propose that early alerts
(potentially with some improvements compared to their current LVK
implementation) can be useful even if they do not precisely localize the event
for the one-shot strategy to be applicable.

For one, even a thousand-degree localization area can define a rough direction
in the sky, allowing a telescope to start slewing. If several early updates
were sent for the same event, the telescope could progressively refine its
pointing as it receives more information, lowering the time required to have
the true location in the field of view once the area becomes small enough.
Further, if a sequence of alerts were issued as updates to the first early
one, some steps of the response may be fast-tracked, further reducing latency.
Even if this were to increase the false positive rate, it would constitute a
minuscule commitment in terms of telescope time: for any early alert, we can
know for sure whether it was a false alarm within a minute at most, and from
O4 we know to expect a rate of a few per year for these. In the most
pessimistic scenario, this would account for a half-hour per year of telescope
time. However, currently the low-latency GW infrastructure will send at most
one early alert per event (technically, per “superevent”, an aggregate of
several alerts from different pipelines arrived within a short interval and
believed to have the same source): we believe that this may become suboptimal
in the next observing run(s).

How do we determine which IACTs are the most suited for this kind of
observation? Surely a fast slew speed and a large field of view are
beneficial. Furthermore, their energy band needs to reach low enough: most
IACTs peak in sensitivity around 1 TeV, while from the GRBs observed so far in
this band we expect the second hump to peak in the 10-300 GeV range;[10, 11]
although, it should be noted that these were long (i.e. likely originating
from core collapsars), and no short GRB was detected in the VHE band yet.

In terms of fast slew speed and energy band, the MAGIC telescope and the CTA-
LST are the best candidates, slewing as fast as 6 degrees per second, and
reaching less than 100 GeV in energy.[8] The ASTRI Mini Array,[9] on the other
hand, slews slower and has a higher energy threshold, but its field of view is
almost three times as large as those of the aforementioned telescopes,
allowing for a higher probability of catching the source with a single
pointing.

We are currently exploring the effect these strategies would have in terms of
increasing the number of expected detections, their practical feasibility, and
the changes they will require from both the gravitational wave side and the
telescope one.

## References

## References

  * [1] Abbott et al., Phys. Rev. Lett. 119, 16 (2017)
  * [2] Ajello et al., The Astrophysical Journal 861, 2 (2018)
  * [3] S. S. Chaudary, A. Toivonen et al., Proc. Natl. Acad. Sci. U.S.A. 121, 18 (2024).
  * [4] S. Sachdev et al., The Astrophysical Journal Letters 905, 2 (2020).
  * [5] L. Singer, L. Price, Phys. Rev. D 93, 2 (2016).
  * [6] V. A. Acciari et al., The Astrophysical Journal 908, 90 (2021)
  * [7] B. Banerjee et al., Astronomy and Astrophysics 678, A126 (2023)
  * [8] H. Abe et al., Astronomy and Astrophysics 680, A66 (2023)
  * [9] S. Vercellone et al., Journal of High Energy Astrophysics 35, 1-42 (2022)
  * [10] M. Tavani et al., The Astrophysical Journal Letters 956, 1 (2023)
  * [11] B. Banerjee et al., preprint, https://arxiv.org/abs/2405.15855 (2024)