Original: Learning Switchable Representation with Masked Decoding and Sparse Encoding
Revised: Learning Switchable Representation with Masked Decoding and Sparse Encoding

Original: Anonymous Authors

Original: Abstract
Revised: Abstract

Original: In this study, we explore the unsupervised learning based on private/shared factor decomposition, which decomposes the latent space into private factors that vary only in a specific domain and the shared factors that vary in all domains.
Revised: In this study, we explore the unsupervised learning based on private/shared factor decomposition, which decomposes the latent space into private factors that vary only in a specific domain and the shared factors that vary in all domains.

Original: We study when/how we can force the model to respect the true private/shared factor decomposition that underlies the dataset.
Revised: We study when/how we can force the model to respect the true private/shared factor decomposition that underlies the dataset.

Original: We show that, when we train a masked decoder and an encoder with sparseness regularization in the latent space, we can identify the correct private/shared decomposition up to mixing within each component.
Revised: We show that, when we train a masked decoder and an encoder with sparseness regularization in the latent space, we can identify the correct private/shared decomposition up to mixing within each component.

Original: We empirically confirm this result and study the efficacy of this training strategy as a representation learning method.
Revised: We empirically confirm this result and study the efficacy of this training strategy as a representation learning method.

Original: 1. Introduction
Revised: 1. Introduction

Original: Finding a few common factors that consistently explain different phenomena is a fundamental challenge for unsupervised learning.
Revised: Finding a few common factors that consistently explain different phenomena is a fundamental challenge for unsupervised learning.

Original: For example, consider two datasets of car images.
Revised: For example, consider two datasets of car images.

Original: One consists of synthetic images of 3D objects generated with realistic optical rules (physically based rendering), and the another consists of real images.
Revised: One consists of synthetic images of 3D objects generated with realistic optical rules (physically based rendering), and the another consists of real images.

Original: Their appearance should be different, for example, in the details of images and the texture of the objects.
Revised: Their appearance should be different, for example, in the details of images and the texture of the objects.

Original: However, their physical properties, such as the positioning of objects and the way the image changes with respect to camera positions, shall all be the same between the synthetic and real datasets.
Revised: However, their physical properties, such as the positioning of objects and the way the image changes with respect to camera positions, shall all be the same between the synthetic and real datasets.

Original: Many of these physical properties explain the structure within each dataset; abstractly, they correspond to the axis A, B and C in the visualization in Figure 1.
Revised: Many of these physical properties explain the structure within each dataset; abstractly, they correspond to the axis A, B and C in the visualization in Figure 1.

Original: Unfortunately, however, the mechanism of finding such commonalities has not been fully established.
Revised: Unfortunately, however, the mechanism of finding such commonalities has not been fully established.

Original: For the above datasets, let us consider applying principal component analy-
Revised: For the above datasets, let us consider applying principal component analy-

Original: [Figure]
Revised: [Figure]

Original: sis (PCA).
Revised: sis (PCA).

Original: PCA extracts the subspace in which the variation of the data in the observation space is maximized.
Revised: PCA extracts the subspace in which the variation of the data in the observation space is maximized.

Original: If the variation across the two datasets is larger than the variation within each dataset, PCA would first extract the axis that is specialized in explaining the difference between the two datasets.
Revised: If the variation across the two datasets is larger than the variation within each dataset, PCA would first extract the axis that is specialized in explaining the difference between the two datasets.

Original: However, this might cause some problems because this principal axis might be highly correlated to important factors of variation in each dataset (A, B, C in Figure 1).
Revised: However, this might cause some problems because this principal axis might be highly correlated to important factors of variation in each dataset (A, B, C in Figure 1).

Original: Thus, simply removing the factor of the largest difference might lose the important structural information within each dataset.
Revised: Thus, simply removing the factor of the largest difference might lose the important structural information within each dataset.

Original: Also, because of this entanglement between the first principal component and A, B, and C, it will not be possible to align the rest of the principal components to these structural factors of variations.
Revised: Also, because of this entanglement between the first principal component and A, B, and C, it will not be possible to align the rest of the principal components to these structural factors of variations.

Original: To obtain the desired latent structure, we have to choose an appropriate inductive bias (Locatello et al., 2019).
Revised: To obtain the desired latent structure, we have to choose an appropriate inductive bias (Locatello et al., 2019).

Original: PCA does not work well in the previous example because the objective of PCA is based on the belief that the axis with the largest variance corresponds to the most important structural factor.
Revised: PCA does not work well in the previous example because the objective of PCA is based on the belief that the axis with the largest variance corresponds to the most important structural factor.

Original: One way to resolve this problem is to introduce the shared-private decomposition of latent space (Bousmalis et al., 2016; Liu et al., 2017; Cao et al., 2018;
Revised: One way to resolve this problem is to introduce the shared-private decomposition of latent space (Bousmalis et al., 2016; Liu et al., 2017; Cao et al., 2018;

Original: Peng et al., 2019; Chattopadhyay et al., 2020; Bui et al., 2021).
Revised: Peng et al., 2019; Chattopadhyay et al., 2020; Bui et al., 2021).

Original: In this framework, the domain-private factors consist of the factors that vary only in one individual domain, whereas the shared-domain factors consist of the factors that vary in all domains.
Revised: In this framework, the domain-private factors consist of the factors that vary only in one individual domain, whereas the shared-domain factors consist of the factors that vary in all domains.

Original: By introducing such a structure in the generating process, this strategy imposes the inductive bias that the set of data-describing structural factors includes the features shared by all domains (e.g., C in Figure 1.)
Revised: By introducing such a structure in the generating process, this strategy imposes the inductive bias that the set of data-describing structural factors includes the features shared by all domains (e.g., C in Figure 1.)

Original: However, even if we introduce such an inductive bias, it is not entirely clear if the trained model would identify the true private/shared structure in the dataset.
Revised: However, even if we introduce such an inductive bias, it is not entirely clear if the trained model would identify the true private/shared structure in the dataset.

Original: In this work, we investigate when we can use the private/shared inductive bias to extract the true private/shared factors so that we can use them for improvement on representation learning and extrapolation tasks.
Revised: In this work, we investigate when we can use the private/shared inductive bias to extract the true private/shared factors so that we can use them for improvement on representation learning and extrapolation tasks.

Original: We provide an empirical and basic theoretical result indicating that, when the decoder uses a domain-specific mask function, we can align the feature space to the private-vs-shared decomposition by encouraging the sparsity in the latent space.
Revised: We provide an empirical and basic theoretical result indicating that, when the decoder uses a domain-specific mask function, we can align the feature space to the private-vs-shared decomposition by encouraging the sparsity in the latent space.

Original: We provide an empirical and basic theoretical result indicating that we can align the feature space to the private-vs-shared decomposition by encouraging the sparsity in the latent space.
Revised: We provide an empirical and basic theoretical result indicating that we can align the feature space to the private-vs-shared decomposition by encouraging the sparsity in the latent space.

Original: 2. Method
Revised: 2. Method

Original: In this paper, we consider the multi-domain setting in which the domain membership of each datum is known.
Revised: In this paper, we consider the multi-domain setting in which the domain membership of each datum is known.

Original: Therefore we assume that each observation is a pair ( x , k ) , where x ∈ R d is the raw observation and k ∈ { 1 , . . .
Revised: Therefore we assume that each observation is a pair ( x , k ) , where x ∈ R d is the raw observation and k ∈ { 1 , . . .

Original: , L } is the domain label from which x was drawn.
Revised: , L } is the domain label from which x was drawn.

Original: Our goal is to recover the latent factors z ∈ R m when the true generating process has private factors and shared factors.
Revised: Our goal is to recover the latent factors z ∈ R m when the true generating process has private factors and shared factors.

Original: We elaborate the assumed generating process below.
Revised: We elaborate the assumed generating process below.

Original: In our setting, the latent factors are separated into L + 1 parts as z =
Revised: In our setting, the latent factors are separated into L + 1 parts as z =

Original: [ z 1 , . .
Revised: [ z 1 , . .

Original: z L , z shared ] where each z k ∈ R m k represents the private factor of the domain k and
Revised: z L , z shared ] where each z k ∈ R m k represents the private factor of the domain k and

Original: z shared ∈
Revised: z shared ∈

Original: R m shared represents the shared factor so that m = m shared + (cid:80)
Revised: R m shared represents the shared factor so that m = m shared + (cid:80)

Original: k m k .
Revised: k m k .

Original: For brevity, let I k represent the set of coordinate indices in R m corresponding to the k th domain, and let I shared represent the set of coordinate indices for the shared factors.
Revised: For brevity, let I k represent the set of coordinate indices in R m corresponding to the k th domain, and let I shared represent the set of coordinate indices for the shared factors.

Original: We prepare such a structure in the latent space so that the features that will be trained by the encoder will not entangle the domain-specific features.
Revised: We prepare such a structure in the latent space so that the features that will be trained by the encoder will not entangle the domain-specific features.

Original: The generative process of each observation in the domain k is then written as
Revised: The generative process of each observation in the domain k is then written as

Original: [Equation]
Revised: [Equation]

Original: Based on this model, we train an autoencoder to seek the underlying structure of the observation that aligns with our decomposition.
Revised: Based on this model, we train an autoencoder to seek the underlying structure of the observation that aligns with our decomposition.

Original: Let f : R d
Revised: Let f : R d

Original: → R m be the encoder and g : R m → R d be the decoder.
Revised: → R m be the encoder and g : R m → R d be the decoder.

Original: We define the objective function as
Revised: We define the objective function as

Original: [Equation]
Revised: [Equation]

Original:  where ⊙ denotes the element-wise multiplication, r ( · ) is a sparsity regularizer for z , and β ≥ 0 controls the sparsity.
Revised: [Figure] where ⊙ denotes the element-wise multiplication, r ( · ) is a sparsity regularizer for z , and β ≥ 0 controls the sparsity.

Original: Here, u k ∈ { 0 , 1 } m is the binary mask such that u ki
Revised: Here, u k ∈ { 0 , 1 } m is the binary mask such that u ki

Original: = 1 whenever i ∈ I
Revised: = 1 whenever i ∈ I

Original: k ∪ I shared .
Revised: k ∪ I shared .

Original: The masks u k are not trainable parameters in our objective.
Revised: The masks u k are not trainable parameters in our objective.

Original: If there are two domains to be studied, then the masks shall therefore look 1 1 10 0 01 1 1 1 10 0 01 1 11 1 1 1 1domaindomain

Original: [Figure]

Original: Decoder

Original: [Figure] 1111 0000⦿ [Figure] like u 1 = [ 1 , 0 , 1 ] and u 2 = [ 0 , 1 , 1 ] where 1 indicates the all-one vector with appropriate dimension (Figure 2).
Revised: If there are two domains to be studied, then the masks shall therefore look like u 1 = [ 1 , 0 , 1 ] and u 2 = [ 0 , 1 , 1 ] where 1 indicates the all-one vector with appropriate dimension (Figure 2).

Original: In other words, the mask u k is used to switch between the different domains.
Revised: In other words, the mask u k is used to switch between the different domains.

Original: For example, if the images are colored in the domain k and monocolor in all other domains, then z k may contain color features.
Revised: For example, if the images are colored in the domain k and monocolor in all other domains, then z k may contain color features.

Original: Mask u k is an integral part of the decoder because z k won’t be used in decoding x from z when x is in the domain k ′ ̸
Revised: Mask u k is an integral part of the decoder because z k won’t be used in decoding x from z when x is in the domain k ′ ̸

Original: = k .
Revised: = k .

Original: We note that the mask is a structure that is included in the decoder but not in the encoder.
Revised: We note that the mask is a structure that is included in the decoder but not in the encoder.

Original: Thus, even for an observation x from domain k , our encoder f may have nonzero outputs on I k ′ (thus, z k ′ is not necessarily 0 ) for k ̸ = k ′ .
Revised: Thus, even for an observation x from domain k , our encoder f may have nonzero outputs on I k ′ (thus, z k ′ is not necessarily 0 ) for k ̸ = k ′ .

Original: However, the sparsity regularizer r ( · ) encourages the encoder to parsimoniously use each dimension in the latent space so that there won’t be much redundancy in the representation.
Revised: However, the sparsity regularizer r ( · ) encourages the encoder to parsimoniously use each dimension in the latent space so that there won’t be much redundancy in the representation.

Original: We will show that an appropriately-designed sparsity regularizer can encourage the use of shared factor.
Revised: We will show that an appropriately-designed sparsity regularizer can encourage the use of shared factor.

Original: Identifiability We show that, when the decoding process in (1) is linear, an autoencoder trained in our approach can identify the true private and shared factors up to linear transformations.
Revised: We show that, when the decoding process in (1) is linear, an autoencoder trained in our approach can identify the true private and shared factors up to linear transformations.

Original: Although our proof is done for L = 2 , the strategy should be extendable to the case with L > 2 .
Revised: Although our proof is done for L = 2 , the strategy should be extendable to the case with L > 2 .

Original: Proposition 2.1.
Revised: Proposition 2.1.

Original: Suppose that the datasets in two domains
Revised: Suppose that the datasets in two domains

Original: [Equation]
Revised: [Equation]

Original: are generated from the latents Z k ∈ R n k × m as(cid:18)
Revised: are generated from the latents Z k ∈ R n k × m as

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: where W ∈ R m × d maps ( A, C 1 ) and ( B, C 2 ) invertibly to
Revised: where W ∈ R m × d maps ( A, C 1 ) and ( B, C 2 ) invertibly to

Original: X 1 and X 2 repsectively , while Ker ( C 1 ) =
Revised: X 1 and X 2 repsectively , while Ker ( C 1 ) =

Original: Ker ( C 2 ) and col ( A ) ⊥ col
Revised: Ker ( C 2 ) and col ( A ) ⊥ col

Original: ( C 1 ) as well as col ( B ) ⊥ col
Revised: ( C 1 ) as well as col ( B ) ⊥ col

Original: ( C 2 ) .
Revised: ( C 2 ) .

Original: Consider training another representation X
Revised: Consider training another representation X

Original: =
Revised: =

Original: ˆ Z ˆ W in the same form by minimizing the reconstruction loss as well as the sparsity loss
Revised: ˆ Z ˆ W in the same form by minimizing the reconstruction loss as well as the sparsity loss

Original: [Equation]
Revised: [Equation]

Original: where ∥ · ∥ nzc is the number of nonzero columns in the matrix.
Revised: where ∥ · ∥ nzc is the number of nonzero columns in the matrix.

Original: Then the trained factors ( ˆ A, ˆ B, ˆ C 1 , ˆ C 2 ) identifies ( A, B, C 1 , C 2 )
Revised: Then the trained factors ( ˆ A, ˆ B, ˆ C 1 , ˆ C 2 ) identifies ( A, B, C 1 , C 2 )

Original: upto mixing within each component.
Revised: upto mixing within each component.

Original: What does this proposition mean in terms of (2)?
Revised: What does this proposition mean in terms of (2)?

Original: We note that the relation x = W ⊤ z describes the nonparametric encoder f : x ∈ R d → z ∈ R m and the linear decoder g : z → W T z .
Revised: We note that the relation x = W ⊤ z describes the nonparametric encoder f : x ∈ R d → z ∈ R m and the linear decoder g : z → W T z .

Original: Thus, in the stacked representation Z ∈ R ( n 1 + n 2 ) × m that appears in the relation (4), each column represents the latent coordinate dimension.
Revised: Thus, in the stacked representation Z ∈ R ( n 1 + n 2 ) × m that appears in the relation (4), each column represents the latent coordinate dimension.

Original: If A ∈ R n 1 × m a , B ∈ R n 2 × m b , C k ∈ R n k × m c so that m a + m b + m c = m , then m a corresponds to the private factor dimension for X 1 , m b corresponds to the private factor dimension for X 2 , and m c corresponds to the shared factor dimension( m shared in the previous section).
Revised: If A ∈ R n 1 × m a , B ∈ R n 2 × m b , C k ∈ R n k × m c so that m a + m b + m c = m , then m a corresponds to the private factor dimension for X 1 , m b corresponds to the private factor dimension for X 2 , and m c corresponds to the shared factor dimension( m shared in the previous section).

Original: Thus, the column sparseness regularization (5) encourages the latent space to be sparse in each factor dimension.
Revised: Thus, the column sparseness regularization (5) encourages the latent space to be sparse in each factor dimension.

Original: We note that, in (4), the mask used in (2) is included the form of Z k itself.
Revised: We note that, in (4), the mask used in (2) is included the form of Z k itself.

Original: Also, the factor orthogonality is analogous to independency.
Revised: Also, the factor orthogonality is analogous to independency.

Original: Therefore, this result suggests a learning strategy that encourages independencies as well as sparseness in the encoded latent space while using a masked decoder that respects the private/shared structure.
Revised: Therefore, this result suggests a learning strategy that encourages independencies as well as sparseness in the encoded latent space while using a masked decoder that respects the private/shared structure.

Original: Implementation by VAE We would like to train an autoencoder based on our finding in Proposition 2.1.
Revised: We would like to train an autoencoder based on our finding in Proposition 2.1.

Original: Unfortunately, however, ∥ · ∥ nzc is a variation of the ℓ 0 norm and its direct optimization is NP-hard (Feng et al., 2018) in general.
Revised: Unfortunately, however, ∥ · ∥ nzc is a variation of the ℓ 0 norm and its direct optimization is NP-hard (Feng et al., 2018) in general.

Original: In our experiment, we propose to kill two birds with one stone by combining (2) with β -VAE (Higgins et al., 2016), which is known to encourage the sparsity as well as factor independencies (Rolinek et al., 2019; Zietlow et al., 2021).
Revised: In our experiment, we propose to kill two birds with one stone by combining (2) with β -VAE (Higgins et al., 2016), which is known to encourage the sparsity as well as factor independencies (Rolinek et al., 2019; Zietlow et al., 2021).

Original: We put the details of VAE implementation in Appendix B.
Revised: We put the details of VAE implementation in Appendix B.

Original:  Related Work
Revised: 3. Related Work

Original: The idea of separating the latent space into private and shared ones has a long history.
Revised: The idea of separating the latent space into private and shared ones has a long history.

Original: Canonical correlation analysis (Hotelling, 1936) captures a linear subspace where two random variables are mostly correlated, and the captured subspace can be seen as the space of shared factors.
Revised: Canonical correlation analysis (Hotelling, 1936) captures a linear subspace where two random variables are mostly correlated, and the captured subspace can be seen as the space of shared factors.

Original: (Salzmann et al., 2010) extended this idea by simultaneously finding the private spaces.
Revised: (Salzmann et al., 2010) extended this idea by simultaneously finding the private spaces.

Original: In the literature on domain adaptation and generalization, a lot of studies have used the shared-private representation, based on various approaches such as linear orthogonality (Bousmalis et al., 2016), adversarial frameworks (Bousmalis et al., 2016; Liu et al., 2017;
Revised: In the literature on domain adaptation and generalization, a lot of studies have used the shared-private representation, based on various approaches such as linear orthogonality (Bousmalis et al., 2016), adversarial frameworks (Bousmalis et al., 2016; Liu et al., 2017;

Original: Cao et al., 2018; Peng et al., 2019), and meta learning (Bui et al., 2021).
Revised: Cao et al., 2018; Peng et al., 2019), and meta learning (Bui et al., 2021).

Original: Zhang et al.
Revised: Zhang et al.

Original: (2020) generalized the VAE framework to cover a broad family of graphical models where the variable relationships are described by the mask.
Revised: (2020) generalized the VAE framework to cover a broad family of graphical models where the variable relationships are described by the mask.

Original: As the closest work, Chattopadhyay et al.
Revised: As the closest work, Chattopadhyay et al.

Original: (2020) adopted the binary mask for the shared-private decomposition.
Revised: (2020) adopted the binary mask for the shared-private decomposition.

Original: Unlike us, they incorporate the space separation by learning the masks.
Revised: Unlike us, they incorporate the space separation by learning the masks.

Revised: Binary mask is often used in the literature of OOD and extrapolation

Revised: (Zhou et al., 2019; Huang et al., 2020; Zhang et al., 2021).

Revised: Zhang et al.

Revised: (2020) generalized VAE to cover a broad family of graphical models where the variable relationships are described by the mask.

Original:  In this study, we empirically and theoretically explore how the sparseness regularization helps the identification of the true private/shared structure in an unsupervised manner.
Revised: Similarly, Yang et al. (2021) used the mask to take into account the causal relationships among the latent factors. In this study, we empirically and theoretically explore how the sparseness regularization helps the identification of the true private/shared structure in an unsupervised manner.

Original: Some recent studies have utilized the sparseness in deep generative models as a form of domain knowledge in datasets such as videos (Klindt et al., 2020) and texts (Moran et al., 2021). 
Revised: Some recent studies have utilized the sparseness in deep generative models as a form of domain knowledge in datasets such as videos (Klindt et al., 2020) and texts (Moran et al., 2021). Locatello et al.

Revised: (2020) developed a weakly supervised framework where each observation is given as a pair ( x 1 , x 2 ) that differs only in a few latent factors, and establish an identifiability result.

Revised: Their assumption related to our study because they essentially assume that the feature of z 1 − z 2 is sparse.

Revised: We are however different in that we do not assume that the dataset is aligned as in Locatello et al.

Original: Khemakhem et al.
Revised: Khemakhem et al.

Original: (2020) show that nonlinear VAE is identifiable when the prior distribution cleverly leverages auxiliary information such as class labels.
Revised: (2020) show that nonlinear VAE is identifiable when the prior distribution cleverly leverages auxiliary information such as class labels.

Original: Our case is related to this study if we see the domain label as auxiliary information.
Revised: Our case is related to this study if we see the domain label as auxiliary information.

Original: Unfortunately, we cannot directly compare their analysis to ours because they assume m to be smaller than the number of domains.
Revised: Unfortunately, we cannot directly compare their analysis to ours because they assume m to be smaller than the number of domains.

Original: 3. Experiments
Revised: 4. Experiments

Original: As a proof of concept, we use two datasets; MNIST as domain 1 and Fashion MNIST ( FMNIST ) as domain 2.
Revised: As a proof of concept, we use two datasets; MNIST as domain 1 and Fashion MNIST ( FMNIST ) as domain 2.

Original: To ensure that they share some common factors, weexpand the image size from 28 × 28 to 64 × 64 and apply random 2D translation to both domains.
Revised: To ensure that they share some common factors, we apply random 2D translation to both domains.

Original: As the private factors, we apply random rotation to MNIST and color randomization to FMNIST .
Revised: As the private factors, we apply random rotation to MNIST and color randomization to FMNIST .

Revised: The data augmentations are summarized in Figure 3.

Original: We train our model for 100 epochs by Adam of batch size
Revised: We train our model for 100 epochs by Adam of batch size

Original: 32 and learning rate 2 e − 4 .
Revised: 32 and learning rate 2 e − 4 .

Original: We set β = 3 and the total number of factors to 30 so that we have 10 for each private and shared factor.
Revised: For β -VAE, we set β = 3 and the total number of factors to 30 so that we have 10 for each private and shared factor.

Original: We employ a three-layer convolution network as the encoder and spatial broadcasting (Watters et al., 2019) as the decoder.[Figure]
Revised: We employ a three-layer convolution network as the encoder and spatial broadcasting (Watters et al., 2019) as the decoder.

Original:  Latent Traversal
Revised: 4.1. Latent Traversal

Original: First, we study the active/inactive coordinates of the learned features for the observations in each domain. 
Revised: First, we study the active/inactive coordinates of the learned features for the observations in each domain. Figure 3 shows [Figure]

Original: The right figure shows the domain-wise standard deviation of z for the test data.
Revised: s t a n d a r d d e v i a t i o n MNIST FashionMNIST MNIST FMNIST rot ✓ color ✓ trans ✓ ✓ the domain-wise standard deviation of z for the test data.

Original: It clearly shows 1) roughly half of the features do not vary on both datasets and 2) the encoder learns the active/inactive patterns enforced by the mask used for the decoder.
Revised: It clearly shows 1) roughly half of the features do not vary on both datasets and 2) the encoder learns the active/inactive patterns enforced by the mask used for the decoder.

Original: Figure 3 depicts the latent traversals over the factors that are relevant to the four augmentations we used to generate the dataset.
Revised: Figure 4 depicts the latent traversals over the factors that are relevant to the four augmentations we used to generate the dataset.

Original: As expected, the 2D translation is captured as the shared factors while rotation and colorization are captured as private factors.
Revised: As expected, the 2D translation is captured as the shared factors while rotation and colorization are captured as private factors.

Original: Next, we compare how the masking affects the learned representation.
Revised: Next, we compare how the masking affects the learned representation.

Original: We train the vanilla β -VAE with the same setting and obtain the latent factors that are the most sensitive in reconstruction.
Revised: We train the vanilla β -VAE with the same setting and obtain the latent factors that are the most sensitive in reconstruction.

Original: To select the features based on sensitivity, we approximate the Jacobian of the decoder by the finite difference method and select the top-two factors in terms of the ℓ 2 norm.
Revised: To select the features based on sensitivity, we approximate the Jacobian of the decoder by the finite difference method and select the top-two factors in terms of the ℓ 2 norm.

Original: Figure 4 compares the two-dimensional latent traversals of an encoded MNIST input for our method and β -VAE.
Revised: Figure 5 compares the two-dimensional latent traversals of an encoded MNIST input for our method and β -VAE.

Original: We see that the variations in the proposed latent subspace are closed within MNIST , while β -VAE inadvertently captures the across-domain axis that connects digits and clothing shapes, exhibiting the same faulty behavior of
Revised: We see that the variations in the proposed latent subspace are closed within MNIST , while β -VAE inadvertently captures the across-domain axis that connects digits and clothing shapes, exhibiting the same faulty behavior of

Original: PCA explained in the introduction.
Revised: PCA explained in the introduction.

Original:  Counterfactual Generation
Revised: 4.2. Counterfactual Generation

Original: If the proposed method succeeds in identifying the private and shared factors, we shall be able to use the private factor of domain 1 to alter a dataset in domain 2 to produce counterfactual data without corrupting the contextual meaning of the domain1-private factor.
Revised: If the proposed method succeeds in identifying the private and shared factors, we shall be able to use the private factor of domain 1 to alter a dataset in domain 2 to produce counterfactual data without corrupting the contextual meaning of the domain1-private factor.

Original: To validate this, we study how the image transforms when we change the mask in the reconstruction process.
Revised: To validate this, we study how the image transforms when we change the mask in the reconstruction process.

Original: For example, let u ms = [ 1 , 0 , 1 ] be the mask for MNIST .
Revised: For example, let u ms = [ 1 , 0 , 1 ] be the mask for MNIST .

Original: If we apply u ms to z of an FMNIST input, the appearance of the reconstruction should look like FMNIST while keeping the position, since position is used as the shared factors in the generation process.
Revised: If we apply u ms to z of an FMNIST input, the appearance of the reconstruction should look like FMNIST while keeping the position, since position is used as the shared factors in the generation process.

Original: The third row of Figure 5 shows the reconstruction results with u ms .
Revised: The third row of Figure 6 shows the reconstruction results with u ms .

Original: As expected, while the private information ( FMNIST shapes and color) is overwritten, the shared information (position) is preserved.
Revised: As expected, while the private information ( FMNIST shapes and color) is overwritten, the shared information (position) is preserved.

Original: Similarly, the fourth and fifth rows show the reconstructions when we use the mask that deletes the shared factors and preserves the private factors.
Revised: Similarly, the fourth and fifth rows show the reconstructions when we use the mask that deletes the shared factors and preserves the private factors.

Original: The positional information is lost in these figures, again confirming our hypothesis.
Revised: The positional information is lost in these figures, again confirming our hypothesis.

Original: Domain Adaptation Finally, we check how the obtained representation is useful in a domain adaptation situation.
Revised: Finally, we check how the obtained representation is useful in a domain adaptation situation.

Original: Here we consider an unsupervised case where we have many data in the source domain but a few in the target domain.
Revised: Here we consider an unsupervised case where we have many data in the source domain but a few in the target domain.

Original: We first train the VAE, and use the trained encoder as the fixed feature extractor.
Revised: We first train the VAE, and use the trained encoder as the fixed feature extractor.

Original: We then evaluate the classification accuracy by a linear probe setting, i.e., add a linear layer and minimize the softmax cross-entropy loss.
Revised: We then evaluate the classification accuracy by a linear probe setting, i.e., add a linear layer and minimize the softmax cross-entropy loss.

Original: We use MNIST as the source and FMNIST as the target, where we apply random augmentations in the form of horizontal translation to MNIST and vertical translation to FMNIST .
Revised: We use MNIST as the source and FMNIST as the target, where we apply random augmen- tations in the form of horizontal translation to MNIST and vertical translation to FMNIST .

Original: We apply random rotation to both datasets.
Revised: We apply random rotation to both datasets.

Original: We conducted a FMNIST classification task on two settings: IID and OOD.
Revised: We conducted a FMNIST classification task on two settings: IID and OOD.

Original: In IID, we only apply the random rotation and random horizontal translation to the test set of FMNIST .
Revised: In IID, we only apply the random rotation and random horizontal translation to the test set of FMNIST .

Original: In OOD, We apply all random transformations ( rotation, horizontal translation, vertical translation) to the test set of FMNIST .
Revised: In OOD, We apply all random transformations ( rotation, horizontal translation, vertical translation) to the test set of FMNIST .

Original: Figure 6 compares the accuracy of the vanilla β -VAE and the proposed method for different proportions of FMNIST used in the training of the VAE.
Revised: Figure 7 compares the accuracy of the vanilla β -VAE and the proposed method for different proportions of FMNIST used in the training of the VAE.

Original: The proposed method outperforms the original VAE in both IID and OOD settings, especially when the number of target domain data is small. 
Revised: The proposed method outperforms the original VAE in both IID and OOD settings, especially when the number of target domain data is small. [Figure] ll

Original: [Figure]
Revised: [Figure]

Original: l ll ll
Revised: ll ll

Original: [Figure]
Revised: [Figure]

Original: The test accuracy on FMNIST is plotted against the data proportion # FMNIST / # MNIST in the training set.
Revised: ll

Original: Conclusion In this paper, we explored the way to enforce the private/shared decomposition in the latent space by using the masked decoder and sparsity-regularized encoder.
Revised: [Figure] 5. Discussion In this paper, we explored the way to enforce the private/shared decomposition in the latent space by using the masked decoder and sparsity-regularized encoder.

Original: The empirical and theoretical results in this paper suggest that there is much more room left to research on the way of enforcing the desired structure in unsupervised learning.
Revised: The empirical and theoretical results in this paper suggest that there is much more room left to research on the way of enforcing the desired structure in unsupervised learning.

Original: Future works include the extension of the theory to a nonlinear decoder and a similar investigation of the other forms of structures that are useful in domain adaptation. 
Revised: Future works include the extension of the theory to a nonlinear decoder and a similar investigation of the other forms of structures that are useful in domain adaptation. We now turn to the experiment in Section 4.3.

Original: A. Proof and  the formal statement of 2.
Revised: The augmentations used in the experiments can be seen as synthetically assigned latent factors, which are correlated with a specific data set, as shown in Figure 3. Specifically, we see correlations between MNIST and rotation, and FMNIST and color. This is a typical example of that a spurious correlation occurs, where two independent factors are captured by the model as a single entangled factor (Tr¨auble et al., 2021). In fact, we can see that in Figure 9 the shape information of FMNIST and color information are mixed in the latent space.

Revised: What will happen if a model learns a latent space that reflects a particular correlation?

Revised: One possible problem is that it does not generalize well in OOD situations where different correlations are applied.

Revised: In the example above, color information is indeed captured in the latent space, but it is highly likely that the model did not learn color in a generic sense that can be applied to various objects, but rather a limited color that is applied only to FMNIST .

Revised: The setting of the experiments in Section 4.3 exactly corresponds to that case, and the results suggest that the proposed method is robust to such OOD problems.

Revised: Acknowledgements

Revised: We thank Daisuke Okanohara, Yarin Gal and Peter Abbeel for helpful discussions.

Revised: A. Proof and the formal statement of 2.

Original: Proof of Proposition 2.1.
Revised: Proof of Proposition 2.1.

Original: Let U
Revised: Let U

Original: + denote the pseudo inverse of matrix U .
Revised: + denote the pseudo inverse of matrix U .

Original: Put W ˆ W + = P so that(cid:18)
Revised: Put W ˆ W + = P so that

Original: [Equation]
Revised: [Equation]

Original: We will show that all but diagonal blocks of P are zero.
Revised: We will show that all but diagonal blocks of P are zero.

Original: Since W ˆ W + is invertible on its restriction to both ( A, C ) and ( B, C ) , 
Revised: Since W ˆ W + is invertible on its restriction to both ( A, C ) and ( B, C ) , [Equation]

Original: [Equation]
Revised: [Equation]

Original: Also, the rank agrees with ∥ · ∥ 0 in its sparsest form for A, B , and this applies to C as well because λ > 0 .
Revised: Also, the rank agrees with ∥ · ∥ 0 in its sparsest form for A, B , and this applies to C as well because λ > 0 .

Original: Thus, by the minimality assumption,
Revised: Thus, by the minimality assumption,

Original: [Equation]

Original: [Equation]

Original: +

Original: [Equation]
Revised: [Equation]

Original: λ

Original: [Equation]
Revised: [Equation]

Original: ˆ C 1 ) .

Original: [Equation]

Original: Put
Revised: Put

Original: [Equation]
Revised: [Equation]

Original: − [Equation] Since rank ( C 2 ) = rank ( C 1 ) , we use ∆ c to denote ∆ c k .
Revised: Since rank ( C 2 ) = rank ( C 1 ) , we use ∆ c to denote ∆ c k .

Original: then note that
Revised: then note that

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: Using this relation we obtain
Revised: Using this relation we obtain

Original: [Equation]
Revised: [Equation]

Original: If λ < 2 , this holds only if ∆ c = 0 .
Revised: If λ < 2 , this holds only if ∆ c = 0 .

Original: By (7), this forces ∆ a = 0 , ∆ b = 0 .
Revised: By (7), this forces ∆ a = 0 , ∆ b = 0 .

Original: Recall that we have
Revised: Recall that we have

Original: [Equation]
Revised: [Equation]

Original: Now, by assumption, Ker ( C 1 ) =
Revised: Now, by assumption, Ker ( C 1 ) =

Original: Ker ( C 2 ) .
Revised: Ker ( C 2 ) .

Original: Thus, for exactly k 0 number of nonzero v , we have
Revised: Thus, for exactly k 0 number of nonzero v , we have

Original: [Equation]
Revised: [Equation]

Original: However this implies that
Revised: However this implies that

Original: [Equation]
Revised: [Equation]

Original: But since col ( A ) ⊥ col
Revised: But since col ( A ) ⊥ col

Original: ( C 1 ) as well as col ( B ) ⊥ col
Revised: ( C 1 ) as well as col ( B ) ⊥ col

Original: ( C 2 ) , this would take place only if
Revised: ( C 2 ) , this would take place only if

Original: [Equation]
Revised: [Equation]

Original: Combining this result with ∆ c = 0 , we therefore have
Revised: Combining this result with ∆ c = 0 , we therefore have

Original: [Equation]
Revised: [Equation]

Original: Because the presence of AP Tac only increases the rank in the (17) by orthogonality, rank ( C 1 P Tcc )
Revised: Because the presence of AP Tac only increases the rank in the (17) by orthogonality, rank ( C 1 P Tcc )

Original: = rank ( C 1 ) implies
Revised: = rank ( C 1 ) implies

Original: AP Tac = 0 .
Revised: AP Tac = 0 .

Original: Most importantly, col ( C 1 ) =
Revised: Most importantly, col ( C 1 ) =

Original: col ( ˆ C 1 ) as spaces.
Revised: col ( ˆ C 1 ) as spaces.

Original: Likewise, BP Tbc = 0 and col ( C 2 ) =
Revised: Likewise, BP Tbc = 0 and col ( C 2 ) =

Original: col ( ˆ C 2 ) .
Revised: col ( ˆ C 2 ) .

Original: Finally, recall that
Revised: Finally, recall that

Original: [Equation]
Revised: [Equation]

Original: Baecause col
Revised: Baecause col

Original: C 1 )
Revised: C 1 )

Original: =
Revised: =

Original: col ( ˆ C 1 ) , the column space of ˆ A cannot intersect with col ( C 1 ) by the orthogonality relation col ( ˆ C 1 )
Revised: col ( ˆ C 1 ) , the column space of ˆ A cannot intersect with col ( C 1 ) by the orthogonality relation col ( ˆ C 1 )

Original: ⊥ col ( ˆ A ) .
Revised: ⊥ col ( ˆ A ) .

Original: Thus C 1 P Tca = 0 .
Revised: Thus C 1 P Tca = 0 .

Original: Likewise, C 2 P Tcb = 0 with symmetrical argument.
Revised: Likewise, C 2 P Tcb = 0 with symmetrical argument.

Original: Alltogether,
Revised: Alltogether,

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation][Equation]
Revised: [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original:  (21)

Original: And the claim follows.
Revised: And the claim follows.

Original: B. Detail of VAE implementation
Revised: B. Detail of VAE implementation

Original: VAE employs the variational distribution q ( z | x ) =
Revised: VAE employs the variational distribution q ( z | x ) =

Original: N ( µ , diag( σ 2 )) as a posterior, where the mean and the variance are obtained by the encoders f µ , f σ 2 as µ = f µ ( x ) and σ 2 = f
Revised: N ( µ , diag( σ 2 )) as a posterior, where the mean and the variance are obtained by the encoders f µ , f σ 2 as µ = f µ ( x ) and σ 2 = f

Original: σ 2 ( x ) .
Revised: σ 2 ( x ) .

Original: By using the reparametrization trick z = µ + σϵ with a random Gaussian ϵ ∼ N (0 , I ) , the β VAE objective with the Gaussian likelihood is given by Eq.
Revised: By using the reparametrization trick z = µ + σϵ with a random Gaussian ϵ ∼ N (0 , I ) , the β VAE objective with the Gaussian likelihood is given by Eq.

Original: 2 where the regularization term corresponds to the KL divergence: r ( z ) =
Revised: 2 where the regularization term corresponds to the KL divergence: r ( z ) =

Original: KL( p ( z ) ∥
Revised: KL( p ( z ) ∥

Original: q ( z | x ))
Revised: q ( z | x ))

Original: = 12 (cid:80) mi =1 ( µ 2 i + σ 2 i − log σ 2 i − 1) .
Revised: = 12 (cid:80) mi =1 ( µ 2 i + σ 2 i − log σ 2 i − 1) .

Original: The KL term measures the distance between the standard Gaussian prior p ( z ) =
Revised: The KL term measures the distance between the standard Gaussian prior p ( z ) =

Original: N (0 , I ) and the variational distribution.
Revised: N (0 , I ) and the variational distribution.

Original: Essentially, we simply use this KL term as the choice of r in (2).
Revised: Essentially, we simply use this KL term as the choice of r in (2).

Original: We do not use the masking structure for the encoder itself because the mask in the decoding process would make sure that the reconstruction loss of the domain k is not affected by z ′ k with k ̸ = k ′ .
Revised: We do not use the masking structure for the encoder itself because the mask in the decoding process would make sure that the reconstruction loss of the domain k is not affected by z ′ k with k ̸ = k ′ .

Original: For example, suppose that z 2 is nonzero for the generation of domain 1.
Revised: For example, suppose that z 2 is nonzero for the generation of domain 1.

Original: In this case, z 2 does not contribute to the reconstruction loss in (2), so the derivative of the reconstruction with respect to z 2 is 0 and q ( z 2 | x ) =
Revised: In this case, z 2 does not contribute to the reconstruction loss in (2), so the derivative of the reconstruction with respect to z 2 is 0 and q ( z 2 | x ) =

Original: p ( z 2 ) .
Revised: p ( z 2 ) .

Original: Thus, optimizing the mean and variance parameter with respect domain 1 alone would automatically prefer µ 2 = 0 and σ 22 = 1 as the optimal solution.
Revised: Thus, optimizing the mean and variance parameter with respect domain 1 alone would automatically prefer µ 2 = 0 and σ 22 = 1 as the optimal solution.

Original: C. Full results of latent traversals
Revised: C. Full results of latent traversals

Original: Figure 7 shows the latent traversals of the proposed method for all the latent factors where the standard deviation of z is larger than 0 .
Revised: Figure 8 shows the latent traversals of the proposed method for all the latent factors where the standard deviation of z is larger than 0 .

Original: Similarly, Figure 8 shows the latent traversals of β -VAE.
Revised: Similarly, Figure 9 shows the latent traversals of β -VAE.