Learning Latent Graph Structures and their Uncertainty

## Abstract

Abstract Within a prediction task, Graph Neural Networks (GNNs) use relational information as an inductive bias to enhance the model’s accuracy. As task-relevant relations might be unknown, graph structure learning approaches have been proposed to learn them while solving the downstream prediction task. In this paper, we demonstrate that minimization of a point-prediction loss function, e.g., the mean absolute error, does not guarantee proper learning of the latent relational information and its associated uncertainty. Conversely, we prove that a suitable loss function on the stochastic model outputs simultaneously grants (i) the unknown adjacency matrix latent distribution and (ii) optimal performance on the prediction task. Finally, we propose a sampling-based method that solves this joint learning task. Empirical results validate our theoretical claims and demonstrate the effectiveness of the proposed approach.

## 1 Introduction

Relational information processing has provided breakthroughs in the analysis of rich and complex data coming from, e.g., social networks, natural language, and biology. This side information takes various forms, from structuring the data into clusters, to defining causal relations and hierarchies, and enables machine learning models to condition their predictions on dependency-related observations.
In this context, predictive models take the form $y=f_{\psi}(x,A)$, where the input-output relation $x\mapsto y$ – modeled by $f_{\psi}$ and its parameters in $\psi$ – is conditioned on relational information encoded in variable $A$.
Graph Neural Networks (GNNs) Scarselli et al., [] are one example of models of this kind that rely on a graph structure represented as an adjacency matrix $A$ and have been demonstrated successful in a plethora of applications, e.g., [Fout et al.,, ; Shlomi et al.,, ].
Throughout this paper, we focus on predictors where $A$ is an adjacency matrix, although the theoretical results we develop are valid for $A$ being any discrete latent random variable.

Indeed, relational information is needed to implement such a relational inductive bias and, in some cases, it is provided at the application design phase. However, more frequently, such topological information is not rich enough to address the problem at hand and – not seldom – completely unavailable.
Therefore, Graph Structure Learning (GSL) emerges as an approach to learn the graph topology Kipf et al., []; Franceschi et al., []; Yu et al., []; Fatemi et al., []; Zhu et al., []; Cini et al., [] alongside the predictive model $f_{\psi}$. This entails formulating a joint learning process that learns the adjacency matrix $A$ – or a parameterization of it – altogether with the predictor’s parameters $\psi$. This can be achieved by optimizing a loss function, e.g., a point prediction measure based on the square or the absolute prediction error.

Different sources of uncertainty affect the graph structure learning process, including epistemic uncertainty in the data and variability inherent in the data-generating process.
Examples include e.g., social interactions where links can intermittently be present, traffic flows affected by road closures and temporary detours, and adaptive communication routing.
It follows that a probabilistic framework is appropriate to accurately capture the uncertainty in the learned relations whenever randomness affects the graph topology.
Probabilistic approaches have been devised in recent years. For instance, research carried out in [Franceschi et al.,, ; Zhang et al.,, ; Elinas et al.,, ; Cini et al.,, ] propose methods that learn a parametric distribution $P^{\theta}_{A}$ over the latent graph structure $A$. However, none of them have studied whether these approaches were able to learn a *calibrated* latent distribution $P^{\theta}_{A}$, properly reflecting the uncertainty associated with the learned topology.

In this paper, we address the joint problem of learning a predictive model yielding optimal point-prediction performance of the output $y$ and, contextually, a calibrated distribution for the latent adjacency matrix $A$. In particular, the novel contributions can be summarized as:

- 1.
We demonstrate that models trained to achieve optimal point predictions do *not* guarantee calibration of the adjacency matrix distribution [Section 4].
- 2.
We provide theoretical conditions on the predictive model and loss function that guarantee both distribution calibration and optimal point-predictions [Section 5].
- 3.
We propose a theoretically-grounded sampling-based learning method to address the joint learning problem
[Section 5].
- 4.
We empirically validate major paper’s theoretical developments and claims and show that the proposed method is indeed able to solve the joint learning task [Section 6].

## 2 Related work

GSL is often employed end-to-end with a predictive model to better solve a downstream task.
Examples include applications within graph deep learning methods for static Jiang et al., []; Yu et al., []; Kazi et al., [] and temporal data Wu et al., [, ]; Cini et al., []; De Felice et al., []; a recent review is provided by Zhu et al., [].

Some approaches from the literature model the latent graph structure as stochastic Kipf et al., []; Franceschi et al., []; Elinas et al., []; Shang et al., []; Cini et al., [], mainly as a way to enforce sparsity of the adjacency matrix.
To operate on discrete latent random variables, Franceschi et al., [] utilize straight-through gradient estimations, Cini et al., [] rely on score-based gradient estimators, while Niepert et al., [] design an implicit maximum likelihood estimation strategy.

To the best of our knowledge, in the context of GSL, no prior work has studied the joint learning problem of calibrating the latent graph distribution while achieving optimal point prediction.
Moreover,
the generality of what is here developed enables its adoption with different probabilistic model architectures too.

## 3 Problem formulation

Consider a set of $N$ interacting entities and the data-generating process

$$ $\begin{cases}y=f^{*}(x,A)\\ A\sim P^{*}_{A}\end{cases}$ (1) $$

where $y\in\mathcal{Y}$ is the system output obtained from input $x\in\mathcal{X}$ through function $f^{*}$ and conditioned on a realization of the latent adjacency matrix $A\in\mathcal{A}\subseteq\{0,1\}^{N\times N}$ drawn from distribution $P^{*}_{A}$; superscript $*$ refers to unknown entities.
Each entry of the adjacency matrix $A$ is a binary value encoding the existence of a pairwise relation between two nodes.
In the sequel, $x$ and $y$ are stacks of $N$ node-level feature vectors of dimension $d_{in}$ and $d_{out}$, respectively, i.e., $\mathcal{X}\subseteq\mathbb{R}^{N\times d_{in}}$ and $\mathcal{Y}\subseteq\mathbb{R}^{N\times d_{out}}$.

Given a training dataset $\mathcal{D}=\{(x_{i},y_{i})\}_{i=1}^{n}$ of $n$ input-output observations from (1), we aim at learning a probabilistic predictive model

$$ $\begin{cases}\hat{y}=f_{\psi}(x,A)\\ A\sim P_{A}^{\theta}\end{cases}$ (2) $$

from $\mathcal{D}$, while learning at the same time distribution $P_{A}^{\theta}$ approximating $P_{A}^{*}$.
The two parameter vectors $\theta$ and $\psi$ are trained to approximate distinct entities in (1), namely the distribution $P^{*}_{A}$ and function $f^{*}$, respectively. We assume

###### Assumption 3.1 .

The family $\{P_{A}^{\theta}\}$ of probability distributions $P_{A}^{\theta}$ parametrized by $\theta$ and the family of predictive functions $\{f_{\psi}\}$ are expressive enough to contain the true latent distribution $P^{*}_{A}$ and function $f^{*}$, respectively.

Assumption 3.1 implies that $f^{*}\in\{f_{\psi}\}$ and $P_{A}^{*}\in\{P_{A}^{\theta}\}$ but does not request uniqueness of the parameters vectors $\psi^{*}$ and $\theta^{*}$ such that $f_{\psi^{*}}=f^{*}$ and $P_{A}^{\theta^{*}}=P_{A}^{*}$. Under such assumption the minimum function approximation error
is null and we can focus on the theoretical conditions requested to guarantee a successful learning, i.e., achieving both optimal point predictions and latent distribution calibration.

###### Assumption 3.1 .

#### Optimal point predictions

Outputs $y$ and $\hat{y}$ of probabilistic model (1) and (2) are random variables following push-forward distributions(^1^11The distribution of $y=f^{*}(x,A)$ originated from $P_{A}^{*}$ and of $\hat{y}=f_{\psi}(x,A)$ originated from $P^{\theta}_{A}$.)
$P^{*}_{y|x}$ and $P_{y|x}^{\theta,\psi}$, respectively. A single point prediction $y_{PP}\in\mathcal{Y}$ can be obtained through an appropriate functional $T[\cdot]$ as

$$ $y_{PP}=y_{PP}(x,\theta,\psi)\equiv T\left[P_{y|x}^{\theta,\psi}\right].$ (3) $$

For example, $T$ can be the expected value or the value at a specific quantile. We then define an optimal predictor as one whose parameters $\theta$ and $\psi$ minimize the expected point-prediction loss

$$ $\mathcal{L}^{point}(\theta,\psi)=\mathbb{E}_{x\sim P_{x}^{*}}\left[\mathbb{E}_ {y\sim P^{*}_{y|x}}\left[\ell\big{(}y,y_{PP}(x,\theta,\psi)\big{)}\right]\right]$ (4) $$

between the system output $y$ and the point-prediction $y_{PP}$, as measured by of a loss function $\ell:\mathcal{Y}\times\mathcal{Y}\to\mathbb{R}_{+}$.

Statistical functional $T$ is coupled with the loss $\ell$ as the optimal functional $T$ to employ given a specific loss $\ell$ is often known Berger, []; Gneiting, [], when $P_{y|x}^{\theta,\psi}$ approximates well $P_{y|x}^{*}$.
For instance, if $\ell$ is the Mean Absolute Error (MAE) the associated functional $T$ is the median, if $\ell$ is the Mean Squared Error (MSE) the associated functional is the expected value.

#### Latent distribution calibration

Calibration of a parametrized distribution requires learning parameters $\theta$, so that distribution $P^{\theta}_{A}$ aligns with $P^{*}_{A}$. Quantitatively, a dissimilarity measure
$\Delta^{cal}:\mathcal{P}_{A}\times\mathcal{P}_{A}\to\mathbb{R}_{+},$
defined over a set $\mathcal{P}_{A}$ of distributions on $\mathcal{A}$, assesses how close two distributions are. The family of $f$-divergences Rényi, [], such as the Kullback-Leibler divergence, and the integral probability metrics Müller, [], such as the maximum mean discrepancy Gretton et al., [] are examples of such dissimilarity measures. In this paper, we are interested in those discrepancies for which
$\Delta^{cal}(P_{1},P_{2})=0\iff P_{1}=P_{2}$ holds.
It follows that the latent distribution $P^{\theta}_{A}$ is calibrated on $P^{*}_{A}$ if it minimizes the latent distribution loss

$$ $\mathcal{L}^{cal}=\mathbb{E}_{x\sim P_{x}^{*}}\left[\Delta^{cal}\left(P^{*}_{A },P^{\theta}_{A}\right)\right],$ (5) $$

or simply $\mathcal{L}^{cal}=\Delta^{cal}\left(P^{*}_{A},P^{\theta}_{A}\right)$, when $A$ and $x$ are independent.

The problem of designing a predictive model (2) that both yields optimal point predictions (i.e., minimizes $\mathcal{L}^{point}$ in (4)) and calibrates the latent distribution (i.e., minimizes $\mathcal{L}^{cal}$ in (5)) is non-trivial for two main reasons. At first, as the latent distribution $P_{A}^{*}$ is unknown (and no samples from it are available), we cannot directly estimate $\mathcal{L}^{cal}$.
Second, as shown in Section 4, multiple sets of $\theta$ parameters may minimize $\mathcal{L}^{point}$ without minimizing $\mathcal{L}^{cal}$.

## 4 Limitations of point-prediction optimization

In this section, we demonstrate that the optimization of a point prediction loss Equation (4) does not generally grant calibration of the latent random variable.

###### Proposition 4.1 .

Consider Assumption 3.1.
Loss function $\mathcal{L}^{point}(\theta,\psi)$ in (4) is minimized by all $\theta$ and $\psi$ s.t. $T\Big{[}P_{y|x}^{\theta,\psi}\Big{]}=T\Big{[}P^{*}_{y|x}\Big{]}$ almost surely on $x$ and, in particular,

$$ $\mathcal{L}^{point}(\theta,\psi)\text{ is minimal}\quad\begin{array}[]{cc} \Longleftarrow\\ \;\;\not\!\!\Longrightarrow\end{array}\quad P_{y|x}^{\theta,\psi}=P^{*}_{y|x}.$ $$

Figure: Figure 1: A data generating model, as in (1), is used to produce a dataset with latent distribution parameter $\theta^{*}$. Outputs are generated for different values $\theta$ as in (2). In red, losses are computed as in (4) with $\ell$ being the MAE. In blue, losses are computed with our approach described further on.
Refer to caption: extracted/5631973/Figures/MAE_and_MMD_vs_theta.png

The proof of the proposition is given in Appendix A.1; we provide a counterexample for which calibration is not granted even when the processing function $f_{\psi}$ is equal to $f^{*}$ in Appendix A.2.

Figure 1 empirically demonstrates that optimizing point prediction losses does not necessarily guarantee distribution calibration.
In particular, we compute different losses between data generated with a ground truth system model (model (1) with optimal parameter $\theta^{*}$) and outputs produced with a different model (model (2), with varying $\theta$ parameters).
In red, the MAE is used as the loss function $\ell$ in the point prediction loss $\mathcal{L}^{point}$ of (4). Since all $\theta\geq 0.725$ produce statistically equivalent losses, this simple experiment demonstrates the inefficacy of minimizing $\mathcal{L}^{point}$ for latent distribution calibration. In blue, we show the loss we propose in the next section, which clearly presents a minimum in $\theta^{*}$. The details of this experiment can be found in Section 6.1.
However, we recommend reading the entire paper first to better understand the experiment’s context and setting.

Given the provided negative result and the impossibility of assessing loss $\mathcal{L}^{cal}$ in (5), in the next section, we propose another optimization objective that, as we will prove, allows us to both calibrate the latent random variable and to have optimal point predictions.

###### Proposition 4.1 .

## 5 Predictive distribution optimization: two birds with one stone

In this section, we show that we can achieve an optimal point predictor (2) and a calibrated latent distribution $P_{A}^{\theta}$ by comparing
push-forward distributions $P_{y|x}^{*}$ and $P_{y|x}^{\theta,\psi}$ of the outputs $y$ conditioned on input $x$. In particular, Theorem 5.2 below proves that, under appropriate conditions, minimization of the *output distribution loss*

$$ $\mathcal{L}^{dist}(\theta,\psi)=\mathbb{E}_{x\sim P_{x}^{*}}\Big{[}\Delta(P_{y |x}^{*},P_{y|x}^{\theta,\psi})\Big{]}$ (6) $$

provides calibrated $P_{A}^{\theta}$, even when $P_{A}^{*}$ is not available; $\Delta:\mathcal{P}_{y}\times\mathcal{P}_{y}\to\mathbb{R}_{+}$ is a dissimilarity measure between distributions over space $\mathcal{Y}$. We assume the following on dissimilarity measure $\Delta$.

###### Assumption 5.1 .

$\Delta(P_{1},P_{2})$ $\geq 0$ for all distributions $P_{1}$ and $P_{2}$ in $\mathcal{P}_{y}$ and $\Delta(P_{1},P_{2})=0$ if and only if $P_{1}=P_{2}$.

Several choices of $\Delta$ meet Assumption 5.1, e.g., $f$-divergences and some integral probability metrics Müller, []; the dissimilarity measure $\Delta$ employed in this paper is discussed in Section 5.1.

###### Theorem 5.2 .

Let $I=\{x:A\mapsto f^{*}(x,A)\text{ is injective}\}\subseteq\mathcal{X}$ be the set of points $x\in\mathcal{X}$ such that map $A\mapsto f^{*}(x,A)$ is injective.
Under Assumptions 3.1 and 5.1,
if $\mathbb{P}_{x\sim P_{x}^{*}}(I)>0$, then

$$ $\displaystyle\mathcal{L}^{dist}(\theta,\psi^{*})=0$ $\displaystyle\implies\begin{cases}\mathcal{L}^{point}(\theta,\psi^{*})\text{ is minimal}\\ \mathcal{L}^{cal}(\theta)=0,\end{cases}$ $$

where $\psi^{*}$ is such that $f_{\psi^{*}}=f^{*}$.

Theorem 5.2 is proven in Appendix A.3.
Under the theorem’s hypotheses, a predictor that minimizes $\mathcal{L}^{dist}$ is both calibrated on the latent random distribution and provides optimal point predictions. This overcomes limits of Proposition 4.1 where optimization of $\mathcal{L}^{point}(\theta,\psi^{*})$ does not grant $\mathcal{L}^{cal}(\theta)=0$.

The hypotheses under which Theorem 5.2 holds are rather mild. In fact, condition $\mathbb{P}_{x\sim P_{x}^{*}}(I)>0$ pertains to the data-generating process and intuitively ensures that,
for some $x$, different latent random variables produce different outputs.
A sufficient condition for $\mathbb{P}_{x\sim P_{x}^{*}}(I)>0$ to hold is the existence of a point $\bar{x}$ in the support of $P_{x}^{*}$ such that $A\mapsto f^{*}(\bar{x},A)$ is injective with $f^{*}$ continuous w.r.t. $\bar{x}$; see Corollary A.1 in Appendix A.3.
Instead, condition $f_{\psi}=f^{*}$ is set to avoid scenarios of different, yet equivalent,(^2^22E.g., $f_{\psi}(A,x)=f_{*}(\mathbf{1}-A,x)$ and $P^{\theta}_{A}$ encoding the absence of edges instead of their presence as in $P^{*}_{A}$.)
representations of the latent distribution.

Assumptions 3.1 and 5.1 can be met with an appropriate choice of model (2) and measure $\Delta$; as such they are controllable by the designer.
Assumption 3.1 ensures that the system model can be approximated with high accuracy.
For instance, by selecting $P_{A}^{\theta}$ as $N\times N$ independent Bernoulli variables allows us to modeling any graph distribution with independent edges.
Assumption 5.1 prevents from obtaining mismatched output distributions when $\mathcal{L}^{dist}(\theta,\psi)=0$.

As mentioned above, popular measures, e.g., the Kullback-Leibler divergence, meet the theorem’s assumptions and therefore can be adopted as $\Delta$. However, here we propose considering the Maximum Mean Discrepancy (MMD) Gretton et al., [] as a versatile alternative that allows Monte Carlo computation without requiring evaluations of the likelihood w.r.t. the output distributions $P^{*}_{y|x}$ and $P^{\theta,\psi}_{y|x}$.

###### Assumption 5.1 .

###### Theorem 5.2 .

### 5.1 Maximum mean discrepancy

Given two distributions $P_{1},P_{2}\in\mathcal{P}_{y}$, MMD can be defined as

$$ $\text{MMD}_{\mathcal{G}}[P_{1},P_{2}]=\sup_{g\in\mathcal{G}}\left\{\mathbb{E}_ {y\sim P_{1}}\big{[}g(y)\big{]}-\mathbb{E}_{y\sim P_{2}}\big{[}g(y)\big{]} \right\},$ (7) $$

i.e., the supremum, taken over a set $\mathcal{G}$ of functions $\mathcal{Y}\to\mathbb{R}$, of the difference
between expected values w.r.t. $P_{1}$ and $P_{2}$.
An equivalent form is derived for a generic kernel function $\kappa(\cdot,\cdot):\mathcal{Y}\times\mathcal{Y}\to\mathbb{R}$:

$$ $\text{MMD}^{2}_{\mathcal{G}_{\kappa}}[P_{1},P_{2}]=\underset{y_{1},y^{\prime}_ {1}\sim P_{1}}{\mathbb{E}}\Big{[}\kappa(y_{1},y^{\prime}_{1})\Big{]}-2 \underset{\begin{subarray}{c}y_{1}\sim P_{1}\\ y_{2}\sim P_{2}\end{subarray}}{\mathbb{E}}\Big{[}\kappa(y_{1},y_{2})\Big{]}+ \underset{y_{2},y^{\prime}_{2}\sim P_{2}}{\mathbb{E}}\Big{[}\kappa(y_{2},y^{ \prime}_{2})\Big{]}$ (8) $$

and it is associated with the unit-ball $\mathcal{G}_{k}$ of functions in the reproducing kernel Hilbert space of $\kappa$; note that (8) is the square of (7). Moreover, when universal kernels are considered (e.g., the Gaussian one), then (8) fulfills Assumption 5.1 (see Theorem 5 of Gretton et al., []).
Dissimilarity in (8) can be conveniently estimated via Monte Carlo (MC) and employed within a deep learning framework. Accordingly, we set $\Delta=\text{MMD}^{2}_{\mathcal{G}_{\kappa}}$ and learn parameter vectors $\psi$ and $\theta$ by minimizing $\mathcal{L}^{dist}(\theta,\psi)$ via gradient-descent methods.

### 5.2 Finite-sample computation of the loss

To compute the gradient of $\mathcal{L}^{dist}(\theta,\psi)=\mathbb{E}_{x\sim P_{x}^{*}}\left[\text{MMD}^{
2}_{\mathcal{G_{\kappa}}}\left[P_{y|x}^{\theta,\psi},P_{y|x}^{*}\right]\right]$ w.r.t. parameter vectors $\psi$ and $\theta$,
we rely on MC sampling to estimate in (6)
expectations over input $x\sim P_{x}^{*}$, target output $y\sim P_{y|x}^{*}$ and model output $\hat{y}\sim P_{y|x}^{\theta,\psi}$ . This amounts to substituhte $\text{MMD}^{2}_{\mathcal{G_{\kappa}}}$ with

$$ $\widehat{\text{MMD}}^{2}(\theta,\psi;x,y)=\frac{\sum_{i=1}^{N_{adj}}\sum_{j=1} ^{i-1}\kappa(\hat{y}_{i},\hat{y}_{j})}{N_{adj}(N_{adj}-1)}-2\frac{\sum_{i=1}^{ N_{adj}}\kappa(y,\hat{y}_{i})}{N_{adj}}$ (9) $$

In (9), $N_{adj}>1$ is the number of adjacency matrices sampled from $P_{A}^{\theta}$ to obtain output samples $\hat{y}_{i}=f_{\psi}(x,A_{i})\sim P_{y|x}^{\theta,\psi}$, whereas the pair $(x,y)$ is a pair from the training set $\mathcal{D}$. We remark that in (9) the third term of (8) – i.e., the one associated with the double expectation with respect to $P_{y|x}^{*}$ – is neglected as it does not depend on $\psi$ and $\theta$.

Gradient $\nabla_{\psi}\mathcal{L}^{dist}(\theta,\psi)$ is computed via automatic differentiation by averaging $\nabla_{\psi}\widehat{\text{MMD}}^{2}(\theta,\psi)$ within a mini-batch of observed data pairs $(x_{i},y_{i})\in\mathcal{D}$.
For $\nabla_{\theta}\mathcal{L}^{dist}(\theta,\psi)$, the same approach is not feasible. This limitation arises because the gradient is computed with respect to the same parameter vector $\theta$ that defines the integrated distribution.
Here, we rely on a score-function gradient estimator (SFE) Williams, []; Mohamed et al., [] which uses the log derivative trick to rewrite the gradient of an expected loss $L$ as
$\nabla_{\theta}\mathbb{E}_{A\sim P^{\theta}}[L(A)]=\mathbb{E}_{A\sim P^{\theta
}}[L(A)\nabla_{\theta}\log P^{\theta}(A)]$, with $P^{\theta}(A)$ denoting the likelihood of $A\sim P^{\theta}$.
Applying the SFE to our problem the gradient of the loss function w.r.t. $\theta$ reads:

$$ $\nabla_{\theta}\mathcal{L}^{dist}(\psi,\theta)=\underset{(x,y^{*})\sim P_{x,y} ^{*}}{\mathbb{E}}\bigg{[}\underset{\hat{y}_{1},\hat{y}_{2}\sim P^{\theta,\psi} _{y|x}}{\mathbb{E}}\left[\kappa(\hat{y}_{1},\hat{y}_{2})\nabla_{\theta}\log \left(P_{y|x}^{\theta,\psi}(\hat{y}_{1})P^{\theta,\psi}_{y|x}(\hat{y}_{2}) \right)\right]\\ -2\underset{\hat{y}\sim P^{\theta,\psi}_{y|x}}{\mathbb{E}}\left[\kappa(y^{*}, \hat{y})\nabla_{\theta}\log P^{\theta,\psi}_{y|x}(\hat{y})\right]\bigg{]}$ (10) $$

An apparent setback of SFEs is their high variance Mohamed et al., [], which we address in Section 5.3 by deriving a variance-reduction technique based on control variates that requires negligible computational overhead.

### 5.3 Variance-reduced loss for SFE

Two natural approaches to reduce the variance of MC estimates of (10) involve (i) increasing the number $B$ of training data points in the mini-batch used for each gradient estimate and (ii) increasing the number $N_{adj}$ of adjacency matrices sampled for each data point in (9).
These techniques act on two different sources of noise. Increasing $B$ decreases the variance coming from the data-generating process, whereas increasing $N_{adj}$ improves the approximation of the predictive distribution $P_{y|x}^{\theta,\psi}$.
Nonetheless, by fixing $B$ and $N_{adj}$, it is possible to further reduce the latter source of variance by employing the *control variates* method Mohamed et al., [] that, in our case, requires only a negligible computational overhead but sensibly improves the training speed (see Section 6).

Consider the expectation $\mathbb{E}_{A\sim P^{\theta}}[L(A)\nabla_{\theta}\log P^{\theta}(A)]$ of the SFE – both terms in (10) can be cast into that form.
With the control variates method, $L(A)$ is replaced by a surrogate function

$$ $\tilde{L}(A)=L(A)-\beta\Big{(}h(A)-\mathbb{E}_{A\sim P^{\theta}}[h(A)]\Big{)}$ (11) $$

that leads to a reduced variance in MC estimator while maintaining it unbiased. In this paper, we set function $h(A)$ to $\nabla_{\theta}\log P^{\theta}(A)$ and show how to compute a near-optimal choice for scalar value $\beta$, often called *baseline* in the literature. As the expected value of $\nabla_{\theta}\log P^{\theta}(A)$ is zero, gradient (10) rewrites as

$$ $\nabla_{\theta}\mathcal{L}^{dist}=\underset{(x,y^{*})\sim P_{x,y}^{*}}{\mathbb {E}}\bigg{[}\underset{A_{1},A_{2}\sim P^{\theta}_{A}}{\mathbb{E}}\left[\left( \kappa(f_{\psi}(x,A_{1}),f_{\psi}(x,A_{2}))-\beta_{1}\right)\;\nabla_{\theta} \log\left(P^{\theta}_{A}(A_{1})P^{\theta}_{A}(A_{2})\right)\right]\\ -2\underset{A\sim P^{\theta}_{A}}{\mathbb{E}}\left[\left(\kappa(y^{*},f_{\psi} (x,A))-\beta_{2}\right)\;\nabla_{\theta}\log P^{\theta}_{A}(A)\right]\bigg{]}.$ (12) $$

In Appendix B, we show that in our setup the best values of $\beta_{1}$ and $\beta_{2}$ are approximated by

$$ $\displaystyle\tilde{\beta}_{1}$ $\displaystyle=\underset{\begin{subarray}{c}x\sim P_{x}^{*}\\ A_{1},A_{2}\sim P^{\theta}_{A}\end{subarray}}{\mathbb{E}}\Big{[}\kappa\big{(}f _{\psi}(x,A_{1}),f_{\psi}(x,A_{2})\big{)}\Big{]},$ $\displaystyle\tilde{\beta}_{2}$ $\displaystyle=\underset{\begin{subarray}{c}(x,y^{*})\sim P_{x,y}^{*}\\ A\sim P^{\theta}_{A}\end{subarray}}{\mathbb{E}}\Big{[}\kappa\big{(}y^{*},f_{ \psi}(x,A)\big{)}\Big{]},$ (13) $$

which can be efficiently computed via MC, as kernel values in (13) are already computed to estimate (12).

### 5.4 Computational complexity

Focusing on the most significant terms, for every data pair $(x,y)$ in the training set, computing the loss $\mathcal{L}^{dist}$ requires $\mathcal{O}(N_{adj}^{2})$ kernel evaluations $\kappa(\hat{y}_{i},\hat{y}_{j})$ in (9), $\mathcal{O}(N_{adj})$ forward passes through the GNN $\hat{y}_{i}=f_{\psi}(x,A_{i})$ in (9) and $\mathcal{O}(N_{adj})$ likelihood computations $P^{\theta}_{A}(A_{i})$ in (12). The computation of baselines $\beta_{1}$ and $\beta_{2}$ in (13) requires virtually no overhead, as commented in previous Section 5.3.

Similarly, computing the loss’s gradients requires $\mathcal{O}(N_{adj}^{2})$ derivatives for what concerns the kernels, $\mathcal{O}(N_{adj})$ gradients $\nabla_{\psi}\hat{y}_{i}$ and $\nabla_{\theta}\log P^{\theta}_{A}(A_{i})$.
Our empirical analyses confirm that the processing within the GNN is the most demanding operation, despite the sparse computational graph in both forward and backward passes.

## 6 Experiments

This section empirically validates the proposed technique and the main paper’s claims.
Section 6.1 demonstrates that the proposed approach can successfully solve the joint learning problem and shows the effectiveness of the proposed variance reduction; Section 6.2 studies the impact of the number $N_{adj}$ of sampled adjacency matrices on calibration and prediction performance.
Experiments consider a synthetic dataset to allow evaluating the discrepancy between the true latent distribution and the learned one; these figures of merit would not be available in real-world applications as the latent distribution is indeed unknown.
We remark that the latent distribution $P_{A}^{*}$ is used *only* to assess performance and does not drive the model training in any way.

#### Dataset and models

Consider data-generating process (1) with latent distribution $P_{A}^{*}=P_{A}^{\theta^{*}}$ producing $N$-node adjacency matrices. $P_{A}^{*}$ is defined by a set of $N\times N$ independent Bernoulli distributions, each of which corresponds to the sampling probability of an edge. Function $f_{*}=f_{\psi^{*}}$ is a generic GNN with node-level readout, i.e., $f_{\psi^{*}}(\cdot,A):\mathbb{R}^{N\times d_{in}}\to\mathbb{R}^{N\times d_{out}}$. In the below experiments, $N$ is set to $12$, while input and output node feature dimensions are $d_{in}=4$ and $d_{out}=1$, respectively. The components $\theta^{*}$ are set to either 0 0 or $3/4$ according to the pattern depicted in Figure 9; the specifics of $f_{\psi^{*}}$ and $P_{x}^{*}$ are detailed in Appendix C. We result in a dataset of $35k$ input-output pairs $(x,y)$, 80% of which are used as training set, 10% as validation set, and the remaining 10% as test set. As predictive model family (2), we follow the same architecture of $f_{\psi^{*}}$ and $P_{A}^{\theta^{*}}$ ensuring that during all the experiments Assumption 3.1 is fulfilled. The model parameters are trained by optimizing the expected squared MMD in (9) with the rational quadratic kernel Bińkowski et al., [].

### 6.1 Graph structure learning & \& & optimal point predictions

To test our method’s ability to both calibrate the latent distribution and make optimal predictions, we train the model minimizing $\mathcal{L}^{dist}$ as described in Section 5.2.

Figure 2 reports the validation losses during training: MMD loss $\mathcal{L}^{dist}$, MAE between the learned parameters $\theta$ and the ground truth $\theta^{*}$ as $\mathcal{L}^{cal}$, and point-prediction loss $\mathcal{L}^{point}$ with $\ell$ being the MAE. The results are averaged over 20 different model initializations and error bars report $\pm 1$ standard deviation from the mean. Results are reported with and without applying the variance reduction (Section 5.3), by training only parameters $\theta$ while freezing $\psi$ to $\psi^{*}$ (same setting of Theorem 5.2), and by joint training of both $\psi$ and $\theta$.

Figure: ((a))
Refer to caption: extracted/5631973/Figures/val_MMD_loss_graphid_AND_joint_2.png

#### Solving the joint learning problem

Figure 2(a) shows that the training succeeded and the MMD loss $\mathcal{L}^{dist}$ approached its minimum (dotted line). Having minimized $\mathcal{L}^{dist}$, from Figure 2(b) we see that also the calibration of latent distribution $P^{\theta}_{A}$ was successful; in particular, the figure shows that the validation MAE ($N^{-2}\lVert\theta^{*}-\theta\rVert_{1}$) approaches zero as training proceeds (MAE $<0.04$). Regarding the point predictions, Figure 2(c) confirms that $\mathcal{L}^{point}$ reached its minimum value; recall that optimal prediction MAE is not 0 0, as the target variable $y$ is random, and note that a learning rate reduction is applied at epoch number $5$.
The optimality of the point-prediction is supported also by the performance on separate test data and with respect to the MSE as point-prediction loss $\ell$.
Moreover, we observe that calibration is achieved regardless of the variance reduction and whether or not parameters $\psi$ are trained.
Lastly, Figure 4 shows the learned parameters $\theta$ of the latent distribution and the corresponding absolute discrepancy resulted from a (randomly chosen) training run.

#### Optimization landscape of ℒ p ⁢ o ⁢ i ⁢ n ⁢ t superscript ℒ 𝑝 𝑜 𝑖 𝑛 𝑡 \mathcal{L}^{point} caligraphic_L start_POSTSUPERSCRIPT italic_p italic_o italic_i italic_n italic_t end_POSTSUPERSCRIPT and ℒ d ⁢ i ⁢ s ⁢ t superscript ℒ 𝑑 𝑖 𝑠 𝑡 \mathcal{L}^{dist} caligraphic_L start_POSTSUPERSCRIPT italic_d italic_i italic_s italic_t end_POSTSUPERSCRIPT

In this experiment, we analyze the values of $\mathcal{L}^{point}(\psi^{*},\theta)$ and $\mathcal{L}^{dist}(\psi^{*},\theta)$ for different values of $\theta$. $\mathcal{L}^{point}$ is computed employing MAE as loss function $\ell$.
Specifically, we let scalar $p$ vary from $1/2$ to $1$ and set all $\theta_{ij}=p$ for $i,j$ where $\theta^{*}_{ij}=3/4$. Figure 1 reports the obtained results, highlighting an almost flat $\mathcal{L}^{point}$ for values $\geq 0.725$. In contrast, $\mathcal{L}^{dist}$ displays a pronounced concave shape with a clear minimum around $\theta^{*}$ which suggests that calibration is easier when we minimize $\mathcal{L}^{dist}$ instead of $\mathcal{L}^{point}$.

Overall, we conclude that our approach is effective in solving the joint learning problem of calibrating the latent variable while producing optimal point predictions.

Figure: Figure 3: The learned parameters for the latent distribution corresponding to the stochastic adjacency matrix.
Refer to caption: extracted/5631973/Figures/graph_id/theta_GRAPH_ID.png

#### Variance reduction effectiveness

Figures 2(a), 2(b) and 2(c) demonstrate that the proposed variance reduction method (Section 5.2) yields notable advantages training speed up (roughly $50\%$ faster). For this reason, the next experiments rely on the variance reduction.

**Table 1: Convergence of the training loss $\mathcal{L}^{dist}$ varying the number of sampled adjacency matrices $N_{adj}$ and learning rates (LRs). X: no convergence; C: convergence, but not the smallest LR; S: smallest LR with convergence. Convergence is achieved when $\mathcal{L}^{dist}$ is below a fixed threshold within a predefined number of GNN calls (see also Section 6.2).**
| $N_{adj}$<br>LR | 0.01 | 0.02 | 0.05 | 0.1 | 0.2 | 0.5 | 1 | 2 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 2 | X | S | C | C | X | X | X | X |
| 4 | X | S | C | C | C | X | X | X |
| 8 | X | X | S | C | C | X | X | X |
| 16 | X | X | S | C | C | X | X | X |
| 32 | X | X | X | S | C | C | X | X |
| 64 | X | X | X | S | C | C | C | X |
| 128 | X | X | X | X | S | C | C | X |

### 6.2 Sensitivity to the number N a ⁢ d ⁢ j subscript 𝑁 𝑎 𝑑 𝑗 N_{adj} italic_N start_POSTSUBSCRIPT italic_a italic_d italic_j end_POSTSUBSCRIPT of sampled adjacency matrices

In this section, we study the relations between the number $N_{adj}$ of sampled adjacency matrices used to estimate $\nabla_{\theta}\mathcal{L}^{dist}$ in (12) via MC, the learning rates (LRs) used for training the models, and the achieved valued of the loss functions $\mathcal{L}^{dist}$, $\mathcal{L}^{point}$, and $\mathcal{L}^{cal}$.

#### Relations between N a ⁢ d ⁢ j subscript 𝑁 𝑎 𝑑 𝑗 N_{adj} italic_N start_POSTSUBSCRIPT italic_a italic_d italic_j end_POSTSUBSCRIPT and LRs

In Table 1, we show whether successful training was achieved within a given time period for different configurations of LR and $N_{adj}$. The time for successful learning is measured here in terms of evaluations of GNN $f_{\psi}$, the most demanding operation as discussed in Section 5.4. A training run is here considered successful if it achieves validation loss $\mathcal{L}^{dist}<-0.903$ (see Figure 2(a)) within 6M GNN evaluations.
As one can expect, too-small LRs do not allow convergence within the predefined lapse of time, whereas too-large LRs bring unstable learning.
In Table 1 we see that a range of LR for every $N_{adj}$ yields successful learning. in particular, larger $N_{adj}$ are associated with larger LRs – a reasonable behavior, if we recall that for every input $x$, $N_{adj}$ GNNs evaluations are requested. Therefore, one advantage of the variance reduction effect of increasing $N_{adj}$ is that larger learning rates (LR) can be employed.

#### Model accuracy as function of N a ⁢ d ⁢ j subscript 𝑁 𝑎 𝑑 𝑗 N_{adj} italic_N start_POSTSUBSCRIPT italic_a italic_d italic_j end_POSTSUBSCRIPT

We evaluate the accuracy of models trained using different values for $N_{adj}$. In particular, for every $N_{adj}$ in Table 1 we select the smallest LR yielding successful training.
Figures 7, 7, and 7 show that the model performance is reasonably stable, except for $N_{adj}=2$. In particular, both $\mathcal{L}^{cal}$ and $\mathcal{L}^{point}$ losses on the considered problem appear equivalent for $N_{adj}\geq 16$, thus suggesting that $N_{adj}$ is not a critical hyperparameter to be tuned.

Figure: Figure 5: Final MAE on the $\theta$ parameters for different values of the hyperparameter $N_{adj}$ used in Equation (9).
Refer to caption: extracted/5631973/Figures/N_adj_vs_theta.png

## 7 Conclusions

Graph structure learning has emerged as a research field focused on learning graph topologies in support of solving downstream predictive tasks.
Assuming stochastic latent graph structures, we are led to a joint optimization objective:
(i) learning the correct distribution of the latent topology while (ii) achieving optimal predictions on the downstream task. In this paper, at first, we prove both positive and negative theoretical results to demonstrate that appropriate loss functions must be chosen to solve this joint learning problem.
Second, we propose a sampling-based learning method that does not require the computation of the predictive likelihood. Our empirical results demonstrate that this approach achieves optimal point predictions on the considered downstream task while also yielding calibrated latent graph distributions.

Finally, we acknowledge that the proposed method requires sampling and processing multiple adjacency matrices for each input and, although the model and prediction accuracy is enhanced, a computation overhead is requested.
We plan future research to explore the applicability of this method to real-world datasets and to other classes of neural networks beyond GNNs; the current study, in fact, focuses on a set of controlled experiments on synthetic data to validate all the theoretical claims.

## Acknowledgments

This research was funded by the Swiss National Science Foundation under grant 204061: *High-Order Relations and Dynamics in Graph Neural Networks*.

## 

- Berger, [1990]
Berger, J. O. (1990).
Statistical decision theory.
In Time Series and Statistics, pages 277–284. Springer.
- Bińkowski et al., [2018]
Bińkowski, M., Sutherland, D. J., Arbel, M., and Gretton, A. (2018).
Demystifying mmd gans.
In International Conference on Learning Representations.
- Cini et al., [2023]
Cini, A., Zambon, D., and Alippi, C. (2023).
Sparse graph learning from spatiotemporal time series.
Journal of Machine Learning Research, 24:1–36.
- De Felice et al., [2024]
De Felice, G., Cini, A., Zambon, D., Gusev, V., and Alippi, C. (2024).
Graph-based Virtual Sensing from Sparse and Partial
Multivariate Observations.
In The Twelfth International Conference on Learning
Representations.
- Elinas et al., [2020]
Elinas, P., Bonilla, E. V., and Tiao, L. (2020).
Variational inference for graph convolutional networks in the absence
of graph data and adversarial settings.
Advances in Neural Information Processing Systems,
33:18648–18660.
- Fatemi et al., [2021]
Fatemi, B., El Asri, L., and Kazemi, S. M. (2021).
Slaps: Self-supervision improves structure learning for graph neural
networks.
Advances in Neural Information Processing Systems,
34:22667–22681.
- Fey and Lenssen, [2019]
Fey, M. and Lenssen, J. E. (2019).
Fast graph representation learning with pytorch geometric.
arXiv preprint arXiv:1903.02428.
- Fout et al., [2017]
Fout, A., Byrd, J., Shariat, B., and Ben-Hur, A. (2017).
Protein interface prediction using graph convolutional networks.
Advances in neural information processing systems, 30.
- Franceschi et al., [2019]
Franceschi, L., Niepert, M., Pontil, M., and He, X. (2019).
Learning discrete structures for graph neural networks.
In International conference on machine learning, pages
1972–1982. PMLR.
- Gneiting, [2011]
Gneiting, T. (2011).
Making and Evaluating Point Forecasts.
Journal of the American Statistical Association,
106(494):746–762.
- Gretton et al., [2012]
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A.
(2012).
A kernel two-sample test.
The Journal of Machine Learning Research, 13(1):723–773.
- Harris et al., [2020]
Harris, C. R., Millman, K. J., Van Der Walt, S. J., Gommers, R., Virtanen, P.,
Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N. J., et al.
(2020).
Array programming with numpy.
Nature, 585(7825):357–362.
- Hunter, [2007]
Hunter, J. D. (2007).
Matplotlib: A 2d graphics environment.
Computing in science & engineering, 9(03):90–95.
- Jiang et al., [2019]
Jiang, B., Zhang, Z., Lin, D., Tang, J., and Luo, B. (2019).
Semi-supervised learning with graph learning-convolutional networks.
In Proceedings of the IEEE/CVF conference on computer vision and
pattern recognition, pages 11313–11320.
- Kazi et al., [2022]
Kazi, A., Cosmo, L., Ahmadi, S.-A., Navab, N., and Bronstein, M. M. (2022).
Differentiable graph module (dgm) for graph convolutional networks.
IEEE Transactions on Pattern Analysis and Machine Intelligence,
45(2):1606–1617.
- Kingma and Ba, [2014]
Kingma, D. P. and Ba, J. (2014).
Adam: A method for stochastic optimization.
arXiv preprint arXiv:1412.6980.
- Kipf et al., [2018]
Kipf, T., Fetaya, E., Wang, K.-C., Welling, M., and Zemel, R. (2018).
Neural relational inference for interacting systems.
In International conference on machine learning, pages
2688–2697. PMLR.
- Mnih et al., [2016]
Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T., Harley, T.,
Silver, D., and Kavukcuoglu, K. (2016).
Asynchronous methods for deep reinforcement learning.
In International conference on machine learning, pages
1928–1937. PMLR.
- Mohamed et al., [2020]
Mohamed, S., Rosca, M., Figurnov, M., and Mnih, A. (2020).
Monte carlo gradient estimation in machine learning.
The Journal of Machine Learning Research, 21(1):5183–5244.
- Müller, [1997]
Müller, A. (1997).
Integral probability metrics and their generating classes of
functions.
Advances in applied probability, 29(2):429–443.
- Niepert et al., [2021]
Niepert, M., Minervini, P., and Franceschi, L. (2021).
Implicit MLE: Backpropagating Through Discrete Exponential
Family Distributions.
In Advances in Neural Information Processing Systems,
volume 34, pages 14567–14579. Curran Associates, Inc.
- 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.
- Rényi, [1961]
Rényi, A. (1961).
On measures of entropy and information.
In Proceedings of the fourth Berkeley symposium on mathematical
statistics and probability, volume 1: contributions to the theory of
statistics, volume 4, pages 547–562. University of California Press.
- Scarselli et al., [2008]
Scarselli, F., Gori, M., Tsoi, A. C., Hagenbuchner, M., and Monfardini, G.
(2008).
The graph neural network model.
IEEE transactions on neural networks, 20(1):61–80.
- Shang et al., [2021]
Shang, C., Chen, J., and Bi, J. (2021).
Discrete graph structure learning for forecasting multiple time
series.
In International Conference on Learning Representations.
- Shlomi et al., [2020]
Shlomi, J., Battaglia, P., and Vlimant, J.-R. (2020).
Graph neural networks in particle physics.
Machine Learning: Science and Technology, 2(2):021001.
- Sutton et al., [1999]
Sutton, R. S., McAllester, D., Singh, S., and Mansour, Y. (1999).
Policy gradient methods for reinforcement learning with function
approximation.
Advances in neural information processing systems, 12.
- Williams, [1992]
Williams, R. J. (1992).
Simple statistical gradient-following algorithms for connectionist
reinforcement learning.
Machine learning, 8:229–256.
- Wu et al., [2020]
Wu, Z., Pan, S., Long, G., Jiang, J., Chang, X., and Zhang, C. (2020).
Connecting the dots: Multivariate time series forecasting with graph
neural networks.
In Proceedings of the 26th ACM SIGKDD international conference
on knowledge discovery & data mining, pages 753–763.
- Wu et al., [2019]
Wu, Z., Pan, S., Long, G., Jiang, J., and Zhang, C. (2019).
Graph wavenet for deep spatial-temporal graph modeling.
In Proceedings of the 28th International Joint Conference on
Artificial Intelligence, pages 1907–1913.
- Yu et al., [2021]
Yu, D., Zhang, R., Jiang, Z., Wu, Y., and Yang, Y. (2021).
Graph-revised convolutional network.
In Machine Learning and Knowledge Discovery in Databases:
European Conference, ECML PKDD 2020, Ghent, Belgium, September 14–18, 2020,
Proceedings, Part III, pages 378–393. Springer.
- Zhang et al., [2019]
Zhang, Y., Pal, S., Coates, M., and Ustebay, D. (2019).
Bayesian graph convolutional neural networks for semi-supervised
classification.
In Proceedings of the AAAI conference on artificial
intelligence, volume 33, pages 5829–5836.
- Zhu et al., [2021]
Zhu, Y., Xu, W., Zhang, J., Liu, Q., Wu, S., and Wang, L. (2021).
Deep graph structure learning for robust representations: A survey.
arXiv preprint arXiv:2103.03036, 14:1–1.

## Appendix A Proofs of the theoretical results

### A.1 Minimizing ℒ p ⁢ o ⁢ i ⁢ n ⁢ t superscript ℒ 𝑝 𝑜 𝑖 𝑛 𝑡 \mathcal{L}^{point} caligraphic_L start_POSTSUPERSCRIPT italic_p italic_o italic_i italic_n italic_t end_POSTSUPERSCRIPT does not guarantee calibration

Proof of Proposition 4.1.

###### Proof.

Recall the definition of $\mathcal{L}^{point}$ in (4) using (3)

$$ $\mathcal{L}^{point}(\psi,\theta)=\mathbb{E}_{x}\Big{[}\mathbb{E}_{y^{*}\sim P^ {*}_{y|x}}\Big{[}\ell\big{(}y^{*},T\big{[}P^{\theta,\psi}_{y|x}\big{]}\big{)} \Big{]}\Big{]}$ $$

Given loss function $\ell$, $T$ is, by definition [Berger,, ; Gneiting,, ], the functional that minimizes

$$ $\mathbb{E}_{y^{*}\sim P^{*}_{y|x}}\Big{[}\ell\big{(}y^{*},T\big{[}P^{*}_{y|x} \big{]}\big{)}\Big{]}$ $$

Therefore, if $P^{\theta,\psi}_{y|x}=P^{*}_{y|x}\implies\mathcal{L}^{point}$ is minimal. If another distribution over $y$, namely, $P^{\psi^{\prime},\theta^{\prime}}_{y|x}$ parametrized by $\theta^{\prime}$ and $\psi^{\prime}$ satisfies:

$$ $T\Big{[}P^{\psi^{\prime},\theta^{\prime}}_{y|x}\Big{]}=T\Big{[}P^{*}_{y|x}\Big {]}$ $$

almost surely on $x$, then,

$$ $\displaystyle\mathcal{L}^{point}(\theta^{\prime},\psi^{\prime})$ $\displaystyle=\mathbb{E}_{x}\Big{[}\mathbb{E}_{y^{*}\sim P^{*}_{y|x}}\Big{[} \ell\big{(}y^{*},T\big{[}P^{\psi^{\prime},\theta^{\prime}}_{y|x}\big{]}\big{)} \Big{]}\Big{]}$ $\displaystyle=\mathbb{E}_{x}\Big{[}\mathbb{E}_{y^{*}\sim P^{*}_{y|x}}\Big{[} \ell\big{(}y^{*},T\big{[}P^{*}_{y|x}\big{]}\big{)}\Big{]}\Big{]}$ $$

Thus, $P^{\psi^{\prime},\theta^{\prime}}_{y|x}$ minimizes $\mathcal{L}^{point}$.

Appendix A.2 discusses graph distributions where $T\big{[}P^{\psi^{\prime},\theta^{\prime}}_{y|x}\big{]}=T\big{[}P^{*}_{y|x}\big
{]}$ but $P^{\psi^{\prime},\theta^{\prime}}_{y|x}\not=P^{*}_{y|x}$. We conclude that reaching the minimum of $\mathcal{L}^{point}(\psi,\theta)$ does not imply $P^{\psi,\theta}_{y|x}=P^{*}_{y|x}$.
∎

###### Proof.

### A.2 Minimizing ℒ p ⁢ o ⁢ i ⁢ n ⁢ t superscript ℒ 𝑝 𝑜 𝑖 𝑛 𝑡 \mathcal{L}^{point} caligraphic_L start_POSTSUPERSCRIPT italic_p italic_o italic_i italic_n italic_t end_POSTSUPERSCRIPT does not guarantee calibration: an example with MAE

In this section, we show that $\mathcal{L}^{point}$ equipped with MAE as $\ell$ admits multiple global minima for different parameters $\theta$, even for simple models and $f_{\psi}=f^{*}$.

Consider a single Bernoulli of parameter $\theta^{*}>1/2$ as latent variable $A$ and a scalar function $f^{*}(x,A)$ such that $f^{*}(x,1)>f^{*}(x,0)$ for all $x$.
Given input $x$ the value of functional $T(P^{*}_{y|x})$ that minimizes

$$ $\mathbb{E}_{y\sim P^{*}_{y|x}}\Big{[}\left|y-T\big{[}P^{*}_{y|x}\big{]}\right| \Big{]}=\theta^{*}\left|f^{*}(x,1)-T\big{[}P^{*}_{y|x}\big{]}\right|+(1-\theta ^{*})\left|f^{*}(x,0)-T\big{[}P^{*}_{y|x}\big{]}\right|$ $$

is $T(P^{*}_{y|x})=f^{*}(x,1)$; this derives from the fact that range of $f^{*}$ is $\{f^{*}(x,0),f^{*}(x,1)\}$ and the likelihood of $f^{*}(x,1)$ is larger than that of $f^{*}(x,0)$.

Note that $T\big{[}P^{*}_{y|x}\big{]}=f^{*}(x,1)$ for all $x$, therefore also $\mathcal{L}^{point}$ is minimized by such $T$. Moreover, $T\big{[}P^{*}_{y|x}\big{]}$ is function of $\theta^{*}$ and equal to $f^{*}(x,1)$ for all $\theta>1/2$. We conclude that for any $\theta\neq\theta^{*}$ distributions $P^{\theta,\psi}_{y|x}$ and $P^{*}_{y|x}$ are different, yet both of them minimize $\mathcal{L}^{point}$ if $\theta>1/2$.

A similar reasoning applies for $\theta^{*}<1/2$.

### A.3 Minimizing ℒ d ⁢ i ⁢ s ⁢ t superscript ℒ 𝑑 𝑖 𝑠 𝑡 \mathcal{L}^{dist} caligraphic_L start_POSTSUPERSCRIPT italic_d italic_i italic_s italic_t end_POSTSUPERSCRIPT guarantees calibration and optimal point predictions.

Proof of Theorem 5.2

###### Proof.

Recall from Equation (6) that

$$ $\mathcal{L}^{dist}(\theta)=\mathbb{E}_{x}\Big{[}\Delta(P_{y|x}^{*},P_{y|x}^{ \theta})\Big{]}$ $$

We start by proving that if $\mathcal{L}^{dist}(\theta,\psi)=0\implies\mathcal{L}^{point}(\theta,\psi)\text
{ is minimal}$.

Note that $\mathcal{L}^{dist}(\theta,\psi)=0$ implies that $\Delta(P_{y|x}^{*},P_{y|x}^{\theta})=0$ almost surely in $x$. Then, by Assumption 5.1, $P_{y|x}^{*}=P_{y|x}^{\psi,\theta}$ almost surely on $x$ and, in particular, $T[P_{y|x}^{*}]=T[P_{y|x}^{\psi,\theta}]$, which leads to $\mathcal{L}^{point}(\psi,\theta)$ being minimal (Proposition 4.1).

We now prove that if $\mathcal{L}^{dist}(\theta,\psi^{*})=0\implies\mathcal{L}^{cal}(\theta)=0$.

From the previous step, we have that $\mathcal{L}^{dist}(\theta,\psi)=0$ implies $P_{y|x}^{*}=P_{y|x}^{\psi,\theta}$ almost surely for $x\in I$. Under the assumption that $f_{\psi}=f_{*}$ and the injectivity of $f_{*}$ in such $x\in I$, for any output $y$ a single $A$ exists such that $f_{*}(x,A)=y$. Therefore, the probability mass function of $y$ equals that of $A$. Accordingly, $P_{y|x}^{*}=P_{y|x}^{\psi,\theta}$ implies $P_{A}^{*}=P_{A}^{\theta}$.

∎

Here, we also prove a corollary of Theorem 5.2.

###### Corollary A.1 .

Under Assumptions 3.1 and 5.1,
if

- 1.
$\exists\bar{x}\in\textit{Supp}(P^{*}_{x})\subseteq\mathcal{X}$ such that $f^{*}({\bar{x}};\cdot)$ is injective,
- 2.
$f^{*}(x,A)$ is continuous in $\bar{x}$ $\forall A\in\mathcal{A}$,

then

$$ $\displaystyle\mathcal{L}^{dist}(\theta,\psi^{*})=0$ $\displaystyle\implies\begin{cases}\mathcal{L}^{point}(\theta,\psi^{*})\text{ is minimal}\\ \mathcal{L}^{cal}(\theta)=0,\end{cases}$ $$

The corollary shows that it is sufficient that $f^{*}$ is continuous in $x$ and there exists one point $\bar{x}$ where $f^{*}(\bar{x},{}\cdot{})$ is injective to meet theorem’s hypothesis $\mathbb{P}_{x\sim P_{x}^{*}}(I)>0$; we observe that, as $\mathcal{A}$ is discrete, the injectivity assumption is not as restrictive as if the domain were continuous.

###### Proof.

As $\mathcal{A}$ is a finite set, the minimum $\bar{\epsilon}=\min_{A,A^{\prime}\in\mathcal{A}}\lVert f^{*}(\bar{x},A)-f^{*}(
\bar{x},A^{\prime})\rVert>0$ exists and, by the injectivity assumption, is strictly positive.

By continuity of $f^{*}({}\cdot{},A)$, for every $\epsilon<\frac{1}{2}\bar{\epsilon}$ there exists $\delta$, such that for all $x\in B(\bar{x},\delta)$ we have $\lVert f^{*}(\bar{x},A)-f^{*}(x,A)\rVert<\epsilon$. It follows that, $\forall x\in B$,

$$ $\displaystyle\lVert f^{*}(x,A)$ $\displaystyle-f^{*}(x,A^{\prime})\rVert$ $\displaystyle\geq\lVert f^{*}({\bar{x}},A)-f^{*}({\bar{x}},A^{\prime})\rVert- \lVert f^{*}({\bar{x}},A)-f^{*}(x,A)\rVert-\lVert f^{*}(\bar{x},A^{\prime})-f^ {*}(x,A^{\prime})\rVert$ $\displaystyle\geq\lVert f^{*}({\bar{x}},A)-f^{*}({\bar{x}},A^{\prime})\rVert-2\epsilon$ $\displaystyle\geq\lVert f^{*}({\bar{x}},A)-f^{*}({\bar{x}},A^{\prime})\rVert- \bar{\epsilon}>0$ $$

Finally, as $\bar{x}\in{\rm Supp}(P_{x}^{*})$ and $B(\bar{x},\delta)\subseteq I$, we conclude that

$$ $\mathbb{P}_{x}(I)\geq\mathbb{P}_{x}(B(\bar{x},\delta))>0,$ $$

therefore, we are in the hypothesis of Theorem 5.2 and can conclude that

$$ $\displaystyle\mathcal{L}^{dist}(\theta,\psi^{*})=0$ $\displaystyle\implies\begin{cases}\mathcal{L}^{point}(\theta,\psi^{*})\text{ is minimal}\\ \mathcal{L}^{cal}(\theta)=0,\end{cases}$ $$

∎

###### Proof.

###### Corollary A.1 .

###### Proof.

## Appendix B Estimation of optimal β 1 subscript 𝛽 1 \beta_{1} italic_β start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT and β 2 subscript 𝛽 2 \beta_{2} italic_β start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT

Here we show that, when reducing the variance of the SFE via control variates in (12), the best $\beta_{1}$ and $\beta_{2}$ can be approximated by

$$ $\displaystyle\tilde{\beta}_{1}$ $\displaystyle=\underset{\begin{subarray}{c}x\sim P_{x}^{*}\\ A_{1},A_{2}\sim P^{\theta}_{A}\end{subarray}}{\mathbb{E}}\Big{[}\kappa\left(f_ {\psi}(x,A_{1}),f_{\psi}(x,A_{2})\right)\Big{]},$ $\displaystyle\tilde{\beta}_{2}$ $\displaystyle=\underset{\begin{subarray}{c}(x,y^{*})\sim P_{x,y}^{*}\\ A\sim P^{\theta}_{A}\end{subarray}}{\mathbb{E}}\Big{[}\kappa\left(y^{*},f_{ \psi}(x,A)\right)\Big{]},$ (14) $$

Consider generic function $L(A)$ depending on a sample $A$ of a parametric distribution $P^{\theta}_{A}(A)$ and the surrogate loss $\tilde{L}(A)$ in (11), i.e.,

$$ $\tilde{L}(A)=L(A)-\beta\Big{(}h(A)-\mathbb{E}_{A\sim P^{\theta}}[h(A)]\Big{)};$ (15) $$

This choice is not new in the literature Sutton et al., []; Mnih et al., [] where $\beta$ is often referred to as *baseline*.
The 1-sample MC approximation of the loss becomes

$$ $\nabla_{\theta}\mathbb{E}_{A\sim P^{\theta}}[L(A)]\approx\tilde{L}(A^{\prime}) \nabla_{\theta}\log P^{\theta}(A^{\prime})=(L(A^{\prime})-\beta)\nabla_{\theta }\log P^{\theta}(A^{\prime}),$ (16) $$

with $A^{\prime}$ sampled from $P^{\theta}_{A}$.
The variance of the estimator is

$$ $\mathbb{V}_{A\sim P^{\theta}}\left[(L(A)-\beta)\nabla_{\theta}\log P^{\theta}( A)\right]=\mathbb{V}_{A\sim P^{\theta}}\left[L(A)\nabla_{\theta}\log P^{\theta }(A)\right]+\\ +\beta^{2}\;\mathbb{E}_{A\sim P^{\theta}}\left[\left(\nabla_{\theta}\log P^{ \theta}(A)\right)^{2}\right]-2\beta\;\mathbb{E}_{A\sim P^{\theta}}\left[L(A) \left(\nabla_{\theta}\log P^{\theta}(A)\right)^{2}\right]$ (17) $$

and the optimal value $\beta$ that minimizes it is

$$ $\tilde{\beta}=\frac{\mathbb{E}_{A\sim P^{\theta}}\left[L(A)\left(\nabla_{ \theta}\log P^{\theta}(A)\right)^{2}\right]}{\mathbb{E}_{A\sim P^{\theta}} \left[\left(\nabla_{\theta}\log P^{\theta}(A)\right)^{2}\right]}$ (18) $$

If we approximate the numerator with $\mathbb{E}[L(A)]\mathbb{E}[(\nabla_{\theta}\log P^{\theta}(A))^{2}]$, we obtain that $\tilde{\beta}\approx\mathbb{E}[L(A)]$. By substituting $L(A)$ with
the two terms of (10) we get the values of $\beta_{1}$ and $\beta_{2}$ in (14).

We experimentally validate the effectiveness of this choice of $\beta$ in Section 6.

## Appendix C Further experimental details

### C.1 Dataset description and models

In this section, we describe the considered synthetic dataset, generated from the system model (1).
The latent graph distribution $P^{*}_{A}$ is a multivariate Bernoulli distribution of parameters $\theta^{*}_{ij}$: $P^{*}_{A}\equiv P_{\theta^{*}}(A)$ = $\prod_{ij}$ $\theta_{ij}^{*A_{ij}}$ $(1-\theta^{*}_{ij})^{(1-A_{ij})}$. The components of $\theta^{*}$ are all null, except for the edges of the graph depicted in Figure 8 which are set to $3/4$.

Figure: Figure 8: The adjacency matrices used in this paper are sampled from this graph. Each edge in orange is independently sampled with probability $\theta^{*}$. In the picture 3 communities of an arbitrarly large graph are shown.
Refer to caption: x1.png

A heatmap of the adjacency matrix can be found in Figure 9.

Figure: Figure 9: $\theta^{*}_{ij}$ parameters for each edge of the latent adjacency matrix. Each square corresponds to an edge, the number inside is the probability of sampling that edge for each prediction.
Refer to caption: extracted/5631973/Figures/true_theta_parameters.png

Regarding the GNN function $f^{*}$, we use the following system model:

$$ $\begin{dcases*}y=f_{\psi^{*}}(A,x)=\text{tanh}\left(\sum_{l=1}^{L}\mathbbm{1}[ A^{l}\centernot=0]x\psi^{*}_{l}\right)\\ A\sim P_{\theta^{*}}(A)\end{dcases*}$ (19) $$

Where $\mathbbm{1}[\cdot]$ is the element-wise indicator function: $\mathbbm{1}[a]=1\iff a$ is true.
$x\in\mathbb{R}^{N\text{ x }d_{in}}$ are randomly generated inputs: $x\sim\mathcal{N}(0,\sigma_{x}^{2}\mathbb{I})$.
$\psi^{*}_{l}\in\mathbb{R}^{d_{out}\text{ x }d_{in}}$ are part of the system model parameters. We summarize the parameters considered in our experiment in Table 2.

**Table 2: Table of the parameters used to generate the synthetic dataset.**
| $\theta^{*}$ | $0.75$ |
| --- | --- |
| $\sigma_{x}$ | $1.5$ |
| $N$ | $12$ |
| $d_{in}$ | $4$ |
| $d_{out}$ | $1$ |
| $\psi_{1}^{*}$ | [$-0.2$, $0.4$, $-0.8$, $0.6$] |
| $\psi_{2}^{*}$ | [$-0.3$, $0.8$, $0.2$, $-0.7$] |

The approximating model family (2) used in the experiment is the same as the data-generating process, with all components of parameter vectors $\theta$ and $\psi$ being trainable.
The squared MMD discrepancy is defined over Rational Quadratic kernel Bińkowski et al., []

$$ $\kappa(y^{\prime},y^{\prime\prime})=\left(1+\frac{\lVert y^{\prime}-y^{\prime \prime}\rVert_{2}^{2}}{2\,\alpha\,\sigma^{2}}\right)^{-\alpha}$ $$

of parameters $\sigma=0.7$ and $\alpha=0.02$.

The model is trained using Adam optimizer Kingma and Ba, [] with parameters $\beta_{1}=0.6$, $\beta_{2}=0.95$.
Where not specified, the learning rate is set to $0.1$ and decreased to $0.01$ after 5 epochs. We grouped data points into batches of size 128. Initial values of $\theta$ are independently sampled from the $\mathcal{U}(0.25,0.35)$ uniform distribution.

### C.2 Description of the experiment in Section 4

In this experiment, we generate 512 data points using the system model described in Appendix C.1.
We construct a model identical to the system model, except that $\theta_{ij}=p$ for all $i,j$ where $\theta_{i,j}^{*}=0.75$ and 0 0 elsewhere. We vary scalar $p$ from $0.5$ to $1$ with steps of $0.025$.
Therefore, only the model with $p=0.75$ is identical to the data-generating model.

For each input $x$ in the dataset, a point prediction is produced by sampling $N_{adj}=32$ adjacency matrices and computing the median. This approach allows to estimate $\mathcal{L}^{point}$ using the MAE as loss function $\ell$, as depicted by the red points in Figure 1, for different values of $\theta$. For comparison purposes, we estimate $\mathcal{L}^{dist}$ using the maximum mean discrepancy as proposed in Section 5.

### C.3 Compute resources and open-source software

The paper’s experiments were run on a workstation with AMD EPYC 7513 processors and NVIDIA RTX A5000 GPUs; on average, a single model training terminates in a few tens of minutes with a memory usage of about 2GB.

The developed code relies on PyTorch Paszke et al., [] and the following additional open-source libraries:
PyTorch Geometric Fey and Lenssen, [], NumPy Harris et al., [] and Matplotlib Hunter, [].

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Why the inindependent bernouli distribution is used in experiment for adjacency matrix A? In reality, the graph structures follows (1) Erdős–Rényi Model (2) Barabási–Albert Model (BA Model). 
2. Maybe I missed somewhere, what is the feature distribution $P(X)$ used in experiments?
soundness: 2
strengths: 1. The theoretical results of the paper are fully-justified although the reviewer didn't check every details of the proof. 
2. The paper proves the feasibility of minimizing joint distribution discrepancy of (A<X) leads to both optimal point predictions and latent distribution calibration
summary: This paper tackles the latent structure learning on graph structured data and demonstrate that minimizing prediction function does not guarantee a calibrated model with rigorous theoretical justifications. The author proposes a sampling-based optimization using Maximum Mean Discrepancy (MMD) between output distributions and shows the effectiveness of joint learning tasks.
weaknesses: 1. The paper does not focus on graph-specific optimization techniques. It seems that no matter how strcuture of {x} is generated, the consluion of the theoretical results always applied. In the proof of injectivity for example, the graph structure is modeled as a linear projection with continuous value 
2. The paper oversimplifies the graph structures, for example, continuous adjacency matrices and bernouli distribution in experiments, which makes the practical impact of the paper quite limited in my opinion.
3. The contribution and proposed algorithm needs to be justified on more data modality with latent structure given that the paper seldomly uses property of the graph structure in its optimization. I would like to see the paper remove the focus on graph structures and verify the effectiveness in different domains such as image/audio etc.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Why MMD? And why not KL or JS divergence? Lines 247-251 are insufficient. Please give reasons why the theory is not feasible or experimental gaps.
soundness: 2
strengths: 1. The uncertainty issue of graph structure learning is important.
2. Theoretical analysis and proofs are provided.
summary: The authors point out the limitation of current point-prediction methods, which cannot guarantee the calibration of the distribution of adjacency matrix $A$.
Therefore, the authors propose a sampling-based learning method for joint optimization.
weaknesses: 1. Section 4 has no analysis specifically related to the adjacency matrix $A$, which contradicts contribution 1. Maybe $A$ is input into $L$ or $x$, but the analysis in Section 4 could be applied to any scenario and cannot bring any insights for the community of graph learning.

2. Similarly, the methodological contribution in Section 5 is also universal and seems unrelated to GSL or GNNs.

3. Lack of experimental results on real datasets compared with GSL baselines.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: Listed in my above comments, namely addressing motivation (if the downstream prediction is the same, why do we care at all about GSL producing high fidelity graph estimates?) and extensions to real data.
soundness: 3
strengths: I believe this work to be the first to consider the calibration aspect of the latent graph structure, which is an interesting point.

A primary supporting argument for GSL methods is that the learned graph is itself useful for interpretability purposes. But most (all?) approaches make little effort to (i) show an instance where this is actually useful beyond a nice visualization or (ii) discuss why we should believe this underlying graph corresponds to some real object.

This work seems to address address (ii).
summary: This paper addresses the Graph Structure Learning (GSL) problem, sometimes referred to as the ‘Latent Graph Learning’ problem, which concerns co-learning the graph structure and GNN weights for downstream prediction. The paper reveals issues with point-predictor losses with regards to the learning of the latent graph structure, and proposes a sampling based approach to overcome such issues.
weaknesses: **Motivation**

As best I can understand, the paper attempts to address (ii), but without addressing (i), i.e. the motivation of the entire effort. Why should anyone care in the first place? 

The authors attempt to provide high-level applications, e.g., “Examples include e.g., social interactions where links can intermittently be present, traffic flows affected by road closures and temporary detours, and adaptive communication routing. It follows that a probabilistic framework is appropriate to accurately capture the uncertainty in the learned relations whenever randomness affects the graph topology.”

But it is not clear to me how GSL is/would be actually used in these applications, why uncertainty over such an inferred latent graph is important in these applications, nor why calibrated uncertainty measures of the latent graph are important.

&nbsp;


**Gaps in Related Work**

There are large gaps in related work on graph structure learning, e.g. using unrollings and/or Bayesian Neural Network. For recent work see Graph Structure Learning with Interpretable Bayesian Neural Networks, which addresses learning graphs with uncertainty estimates. Link at the bottom.

&nbsp;


**Extension to Real Data**

As the underlying network topology is rarely ever observed, how do we evaluate whether our estimate of it (along with the corresponding uncertainty/calibration) is good in a real data setting? The only way I can think of is by using it to make downstream predictions on labels we do observe, i.e. marginalizing it out and evaluating the resulting posterior predictive.


&nbsp;


**Minor Comments**

In Eqn (1) and (2) it may make sense to place A \sim P^{\theta)_A above \hat{y} = f_{\psi}(x, A). This is more standard in the probabilistic community which often views this model as a data-generating process:  A must be sampled first before it can be used as an input for f.

&nbsp;

Graph Structure Learning with Interpretable Bayesian Neural Networks, https://openreview.net/forum?id=2noXK5KBbx

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: See weakness.
soundness: 3
strengths: - This paper first explores calibrating graph structure learning by minimizing the dissimilarity measure between two distributions.
- The theoretical analysis shows how the proposed method can be beneficial in cases, such as when  $f^{\ast} = f_{\psi}$, where the proposed method succeeds while conventional methods fail.
- The authors show that proposed method can be useful by incorporating variance reduction techniques and addressing the complexity.
- The experiments validate the theorems from several perspectives.
summary: This paper aims to calibrate the latent distribution of graphs for predicting outcomes of interacting entities. To this end, the paper considers a dissimilarity measure, which can be f-divergence, Stein variational gradients, or maximum mean discrepancy, between two predictive distributions. Then, the paper demonstrates that minimizing this measure ensures success of calibration under certain conditions where conventional point-wise approaches fail. For practical implementation, the dissimilarity measure is considered with maximum mean discrepancy computed over finite samples. The experiments focus on validating the proposed theorem.
weaknesses: I have no concerns about the methods or their theorem. However, my overall concerns stem from the experiments.

- **Lack of baselines.** The experiments only provide the outputs of the proposed method and validate the theorem. However, they do not provide experiments for showing how conventional approaches fail in contrast, such as their shortcomings in estimating the underlying graph distribution. This limits the emphasis on the advantages of the proposed method.

- **Lack of benchmarks.** The experiments focus on synthetic benchmarks with limited scope, which limits highlighting the benefits of the proposed method. Are there no conventional benchmarks, such as node classification benchmarks, that could better highlight the benefits in practice?

- **Lack of uncertainty analysis.** In my understanding, the ultimate goal is to model output $y$ uncertainty by calibrating the latent distribution for graphs. However, the experiments lack analysis of the capability to model this uncertainty.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Why is the problem of learning the distribution of a latent random variable (the adjacency matrix) relevant in cases where the inference task can be perfectly solved? Some settings where sampling the graph, evaluating edge probability or explainability should be put forward to better ground the contribution.

2. Where does $P_x*$ comes from? There is no discussion around $X$ being a random variable in the problem setting, and it's relation (if any) to $p(A)$. Currently, it may seems like $A$ and $X$ are assumed to be independent, which seems unlikely?

3. Why is there no baselines presented in the experiment section? For example, Anees Kazi et al. (2022) could have been presented to highlight the importance of the proposed joint learning formulation. 

4. Aren't most GNN architecture not injective and would violate the injective constraint from Theorem 5.2?
soundness: 4
strengths: 1. The paper is well-structured and clearly written, with a logical flow that guides the reader through the problem, methodology, and results on the synthetic experiment.

2. The paper provides a comprehensive treatment of the problem, carefully addressing its assumptions and hypotheses.  It thoroughly discusses the implications and limitations of each assumption and hypothesis, and it empirically verifies what happens when the assumptions are violated. 

3. The proposed loss function is principled, built on a solid theoretical foundation.
summary: (After reading the revised version, I increased my score from 5 to 6.)

This paper addresses the problem of jointly learning the latent adjacency matrix $A$ and a node inference task via a function $f(A, x)$. It considers a family of predictive models $P_{y|x}^{\theta, \psi}$, where one component learns the distribution $P(A)$ and another models the inference task $f$. The paper shows that optimizing a pointwise loss does not necessarily ensure proper learning of the distribution over $P(A)$. To address this, the authors propose a new loss formulation based on a distribution dissimilarity metric that guarantees learning both the latent graph structure and the inference function for the considered class of functions and settings. The methodology provides one example of how to optimize such a loss by choosing Maximum Mean Discrepancy (MMD) as the dissimilarity metric. The authors then develop a method to optimize the proposed MMD loss and provide empirical experiments to show that their approach successfully learns both the distribution $P(A)$ and the inference function $f$
weaknesses: 1. Motivation. 
The motivation for the joint learning task is not clearly articulated. The paper’s main result highlights cases where successfully solving one task does not imply solving an unrelated task (i.e., learning the distribution of a latent random variable). However, the primary cited motivation for learning the graph distribution is to improve performance on a downstream task. If the downstream task can already be solved, it is unclear why learning the distribution over a latent graph would then be necessary. In summary, the authors should provide a practical setting where their proposed joint learning task and loss function would be called for.

2. Contribution. The paper's main theoretical contribution appears to be the observation that optimizing for task $X$ (achieving optimal point prediction) does not inherently lead to solving subtask $Y$ (learning the distribution over the latent adjacency matrix). This outcome is somewhat expected, given that the problem formulation lacks any dependencies linking subtask $Y$ to task $X$. A more specific formulation that establishes an explicit connection between the two tasks might provide a stronger foundation for the contribution.

3. Scope of the experiment. The experiments are limited to a small synthetic dataset with a small number of nodes ($N = 12$), a fixed feature dimension ($d = 4$), and fixed graph sparsity. The paper does include an experiment with more nodes ($N = 116$); however, the authors note that beyond this number, the number of free parameters becomes prohibitive. A more extensive empirical section showing how training with the proposed loss scales with the number of nodes, the feature dimension, and the graph sparsity (of the ground truth) would strengthen the presentation and provide more insight into the limits and applications of the proposed loss. 

Relevant citations
Xingyue Pu, Tianyue Cao, Xiaoyun Zhang, Xiaowen Dong, and Siheng Chen. Learning to learn
graph topologies. In Advances in Neural Information Processing Systems, pages 4249–4262,
2021.
Ruoyu Li, Sheng Wang, Feiyun Zhu, and Junzhou Huang. Adaptive graph convolutional neural
networks. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018.

Antonio Ortega, Pascal Frossard, Jelena Kovacevic, José MF Moura, and Pierre Vandergheynst. 
Graph signal processing: Overview, challenges, and applications. Proceedings of the IEEE,
106(5):808–828, 2018.

Minor
 
1. (68) typo 
2. (128) Missing end of sentence 
3. (275) typo substituhte

Meta-review: The paper addresses problem of graph structure learning problem. It mentions limitations of current point-prediction methods which cannot guarantee the calibration of the distribution of adjacency matrix. A sampling-based optimization method using Maximum Mean Discrepancy (MMD) to address this issue. The paper presents a theoretical analysis and an empirical evaluation.  

Strengths: 
- theoretical foundation, 
- minimization of the joint distribution discrepancy for calibrating the graph structure, 
- interpretability potential. 

Weaknesses: 
- incorporation the graph structure is not clear, graph structure oversimplified and the methodological contribution appears not specifically dedicated to graphs,
- more modality should be taken into account, there is lacks of experiments on real datasets, lacks of baselines, lack of uncertainty analysis, 
- motivation can be improved, 
- some related works are missing.

During rebuttal, authors have addressed concerns from the reviewers in particular relatively to the motivation and the related work issues.  
During the discussion, reviewers have recognized the merits of the paper, but the issue on the lack of graph structure integration, real-world experiments and some contextualization.  
There was no consensus towards acceptance.  

I propose then rejection.

Decision: Reject