Training Energy-Based Normalizing Flow with Score-Matching Objectives

Summary

The authors consider normalizing flows with linear transformations. Computing the determinant of their Jacobian is expensive, limiting their applicability when training via maximum likelihood. Therefore, the authors introduced an energy-based method to train normalizing flows with linear transformation, which circumvents the need of computing the determinant of the Jacobian. It only has to be computed once for inference, which leads to a significant speed-up in training. Compared to the maximum likelihood-based training, the authors achieve similar results while getting a significant speed-up.

Strengths

The idea of combining energy-based training with normalizing flows by phrasing the flow as an energy-based model is novel and interesting. It circumvents the need of having to compute the determinant of the Jacobian of some layers, which can save a lot of compute. It still needs to be computed for inference; however, since it is constant in the cases the authors considered, it has to be computed only once and can be stored for further use.

Weaknesses

My biggest concern is the relevance of this work. Normalizing flows typically use transformations that are designed in such a way that the Jacobian determinant is easy to compute. For that reason, linear transformations are not very popular, also because they are not very expressive. The claim of the authors, that they are used in Glow and subsequent flow models for images, is true, but the linear transformation was only applied across channels, and since the number of channels is much smaller than the total number of input dimensions, i.e. the number of all pixels across all channels, they are not very costly in terms of evaluating their Jacobian's determinant. The authors apply their method to MNIST and CIFAR10, roughly matching the performance of their baseline with their method. However, both perform significantly worse than Glow and related methods. Hence, it is not clear whether their method gives them any advantage over these methods, either in terms of performance or speed. There are also procedures to estimate the Jacobian's determinant to cut down cost, such as the Skilling-Hutchinson estimator used by residual flows. The authors should have compared their method with using such an estimator. Unfortunately, their method cannot be used for residual flows, as the determinant of the Jacobian cannot depend on the input to be treated as a normalization constant in an energy-based method.

Questions

* How does your method compare against Glow in terms of runtime? * How does using an estimator for the Jacobian's determinant compare against your method?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

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

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

Limitations were mentioned in the weaknesses section of the review.

Summary

Normalizing flows are generative models trained by maximizing the log-likelihood of the dataset given the model, which is expressed as the sum of the base probability of a transformed variable plus a series of log-Jacobian terms that keep track of the volume distortion factors. Evaluating the log-Jacobian is usually the main bottleneck of normalizing flow architectures, so much so that the most popular architectures are designed so as to keep the log-Jacobian computations tractable. This paper introduces the use of techniques from the energy-based models literature in order to more efficiently train normalizing flow architectures. In a normalizing flow architecture, it is common to have both elementwise non-linear layers, which give data dependent but tractable log-Jacobian terms, and linear layers whose log-Jacobian is data-independent but expensive to compute. The main idea introduced in this paper is to split the log-likelihood of a normalizing flow into the product of two terms: 1) an "energy function" that contains all the data-dependent terms and 2) a normalization constant that collects all the data-independent log-Jacobians of the linear layers. This mirrors the split in energy-based models, where the energy function is tractable but the normalization function is intractable and it is not directly evaluated. Consequently, the authors apply a series of training objectives developed in order to train energy-based models to this new setting, obtaining an efficient way of training normalizing flows with expensive linear layers.

Strengths

The main idea is clever and original. It follows from the realization that the intractable terms of the log-Jacobian are independent from the data. This offers an interesting and well-motivated reason for borrowing techniques from the energy-based literature. The experiment section is adequate, with experiments both on toy data and on real high-dimensional datasets. While the results are not spectacular, they certainly support the claim that the energy-based training can be used to train normalizing flows with a substantial speed-up and without a major loss of performance. The paper is well-written and the relevant literature is properly referenced.

Weaknesses

Major points: In general, I am not convinced that this approach solves a major problem in the current literature. In fact, most successful flow architectures are designed to have fast log-Jacobian evaluations. Moreover, both continuous flows and the more recent flow-matching models can be used to train architectures with intractable Jacobian terms. The experiments only compare with the normalizing flow baseline and therefore they do not provide evidence that this approach has a competitive advantage over continuous flows and diffusion-like models such as flow-matching models. Minor points: I find the presentation somehow difficult to read, with several important equations and concepts scattered all over the paper. In particular, I think it would be useful to the reader to explicitly discuss the score-matching loss associated with Eq.9, for example, by writing the concrete form of Eq.5 as applied to Eq.9 in the main text.

Questions

- Could you clarify in which situations this model should be preferred over continuous flows and score-based diffusion models?

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

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The limitations of the approach have not be extensively discussed. I do not see any direct negative social impact.

Summary

The authors use a connection between flows and EBMs to devise a new training method for flows via score-matching. Their training procedure improves efficiency via avoiding computation of the Jacobian determinants that are independant to values of the samples (i.e. only contribute to the normalizing constant of the PDF). Additionally they identify two methods for improving training of the flows via score-matching, namely (1) match-after-preprossing whereby the score-matching occurs to the pre-processed variables avoiding numerical instability from the pre-processing layers, and (2) using an exponential moving average for the weight updates. In their experiments, they show their approach significantly improves training efficiency.

Strengths

- The paper identifies that computing the unnormalized PDF of a flow is significantly cheaper than the normalized PDF for linear flow layers, O(D^2) instead of O(D^3), noting that computing the Jacobian determinant is not necessary for the unnormalized PDF. This contribution is novel and valuable. - The experiments of the paper are strong, and demonstrate the utility of their approach.

Weaknesses

For sample generation by inverting the flow, the inverse matrix for the linear flow layers will cost $D^3$ - this is not mentioned in the paper. This inversion may be performed once, and then re-used for sample generation but for large $D$ this could be prohibitive.

Questions

This is the first time I've seen imputation done with flows - is this a novel contribution from the paper? How would sampling work for large $D$?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

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

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

There is no discussion around inverting the flow when $D$ is large (see Weaknesses/Questions).

Summary

This paper proposes a new flow based approach to generative modeling called EBFlow. EBFlow uses training objectives from EBMs to reduce training cost for flow based methods unlike previous approaches that mostly achieved this via different architectures or biased estimation methods. The main insight for EBFlow is that the change of variables formula for a series of transformations for a NF is broken in to two parts: one containing only linear flows term and the other containing terms from non-linear transformations. The linear flows term is then interpreted as the normalizing constant and this interpretation aids in faster training strategies for flow based methods.

Strengths

The motivation for reducing training costs for NFs is discussed well and is a very relevant and difficult problem. Several studies have been conducted to address this and the authors here propose a nice and novel methodology to address this. The exploration and discussion of the method as well as the motivation is thorough. The empirical analysis considered is detailed.

Weaknesses

I am not sure how to interpret the empirical results. It is obvious from the plots in Fig 2 that the proposed method is much faster than ML based objective. However, the performance in all the tasks for the ML based objective is usually better or very close to EBFlow. I was curious why this is? Also, for the results reported in Tables 1 and 2, were the methods trained for the same wall clock time?

Questions

See above.

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

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes.

Summary

This paper has proposed an energy-based normalizing flow model, where the computaiton of Jacobin determinants for linear transformations can be skiped with score-matching objectives. This could enable deeper layers of linear transformation and make the flow model more efficient.

Strengths

This paper has proposed an interesting idea about the relationship between flow model and energy-based model. The flow model could be interpreted as an energy-based model with a tractable normalizaiton term and the proposed score-matching objective could significantly improve the efficiency.

Weaknesses

1. The experiments lack comparison with other energy-based flow model like [1][2][3] and recent flow models like [4][5][6]. 2. Is there any theoretic analysis for the complexity between different objective function? 3. The author misses citations for the very first deep-learning-based energy-based models like [7][8]. Writing weakness: 1. The background in section 2.2 could break up into two sections, i.e., energy-based model and score-matching model. 2. It is hard to see the difference between results in the first two rows in Figure 1. 3. In line 51, even the tranformation function in flow model is invertible, the author should specify that $g()$ is the inverse transformation function. 4. In line 51, using $g_i(\cdot;\theta)$ is misleading. Does different transformation functions share the same parameters? 5. I think the author intends to model $g(\theta)$. Therefore, it is better to avoid using $E(x;\theta)$ to avoid confusion as it is usually used for representing a parametric energy model. 6. For Table 2, the author could change the unit to batch/second for easier comparison. Reference: [1] Xie, Jianwen, et al. "A Tale of Two Latent Flows: Learning Latent Space Normalizing Flow with Short-run Langevin Flow for Approximate Inference." arXiv preprint arXiv:2301.09300 (2023). [2] Xie, Jianwen, et al. "A tale of two flows: Cooperative learning of langevin flow and normalizing flow toward energy-based model." arXiv preprint arXiv:2205.06924 (2022). [3] Nijkamp, Erik, et al. "Learning energy-based model with flow-based backbone by neural transport mcmc." arXiv preprint arXiv:2006.06897 2 (2020). [4] Chen, Ricky TQ, et al. "Residual flows for invertible generative modeling." Advances in Neural Information Processing Systems 32 (2019). [5] Maaløe, Lars, et al. "Biva: A very deep hierarchy of latent variables for generative modeling." Advances in neural information processing systems 32 (2019). [6] Vahdat, Arash, Karsten Kreis, and Jan Kautz. "Score-based generative modeling in latent space." Advances in Neural Information Processing Systems 34 (2021): 11287-11302. [7] Xie, Jianwen, et al. "A theory of generative convnet." International Conference on Machine Learning. PMLR, 2016. [8] Xie, Jianwen, et al. "Cooperative learning of energy-based model and latent variable model via mcmc teaching." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 32. No. 1. 2018.

Questions

1. Could the author explain in detail why the complexity of calcuating Jacobin determinants is $O(D^3 \cdot L)$ for non-linear transformation? Is it general in all kinds of non-linear functions? 2. What does the sentence in line 62-63 mean? 3. What is the FID score in generation tasks? 4. What is the distribution of $p_{u}$?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

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

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

This paper has propose an interesting view of taking flow model as an eneryg-based model with score matching objectives. However, the expriments lack comparison of many baseline models as listed in 'weakness' part and other eval metrics in generation. Additional, the author should rephrase some sentences and formulars to avoid misleading.

Reply to the authors

I'm still not completely convinced by the author's arguments, but I acknowledge their position. Given the mostly positive feedback from the over reviewers, I'm not opposed to accepting this paper and increase my score to a borderline accept.