Year: 2021

Abstract:
Disentangled feature representation is essential for data-efficient learning. The feature space of deep models is inherently compositional. Existing β-VAE-based methods, which only apply disentanglement regularization to the resulting embedding space of deep models, cannot effectively regularize such compositional feature space, resulting in unsatisfactory disentangled results. In this paper, we formulate the compositional disentanglement learning problem from an informationtheoretic perspective and propose a recursive disentanglement network (RecurD) that propagates regulatory inductive bias recursively across the compositional feature space during disentangled representation learning. Experimental studies demonstrate that RecurD outperforms β-VAE and several of its state-of-the-art variants on disentangled representation learning and enables more data-efficient downstream machine learning tasks.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 COMPOSITIONAL DISENTANGLEMENT LEARNING
text: 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.

heading: 2.1 DISENTANGLEMENT LEARNING OBJECTIVE
text: 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.

heading: 2.2 COMPOSITIONAL OBJECTIVE
text: 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.

heading: 3 RECURSIVE DISENTANGLEMENT NETWORK
text: We now describe a learning method with the goal of optimizing the compositional disentanglement learning objective. This method, called Recursive Disentanglement Network (RecurD), propagates inductive bias (disentanglement) recursively across the compositional feature space.

heading: 3.1 MODEL ARCHITECTURE
text: As shown in Figure 2, RecurD contains an encoder and a decoder to learn the disentangled representation z of data x and to reconstruct x̂, where the encoder contains multiple Recursive Modules and the decoder is a Deconvolutional Neural Network (Zeiler & Fergus, 2014).
The first Recursive Module of the encoder is implemented by a multi-channel convolutional network (Lawrence et al., 1997) to encode the original image x. Following the notation in Section 2.2, the output of the 1-st Recursive Module is denoted as m2, which is also the input of 2-nd Recursive Module. As for the l-th (l ≥ 2) Recursive Module, it contains (1) a Router R to learn a composition matrix wl from the input features ml to decompose ml into subsets; and (2) a Groupof-Encoders (GoE) layer consisting of n encoders to induce the output feature ml+1.
Inspired by the Gate of Mixture-of-Experts (Shazeer et al., 2017; Fedus et al., 2021) on parameter selection for each input, we present a Router with TopK to learn the compositional relation. In
detail, the Router takes ml as input and compute composition matrix wl via learning similarity:
wl = softmax ( TopK ( R ( ml ) , k )) , R ( ml ) = softmax ([ ml ] [ ml ]T) ◦ v and
TopK ( R ( ml ) , k ) = { R ( ml ) ij , if R ( ml ) ij is in the top k elements of R ( ml ) ·j and
−∞, otherwise.
(5)
Here, R (m) denotes the similarity matrix and v is a learning matrix learned by a linear layer. ◦ denotes Hadamard product. TopK is the compositional strategy to determine whether input feature mlij belongs to the input feature set of m l+1 j , and k is a hyperparameter. The GoE layer consists
of dl+1 parallel encoders { Enc1, . . . , Encdl+1 } to generate the output features ml+1, where each encoder is implemented by a convolutional neural network with specific parameters:
ml+1j = Encj ( ml ×wlj ) . (6)
Then, ml+1 is obtained by concatenating the outputs from each encoder, which is converted directly to the input of the l + 1-th Recursive Module. Note that the output of L-th Recursive Module is the learned disentangled representation z. Finally, the Decoder Dec takes z as input to obtain the reconstructed input x̂: x̂ = Dec (z).

heading: 3.2 LEARNING OF RECURD
text: As mentioned in Equation 4, the learning objective of RecurD governs the recursive disentanglement learning across the compositional feature space. As shown by related works, a precise estimation of mutual information in the high dimensional space is important for accurately estimating the loss. According to the recent progresses, the mutual information between two representations can be maximized by using any sample-based differentiable mutual information lower bound. Similar to the work of Federici et al. (2020); Hjelm et al. (2019); Wen et al. (2020), we utilize MINE estimators (Belghazi et al., 2018) to estimate the mutual information. This approach introduces an auxiliary parametric model Estimator Net which is jointly optimized during the training procedure using re-parametrized samples from the posterior distribution.

heading: 4 RELATED WORK
text: β-VAE-based methods introduce various regularization terms and directly apply them on the resulting embedding space. Higgins et al. (2017) proposed β-VAE method by introducing a constraint β over the KL divergence between the inferred distribution qφ(z|x) and its prior–an isotropic unit Gaussian. Burgess et al. (2018) found that the network specializes in the factor that contributes most to a small reconstruction error, with a limited channel capacity. Thus, they proposed the AnnealedVAE, an extension of β-VAE, which gradually adds more latent encoding capacity by enforcing the KL divergence to be at a controllable value C. Kim & Mnih (2018) analyzed the disentanglement performance of β-VAE by breaking down the regularization term into two components. One is Total Correlation (TC) (Watanabe, 1960), KL(q(z)||p(z)), which encourages the marginal distribution of representations to be factorial. The other is mutual information I(x; z), which reduces the amount of information about x stored in z. Based on the decomposition, they proposed FactorVAE to relax the regularization of I(x; z) while directly encouraging independence in the variational distribution
by heavily penalizingKL(q(z)||p(z)). Similarly, Chen et al. (2018) proposed β-TCVAE by decomposing the TC penalty as the dependence among variables (KL(q(z)‖ ∏ j q(zj))) and the distance
between each variable’s posterior and prior ( ∑ j KL(q(zj)‖p(zj))). Since TC is intractable, both FactorVAE and β-TCVAE use approximation methods. However, in unsupervised learning, it is challenging to discard only superfluous information from x, thus high values of β will lead to poor reconstruction and unsatisfactory disentanglement. Another TC-related approach is DIP-VAE (Kumar et al., 2017), whose designers argued that the estimation of TC requires additional parameters and suffers from vanishing gradients. Therefore, DIP-VAE optimizes the moments distance between the aggregated posterior and a factorized prior instead of estimating TC. The above methods improved over β-VAE by applying inductive biases directly to the high-level latent variable space. However, the feature space of deep models is compositional in nature. Existing variants of β-VAE cannot effectively apply disentanglement regularization across such compositional feature space, and yield inferior disentanglement in representation learning. Our approach diverges from prior works by taking a principled information-theoretic approach to formulate and analyze the compositional disentanglement feature structure. We also describe RecurD, which propagates regulatory inductive bias recursively across the compositional feature space during disentangled representation learning.
Recent works on hierarchical VAEs introduce layer-wise disentanglement regularization to learn conditioning structures across multi-layer latent variables, such as VampPrior (Tomczak & Welling, 2018), Ladder VAEs (Sønderby et al., 2016), and NVAE (Vahdat & Kautz, 2020). In these hierarchical model structures, e.g., the cross-layer residual connection structure in NVAE and VampPrior, inter-layer regularization is less of a focus. The latent variables of a preceding layer serve as shared inputs to the next layer, which introduces information redundancy and hence impairs representation disentanglement. In contrast, the proposed compositional objective optimizes the statistical independence of inter-layer and intra-layer latent variables simultaneously, thereby minimizing the information redundancy of inter-layer information sharing and improving disentanglement quality.

heading: 5 EXPERIMENTS
text: This section presents both quantitative and qualitative experiments to evaluate RecurD in terms of disentanglement quality and data efficiency on downstream tasks.

heading: 5.1 PERFORMANCE OF DISENTANGLEMENT LEARNING
text: We compare the disentanglement learning performance of RecurD with β-VAE and its state-of-theart variants, including β-VAE (Higgins et al., 2017), Annealed-VAE (Burgess et al., 2018), FactorVAE (Kim & Mnih, 2018), β-TCVAE (Chen et al., 2018), DIP-VAE (Kumar et al., 2017).
Datasets: We consider two datasets in which each image is obtained by a deterministic function of ground-truth factors: dSprites (Matthey et al., 2017) and 3DShapes (Burgess & Kim, 2018). dSprites contains 737, 280 binary 64 × 64 images of 2D shapes with 5 ground truth factors, i.e., 3 shapes, 6 scales, 40 orientations, 32 x-positions, 32 y-position. 3DShapes contains 480, 000 RGB 64×64×3 images of 3D shapes with 6 ground truth factors, i.e., 4 shapes, 8 scales, 15 orientations, 10 floor hues, 10 wall hues, 10 object hues. For all experiments, we use a 9:1 training to testing data ratio, following earlier work (Kumar et al., 2017; Locatello et al., 2019). More details on network architecture and hyperparameter settings are included in Appendix C.
Evaluation Metrics for Disentanglement: There is no standard metric for evaluating disentanglement (Zhou et al., 2021; Ridgeway & Mozer, 2018), and most existing metrics involve the estimation of a variable-factor matrix relating the factors of variation to the learned representations. In the experiments, we consider three widely used metrics: Separated Attribute Predictability (SAP) (Kumar et al., 2017), Mutual Information Gap (MIG) (Chen et al., 2018) and Disentanglement, Completeness, and Informativeness (DCI) (Eastwood & Williams, 2018). DCI contains three metrics for disentanglement (DCI-D), Completeness (DCI-C) and Informativeness (DCI-I), respectively. Due to the space limitation, we present the results of DCI-D and include the results on DCI-C and DCII in the Appendix D. Overall, these metrics can comprehensively evaluate RecurD from different disentanglement measurements.

heading: 5.1.1 QUANTITATIVE RESULTS
text: For quantitative analysis, we conduct three sets of experiments: 1) evaluating the average performance of disentanglement score via three evaluation metrics; 2) measuring the trade-off among reconstruction, minimal sufficiency and disentanglement with mini-batch samples; and 3) analyzing the influence of compositional objective for disentanglement learning via the three properties.
Table 2 compares the reconstruction error and three widely-used disentanglement metrics of RecurD and five other methods on the dSprites and 3DShapes datasets. Compared with all the baselines, RecurD achieves much lower reconstruction error as well as higher SAP, MIG, and DCI scores in most of the cases. It is important to point out that the reconstruction error of β-VAE increases as β increases (stronger disentanglement regularization), which indicates that β-VAE comprises reconstruction error for disentanglement. However, RecurD achieves better in both reconstruction and disentanglement, demonstrating the superiority of the proposed compositional learning objective and the recursive disentanglement network.
Figure 3 shows scatter plots relating Minimal Sufficiency and Disentanglement to reconstruction error (Sufficiency), where each point represents a mini-batch of data. Here, we compute the minimal sufficiency score using I(x; z) and present it by the area of each point. In all baseline methods, we observe that smaller scatters are correlated with higher reconstruction errors and higher disentanglement scores. The reason is that β-VAE and its variants only place disentanglement regularization on the embedding space of z which imposes a limit on the capacity of the information channel (Higgins et al., 2017), so that features which are important for reconstruction but harmful for disentanglement may lose during training, leading to compromised reconstruction quality. However, RecurD reveals the necessity of disentanglement during feature composition in feed-forward networks and induces effective information compression. Therefore, compared to all the baselines, the representations from RecurD obtain better disentanglement and minimal sufficiency without degrading reconstruction quality.
Figure 4 shows the impact of the number of Recursive Modules in disentanglement learning. We evaluate three RecurD variants on the principled properties, including RecurD 0, RecurD 1 and RecurD 2 with 1, 2 and 3 Recursive Modules, respectively. And we report the values correspond-
ing to the three terms in our compositional learning objective on the training sets. We observe that RecurD 1 and RecurD 2 perform much better than RecurD 0 on the optimization of minimal sufficiency and disentanglement. In addition, during the early stage of training, the performance of RecurD 2 improves much faster than RecurD 1. This experiment demonstrates that the recursive propagation of inductive bias through the feed-forward network improves disentangled representation learning.

heading: 5.1.2 QUALITATIVE RESULTS
text: For qualitative analysis, we present the latent traversals of RecurD on two datasets in Figure 5, in which we vary a single variable learned by an encoder in GoE while keeping all others fixed. For images in 3DShapes, the latent traversals show that RecurD is able to successfully capture all the six factors of variation. For images in dSprites, we observe that RecurD is able to discover x-position, y-position and scale (continuous variables). More importantly, RecurD can, to some extent, discover shape and orientation (discrete variables), which have been proved to be struggling for many other methods (Kumar et al., 2017; Kim & Mnih, 2018;
Locatello et al., 2019). The reason is that β-VAE and its variants encourage independence among variables via controlling the information capacity of z and matching the posterior p(z|x) to an isotropic unit Gaussian. However, learning discrete variables would require using a discrete prior instead of Gaussian (Kim & Mnih, 2018). On the other hand, RecurD is able to model both discrete and continuous factors by directly disentangling based on the three principled properties on the entire compositional feature space, leading to stronger representation capability.

heading: 5.1.3 ABLATION AND PARAMETER DEPENDENCE STUDY
text: In this section, we evaluate the impacts of hyperparameters in both learning objective and model architecture. First, we study the impact of compositional learning objective with varying regularization coefficients (λ1 and λ2) of minimal sufficiency and disentanglement. Then, we evaluate the influence of hyperparameter k — the group size in the GoE of Recursive Module. Figure 6 (RecurD with varying λ1, λ2 and k) shows the scatter plots of disentanglement score (SAP) along with reconstruction error on the test sets of dSprites and 3DShapes.
As shown in Figure 6, larger penalties on both minimal sufficiency and disentanglement yield higher disentanglement score and lower reconstruction error, demonstrating the importance of the minimal sufficiency term and the disentanglement term in Equation 1. Note that RecurD with λ2 = 0 is
reduced to a standard β-VAE with the compositional architecture (as I(x; z) is an lower bound of KL(p(z|x)‖p(z))), which underperforms RecurD, indicating that an explicit penalty on disentanglement is important for disentangled representation learning. As for the group size in the GoE — k, it is not surprising that when k increases, RecurD has lower reconstruction errors, with a slightly inferior disentanglement performance. The reason is that a dense composition of feature space can help the model maintain sufficient information but make it difficult to disentangle latent variables. RecurD with k = 1 results in poor performance on both reconstruction and disentanglement, indicating that over-emphasizing decomposition in the feature space may fail to preserve sufficient information. More results are presented in the Appendix D.

heading: 5.2 PERFORMANCE OF DOWNSTREAM TASKS
text: In this section, we compare the performance of RecurD and five baseline methods by measuring their data efficiency on two downstream tasks: a standard classification task on the MNIST dataset and a domain generalization task on the MNIST-Rotation dataset (Ghifary et al., 2015). MNIST-Rotation is a synthetic dataset consisting of 6 domains, each containing 1, 000 images of the 10 digits randomly selected from the training set of MNIST, with 6 rotation degrees: 0◦, 15◦, 30◦, 45◦, 60◦ and 75◦. For both tasks, we set the training sets of different proportions: 20%, 40%, 60%, 80%, and 100%, to evaluate the classification performance of RecurD and the baselines with different amounts of training data. For the domain generalization problem, we follow the previous
works (Li et al., 2017; Balaji et al., 2018; Du et al., 2020) with the same train-test split strategy and the leave-one-domain-out strategy, i.e., we take the samples from one domain as the target domain for testing, and the samples from the remaining domains as the source domain for training.
Figure 7 reports the average accuracy of different methods on the same test set. We can observe that all methods achieve decent data efficiency on both tasks, i.e., without much performance degradation even with 20% of training data, suggesting that the learning process of disentangled representation can more effectively capture information from the inputs. However, β-VAE and its variants do not achieve satisfactory classification accuracy for either task compared to a single-layer neural network. The reason is that β-VAE and its variants obtain the disentangled representation by limiting the capacity of information channels, thus they tend to only maintain features that contribute more to disentanglement, which may lose informative features that contribute more to classification. Compared to the baselines, RecurD achieves consistently better performance in both tasks, especially on the harder domain generalization task. The reason is that RecurD can learn disentangled representations without sacrificing the reconstruction performance, confirming the hypothesis that compositional disentanglement learning yields better generalization and more data-efficient representations. We believe that the informative disentangled representations emerge when the right balance is achieved between sufficient information preservation and minimal sufficient information learned in a disentangled manner.

heading: 6 CONCLUSION
text: This paper has described a solution to the compositional disentangled representation learning problem. We first presented a general information-theoretic formulation of disentanglement representation learning, and then extended it to the compositional feature space. We then described RecurD, a recursive disentanglement network, which propagates regulatory inductive bias recursively across the compositional feature space. RecurD outperforms β-VAE and its state-of-the-art variants on disentangled representation learning and achieves more data-efficient learning in downstream tasks.

heading: A PROPERTY OF MARKOV CHAIN AND MUTUAL INFORMATION
text: This section lists the related properties of Markov chain and mutual information used in this work.
Markov Chain
Consider three random variables a, b and c coming from a joint distribution p(a,b, c), if the conditional distribution of c depends only on a and is conditionally independent of b, the variables can form a Markov chain in the order denoted as: b → a → c. Specially, the joint distribution can be written as (Cover, 1999):
p(a,b, c) = p(a)p(b|a)p(c|a). (7)
Markovity implies a conditional independence between c and b when a is observed, the reason is:
p(c,b|a) = p(a,b, c) p(a) = p(c,a)p(b|a) p(a) = p(c|a)p(b|a). (8)
By rewriting the conditional independence, the Markov chain of b→ a→ c also implies c→ a→ b as:
p(b, c|a) = p(a,b, c) p(a) = p(b,a)p(c|a) p(a) = p(c|a)p(b|a). (9)
Mutual Information
1. Positivity: I(a;b) ≥ 0; I(a;b|c) ≥ 0. 2. Chain Rule: I(a;b, c) = I(a;b) + I(a; c|b) 3. I(a;b) = H(a)−H(a|b) = 0 if and only if a and b are independent.
I(a;b) = H(a)−H(a | b) = E [ log 1
p(a)
] − E [ log 1
p(a|b) ] = E [ log
p(a|b) p(a) p(b) p(b) ] = E [ log p(a,b)
p(a)b ] = D (p(a,b)‖p(a))× p(b)) ≥ 0
(10)
4. Data-Processing Inequality (DPI): If three random variables a, b and c coming from a joint distribution p(a,b, c) can form a Markov chain in the order denoted as: b→ a→ c, then I(b;a) ≥ I(b; c). (Theorem 2.8.1 in Cover (1999))

heading: B LEARNING OBJECTIVES DECOMPOSITION
text: In this section, we decompose the proposed learning objective to analyze the relationship of learning objectives between Equation 1 and β-VAE-based methods. Let p(x) denote the true distribution of the data, and pφ(z|x) and pθ(x|z) denote the unknown distributions that we need to estimate, parametrized by an encoder with φ and a decoder with θ.
B.1 DECOMPOSITION OF MINIMAL SUFFICIENCY
Minimal Sufficiency is defined as z can encode the minimum amount information of x required to reconstruct x̂, optimized by minimize I(x; z). Inspired from the work of Chen et al. (2018), we can
decompose the minimal sufficiency by assuming the prior p(z) as a factorized Gaussian as:
I(x; z) = Eq(x,z) [KL(q(x, z)‖q(z)p(x))] = Eq(x,z) [ log q(x, z)
q(z)p(x)
] = Eq(x,z) [ log
q(z|x)p(x) q(z)p(x) ] = Ep(x)[Eq(z|x)[log q(z|x)− log q(z) + log p(z)− log p(z) + log
∏ j q(zj)− log ∏ j q(zj)]]
= Eq(x,z)[log q(z|x) q(z) ]− Eq(z) ∑ j log q(zj) p(zj) − Eq(z) [ log q(z)∏ j q(zj) ]
= Ep(x)[KL(q(z|x)‖p(z))]− ∑ j KL(q(zj)‖p(zj))−KL(q(z)‖ ∏ j q(zj))
(11)
The first term is KL divergence between inferred posterior and prior, usually as a penalty term of β-VAEs for disentangling. The second term is the dimension-wise KL divergence, which represents the distance between each latent dimension to the prior. The third term is the Total Correlation referring to the independence among variables.
B.2 ESTIMATION OF DISENTANGLEMENT
Disentanglement is defined as the independence between any two variables, optimized by minimizing ∑ i 6=j I(zi; zj), which is an lower bound of Total Correlation.
The concept of Total Correlation (TC) was described by McGill (1954) and formally formulated by Watanabe (1960) to evaluate the mutual independence of multi-variant variables. Assume a set of z of random variables z1, . . . , zn. The TC in z is expressed as:
C(z1, . . . , zn) = n∑ i=1 H (zi)−H (z1, z2, . . . , zn) , (12)
where H (z1, z2, . . . , zn) denotes the joint entropy. Furthermore, according to the work of Te Sun (1980; 1975), the relation between TC and mutual information can be described as:
C(z) = n∑ i=1 H(zi)−H(z) = − ∑ n≥2 ∆H(z) (13)
The general ∆H(z) is defined as Fano’s multiple mutual information among variables. In the case of n = 2, it reduces to Shannon’s mutual information and in the case of n = 3, it coincides with McGill’s mutual information. The formulation of TC can be simplified as:
C(z1, . . . , zn) = ∑ j 6=i I(zi; zj) + ∑ k 6=j 6=i ∆H(zi; zj ; zk) + . . .+ ∆H(z1; . . . ; zn) (14)
In VAEs, to measure dependence for multiple latent variables, TC is computed as DKL(q(z)‖ ∏ j q(zj)). Therefore, the proposed disentanglement ∑ i 6=j I(zi; zj) is a lower bound of TC.
B.3 RELATION WITH OTHER WORKS
After decomposing the minimal sufficiency and estimating the disentanglement by TC, we can combine two regularizers as:
λ1I(x; z) + λ2 ∑ j 6=i I(zi; zj)
≤ λ1Ep(x)[KL(q(z|x)‖p(z))]− λ12 ∑ j KL(q(zj)‖p(zj))− λ13 KL(q(z)‖ ∏ j q(zj))
+ λ2 ∑ j KL(q(zj)‖p(zj))
= λ1Ep(x)[KL(q(z|x)‖p(z))] + (λ2 − λ12) ∑ j KL(q(zj)‖p(zj))− λ13 KL(q(z)‖ ∏ j q(zj))
= λaEp(x)[KL(q(z|x)‖p(z))] + λb ∑ j KL(q(zj)‖p(zj)) + λc KL(q(z)‖ ∏ j q(zj))
(15) Therefore, by assigning different weights to the decomposed term, we can establish the relationship between the proposed information-theoretic objective and existing β-VAE variants, which will provide insights regarding the capabilities and limitations of existing methods, and further motivate the proposed work.
Specially, the objective function of β-TCVAE is I(x, z + KL(q(z)‖ ∏ j q(zj) +∑
j KL(q(zj)‖p(zj)). if we decompose the I(x; z) of β-TCVAE as: KL(q(z|x)‖p(z)) − KL(q(z)‖ ∏ j q(zj) − ∑ j KL(q(zj)‖p(zj)), the regularizer term of β-TCVAE can be written
as: αKL(q(z|x)‖p(z)) + (β − α)KL(q(z)‖ ∏ j q(zj) + (γ − α) ∑ j KL(q(zj)‖p(zj)). Since α = γ = 1 in β-TCVAE, therefore when λc = 0, the last line of Equation 15 is equivalent to β-TCVAE.

heading: C EXPERIMENTAL DETAILS
text: C.1 IMPLEMENTATION
For all baselines, we use a Convolutional Neural Network for the encoder, and a Deconvolutional Neural Network for the decoder. Specially, for Factor-VAE, we use a 6-layer Multi-Layer Perceptron for the discriminator, with the leaky ReLU as the activation on per layer, and we set γ as 6.4. For β-VAE, we use a set of β = {4, 16, 60}. Annealed-VAE uses γ = 1000 with a linearly increasing C from 0.5 nats to a 25.0 nats). β-TCVAE uses α = γ = 1 and β = 4. DIP-VAE is implemented by the parameters λod = 100 and λd = 10.
For RecurD, we implement RecurD 0 as the same encoder/decoder architecture with all baselines, as shown in Table 3. As for RecurD 1 and RecurD 2, we implement the same decoder architecture with all baselines, and the details about the encoder of RecurD 1 and RecurD 2 are shown in Table 4 and Table 5. Following the previous work, we use negative cross-entropy to compute reconstruction error, which represents sufficiency in our work. As for the computation of minimal sufficiency and disentanglement, we implement the MINE estimator by 4-layer Multi-Layer Perceptron, similar with the work of Federici et al. (2020). As for the hyperparameter in our model, we vary λ1 and λ2 in the set {0.1, 0.2, 0.5, 1, 2, 5, 10, 50} while fixing λ1 = 1 and λ1 = 2 for both datasets. During training, we use Adam optimiser with learning rate 1e − 4, β1 = 0.9, β2 = 0.999 for parameter updates. Specially, we utilize RecurD 1 on dSprites, 3DCars and 3DShapes, and RecurD 2 on CelebA.

heading: D ADDITIONAL EXPERIMENTS
text: This section reports additional experiments, which draw similar conclusions as that of Section 5.
D.1 QUANTITATIVE RESULTS
The first set of quantitative results is the supplementary results of disentanglement score DCI-C and DCI-I on dSprites and 3DShapes. Table 6, Figure 8 and Figure 9 show the results of disentanglement score with reconstruction error, which also show RecurD can achieve a better trade-off between reconstruction and disentanglement.
The second set is the performance comparison results with three additional baselines, i.e, InfoGAN (Chen et al., 2016), Control-VAE and NVAE. InfoGAN (Chen et al., 2016) maximizes the mutual information between the small subset of the latent variables and the observations to increase the interpretability of the latent representation. Control-VAE (Shao et al., 2020) dynamically tunes the weight β on the KL term to achieve a good trade-off between disentanglement and reconstruction quality. NVAE (Vahdat & Kautz, 2020) optimizes high-quality image generation via global correlation capturing across multi-layer latent variables.
We compare with Info-GAN Control-VAE and NVAE on three datasets, including dSprites, 3Dshapes and 3Dcars. The 3DCars (Reed et al., 2015) exhibits different car models from Fidler et al. (2012) under different camera viewpoints. The evaluation method is the disentanglement score MIG. Table 7 show that NVAE and VampPrior can achieve comparable reconstruction error compared to RecurD. However, their disentanglement qualities are not as good as RecurD, because RecurD additionally regularizes the inter-layer information sharing and alleviates the information redundancy of multiple layers. This experiment demonstrates the superiority of the proposed RecurD method in disentanglement compared to existing hierarchical VAEs.
D.2 ABLATION STUDY
In this section, we report the supplementary results of ablation study. The first is to study the impact of compositional learning objective with varying regularization coefficients (λ1 and λ2) of minimal sufficiency and disentanglement. Figure 10 and Figure 11 show the scatter plots of disentanglement scores along with reconstruction error on the test sets of dSprites, with varying λ1, λ2. Figure 12 and Figure 13 show the results on the 3DShapes.
The second set is to evaluate the influence of hyperparameter k — the group size in the GoE of Recursive Module. Figure 14 and Figure 15 show the influence of varying k for disentanglement scores and reconstruction error on the test sets of dSprites. Figure 16 and Figure 17 show the results on the 3DShapes.
The third set is to evaluate the disentanglement metrics on the preceding layers. We compute MIG on mL−1 and mL−2 of RecurD on 3DShapes. As shown in Table 8, high-level representations have higher MIG than that of low-level representations (z > mL−1 > mL−2). These results confirm that
the recursive propagation of inductive bias through the feed-forward network improves disentangled representation learning.
The last set is to evaluate the impact of architecture of Gate-of-Encoders (GoE), we conduct three variant GoEs, including Linear: GoE is implemented as a linear layer with softmax; Fix: GoE is implemented as fix assignment, i.e., we split ml in d+ 1 equal slices; and Att: GoE is implemented as a multi-head attention layer (the head is fixed as 8). As shown in Table 9, this study demonstrates that GoE indeed fits well with the disentanglement process.
D.3 QUALITATIVE RESULTS
In this section, we report the qualitative samples of traversal images on three datasets, including dSprites(Figure 18), 3DShapes(Figure 19) and CelebA(from Figure 22 to Figure 29).
Specifically, on CelebA, we tentatively increase the dimensionality of latent variables of RecurD (from 16 to 32) by doubling the output of the last layer of each encoder from a single latent variable to a 2-dimensional variable. In this setting, our model with 32-dimensional latent variables can produce clear traversal examples on CelebA, as shown in Figure 21.
Furthermore, to further illustrate the efficacy of recursive disentanglement, we conduct another study by further reducing the dimensionality of the latent variables of Beta-VAE and RecurD (from 16 to 8) and comparing the traversal samples generated by the two models. The rationale behind this experimental setup is that a method with better disentanglement capability can more effectively eliminate information redundancy, and pass more non-redundant information across the information bottleneck. As shown in Figure 20, as the dimensionality of the latent variable decreases, β-VAE can only disentangle four factors while RecurD can disentangle six factors with two additional factors.
In addition, our model with 16-dimensional latent variables can disentangle eight factors, including old (Figure 22), azimuth (Figure 23), beard (Figure 24), bangs (Figure 25), face width (Figure 26), brightness (Figure 27), glasses (Figure 28) and skin color (Figure 29).
D.4 EVALUATION OF COMPUTATIONAL COSTS
In this section, we evaluate the computational complexity of our model on 3DShapes and CelebA. The evaluation metrics consist of multiply-accumulate operation (MACs), the model parameters (Params), evaluation time (Eva. time) and converge time (all models converge to the same reconstruction error as β-VAE, denoted as Con. time).
In terms of parameter efficiency, as shown in Table 10, the recursive disentanglement network itself only contains 0.826 million parameters (RecurD 2 w/o MINEs), which are comparable to Beta-VAE. Most of the parameters of RecurD 2 are contributed by using MINE estimator. Specifically, in the initial implementation (RecurD 2/specific MINEs), each pair of outputs of encoders is equipped with a specific MINE model, which contributes 7.214 million parameters. We further optimize the design by using shared MINE within the same feature category, which can effectively reduce the total number of parameters down to 1.325 million (RecurD 2/shared MINEs).

Review summaries:
This paper formulates the disentanglement problem from an information theoretic perspective, but focusing on an objective that encourages a compositional disentangled feature space on the layers that precede the final latents.  With objective, the authors describe a new method using Gate of Mixture-of-Experts to implement the compositional disentangled reconstruction objective. Some of the terms require mutual information estimation, for which they use MINE estimators.  They run experiments across dSprites and 3DShapes and look into reconstruction error and different disentanglement metrics, observing that they method outperform existing beta-VAE-like baselines, without any compositional incentives. They also analyse the loss components with different architectures and observe that degrees of compositionally in the architecture yields better disentanglement.  Finally, they look into some ablations of the regularisation pressure and into data efficiency in downstream tasks.   

The paper proposes a new approach for learning disentangled variational autoencoders. In addition to pushing the sufficiency, minimal sufficiency, and disentanglement of the latent representation, the paper proposes to also regularize those on earlier features in the network. Experiments demonstrate promising results.

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.