R ECURSIVE D ISENTANGLEMENT N ETWORK
A BSTRACT
Disentangled feature representation is essential for data-efﬁcient 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-efﬁcient downstream machine learning tasks.
I NTRODUCTION
Recent progress in machine learning demonstrates the ability to learn disentangled representations is essential for data-efﬁcient 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 𝐦 𝐳 𝐱 𝐱 (cid:3548)
Encoder
Decoder [Figure] 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 * .
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 beneﬁt 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: sufﬁciency, minimal sufﬁciency, and disentanglement.
Theoretical analysis shows that the 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 sufﬁcient 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-efﬁcient learning in downstream machine learning tasks.
C OMPOSITIONAL D ISENTANGLEMENT L EARNING
In this section, we ﬁrst 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. D ISENTANGLEMENT L EARNING O BJECTIVE
The challenge of representation learning can be formulated as ﬁnding a distribution p ( z | x ) that maps original data x ∈ X into a representation z with ﬁxed amount of variables z
= { z 1 , . . .
, z n } (Bengio et al., 2013).
The key intuition of z is to capture minimal sufﬁcient 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 deﬁned as follows:
Deﬁnition 1.
Sufﬁciency: a representation z of x for ˆx is sufﬁcient if I ( x , ˆx )
=
I ( z , ˆx ) .
Whenever NMI ( c i
c j )
=
I ( c i
c j ) /H
( c ) = c and c For the reconstruction task, z is sufﬁcient 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 sufﬁcient property can be optimized by minimizing H ( ˆx
| z ) (Federici et al., 2020; Dubois et al., 2020).
Deﬁnition 2.
Minimal Sufﬁciency: a representation z of x is minimal sufﬁcient if I ( x ; z )
=
I ( z ; ˆx ) .
A minimal sufﬁcient 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 sufﬁcient, 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 sufﬁciency property can be optimized by minimizing I ( x ; z ) .
Deﬁnition 3.
Disentanglement: a representation denoted as z = { z 1 , . . .
, z n } is disentangled if (cid:80) j (cid:54) =
i I ( z i ; z j )
= 0 .
From the deﬁnition of mutual information, I ( z i ; z j ) =
H ( z i ) − H
( z i | z j ) denotes the reduction of uncertainty in z i when z j is observed (Cover, 1999).
If any two components z i and z j are disentangled, changes to z i have no inﬂuence on z j , which means I ( z i ; z j ) = 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 sufﬁciency property.
The principled objective of disentanglement learning is to minimize the following objective:
[Equation]
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 ﬁxed 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 (cid:80) j (cid:54)
=
i I ( z i , z j ) (Te Sun, 1980; 1975), then we assign different weights as follows:
[Equation]
As shown in Table 1, the learning objectives of β -VAE and its four variants can be regarded as speciﬁc 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.
Method
[Equation]
FactorVAE
[Equation]
[Table]
Weight
[Equation]
Relation
[Equation]
2. C OMPOSITIONAL O BJECTIVE
Considering an encoder with L layers to encode original data x into disentangled representation z .
Let us denote m l as the input features of the l -layer, which are divided into groups of features, i.e., m l = ∪ j m l j , where m l j is the j -th feature subset.
We formulate the compositional relation between features from two consecutive layers as follows: m l +1 j = Layer (cid:0) m l × w lj (cid:1) .
Layer can be applicable to commonly used neural network layers, e.g., the convolution layer in computer vision tasks.
The compositional relation is achieved by a composition matrix w l ∈ R d l × d l +1 , so that features in m l are divided into d l +1 groups after passing through all compositional vectors ( w lj s).
Note that m l +
j is only related to the subset of features from m l selected by w lj .
Similar to Section 2.1, we assume that the learning process of m l +1 depends on m l through original data x , denoted as the Markov chain:
m l → x → m l +1 .
It can be written as m l +1 → x → m l according to the conditional independence implied by the Markovity (Cover, 1999).
Consider two output features m l +
i , m l +1 j ∈ m l +1 and their corresponding input feature subsets m li , m lj ∈
m l , we deﬁne key notions as follows:
Deﬁnition 4.
Compositional Disentanglement: m li and m lj are disentangled if I (cid:0) m li ; m lj (cid:1)
= 0 .
A disentangled representation of m li and m lj may improve the disentanglement quality between m l +1 i and m l +1 j .
Similar to Deﬁnition 3, we can achieve compositional disentanglement by minimizing I (cid:0) m li ; m lj (cid:1) .
Deﬁnition 5.
Compositional Minimal Sufﬁciency: Assume that the learning process of m l +
j is denoted by the Markov chain: m l +1 j → x → (cid:0) m li , m lj (cid:1) .
Given the original data x , an input feature set m lj for the output feature m l +1 j is minimal sufﬁcient if I (cid:0) x ;
[Equation]
For the output feature m l +1 j , the input feature set m lj is sufﬁcient and another input feature set m li is superﬂuous when m lj 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:
[Equation]
where the difference between I (cid:0) x ; m l +1 j (cid:1) and I (cid:0) m l +1 j ; m lj (cid:1) is equivalent to the difference between I (cid:0) m l +1 j ; m li (cid:1) and I (cid:0) m li ; m lj (cid:1) .
Therefore, matching I (cid:0) m l +1 j ; m li (cid:1) to I (cid:0) m li ; m lj (cid:1) can yield a minimal sufﬁcient representation m lj for m l +1 j .
Based on the deﬁnition of compositional disentanglement, we can optimize the minimal sufﬁciency by forcing I (cid:0) m l +1 j ; m li (cid:1) to be 0 .
More details can be found in Appendix B.
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 deﬁned as follows:
[Equation]
(cid:124) (cid:123)(cid:122) (cid:125) sufﬁcient + λ 
L (cid:88) l = d l +1 (cid:88)
j (cid:54)
= i I (cid:0) m li ; m l +1 j (cid:1) (cid:124) (cid:123)(cid:122) (cid:125) minimal sufﬁcient + λ  L +1 (cid:88) l
= d l +1 (cid:88) j (cid:54) =
i
I (cid:0) m li , m lj (cid:1) (cid:124) (cid:123)(cid:122) (cid:125) disentangled , (4) where m L +1 denotes the ﬁnal disentangled representation z .
Our intuition is that disentangled learning for compositional feature space could beneﬁt the disentanglement learning for high-level representations.
R ECURSIVE D ISENTANGLEMENT N ETWORK
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.
3. M ODEL ARCHITECTURE
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 ﬁrst 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 m 2 , 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 w l from the input features m l to decompose m l into subsets; and (2) a Groupof-Encoders (GoE) layer consisting of n encoders to induce the output feature m l +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 m l as input and compute composition matrix w l via learning similarity:
[Equation]
(cid:1)
[Equation]
(cid:1)(cid:1)
[Equation]
(cid:0)
[Equation] (cid:1) = [Equation] (cid:16)(cid:2)
[Equation]
(cid:3) (cid:2) m l (cid:3) T (cid:17) ◦ v , and TopK [Equation] (cid:1)
[Equation]
(cid:0)
[Equation] (cid:1) [Equation] is in the top k elements of R (cid:0) m l (cid:1) · j −∞ , otherwise.
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 m lij belongs to the input feature set of m l +1 j , and k is a hyperparameter.
The GoE layer consists of d l +1 parallel encoders (cid:8) Enc 1 , . . .
, Enc d l +1 (cid:9) to generate the output features m l +1 , where each encoder is implemented by a convolutional neural network with speciﬁc parameters:
[Equation]
Then, m l +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 ) .
3. L EARNING OF R ECUR D
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.
; Hjelm et al.
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.
R ELATED W ORK
β -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), 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 penalizing the TC term.
Similarly, Chen et al.
(2018) proposed β -TCVAE by decomposing the TC penalty as the dependence among variables and the distance between each variable’s posterior and prior.
Since TC is intractable, both FactorVAE and β -TCVAE use approximation methods.
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.
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.
E XPERIMENTS
This section presents both quantitative and qualitative experiments to evaluate RecurD in terms of disentanglement quality and data efﬁciency on downstream tasks.
5. P ERFORMANCE OF D ISENTANGLEMENT L EARNING
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).
More experimental results can be found in Appendix E.
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.,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 ﬂoor 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 D.
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 E. Overall, these metrics can comprehensively evaluate RecurD from different disentanglement measurements.
5.1. Q UANTITATIVE R ESULTS
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 sufﬁciency and disentanglement with mini-batch samples; and 3) analyzing the inﬂuence 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 ﬁve 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 Sufﬁciency and Disentanglement to reconstruction error ( Sufﬁciency ), where each point represents a mini-batch of data.
Here, we compute the minimal sufﬁciency 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 sufﬁciency 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-
[Table] [Table]
[Figure]
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 sufﬁciency 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.
5.1. Q UALITATIVE R ESULTS
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 ﬁxed.
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.
5.1. A BLATION AND P ARAMETER D EPENDENCE S TUDY
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 coefﬁcients ( λ 1 and λ 2 ) of minimal sufﬁciency and disentanglement.
Then, we evaluate the inﬂuence 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 sufﬁciency and disentanglement yield higher disentanglement score and lower reconstruction error, demonstrating the importance of the minimal sufﬁciency 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 ) (cid:107) 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 sufﬁcient information but make it difﬁcult 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 sufﬁcient information.
More results are presented in the Appendix E.
5. P ERFORMANCE OF D OWNSTREAM T ASKS
In this section, we compare the performance of RecurD and ﬁve baseline methods by measuring their data efﬁciency on two downstream tasks: a standard classiﬁcation 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, withrotation 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 classiﬁcation 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 efﬁciency 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 classiﬁcation 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 classiﬁcation.
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 sacriﬁcing the reconstruction performance, conﬁrming the hypothesis that compositional disentanglement learning yields better generalization and more data-efﬁcient representations.
We believe that the informative disentangled representations emerge when the right balance is achieved between sufﬁcient information preservation and minimal sufﬁcient information learned in a disentangled manner.
C ONCLUSION
This paper has described a solution to the compositional disentangled representation learning problem.
We ﬁrst 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-efﬁcient learning in downstream tasks.
A P ROPERTY OF M ARKOV C HAIN AND M UTUAL I NFORMATION
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):
[Equation]
Markovity implies a conditional independence between c and b when a is observed, the reason is:
[Equation]
By rewriting the conditional independence, the Markov chain of b → a → c also implies c → a → b as:
[Equation]
Mutual Information
Positivity: I ( a ; b ) ≥ 0 ; I ( a ; b | c ) ≥ 0 .
Chain Rule: I ( a ; b , c ) =
I ( a ; b ) + I ( a ; c | b ) 3.
I ( a ; b ) =
H ( a ) − H ( a | b ) =
if and only if a and b are independent.
[Equation]
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)
B L EARNING O BJECTIVES D ECOMPOSITION
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. D ECOMPOSITION OF M INIMAL S UFFICIENCY
Minimal Sufﬁciency is deﬁned 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 sufﬁciency by assuming the prior p ( z ) as a factorized Gaussian as: [Equation]
[Equation]
The ﬁrst 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. E STIMATION OF D ISENTANGLEMENT
Disentanglement is deﬁned as the independence between any two variables, optimized by minimizing (cid:80) i (cid:54) =
j I ( z i ; z j ) , 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 z 1 , . .
. , z n .
The TC in z is expressed as:
[Equation]
where H ( z 1 , z 2 , . . . , z n ) 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:
[Equation]
The general ∆ H ( z ) is deﬁned 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 simpliﬁed as:
[Equation] [Equation]
In VAEs, to measure dependence for multiple latent variables, TC is computed as
[Equation]
) (cid:17) .
Therefore, the proposed disentanglement (cid:80) i (cid:54) =
j I ( z i ; z j ) is a lower bound of TC.
B. R ELATION WITH O THER W ORKS
After decomposing the minimal sufﬁciency and estimating the disentanglement by TC, we can combine two regularizers as:
[Equation]
Therefore, by assigning different weights to the decomposed term
[Equation]
we can establish the relationship between the proposed information-theoretic objective and existing β -VAE variants
[Equation]
which will provide insights regarding the capabilities and limitations of existing methods
[Equation]
and further motivate the proposed work.
Specially, the objective function of β -TCVAE is I ( x , z ) + KL (cid:16) q
( z ) (cid:107) (cid:81)
j q ( z j ) (cid:17) + (cid:80) j KL ( q ( z j ) (cid:107) p ( z j )) .
If we decompose the I ( x ; z ) of β -TCVAE as: KL ( q ( z | x ) (cid:107) p ( z )) −
KL (cid:16) q ( z ) (cid:107) (cid:81)
j q ( z j ) (cid:17) − (cid:80) j KL ( q ( z j ) (cid:107) p ( z j )) , the regularizer term of β -TCVAE can be written as: αKL ( q ( z | x ) (cid:107) p ( z ))
+( β − α ) KL (cid:16) q ( z ) (cid:107) (cid:81)
j q ( z j ) (cid:17) +( γ − α ) (cid:80) j KL ( q ( z j ) (cid:107) p ( z j )) .
Since α = γ = 1 in β -TCVAE, therefore when λ c = 0 , the last line of Equation 15 is equivalent to β -TCVAE.
C P ROOF OF C OMPOSITIONAL M INIMAL S UFFICIENCY
In this section, we prove the statements reported in the section 2.2 of the paper.
Property: ( P 1 ) : Data-processing inequality in the markov chain: considering the markov chain b → a → c , then I ( b , a ) ≥ I ( b ; c ) .
(Theorem 2.8.1 (Cover, 1999).)
( P 2 ) :
Chain rule for mutual information: I ( a ; b , c ) =
I ( a ; b )+ I ( a ; c | b ) .
(Theorem 2.5.1 (Cover, 1999).)
( P 3 ) :
Decompositon of conditional mutual information.
(Proposition B.1.
(Federici et al., 2020).)
Conditional independence assumptions:
{+The Markov Chain m l +
+}
j
→ x → (cid:0) m li , m lj (cid:1) , implies the conditional independence between m l +1 j and (cid:0) m li , m lj (cid:1) when x is observed.
Proof :
[Equation]
[Equation]
D E XPERIMENTAL D ETAILS
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
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 sufﬁciency in our work.
As for the computation of minimal sufﬁciency and disentanglement, we implement the MINE estimator by 4 -layer Multi-Layer Perceptron, similar with the work of Federici et al.
As for the hyperparameter in our model, we vary λ 1 and λ 2 in the set { 0 . 1 , 0 .
5 , 1 , 2 , 5 , 10 , 50 } while ﬁxing λ 1 = 1 and λ 1 = 2 for both datasets.
During training, we use Adam optimiser with learning rate 1 e − 4 , β 1 = 0 .
9 , β 2 = 0 .
999 for parameter updates.
Specially, we utilize RecurD 1 on dSprites, 3DCars and 3DShapes, and RecurDon CelebA.
Input dSprites (3DShapes): 64 × 64 × 1(3)
Images
Encoder Conv: k=4, s=2, p=1, channel=32, ReLU
Conv: k=4, s=2, p=1, channel=32, ReLU Conv: k=4, s=2, p=1, channel=64, ReLU
Conv: k=4, s=2, p=1, channel=64, ReLU
FC: 128 (256)
FC: 2 × Decoder Input: z ∈ R FC: 128, ReLU
FC: 4 × 4 × 64 , ReLU Deconv: k=4, s=2, p=1, channel=64, ReLU Deconv: k=4, s=2, p=1, channel=32, ReLU Deconv: k=4, s=2, p=1, channel=32, ReLU
Deconv: k=4, s=2, p=1, channel=1(3), ReLU Output: 64 × 64 × 1 (3) Reconst. Images [Table]
Router( v ) -
GoE 4 ×
[Conv: k=4, s=2, p=1, channel=32, ReLU]
Conv: k=4, s=2, p=1, channel=64, ReLU
[Table]
Router( v )
FC, 8, ReLU
GoE 8 × (cid:20)
Conv : k= 4 , s= 2 , p= 1 , channel= 16 , ReLU
FC , 16; FC , 2 ×(cid:21)
[Table]
E A DDITIONAL E XPERIMENTS
This section reports additional experiments, which draw similar conclusions as that of Section 5.
E. Q UANTITATIVE R ESULTS
The ﬁrst 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 [Table] 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. [Table]
[Table]
E. A BLATION S TUDY
In this section, we report the supplementary results of ablation study.
The ﬁrst is to study the impact of compositional learning objective with varying regularization coefﬁcients ( λ 1 and λ 2 ) of minimal sufﬁciency and disentanglement.
Figure E.2 and Figure E.2 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 inﬂuence of hyperparameter k — the group size in the GoE of
Recursive Module.
Figure E.2 and Figure E.2 show the inﬂuence of varying k for disentanglement scores and reconstruction error on the test sets of dSprites.
Figure E.2 and Figure E.2 show the results on the 3DShapes.
The third set is to evaluate the disentanglement metrics on the preceding layers.
We compute MIG on m L − 1 and m L − 2 of RecurD on 3DShapes.
As shown in Table 8, high-level representations have higher MIG than that of low-level representations ( z > m L − 1 > m L − 2 ).
These results conﬁrm 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 ﬁx assignment, i.e., we split m l in d + 1 equal slices; and Att : GoE is implemented as a multi-head attention layer (the head is ﬁxed as 8 ).
As shown in Table 9, this study demonstrates that GoE indeed ﬁts well with the disentanglement process. [Table]
[Table]
E. Q UALITATIVE R ESULTS
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 20 to Figure 25).
Speciﬁcally, 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.
E
E VALUATION OF C OMPUTATIONAL C OSTS
In this section [Table] 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 efﬁciency, as shown in Table
[Table]
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.
Speciﬁcally, in the initial implementation (RecurD 2/speciﬁc MINEs), each pair of outputs of encoders is equipped with a speciﬁc 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).
[Table]
RecurD3.
RecurD3. RecurD 2 w/o MINEs 2. RecurD 2/shared MINEs 3. beta-VAE 3.
Factor-VAE 3.