RG-Flow: A hierarchical and explainable flow model based on renormalization group and sparse prior

1 INTRODUCTION . One of the most important unsupervised learning tasks is to learn the data distribution and build generative models . Over the past few years , various types of generative models have been proposed . Flow-based generative models are a particular family of generative models with tractable distributions ( Dinh et al. , 2017 ; Kingma & Dhariwal , 2018 ; Chen et al. , 2018b ; 2019 ; Behrmann et al. , 2019 ; Hoogeboom et al. , 2019 ; Brehmer & Cranmer , 2020 ; Rezende et al. , 2020 ; Karami et al. , 2019 ) . Yet the latent variables are on equal footing and mixed globally . Here , we propose a new flow-based model , RG-Flow , which is inspired by the idea of renormalization group in statistical physics . RGFlow imposes locality and hierarchical structure in bijective transformations . It allows us to access information at different scales in original images by latent variables at different locations , which offers better explainability . Combined with sparse priors ( Olshausen & Field , 1996 ; 1997 ; Hyvärinen & Oja , 2000 ) , we show that RG-Flow achieves hierarchical disentangled representations . Renormalization group ( RG ) is a powerful tool to analyze statistical mechanics models and quantum field theories in physics ( Kadanoff , 1966 ; Wilson , 1971 ) . It progressively extracts more coarse-scale statistical features of the physical system and decimates irrelevant fine-grained statistics at each scale . Typically , the local transformations used in RG are designed by human physicists and they are not bijective . On the other hand , the flow-based models use cascaded invertible global transformations to progressively turn a complicated data distribution into Gaussian distribution . Here , we would like to combine the key ideas from RG and flow-based models . The proposed RG-flow enables the machine to learn the optimal RG transformation from data , by constructing local invertible transformations and build a hierarchical generative model for the data distribution . Latent representations are introduced at different scales , which capture the statistical features at the corresponding scales . Together , the latent representations of all scales can be jointly inverted to generate the data . This method was recently proposed in the physics community as NeuralRG ( Li & Wang , 2018 ; Hu et al. , 2020 ) . Our main contributions are two-fold : First , RG-Flow can separate the signal statistics of different scales in the input distribution naturally , and represent information at each scale in its latent vari- ables z . Those hierarchical latent variables live on a hyperbolic tree . Taking CelebA dataset ( Liu et al. , 2015 ) as an example , the network will not only find high-level representations , such as the gender factor and the emotion factor for human faces , but also mid-level and low-level representations . To visualize representations of different scales , we adopt the concept of receptive field from convolutional neural networks ( CNN ) ( LeCun , 1988 ; LeCun et al. , 1989 ) and visualize the hidden structures in RG-flow . In addition , since the statistics are separated into a hierarchical fashion , we show that the representations can be mixed at different scales . This achieves an effect similar to style mixing . Second , we introduce the sparse prior distribution for latent variables . We find the sparse prior distribution is helpful to further disentangle representations and make them more explainable . The widely adopted Gaussian prior is rotationally symmetric . As a result , each of the latent variables in a flow model usually does not have a clear semantic meaning . By using a sparse prior , we demonstrate the clear semantic meaning in the latent space . 2 RELATED WORK . Some flow-based generative models also possess multi-scale latent space ( Dinh et al. , 2017 ; Kingma & Dhariwal , 2018 ) , and recently hierarchies of features have been utilized in Schirrmeister et al . ( 2020 ) , where the top-level feature is shown to perform strongly in out-of-distribution ( OOD ) detection task . Yet , previous models do not impose hard locality constraint in the multi-scale structure . In Appendix C , the differences between globally connected multi-scale flows and RG-Flow are discussed , and we see that semantic , meaningful receptive fields do not show up in the globally connected cases . Recently , other more expressive bijective maps have been developed ( Hoogeboom et al. , 2019 ; Karami et al. , 2019 ; Durkan et al. , 2019 ) , and those methods can be incorporated into the proposed structure to further improve the expressiveness of RG-Flow . Some other classes of generative models rely on a separate inference model to obtain the latent representation . Examples include variational autoencoders ( Kingma & Welling , 2014 ) , adversarial autoencoders ( Makhzani et al. , 2015 ) , InfoGAN ( Chen et al. , 2016 ) , and BiGAN ( Donahue et al. , 2017 ; Dumoulin et al. , 2017 ) . Those techniques typically do not use hierarchical latent variables , and the inference of latent variables is approximate . Notably , recent advances suggest that having hierarchical latent variables may be beneficial ( Vahdat & Kautz , 2020 ) . In addition , the coarseto-fine fashion of the generation process has also been discussed in other generative models , such as Laplacian pyramid of adversarial networks ( Denton et al. , 2015 ) , and multi-scale autoregressive models ( Reed et al. , 2017 ) . Disentangled representations ( Tenenbaum & Freeman , 2000 ; DiCarlo & Cox , 2007 ; Bengio et al. , 2013 ) is another important aspect in understanding how a model generates images ( Higgins et al. , 2018 ) . Especially , disentangled high-level representations have been discussed and improved from information theoretical principles ( Cheung et al. , 2015 ; Chen et al. , 2016 ; 2018a ; Higgins et al. , 2017 ; Kipf et al. , 2020 ; Kim & Mnih , 2018 ; Locatello et al. , 2019 ; Ramesh et al. , 2018 ) . Apart from the high-level representations , the multi-scale structure also lies in the distribution of natural images . If a model can separate information of different scales , then its multi-scale representations can be used to perform other tasks , such as style transfer ( Gatys et al. , 2016 ; Zhu et al. , 2017 ) , face mixing ( Karras et al. , 2019 ; Gambardella et al. , 2019 ; Karras et al. , 2020 ) , and texture synthesis ( Bergmann et al. , 2017 ; Jetchev et al. , 2016 ; Gatys et al. , 2015 ; Johnson et al. , 2016 ; Ulyanov et al. , 2016 ) . Typically , in flow-based generative models , Gaussian distribution is used as the prior for the latent space . Due to the rotational symmetry of Gaussian prior , an arbitrary rotation of the latent space would lead to the same likelihood . Sparse priors ( Olshausen & Field , 1996 ; 1997 ; Hyvärinen & Oja , 2000 ) was proposed as an important tool for unsupervised learning and it leads to better explainability in various domains ( Ainsworth et al. , 2018 ; Arora et al. , 2018 ; Zhang et al. , 2019 ) . To break the symmetry of Gaussian prior and further improve the explainability , we introduce a sparse prior to flow-based models . Please refer to Figure 12 for a quick illustration on the difference between Gaussian prior and the sparse prior , where the sparse prior leads to better disentanglement . Renormalization group ( RG ) has a broad impact ranging from particle physics to statistical physics . Apart from the analytical studies in field theories ( Wilson , 1971 ; Fisher , 1998 ; Stanley , 1999 ) , RG has also been useful in numerically simulating quantum states . The multi-scale entanglement renormalization ansatz ( MERA ) ( Vidal , 2008 ; Evenbly & Vidal , 2014 ) implements the hierarchical structure of RG in tensor networks to represent quantum states . The exact holographic mapping ( EHM ) ( Qi , 2013 ; Lee & Qi , 2016 ; You et al. , 2016 ) further extends MERA to a bijective ( unitary ) flow between latent product states and visible entangled states . Recently , Li & Wang ( 2018 ) ; Hu et al . ( 2020 ) incorporates the MERA structure and deep neural networks to design a flow-base generative model that allows machine to learn the EHM from statistical physics and quantum field theory actions . In quantum machine learning , recent development of quantum convolutional neural networks also ( Cong et al. , 2019 ) utilize the MERA structure . The similarity between RG and deep learning has been discussed in several works ( Bény , 2013 ; Mehta & Schwab , 2014 ; Bény & Osborne , 2015 ; Oprisa & Toth , 2017 ; Lin et al. , 2017 ; Gan & Shu , 2017 ) . The information theoretic objective that guides machine-learning RG transforms are proposed in recent works ( Koch-Janusz & Ringel , 2018 ; Hu et al. , 2020 ; Lenggenhager et al. , 2020 ) . The meaning of the emergent latent space has been related to quantum gravity ( Swingle , 2012 ; Pastawski et al. , 2015 ) , which leads to the exciting development of machine learning holography ( You et al. , 2018 ; Hashimoto et al. , 2018 ; Hashimoto , 2019 ; Akutagawa et al. , 2020 ; Hashimoto et al. , 2020 ) . 3 METHODS . Flow-based generative models . Flow-based generative models are a family of generative models with tractable distributions , which allows efficient sampling and exact evaluation of the probability density ( Dinh et al. , 2015 ; 2017 ; Kingma & Dhariwal , 2018 ; Chen et al. , 2019 ) . The key idea is to build a bijective map G ( z ) = x between visible variables x and latent variables z . Visible variables x are the data that we want to generate , which may follow a complicated probability distribution . And latent variables z usually have simple distribution that can be easily sampled , for example the i.i.d . Gaussian distribution . In this way , the data can be efficiently generated by first sampling z and mapping them to x through x = G ( z ) . In addition , we can get the probability associated with each data sample x , log pX ( x ) = log pZ ( z ) log @ G ( z ) @ z . ( 1 ) The bijective map G ( z ) = x is usually composed as a series of bijectors , G ( z ) = G1 G2 · · · Gn ( z ) , such that each bijector layer Gi has a tractable Jacobian determinant and can be inverted efficiently . The two key ingredients in flow-based models are the design of the bijective map G and the choice of the prior distribution pZ ( z ) . Structure of RG-Flow networks . Much of the prior research has focused on designing more powerful bijective blocks for the generator G to improve its expressive power and to achieve better approximations of complicated probability distributions . Here , we focus on designing the architecture that arranges the bijective blocks in a hierarchical structure to separate features of different scales in the data and to disentangle latent representations . 1 Our design is motivated by the idea of RG in physics , which progressively separates the coarsegrained data statistics from fine-grained statistics by local transformations at different scales . Let x be the visible variables , or the input image ( level-0 ) , denoted as x ( 0 ) ⌘ x . A step of the RG transformation extracts the coarse-grained information x ( 1 ) to send to the next layer ( level-1 ) , and splits out the rest of fine-grained information as auxiliary variables z ( 0 ) . The procedure can be described by the following recursive equation ( at level-h for example ) , x ( h+1 ) , z ( h ) = Rh ( x ( h ) ) , ( 2 ) which is illustrated in Fig . 1 ( a ) , where dim ( x ( h+1 ) ) + dim ( z ( h ) ) = dim ( x ( h ) ) , and the RG transformation Rh can be made invertible . At each level , the transformation Rh is a local bijective map , which is constructed by stacking trainable bijective blocks . We will specify its details later . The split-out information z ( h ) can be viewed as latent variables arranged at different scales . Then the inverse RG transformation Gh ⌘ R 1h simply generates the fine-grained image , x ( h ) = R 1h ( x ( h+1 ) , z ( h ) ) = Gh ( x ( h+1 ) , z ( h ) ) . ( 3 ) The highest-level image x ( hL ) = GhL ( z ( hL ) ) can be considered as generated directly from latent variables z ( hL ) without referring to any higher-level coarse-grained image , where hL = log2 L log2 m , for the original image of size L⇥L with local transformations acting on kernel size m⇥m . Therefore , given the latent variables z = { z ( h ) } at all levels h , the original image can be restored by the following nested maps , as illustrated in Fig . 1 ( b ) , x ⌘ x ( 0 ) = G0 ( G1 ( G2 ( · · · , z ( 2 ) ) , z ( 1 ) ) , z ( 0 ) ) ⌘ G ( z ) , ( 4 ) where z = { z0 , · · · , zhL } . RG-Flow is a flow-based generative model that uses the above composite bijective map G as the generator . To model the RG transformation , we arrange the bijective blocks in a hierarchical network architecture . Fig . 2 ( a ) shows the side view of the network , where each green or yellow block is a local bijective map . Following the notation of MERA networks , the green blocks are the disentanglers , which reparametrize local variables to reduce their correlations , and the yellow blocks are the decimators , which separate the decimated features out as latent variables . The blue dots on the bottom are the visible variables x from the data , and the red crosses are the latent variables z . We omit color channels of the image in the illustration , since we keep the number of color channels unchanged through the transformation . Fig . 2 ( b ) shows the top-down view of a step of the RG transformation . The green/yellow blocks ( disentanglers/decimators ) are interwoven on top of each other . The covering area of a disentangler or decimator is defined as the kernel size m⇥m of the bijector . For example , in Fig . 2 ( b ) , the kernel size is 4 ⇥ 4 . After the decimator , three fourth of the degrees of freedom are decimated into latent variables ( red crosses in Fig . 2 ( a ) ) , so the edge length of the image is halved . As a mathematical description , for the single-step RG transformation Rh , in each block ( p , q ) labeled by p , q = 0 , 1 , . . . , L2hm 1 , the mapping from x ( h ) to ( x ( h+1 ) , z ( h ) ) is given by n y ( h ) 2h ( mp+m2 +a , mq+m2 +b ) o ( a , b ) 2⇤1m = Rdish ✓n x ( h ) 2h ( mp+m2 +a , mq+m2 +b ) o ( a , b ) 2⇤1m ◆ n x ( h+1 ) 2h ( mp+a , mq+b ) o ( a , b ) 2⇤2m , n z ( h ) 2h ( mp+a , mq+b ) o ( a , b ) 2⇤1m/⇤2m =Rdech ✓n y ( h ) 2h ( mp+a , mq+b ) o ( a , b ) 2⇤1m ◆ , ( 5 ) where ⇤km = { ( ka , kb ) | a , b = 0 , 1 , . . . , mk 1 } denotes the set of pixels in a m ⇥ m square with stride k , and y is the intermediate result after the disentangler but not the decimator . The notation x ( h ) ( i , j ) stands for the variable ( a vector of all channels ) at the pixel ( i , j ) and at the RG level h ( similarly for y and z ) . The disentanglers Rdish and decimators R dec h can be any bijective neural network . Practically , We use the coupling layer proposed in the Real NVP networks ( Dinh et al. , 2017 ) to build them , with a detailed description in Appendix A . By specifying the RG transformation Rh = Rdech Rdish above , the generator Gh ⌘ R 1 h is automatically specified as the inverse transformation . Training objective . After decomposing the statistics into multiple scales , we need to make the latent features decoupled . So we assume that the latent variables z are independent random variables , described by a factorized prior distribution pZ ( z ) = Y l p ( zl ) , ( 6 ) Under review as a conference paper at ICLR 2021 x Visible variables z Latent variables ( c ) G en er at io n C au sa lC on e x Visible variables z Latent variables ( d ) In fe re nc e C au sa lC on e z Latent variables ( a ) ( b ) where l labels every element in z , including the RG level , the pixel position and the channel . This prior gives the network the incentive to minimize the mutual information between latent variables . This minimal bulk mutual information ( minBMI ) principle was previously proposed to be the information theoretic principle that defines the RG transformation ( Li & Wang ( 2018 ) ; Hu et al . ( 2020 ) ) . Starting from a set of independent latent variables z , the generator G should build up correlations locally at different scales , such that the multi-scale correlation structure can emerge in the resulting image x to model the correlated probability distribution of the data . To achieve this goal , we should maximize the log likelihood for x drawn from the data set . The loss function to minimize reads L = Ex⇠pdata ( x ) log pX ( x ) = Ex⇠pdata ( x ) ✓ log pZ ( R ( x ) ) + log @ R ( x ) @ x ◆ , ( 7 ) where R ( x ) ⌘ G 1 ( x ) = z denotes the RG transformation , which contains trainable parameters . By optimizing the parameters , the network learns the optimal RG transformation from the data . Receptive fields of latent variables . Due to the nature of local transformations in our hierarchical network , we can define the generation causal cone for a latent variable to be the affected area when that latent variable is changed . This is illustrated as the red cone in Fig . 2 ( c ) . To visualize the latent space representation , we define the receptive field for a latent variable zl as RFl = Ez⇠pZ ( z ) @ G ( z ) @ zl c , ( 8 ) where | · |c denotes the 1-norm on the color channel . The receptive field reflects the response of the generated image to an infinitesimal change of the latent variable zl , averaged over pZ ( z ) . Therefore , the receptive field of a latent variable is always contained in its generation causal cone . Higher-level latent variables have larger receptive fields than those of the lower-level ones . Especially , if the receptive fields of two latent variables do not overlap , which is often the case for lower-level latent variables , they automatically become disentangled in the representation . Image inpainting and error correction . Another advantage of the network locality can be demonstrated in the inpainting task . Similar to the generation causal cone , we can define the inference causal cone shown as the blue cone in Fig . 2 ( d ) . If we perturb a pixel at the bottom of the blue cone , all the latent variables within the blue cone will be affected , whereas the latent variables outside the cone can not be affected . An important property of the hyperbolic tree-like network is that the higher level contains exponentially fewer latent variables . Even though the inference causal cone is expanding as we go into higher levels , the number of latent variables dilutes exponentially as well , resulting in a constant number of latent variables covered by the inference causal cone on each level . Therefore , if a small local region on an image is corrupted , only O ( log L ) latent variables need to be modified , where L is the edge length of the entire image . While for globally connected networks , all O ( L2 ) latent variables have to be varied . Sparse prior distribution . We have chosen to hard-code the RG information principle by using a factorized prior distribution , i.e . pZ ( z ) = Q l p ( zl ) . The common practice is to choose p ( zl ) to be the standard Gaussian distribution , which is spherical symmetric . If we apply any rotation to z , the distribution will remain the same . Therefore , we can not avoid different features from being mixed under the arbitrary rotation . To overcome this issue , we use an anisotropic sparse prior distribution for pZ ( z ) . In our implementation , we choose the Laplacian distribution p ( zl ) = 12b exp ( |zl|/b ) , which is sparser compared to Gaussian distribution and breaks the spherical symmetry of the latent space . In Appendix E , we show a two-dimensional pinwheel example to illustrate this intuition . This heuristic method will encourage the model to find more semantically meaningful representations by breaking the spherical symmetry .

Summary: The paper proposes a method, named as RG-flow, which combines the ideas of Renormalization group (RG) and flow-based models. The RG is applied to separate signal statistics of different scales in the input distribution and flow-based idea represents each scale information in its latent variables with sparse prior distribution. Inspired by receptive field from CNNs, the authors visualize the latent space representation, which reveals the progressive semantics at different levels as the instinctive expectation.