Fast Flow Reconstruction via Robust Invertible nxn Convolution

Flow-based generative models have recently become one of the most efficient approaches to model data generation. Indeed, they are constructed with a sequence of invertible and tractable transformations. Glow first introduced a simple type of generative flow using an invertible $1 \times 1$ convolution. However, the $1 \times 1$ convolution suffers from limited flexibility compared to the standard convolutions. In this paper, we propose a novel invertible $n \times n$ convolution approach that overcomes the limitations of the invertible $1 \times 1$ convolution. In addition, our proposed network is not only tractable and invertible but also uses fewer parameters than standard convolutions. The experiments on CIFAR-10, ImageNet and Celeb-HQ datasets, have shown that our invertible $n \times n$ convolution helps to improve the performance of generative models significantly.

There are two types of generative models extensively deployed in recent years, including likelihood-based methods [29,30,32,33] and Generative Adversarial Networks (GANs) [34]. Likelihood-based methods have three main categories: Autoregressive models [30], variational autoencoders (VAEs) [35], and flow-based models [29,32,33]. The flow-based generative model is constructed using a sequence of invertible and tractable transformations, the model explicitly learns the data distribution and therefore the loss function is simply a negative log-likelihood.

arXiv:1905.10170v3 [cs.CV] 7 Aug 2022
The flow-based model was first introduced in [32] and later extended in RealNVP [33]. These methods introduced an affine coupling layer that is invertible and tractable based on Jacobian determinant. As the design of the coupling layers, at each stage, only a subset of data is transformed while the rest is required to be fixed. Therefore, they may be limited in flexibility. To overcome this limitation, coupling layers are alternated with less complex transformations that manipulate on all dimensions of the data. In RealNVP [33], the authors use a fixed channel permutation using fixed checkerboard and channel-wise masks. Kingma et al. [29] simplifies the architecture by replacing the reverse permutation operation on the channel ordering with invertible 1 × 1 convolutions.
However, the 1 × 1 convolutions are not flexible enough in these scenarios. It is extremely hard to compute the inverse form of the standard n × n convolutions, and this step usually produces high computational costs. There are prior approaches that design the invertible n × n convolutions by using emerging convolution [36], periodic convolutions [36], autoregressive flow [37] or stochastic approximation [38][39][40]. In this paper, we propose an approach to generalize an invertible 1 × 1 convolution to a more general form of n × n convolution. Firstly, we reformulate the standard convolution layer by shifting the inputs instead of the kernels. Then, we propose an invertible shift function that is a tractable form of Jacobian determinant. Through the experiments on CIFAR-10 [41], ImageNet [42] and Celeb-HQ [43] datasets, we prove that our proposals are significant and efficient for highdimensional data. Figure 1 illustrates the advantages of our approach with high-resolution synthesized images. Contributions: This work generalizes the invertible 1 × 1 convolution to an invertible n × n convolution by reformulating the convolution layer using our proposed invertible shift function. Our contributions can be summarized as follows: • Firstly, by analyzing the standard convolution layer, we reformulate its equation into a form such that, rather than shifting the kernels during the convolution process, shifting the input provides equivalent results. • Secondly, we propose a novel invertible shift function that mathematically helps to reduce the computational cost of the standard convolution while keeping the range of the receptive fields. The determinant of the Jacobian matrix produced by this shift function can be computed efficiently. • Thirdly, evaluations of several datasets on both objects and faces have shown the generalization of the proposed n × n convolution using our proposed novel invertible shift function.

Related Work
The generative models can be divided into two groups, that is, Generative Adversarial Networks and Flow-based Generative Models. In the first group, Generative Adversarial Networks [34] provide an appropriate solution to model the data generation. The discriminative model learns to distinguish the real data from the fake samples produced using a generative model. Two models are trained as they are playing a mini-max game. Meanwhile, in the second group, the Flow-based Generative Models [29,32,33] are constructed using a sequence of invertible and tractable transformations. Unlike GAN, the model explicitly learns the data distribution p(x) and therefore the loss function is efficiently employed with the log-likelihood.
In this section, we discuss several types of flow-based layers that are commonly used in flow-based generative models. An overview of several invertible functions is provided in the Table 1. In particular, all functions easily obtain the reverse function and tractability of a Jacobian determinant. The symbols , / denote element-wise multiplication and division. h, w denotes the height and width of the input/output. c, i, j are the depth channel index and spatial indices, respectively.

Description Function Reverse Function Log-Determinant
ActNorm [29] Coupling Layers: NICE [32] and RealNVP [33] presented coupling layers with a normalizing flow by stacking a sequence of invertible bijective transformation functions. The bijective function is designed as an affine coupling layer, which is a tractable form of Jacobian determinant. RealNVP can work in a multi-scale architecture to build a more efficient model for large inputs. To further improve the propagation step, the authors applied batch normalization and weight normalization during training. Later, Ho et. al. [44] presented a continuous mixture cumulative distribution function to improve the density modeling of coupling layers. In addition to improving the expressiveness of transformations of coupling layers, [44] utilized multi-head self-attention layers [45] in the transformations.
Inverse Autoregressive Convolution: Germain et al. [46] introduced autoregressive autoencoders by constructing an extension of a non-variational autoencoder that can estimate distributions and is straightforward in computing their Jacobian determinant. Masked autoregressive flow [37] is a type of normalizing flow, where the transformation layer is built as an autoregressive neural network. Inverse autoregressive flow [30] formulates the conditional probability of the target variable as an autoregressive model. Invertible 1 × 1 Convolution: Kingma et al. [29] proposed simplifying the architecture via invertible 1 × 1 convolutions. Learning a permutation matrix is a discrete optimization that is not amenable to gradient ascent. However, the permutation operation is simply a special case of a linear transformation with a square matrix. We can pursue this work with convolutional neural networks, as permuting the channels is equivalent to a 1 × 1 convolution operation with an equal number of input and output channels. Therefore, the authors replace the fixed permutation with learned 1 × 1 convolution operations. Activation Normalization: [29] performs an affine transformation using scale and bias parameters per channel. This layer simply shifts and scales the activations with datadependent initialization that normalizes the activations given an initial minibatch of data. This allows the scaling down of the minibatch size to 1 (for large images) and the scaling up of the size of the model.
Invertible n × n Convolution: Since the invertible 1 × 1 convolution is not flexible, Hoogeboom et al. [36] proposed an invertible n × n convolution generalized from the 1 × 1 convolutions. The authors presented two methods to produce the invertible convolutions: (1) Emerging Convolution and (2) Invertible Periodic Convolutions. Emerging Convolution is obtained by chaining specific invertible autoregressive convolutions [30] and speeding up this layer through the use of an accelerated parallel inversion module implemented in Cython. Invertible Periodic Convolutions transform data to the frequency domain via Fourier transform; this alternative convolution has a tractable Jacobian determinant and inverse. However, these invertible n × n convolutions require more parameters; therefore, these have an additional computational cost compared to our proposed method.
Lipschitz Constant: Behrmann et al. [38] developed a theory that any residual blocks satisfying the Lipschitz Constant can be invertible. Hence, Behrmann et al. proposed an invertible residual network (i-ResNet) as a normalizing flow-based model. Similar to [29,32,33,36], i-ResNet is learned by optimizing the negative log-likelihood in which the inverse flow and Jacobian determinant of the residual block can be efficiently approximated by the stochastic methods. Inheriting the success of Lipschitz theory, Kim et al. [39] proposed an L 2 self-attention that allows the self-attention of the Transformer networks [45] to be invertible.

Flow-Based Generative Model
Let x be a high-dimensional vector with unknown true distribution x ∼ p X (x), x ∈ X , a simple prior probability distribution p Z on a latent variable z ∈ Z, a bijection f : X → Z, the change of variable formula defines a model distribution on X as shown in Equation (1). where ∂x is the Jacobian of f at x. The log-likelihood objective is then equivalent to minimizing: Since the data x are discrete data, we add a random uniform noise u ∈ U (0, a), where a is determined by the discretization level of the data, to make x be continuous data. The generative process can be defined as Equation (3).
The bijection function f is constructed from a sequence of invertible and tractable Jacobian determinant transformations: . Such a sequence of invertible transformations is also called a normalizing flow. Here, Equation (2) can be written as in Equation (4). where

Standard n × n Convolution
In this section, we revisit the standard n × n convolution. Let X be an C × H × W input; W is a D × C × K kernel, and the convolution can be expressed as follows: where X k :,:,: is a C × H × W matrix that represents a spatially shifted version of input matrix X with shift amount (i k , j k ), . W :,:,k represents the D × C matrix corresponding to the kernel index k, the symbol denotes a convolution operator.
In Equation (5), the standard convolution is simply a sum of 1 × 1 convolutions on shifted inputs. The function S k maps the input X to the corresponding shifted input X k :,:,: . The standard convolution uses the common shifted input with integer-valued shift amounts for index k. Figure 2 illustrates our reformulated n × n convolution, if we can share the shifted inputs regardless of the kernel index, especially S k (X) = S(X), the standard convolution will be simplified as the 1 × 1 convolution as shown in Equation 6. In this paper, we propose a shift function S, which is an invertible and tractable form of the Jacobian determinant.   Reformulating n × n convolution. We propose to shift inputs instead of kernels. The proposed invertible n × n convolution will be simplified as a combination of the invertible shift function S and the invertible 1 × 1 convolution.

Invertible n × n Convolution
In this section, we first introduce our proposed Invertible Shift Function and then present invertible n × n convolution in details.

Invertible Shift Function
The shift function S will approximate all shifted input X k :,:,: (1 ≤ k ≤ K). Here, we propose to design S as a linear transformation per channel; specifically, we have learnable variables α c , β c ; 1 ≤ c ≤ C are scale and translation parameters for each channel, respectively. The shift function S can be formulated as follows: where c, i, j are the depth channel index and spatial indices, respectively. The reverse function of S can be easy to obtain: Thanks to Equation (7), the value of S(X c,i,j ) only depends on X c,i,j and the Jacobian matrix will be in the form of the diagonal matrix as follows: Therefore, the determinant of Equation (9) is the product of all elements in the diagonal of the matrix J as in Equation (10).

Invertible n × n Convolution
Kingma [29] proposed invertible 1 × 1 convolution as the smart way to learn the permutation matrix instead of the fixed permutation [32,33]. However, the 1 × 1 suffers from limited flexibility compared to the standard convolution. In particular, the receptive fields of 1 × 1 convolution is limited. When the network goes deeper, the receptive fields of 1 × 1 convolutions are still small areas; these, therefore, cannot generalize or model large objects of high-dimensional data. However, the 1 × 1 convolution has its own advantages compared to the standard convolution. First, the 1 × 1 convolution allows the network to compress the data of the input volume to be smaller. Second, 1 × 1 suffers less overfitting due to small kernel sizes. Therefore, in our proposal, we still take advantages of the 1 × 1 convolution. Specifically, we adopt the successfully invertible 1 × 1 convolution of Glow [29] in our design.
In the previous subsection, we proved that the shift function S is invertible and proved the tractability of the Jacobian determinant. In Section 3.2, we indicated that if we can share shifted inputs regardless of the kernel index via the shift function S, we can simplify the standard n × n convolution to the composition of the S and 1 × 1 convolution. Therefore, the invertible n × n convolution will be equivalent to the combination of the invertible shift function S and the invertible 1 × 1 convolution. Specifically, the input will first be forwarded to the shift function S and then convoluted with the 1 × 1 filter. Algorithm 1 illustrates the pseudo code of the invertible n × n convolution.

Algorithm 1: Invertible n × n Convolution
Input: An input X ∈ R N×H×W×C Result: An output of invertible n × n convolution and the log Jacobian determinant Initialize α, β ∈ R C for the invertible shift function; Initialize W ∈ R C×C as a rotation matrix for the invertible 1 × 1 convolution function; logdet = 0.0; The invertible shift function; W); The inverse will be Y = Conv(Z, W −1 ); logdet = logdet + log(det(W)) * H * W; Return Z and logdet; Figure 3a illustrates our one step of flow. We adopt the common design of a flow step [29,31,36] in our design. Our proposal can be easily integrated to the multi-scale architecture designed by Dinh et al. [33] (Figure 3b). From our proposal, we can generalize the invertible 1 × 1 convolution to the invertible n × n convolution through the shift function S. It can help to encourage the filters to learn a more efficient data representation and embed more useful latent features than the invertible 1 × 1 convolution used in Glow [29].
Besides, we use fewer parameters and have less inference time compared to the standard n × n convolutions.

Experiments
In this section, we present our experimental results on CIFAR-10, ImageNet and Celeb-HQ datasets. Firstly, in Section 5.1, we compare log-likelihood against the previous flow-based models, that is, RealNVP [33], Glow [29] and Emerging Convolution [36]. Finally, in Section 5.2, we show our qualitative results trained on the Celeb-HQ dataset.

Quantitative Experiments
Datasets and Metric: We evaluate our invertible n × n convolution on CIFAR-10 ( Figure 4a) and ImageNet (Figure 4b) with 32 × 32 and 64 × 64 image sizes. We use bits per dimension as the criteria with which to evaluate models. We compare our method against RealNVP [33], Glow [29] and Emerging Convolution [36]. We adopt the network structures of Glow and replace all invertible 1 × 1 convolutions of Glow with our invertible n × n convolutions. For the data preprocessing, we follow the same process as in RealNVP [33]. Network Configurations: In the CIFAR experiment, the depth of flow K and the number of levels L are set to 32 and 3, respectively. Meanwhile, the depth of flow in ImageNet experiments is set to 48, the numbers of levels of ImageNet 32 × 32 and ImageNet 64 × 64 experiments are set to 3 and 4, respectively. We use the Adam optimizer [47] to optimize the networks in which batch size and learning rate are set to 64 (per GPU) and 0.001, respectively. We choose Normal Distribution as the prior distribution p Z (z) ∼ N (z; 0, I) in all experiments.
The shift function S will be not inverse if the α c = 0 (∃ c ∈ [1...C]). Hence, in the training process, we will first initialize α c = 1 and β c = 0 (1 ≤ c ≤ C). During the learning processing, we keep α c (1 ≤ c ≤ C) as a different 0 to guarantee that the shift function S is inverse and to guarantee the tractability of the Jacobian determinant. Training models on high-dimensional data requires large memory. To be able to train with a large batch size, we simultaneously and distributively trained the models on four GPUs via Horovod (https://github.com/horovod/horovod) and TensorFlow (https://tensorflow.org) frameworks. Results: Table 2 shows our experimental results. In particular, our proposal helps to improve the generative models on ImageNet 32 × 32 and ImageNet 64 × 64 datasets, which are more challenging than CIFAR-10. In particular, our proposed method achieves a stateof-the-art performance in which the bit per dimension results in ImageNet 32 × 32, and ImageNet 64 × 64 is 3.96 and 3.74, respectively. In comparison, the Emerging Convolution [36] and Glow achieve similar results in both ImageNet 32 × 32 and ImageNet 64 × 64 benchmarks, which are 4.09 and 3.81, respectively. Meanwhile, the corresponding results of RealNVP on these benchmarks are 4.28 and 3.98, respectively. As shown by the results, our proposed invertible n × n convolution provides a better generative capability than the stand-alone invertible 1 × 1 convolution. Since Emerging Convolution uses invertible auto-regressive convolution, our proposal is, therefore, less complicated and has faster inference than Emerging Convolution. In the CIFAR-10 benchmark, although our model does not perform as well as Glow [29] and Emerging Convolution [36], we find it interesting that our method gains competitive results with a small number of modifications. The gap in performance is partially caused by the small amount of CIFAR-10 data that is inefficient for training the well-generalized convolution.

Qualitative Experiments
The CelebA-HQ dataset [43] was selected to train the model using the architectures defined in the previous section with a higher resolution (256 × 256 image sizes). The depth of flow K and the number of levels L were set to 32 and 6, respectively. Since highdimensional data requires large memory, we reduced the batch size to 1 (per GPU) and trained on eight GPUs. The qualitative experiment aims to study the efficiency of the model when it scales up to the high-resolution images, synthesizes realistic images, and provides the meaningful latent space. Figure 4c shows the examples of Celeb-HQ datasets. We trained our model on 5-bit images in order to improve visual quality with a slight trade-off of color fidelity. As shown by the synthetic images in Figure 5, our model can generalize realistic images in high dimensional data.

Conclusions and Future Work
This paper has presented a novel invertible n × n convolution approach. By reformulating the convolution layer, we propose to use the shift function to shift inputs instead of kernels. We prove that our shift function is invertible and tractable in terms of calculating the Jacobian determinant. The method leverages the shift function and the invertible 1 × 1 convolution to generalize to the invertible n × n convolution. Through experiments, our proposal has achieved state-of-the-art results in quantitative measurement and is able to generate realistic images with high-resolution.
There are several challenges that remain to be addressed in future work. In particular, when the model scales up to the high-resolution images, it requires a large amount of GPU memory during the training process, that is, the back-propagation process. Maintaining the rotation matrix property for the invertible 1 × 1 convolution when training the model on a large dataset is also a challenging task, since the model easily falls into the non-inverse matrix due to the stochastic gradient update of the back-propagation algorithm. That issue is interesting work and should be improved in the future.