Optimal Transport for Probabilistic Circuits

Thank you for your feedback. We have revised the paper based on your review, including: rewriting Algorithm 1 and the following description to be more clear, clarifying the parallels between EM and our algorithm WM, and adding details about future work. Thank you for the pointers to related works. We have revised the introduction to discuss related work regarding the TV distance.

… Equation (4): Why is the second equality correct?

The second equality in Equation 4 makes use of how the empirical distribution \hat{Q} is defined; we have updated the notation of this section and clarified this in the paper.

Along with the above changes, we have also clarified some sections and modified some formatting according to your comments. Thank you again for your feedback.

We hope that we have addressed your comments in the latest revision of our paper. As the revision period is coming to an end soon, please let us know if you have any unaddressed questions or suggestions for us to improve the paper.

Thank you for your feedback on the paper; we have made some revisions, and will follow up with another revision addressing the remaining feedback shortly.

Q1. Could you provide additional clarification on the proof of Theorem 1 and 2? I think it's better to add some graph in the proof.

We will include figures in the appendix to clarify the proofs of Theorems 1 and 2 shortly.

Q2. Have you evaluated CWp across a broader variety of PCs, as its application might be limited?

We are working on experiments that utilize optimal transport maps between PCs for color transfer between images to showcase an interesting and practical application of our algorithm, as well as computing the optimal transport distance between larger circuits learned on high-dimensional image data.

Q3. Could you elaborate on the computational complexity of your approach?

We compute CW2 and the associated transport map in $\mathcal{O}(mn)$-time, where $m$ and $n$ are the number of edges in each original circuit (see the last paragraph of Section 3.2 for more details). We show that computing the circuit parameters that minimize the empirical Wasserstein distance is NP-hard (see Theorem 2), but our proposed iterative algorithm where each step runs in $O(n)$-time in the number of circuit edges is guaranteed to converge to a local minimum (see the second-to-last paragraph of Section 4.1 for more details).

This paper addresses only a restricted set of cases within the broader context of probabilistic circuits.

Within the framework of probabilistic circuits, the tractability of certain queries is guaranteed by imposing constraints on the circuit structure (smoothness, decomposability, and compatibility are required in our case). We also note that the constraints we impose for our algorithm are the same constraints required by all existing tractable algorithms for pairwise queries (such as KL-divergence (Vergari et al., 2021)) between PCs. Crucially, enforcing these structural properties do not restrict the PCs’ expressivity (i.e., they can still represent any distribution), but may limit their expressive efficiency (i.e., the circuit may need to be exponentially large). Our algorithm is thus applicable to any two PCs - although we incur a possibly exponential increase in the size of the circuits by making them compatible. We appreciate your comment, and have noted this in the paper.

Additionally to the changes mentioned above, we have fixed the typo on line 786. Thank you again for your feedback.

We hope that we have addressed your comments in the latest revision of our paper. As the revision period is coming to an end soon, please let us know if you have any unaddressed questions or suggestions for us to improve the paper.

Thank you for your feedback on the paper; we have updated the existing example of a coupling circuit and will be adding example figures for some of the proofs shortly.

Q1. Why wasn’t CW compared with W?

To the best of our knowledge, there are no existing algorithms that can be run for even very small circuits in a reasonable amount of time that exactly compute W. Thus, such a comparison is unfortunately quite difficult.

Q2. How much error is considered acceptable?

To verify the applicability of our approach to real-world problems, we are working on experiments that utilize optimal transport maps between PCs for color transfer between images to showcase a practical application of our algorithm.

Using simple examples to illustrate definitions and results could make the paper easier to read and follow.

We have updated Figure 1 to more clearly explain a coupling circuit.

The proposed method performs well with a small set of variables; however, runtime challenges typically arise in large-scale systems with many variables.

We first note that Appendix C.1 contains experimental results for computing the transport map between circuits that are two orders of magnitude larger than those mentioned in the body of the paper, which were limited in size to be able to compute MW2. We are also working on experiments that compute the optimal transport distance between two circuits learned on high-dimensional image data.

While the proposed metrics could significantly reduce runtime, they also lead to an increase in error. How much error is considered acceptable? There is no analytical approach or numerical result provided to show the impact of this error.

In Section 5.2, we perform experiments to quantify the error between CW2 and MW2 as the ratio between the two quantities. Unfortunately, we are unable to include more data points in these figures, as computing MW2 becomes impractical for circuits larger than those plotted. We are open to suggestions on additional numerical results that we could provide to showcase the gap between MW2 and CW2.

Thank you for your valuable feedback on the paper; we appreciate the time and effort you put into reviewing our work.

Q1. The notion of compatibility seems really quite strong. Are there any ways you can see of weakening it…? It would be particularly nice if the method were to degrade naturally to an intractable (exponentially-sized) problem in the worst case when the two PCs have very different structures…

We would like to first clarify that compatibility does not require that two circuits have identical structure; it only requires that two corresponding product nodes with same scopes decompose the scopes in the same way into children (up to a bijection in our case). In other words, two compatible circuits have the same hierarchical scope partitioning, but they can have different structures. Furthermore, compatibility between two circuits is necessary for all pairwise queries (such as KL-divergence or joint entropy (Vergari et al., 2021)) with known tractable algorithms for PCs as far as we can tell. Crucially, enforcing these structural properties do not restrict the PCs’ expressivity (i.e., they can still represent any distribution), but may limit their expressive efficiency (i.e., the circuit may need to be exponentially large). Thus, our algorithm can already be applied to arbitrary pairs of PCs with very different structures, although we incur a possibly exponential increase in the size of the circuits by making them compatible. We appreciate your comment, and have clarified this in Section 3.1 and Figure 1 in the paper.

Q2. …W1 and W2 are the most commonly-used in practice … if the authors think that it might be possible to find a tractable algorithm for W1 or W2, then I feel strongly that it would be better for the research community to delay publication until this question is resolved.

While W1 and W2 are the most common in practice, they are generally approximated. Moreover, even though there are various approximation methods for W2 between GMMs, the question of whether exactly computing W1 or W2 between GMMs (which are a special case of PCs) is NP-hard is still an open question to the best of our knowledge. Thus, we consider this result outside the scope of this work–introducing the first algorithm to compute and optimize Wasserstein-type distances for PCs–and respectfully disagree that publication of this work should be delayed until we can resolve this question.

Q3. … In what contexts are MW used? Has it been experimentally validated? Might you be able to show a computational benefit in prior tasks by drawing from that literature?

Past work has experimentally validated MW and associated transport maps when applied to the task of color transfer (Delon & Desolneux, 2020). We are currently working on applying our optimal transport approach to the task of color transfer, which will provide a direct comparison between CW and MW for a real-world application. This will complement our existing results showing the computational benefit of CW over MW between synthetic PCs. We hope to follow up soon with the results.

Q4. …significance of [the Figure 4] visualization…

We include figure 4 to a) demonstrate that we can easily get the transport plan from our CW algorithm and b) show that the transport plan matches our intuition for what the optimal transport plan should be, despite it being an approximation. However, in the interest of space in the body of the paper, we have moved the figure to the appendix.

Q5. I assume that the points in Figure 3 represent pairs randomly sampled circuits. Are they all with the same architecture, or do only pairs have the same architecture? Why are there so few points? Figure 3 suggests to me that your method may work well on circuits of larger depths. Have you tried to quantify this across a broader set of distributions over {P,Q}, or theoretically show that this must be the case? These might be interesting avenues to explore.”

We apologize for the confusion regarding our experiment setup, and have clarified in the paper that each point represents an average over 100 randomly-initialized circuits with a fixed dimensionality and sum node branching factor. We have also provided justification on why our approach seems to work well for circuits with a higher depth; in short, since we only plot the number of circuit edges rather than the number of learnable parameters in the circuit (which are different since product node edges are unweighted), it is possible for a higher-depth circuit to have more edges but less learnable parameters than a lower-depth circuit and thus incur less error. Lastly, we are unable to include more data points in these figures, as computing MW2 becomes impractical for circuits larger than those plotted.

Due to the OpenReview character limit, we answer the remaining comments in our second reply.

We answer the remaining comments below.

on the experimental side, the examples are all small-scale synthetic toys.

We first note that experimental comparisons for distance computation in the main paper (Figs 2 & 3) were done on smaller circuits because the baseline (MW) could not scale to larger circuits. Appendix C.1 contains experimental results for computing the transport map between circuits that are two orders of magnitude larger, albeit still synthetic. Moreover, the parameter learning experiments trained circuits with up to millions of parameters using MNIST benchmark dataset. Nevertheless, we agree that additional experiments with real-world applications would be valuable. We are currently exploring utilizing optimal transport maps between PCs for color transfer between images, as well as computing the optimal transport distance between two circuits learned on high-dimensional image data.

The authors correctly identify that their metric is closely related to EM, and compare against it as baseline. However, their method does significantly worse.

While our approach achieves worse likelihoods than EM (which explicitly maximizes likelihood), we note that our approach outperforms EM with regard to learning circuits with a lower Wasserstein distance to the empirical data distribution (experimental results in Appendix C.2). We are currently investigating a slight variation of our current algorithm that appears to yield circuits with similar likelihoods to our current approach but significantly lower Wasserstein distances to the empirical data distribution, which will be included in the paper shortly.

Lastly, thank you for your detailed comments about definitions and notations. We mostly followed the standard definitions for PCs (Choi et al., 2020) but agree that they could be made more precise and have revised the paper accordingly.