Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The authors propose an interpretation of the VAE latent space as a Riemannian manifold and leverage this structure to design an improved sampling strategy. To construct a Riemannian manifold, the authors interpret the covariance of the posterior of each latent as a local Riemannian metric. The local metrics are interpolated to produce the latent manifold. To generate samples, the authors use HMC on the manifold specific uniform distribution constructed using local metrics.
main_review: #### Strengths
- The stated contributions of the paper are substantiated by the experiments.
- The proposed work is competitive and often superior to related methods on several datasets.

#### Weaknesses
- The use of non-standard English makes reading somewhat difficult.
- I would have liked to see some discussion of the cost incurred from HMC sampling
- I would have liked to see some discussion of the relationship of the proposed approach to [1-2].

[1] Rakowski, Alexander, and Christoph Lippert. "Disentanglement and Local Directions of Variance." Joint European Conference on Machine Learning and Knowledge Discovery in Databases. Springer, Cham, 2021.

[2] Rolinek, Michal, Dominik Zietlow, and Georg Martius. "Variational autoencoders pursue pca directions (by accident)." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2019.

#### Possible typos:
- “...using either other regularization…” → using other regularization (p.6)
- “...4 well known databases” → “four well known datasets” (p.6)
- “visually speaking” → “qualitatively” (p.6)
- “...same models and data sets…” → same models and datasets (p.8)
- “... quite robust to the data set” → quite robust to the dataset (p.6)
- "According to us, these metrics" --> These metrics (p.9)
- "A d-dimensional manifold... d-dimensional Euclidean" --> A $d$-dimensional manifold... $d$-dimensional Euclidean (p.3)
- There are parts of figure 1 that aren’t fully described in the caption (e.g., part (a) and part (b) in top/bottom)

summary_of_the_review: The authors introduce a novel sampling procedure for the classical VAE that produces quantitatively and qualitatively strong results. The sampling approach leverages a novel theoretically-based interpretation of the VAE latent space structure. The paper itself is intuition building (e.g., toy example) but difficult to read in some places. It also appears the authors may have neglected a related body of work.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper introduces a geometric perspective on VAEs trained with a Gaussian mixture prior (standard multivariate or learned mixture of Gaussian): the covariance matrix encoded for a given input image in a VAE corresponds to a Riemannian metric (inverse covariance) around the local neighborhood of the data point encoded.
The authors suggest sample generation in VAE  should be based on this covariance-induced distance (Mahalanobis). To obtain a single continuous metric at the end of the training, the authors suggest an interpolation based on available data and then generating data using a Monte Carlo (HMC) sampler based on this interpolated metric. 
Experiments using samples generated using this method shows improvements in FID and precision-recall measures in MNIST, SVHN, CIFAR10, and CelebA datasets. The authors claim that the proposed sample generation produces good samples even with datasets of small size.

main_review: The paper provides a way of incorporating geometry into the VAE setup for generating samples that belong to this manifold. While geometry-aware sampling strategies are available in prior literature, the authors claim that these processes are time-consuming and complex. To better position this statement, I think the runtime and complexity of the proposed method need to be discussed and substantiated.

I assume the authors do not change the training process and are only proposing a better sample generation strategy - it is not clear from the current text if this is the case and if so, it is unclear what training procedure they are resorting to. The authors review Riemannian geometry concepts that are relevant to the method proposed and it is easy to follow through this part of the paper.
The toy example shows the advantage of using the proposed sample generation procedure in comparison to vanilla VAE sample generation - the proposed method generates relevant samples that maintain the simple geometric structure as in the input data.
Experiments show improvements over RAE’s, VAMP, and RHVAE in terms of FID score and classification with a VAE generated set of data. The performance of the proposed method in a low data regime looks interesting and would be interesting to have a discussion on why the proposed method is able to achieve this. 

My major concern is the variability of the samples generated. I wonder if the samples generated in real dataset experiments are biased towards resampling available input data. If \rho variable used for the calculation of the metric in equation (7) is very large then the sampling process is biased towards just picking the closest input data point. I think the authors need to better support the claim on the variability of generated samples - interpolation scenarios in real images, random samples (not clear if the generated samples shown are random or cherry-picked), or even reconstructions for a given input (this would matter if the proposed method is modifying the training process as well). From the current set of visualizations and results, it is not clear if the images generated are new or just resampled input data.

Other concerns (in no particular order) are listed below.
- There is no clarity on the modification of the loss function (if any) used to train the VAE model. I assume the model is learning a mean vector and a covariance matrix for each input image - but are there any additional constraints to ensure that the matrix learned is a valid covariance matrix? More critically, what if the covariance matrix is not invertible? How are samples generated during training?

- There is no discussion on hyperparameters (\lambda, \rho) and the sensitivity of the experiments to these parameters. The supplementary material says rho is assigned by a max-min distance formulation - is this based on the entire dataset or a subset? Any reason why this is the right choice? 

- In the OASIS dataset experiment, the generated dataset is made class balanced and evaluated. But the “original” classifier is trained with an unbalanced set - Is it possible to evaluate by resampling or reweighting the original training set to make them class balanced? I understand this introduces bias in the classifier but so does the experimental setup proposed. 

- Any reason for dropping the RHVAE evaluation in Table 2?

- I'm not sure FID comparison with methods using prior is a fair setting - I believe it gives methods with “ex-post density estimate” an unfair advantage in reducing the gap between the posterior and the prior. This is by no means specific to this paper but I would be interested in knowing the author’s perspective on this.

The ideas and experiments by [1,2] are relevant to the proposed method and the current paper can be improved by incorporating and differentiating with respect to these recent works.

[1] Kalatzis, Dimitrios, et al. "Variational Autoencoders with Riemannian Brownian Motion Priors." International Conference on Machine Learning. PMLR, 2020.

[2] Connor, Marissa, Gregory Canal, and Christopher Rozell. "Variational Autoencoder with Learned Latent Structure." International Conference on Artificial Intelligence and Statistics. PMLR, 2021.
summary_of_the_review: Overall, I believe the paper needs revision to better support the claims in the paper. The idea of incorporating geometry in the VAE generation procedure is interesting and has potential. I believe the authors need to improve clarity on experimental setup and evaluation, and better support the variability of samples generated. The current form of paper has a large room for improvement and I hope the authors are able to incorporate the concerns raised and better position their work.

----
Update: I appreciate the authors for answering the concerns raised and am increasing my recommendation to 6. The additional experiments and details on the algorithm in the appendix improve the clarity of the paper. The proposed sampling strategy shows improvement (FID metric) in conjunction with a simple VAE model. It might be interesting to see if the proposed sampling can be used to improve other VAE models. The experimental setup with OASIS is not as compelling (my suggestion of resampling is only a simple strategy and I believe there are better class-balancing classification methods), but highlights the shortcoming of using earlier VAE methods in this setup.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: It is a theoretical paper that discusses the Riemannian structure of the latent space of a VAE.
The associated measure on a point \mu(x) derives from \Sigma^-1(x), which are the parameters of the posterior p(z|x) or simply the outputs of the encoder.
The measure allows to define a Riemannian normal distribution on that space/manifold, to compute geodesics but also to "better" sample.
main_review: 
The idea is to equip the latent space with a Riemannian structure, the measure of which depends on the inverse of covariance of the posterior.
Before sampling, the distribution needs to be approximated using MCMC.
The results are ok.

Pros: 
* The article is well written and smoothly walks the reader through the mathematics behind variational autoencoders (VAE) and Riemannian geometry. 
* The idea is clear and straightforward. The theory is robust. However, I am doubtful of the practical outcomes. 

Cons: 
* The analyses present several shortcomings.
* The baselines are chosen accordingly to the theoretical line of the paper. That is to improve the samples' quality without touching (too much) the loss nor the architecture.
That is fair, although I was expecting some Energy-based approaches such as VAEBM (Xiao et al. 2020) or Pang et al. (which is in the references but not cited)
* The authors overemphasize sometimes. 
** "Now that the VAE has learned and capture the intrinsic geometry of the data within the latent space seen as a Riemannian manifold," 
**    Why is the variance-based measure necessarily the one that captures the "intrinsic geometry of the data"?
    Besides, what if the data is made of several connected components? Such a situation is not allowed by the theory developed here.
**  VAE+GMM is almost as good as the proposed method in Table 2.

summary_of_the_review: The document is good, easy to read, and fits within the scope of the conference.
The Riemannian point of view on VAE surely brings interesting insights that might lay the foundation of other works. 
Nevertheless, I have doubts about the importance of the contributions and their impact on the field. 
Moreover, the analysis is pretty basic. Maybe because there is too much to say/study for a conference paper.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes to analyse the VAE framework using differential geometry tools. More specifically, they propose to impose a (Riemannian) metric in the latent space of a VAE. Such a Riemanian metric then enables data generation based on the (Riemanian) Uniform distribution on the learned latent space, improving the quality of the generated samples.

main_review: This work proposes a novel interpretation of variational encoding by the prism of differential geometry.  In particular the data $x_i$ provides two main elements, the “position” $\mu(x_i)$ and instead of interpreting the estimated $\Sigma(x_i)$ as the variance of the normal law i.e of p(z|x) it is interpreted as the value of the Riemanian metric at $\mu(x_i)$. In that perspective, the authors propose to sample not from an isotropic Gaussian of $\mathbb{R}^d$, but rather using uniformly from the learned manifold.
The experiments are interesting and well-conducted.

Remarks and Question:
1. On the interpretation of the Geometric interpretation of the VAE. Section 4.3 is a little confusing. Indeed, the initial assumption that the inverse variance corresponds to the local metric tensor is acceptable (the lower the variance the less the sample is allowed to move from the mean). However, the rest of the section is in my opinion confusing. Indeed, because the authors seem to suggest that the original-VAE is in fact an approximation of a Riemannian VAE, can the authors link the Riemannian VAE penalties and how they are approximated in the original VAE framework? For clarity, I highly recommend including such discussion in the final version.
2. Can the authors comment on the link between their work and the desire of the VAE-community to rather enrich the expressivity of p(z|x) using normalizing flows or Hamiltonian scheme [1, 2]
3. I can’t relate eq.7 to the interpretation of the inverse variance as a local Riemannian metric (i.e $G(\mu(x_i) = \Sigma^{-1}(x_i)$). Indeed as of eq.7 $G(\mu(x_i)) \neq \Sigma^{-1}$, depending on how $\mu(x_i)$ is close to the other $\mu_i$.
4. In the same line can the author comment on the need for the regularization term $+ \lambda I_d$ ? How to choose $\lambda$ in practice ?

On the Experiments:
1. What is the influence of the number of selected points in the sampling based on 7). For example can the author provide a FID curve somehow like Fig.4 with varying number of “support vectors” $(\mu_i, \Sigma^{-1}_i)$ ? Indeed, the more complex the latent space geometry, the more “support-vectors” it may require for an acceptable decoding.
2. Because the sampling method in the latent space is dependent on the chosen support vectors, did the authors observe any variance in the quality of the samples depending on the choice of the K-centroids (something like variance of the FID at a defined number of support vectors)?
3. On the Decoding processes : One drawback of such sampling scheme (due to the choice of the centroids) is the fact that in the decoding processes, the decoder simply provides an interpolation between the decoded K-support vectors. Did the authors observe such an issue in practice ?

I would be glad to increase my score if the major concerns are answered.

[1]: Hamiltonian Variational Auto-Encoder,  Anthony L. Caterini, Arnaud Doucet, Dino Sejdinovic  

[2]: Variational Inference with Normalizing Flows, Danilo Jimenez Rezende, Shakir Mohamed 

summary_of_the_review: The geometric interpretation of the VAE is interesting. Moreover, the author's proposition for resampling seems sound in the light of their proposition and the experiments are well conducted. Yet, the theoretical link between the VAE and the Riemannian distance based VAE requires a more thorough discussion.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes to generalize the vanilla VAE by formulating the conditional distribution and the marginal distribution using the Riemannian geometry. The authors proposed a generalized perspective on various VAEs proposed so far. The paper is rather easy to follow, however, it is unclear what precisely is used in the experiments and how the proposed approach works.
main_review: Strengths:
- The proposed approach, i.e., using the Riemannian geometry to formulate distributions in the VAE, is interesting.
- The paper proposes a general perspective on VAEs through a geometrical perspective. It generalizes some papers on VAEs with various latent manifolds (e.g., hyperspherical VAE, VAEs with Poincare disk).

Weaknesses:
- Remark 1 is not necessarily true for two reasons. First, this is not true that $p_{\theta}(\mathbf{x}|\mathbf{z})$ is often taken as a multivariate Gaussian. The conditional distribution must be chosen according to the problem at hand. If data is binary, the Bernoulli distribution is taken; if data is continuous, we can pick from a plethora of distributions. Second, a very common choice to model RGB images is a discretized logistic distribution rather than a Gaussian distribution.
- The example in Sec. 4.5 is a bit vague. First, the setup is not well explained. Second, Figure 1 is not properly explained: there are three times (a), and one (b), but they are neither referred to in the caption nor in the text. Moreover, the font used for the top (a) and top (b) are so small that it is hard to read.
- The paper explains the idea, however, it lacks a more concrete statement of the proposed approach. The authors mention that they use Eq. (7) as the Riemannian metric and the HMC sampler for sampling, however, it is not necessarily obvious what they precisely use later on in the experiments. It would be highly beneficial to provide a specific instantiation of their approach, e.g., how the metric is calculated, what is the final objective, what is the sampling procedure, what is the training algorithm. At the moment, the paper presents the ideas at a rather general level.


summary_of_the_review: I like the paper and the idea of using Riemannian geometry in VAEs. However, the paper presents the idea at a pretty general level and it is hard to see the details that matter from the computational perspective. I doubt it would be easy and straightforward to code the proposed approach up. This is my main problem with the paper, namely, the reproducibility. And the follow-up problem is about the computational burden of the proposed approach. Since it is unclear how to sample and how to calculate various components of the VAE, it is hard to assess how “heavy” the method is and whether it could be scaled up. As we see nowadays, models that are easy to scale up (e.g., hierarchical VAEs like NVAE or diffusion-based models) allow achieving SOTA performance.

=== UPDATE ===
I appreciate the rebuttal. Overall, I still doubt whether the paper is a strong contribution. Nevertheless, the new additions (especially Appendix C) make the paper better. Therefore, I decided to increase my score to 6.
review_title: null

Meta-review: The paper proposes to use covariance of the approximate posterior to induce a metric on the latent space of the VAE and use it to sample from the Riemannian manifold learned by a VAE. Experiments on MNIST and CelebA show the method outperforms vanilla VAE in terms of sample quality (FID and PR scores). It is also shown to work better than baseline VAE models on a medical imaging classification task. While the reviewers have acknowledged the contributions of the paper, the novelty in the contributions and their importance/impact was seen to be rather limited. The main concern from the reviewers is -- while the paper is mainly based on the use of inverse covariance as the metric for manifold, it doesn't give a reasonable theoretical justification on it is a sensible metric that captures the intrinsic geometry of data. Authors in their response justify it as -- since the covariance matrices are learned from the data and favor through the posterior sampling some direction in the latent space, it is a natural choice as metric. This is not very convincing. A more technical justification for this will certainly make the paper more convincing. I suggest the authors to look at "Kumar, Abhishek, and Ben Poole. "On Implicit Regularization in $ β $-VAEs." International Conference on Machine Learning. PMLR, 2020" which theoretically connects inverse covariance and the Riemannian metric in Sec 5.2, and see if it can be adapted in their context.
Decision: Reject