Self Normalizing Flows

Given an observation ${\bm{x}}\in\mathbb{R}^{D}$, it is assumed that ${\bm{x}}$ is generated from an underlying real vector ${\bm{z}}\in\mathbb{R}^{D}$ through an invertible and differentiable transformation $g$, and that $f=g^{-1}$ is also differentiable (i.e. $g$ is a diffeomorphism). It is further assumed that ${\bm{z}}$ is a sample from a simple known underlying distribution $p_{{\bm{Z}}}$ such as a standard Gaussian. Then, the probability density $p_{{\bm{X}}}$ can be computed exactly using the change of variables formula:

where the change of volume term $\left|{\bm{J}}_{f}\right|=\left|\frac{\partial f({\bm{x}})}{\partial{\bm{x}}}\right|$ is the determinant of the Jacobian of the transformation between ${\bm{z}}$ and ${\bm{x}}$, evaluated at ${\bm{x}}$.

Typically, the functions $f$ and $g$ are defined as compositions of diffeomorphisms themselves, i.e. $g=g_{0}\circ g_{1}\circ\dots g_{K}$ and $f=f_{K}\circ f_{K-1}\circ\dots f_{0}$ where ${g_{k}}^{-1}=f_{k}$. This formulation takes advantage of the fact that a composition of diffeomorphic functions is also a diffeomorphism, meaning that if each $f_{k}$ is invertible and differentiable, then so is the composition $f$ and the change of variables formula in Equation 1 still holds.

Most approaches then propose defining and parameterizing only one direction of the flow, for example the ’forward’ functions $f_{k}$, and compute the inverses $g_{k}$ exactly when needed. The log-likelihood of the observations is then simultaneously maximized, with respect to a given $f_{k}$’s vector of parameters $\bm{\theta}_{k}$, for all $k$, requiring the gradient. Using the identity $\frac{\partial}{\partial{\bm{J}}}\log|{\bm{J}}|={{\bm{J}}^{-T}}$ we obtain the following gradient of the loss with respect to a given layer $k$’s parameters:

Following the conventions of Magnus (2010) for matrix derivatives, we make use of the vectorization operator $\mathrm{vec}$ which maps from $m\times n$ matricies to $mn\times 1$ column vectors by stacking the columns of the matrix sequentially, and formulate the parameters $\bm{\theta}_{k}$ as column vectors.

In order to avoid the inverse Jacobian in the gradient, we instead propose to define and parameterize both the forward and inverse functions $f_{k}$ and $g_{k}$ with parameters $\bm{\theta}_{k}$ and $\bm{\gamma}_{k}$ respectively. We then constrain the parameterized inverse $g_{k}$ to be approximately equal to the true inverse $f^{-1}_{k}$ through a layer-wise reconstruction loss. We can thus define our maximization objective as the mixture of the log-likelihoods induced by both models minus the reconstruction penalty:

where $p^{f}_{{\bm{X}}}$ and $p^{g}_{{\bm{X}}}$ now denote the densities induced by both the forwards and inverse transformations separately, and ${\bm{h}}_{k}=\mathrm{gradient\_stop}(f_{k-1}\circ...f_{0}({\bm{x}}))$ is the output of function $f_{k-1}$ with the gradients blocked such that only $g_{k}$ and $f_{k}$ receive gradients from the reconstruction loss at layer $k$. We see that when $f=g^{-1}$ exactly, this is equivalent to the traditional normalizing flow framework.

By the inverse function theorem, we know that the inverse of the Jacobian of an invertible function is given by the Jacobian of the inverse function, i.e. ${\bm{J}}_{f}^{-1}({\bm{x}})={\bm{J}}_{f^{-1}}({\bm{z}})$. Therefore, we see that with the above parameterization and constraint, we can approximate both the change of variables formula, and the gradients for both functions, in terms of the Jacobians of the respective inverse functions. Explicitly:

where Equation 5 follows from the derivation of Equation 4 and the application of the derivative of the inverse. We note that the above approximation requires that the Jacobians of the functions are approximately inverses in addition to the functions themselves being approximate inverses. For the models presented in this work, this property is obtained for free since the Jacobian of a linear mapping is the matrix representation of the map itself. However, for more complex mappings, this may not be exactly the case and should be constrained explicitly. We suggest this could be efficiently implemented with an additional loss analogous to a reconstruction loss but employing Jacobian vector products instead of matrix vector products (see Section A.4).

Although there are no known convergence guarantees for such a method, we observe in practice that, with sufficiently large values of $\lambda$, most models quickly converge to solutions which maintain the desired constraint. This phenomenon is demonstrated in Figure 3 where it is shown that the angle between the true gradient and the approximate gradient at each training iteration decreases to less than 1 degree over the course of training for both models tested. As a visual example of the quality of the inverse approximation, Figure 2 shows samples from the base distribution $p_{{\bm{Z}}}$ passed through both the true inverse $f^{-1}$ (top) and the learned approximate inverse $g$ (bottom) to generate samples from $p_{{\bm{X}}}$. As demonstrated by the nearly identical samples, the approximate inverse appears to be a very close match to the true inverse. The details of the models which generated these samples are in Section 5. Further mitigation strategies for potential optimization difficulties are discussed in Sections 6 and A.4. Finally, in Figure 4 we see the efficiency improvements of our method compared with the exact gradient. As expected, we see the self normalizing method scales much more favorably with input dimensionality than the exact gradient method.

The above framework proposes approximate gradients which allow for the training of normalizing flow architectures without having to compute expensive determinants or matrix inverses. However, to compute the exact log-likelihood of an observation (i.e. perform inference), the exact log Jacobian determinant still has to be computed for each input, which may be expensive.

To alleviate this limitation we take advantage of the compositional formulation of $f$, and the multiplicative property of the determinant. In detail, the determinant of the product of square matrices is equal to the product of their determinants. Therefore, assuming our flow is composed of square transformations $f_{k}$, the Jacobian is given by ${\bm{J}}_{f}=\prod_{k}{\bm{J}}_{f_{k}}$, and the log determinant of the Jacobian is given by:

For neural network architectures composed of sequential linear and nonlinear transformations, this allows us to separate the components of the Jacobian which are data-independent (e.g. those of linear layers) from those that are data-dependant (e.g. those of the activations). For example, given a network composed of $L$ layers of the form ${\bm{h}}_{l}=\sigma_{l}({\bm{W}}_{l}{\bm{h}}_{l-1})$ where $\sigma_{l}$ is a point-wise non-linearity, combining Equations 1 and 6 yields:

where ${\bm{\Sigma}}_{k}({\bm{x}})$ is the Jacobian of the activation function at layer $k$, evaluated at the sample ${\bm{x}}$. Importantly, we observe that the data-independent terms, $\sum_{k}\log|{\bm{W}}_{k}|$, are the expensive part of inference, and the data-dependent terms, $\sum_{k}\log|{\bm{\Sigma}}_{k}({\bm{x}})|$, are frequently cheap to compute analytically given the derivative of the activation. Therefore, as can be seen, once trained, the expensive data-independent terms must only be computed once and their values can then be reused for all future examples, effectively amortizing the cost of the determinant computation.

As a specific case of the self normalizing framework, we introduce a single fully connected self normalizing layer, as exemplified in Figure 1. Let ${\bm{W}},{\bm{R}}\in\mathbb{R}^{D\times D}$, with $f({\bm{x}})={\bm{W}}{\bm{x}}={\bm{z}}$, and $g({\bm{z}})={\bm{R}}{\bm{z}}$, such that ${\bm{W}}^{-1}\approx{\bm{R}}$. We additionally denote the layer-wise reconstruction loss as $\mathcal{E}({\bm{x}})=||{\bm{R}}{\bm{W}}{\bm{x}}-{\bm{x}}||^{2}_{2}$, but leave the derivation of the gradients of this term for the appendix (see Section A.1) since these are efficiently computed by standard deep learning frameworks and require no approximations.

Taking the gradients of Equation 3 with respect to all parameters ${\bm{W}}$ and ${\bm{R}}$, we get the following exact gradients:

where $\bm{\delta}^{f}_{{\bm{z}}}=\frac{\partial\log p_{{\bm{Z}}}({\bm{W}}{\bm{x}})}{\partial{\bm{W}}{\bm{x}}}$ and $\bm{\delta}^{g}_{{\bm{z}}}=\frac{\partial\log p_{{\bm{Z}}}({\bm{R}}^{-1}{\bm{x}})}{\partial{\bm{R}}^{-1}{\bm{x}}}$ can be computed by standard backpropagation. To avoid computing matrix inverses, using our framework we substitute ${\bm{W}}^{-1}$ with ${\bm{R}}$ in Equation 8, and symmetrically, all instances of ${\bm{R}}^{-1}$ with ${\bm{W}}$ in Equation 9. Additionally, we approximate $\bm{\delta}^{f}_{{\bm{z}}}\approx\bm{\delta}^{g}_{{\bm{z}}}$ in Equation 9 such that we do not need to ‘forward propagate’ through the inverse model, and can re-use the backpropagated error from the forward model. This corresponds to assuming that ${\bm{R}}^{-1}{\bm{x}}\approx{\bm{W}}{\bm{x}}$ for all ${\bm{x}}$ which follows directly from ${\bm{R}}^{-1}\approx{\bm{W}}$. Ultimately this results in the following approximations to gradient:

where $\bm{\delta}^{f}_{{\bm{x}}}=\frac{\partial\log p_{{\bm{Z}}}({\bm{W}}{\bm{x}})}{\partial{\bm{x}}}={\bm{W}}^{T}\delta^{f}_{{\bm{z}}}$. We see that by using such a self normalizing layer, the gradient of the log-determinant of the Jacobian term, which originally required an expensive matrix inversion at each iteration, is approximately given by the weights of the inverse transformation – sidestepping computation of both the Jacobian and the inverse. Additionally, we see the $\bm{\delta}^{f}_{{\bm{x}}}$ term required for Equation 11 is already computed by traditional backpropagation, making the update efficient. Finally, we note the above gradients can trivially be extended to compositions of such layers, combined with non-linearities, by substituting the appropriate deltas for each layer, and the corresponding layer inputs and outputs for ${\bm{x}}$ and ${\bm{z}}$ respectively.

To construct a self normalizing convolutional layer, we consider the setting where both the forward transformation $f$, and the inverse transformation $g$ are convolutional with the same kernel size. We note, importantly, that the inverse of a convolution operation is not necessarily another convolution. However, for sufficiently large $\lambda$, we observe that $f$ is regularized such that it is restricted to the class of convolutions which is approximately invertible by a convolution.

We define the parameters of $f$ to be the kernel ${\bm{w}}$ and similarly, the parameters of $g$ to be the kernel ${\bm{r}}$. Then, letting $f({\bm{x}})={\bm{w}}\star{\bm{x}}={\bm{z}}$ and $g({\bm{z}})={\bm{r}}\star{\bm{z}}$, such that $f^{-1}\approx g$ with ${\bm{x}},{\bm{z}}\in\mathbb{R}^{D}$, we proceed to derive the exact gradients of the log-likelihood. Again, we ignore the reconstruction term for simplicity as it requires no approximations.

To make the derivation easier, we note that the convolution operation is a linear operation, and can therefore be represented in matrix form. We define a transformation ${\bm{W}}=\mathcal{T}({\bm{w}})$, which maps between the convolutional kernel ${\bm{w}}$ and the corresponding matrix form of the convolution:

Letting ${\bm{w}}$ be a column vector and again making use of the vectorization operator $\mathrm{vec}$, we compute the exact gradient of the log-likelihood with respect to the kernel weights ${\bm{w}}$:

where we see that the first term $\frac{\partial(\mathrm{vec}\ \mathcal{T}({\bm{w}}))^{T}}{\partial{\bm{w}}}\left(\mathrm{vec}\ \left[\bm{\delta}^{f}_{{\bm{z}}}{\bm{x}}^{T}\right]\right)$ is given by the convolution $\bm{\delta}^{f}_{{\bm{z}}}\star{\bm{x}}$, as is usually done with backpropagation in convolutional neural networks. Then, given our soft constraint $f^{-1}\approx g$, we can approximate $\mathcal{T}({\bm{w}})^{-T}$ with $\mathcal{T}({\bm{r}})^{T}$ giving us:

To simplify the second term we note two points. First, the transpose of a convolution can similarly be achieved by standard convolution with a transformed kernel. We call this transformed kernel $\mathrm{flip}({\bm{r}})$ such that $\mathcal{T}(\mathrm{flip}({\bm{r}}))=\mathcal{T}({\bm{r}})^{T}$. Succinctly, $\mathrm{flip}(\cdot)$ is implemented by swapping the input and output axes, and mirroring the spatial (height and width) dimensions of the kernel. The second point is that the partial derivative $\frac{\partial(\mathrm{vec}\ \mathcal{T}({\bm{w}}))^{T}}{\partial{\bm{w}}}$ is given by a rectangular matrix, which, when multiplied by the vectorized form of the convolution matrix $\mathcal{T}({\bm{w}})$ yields a constant multiple ${\bm{m}}$ element-wise multiplied with the kernel ${\bm{w}}$. Each multiple element $m_{i}$ is given by the number of times the associated kernel element $w_{i}$ is shared across the matrix $\mathcal{T}({\bm{w}})$. We provide the derivation of this more thoroughly in the appendix, as well as an efficient method for calculating ${\bm{m}}$ for arbitrary convolutions (see Section A.2). In combination, we arrive at the approximate gradient of the likelihood with respect to the kernel ${\bm{w}}$:

We see that the symmetric derivation can be obtained for the gradient with respect to the kernel ${\bm{r}}$, as outlined below:

On the MNIST dataset we train three classes of models: 2-layer fully connected (FC) models with smooth-leaky-ReLU activations (Gresele et al., 2020), 9-layer convolutional (CNN) models with spline activations (Durkan et al., 2019), and 32-layer Glow models composed of affine coupling layers and 1x1 convolutional mixing layers (Kingma & Dhariwal, 2018). In all cases, we compare a self normalizing version with its exact gradient baseline.

As can be seen in Table 1 the models composed of self normalizing flow layers are nearly identical in performance to their exact gradient counterparts on the MNIST dataset. We see that the fully connected self normalizing model drastically outperforms the relative gradient method of Gresele et al. (2020), reaching the same performance as the exact gradient method. Additionally, we see that the self normalizing convolutional layer outperforms its constrained counterpart from Hoogeboom et al. (2019), likely due to the fact that the emerging convolution is unable to represent the 1x1 convolution explicitly. We hypothesize that the convolutional self normalizing flow model slightly under-performs the exact gradient method due to the convolutional-inverse constraint. We propose this constraint can be relaxed by using more complex inverse functions, potentially composed of multiple layers and non-linearities (see Section A.4), but leave this to future work. All training details can be found in the appendix (see Section A.3).

In Figure 5 the plot on the right shows that the qualitative convergence properties of the approximate gradient methods are very similar to those of the exact gradient, eventually converging to nearly the same validation likelihood. However, as can be seen in the plot on the left, due to the reduced computational complexity, the approximate gradient method trains in less than half the time, and even more quickly to approximate convergence. This timing comparison is demonstrated more exactly in Table 3 where the time per training batch and time per sample is computed for all models presented in this work. As can be seen, the self normalizing flow models are the fastest of all presented models, with the exception of the relative gradient method which appears to lag behind in likelihood performance. From this table we also see that the relative improvements to speed are directly related to the portions of the network which are replaced with self normalizing components. Since only the 1x1 convolution is replaced in the Glow framework, there is only a slight speed increase to be had.

Finally, we quantitatively measure the quality of the gradient approximation by measuring the alignment of directions of the approximate gradient and the exact gradient in Figure 3. Specifically, for models trained in the self normalizing framework, we measure the ‘angular error’ between the approximate gradient and the true gradient at each training iteration, for each layer, and plot the average angular error over all layers for each training epoch. The shaded area denotes the standard deviation in angles. We make two observations from this figure. First, both the 2-layer FC model and the 9-Layer CNN appear to have an initial slight divergence from the true gradient, but then quickly align to less than 1-degree of error, suggesting they are close approximations to the true gradient direction. Second, we observe that the average angular error of the approximate gradient is significantly larger for the convolutional model than it is for the fully connected model. When comparing this with the results in Table 1 we see that this could be contributing to the slightly lower performance of the self normalizing model when compared with the exact gradient methods. We again hypothesize this could be due to convolutional inverse constraint, making the inverse approximation more challenging. We leave further exploration of this topic to future work but believe that the performance could be ameliorated with improved inverse approximations, potentially achieved through more complex constrained optimization techniques.

For large scale experiments, we incorporate self normalizing flow layers into the Glow framework and train the models on CIFAR-10 and Imagenet 32x32. Specifically, we use the same model architectures as those proposed in Kingma & Dhariwal (2018), with some slight changes to the optimization parameters as detailed in Section A.3.

We observe that the self normalizing flow models are able to achieve competitive bits per dimension on both datasets (as seen in Table 2), while simultaneously training slightly faster than their exact gradient counterparts (as seen in Table 3). Importantly, we experimented with values of $\lambda$ in the set $\{1,10,100,1000\}$, and chose $\lambda=1000$ for all CIFAR-10 and Imagenet models due to increased stability during training. We observed a slight reduction in final likelihood performance as a result of such a large reconstruction weight and believe this is a factor in the performance gap between the self normalizing and exact gradient models. We believe that with more tuning, or with a dynamically weighted constrained optimization method such as that presented in Platt & Barr (1988), the self normalizing model is likely to match the exact gradient model even more closely. Preliminary experiments in this direction are shown in Section A.5.

Although it is clear that the glow framework is not the optimal setting for the application of self normalizing layers, given the determinant calculation of the 1x1 convolution is relatively quick to evaluate already, we present this work as a proof of scalability of our framework to models with greater than 100 steps of flow (as in the ImageNet 32x32 case), and to larger scale images.