Paper: Variational Inference with Singularity-Free Planar Flows (ICLR-2024)

Reviewer comment (Novelty):
Planar flow, the focus of this work, is one of the earliest methods introduced in the normalizing flow literature, and numerous advanced normalizing flows such as affine coupling flows, neural spline flows, etc., have been developed, showcasing superior performance. Given this context, the contribution of this work appears quite limited. To justify its contribution, it would be imperative to demonstrate that the modifications introduced significantly enhance the performance of planar flow. Ideally, it should now be on par with, or outperform, the more recent flow methods.

Author response:
We agree that the more recent flow methods are powerful in specific tasks. However, the expressiveness and scalability of the these flow methods come with a price — they are more intricately designed and require extra attention on parameter tuning/training, which can be inefficient in solving some variational inference tasks. This inefficiency is more severe when they are run on CPU-only devices. 

While we agree that intricate models with fine-tuned parameters are powerful in specific tasks, we believe that simplicity can be a strength. Our objective was to provide a clear and accessible solution to general variational problems, where ease of implementation and computational efficiency are crucial, e.g., sampling from a customize distribution, posterior inference in Bayesian statistical regression, and low-dimensional VAE.

Reviewer comment (Experiments):
The empirical experiments conducted in the study lack a comprehensive comparison with a variety of recent flow methods. Moreover, the paper only delves into variational inference applications, omitting performance evaluations on density estimation and the quality of synthetic data generation. To solidify the work's standing and importance, an expansion of the experimental section to include these aspects is necessary.

Author response:
Thank you for your suggestion. We have conducted an extra comparison with a recent flow methods, neural spline flows (NSFs). The updated comparison reaffirms the parameter efficiency of the new planar flows. We have included the results of the NSFs experiment in Figure 6, and a table summarizing the parameter counts is now available in Appendix D in the updated version.

As indicated by our title, *Variational Inference with Singularity-Free Planar Flows*, our focus is solely on variational inference. One reason for this focus is that the inversion of the planar flow does not have a closed form, making it less attractive for density estimation and data generation.

Reviewer comment (Evaluation):
This work falls short in providing a thorough and rigorous exploration of the singularity issue. The connection between the presence of a singularity and its detrimental impact on the performance of planar flow remains unclear to readers. The paper asserts that the proposed approach enhances the stability of model training, setting the expectation for a detailed analysis.

Author response:
Thank you for your thorough comment. We should have explained more about the singularity issue. While we don't have a formal theoretical proof, we believe our approach is logically grounded. We would like to highlight the logical coherence in our method:

The training instability and suboptimal convergence arising from the singularity are predominantly numerical issues. The presence of an unstable gradient around the singularity renders certain distributions within the variational family practically unattainable. Consequently, this limitation restricts the expressiveness of planar flows. By eliminating the singularity, we can explore all distributions during training, which leads to better approximation in general. We have incorporated this argument in the updated version.

Reviewer comment (Experiments):
To meet the expectation of analyzing singularity, the paper should ideally include a comparative analysis of the gradient variance in planar flow, employing both the original and the new parameterization schemes. Additionally, a theoretical examination of the training dynamics in a deep planar flow setup would contribute to a more comprehensive understanding. The paper should also offer a conclusive answer to whether issues such as gradient explosion or vanishing are effectively circumvented with the proposed modifications. Without these elements, the paper's exploration of the singularity issue feels incomplete, leaving readers with unanswered questions and a lack of clarity on the significance of the proposed solution.

Author response:
We trace the gradients of both flows during the training. Unfortunately, we did not find enough robust observations that can be included as a solid evidence. The gradients behavior of the old planar flow is not necessarily noise in all cases. We suspect that some regions of the variational family are not explored during training due to the singularity, thereby limiting the overall quality of the approximation. 

We made every effort to design a control experiment for comparing the original and the new planar flows, where all confounding factors are appropriately controlled (e.g., the initial parameters and training data). Hence, the only difference between the two flows is the reparameterization. We hope that our control experiments, conducted with a large number of replicates, can convincingly demonstrate that the improvement is indeed the result of eliminating the singularity.

Reviewer comment (Experiments):
The experiments are all in relatively low dimensional space (the VAE example uss a latent dimension D=20). It is not clear if the benefit of the proposed method would extend to higher dimension problems.

Author response:
**A2**: We appreciate your insight into the potential limitations of low-dimensional experiments. We took your suggestion and run the VAE experiment with a higher dimension $D=40$. Similar results are observed. This consistency supports the robustness of our findings, suggesting that the observations are not merely circumstantial but hold in higher-dimensional spaces as well. We have included this experiment in Appendix G. 

Note that we did not run the IAF and NSF for $D=40$ as changing to a higher dimension necessitates a corresponding adjustment in the neural network size of IAF and NSF. Unfortunately, we did not have sufficient time to thoroughly test the neural network size to ensure comparable performance.

Reviewer comment (Reproducibility):
It seems that SNF has similar issue in its parameterization. How did you implement it in your experiments?

Author response:
**A3**: Yes, the SNF has similar constraints on its parameters, specifically, the diagonal entries of the two triangular matrices: $r_{ii} \tilde{r}_{ii} > -1$. We followed the implementation by its proposers, van den Berg et al. (2018) — applying $tanh()$ to all diagonal entries. We note that it is possible to adapt our patch for the Sylvester flow to improve its expressivity/trainability.

Reviewer comment (Experiments):
Figure 6 shows an interesting result. First, it is an unfair comparison between planar flows and IAF in terms of number of layers since each layer in IAF is more complicated and requires more parameters, and this makes the performance of planar flows even better. Second, it seems that the more powerful IAF does not do well in terms of posterior estimation in this case. Is it because the training is not long enough? Was the KL reported in the right plot the average over all test images? Can the author provide more direct posterior comparison (e.g., scatter plot of samples) to a ground truth MCMC run?

Author response:
**A4**: Thank you for raising this question. We should have explained more on the results of IAFs. First, we answer the concern about the training and metrics reported.
  - The training is long enough for all flows. 
  - The KL reported in Figure 6 is the difference between $\log p(x)$ and ELBO, which are averaged over all 10,000 test images. For each test images, we estimate with 500 importance samples. We made every effort to ensure the accuracy of the metrics we reported.
  - Unfortunately, we are not able to provide a more direct posterior comparison to a ground truth MCMC run since the true posteriors are unknown.

We have conducted an extra experiment with neural spline flows (NSFs). The updated comparison reaffirms the parameter efficiency of the new planar flows. 

*Reason why IAFs and NSFs do not perform well in our study*: While the IAFs and NSFs employed in this study have much larger numbers of parameters compared to the planar flows, it is essential to note that their sizes remain relatively small in the context of the neural network architecture. Also, training of such intricate models with deep layers may require extra attention. Consequently, the results do not reflect their optimal performance and are solely used for comparison in terms of parameter efficiency. We have incorporated this comment in the revised version (page 9).

Reviewer comment (Evaluation):
It seems that the advantage of the new parameterization decreases as the number layers increase, any explanations? Also, for planar flows, a large number of layers are often used in practice which makes the benefit of new parameterization a bit less attractive.

Author response:
**A5**: The advantage of the new parameterization does not always decrease as the number layers increases. For example, Figure 6 shows that the improvement is larger as the number of layers increases. In general, the expressiveness of the planar flow increases as the number of layers increases, which might cover the issue arose from the singularity. However, the unstableness makes it all unpredictable. This numerical issue might be more severe as the number of layers increases. Hence, we suggest to use the singularity-free planar flow to avoid any undesirable complications.

Reviewer comment (Experiments):
The particular choice of the $m(\cdot)$ function seems a bit arbitrary. As argued in the paper, all we need is $m(x) = \mathcal{O}(x)$ when $x \to 0$. There are many functions that satisfy this condition. To confirm that the issue of the original parameterization is indeed the singularity, the authors should pick a few other $m(\cdot)$ functions satisfying $m(x) = \mathcal{O}(x)$ and check if they share similar performance.

Author response:
**A1**: We regret that we did not make it clear how we chose the function $m(\cdot)$. The sufficient condition we stated was in the context of removing the singularity. There are other implicit conditions as well. The main role of $m(\cdot)$ is to squeeze the unconstrained dot product to the required region $(-1, \infty)$, which in turn reparameterizes the unconstrained $v'$ to the feasible values. 

We chose $m(x)=x$ if $x \ge 0$ and a simple continuously differentiable extension to $x<0$, so that a minimal reparameterization is imposed on the unconstrained $v'$. This choice aligns with our belief in simplicity and minimizes the computation required. 

Certainly, there are other more complicate functions satisfying all the conditions. However, they might just introduce unnecessary computation. We appreciate the reviewer for raising this question. We have modified the corresponding section to enhance clarity in the updated version (page 5).

Reviewer comment (Experiments):
One advantage of the planar flow in this paper is that its performance is comparable to the IAFs and SNFs but the planar flow has fewer parameters. One experiment to further strengthen this point is to compare the wall-clock running time of each flow.

Author response:
**A3**: Thank you for your valuable suggestion. The planar flow is indeed faster and requires less computation resource than other flows due to its simplicity, especially with the minimal reparameterization we proposed. Based on our experiments, the new planar flows are approximately faster by a few hours than IAFs and SNFs for completing a replicate of the VAE experiment. We did not report this in the paper since the wall-clock running time we collected is not robust enough due to factors beyond our control.

We also note that the wall-clock running time is not the only metric and does not fully demonstrate the strength of the planar flow. For example, the planar flow is also trained on GPUs (to facilitate a fair control experiment), while it only requires small GPU usage, hence low power consumption. Additionally, it is possible to train planar flows efficiently on CPU-only devices.

Reviewer comment (Experiments):
Have the authors generated images from the trained VAEs on the MNIST as a sanity check?

Author response:
**A4**: Thank you for pointing it out. We have performed the sanity check as you suggested. The images generated from the trained VAEs are reasonable, which indicates that our models are properly trained. We have included the generated images in Appendix F.

Reviewer comment (Experiments):
Have the authors tried other function forms of $m(\cdot)$ that satisfy the condition $m(x) = \mathcal{O}(x)$?

Author response:
**A1**: We regret that we did not make it clear how we chose the function $m(\cdot)$. The sufficient condition we stated was in the context of removing the singularity. There are other implicit conditions as well. The main role of $m(\cdot)$ is to squeeze the unconstrained dot product to the required region $(-1, \infty)$, which in turn reparameterizes the unconstrained $v'$ to the feasible values. 

We chose $m(x)=x$ if $x \ge 0$ and a simple continuously differentiable extension to $x<0$, so that a minimal reparameterization is imposed on the unconstrained $v'$. This choice aligns with our belief in simplicity and minimizes the computation required. 

Certainly, there are other more complicate functions satisfying all the conditions. However, they might just introduce unnecessary computation. We appreciate the reviewer for raising this question. We have modified the corresponding section to enhance clarity in the updated version (page 5).

Reviewer comment (Novelty):
The simplicity of the proposal might indicate that it is not necessarily quite intellectually rich. While the related works point to many papers that choose different parameterizations for the coupling functions in normalizing flows, changing a parameterization in the way done here may not be conceptually very different from proposing other transformations that avoid the original issue in planar flows.

Author response:
**A1**: We acknowledge that our approach is intentionally simple, and we understand the concern about its conceptual depth. In choosing simplicity, our intention was to provide a clear and accessible solution to general variational problems, particularly in scenarios where ease of implementation and computational efficiency are crucial.

Reviewer comment (Novelty):
The authors should include related works that lay the groundwork for flow-based variational inference in scientific computing, specifically:  
- *Boltzman Generators*, Frank Noé, Simon Olsson, Jonas Köhler, Hao Wu, 2019  
- *Flow-based generative models for Markov chain Monte Carlo in lattice field theory*, Michaels S. Albergo, Gurtej Kanwar, Phiala E. Shanahan, 2019.

Author response:
**A2**: Thank you for your valuable suggestion. We have diligently reviewed the papers you recommended, and incorporated references to these works in the updated version (page 2).

Reviewer comment (Novelty):
Can the authors explain if their patch is relevant to the generalization of planar flows, namely Sylvester flows?

Author response:
**A3**: The Sylvester flow is considered as the generalization of the planar flow since it has the same form, $f(z) = z + A h(Bz + b)$, but allows for matrix weights. Our patch to the planar flow does not change the dimension of the weight parameters, so it is different from the approach of the Sylvester flow. However, there are similar constraints on the parameters of the Sylvester flow, $r_{ii} \tilde{r}_{ii} > -1$. Hence, it is possible to adapt our patch for the Sylvester flow to improve its expressivity/trainability.

Reviewer comment (Experiments):
Are there any higher-dimensional experiments that the authors could run? The downside of low-dimensional experiments is that it is hard to determine if the observations are generic or circumstantial.

Author response:
**A5**: We appreciate your insight into the potential limitations of low-dimensional experiments. We took your suggestion and run the VAE experiment with a higher dimension $D=40$. Similar results are observed. This consistency supports the robustness of our findings, suggesting that the observations are not merely circumstantial but hold in higher-dimensional spaces as well. We have included this experiment in Appendix G. 

Note that we did not run the IAF and NSF for $D=40$ as changing to a higher dimension necessitates a corresponding adjustment in the neural network size of IAF and NSF. Unfortunately, we did not have sufficient time to thoroughly test the neural network size to ensure comparable performance.