SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chen-Hao Chao $^{1}$ , Wei-Fang Sun $^{1,2}$ , Yen-Chang Hsu $^{3}$ , Zsolt Kira $^{4}$ , and Chun-Yi Lee $^{1*}$
<sup>1</sup> Elsa Lab, National Tsing Hua University, Hsinchu City, Taiwan
$^{2}$  NVIDIA AI Technology Center, NVIDIA Corporation, Santa Clara, CA, USA
$^{3}$  Samsung Research America, Mountain View, CA, USA
<sup>4</sup> Georgia Institute of Technology, Atlanta, GA, USA
In this paper, we establish a connection between the parameterization of flow-based and energy-based generative models, and present a new flow-based modeling approach called energy-based normalizing flow (EBFlow). We demonstrate that by optimizing EBFlow with score-matching objectives, the computation of Jacobian determinants for linear transformations can be entirely bypassed. This feature enables the use of arbitrary linear layers in the construction of flow-based models without increasing the computational time complexity of each training iteration from  $\mathcal{O}(D^2 L)$  to  $\mathcal{O}(D^3 L)$  for an  $L$ -layered model that accepts  $D$ -dimensional inputs. This makes the training of EBFlow more efficient than the commonly-adopted maximum likelihood training method. In addition to the reduction in runtime, we enhance the training stability and empirical performance of EBFlow through a number of techniques developed based on our analysis of the score-matching methods. The experimental results demonstrate that our approach achieves a significant speedup compared to maximum likelihood estimation while outperforming prior methods with a noticeable margin in terms of negative log-likelihood (NLL).
Parameter estimation for probability density functions (pdf) has been a major interest in the research fields of machine learning and statistics. Given a  $D$ -dimensional random data vector  $\mathbf{x} \in \mathbb{R}^D$ , the goal of such a task is to estimate the true pdf  $p_{\mathbf{x}}(\cdot)$  of  $\mathbf{x}$  with a function  $p(\cdot; \theta)$  parameterized by  $\theta$ . In the studies of unsupervised learning, flow-based modeling methods (e.g., [1-4]) are commonly-adopted for estimating  $p_{\mathbf{x}}$  due to their expressiveness and broad applicability in generative tasks.
Flow-based models represent  $p(\cdot ;\theta)$  using a sequence of invertible transformations based on the change of variable theorem, through which the intermediate unnormalized densities are re-normalized by multiplying the Jacobian determinant associated with each transformation. In maximum likelihood estimation, however, the explicit computation of the normalizing term may pose computational challenges for model architectures that use linear transformations, such as convolutions [4, 5] and fully-connected layers [6, 7]. To address this issue, several methods have been proposed in the recent literature, which includes constructing linear transformations with special structures [8-12] and exploiting special optimization processes [7]. Despite their success in reducing the training complexity, these methods either require additional constraints on the linear transformations or biased estimation on the gradients of the objective.
Motivated by the limitations of the previous studies, this paper introduces an approach that reinterprets flow-based models as energy-based models [13], and leverages score-matching methods [14-17] to
*Corresponding author. Email: cylee@cs.nthu.edu.tw
optimize  $p(\cdot ;\theta)$  according to the Fisher divergence [14, 18] between  $p_{\mathbf{x}}(\cdot)$  and  $p(\cdot ;\theta)$ . The proposed method avoids the computation of the Jacobian determinants of linear layers during training, and reduces the asymptotic computational complexity of each training iteration from  $\mathcal{O}(D^3 L)$  to  $\mathcal{O}(D^2 L)$  for an  $L$ -layered model. Our experimental results demonstrate that this approach significantly improves the training efficiency as compared to maximum likelihood training. In addition, we investigate a theoretical property of Fisher divergence with respect to latent variables, and propose a Match-after-Preprocessing (MaP) technique to enhance the training stability of score-matching methods. Finally, our comparison on the MNIST dataset [19] reveals that the proposed method exhibits significant improvements in comparison to our baseline methods presented in [17] and [7] in terms of negative log likelihood (NLL).
In this section, we discuss the parameterization of probability density functions in flow-based and energy-based modeling methods, and offer a number of commonly-used training methods for them.
2.1 Flow-based Models
Flow-based models describe  $p_{\mathbf{x}}(\cdot)$  using a prior distribution  $p_{\mathbf{u}}(\cdot)$  of a latent variable  $\mathbf{u} \in \mathbb{R}^D$  and an invertible function  $g = g_L \circ \dots \circ g_1$ , where  $g_i(\cdot; \theta): \mathbb{R}^D \to \mathbb{R}^D, \forall i \in \{1, \dots, L\}$  and is usually modeled as a neural network with  $L$  layers. Based on the change of variable theorem and the distributive property of the determinant operation  $\operatorname*{det}(\cdot), p(\cdot; \theta)$  can be described as follows:
$$
p (\boldsymbol {x}; \theta) = p _ {\mathbf {u}} (g (\boldsymbol {x}; \theta)) | \det  (\mathbf {J} _ {g} (\boldsymbol {x}; \theta)) | = p _ {\mathbf {u}} (g (\boldsymbol {x}; \theta)) \prod_ {i = 1} ^ {L} | \det  (\mathbf {J} _ {g _ {i}} (\boldsymbol {x} _ {i - 1}; \theta)) |, \tag {1}
$$
where  $\pmb{x}_i = g_i\circ \dots \circ g_1(\pmb {x};\theta)$ ,  $\pmb{x}_0 = \pmb{x}$ ,  $\mathbf{J}_g(\pmb {x};\theta) = \frac{\partial}{\partial\pmb{x}} g(\pmb {x};\theta)$  represents the Jacobian of  $g$  with respect to  $\pmb{x}$ , and  $\mathbf{J}_{g_i}(\pmb{x}_{i - 1};\theta) = \frac{\partial}{\partial\pmb{x}_{i - 1}} g_i(\pmb{x}_{i - 1};\theta)$  represents the Jacobian of the  $i$ -th layer of  $g$  with respect to  $\pmb{x}_{i - 1}$ . This work concentrates on model architectures employing linear flows [20] to design the function  $g$ . These model architectures primarily utilize linear transformations to extract crucial feature representations, while also accommodating non-linear transformations that enable efficient Jacobian determinant computation. Specifically, let  $\mathcal{S}_l$  be the set of linear transformations in  $g$ , and  $\mathcal{S}_n = \{g_i\mid i\in \{1,\dots ,L\} \} \setminus \mathcal{S}_l$  be the set of non-linear transformations. The general assumption of these model architectures is that  $\prod_{i = 1}^{L}|\det (\mathbf{J}_{g_i})|$  in Eq. (1) can be decomposed as  $\prod_{g_i\in S_n}|\det (\mathbf{J}_{g_i})|\prod_{g_i\in S_l}|\det (\mathbf{J}_{g_i})|$ , where  $\prod_{g_i\in S_n}|\det (\mathbf{J}_{g_i})|$  and  $\prod_{g_i\in S_l}|\det (\mathbf{J}_{g_i})|$  can be calculated within the complexity of  $\mathcal{O}(D^2 L)$  and  $\mathcal{O}(D^3 L)$ , respectively. Previous implementations of such model architectures include Generative Flows (Glow) [4], Neural Spline Flows (NSF) [5], and the independent component analysis (ICA) models presented in [6, 7].
Given the parameterization of  $p(\cdot ;\theta)$ , a commonly used approach for optimizing  $\theta$  is maximum likelihood (ML) estimation, which involves minimizing the Kullback-Leibler (KL) divergence  $\mathbb{D}_{\mathrm{KL}}[p_{\mathbf{x}}(\boldsymbol {x})\| p(\boldsymbol {x};\theta)] = \mathbb{E}_{p_{\mathbf{x}}(\boldsymbol {x})}\left[\log \frac{p_{\mathbf{x}}(\boldsymbol{x})}{p(\boldsymbol{x};\theta)}\right]$  between the true density  $p_{\mathbf{x}}(\boldsymbol {x})$  and the parameterized density  $p(\boldsymbol {x};\theta)$ . The ML objective  $\mathcal{L}_{\mathrm{ML}}(\theta)$  is derived by removing the constant term  $\mathbb{E}_{p_{\mathbf{x}}(\boldsymbol {x})}[\log p_{\mathbf{x}}(\boldsymbol {x})]$  with respect to  $\theta$  from  $\mathbb{D}_{\mathrm{KL}}[p_{\mathbf{x}}(\boldsymbol {x})\| p(\boldsymbol {x};\theta)]$ , and can be expressed as follows:
$$
\mathcal {L} _ {\mathrm {M L}} (\theta) = \mathbb {E} _ {p _ {\mathbf {x}} (\boldsymbol {x})} [ - \log p (\boldsymbol {x}; \theta) ]. \tag {2}
$$
The ML objective explicitly evaluates  $p(\pmb{x}; \theta)$ , which involves the calculation of the Jacobian determinant of the layers in  $S_l$ . This indicates that certain model architectures containing convolutional [4, 5] or fully-connected layers [6, 7] may encounter training inefficiency due to the  $\mathcal{O}(D^3 L)$  cost of evaluating  $\prod_{g_i \in S_l} |\det(\mathbf{J}_{g_i})|$ . Although a number of alternative methods discussed in Section 3 can be adopted to reduce their computational cost, they either require additional constraints on the linear transformation or biased estimation on the gradients of the ML objective.
2.2 Energy-based Models
Energy-based models are formulated based on a Boltzmann distribution, which is expressed as the ratio of an unnormalized density function to an input-independent normalizing constant. Specif-
ically, given a scalar-valued energy function  $E(\cdot ;\theta):\mathbb{R}^{D}\to \mathbb{R}$ , the unnormalized density function is defined as  $\exp (-E(\pmb {x};\theta))$ , and the normalizing constant  $Z(\theta)$  is defined as the integration  $\int_{\pmb{x}\in \mathbb{R}^D}\exp (-E(\pmb {x};\theta))dx$ . The parameterization of  $p(\cdot ;\theta)$  is presented in the following equation:
$$
p (\boldsymbol {x}; \theta) = \exp (- E (\boldsymbol {x}; \theta)) Z ^ {- 1} (\theta). \tag {3}
$$
Optimizing  $p(\cdot ;\theta)$  in Eq. (3) through directly evaluating  $\mathcal{L}_{\mathrm{ML}}$  in Eq. (2) is computationally infeasible, since the computation requires explicitly calculating the intractable normalizing constant  $Z(\theta)$ . To address this issue, a widely-used technique [13] is to reformulate  $\frac{\partial}{\partial\theta}\mathcal{L}_{\mathrm{ML}}(\theta)$  as its sampling-based variant  $\frac{\partial}{\partial\theta}\mathcal{L}_{\mathrm{SML}}(\theta)$ , which is expressed as follows:
$$
\mathcal {L} _ {\mathrm {S M L}} (\theta) = \mathbb {E} _ {p _ {\mathbf {x}} (\boldsymbol {x})} [ E (\boldsymbol {x}; \theta) ] - \mathbb {E} _ {\operatorname {s g} (p (\boldsymbol {x}; \theta))} [ E (\boldsymbol {x}; \theta) ], \tag {4}
$$
where  $\mathrm{sg}(\cdot)$  indicates the stop-gradient operator. Despite the fact that Eq. (4) prevents the calculation of  $Z(\theta)$ , sampling from  $p(\cdot; \theta)$  typically requires running a Markov Chain Monte Carlo (MCMC) process (e.g., [21, 22]) until convergence, which can still be computationally expensive as it involves evaluating the gradients of the energy function numerous times. Although several approaches [23, 24] were proposed to mitigate the high computational costs involved in performing an MCMC process, these approaches make use of approximations, which often cause training instabilities in high-dimensional contexts [25].
Another line of researches proposed to optimize  $p(\cdot ;\theta)$  through minimizing the Fisher divergence  $\mathbb{D}_{\mathrm{F}}[p_{\mathbf{x}}(\boldsymbol {x})||p(\boldsymbol {x};\theta)] = \mathbb{E}_{p_{\mathbf{x}}(\boldsymbol {x})}\left[\frac{1}{2}\| \frac{\partial}{\partial\boldsymbol{x}}\log \left(\frac{p_{\mathbf{x}}(\boldsymbol{x})}{p(\boldsymbol{x};\theta)}\right)\| ^2\right]$  between  $p_{\mathbf{x}}(\boldsymbol {x})$  and  $p(\boldsymbol {x};\theta)$  using the score-matching (SM) objective  $\mathcal{L}_{\mathrm{SM}}(\theta) = \mathbb{E}_{p_{\mathbf{x}}(\boldsymbol {x})}\left[\frac{1}{2}\| \frac{\partial}{\partial\boldsymbol{x}} E(\boldsymbol {x};\theta)\| ^2 -\mathrm{Tr}(\frac{\partial^2}{\partial x^2} E(\boldsymbol {x};\theta))\right]$  [14] to avoid the explicit calculation of  $Z(\theta)$  as well as the sampling process required in Eq. (4). Several computationally efficient variants of  $\mathcal{L}_{\mathrm{SM}}$ , including sliced score matching (SSM) [16], finite difference sliced score matching (FDSSM) [17], and denoising score matching (DSM) [15], have been proposed.
SSM is derived directly based on  $\mathcal{L}_{\mathrm{SM}}$  with an unbiased Hutchinson's trace estimator [26]. Given a random projection vector  $\mathbf{v} \in \mathbb{R}^D$  drawn from  $p_{\mathbf{v}}$  and satisfying  $\mathbb{E}_{p_{\mathbf{v}}(\mathbf{v})}[v^T v] = I$ , the objective function denoted as  $\mathcal{L}_{\mathrm{SSM}}$ , is defined as follows:
$$
\mathcal {L} _ {\mathrm {S S M}} (\theta) = \frac {1}{2} \mathbb {E} _ {p _ {\mathbf {x}} (\boldsymbol {x})} \left[ \left\| \frac {\partial E (\boldsymbol {x} ; \theta)}{\partial \boldsymbol {x}} \right\| ^ {2} \right] - \mathbb {E} _ {p _ {\mathbf {x}} (\boldsymbol {x}) p _ {\mathbf {v}} (\boldsymbol {v})} \left[ \boldsymbol {v} ^ {T} \frac {\partial^ {2} E (\boldsymbol {x} ; \theta)}{\partial \boldsymbol {x} ^ {2}} \boldsymbol {v} \right]. \tag {5}
$$
FDSSM is a parallelizable variant of  $\mathcal{L}_{\mathrm{SSM}}$  that adopts the finite difference method [27] to approximate the gradient operations in the objective. Given a uniformly distributed random vector  $\varepsilon$ , it accelerates the calculation by simultaneously forward passing  $E(\pmb{x};\theta)$ ,  $E(\pmb{x} + \varepsilon; \theta)$ , and  $E(\pmb{x} - \varepsilon; \theta)$  as follows:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {F D S S M}} (\theta) = 2 \mathbb {E} _ {p _ {\mathbf {x}} (\boldsymbol {x})} [ E (\boldsymbol {x}; \theta) ] - \mathbb {E} _ {p _ {\mathbf {x}} (\boldsymbol {x}) p _ {\xi} (\boldsymbol {\varepsilon})} [ E (\boldsymbol {x} + \boldsymbol {\varepsilon}; \theta) + E (\boldsymbol {x} - \boldsymbol {\varepsilon}; \theta) ] \\ + \frac {1}{8} \mathbb {E} _ {p _ {\mathbf {x}} (\boldsymbol {x}) p _ {\xi} (\varepsilon)} \left[ (E (\boldsymbol {x} + \varepsilon ; \theta) - E (\boldsymbol {x} - \varepsilon ; \theta)) ^ {2} \right], \tag {6} \\ \end{array}
$$
where  $p_{\xi}(\varepsilon) = \mathcal{U}(\varepsilon \in \mathbb{R}^{D}|\| \varepsilon \| = \xi)$ , and  $\xi$  is a hyper-parameter that usually assumes a small value.
DSM approximates the true pdf through a surrogate that is constructed using the Parzen density estimator  $p_{\sigma}(\tilde{\boldsymbol{x}})$  [28]. The approximated target  $p_{\sigma}(\tilde{\boldsymbol{x}}) = \int_{\boldsymbol{x} \in \mathbb{R}^D} p_{\sigma}(\tilde{\boldsymbol{x}} | \boldsymbol{x}) p_{\boldsymbol{x}}(\boldsymbol{x}) d\boldsymbol{x}$  is defined based on an isotropic Gaussian kernel  $p_{\sigma}(\tilde{\boldsymbol{x}} | \boldsymbol{x}) = \mathcal{N}(\tilde{\boldsymbol{x}} | \boldsymbol{x}, \sigma^2 \boldsymbol{I})$  with a variance  $\sigma^2$ . The objective  $\mathcal{L}_{\mathrm{DSM}}$ , which excludes the Hessian term in  $\mathcal{L}_{\mathrm{SSM}}$ , is written as follows:
$$
\mathcal {L} _ {\mathrm {D S M}} (\theta) = \mathbb {E} _ {p _ {\boldsymbol {x}} (\boldsymbol {x}) p _ {\sigma} (\tilde {\boldsymbol {x}} | \boldsymbol {x})} \left[ \frac {1}{2} \left\| \frac {\partial E (\tilde {\boldsymbol {x}} ; \theta)}{\partial \tilde {\boldsymbol {x}}} + \frac {\boldsymbol {x} - \tilde {\boldsymbol {x}}}{\sigma^ {2}} \right\| ^ {2} \right]. \tag {7}
$$
To conclude,  $\mathcal{L}_{\mathrm{SSM}}$  is an unbiased objective that satisfies  $\frac{\partial}{\partial\theta}\mathcal{L}_{\mathrm{SSM}}(\theta) = \frac{\partial}{\partial\theta}\mathbb{D}_{\mathrm{F}}[p_{\mathbf{x}}(\boldsymbol {x})\| p(\boldsymbol {x};\theta)]$  [16], while  $\mathcal{L}_{\mathrm{FDSSM}}$  and  $\mathcal{L}_{\mathrm{DSM}}$  require careful selection of hyper-parameters  $\xi$  and  $\sigma$ , since  $\frac{\partial}{\partial\theta}\mathcal{L}_{\mathrm{FDSSM}}(\theta) = \frac{\partial}{\partial\theta} (\mathbb{D}_{\mathrm{F}}[p_{\mathbf{x}}(\boldsymbol {x})\| p(\boldsymbol {x};\theta)] + o(\xi))$  [17] contains an approximation error  $o(\xi)$ , and  $p_{\sigma}$  in  $\frac{\partial}{\partial\theta}\mathcal{L}_{\mathrm{DSM}}(\theta) = \frac{\partial}{\partial\theta}\mathbb{D}_{\mathrm{F}}[p_{\sigma}(\tilde{\boldsymbol{x}})\| p(\tilde{\boldsymbol{x}};\theta)]$  may bear resemblance to  $p_{\mathbf{x}}$  only for small  $\sigma$  [15].
3.1 Accelerating Maximum Likelihood Training of Flow-based Models
A key focus in the field of flow-based modeling is to reduce the computational expense associated with evaluating the ML objective [7-12, 29]. These acceleration methods can be classified into two categories based on their underlying mechanisms.
Specially Designed Linear Transformations. A majority of the existing works [8-12, 29] have attempted to accelerate the computation of Jacobian determinants in the ML objective by exploiting linear transformations with special structures. For example, the authors in [8] proposed to constrain the weights in linear layers as lower triangular matrices to speed up training. The authors in [9, 10] proposed to adopt convolutional layers with masked kernels to accelerate the computation of Jacobian determinants. The authors in [29] leveraged orthogonal transformations to bypass the direct computation of Jacobian determinants. More recently, the authors in [12] proposed to utilize linear operations with special butterfly structures [30] to reduce the cost of calculating the determinants. Although these techniques avoid the  $\mathcal{O}(D^3 L)$  computation, they impose restrictions on the learnable transformations, which potentially limits their capacity to capture complex feature representations, as discussed in [7, 31, 32]. Our experimental findings presented in Appendix A.5 support this concept, demonstrating that flow-based models with unconstrained linear layers outperform those with linear layers restricted by lower / upper triangular weight matrices [8] or those using lower-upper (LU) decomposition [4].
Specially Designed Optimization Process. To address the aforementioned restrictions, a recent study [7] proposed the relative gradient method for optimizing flow-based models with arbitrary linear transformations. In this method, the gradients of the ML objective are converted into their relative gradients by multiplying themselves with  $\mathbf{W}^T\mathbf{W}$ , where  $\mathbf{W} \in \mathbb{R}^{D\times D}$  represents the weight matrix in a linear transformation. Since  $\frac{\partial}{\partial\mathbf{W}}\log |\det (\mathbf{W})|\mathbf{W}^T\mathbf{W} = \mathbf{W}$ , evaluating relative gradients is more computationally efficient than calculating the standard gradients according to  $\frac{\partial}{\partial\mathbf{W}}\log |\det (\mathbf{W})| = (\mathbf{W}^T)^{-1}$ . While this method reduces the training time complexity from  $\mathcal{O}(D^3 L)$  to  $\mathcal{O}(D^2 L)$ , a significant downside to this approach is that it introduces approximation errors with a magnitude of  $o(\mathbf{W})$ , which can escalate relative to the weight matrix values.
3.2 Training Flow-based Models with Score-Matching Objectives
The pioneering study [14] is the earliest attempt to train flow-based models by minimizing the SM objective. Their results demonstrate that models trained using the SM loss are able to achieve comparable or even better performance to those trained with the ML objective in a low-dimensional experimental setup. More recently, the authors in [16] and [17] proposed two efficient variants of the SM loss, i.e., the SSM and FDSSM objectives, respectively. They demonstrated that these loss functions can be used to train a non-linear independent component estimation (NICE) [1] model on high-dimensional tasks. While the training approaches of these works bear resemblance to ours, our proposed method places greater emphasis on training efficiency. Specifically, they directly implemented the energy function  $E(\pmb{x};\theta)$  in the score-matching objectives as  $-\log p(\pmb{x};\theta)$ , resulting in a significantly higher computational cost compared to our method introduced in Section 4. In Section 5, we further demonstrate that the models trained with the methods in [16, 17] yield less satisfactory results in comparison to our approach.
In this section, we introduce a new framework for reducing the training cost of flow-based models with linear transformations, and discuss a number of training techniques for enhancing its performance.
4.1 Energy-Based Normalizing Flow
Instead of applying architectural constraints to reduce computational time complexity, we achieve the same goal through adopting the training objectives of energy-based models. We name this approach as Energy-Based Normalizing Flow (EBFlow). A key observation is that the parametric density function of a flow-based model can be reinterpreted as that of an energy-based model through identifying
the input-independent multipliers in  $p(\cdot ;\theta)$ . Specifically,  $p(\cdot ;\theta)$  can be explicitly factorized into an.   
unnormized density and a corresponding normalizing term as follows:
$$
\begin{array}{l} p (\boldsymbol {x}; \theta) = p _ {\mathbf {u}} (g (\boldsymbol {x}; \theta)) \prod_ {i = 1} ^ {L} | \det  \left(\mathbf {J} _ {g _ {i}} (\boldsymbol {x} _ {i - 1}; \theta)\right) | \\ = \underbrace {p _ {\mathbf {u}} \left(g \left(\boldsymbol {x} ; \theta\right)\right) \prod_ {g _ {i} \in S _ {n}} \left| \det \left(\mathbf {J} _ {g _ {i}} \left(\boldsymbol {x} _ {i - 1}; \theta\right)\right) \right|} _ {\text {U n n o r m a l i z e d D e n s i t y}} \underbrace {\prod_ {g _ {i} \in S _ {l}} \left| \det \left(\mathbf {J} _ {g _ {i}} (\theta)\right) \right|} _ {\text {N o r m . C o n s t .}} \triangleq \underbrace {\exp \left(- E \left(\boldsymbol {x} ; \theta\right)\right)} _ {\text {U n n o r m a l i z e d D e n s i t y}} \underbrace {Z ^ {- 1} (\theta)} _ {\text {N o r m . C o n s t .}} \tag {8} \\ \end{array}
$$
where the energy function  $E(\cdot ;\theta)$  and the normalizing constant  $Z^{-1}(\theta)$  are selected as follows:
$$
E (\boldsymbol {x}; \theta) \triangleq - \log \left(p _ {\mathbf {u}} (g (\boldsymbol {x}; \theta)) \prod_ {g _ {i} \in \mathcal {S} _ {n}} | \det  \left(\mathbf {J} _ {g _ {i}} \left(\boldsymbol {x} _ {i - 1}; \theta\right)\right) |\right), Z ^ {- 1} (\theta) = \prod_ {g _ {i} \in \mathcal {S} _ {l}} | \det  \left(\mathbf {J} _ {g _ {i}} (\theta)\right) |. \tag {9}
$$
The detailed derivations of Eqs. (8) and (9) are elaborated in Lemma A.11 of Section A.1.2. By isolating the computationally expensive term in  $p(\cdot ;\theta)$  as the normalizing constant  $Z(\theta)$ , the parametric pdf defined in Eqs. (8) and (9) becomes suitable for the training methods of energy-based models. In the subsequent paragraphs, we discuss the training, inference, and convergence property of EBFlow.
Training Cost. Based on the definition in Eqs. (8) and (9), the score-matching objectives specified in Eqs. (5)-(7) can be adopted to prevent the Jacobian determinant calculation for the elements in  $S_{l}$ . As a result, the training complexity can be significantly reduced to  $\mathcal{O}(D^2 L)$ , as the  $\mathcal{O}(D^3 L)$  calculation of  $Z(\theta)$  is completely avoided. Such a design allows the use of arbitrary linear transformations in the construction of a flow-based model without posing computational challenge during the training process. This feature is crucial to the architectural flexibility of a flow-based model. For example, fully-connected layers and convolutional layers with arbitrary padding and striding strategies can be employed in EBFlow without increasing the training complexity. EBFlow thus exhibits an enhanced flexibility in comparison to the related works that exploit specially designed linear transformations.
Inference Cost. Although the computational cost of evaluating the exact Jacobian determinants of the elements in  $S_{l}$  still requires  $\mathcal{O}(D^3 L)$  time, these operations can be computed only once after training and reused for subsequent inferences, since  $Z(\theta)$  is a constant as long as  $\theta$  is fixed. In cases where  $D$  is extremely large and  $Z(\theta)$  cannot be explicitly calculated, stochastic estimators such as the importance sampling techniques (e.g., [33, 34]) can be used as an alternative to approximate  $Z(\theta)$ . We provide a brief discussion of such a scenario in Appendix A.3.