Recursive Disentanglement Network

1 INTRODUCTION . Recent progress in machine learning demonstrates the ability to learn disentangled representations is essential to data-efficient learning , such as controllable image generation , image manipulation , and domain adaptation ( Suter et al. , 2019 ; Zhu et al. , 2018 ; Peng et al. , 2019 ; Gabbay & Hoshen , 2021 ; 2019 ) . β-VAE ( Higgins et al. , 2017 ) and its variants are the most investigated approaches for disentangled representation learning . Recent works on β-VAE-based methods introduce various inductive biases as regularization terms and directly apply them on the resulting embedding space of deep models , such as the bottleneck capacity constraint ( Higgins et al. , 2017 ; Burgess et al. , 2018 ) , total correlation among variables ( Kim & Mnih , 2018 ; Chen et al. , 2018 ) , and the mismatch between aggregated posterior and prior ( Kumar et al. , 2017 ) , aiming to balance among representation capacity , independence constraints , and reconstruction accuracy . Indeed , as demonstrated by Locatello et al . ( 2020 ; 2019 ) , unsupervised disentanglement is fundamentally impossible without explicit inductive biases on models and data sets . However , our study shows that existing β-VAE-based methods may not be able to learn satisfactory disentangled representations even for fairly trivial cases . This is due to the fact that the feature spaces of deep models have inherently compositional structures , i.e. , each complex feature is a composition of primitive features . However , existing methods with regularization terms solely applied to the resulting embedding space can not effectively propagate disentangled regularization across such compositional feature space . As shown in Figure 1 , applying the standard β-VAE to the widely used dataset dSprites ( Matthey et al. , 2017 ) , we visualize the resulting representation z , as well as its compositional low-level representations m extracted from the previous layer ( as shown in Figure 1 ( a ) ) , and evaluate the independence between each pair of m and each pair of z , respectively 1 . Figure 1 ( b ) and Figure 1 ( c ) show that the disentanglement quality of low-level features m may impact the resulting representation z in terms of disentanglement quality . This study demonstrates the potential benefit to regularize the compositional feature space of deep models during disentangled representation learning . This work aims to tackle the compositional disentanglement learning problem . First , we formulate disentangled representation learning from an information-theoretic perspective , and introduce a new learning objective covering three essential properties for learning disentangled representations : sufficiency , minimal sufficiency , and disentanglement . Theoretical analysis shows that the 1The independence between two components ci and cj is measured by the normalized mutual information ( Chen et al. , 2018 ) . Whenever NMI ( ci ; cj ) = I ( ci ; cj ) /H ( c ) = 0 , ci and cj are independent ( disentangled ) . proposed learning objective is a general form of β-VAE and several of its state-of-the-art variants . Next , we extend the proposed learning objective to cover the disentanglement representation learning problem in the compositional feature space . Governed by the proposed learning objective , we present Recursive Disentanglement Network ( RecurD ) , a compositional disentanglement learning method , which directs the disentanglement learning process across the compositional feature space by applying regulatory inductive bias recursively through the feed-forward network . We argue that the recursive propagation of inductive bias through the feed-forward network imposes a sufficient condition of disentangled representation learning . Empirical studies demonstrate that RecurD outperforms β-VAE ( Higgins et al. , 2017 ) and several other variants of VAE ( Burgess et al. , 2018 ; Kim & Mnih , 2018 ; Chen et al. , 2018 ; Kumar et al. , 2017 ) on disentangled representation learning and achieves more data-efficient learning in downstream machine learning tasks . 2 COMPOSITIONAL DISENTANGLEMENT LEARNING . In this section , we first formulate disentanglement learning from the information-theoretic perspective by introducing three key properties and show such formulation is a general form of the optimization objectives of β-VAE and several of its variants . Next , we extend the principled objective to the compositional feature space to tackle the compositional disentanglement learning problem . 2.1 DISENTANGLEMENT LEARNING OBJECTIVE . The challenge of representation learning can be formulated as finding a distribution p ( z|x ) that maps original data x ∈ X into a representation z with fixed amount of variables z = { z1 , . . . , zn } ( Bengio et al. , 2013 ) . The key intuition of z is to capture minimal sufficient information in a disentangled manner , given the reconstruction task x ≈ x̂ . We denote the representation learning process as a Markov Chain for which x̂→ x→ z , which means z depends on x̂ only through x , i.e. , p ( z|x ) = p ( z|x , x̂ ) ( see also , ( Cover , 1999 ; Achille & Soatto , 2018 ) ) . The principled properties of z are defined as follows : Definition 1 . Sufficiency : a representation z of x for x̂ is sufficient if I ( x , x̂ ) = I ( z , x̂ ) . For the reconstruction task , z is sufficient if it can successfully reconstruct x by x̂ . The difference between I ( x ; x̂ ) and I ( z ; x̂ ) is computed as follows : I ( x ; x̂ ) − I ( z ; x̂ ) = I ( x ; x̂|z ) = H ( x̂|z ) − H ( x̂|x ) . Given the reconstruction task x ≈ x̂ , H ( x̂|x ) is constant and independent to z , so the sufficient property can be optimized by minimizing H ( x̂|z ) ( Federici et al. , 2020 ; Dubois et al. , 2020 ) . Definition 2 . Minimal Sufficiency : a representation z of x is minimal sufficient if I ( x ; z ) = I ( z ; x̂ ) . A minimal sufficient z encodes the minimum amount of information about x required to reconstruct x̂ ( Cover , 1999 ; Achille & Soatto , 2018 ) . Since I ( z ; x̂ ) equals to I ( x ; x̂ ) when z is sufficient , the difference is computed as I ( x ; z ) − I ( z ; x̂ ) = I ( x ; z ) − I ( x ; x̂ ) . Given the reconstruction task x ≈ x̂ , I ( x ; x̂ ) is constant and independent to z , so the minimal sufficiency property can be optimized by minimizing I ( x ; z ) . Definition 3 . Disentanglement : a representation denoted as z = { z1 , . . . , zn } is disentangled if∑ j 6=i I ( zi ; zj ) = 0 . From the definition of mutual information , I ( zi ; zj ) = H ( zi ) − H ( zi|zj ) denotes the reduction of uncertainty in zi when zj is observed ( Cover , 1999 ) . If any two components zi and zj are disentangled , changes to zi have no influence on zj , which means I ( zi ; zj ) = 0 . A representation satisfying all these properties can be found by introducing two Lagrange multipliers λ1 and λ2 for two constrained expected properties with respect to the fundamental sufficiency property . The principled objective of disentanglement learning is to minimize the following objective : L = H ( x̂|z ) + λ1I ( x ; z ) + λ2 ∑ j 6=i I ( zi , zj ) . ( 1 ) The above objective can be interpreted as the reconstruction error , plus two regularizers that yield an optimally disentangled representation . The principled objective also helps us analyze and understand the success of recently developed β-VAE-based methods . These methods operate with an encoder with parameter φ and a decoder with parameter θ , to induce the joint distributions q ( x , z ) = qφ ( z|x ) q ( x ) and p ( x , z ) = pθ ( x|z ) p ( z ) , respectively , where p ( z ) is a fixed prior distribution . The learning objective of β-VAE contains the reconstruction error and KL divergence between the variational posterior and prior . To understand the relationship of learning objectives between Equation 1 and β-VAE-based methods , we decompose I ( x ; z ) ( Kim & Mnih , 2018 ) and estimate an upper bound for ∑ j 6=i I ( zi , zj ) ( Te Sun , 1980 ; 1975 ) , then we assign different weights as follows : λ1I ( x ; z ) + λ2 ∑ j 6=i I ( zi , zj ) ≤ λaEx [ KL ( q ( z|x ) ‖p ( z ) ) ] + λbKL ( q ( z ) ‖ n∏ j=1 q ( zj ) ) + λc n∑ j=1 KL ( q ( zj ) ‖p ( zj ) ) . ( 2 ) As shown in Table 1 , the learning objectives of β-VAE and its four variants can be regarded as specific cases of Equation 1 , i.e. , they assign different weights to our regularization terms , which can balance among latent variables capacity , independence constraints , and reconstruction accuracy , leading to successful disentangled representation learning ( Zhao et al. , 2017 ; Li et al. , 2020 ) . More details can be found in Appendix B . However , in these works , the inductive bias toward disentanglement is only applied to the embedding space of z , ignoring the need of disentanglement during feature composition in feed-forward networks . 2.2 COMPOSITIONAL OBJECTIVE . Considering an encoder with L layers to encode original data x into disentangled representation z . Let us denote ml as the input features of the l-layer , which are divided into groups of features , i.e. , ml = ∪jmlj , where mlj is the j-th feature subset . We formulate the compositional relation between features from two consecutive layers as follows : ml+1j = Layer ( ml ×wlj ) . Layer can be any neural network layers , e.g. , the commonly used convolution layer in computer vision tasks . The compositional relation is achieved by a composition matrix wl ∈ Rdl×dl+1 , so that features in ml are divided into dl+1 groups after passing through all compositional vectors ( wljs ) . Note that m l+1 j is only related to the subset of features from m l selected by wlj . Similar to Section 2.1 , we assume that the learning process of ml+1 depends on ml through original data x , denoted as the Markov chain : ml → x → ml+1 . It can be written as ml+1 → x → ml according to the conditional independence implied by the Markovity ( Cover , 1999 ) . Consider two output features ml+1i , m l+1 j ∈ ml+1 and their corresponding input feature subsets mli , mlj ∈ ml , we define key notions as follows : Definition 4 . Compositional Disentanglement : mli and mlj are disentangled if I ( mli ; m l j ) = 0 . A disentangled representation of mli and m l j may improve the disentanglement quality between ml+1i and m l+1 j . Similar to Definition 3 , we can achieve compositional disentanglement by mini- mizing I ( mli ; m l j ) . Definition 5 . Compositional Minimal Sufficiency : Assume that the learning process of ml+1j is denoted by the Markov chain : ml+1j → x→ ( mli , m l j ) . Given the original data x , an input feature set mlj for the output feature m l+1 j is minimal sufficient if I ( x ; ml+1j ) = I ( mlj ; m l+1 j ) . For the output feature ml+1j , the input feature set m l j is sufficient and another input feature set m l i is superfluous when mlj is able to capture all information of m l+1 j as well as the original data x . Furthermore , according to Data-Processing Inequality ( DPI ) ( Cover , 1999 ; Achille & Soatto , 2018 ) in the Markov chain , there exists an inequality that : I ( x ; ml+1j ) ≥ I ( ml+1j ; m l i , m l j ) = I ( ml+1j ; m l j ) + I ( ml+1j ; m l i|mlj ) = I ( ml+1j ; m l j ) + I ( ml+1j ; m l i ) − I ( mli ; m l j ) , ( 3 ) where the difference between I ( x ; ml+1j ) and I ( ml+1j ; m l j ) is equivalent to the difference be- tween I ( ml+1j ; m l i ) and I ( mli ; m l j ) . Therefore , matching I ( ml+1j ; m l i ) to I ( mli ; m l j ) can yield a minimal sufficient representation mlj for m l+1 j . Based on the definition of compositional disen- tanglement , we can optimize the minimal sufficiency by forcing I ( ml+1j ; m l i ) to be 0 . To learn disentangled representation via effectively regularizing the compositional feature space , we augment the principled learning objective ( Equation 1 ) with compositional regularizers . Therefore , the compositional learning objective for disentangled representation is defined as follows : L = H ( x̂|mL+1 ) ︸ ︷︷ ︸ sufficient +λ1 L∑ l=2 dl+1∑ j 6=i I ( mli ; m l+1 j ) ︸ ︷︷ ︸ minimal sufficient +λ2 L+1∑ l=2 dl+1∑ j 6=i I ( mli , m l j ) ︸ ︷︷ ︸ disentangled , ( 4 ) where mL+1 denotes the final disentangled representation z . Our intuition is that disentangled learning for compositional feature space could benefit the disentanglement learning for high-level representations .

Summary: The paper presents a VAE variant that disentangles the features of the inference network at every layer, with disentanglement defined in terms of mutual information between features. The approach is implemented as "recursive disentanglement network" based on a switch network (aka Mixture-of-Experts gate, introduced in Shazeer 2017 and used in the switching transformers). The results in dSprites and 3DShapes dataset suggest this variant performs better than well-known disentanglement VAE networks (from few years back) in dSprites and well in 3DShapes (though not the best in all measures). In terms of VAE loss, the approach is presented as a generalization of various other disentangling VAEs.