Flowification: Everything is a Normalizing Flow
Anonymous Author(s)
AffiliationAddress email
Abstract
We develop a method that can be used to turn any multi-layer perceptron or convolutional network into a normalizing flow.
In some cases this requires the additionof uncorrelated noise to the model but in the simplest case no additional parameters.
The technique we develop can be applied to a broad range of architectures, allowingthem to be used for a wide range of tasks.
Our models also allow existing densityestimation techniques to be combined with high performance feature extractors.
Incontrast to standard density estimation techniques that require specific architecturesand specialized knowledge, our approach can leverage design knowledge fromother domains and is a step closer to the realization of general purpose architectures.
We investigate the efficacy of linear and convolutional layers for the task of densityestimation on standard datasets.
Our results suggest standard layers lack somethingfundamental in comparison to other normalizing flows
Introduction
Density estimation techniques have proven effective on a wide variety of downstream tasks suchas sample generation and anomaly detection [1–6].
Normalizing flows and autoregressive modelsperform very well at density estimation but do not easily scale to large dimensions [2, 7, 8] and haveto satisfy strict design constraints to ensure efficient computation of their Jacobians and inverses.
Advances in other areas of machine learning cannot be utilized as flow architectures because they are nottypically seen as being invertible; this restricts the application of highly optimized architectures frommany domains to density estimation, and the use of the likelihood for diagnosing these architectures.
Methods using standard convolutions and residual layers for density estimation have been developedfor architectures with specific properties [9–12].
These methods do not provide a recipe for convertinggeneral architectures into flows.
There is no known correspondence between normalizing flows andthe operations defined by linear and convolutional layers.
In this paper we show that a large proportion of machine learning models can be trained as normalizingflows.
The forward pass of these models remains unchanged apart from the possible addition ofuncorrelated noise.
To demonstrate our formulation works we apply it to fully connected layers, convolutions and residual connections.
The code used for these experiments will be made publicly available.
The contributions of this paper include
• In §3.1 we show that linear layers induce densities as augmented normalizing flows [13] with the multi-scale architecture used in RealNVPs [3].
We also show how these layers canbe viewed as funnels [14] to increase their expressivity.
We term this process flowification .
• In §3.2 we argue that most ML architectures can be decomposed into simple buildingblocks that are easy to flowify.
As an example, we derive the specifics for 2-dimensionalconvolutional layers and residual blocks.• In §4 we flowify convolutional networks and train them as normalizing flows using thelikelihood.
This demonstrates that models built from standard layers can be used for densityestimation directly.
Background
Flows A complex probability distribution { p ( x ) | x ∈ X } can be constructed from a simple basedensity 1 { p ( z ) | z ∈ Z } and a diffeomorphism between X and Z , f : X → Z .
The likelihood ofany given x is then given by p ( x ) = p ( f ( x )) | det ( J fx ) | , where J fx is the Jacobian of f evaluatedat x .
Thus, the log-likelihoods of the two densities are related by an additive term, which will bereferred to as the likelihood contribution V ( x, z ) [15].
Normalizing flows [1] chain many invertiblemaps that are tuned to maximize the (log-)likelihood of the data.
A lot of development has gone intoconstructing flexible invertible layers with easy to calculate Jacobians where both the forward andinverse passes are fast to compute [2–4, 16, 17].
As the function f must be invertible it is required to preserve the dimension of the data.
This limitsthe expressivity of f and makes it expensive to model high dimensional data distributions.
To reducethese issues several works have studied dimension altering variants of flows [13, 15, 18–21].
Reducing the dimensionality A simple method for altering the dimension of a flow is totake the output of an intermediate layer z ′ and split it into two pieces z ′
= { z ′ 1 , z ′ 2 } .
Multiscale architectures [3] match one of these pieces directly to the base density and apply furthertransformations to the other.
Funnels [14] model z ′ 2 as being dependent on z ′ 1 and work withp ( z ′ )
=
p ( f ′ ( z ′ 1 )) p ( z ′ 2 | f ′ ( z 1 ))
| det J f ′ z ′ 1 | where f ′ can be any computational graph for which thelikelihood contribution can be calculated.
The conditional distribution can be a trainable densityand z ′ 1 is modelled with additional layers.
It is useful to think of this factorization as a dimensionreducing mechanism from dim( z ′ ) to dim( z ′ 1 ) .
Increasing the dimensionality Dimension increasing flow layers can improve a models flexibility,as demonstrated by augmented normalizing flows [13].
To increase the dimensionality, x is embeddedinto a larger dimensional space and data independent noise u is added to the embedding to obtain adistribution with support of nonzero measure.
This noise addition is similar to dequantization [3, 22],but is orthogonal to the distribution of x and increases its dimension.
It is important to note thatthe augmentation x
(cid:55)→ ( x, u ) increases the dimension from dim x to dim x + dim u .
Under anaugmentation the likelihood of x can be estimated using
[Equation]
In practice, we estimate this expectation value by sampling u everytime a datapoint is passed throughthe network.
This means the integral is estimated with a single sample as in surVAEs [15].
Flowification
Suppose A is a network architecture with parameter space Θ .
Then for any choice of θ ∈ Θ thenetwork with parameters θ realizes a function A θ : R D → R C for some D and C .
Similarly, anormalizing flow model F is a parametric distribution on some R E , where for any choice of γ fromtheir parameter space Γ they define a density function F γ on R E .
In this work we show that a largeclass of neural network architectures can be thought of as a flow model by constructing a map
[Equation]
The embedding of A to its flowification F A results in a flow model that can realize density functionson the augmented space R D
× R N for some N ≥ 0 , which in turn induces a density on R D byintegrating out the component on R N .
The parameter space of F A factorises as Θ × Φ
where Θ isthe parameter space of A and also that of the forward pass of F A , while Φ parametrises the inversepass of F A . In the simplest case Φ = ∅ , i.e. flowification does not require additional parameters.
It isin this sense that we claim that a large fraction of machine learning models are normalizing flows.
Terminology In what follows we work with conditional distributions such as p ( z | x ) and it willbe practical to think of them as “stochastic functions” p : x (cid:55)→ z , that take an input x and producean output z ∼ p ( z | x ) .
Conversely, we think of a function f : x (cid:55)→ z as the Dirac δ -distributionf
( z | x ) =
δ ( z − f ( x )) .
These definitions allow us to have a unified notation for deterministic andstochastic functions such that we can talk about them in the same language.
Consequently, whenwe say "stochastic function", it will include deterministic functions as a corner case.
Depending onwhether f and f − 1 are deterministic or stochastic, we talk about left, right or two-sided inverses.
Wewill be careful to be precise about this.
Method In the following we consider the standard building blocks of machine learning architecturesand enrich them by defining (stochastic-)inverse functions and calculating the likelihood contributionof each layer.
Treating each layer separately allows density estimation models to be built throughcomposition [15].
The stochastic inverse can use the funnel approach, which increases the parametercount, or the multi-scale approach, which does not.
For simplicity we will only consider conditionaldensities in the inverse as this is more general, though it is not required.
We will refer to this processas flowification and the enriched layers as flowified ; non-flowified layers will be called standardlayers .
Flowified layers can then be seen as simultaneously being
• Flow layers that are invertible, their likelihood contribution is known and therefore can be used to train the model to maximize the likelihood.
• Standard layers that can be trained with losses other than the likelihood, but for which the
likelihood can be calculated after this training with fixed weights in the forward direction.
3.1 Linear Layers
Let L W,b : R n → R m denote the the linear layer of a neural network with parameters defined by aweight matrix W ∈ R m × n and bias b ∈ R m .
Formally, L W,b is defined as the affine function
[Equation]
Definition 1.
Let ϕ ( z | x ) : R n → R m be a stochastic function.
We say that ϕ is linear in expectationif there exists W ∈ R m × n and b ∈ R m such that for any x ∈ R n the expected value of ϕ coincideswith the application of L W,b E z ∼ ϕ ( z | x ) [ z ] =
[Equation]
W [Equation] b ( x ) . Similarly, we say that a stochastic function ψ ( z | x ) is convolutional in expectation if the expectation 107 E z ∼ ψ ( z | x ) [ z ] is given by a convolutional layer.
In this section we flowify linear layers, by which we mean we construct a pair of stochastic functions,a forward L ( z | x ) : R n → R m and an inverse L − 1 ( x | z ) : R m → R n such that the forward is linearin expectation and is compatible with the inverse in a way that will be made precise in the followingparagraphs.
SVD parametrization
To build a flowified linear layer, the first step is to parametrize the weight matrices by the singularvalue decomposition (SVD)[23].
This involves writing W ∈ R m × n as a product W =
V Σ U , where
U ∈ R n × n is orthogonal, Σ ∈ R m × n is diagonal and V ∈
R m × m is orthogonal.
This parametrizationis particularly useful for our purposes because the orthogonal transformations are easily invertibleand do not contribute to the likelihood, and the non-invertible piece of the transformation is localizedto Σ .
Parametrizing U and V
We generate elements of the special orthogonal group SO ( d ) by applyingthe matrix-exponential to elements of the Lie algebra so ( d ) of skew-symmetric matrices.
Weparametrize so ( d ) and perform gradient descent there.
As the Lie-algebra is a vector space, this issignificantly easier than working directly with SO ( d ) .
See Appendix G for details.
Parametrizing Σ
The matrix Σ is of shape m × n containing the singular values on the maindiagonal.
We ensure maximal rank of Σ , by parameterizing the logarithm of the main diagonal, thisway all singular values are greater than 0.
It is important to note that this parametrization is not without loss of generality.
In particular,it does not include matrices of non-maximal rank nor orientation reversing ones, where either
U ∈ O ( n ) \ SO ( n ) or
V ∈ O ( m ) \ SO ( m ) .
This implementation detail does not change the generalperspective we provide of linear layers as normalizing flows, but instead simplifies the implementationof flowified layers.
Reducing the dimensionality
Definition 2.
We call the tuple ( L ( z | x ) , L − 1 ( x | z ))
a dimension decreasing flowified linear layer if
L is dimension decreasing, linear in expectation and the following conditions are satisfied
(i)
The forward is deterministic, given by L ( z | x ) =
L W,b ( x ) ,
(ii)
The layer is right-invertible, L ◦ L − 1 = id z , (iii)
The likelihood contribution of L can be exactly computed.
To flowify dimension decreasing linear layers, we define the forward function L as a standard linearlayer with parameters W and b ,
[Equation]
Since W is parametrized by the SVD decomposition, W = V Σ U , we need to invert V, U and Σseparately.
As V and U are rotations, they are invertible in the usual sense.
To construct a stochasticinverse to Σ , we think of it as a funnel [14] and use a neural network p inv (( Ux ) ( m :) | Σ
Ux ) thatmodels the n − m dropped coordinates as a function of the m non-dropped coordinates.
Again, thisis not required to calculate the likelihood under the model, even a fixed distribution could be used,but introducing some trainable parameters significantly improves the performance of the flow thatis defined by the layer.
We use Σ − 1 to denote this stochastic inverse to Σ .
The stochastic inversefunction L − 1 can then be written as
[Equation]
Since the rotations don’t contribute to the log-likelihood, the likelihood of data under a dimensiondecreasing flowified linear layer is
[Equation]
where log Σ denotes the sum of the logarithms of the diagonal elements of Σ . [Figure]
Theorem 3.
The above choices for L and L − 1 define a dimension decreasing flowified linear layer.
Sketch of proof.
The definition of the forward pass (5) makes the forward pass linear in expectationand satisfies (i) by definition.
Unpacking the definitions and decomposing W into its SVD formyields right-invertibility (ii) which in turn implies that the likelihood contribution can be exactlycomputed (iii).
When the inverse density is not made to be conditional the above ideas can be visualized as a standardmulti-scale flow architecture [3] as shown in Fig.
Increasing the dimensionality
Definition 4.
We define the Moore-Penrose pseudoinverse L + W,b of a linear layer L W,b : R n → R mas the affine transformation R m → R n
[Equation]
(8) where W + denotes the Moore-Penrose pseudoinverse of the matrix W .
Definition 5.
We call the tuple ( L ( z | x ) , L − 1 ( x | z ))
a dimension increasing flowified linear layer if
L is dimension increasing, linear in expectation and the following conditions are satisfied
(iv) The inverse L − 1 is deterministic, given by L − 1 ( x | z )
=
L + W,b ( z ) ,
(v) The layer is left-invertible, L − 1 ◦ L
= id x ,
(vi) The likelihood contribution of L can be bounded from below.
To construct dimension increasing flowified linear layers, we rely again on the SVD parametrizationwhere the only nontrivial component is Σ .
In this case Σ is a dimension increasing operation and wethink of it as an augmentation step [13] composed with diagonal scaling.
To augment, we samplem − n coordinates from a distribution U with zero mean and then apply a scaling in m dimensions.
The likelihood contribution is then given by
[Equation]
where log Σ denotes the sum of the logarithms of the m scaling parameters.
The inverse function
L − 1 is the composition of the inverse rotations, the inverse scaling and the dropping of the sampledcoordinates.
Theorem 6.
The above choices for L and L − 1 define a dimension increasing flowified linear layer.
Sketch of proof.
The augmentation step [13] results in a lower bound on the likelihood contribution,implying (vi).
Since E u ∼ U = 0 , the augmentation does not influence the expected value of z , i.e. theforward pass is linear in expectation.
Simple calculations using the SVD decomposition then implyboth (iv) and (v).
Preserving the dimensionality
Dimension preserving layers are a corner case of both of the above scenarios, where padding andsampling are not needed in either direction and the layer is non-stochastically invertible.
All thisimplies (i),(ii),(iii),(iv) and (v) are satisfied.
3.2 Convolutional layers
Convolutions can be seen as a dimension increasing coordinate repetition followed by a matrixmultiplication with weight sharing.
In the previous section we derived the specifics of matrixmultiplication.
We begin this section with the details of coordinate repetition after which we put thepieces together to build a flowified convolutional layer.
In Appendix H we describe an alternativeapproach relying on the Fourier transform.
Repeating coordinates In this paragraph we focus on the N -fold repetition of a single scalarcoordinate x .
This significantly simplifies notation, but the technique generalizes in an obvious way.
Intuitively, the idea is to expand the one dimensional volume to N dimensions by first embeddingand then increasing the volume of the embedding such that the volume in the N − 1 directionscomplementary to the embedding can be controlled.
We have seen in §2 that the operation
[Equation]
has likelihood contribution E u [ − log p ( u )] .
Now, we can apply any N -dimensional rotation R Nwhich maps (1 , 0 ) to 1 √ N (1 , 1 ) to obtain
[Equation]
Note that this rotation does not contribute to the likelihood.
Finally, we apply a diagonal scaling in Ndimensions with factor √ N such that [Equation]
[Equation]
where the final scaling has likelihood contribution N log( √ N ) =
( N/ 2) log N and x is now repeated
N times.
The overall contribution to the likelihood of the embedding (12) is
[Equation]
By construction, the padding distribution R N (0 , √ N u ) is orthogonal to the diagonal embedding
[Equation]
( x, ..., x ) of the data distribution.
The inverse function is given by the projection to the diagonalembedding,
[Equation]
General architectures Now that the likelihood contribution of arbitrary linear layers and coordinaterepetition has been computed it is possible to flowify more general architectures such as convolutionsand residual connections.
It is important to note that just because an architecture works well forcertain tasks, it is not clear if its flowified version will perform well at density estimation.
Decomposing convolutional layers To flowify convolutional layers, we decompose it as a sequenceof building blocks that are easy to flowify separately.
A standard convolutional layer performs thefollowing sequence of steps:
Padding of the input image with zeros to increase its size.
Unfolding of the padded image into tiles.
This step replicates the data according to the
kernel size and stride.
Applying a linear layer.
Finally, we apply the same linear layer to each of the tiles produced
in the previous step.
The outputs then correspond to the pixels of the output image.
Flowification Steps 1 and 3 are already flowified, i.e. their likelihood contribution is computed and an inverse is constructed, in §3.1.
We denote their flowification with Pad and Linear , respectively.
Step 2 fits into the discussion of the previous paragraph of repeating coordinates, where both its inverse (14) and its likelihood contribution (13) are given.
We will denote this operation by Unfold .
Definition 7.
Let Linear , Unfold and Pad be as above and define C and C − 1 be the following stochastic functions
[Equation]
call the resulting layer ( C , C − 1 ) a flowified convolutional layer .
Augment
SVD SVD
A flowified convolutional layer ( C , C − 1 ) is then convolutional in expectation (Definition 1), i.e. thereexists a convolutional layer C θ with parameters θ such that
[Equation]
[Equation]
The flowification of a convolution without padding can be seen in Fig.
The Unfold operation isimplemented as coordinate duplication and Linear is a flowified linear layer parameterized by the
SVD.
Activation functions Functions that are surjective onto R and invertible fit well in our frameworkas they can be used out of the box without any modifications.
In our experiments we use LeakyReLUand rational-quadratic splines [4] as activations functions.
Non-invertible activations can also be usedwhen equipped with additional densities [15].
Residual connections Residual connections can be seen as coordinate duplication followed bytwo separate computational graphs { f 1 , f 2 } who’s output is recombined with a sum.
The sum canbe inverted by defining a density over one of the summands p ( f 1 ( x + u ) | f 1 ( x + u ) , f 2 ( x 2 )) andsampling from this density, which will also define the likelihood contribution.
Then, if the likelihoodcontribution can be calculated for each individual computational graph, the likelihood of the totaloperation can be calculated [15].
Experiments
To test the constructions described in the previous section we flowify multilayer perceptrons andconvolutional architectures and train them to maximize the likelihood of different datasets. Tabular data In this section we study a selection of UCI datasets [24] and the BSDS 300 collectionof natural images [25] using the preprocessed dataset used by masked autoregressive flows [17, 26]. We compare the performance with several baselines for comparison in Table 1. We see that theflowified models have the right order of magnitude for the likelihood but are not competitive.
Image Data In this section we use the MNIST [27] and CIFAR10 [28] datasets with the standardtraining and test splits.
The data is uniformly dequantized as required to train on image data [29, 30].
For both datasets we trained networks consisting only of flowified linear layers ( F MLP) and alsonetworks consisting of convolutional layers followed by dense layers ( F C ONV 1).
To minimize thenumber of augmentation steps that occur in each model we define additional architectures with similarnumbers of parameters but with non-overlapping kernels in the convolutional layers ( F C ONV 2).
Theexact architectures can be found in Appendix F.1. The flowified layers sample from N (0 , a ) for [Table] dimension increasing operations, where a is a per-layer trainable parameter.
We use RQ-splines [4]with 8 knots and a tail bound of 2 as activation functions, where the same function is applied peroutput node.
Importantly, in this section we do not use existing flow layers in the forward pass of ourlayers, nor any data-dependent operations, such as coupling layers.
We also ran experiments withrational quadratic splines [4] are mixed into F C ONV 2.
The improved performance of these modelssuggests that flowified layers do not mismodel the density of the data, but they do lack the capacity tomodel it well.
Samples from these models are shown in Appendix C.
[Table]
The results of the density modelling can be seen in Table.
The images in the first and third columnsof Fig.
3 are generated by sampling from a standard gaussian in the latent space and taking theexpected output of the inverse in every layer.
The images in the second and fourth columns use thesame latent samples as the left column but also sample from the distribution defined by the inversepass of the layers.
Samples from the F MLP models can be found in the Appendix.
As seen in Fig.
2 the F MLP models are outperformed by the F C ONV models and the convolutionalmodels with non-overlapping kernels achieve better results than the ones with overlapping kernels.
This suggests both that the inductive bias of the convolution is useful for modelling distributions ofimages and that the augmentation step costs more in terms of likelihood than it provides in terms ofincreased expressivity.
Related Work
Several works have developed methods that allow standard layers to be made invertible, but theseapproaches restrict the space of the models, whereas we consider networks in their full generality.
Invertible ResNets [9, 10] require that each residual block has a Lipschitz constant less than one, buteven with this restriction they attain competitive results on both classification and density estimation.
The same Lipschitz constraint can also be applied to other networks [11].
In these architectures themulti-scale architecture used in RealNVPs [3] is not leveraged, and so no information is discardedby the model.
It is unclear why this approach outperforms flowified layers, as seen in Table.
2,but it could be due to the preservation of information through the model, the very large numberof parameters that are used in these approaches, the restricted subspace of the models, or somecombination of these three.
It can also be shown that convolutions with the same number of input and output channels can be madeinvertible [12].
These layers perform poorly at the task of density estimation and are outperformedby flowified layers.
This is likely due to the increased expressivity that comes from consideringa larger space of architectures.
There have been several works that develop convolution inspiredinvertible transformations [31–33], but these architectures consider restricted transformations tomaintain invertibility.
Future Work and Conlusion
Our experiments suggest that flowified convolutional networks do not match the density estimationperformance of similarly sized normalizing flows.
A possible explanation is that the dimensionreducing steps discard information and more expressive encoding layers are necessary to transformthe distributions before reducing the dimensionality.
This is supported by the experiments using NSFlayers in-between the flowified layers (see App. C).
The addition of NSF layers leads to improvedperformance both in terms of visual quality and also BPD values.
Possibly the main limitation of anetwork consisting purely of flowified layers is the fact that, unlike flow layers, the forward passes ofstandard linear and convolutional layers are not data-dependent.
This is reinforced by the fact that theentanglement capability typically used in flows also appears in attention mechanisms, which havebeen shown to excel at capturing complex statistical structures [34–36].
With further development – such as increased capacity given to the inverse density, or data dependentparameters in the forward pass – standard architectures could become competitive density estimatorsin their own right and allow for general purpose models to be developed.
The focus of this work wason employing standard layers for density estimation, but it is possible that designing data dependentvariants of standard layers that are more flow-like could improve their performance on tasks suchas classification and regression.
The flowification procedure provides a useful means for designingsuch models, and demonstrates that standard architectures can be considered a subset of normalizingflows, a correspondence that has not previously been demonstrated.
A SurVAE
In this section we describe how the flowification procedure fits into the surVAE framework [15].
[-[Yes] See 
 (c) Did you discuss any potential negative societal impacts-]The surVAE framework can be used to compose bijective and surjective transformations such thatthe likelihood of the composition can be calculated exactly or bounded from below.
To define alikelihood for surjections both a forward and an inverse transformation is defined, where one of thedirections is necessarily stochastic.
[-[Yes] See 
 (b) Did you include complete proofs of all theoretical results? [No] We provided sketches of-]A transformation that is stochastic in the forward direction is called a generative surjection, and thelikelihood of such transformations can only be bounded from below.
Flowified dimension increasinglinear layers, residual connections, and coordinate repetition are examples of generative surjections.
A transformation that is stochastic in the inverse direction is called an inference surjection, where thelikelihood can be calculated exactly if the layer satisfies the right inverse condition [15]. A flowifieddimension reducing linear layer is an inference surjection that meets the right inverse condition.
In Fig.
4 we show a schematic of a flowified dimension increasing linear layer.
In this schematic theaugmentation step is a generative surjection as defined by the surVAE framework [15], while everyother transformation is a bijection.
[Figure]
The augmentation gap Dimension increasing operations (generative surjections) introduce anaugmentation gap in the likelihood calculations as derived by Huang et al.
After this augmentation step the subsequent layers model the joint distribution p ( x, u ) of the data x and the augmentingnoise u .
In this sense, the dimension increasing operations of this work are only normalizing flowsto the degree to which augmented normalizing flows are normalizing flows.
The way we performconvolutions (pad, unfold, apply kernel) requires a large number of dimension expanding steps(unfold) and it is inevitable to work with this (weakened) version of flows.
If a convolutional layeris such that the dimensionality (channels times height times width) of its input is less than or equalthan that of its output, the overall convolution operation is still in the dimension reducing regime,theoretically allowing to compute the exact likelihood.
We investigated this approach by diagonalizingthe convolution using the Fourier transform, which unfortunately turned out to be too expensivecomputationally to be used in practice.
See Appendix H for details.
B The choice of padding in dimension increasing operations
When flowifying a dimension increasing operation the choice of noise to use in the augmentation [13]plays an important role in the performance of the model.
In this section we compare sampling from auniform distribution to sampling from a normal distribution.
In all cases we observe that samplingfrom a normal distribution is significantly better than sampling from a uniform distribution as seenin Table 3.
This is likely due to the normal distribution inverse that is used as the inverse in thedimension reducing layers as well as the normal base distribution.
[Table]
C Flowification + NSF
Below are samples from an architecture combining flowified standard layers and NSF layers. [Figure]
[Figure]
D Complete proofs of Theorems 3 and 6
Proof of Theorem 3.
The missing part of the proof is the calculation showing right invertibility.
Toprove this, first we let z ∈ R
m be arbitrary and then we calculate
[Equation] [Equation] Proof of Theorem 6. The missing part of the proof is showing that L − 1 =
L + W,b and that the resultinglayer is left invertible.
To prove L − 1 = L + W,b we let z ∈ R m , [Equation]
Finally, left invertibility is proven by
[Equation] [Equation] E Samples from F MLP, F C ONV 1 and F C ONV 2 F Experimental details
Everything was implemented in PyTorch [37] and executed on a RTX 3090 GPU.
We also used weights and biases for experiment tracking [38] and the neural spline flows [39] packagefor the rational quadratic activation functions.
Both of these tools are available under an MIT license.
Tabular data
All training hyper parameters were set to be the same as those used in the neural spline flowexperiments [4].
All networks in this section are flowified MLPs.
Several different architectures weretested with different design choices.
Dimension preserving layers of different depth were consideredwith the number of layers set to 3 , 5 , 7 , 9 , 21 . Fixed depth layers were also considered with widthsfrom 256 , 128 , 64 , 32 , 16 and depths from 3 , 5 , 7 . We also tried hand designed widths with depthsfrom [128 , 128 , 128 , 64 , 32] , [128 , 64 , 32 , 16 , 8] , [8 , 16 , 32 , 64 , 128] , [32 , 64 , 128 , 128 , 128] .
In allcases the second architecture was used, except for BSDS300 where the first architecture was used. Leaky ReLU activations were used with a negative slope of 0.5.
Image data
The models on CIFAR- 10 were trained for 1000 epochs with a batch size of 256 and lasted 36-48hours. The models on MNIST were trained for 200 epochs with the same batch size and lastedapproximately 8-12 hours.
The learning rate in all experiments was initialized at 5 × 10 − 4 andannealed to 0 following a cosine schedule [40].
F.1 Architectures
1 import torch.nn as nn 2 import flowification as ffn 3 4 model = nn.
Sequential( ffn.
Flatten(), ffn. Linear(784, 512), RqSpline(), ffn. Linear(512, 256), RqSpline(), ffn.
Linear(256, 128), RqSpline(), ffn [Figure]
[Figure] 1 import torch.nn as nn 2 import flowification as ffn 3 4 model = nn.
Sequential( ffn.
Conv2d(3, 27, kernel_size=3, stride =2), RqSpline(), # ffn. Conv2d(27, 64, kernel_size=3, stride =2), RqSpline(), # ffn.
Conv2d(64, 100, kernel_size=3, stride =2), RqSpline(), # ffn.
Conv2d(100, 112, kernel_size=3, stride =2), RqSpline(), # ffn. Conv2d(112, 128, kernel_size=2, stride =2), RqSpline(), # ffn.
Flatten(), ffn.
Linear(128, 128), RqSpline(), ffn. Linear(128, 128), RqSpline(), ffn.
Linear(128, 128), RqSpline(), ffn. Linear(128, 128), RqSpline(), ffn. Linear(128, 128), RqSpline(), ffn.
Linear(128, 128), RqSpline(), ffn. Linear(128, 128), RqSpline(), ffn.
Linear(128, 128), RqSpline(), ffn. Linear(128, 128), RqSpline(), ffn.
Linear(128, 128), RqSpline() )
[Figure]
torch.nn as nn
[Figure]
flowification as ffn 3 4 model = nn.
Sequential( ffn.
Conv2d(1, 16, kernel_size=3, stride =2), RqSpline(), # ffn.
Conv2d(16, 24, kernel_size=2, stride =2), RqSpline(), # ffn. Conv2d(24, 32, kernel_size=3, stride =2), RqSpline(), # ffn.
Conv2d(32, 48, kernel_size=2, stride =1), RqSpline(), # ffn. Conv2d(48, 64, kernel_size=2, stride =1), RqSpline(), # ffn.
Flatten(), ffn.
Linear(64, 64), RqSpline(), ffn. Linear(64, 64), RqSpline(), ffn. Linear(64, 64), RqSpline(), ffn. Linear(64, 64), RqSpline(), ffn.
Linear(64, 64), RqSpline(), ffn. Linear(64, 64), RqSpline(), ffn. Linear(64, 32), RqSpline(), ffn. Linear(32, 32), RqSpline(), ffn.
Linear(32, 32), RqSpline(), ffn.
Linear(32, 32), RqSpline(), ffn. Linear(32, 32), RqSpline(), ffn. Linear(32, 32), RqSpline(), ffn. Linear(32, 32), RqSpline(), ffn.
Linear(32, 8), RqSpline() )
G Parametrizing rotations via the Lie algebra so ( d )
The Lie algebra so ( d ) of the rotation group SO ( d ) is the space of d × d skew-symmetric matrices [Equation] Moreover, since SO ( d ) is connected and compact, the (matrix-)exponential map so ( d ) → SO ( d ) issurjective [41, Corollary 11.10], i.e. every rotation
U ∈
SO ( d ) can be written as the exponential ofsome skew-symmetric X ∈ so ( d ) .
Motivated by this result, we propose to parametrize the Lie algebra so ( d ) , and use the exponentialmap to obtain rotations.
Since the exponential map is differentiable, we can propagate gradients allthe way back to the Lie algebra and do gradient based optimization there.
Note that is significantly eases the optimization process.
The reason for this is that the Lie group ofunitary transformations has non-trivial geometry and doing gradient descent is cumbersome on suchspaces.
On the other hand, the Lie algebra is a vector space, where gradient descent can shine.
Since PyTorch [37] has built-in support for the matrix exponential, random orthogonal transformationscan be generated with just a few lines of Python code:1 d
=# size of matrix 2 x = torch.randn(d,d) # random (d,d) matrix 3 S = x - x.T() # S is skew -symmetric 4 U = torch.matrix_exp(S)
# U is a rotation
Comparison with Householder transformations The matrix exponential has the advantage ofgenerating the whole rotation group, while many Householder transformations need to be composed tohave the same level of expressivity.
This, of course, comes with an increased ( O ( d 2 ) ) computationalcost.
We did not explore Householder transformations, but we expect that for linear layers with largeinput or output dimensionality, the matrix exponential will simply become too expensive and it willbecome necessary to rely on Householder transformations.
H Flowifying convolutions through the FFT
An alternative way of flowifing convolutional layers is through the convolution theorem which statesthat the Fourier-transformation diagonalizes convolutions.
In this section we sketch the idea anddiscuss the issues with it.
Fig.12 displays the Pytorch code for performing convolutions through the convolution theorem.
Line 12 of the code shows that in the Fourier basis, we can perform convolution as follows; at eachfrequency we form the vector of dimension (# of in_channels) and apply a weight matrix of size (# ofin_channels) ×
(# of out_channels) to it.
If this weight matrix is parametrized by the SVD, we canguarantee its invertibility.
Since the FFT is also invertible, this construction should be a flowificationof the convolutional layer if we can also solve the following two issues:
If the weight matrix is parametrized in frequency domain,
(1) it’s inverse transform might not be real.
(2) it’s inverse transform might not be of the appropriate size (i
e. nonzero outside of the first(kernel_size)-many pixels).
The first of these is easy to solve with the symmetric parametrization induced by k is real ⇔ FFT ( k ) ∗ i = FFT ( k ) N − i where N is the length of the signal.
To deal with the second condition, a auxiliary loss term isintroduced that minimizes the terms that should be 0 .
Although the idea of FFT seems promising,this approach turns out be slower (probably due to the FFT), worse performing and more cumbersomethan the flowificaton presented in §3.2.
Fig.
13 shows some samples from a Fourier-flowifiedconvolutional network trained on MNIST.
[-in_channels , out_channels = 4, signal_length , kernel_size = 32, batch_size =
-]
[Figure]
[-z2 = z
-]