We propose a Gaussian manifold variational auto-encoder (GM-VAE) whose latent space consists of a set of Gaussian distributions. It is known that the set of the univariate Gaussian distributions with the Fisher information metric form a hyperbolic space, which we call a Gaussian manifold. To learn the VAE endowed with the Gaussian manifolds, we propose a pseudo-Gaussian manifold normal distribution based on the Kullback-Leibler divergence, a local approximation of the squared Fisher-Rao distance, to define a density over the latent space. In experiments, we demonstrate the efficacy of GM-VAE on two different tasks: density estimation of image datasets and environment modeling in model-based reinforcement learning. GM-VAE outperforms the other variants of hyperbolic- and Euclidean-VAEs on density estimation tasks and shows competitive performance in model-based reinforcement learning. We observe that our model provides strong numerical stability, addressing a common limitation reported in previous hyperbolic-VAEs.
Paper
Similar papers
Peer review
Summary
The paper proposes a pseudo-Gaussian distribution over hyperbolic spaces, which is suitable for constructing VAEs with hyperbolic latent spaces. This pseudo-Gaussian distribution arrives by replacing the usual Euclidean distance with a KL-divergence (this locally approximates the hyperbolic geodesic distance). The required properties of this distribution (sampling, normalization constant, KL-divergences) are derived and shown to be easily computable. Empirical results show promise.
Strengths
1. The paper is clearly written and easy to follow. 2. The paper provides a novel (to the best of my knowledge) construction of a hyperbolic distribution that seemingly solves a practical problem (numerical instability) in hyperbolic VAEs. 3. I find it elegant how the proposed distribution factorizes into a product of known distributions (this was not obvious to me from the KL construction).
Weaknesses
1. A claimed benefit of the proposed method is that it is numerically more stable than previous methods. While I do not doubt the correctness of this statement, I had wished for an empirical demonstration of the claim. As far as I can tell, numerical stability is the main practical gain over existing methods (note, this is an important gain), so I think it should be studied more extensively experimentally. In short: show me numbers that measure stability. 2. The paper mentions tuning the curvature of the hyperbolic metric, but I did not find it clear how and when such tuning takes place. Is the $c$ parameter tuned by hand? 3. Similarly, the paper mentions using a $\beta$ parameter in the VAE, but it is not clear how such a parameter is chosen in practice. Is it manually tuned? If so, does that not suggest overfitting could be a problem? 4. Since the geodesic distance is replaced by a local approximation (in the form of the KL divergence), the method is only somewhat geometric. It seems to be good enough, but I suppose that none of the established theoretical results on tree embeddings hold under the approximated metric. If so, it would be good to explicate. *Minor comment:* Line 126: The KL divergence is really not a "distance" and should not be described as such. Especially in a geometric paper, such confusions should be avoided.
Questions
1. The KL is a local approximation to the hyperbolic metric. This is used as a distance measure in a pseudo-Gaussian distribution. I wonder how similar this construction is to a tangent space normal distribution. Both approaches rely on local approximations of the metric. Do you have comments on this link? 2. In Table 1, what does $\circ$ and $\times$ mean? I'm struggling to make sense of those symbols.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
4 excellent
Presentation
3 good
Contribution
3 good
Limitations
The paper does not provide much of a discussion of the limitations of the proposed method.
Summary
This paper proposes to utilize a set of Gaussian distributions as the latent space of a VAE model. One benefit of doing so is that this forms a hyperbolic space under the Fisher information metric. A key insight is that this Gaussian manifold corresponds to the Poincare half-plane model of hyperbolic geometry. Using this construction the authors construct a Gaussian manifold VAE which locally approximates the metric and allows for fast calculation of the KL divergence and fast sampling. Empirical results demonstrate the efficacy of the proposed approach on image datasets as well as RL benchmarks.
Strengths
This paper proposes an interesting way to enforce a hyperbolic latent space in a VAE setting. The connection of the FIM and Poincare half-plane using a set of univariate Gaussian distributions is interesting. The definition of the PGM distribution is certainly novel and I especially enjoyed the construction as it has the benefits of exact KL computation while enabling fast sampling. Finally, the overall flow and presentation of the paper at a high level are clear, but certain aspects could be improved to increase readability.
Weaknesses
There are a few weaknesses in the paper which I outline below: **Novelty:** 1. While the use of the PGM distribution is novel in a VAE setting the use of the Poincare Half plane model due to its enhanced numerical stability properties is not. In fact, this paper does not cite a key work in this space "Representing Hyperbolic Space Accurately using Multi-Component Floats" Tao and De Sa 2021. I believe their usage of geometry shares all the same benefits of numerical stability (and perhaps even more) than the one presented here. I think this paper should attempt to implement a VAE with the Poincare Half-plane latent space as proposed in this paper for a fair comparison between approaches. 2. The authors comment that current hyperbolic normal distributions suffer from a variety of drawbacks such as numerical instability, lack of a closed form KL, and high computational cost in sampling. However, these criticisms seem mostly towards the Poincare and Riemannian Normal distribution. The wrapped normal does not have high sampling costs for instance, so the statement should be refined a bit. Furthermore, other hyperbolic distributions can be defined on higher dimensional hyperbolic spaces without resorting to product spaces. This is a limitation of the Poincare half-plane model itself, so I feel there should be more commentary about this and how it may affect experiments. **Presentation weaknesses:** 1. The preliminaries section could be significantly improved. The introduction to Hyperbolic geometry is non-standard and the authors could lean on prior work to improve the exposition. 2. The definition of the Poincare half-plane model is insufficient here. You need to be much more concrete in mathematical terms on how it is defined. Things that one could add is what is the exact form of standard operations using this geometry, e.g. distance, exp map, log map, etc ... Furthermore, a visualization could improve readability. **Experimental weaknesses:** 1. The experiments section has a few weaknesses. First, it seems that the bolded numbers in the tables have std's that overlap with other table entries. In this case, both table entries should be bolded. 2. I believe it is important to test this approach on actual tree-like datasets rather than image datasets. These include the Branching Diffusion Process dataset found in Nagano et. al 2019, as well as the lobster graphs dataset in Bose et. al 2020. A hyperbolic VAE should show more visual benefits in these settings. 3. Please include numerical error analysis in the same vein as Tao and De Sa 2021 Fig 1 and Fig 2 to substantiate the claims that the proposed approach has less numerical instability. **Minor:** 1. The authors say "the tree-structured data can be embedded with arbitrary low distortion in hyperbolic space". This fact is only true for 2D hyperbolic spaces. The cited paper does not extend this to higher dimensions
Questions
1. Past hyperbolic VAE's have made the curvature constant learnable which leads to improvements. Can this also be leveraged here? 2. In Table 1. the authors mark the HWN as not geometry aware. Can the authors explain why this is the case? The HWN explicitly uses the metric in all of its computations. 3. Similarly, the HWN has been used in VAE settings, in which the case KL divergence is needed and is not difficult to calculate. Can the authors explain why they feel this is not the case?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.
Soundness
3 good
Presentation
2 fair
Contribution
3 good
Limitations
Yes.
Summary
In this paper, the authors propose a novel Gaussian manifold generative autoencoder (GM-VAE) whose latent space is a statistical manifold formed by univariate Gaussian distributions. They also introduce a pseudo-Gaussian manifold normal distribution on the latent, which is easy to sample from and improves training stability due to a closed form of KL divergence. The experiments conducted show that GM-VAE outperforms the considered state-of-the-art baselines in the density estimation task on a few image datasets, and gives comparable results to DreamerV2 in the model-based RL task.
Strengths
1. The proposed GM-VAE model is a non-trivial improvement of the VAE architecture with a solid theoretical background. 2. Experimental results on the CUB, Food101, Oxford 102, and (partially) Atari2600 Brekout datasets demonstrate the superiority of GM-VAE over its competitors ($\mathcal{E}$-VAE, $\mathcal{L}$-VAE, and $\mathcal{P}$-VAE). 3. The proposed model has superior computational performance concerning the other hyperbolic VAEs considered. 4. The paper is well written.
Weaknesses
1. The experimental results of the GM-VAE in the model-based RL task are only comparable to those of competitors. 2. Rather small-scale image datasets were used in experiments.
Questions
1. It would be interesting (and thus influence my opinion) to see the results of experiments performed on large-scale datasets (like CelebA). Also, the FID score would be reported (besides NLL). 2. Minor comments: Tab. 2 caption: N/A appears only in the standard deviation (or perhaps confidence interval?) indices.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
2 fair
Presentation
3 good
Contribution
2 fair
Limitations
Generally yes, but for clarity I would suggest summarizing the limitations in a separate paragraph.
Summary
Some VAE methods have been proposed to model the latent space as a hyperbolic space for modeling the structured representation of data. This paper proposes a hyperbolic distribution by using KL divergence to measure the geodesic distribution of Poincare normal, by using KL divergence as the local approximation of the geodesic distance. It makes the proposed method easy to sample and improve the numerical stability. Experiments on image density estimation and model-based RL are conducted.
Strengths
+ The proposed method is with a clear motivation and relies on a simple assumption (in Eq. (2)) to make the analytic operations, which is straightforward. + Experiments are conducted on two different applications. + As shown in the reported results, the proposed method can perform better than or comparable with the related works.
Weaknesses
- It is not clearly studied and discussed how the proposed pseudo Gaussian manifold normal distribution can model the hyperbolic distribution. Specifically, it is unclear how the introduced approximation may influence the modeling of the manifold. Although tractable KL is favorable, how the approximation introduced by the proposed method may influence the modeling capability? - The authors may visualize the manifold of the proposed method and compare that with other methods, such as [27]. - The sampling results of the VAE models should be visualized and compared with previous methods, such as [27] and [25], especially considering that numerical improvements shown in tables are insignificant. - As the claimed main benefit of the proposed method is numerical stability, the analysis and validation of the stability should be strengthened. How and why the previous works may introduce instability? How can the stability of the methods be measured and validated with empirical studies, numerical evaluation, and visualization?
Questions
Please check other questions and suggestions pointed out with the weakness points.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
3 good
Presentation
3 good
Contribution
2 fair
Limitations
The limitations are not explicitly discussed. The authors may check the weak points and consider discussing the potential limitation of the proposed method, such as the potential influence on the modeling capability that may be caused by the approximation.
Thank you for the rebuttal
I will take the provided update into account when discussing the paper with the other reviewers.
Response to reviewer Darj
We appreciate responding to the rebuttal and paying attention to the contents for the further discussion. Feel free to ask any questions in the remaining discussion period!
Thank you for the response
Thank you for the effort in preparing the rebuttal. I am open to updating my rating depending on the conclusions of the Reviewers-AC discussion phase.
Response to reviewer ACLv
We thank for responding to the rebuttal and being open to updating the score. Feel free to ask any questions in the remaining discussion period!
Re:Rebuttal
I thank the authors for their rebuttal. I appreciate the new experiments provided in the 1-page PDF. Q1. The numerical stability of the proposed PGM model is not in doubt. The new Fig 1 shows this and should be included in the paper. > Firstly, the computation of the normalizing constant and density function is necessary For Wrapped distribution in the Lorentz model, this is known to us. Did you mean using a wrapped distribution in the PGM model? Q2. > HWN suffers from numerical instability due to the log map of the Lorentz model This is true. In practice, it is customary to clamp to an absolute value of $\approx 40$ as $\approx e^{40}$ is within floating point precision. Can you see a degradation in performance after clamping in your experiments for the rebuttal? > Secondly, the sampling speed is slow. As shown in Table 2, the sampling speed of HWN is 8.39x slower than the PGM normal due to additional mappings between the Euclidean space and the Lorentz model. What do you mean here? The projection is trivial. You sample a point in $$\mathbb{R}^{n-1}$$ and prepend a $0$ which makes it a point on the tangent space at the identity which is on the Lorentz model. After that, you can parallel transport and exp map to the manifold at a different point. Are you using this projection? Because this is not expensive. > This also affects the speed of the KL divergence computation which is approximated with Monte Carlo sampling. This I don't buy is slow. In general, this is super fast. An analytic KL in this case is not a big benefit for two reasons 1.) The main benefits of hyperbolic space are in low dimensional latents and for this is not a bottleneck 2.) In non-Gaussian VAEs, we use Monte Carlo all the time with no problems. Q3. I appreciate the answer but my original concern still remains. At the moment I am still unsure and need the authors to clarify my questions above. Based on these clarifications I may adjust my score.
Re:Re:Rebuttal
We appreciate participating in the discussion. We address the remaining concerns in the following response. --- **Q1. Further details for the HWN on the Poincare half-plane model.** We thank for the suggestion of quantitative analysis on numerical stability which strengthens our work. We will add Figure 1 in the future revision. We would like to clarify our understanding about the original review first. We thought that [1] shows that *the Poincare half-plane model* is a good choice for numerical stability *because* it is implemented on the Poincare half-plane model. Hence, the other baseline distributions, e.g, HWN and Poincare normal, *should be implemented on the Poincare half-plane model* for fair comparison on numerical stability. And this is why we answered “the computation of the normalizing constant and density function is necessary”. If this is the suggested direction of the original review (correct us if we misunderstood), we need to think about the details of the implementation. Let’s consider the HWN in the Poincare half-plane model. The first way to implement HWN on the half-plane is by transporting the sample of the HWN defined on the Lorentz model via isometry. However, the MCF of [1] is not validated on the Lorentz model and the isometries. The transformation also requires additional computation of the density through the change of variables. The second way is to re-define HWN on the Poincare half-plane model as done in [2, 4], as we mentioned in the previous rebuttal. The computation of density function and normalizing constant should be newly done in this case because the equations of the Riemannian operations change. --- **Q2. Results on clipped HWN.** We conduct an experiment on the effect of clipping in HWN. Before moving the encoder output to the Lorentz model to decide the mean value of the HWN, we clip the norm of the encoder output with the suggested bound. In some settings, there is a slight degradation but the confidence intervals overlap. ||CUB|Food101|Oxford102| |---|---|---|---| |latent dimension|||| |50|$991.37_{\pm 2.83}$ | $1298.80_{\pm 8.74}$|$1298.75_{\pm 1.71}$| |60|$971.80_{\pm 2.83}$|$1228.73_{\pm 6.81}$|$1255.94_{\pm 2.27}$| |70|$954.00_{\pm 3.63}$|$1165.21_{\pm 4.57}$|$1232.40_{\pm 2.31}$| --- **Q3. Sampling speed of HWN.** We would like to emphasize the computation of MC-based KL is *relatively* slower than our closed-form solution, but in practice, as you suggested, this would not be a great problem. Here, we wanted to highlight the empirical difference between two methods but tone the presentation down in a future revision. We provide a detailed analysis on the sampling speed of the HWN and PGM normal. Both the HWN and PGM normal starts with sampling from the Euclidean normal. The HWN then transports the Euclidean sample to the Lorentz model via the projection mentioned in the re-response. The PGM normal finishes the sampling by concatenating the Euclidean sample with the Gamma distribution sample. So the difference between the two sampling procedures is the projection and the Gamma distribution sampling. Although there is no obvious difference between the two methods, the empirical results reveal that the Gamma distribution sampling is faster than the projection. However, in practice, we agree that the gap would not be noticeable during the training. We finally summarize our findings on the training speed of the HWN. Although the additional projection is just a combination of arithmetic operations and seems to have less effect on training VAE, we observe that there exists *a relative difference in the sampling speed* between PGM normal. Furthermore, while hyperbolic VAEs benefit more in low latent dimensions, the batch size is another important factor that affects the training speed. The below table shows that the training time does not vary on the latent dimension, while the relative difference is preserved (the absolute values differ from the paper due to different GPU). ||$\mathcal{E}$-VAE|$\mathcal{L}$-VAE|$\mathcal{P}$-VAE|GM-VAE| |---|---|---|---|---| |latent dimension|||| |2|16.78|27.54|36.55|17.39| |4|16.42|28.74|38.15|17.55| |8|16.75|28.52|38.58|17.30|
Thanks for the response. Can the authors produce the visualization of the manifold similar to (and in comparison with) Figure 8 in [27]?? As shown in Figure 2 in the response, the structure of the proposed method seems more distorted than other methods, which is not a good sign.
Response to reviewer 1SSZ
We thank for the response and additional feedback. We add the visualization of PGM normal at Figure 3 of our one-page PDF supplementary. Due to the page limit, we only differed the $\alpha, \beta$ value with fixed $\gamma$. We will add various visualizations of PGM normal in the future revision as suggested by the response. We apologize for occuring confusion due to the different visualization manifold. The latents of GM-VAE show unique behavior among the hyperbolic VAEs which looks distorted. We expect such behavior appears because of the different manifold, i.e., the Poincare half-plane model, but not because of the quality of embedding as reviewer concerns. We will add the visualization of the latents in the Poincare disk model as the other hyperbolic VAEs in the future revision. We conjecture that the visualization would be improved in a more intuitive manner, because the generation performance is better than $\mathcal{L}$-VAE as shown in Table 1.
Decision
Accept (poster)