Learning Directed Graphical Models with Optimal Transport

Much as we deeply appreciate the reviewer's time, we are perplexed by the reviewer's difficulty in understanding the contribution of our paper, which we believe has been clearly elucidated in the very first lines of the abstract and introduction. Let us take this opportunity to clarify the research problem and contribution.

1. Contribution

We here address the fundamental issue of learning parameters in directed graphical models from incomplete data. Referring to the historical development of graphical learning (summarized in Figure 1), methods can be characterized towards two extremes:

(1) At one extreme, one must make restrictive assumptions about the tractability of the posterior distribution (e.g., EM), the structure (e.g., mean-field approximations) or the model class (e.g., using conjugate exponential families).

(2) At the other extreme, there are black-box inference methods, notably VI. Recall that VI is an inference method where learning is treated as a by-product and model parameters (denoted by $\theta$) are cast as global latent variables, along with the local latent variables (denoted by $z$). Given observed data $\mathcal{D}$, the posterior $p(z, \theta \vert \mathcal{D})$ is generally intractable, thus less straightforward in a complex graph without making mean-field assumptions. A common strategy is to marginalize over local latent variables, often done analytically and not always feasible. Another strategy is to collapse all the parameters into a single latent variable to apply amortized inference, which makes little use of the graphical structure.

It is well known that existing methods, including EM and VI are fundamentally based on likelihood maximization, which is equivalent to minimizing the KL divergence between the model and empirical data distribution. Our method indeed contrast OT with maximum likelihood by minimizing the Wasserstein distance between the model and data distribution, as opposed to KL divergence. This is what we meant by claiming that this is an alternative line of thinking that views parameter learning as an OT problem where the Wasserstein formulation follows suite. Given the above limitations of existing methods, we would like to contribute a versatile learning algorithm that can be straightforwardly applied to any graphical structure and variable types.

2. Methodology

To the comment "Isn't this result already in the literature in multiple places? (It is fairly straightforward to show.)":

A related problem formulation can be found in the WAE paper that presents the theoretical result for a simple graph of 1 observed node and 1 hidden node. We generalize this result to learning an arbitrary graph of multiple nodes and our formulation thus inherits some good properties, as discussed in Appendix D. While the theoretical result might be straightforward to some, our contribution also involves transforming it into a practical and scalable algorithm that has been empirically shown to work effectively on a wide range of problems. To the best of our knowledge, the generalization result for graphical learning has not been presented elsewhere.

3. Experimentation

To the comment "Any informative experimental setup should be chosen to expose interesting contrasts with baselines. The logic doesn't make sense.".

We find this comment quite confusing, since our experimental setup indeed aims to expose the contrasts with the baselines. Concretely, we evaluate the effectiveness of our method on both parameter recovery and downstream tasks.

Our evaluation has been conducted on a wide range of models with different types of latent variables (continuous and discrete) and for different types of data (texts, images, and time series). We begin with a simple and popular task of topic modeling where traditional methods like EM or SVI can solve. We then consider learning HMMs, which remains fairly challenging, with known optimization/inference algorithms (e.g., Baum-Welch algorithm) often too computationally costly to be used in practice. For learning discrete representations, despite the simple graph, the true generative function is unknown and it is often approximated with a deep neural network with a number of parameters that can scale up to millions. Solving those problems with EM or MAP would be both expensive and generally intractable.

We believe that the models used in our experiments are representative and demonstrate our frameworks' applicability well. With the versatility of our framework, it has great potential to be applied to more complex models.

We truly appreciate the reviewer's time and interest in our work. We address the reviewer's questions in the following.

Page 2 line 8, "where the data distribution is the source and the true model distribution is the target".

It is precisely between the data distribution $P_d(X)$ and the model distribution $P_{\theta}(X)$.

But VI is minimizing the KL divergence between the two posterior distributions, which is not the data distribution. These sentences are a bit confusing.

VI is indeed minimizing the KL divergence between the two posterior distributions. Referring to Eq. (10) (Appendix D, page 19), if $\mathcal{Q}$ contains all conditional probability distributions $\phi(Z \vert X)$, the VAE objective concurs with the negative marginal log-likelihood $-\mathbb{E}\_{P_{d}(X)} [\log P_{\theta}(X)]$. To make the KL term tractable, the vanilla VAE uses a standard normal $P(Z)$ and restricts the variational distribution to a class of Gaussian distributions. Consequently, VAE is minimising an upper bound on the negative log-likelihood or, equivalently, on the KL divergence between the data and model distribution. Our claims are concerned with this connection.

Application 4.1

The following table reports the mean absolute error ($\times 10^2$) of the estimated parameters w.r.t. the true parameters (best/second-best methods are highlighted in bold/italic). This additional result further demonstrates the consistent performance of our method. However, it can be seen that the evaluation in terms of distribution divergence/distance as reported in the main paper provides a more conclusive and numerically significant comparison of methods.

Application 4.2

Yes, our model can learn the transitional probability as well. Following the reviewer's comment, we conduct an additional experiment in which we create synthetic data with the hidden states now generated from the HMM described in Section 4.2. Here we treat transition probability $p$ as a learnable parameter in addition to Poisson rates $\lambda_{1:4}$.

We further impose on the parameters the uniform distributions where

• $\lambda_1 \sim U(10,20)$
• $\lambda_2 \sim U(30,40)$
• $\lambda_3 \sim U(50,60)$
• $\lambda_4 \sim U(80,90)$
• $p \sim U(0.80, 0.90)$

We randomly generate $200$ datasets of $50,000$ observations each. For each dataset, we train the models for $50$ iterations with learning rate of $0.05$ at $5$ different initializations, while keeping the same experimental configuration as reported in Appendix E.2.

We report mean error of the estimates of the parameters in the following.

• Poisson rates: | $\lambda_1$: $0.380 \pm 1.107$ | $\lambda_2$: $0.784 \pm 0.770$ | $\lambda_3$ $1.653 \pm 1.243$ | $\lambda_4$: $0.930 \pm 1.084$

• Transition probability $p$: $0.019 \pm 0.011$.

We additionally visualize the distribution of the estimates to show the alignment with the generative uniform distributions. The figure can be anonymously viewed here: https://anonymous.4open.science/r/OTP-7944/hmm_estimates.png

The code has also been added to our anonymous public repository.

We hope to have addressed your concerns. If there are any more questions, we would love to receive your feedback.

We thank the reviewer for the feedback. Below are our responses to the reviewer's question.

Summing over all the parent nodes seems very costly.

Our algorithm is not costly, because (1) we do not perform inference during training, rather directly learn the parameters; (2) we support mini-batch training by leveraging amortized optimisation. Empirically our training time is less than half of that of batch methods like EM. In our experiment with hidden Markov models specifically, the model converges quickly after 50 iterations across most synthetic datasets, given a relatively long sequence of $200$ steps. This is more significantly efficient than the forwards-backwards (or Baum-Welch) algorithm (a form of message passing).

How does the proposed gradient descent compare in terms of complexity with belief propagation?

While our algorithm OTP-DAG does share the same ``forwards-backwards'' nature with belief propagation (BP), it is worth noting that our method focuses on parameter estimation (i.e., learning) rather than inference.

In terms of complexity, the main difference is that OTP-DAG optimises the backward and forward distributions simultaneously via gradient descent, whereas standard BP computes messages recursively. The running time of message passing algorithms is generally exponential in the tree-width of the graph. Meanwhile, OTP-DAG is much less costly by characterising the local densities via the backward maps from the observed nodes to their parents, through which computation is localised and can be done in parallel.

We thank the reviewer for the feedback.

1. To the reviewer's comment: "BP serves, in a similar fashion to what authors propose, as one step for learning the parameters: it is the E step of the EM algorithm for graphical models."

Our algorithm is fundamentally different from BP.

First and most importantly, the focus of our work is on parameter estimation of directed graphical models, NOT inference.

For the parameter learning task, it is true that one could use EM and perform BP in the E step (perhaps at the expense of computational cost due to the recursive nature of the algorithm). With this being said, it should essentially be comparing OTP-DAG with EM, not with BP.

2. The methodological differences with EM and other existing methods have been described in the paper. Let us provide another high-level perspective to once again clarify our novelty.

It is seen that when estimating the parameters, existing algorithms involve doing inference in some manner, e.g., explicitly via VI (as in the name) or iteratively in the E step of EM (either done analytically or with BP).

This is what sets us apart. OTP-DAG does not perform inference. Through the OT formulation, our method directly learns the parameters. We do not explicitly compute the posterior marginal or conditional probabilities. Each backward distribution is constructed in the way that (1) satisfies the push-forward constraints and (2) we can sample from it to evaluate the reconstruction. Our algorithm thus avoids the computational complexity introduced by inference, which allows us to estimate the parameters more efficiently.

Finally, it is natural to ask how OTP-DAG can do inference. This is indeed an interesting question that requires future exploration. As big fans of graphical models, we hope that our approach could potentially set the stage for solving large-scale inference and discovery problems.