Title: Flow Factorized Representation Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Generative modelling

latent disentanglement

variational autoencoders
Abstract: A prominent goal of representation learning research is to achieve representations which are factorized in a useful manner with respect to the ground truth factors of variation. The fields of disentangled and equivariant representation learning have approached this ideal from a range of complimentary perspectives; however, to date, most approaches have proven to either be ill-specified or insufficiently flexible to effectively separate all realistic factors of interest in a learned latent space. In this work, we propose an alternative viewpoint on such structured representation learning which we call Flow Factorized Representation Learning, and demonstrate it to learn both more efficient and more usefully structured representations than existing frameworks. Specifically, we introduce a generative model which specifies a distinct set of latent probability paths that define different input transformations. Each latent flow is generated by the gradient field of a learned potential following dynamic optimal transport. Our novel setup brings new understandings to both \textit{disentanglement} and \textit{equivariance}. We show that our model achieves higher likelihoods on standard representation learning benchmarks while simultaneously being closer to approximately equivariant models. Furthermore, we demonstrate that the transformations learned by our model are flexibly composable and can also extrapolate to new data, implying a degree of robustness and generalizability approaching the ultimate goal of usefully factorized representation learning.

Flow Factorized Representation Learning

Abstract

A prominent goal of representation learning research is to achieve representations which are factorized in a useful manner with respect to the ground truth factors of variation. The fields of disentangled and equivariant representation learning have approached this ideal from a range of complimentary perspectives; however, to date, most approaches have proven to either be ill-specified or insufficiently flexible to effectively separate all realistic factors of interest in a learned latent space. In this work, we propose an alternative viewpoint on such structured representation learning which we call Flow Factorized Representation Learning, and demonstrate it to learn both more efficient and more usefully structured representations than existing frameworks. Specifically, we introduce a generative model which specifies a distinct set of latent probability paths that define different input transformations. Each latent flow is generated by the gradient field of a learned potential following dynamic optimal transport. Our novel setup brings new understandings to both disentanglement and equivariance . We show that our model achieves higher likelihoods on standard representation learning benchmarks while simultaneously being closer to approximately equivariant models. Furthermore, we demonstrate that the transformations learned by our model are flexibly composable and can also extrapolate to new data, implying a degree of robustness and generalizability approaching the ultimate goal of usefully factorized representation learning.

1 Introduction

Developing models which learn useful representations of data has become an increasingly important focus in the machine learning community [5, 55]. For example, Large Language Models such as GPT [9] rely on an extensive pre-training phase to learn valuable representations, enabling quick finetuning on a diversity of tasks. However, a precise definition of what makes an ideal representation is still debated. One line of work has focused on 'disentanglement' of the underlying ground truth generative factors [5, 35, 13]. In general, the definition of 'disentanglement' often refers to learning and controlling statistically independent factors of variation [5, 36]. Over the years, many disentanglement methods have been proposed, including axis-aligned single-dimensional manipulation [35, 13], linear multi-dimensional traversals [78, 77, 90, 66], and, more recently, dynamic non-linear vector-based traversals [84, 79]. Although these methods have been met with significant success (and even linked to single-neuron brain activity [37, 91]), there are known theoretical limitations which make them ill-specified, including the presence of topological defects [7]. This has limited their deployment beyond toy settings.

Another line of work has focused on developing representations which respect symmetries of the underlying data in their output space [15, 36]. Specifically, equivariant representations are those for which the output transforms in a known predictable way for a given input transformation.

They can be seen to share many similarities with disentangled representations since an object undergoing a transformation which preserves its identity can be called a symmetry transformation [36]. Compared with disentanglement methods, equivariant networks are much more strictly defined, allowing for significantly greater control and theoretical guarantees with respect to the learned transformation [16, 50, 73, 20, 39]. However, this restriction also limits the types of transformations to which they may be applied. For example, currently only group transformations are supported, limiting real-world applicability. To avoid this caveat, some recent attempts propose to learn general but approximate equivariance from disentangled representations [49, 45, 79].

In this work, we provide an alternative viewpoint at the intersection of these two fields of work which we call Flow Factorized Representation Learning. Fig. 1 depicts the high-level illustration of our method. Given k different transformations p k ( x t | x 0 ) in the input space, we have the corresponding latent probabilistic path ∫ z 0 , z t q ( z 0 | x 0 ) q k ( z t | z 0 ) p ( x t | z t ) for each of the transformations. Each latent flow path q k ( z t | z 0 ) is generated by the gradient field of some learned potentials ∇ u k following fluid mechanical dynamic Optimal Transport (OT) [4]. Our framework allows for novel understandings of both disentanglement and equivariance . The definition of disentanglement refers to the distinct set of tangent directions ∇ u k that follow the OT paths to generate latent flows for modeling different factors of variation. The concept of equivariance in our case means that the two probabilistic paths, i.e., p k ( x t | x 0 ) in the image space and ∫ z 0 , z t q ( z 0 | x 0 ) q k ( z t | z 0 ) p ( x t | z t ) in the latent space, would eventually result in the same distribution of transformed data.

We build a formal generative model of sequences and integrate the above latent probability evolution as condition updates of the fac-

{z

p(xt

Z0|xo)

torized sequence distribution. Based on the continuity equation, we derive a proper flow of probability density for the time evolution of both the prior and posterior. To perform inference, we approximate the true posterior of latent variables and train the parameters as a Variational Autoencoder (VAE) [47]. When the transformation type k is not observed ( i.e., available as a label), we treat k as another latent variable and incorporate its posterior into our framework by learning it from sequences. Extensive experiments and thorough analyses have been conducted to show the effectiveness of our method. For example, we demonstrate empirically that our representations are usefully factorized, allowing flexible composability and generalization to new datasets. Furthermore, we show that our methods are also approximately equivariant by demonstrating that they commute with input transformations through the learned latent flows. Ultimately, we see these factors combine to yield the highest likelihood on the test set in each setting. Code is publicly available at https://github.com/KingJamesSong/latent-flow .

2 The generative model

In this section, we first introduce our generative model of sequences and then describe how we perform inference over the latent variables of this model in the next section.

2.1 Flow factorized sequence distributions

The model in this work defines a distribution over sequences of observed variables. We further factorize this distribution into k distinct components by assuming that each observed sequence is generated by one of the k separate flows of probability mass in latent space. Since in this work we

k

k

Zo

Z1

ZT

Zo

Z1

ZT

X1

X2

XT

X1

X2

XT

N

N

model discrete sequences of observations ¯ x = { x 0 , x 1 . . . , x T } , we aim to define a joint distribution with a similarly discrete sequence of latent variables ¯ z = { z 0 , z 1 . . . , z T } , and a categorical random variable k describing the sequence type (observed or unobserved). Explicitly, we assert the following factorization of the joint distribution over T timesteps:

Here p ( k ) is a categorical distribution defining the transformation type, p ( x t | z t ) asserts a mapping from latents to observations with Gaussian noise, and p ( z 0 ) = N (0 , 1) . A plate diagram of this model is depicted through the solid lines in Fig. 2.

2.2 Prior time evolution

To enforce that the time dynamics of the sequence define a proper flow of probability density, we compute the conditional update p ( z t | z t -1 , k ) from the continuous form of the continuity equation: ∂ t p ( z ) = -∇ · ( p ( z ) ∇ ψ k ( z )) , where ψ k ( z ) is the k 'th potential function which advects the density p ( z ) through the induced velocity field ∇ ψ k ( z ) . Considering the discrete particle evolution corresponding to this density evolution, z t = f ( z t -1 , k ) = z t -1 + ∇ z ψ k ( z t -1 ) , we see that we can derive the conditional update from the continuous change of variables formula [69, 11]:

In this setting, we see that the choice of ψ ultimately determines the prior on the transition probability in our model. As a minimally informative prior for random trajectories, we use a diffusion equation achieved by simply taking ψ k = -D k log p ( z t ) . Then according to the continuity equation, the prior evolves as:

where D k is a constant coefficient that does not change over time. The density evolution of the prior distribution thus follows a constant diffusion process. We set D k as a learnable parameter which is distinct for each k .

3 Flow factorized variational autoencoders

To perform inference over the unobserved variables in our model, we propose to use a variational approximation to the true posterior, and train the parameters of the model as a VAE. To do this, we parameterize an approximate posterior for p ( z 0 | x 0 ) , and additionally parameterize a set of K

functions u k ( z ) to approximate the true latent potentials ψ ∗ . First, we will describe how we do this in the setting where the categorical random variable k is observed (which we call the supervised setting), then we will describe the model when k is also latent and thus additionally inferred (which we call the weakly supervised setting).

3.1 Inference with observed k (supervised)

When k is observed, we define our approximate posterior to factorize as follows:

We see that, in effect, our approximate posterior only considers information from element x 0 ; however, combined with supervision in the form of k , we find this is sufficient for the posterior to be able to accurately model full latent sequences. In the limitations section we discuss how the posterior could be changed to include all elements { x t } T 0 in future work.

Combing Eq. (4) with Eq. (1), we can derive the following lower bound to model evidence (ELBO):

Substituting and simplifying, Eq. (5) can be re-written as

We thus see that we have an objective very similar to that of a traditional VAE, except that our posterior and our prior now both have a time evolution defined by the conditional distributions.

3.2 Inference with latent k (weakly supervised)

When k is not observed, we can treat it as another latent variable, and simultaneously perform inference over it in addition to the sequential latent ¯ z . To achieve this, we define our approximate posterior and instead factorize it as

Following a similar procedure as in the supervised setting, we derive the new ELBO as

We see that, compared with Eq. (5), only one additional KL divergence term D KL [ q γ ( k | ¯ x ) || p ( k )] ] is added. The prior p ( k ) is set to follow a categorical distribution, and we apply the Gumbel-SoftMax trick [43] to allow for categorical re-parameterization and sampling of q γ ( k | ¯ x ) .

3.3 Posterior time evolution

As noted, to approximate the true generative model which has some unknown latent potentials ψ k , we propose to parameterize a set of potentials as u k ( z , t ) = MLP ([ z ; t ]) and train them through the ELBOs above. Again, we use the continuity equation to define the time evolution of the posterior, and thus we can derive the conditional time update q ( z t | z t -1 , k ) through the change of variables formula. Given the function of the sample evolution z t = g ( z t -1 , k ) = z t -1 + ∇ z u k , we have:

Converting the above continuous equation to the discrete setting and taking the logarithm of both sides gives the normalizing-flow-like density evolution of our posterior:

The above relation can be equivalently derived from the continuity equation ( i.e., ∂ t q ( z ) = -∇· ( q ( z ) ∇ u k ) ). Notice that we only assume the initial posterior q ( z 0 | x 0 ) follows a Gaussian distribution. For future timesteps, we do not pose any further assumptions and just let the density evolve according to the sample motion.

3.4 Ensuring optimal transport of the posterior flow

As an inductive bias, we would like each latent posterior flow to follow the OT path. To accomplish this, it is known that when the gradient ∇ u k satisfies certain PDEs, the evolution of the probability density can be seen to minimize the L 2 Wasserstein distance between the source distribution and the distribution of the target transformation. Specifically, we have:

Theorem 1 (Benamou-Brenier Formula [4]) . For probability measures µ 0 and µ 1 , the L 2 Wasserstein distance can be defined as

where the density ρ and the velocity v satisfy:

The optimality condition of the velocity is given by the generalized Hamilton-Jacobi (HJ) equation ( i.e., ∂ t u + 1 / 2 ||∇ u || 2 ≤ 0 ). The detailed derivation is deferred to the supplementary. We thus encourage our potential to satisfy the HJ equation with an external driving force as

Here we use another MLP to parameterize the external force f ( z , t ) and realize the negativity constraint by setting f ( z , t ) = -MLP ([ z ; t ]) 2 . Notice that here we take the external force as learnable MLPs simply because we would like to obtain a flexible negativity constraint. The MLP architecture is set the same for both u ( z , t ) and f ( z , t ) . To achieve the PDE constraint, we impose a Physics-Informed Neural Network (PINN) [67] loss as

where the first term restricts the potential to obey the HJ equation, and the second term limits u ( z t , t ) to return no update at t =0 , therefore matching the initial condition.

4 Experiments

This section starts with the experimental setup, followed by the main qualitative and quantitative results, then proceeds to discussions about the generalization ability to different composability and unseen data, and ends with the results on complex real-world datasets.

22

3:

1大

333323m

3

333333

555

6

8

8

4.1 Setup

Datasets. We evaluate our method on two widely-used datasets in generative modeling, namely MNIST [54] and Shapes3D [10]. For MNIST [54], we manually construct three simple transformations including Scaling, Rotation, and Coloring. For Shapes3D [10], we use the self-contained four transformations that consist of Floor Hue, Wall Hue, Object Hue, and Scale.

Besides these two common benchmarks, we take a step further to apply our method on Falcol3D and Isaac3D [61], two complex large-scale and real-world datasets that contain sequences of different transformations. Falcol3D consists of indoor 3D scenes in different lighting conditions and viewpoints, while Isaac3D is a dataset of various robot arm movements in dynamic environments.

Baselines. We mainly compare our method with SlowVAE [49] and Topographic VAE (TVAE) [45]. These two baselines could both achieve approximate equivariance. Specifically, TVAE introduces some learned latent operators, while SlowVAE enforces the Laplacian prior p ( z t | z t -1 ) = ∏ αλ / 2Γ(1 /α ) exp( -λ | z t,i -z t -1 ,i | α ) to sequential pairs. Within the disentanglement literature, our method is compared with the supervised PoFlow [79] which adopts a wave-like potential flow for sample evolution, and the unsupervised β -VAE [35] and FactorVAE [46] which encourage independence between single latent dimensions. Finally, the vanilla V AE is used as a controlled baseline.

Metrics. We use the approximate equivariance error E k and the log-likelihood of transformed data log p ( x t ) as the evaluation protocols. The equivariance error is defined as E k = ∑ T t =1 | x t -Decode ( z t ) | where z t = z 0 + ∑ T t =1 ∇ z u k . For TVAE, the latent operator is changed to Roll ( z 0 , t ) . For unsupervised disentanglement baselines [35, 46] and SlowVAE [49], we carefully select the latent dimension and tune the interpolation range to attain the traversal direction and range that correspond to the smallest equivariance error. Since the vanilla V AE does not have the corresponding learned transformation in the latent space, we simply set ∇ z u k = 0 and take it as a lower-bound baseline. For all the methods, the results are reported based on 5 runs.

Notice that the above equivariance error is defined in the output space. Another reasonable evaluation metric is instead measuring error in the latent space as E k = ∑ T t =1 | Encode ( x t ) -z t | . We see the first evaluation method is more comprehensive as it further involves the decoder in the evaluation.

4.2 Main Results

Qualitative results. Fig. 3 and 4 display decoded images of the latent evolution on MNIST [54] and Shapes3D [10], respectively. On both datasets, our latent flow can perform the target transformation precisely during evolution while leaving other traits of the image unaffected. In particular, for the weakly-supervised setting, the decoded images ( i.e., the bottom rows of Fig. 3 and 4) can still reproduce the given transformations well and it is even hard to visually tell them apart from the generated images under the supervised setting. This demonstrates the effectiveness of the weaklysupervised setting of our method, and implies that qualitatively our latent flow is able to learn the sequence transformations well under both supervised and weakly-supervised settings.

Quantitative results. Tables 1 and 2 compare the equivariance error and the log-likelihood on MNIST [54] and Shapes3D [10], respectively. Our method learns the latent flows which model the transformations precisely, achieving the best performance across datasets under different supervision settings. Specifically, our method outperforms the previous best baseline by 69 . 74 on average in the equivariance error and by 32 . 58 in the log-likelihood on MNIST. The performance gain is also consistent on Shapes3D: our method surpasses the second-best baseline by 291 . 70 in the average equivariance error and by 120 . 42 in the log-likelihood. In the weakly-supervised setting, our method also achieves very competitive performance, falling behind that of the supervised setting in the average equivariance error slightly by 6 . 22 on MNIST and by 67 . 88 on Shapes3D.

4.3 Discussion

Extrapolation: switching transformations. In Fig. 5 we demonstrate that, empowered by our method, it is possible to switch latent transformation categories mid-way through the latent evolution

Wall Hue

Object Hue

Wall Hue

Rotation

Coloring

Rotation

Scaling

Coloring

Scaling

Object Hue

a

a

̸

and maintain coherence. That is, we perform z t = z t -1 + ∇ z u k for t ≤ T / 2 and then change to z t = z t -1 + ∇ z u j where j = k for t > T / 2 . As can be seen, the factor of variation immediately changes after the transformation type is switched. Moreover, the transition phase is smooth and no other attributes of the image are influenced.

All Transformations

Scale + Object Hue

Wall Hue + Floor Hue

Wall Hue + Object Hue

Extrapolation: superposing transformations. Besides switching transformations, our method also supports applying different transformations simultaneously, i.e., consistently performing z t = z t -1 + ∑ K k ∇ z u k during the latent flow process. Fig. 6 presents such exemplary visualizations of superposing two and all transformations simultaneously. In each case, the latent evolution corresponds to simultaneous smooth variations of multiple image attributes. This indicates that our method also generalizes well to superposing different transformations.

Notice that we only apply single and separate transformations in the training stage. Switching or superposing transformations in the test phase can be thus understood as an extrapolation test to measure the generalization ability of the learned equivariance to novel compositions.

Equivariance generalization to new data. We also test whether the learned equivariance holds for Out-of-Distribution (OoD) data. To verify this, we validate our method on a test dataset that is different from the training set and therefore unseen to the model. Fig. 7 displays the exemplary visualization results of the V AE trained on MNIST [54] but evaluated on dSprites [59]. Although the reconstruction quality is poor, the learned equivariance is still clearly effective as each transformation still operates as expected: scaling, rotation, and coloring transformations from top to bottom respectively.

4.4 Results on Complex Real-world and Large-scale Datasets

Table 3 and 4 compare the equivariance error of our methods and the representative baselines on Falcol3D and Isaac3D, respectively. Notice that the values are much larger than previous datasets due to the increased image resolution. Our method still outperforms other baselines by a large margin

OoD Input

Output

Encoder

Decoder

Traversal Results

OoD Input

Output

Encoder

Decoder

Traversal Results

OoD Input

Output

Encoder

Decoder

Traversal Results

OoD Input

Output

Encoder

Decoder

Traversal Results

OoD Input

Output

Encoder

Decoder

Traversal Results

OoD Input

Output

Encoder

Decoder

Traversal Results

and achieves reasonable equivariance error. Fig. 8 displays the qualitative comparisons of our method against other baselines. Our method precisely can control the image transformations through our latent flows. Overall, the above results demonstrate that our method can go beyond the toy setting and can be further applied to more complex real-world scenarios.

More visualization results of exemplary latent flows are kindly referred to in the supplementary.

5 Related work

Disentangled representation learning. The idea of learning disentangled representation dates back to factorizing non-redundant input patterns [74] but is recently first studied by InfoGAN [13] and β -VAE [35]. InfoGAN [13] achieves disentanglement by maximizing the mutual information between a subset of latent dimensions and observations, while β -VAE [35] induces the factorized posterior q ( z ) by penalizing the Total Correlation (TC) through an extra hyper-parameter β> 1 controlling the strength of the KL divergence. Following infoGAN, many attempts have been made to facilitate the discovery of semantically meaningful traversal directions through regularization [33, 42, 89, 34, 100, 66, 77, 90, 98, 84, 99, 78, 62]. The follow-up research of β -VAE mainly explored different methods to factorize the aggregated posterior [22, 25, 52, 46, 12, 44, 96, 23, 76, 58, 80, 28]. More recently, some works proposed to discover meaningful directions of diffusion models in the bottleneck of denoising networks [53, 64, 95, 41]. The previous literature mainly considers disentanglement as learning different transformations per dimension or per linear direction. Our method generalizes this concept to learning a distinct tangent bundle ∇ u k that moves every latent sample via dynamic OT.

We see the most similar method to ours is the work of [79]. In [79], the authors also apply the gradient of a potential function to move the latent code. However, their potentials are restricted to obey the wave equations, which do not really correspond to the OT theory. Also, they do not consider the posterior evolution but instead use the loss || z t -Encode ( x t ) || 2 to match the latent codes. By contrast, we propose a unified probabilistic generative model that encompasses the posterior flow that follows dynamic OT, the flow-like time evolution, and different supervision settings.

Equivariant neural networks. Afunction is said to be an equivariant map if it commutes with a given transformation, i.e., T ′ [ f ( x )] = f ( T [ x ]) where T and T ′ represent operators in different domains. Equivariance has been considered a desired inductive bias for deep neural networks as this property can preserve geometric symmetries of the input space [38, 75, 56, 57, 1]. Analytically equivariant

Wall Color

TVAE

PoFlow

Robot X-move

PoFlow

TVAE

Lightning Intensity

Camera X-Pos

Camera Z-Pos

Lightning Intensity

Ours

Ours

Ours

networks typically enforce explicit symmetry to group transformations in neural networks [16, 17, 68, 93, 92, 85, 31, 39]. Another line of research proposed to directly learn approximate equivariance from data [21, 18, 49, 20, 45]. Our framework re-defines approximate equivariance by matching the latent probabilistic flow to the actual path of the given transformation in the image space.

Optimal transport in deep learning. There is a vast literature on OT theory and applications in various fields [87, 88]. Here we mainly highlight the relevant applications in deep learning. The pioneering work of [19] proposed a light-speed implementation of the Sinkhorn algorithm for fast computation of entropy-regularized Wasserstein distances, which opened the way for many differentiable Sinkhorn algorithm-based applications [32, 29, 14, 27, 51]. In generative modeling, the Wasserstein distance is often used to minimize the discrepancy between the data distribution and the model distribution [2, 81, 72, 65]. Inspired by the fluid mechanical interpretation of OT [4], some normalizing flow methods [69, 24, 48] considered regularizing the velocity fields to satisfy the HJ equation, thus matching the dynamic OT plan [94, 30, 83, 63, 60]. Our method applies PINNs [67] to directly model generalized HJ equations in the latent space and uses the gradient fields of learned potentials to generate latent flows, which also aligns to the theory of dynamic fluid mechanical OT.

6 Conclusion

In this paper, we introduce Flow Factorized Representation Learning which defines a set of latent flow paths that correspond to sequences of different input transformations. The latent evolution is generated by the gradient flow of learned potentials following dynamic optimal transport. Our setup reinterprets the concepts of both disentanglement and equivariance . Extensive experiments demonstrate that our model achieves higher likelihoods on standard representation learning benchmarks while simultaneously achieving smaller equivariance error. Furthermore, we show that the learned latent transformations generalize well, allowing for flexible composition and extrapolation to new data.

7 Limitations

For flexibility and efficiency, we use PINN [67] constraints to model the HJ equation. However, such PDE constraints are approximate and not strictly enforced. Other PDE modeling approaches include accurate neural PDE solvers [40, 8, 70] or other improved PINN variants such as competitive PINNs [97] and robust PINNs [3]. Also, when infering with observed k , we change the posterior from q (¯ z | ¯ x , k ) to q (¯ z | x 0 , k ) because we assume k contains sufficient information of the whole sequence. To keep the posterior definition of q (¯ z | ¯ x , k ) , we need to make q ( z t ) also a function of x t . This can be achieved either by changing the potential to u ( z t -1 , x t , t -1) or modifying the external driving force to f ( z t -1 , x t , t -1) . Nonetheless, we see these modifications would make the model less flexible than our current formulations as the element x t might be needed during inference.

Acknowledgments and Disclosure of Funding

This work was supported by the MUR PNRR project FAIR (PE00000013) funded by the NextGenerationEU, by the PRIN project CREATIVE (Prot. 2020ZSL9F9), by the EU H2020 project AI4Media (No. 951911), and by the Bosch Center for Artificial Intelligence.

References

Pulkit Agrawal, Joao Carreira, and Jitendra Malik. Learning to see by moving. In ICCV , 2015.

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In ICML , 2017.

Chandrajit Bajaj, Luke McLennan, Timothy Andeen, and Avik Roy. Recipes for when physics fails: Recovering robust learning of physics informed neural networks. Machine Learning: Science and Technology , 2023.

Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the mongekantorovich mass transfer problem. Numerische Mathematik , 2000.

Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE T-PAMI , 2013.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013.

Diane Bouchacourt, Mark Ibrahim, and Stéphane Deny. Addressing the topological defects of disentanglement via distributed operators. arXiv preprint arXiv:2102.05623 , 2021.

Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. ICLR , 2022.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. NeurIPS , 2020.

Chris Burgess and Hyunjik Kim. 3d shapes dataset. https://github.com/deepmind/3dshapes-dataset/, 2018.

Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations, 2019.

Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. Isolating sources of disentanglement in variational autoencoders. NeurIPS , 2018.

Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. NeurIPS , 2016.

Lenaic Chizat, Pierre Roussillon, Flavien Léger, François-Xavier Vialard, and Gabriel Peyré. Faster wasserstein distance estimation with the sinkhorn divergence. NeurIPS , 2020.

Taco S Cohen and Max Welling. Transformation properties of learned visual representations. ICLR , 2015.

Taco S Cohen and Max Welling. Group equivariant convolutional networks. In ICML . PMLR, 2016.

Taco S Cohen and Max Welling. Steerable cnns. ICLR , 2017.

Marissa Connor, Gregory Canal, and Christopher Rozell. Variational autoencoder with learned latent structure. In AISTATS . PMLR, 2021.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. NeurIPS , 2013.

Neel Dey, Antong Chen, and Soheil Ghafurian. Group equivariant generative adversarial networks. ICLR , 2021.

Nichita Diaconu and Daniel Worrall. Learning to convolve: A generalized weight-tying approach. In ICML . PMLR, 2019.

Nat Dilokthanakul, Pedro AM Mediano, Marta Garnelo, Matthew CH Lee, Hugh Salimbeni, Kai Arulkumaran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture variational autoencoders. ICLR , 2016.

Zheng Ding, Yifan Xu, Weijian Xu, Gaurav Parmar, Yang Yang, Max Welling, and Zhuowen Tu. Guided variational autoencoder for disentanglement learning. In CVPR , 2020.

Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. ICLR , 2017.

Emilien Dupont. Learning disentangled joint continuous and discrete representations. NeurIPS , 2018.

Cian Eastwood and Christopher KI Williams. A framework for the quantitative evaluation of disentangled representations. In ICLR , 2018.

Marvin Eisenberger, Aysim Toker, Laura Leal-Taixé, Florian Bernard, and Daniel Cremers. A unified framework for implicit sinkhorn differentiation. In CVPR , 2022.

Benjamin Estermann and Roger Wattenhofer. Dava: Disentangling adversarial variational autoencoder. ICLR , 2023.

Jean Feydy, Thibault Séjourné, François-Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In AISTATS , 2019.

Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ode: the world of jacobian and kinetic regularization. In ICML . PMLR, 2020.

Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In ICML . PMLR, 2020.

Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio. Learning with a wasserstein loss. NeurIPS , 2015.

Lore Goetschalckx, Alex Andonian, Aude Oliva, and Phillip Isola. Ganalyze: Toward visual definitions of cognitive image properties. In ICCV , 2019.

Erik Härkönen, Aaron Hertzmann, Jaakko Lehtinen, and Sylvain Paris. Ganspace: Discovering interpretable gan controls. NeurIPS , 2020.

Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. ICLR , 2016.

Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexander Lerchner. Towards a definition of disentangled representations. arXiv preprint arXiv:1812.02230 , 2018.

Irina Higgins, Le Chang, Victoria Langston, Demis Hassabis, Christopher Summerfield, Doris Tsao, and Matthew Botvinick. Unsupervised deep learning identifies semantic disentanglement in single inferotemporal face patch neurons. Nature communications , 2021.

Geoffrey E Hinton, Alex Krizhevsky, and Sida D Wang. Transforming auto-encoders. In ICANN . Springer, 2011.

Emiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In ICML . PMLR, 2022.

Jun-Ting Hsieh, Shengjia Zhao, Stephan Eismann, Lucia Mirabella, and Stefano Ermon. Learning neural pde solvers with convergence guarantees. ICLR , 2019.

Vincent Tao Hu, David W Zhang, Meng Tang, Pascal Mettes, Deli Zhao, and Cees GM Snoek. Latent space editing in transformer-based flow matching. In ICML Workshop , 2023.

Ali Jahanian, Lucy Chai, and Phillip Isola. On the" steerability" of generative adversarial networks. ICLR , 2020.

Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. ICLR , 2017.

Yeonwoo Jeong and Hyun Oh Song. Learning discrete and continuous factors of data via alternating disentanglement. In ICML , 2019.

T Anderson Keller and Max Welling. Topographic vaes learn equivariant capsules. NeurIPS , 2021.

Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In ICML , 2018.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. ICLR , 2014.

Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. NeurIPS , 2018.

David Klindt, Lukas Schott, Yash Sharma, Ivan Ustyuzhaninov, Wieland Brendel, Matthias Bethge, and Dylan Paiton. Towards nonlinear disentanglement in natural data with temporal sparse coding. ICLR , 2021.

Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In ICML . PMLR, 2020.

Soheil Kolouri, Navid Naderializadeh, Gustavo K Rohde, and Heiko Hoffmann. Wasserstein embedding for graph learning. ICLR , 2021.

Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. Variational inference of disentangled latent concepts from unlabeled observations. ICLR , 2018.

Mingi Kwon, Jaeseok Jeong, and Youngjung Uh. Diffusion models already have a semantic latent space. ICLR , 2023.

Yann LeCun. The mnist database of handwritten digits. 1998. URL http://yann.lecun.com/exdb/ mnist/ .

Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature , 2015.

Chen-Yu Lee, Saining Xie, Patrick Gallagher, Zhengyou Zhang, and Zhuowen Tu. Deeply-supervised nets. In AISTATS . PMLR, 2015.

Karel Lenc and Andrea Vedaldi. Understanding image representations by measuring their equivariance and equivalence. In CVPR , 2015.

Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-supervised disentanglement without compromises. In ICML . PMLR, 2020.

Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dsprites: Disentanglement testing sprites dataset, 2017. URL https://github.com/deepmind/dsprites-dataset/ .

Kirill Neklyudov, Daniel Severo, and Alireza Makhzani. Action matching: A variational method for learning stochastic dynamics from samples. ICML , 2023.

Weili Nie, Tero Karras, Animesh Garg, Shoubhik Debnath, Anjul Patney, Ankit B Patel, and Anima Anandkumar. Semi-supervised stylegan for disentanglement learning. In ICML , 2020.

James Oldfield, Christos Tzelepis, Yannis Panagakis, Mihalis A Nicolaou, and Ioannis Patras. Panda: Unsupervised learning of parts and appearances in the feature maps of gans. ICLR , 2023.

Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. Ot-flow: Fast and accurate continuous normalizing flows via optimal transport. In AAAI , 2021.

Yong-Hyun Park, Mingi Kwon, Junghyo Jo, and Youngjung Uh. Unsupervised discovery of semantic latent directions in diffusion models. arXiv preprint arXiv:2302.12469 , 2023.

Giorgio Patrini, Rianne van den Berg, Patrick Forre, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. In UAI , 2020.

William Peebles, John Peebles, Jun-Yan Zhu, Alexei Efros, and Antonio Torralba. The hessian penalty: A weak prior for unsupervised disentanglement. In ECCV , 2020.

M. Raissi, P. Perdikaris, and G.E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics , 2019.

Siamak Ravanbakhsh, Jeff Schneider, and Barnabas Poczos. Equivariance through parameter-sharing. In ICML . PMLR, 2017.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In ICML . PMLR, 2015.

Jack Richter-Powell, Yaron Lipman, and Ricky TQ Chen. Neural conservation laws: A divergence-free perspective. NeurIPS , 2022.

Karl Ridgeway and Michael C Mozer. Learning deep disentangled embeddings with the f-statistic loss. NeurIPS , 2018.

Tim Salimans, Han Zhang, Alec Radford, and Dimitris Metaxas. Improving gans using optimal transport. ICLR , 2018.

Victor Garcia Satorras, Emiel Hoogeboom, Fabian B Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows. NeurIPS , 2021.

Jürgen Schmidhuber. Learning factorial codes by predictability minimization. Neural computation , 1992.

Uwe Schmidt and Stefan Roth. Learning rotation-aware features: From invariant priors to equivariant descriptors. In CVPR , 2012.

Huajie Shao, Shuochao Yao, Dachun Sun, Aston Zhang, Shengzhong Liu, Dongxin Liu, Jun Wang, and Tarek Abdelzaher. Controlvae: Controllable variational autoencoder. In ICML . PMLR, 2020.

Yujun Shen and Bolei Zhou. Closed-form factorization of latent semantics in gans. In CVPR , 2021.

Yue Song, Nicu Sebe, and Wei Wang. Orthogonal svd covariance conditioning and latent disentanglement. IEEE T-PAMI , 2022.

Yue Song, Andy Keller, Nicu Sebe, and Max Welling. Latent traversals in generative models as potential flows. In ICML . PMLR, 2023.

Chang-Yu Tai, Ming-Yao Li, and Lun-Wei Ku. Hyperbolic disentangled representation for fine-grained aspect extraction. In AAAI , 2022.

Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto-encoders. ICLR , 2018.

Jakub Tomczak and Max Welling. Vae with a vampprior. In AISTATS . PMLR, 2018.

Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In ICML . PMLR, 2020.

Christos Tzelepis, Georgios Tzimiropoulos, and Ioannis Patras. WarpedGANSpace: Finding non-linear rbf paths in GAN latent space. In ICCV , 2021.

Elise Van der Pol, Daniel Worrall, Herke van Hoof, Frans Oliehoek, and Max Welling. Mdp homomorphic networks: Group symmetries in reinforcement learning. NeurIPS , 2020.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. NeurIPS , 2017.

Cédric Villani. Optimal transport: old and new , volume 338. Springer, 2009.

Cédric Villani. Topics in optimal transportation , volume 58. American Mathematical Soc., 2021.

Andrey Voynov and Artem Babenko. Unsupervised discovery of interpretable directions in the gan latent space. In ICML , 2020.

Yuxiang Wei, Yupeng Shi, Xiao Liu, Zhilong Ji, Yuan Gao, Zhongqin Wu, and Wangmeng Zuo. Orthogonal jacobian regularization for unsupervised disentanglement in image generation. In ICCV , 2021.

James CR Whittington, Will Dorrell, Surya Ganguli, and Timothy Behrens. Disentanglement with biological constraints: A theory of functional cell types. In ICLR , 2023.

Daniel Worrall and Max Welling. Deep scale-spaces: Equivariance over scale. NeurIPS , 2019.

Daniel E Worrall, Stephan J Garbin, Daniyar Turmukhambetov, and Gabriel J Brostow. Harmonic networks: Deep translation and rotation equivariance. In CVPR , 2017.

Liu Yang and George Em Karniadakis. Potential flow generator with l 2 optimal transport regularity for generative models. IEEE TNNLS , 2020.

Tao Yang, Yuwang Wang, Yan Lv, and Nanning Zh. Disdiff: Unsupervised disentanglement of diffusion probabilistic models. arXiv preprint arXiv:2301.13721 , 2023.

Cagatay Yildiz, Markus Heinonen, and Harri Lahdesmaki. Ode2vae: Deep generative second order odes with bayesian neural networks. NeurIPS , 2019.

Qi Zeng, Spencer H Bryngelson, and Florian Schäfer. Competitive physics informed networks. ICLR , 2023.

Jiapeng Zhu, Ruili Feng, Yujun Shen, Deli Zhao, Zheng-Jun Zha, Jingren Zhou, and Qifeng Chen. Low-rank subspaces in gans. NeurIPS , 2021.

Jiapeng Zhu, Yujun Shen, Yinghao Xu, Deli Zhao, and Qifeng Chen. Region-based semantic factorization in gans. ICML , 2022.

Xinqi Zhu, Chang Xu, and Dacheng Tao. Learning disentangled representations with latent variation predictability. In ECCV , 2020.

A Supplementary Material

A.1 Pseudo codes

Fig. 9 displays the Pytorch implementation for training our flow-factorized VAE under different supervision settings. Here we omit the computation of HJ PDEs for concisity.

A.2 Implementation details

Common settings. During the training stage, we randomly sample one single transformation at each iteration. The batch size is set to 128 for both datasets. We use Adam optimizer and the learning rate is set as 1 e -4 for all the parameters. The encoder consists of four stacked convolution layers with the activation function ReLU, while the decoder is comprised of four stacked transposed convolution layers. For the prior evolution, the diffusion coefficient D k is initialized with 0 and we set it as a learnable parameter for distinct k . For MLPs that parameterize the potential u ( z , t ) and the force f ( z , t ) , we use the sinusoidal positional embeddings [86] to embed the timestep t , and use linear layers for embedding the latent code z . Tanh gates are applied as the activation functions of the MLPs . All the experiments are run on a single NVIDIA Quadro RTX 6000 GPU.

MNIST. The input images are of the size 28 × 28 . The sequence of each transformation contains 9 states of variations. The scaling transformation scales the image from 1 . 0 up to 1 . 8 times. The rotation transformation rotates the object by maximally 80 degrees, and the coloring transformation adjusts the image hue from 0 to 340 degrees. The model is trained for 90 , 000 iterations.

Shapes3D. The input images are resized to 64 × 64 . Each transformation sequence consists of 8 images. The model is also trained for 90 , 000 iterations.

Falcol3D and Isaac3D. The input images are in the resolution of 128 × 128 . We use the selfcontained transformations of the datasets, which mainly comprise variations of lighting conditions and viewpoints in indoor 3D scenes for Falcolr3D, and different robot arm movements in dynamic environments for Isaac3D.

Weakly-supervised setting. For the Gumbel-Softmax trick, we re-parameterize q γ ( k | ¯ x ) by

where x i is the category prediction, g i is the sample drawn from Gumbel distributions, and τ is the small temperature to make softmax behave like argmax . We take the 'hard' binary prediction in the forward pass and use the straight-through gradient estimator [6] during backpropagation. The temperature τ is initialized with 1 and is gradually reduced to 0 . 05 with the annealing rate 3 e -5 .

Baselines. For the disentanglement methods, we largely enrich the original MNIST dataset by adding the transformed images of the whole sequence. This makes it possible for both β -VAE and FactorVAE to learn the given transformations in an unsupervised manner. For tuning the interpolation range, we start from the initial value z i and traverse till the appropriate bound which is selected from the range [ -5 , 5] with the interval of 0 . 1 .

A.3 Disentanglement metrics

There are many traditional disentanglement metrics [71, 26, 12], but they are designed for singledimension traversal methods. These metrics assume and require that each latent dimension is responsible for one semantic and manipulating single dimensions of the latent variable would involve distinct output transformations. However, for the recent disentanglement methods including ours [77, 84, 79], there emerges a more realistic disengagement setting: all the latent dimensions are perturbed by vectors for meaningful output variations. When it comes to these vector-based disentanglement methods, their scores of disentanglement metrics would drop considerably and cannot be compared with those single-dimension baselines.

Nonetheless, certain disentanglement metrics such as VP scores [100] can be leveraged as they do not pose any assumptions on the latent space but only require image pairs [ x 0 , x T ] of different transformations for evaluation. The VP metric adopts the few-shot learning setting (using 1% or 10% of the dataset as the training set) and takes a lightweight neural network for learning to classify image pairs [ x 0 , x T ] of different attributes. The generalization ability ( i.e., validation accuracy) can be thus regarded as a reasonable surrogate for the disentanglement ability. Table 5 and 6 present the VP scores of all the baseline methods on MNIST and Shapes3D. To ensure a fair comparison, for FactorVAE and β -VAE, we choose the dimensions with the lowest equivariance errors to generate image pairs of different transformations. Our method outperforms the previous disentanglement baselines and achieves superior performance on the VP scores. This indicates that our flow-factorized VAE has better disentanglement ability.

A.4 Ablation studies

Impact of different priors. We use diffusion equations to model the prior evolution as random particle movement. It would also be intriguing to choose other priors commonly used in the V AE literature, such as Standard Gaussian (SG) priors N (0 , 1) , mixture of Gaussian (MoG) priors ∑ w i N ( µ i , σ 2 i ) , and VAMP priors [82] which average aggregated posterior of N pseudo-inputs as 1 / N ∑ n q ( z n ) . Table 7 presents the equivariance error of different priors on MNIST. Among these priors, our diffusion equations achieve the best performance. This meets our assumption that modeling the prior evolution

as a diffusion process suits more the random motion. Nonetheless, we see that the performance gap between each baseline is narrow, which somehow implies that the impact of different priors is limited.

Impact of different PDEs. We apply the generalized HJ (GHJ) equation as the PINN constraint in order to achieve dynamic OT. It would be also interesting to try other commonly used PDEs. We compare our GHJ with the ordinary HJ (OHJ) equation, the Fokker Planck (FP) equation, and the heat equation. Table 8 compares the equivariance error of PDEs on MNIST. Our GHJ and OHJ equations achieve the best performance as they both satisfy the condition of dynamic OT. This empirical evidence indicates that the OT theory can indeed model better latent flow paths. Moreover, our GHJ outperforms the OHJ by a slight margin. We attribute this advantage to the external driving force f ( z , t ) which gives us more flexibility and dynamics in modeling the velocity fields ∇ u k .

Impact of different K . We conduct an ablation study on the impact of the number of transformations on MNIST and present the evaluation results in Table 9. As indicated above, in general, the performance is not affected by the number of transformations being applied. The fluctuation of the results when K varies can be sufficiently negligible. We expect that this is because the transformations are learned by distinct potentials (which are implemented as K different MLPs). Each flow evolves along with the gradient field ∇ u on the potential landscape u ; having multiple latent flows defined on different potential landscapes therefore does not interfere with each other.

Impact of different sequence lengths. Regarding the impact of sequence lengths, if the sequence is longer and has larger variations, generally the equivariance error would be worse. To better illustrate this point, we perform an ablation study on MNIST and present the results in Table 10. Specifically, we change the sequence length from 9 to 12 , which increases the extent of scaling from maximally 1 . 8 times to maximally 2 . 1 times, and increases the rotation angle from maximally 80 degrees to maximally 110 degrees. As can be observed, the equivariance error gets larger when the sequence becomes longer and the variations are larger. Notice that even for the longer sequence, our method still outperforms other baselines with shorter sequences.

A.5 HJ equations as dynamic optimal transport

We now turn to introduce why HJ equations could minimize the Wasserstein distance. As stated in [4], the L 2 Wasserstein distance can be re-formulated in the fluid mechanical interpretation as

where the density satisfies the continuity equation ( ∂ t ρ = -∇· ( ρ ( x, t ) v ( x, t ) ). If we introduce the momentum m ( x, t ) = ρ ( x, t ) v ( x, t ) and two Lagrange multipliers u and λ , the Lagrangian function of the Wasserstein distance would be:

Scaling

Rotation

Coloring

h

5

5

5

5

3

3

3

6

6

06

5

5

5

8

8

7

q

q

where the second term is the equality constraint, and the third term is an equality constraint with a slack variable s . Using integration by parts formula, the above equation can be re-written as

Based on the set of Karush-Kuhn-Tucker (KKT) conditions ( ∂ m L = 0 , ∂ u L = 0 , ∂ ρ L = 0 , and λ ≥ 0 ), we would have:

where the first condition indicates that the gradient ∇ u acts as the velocity field, and the third condition implies the optimal solution is given by the generalized HJ equation:

We thus apply the generalized HJ equation ( i.e., ∂ t u + 1 2 ||∇ u || 2 ≤ 0 ) as the constraints. We further use an extra negative force because this would give more dynamics for modeling the posterior flow.

A.6 More visualizations

Fig. 10, 11, and 12 display more visualization results of the latent evolution on MNIST, Shapes3D, Falcol3D and Isaac3D, respectively. Across all the datasets, our method presents precise control of the given transformations. Fig. 13 and 14 show more latent evolution results of switching transformations (top) and combining transformations (bottom) on MNIST and Shapes3D, respectively. Fig. 15 also visualizes a few examples of superposing and switching transformation on Falcol3D and Isaac3D. Our latent flows learn to compose or switch different transformations precisely and flexibly.

Wal Hue

Object Hue

Scale

Lightning Intensity

Camera Y-Pos

Camera Z-Pos

Lightning X-Dir

Lightning Z-Dir

Camera X-Pos

Robot Y-Move

Object Color

Lightning Y-Dir

Wall Color

Lightning Intensity

正

P

Rotation

Scaling

Scaling

Rotation

Rotation

Coloring

Coloring

Rotation

Scaling

Rotation

Scaling

Coloring

Coloring

Scaling

Rotation

Coloring

Coloring

Rotation

Scaling

Coloring

Scaling

Rotation

Rotation

Scaling

Scaling

Coloring

Coloring

Rotation

Scaling

Rotation

Scaling + Rotation

Scaling + Coloring

Rotation + Coloring

88

3

333

333

000

0

00

222

MMM

3DM

686

868

888

5555554

55

5

S

4

4

4

4

?

D

Q

B

9

G

Object Hue

Wall Hue

Wall Hue

Object Hue

Object Hue

Scale

Wall Hue

Scale

Object Hue

Scale

Wall Hue

Wall Hue

Scale

Wall Hue

Object Hue

Wall Hue

Object Hue

Scale

Wall Hue

Wall Hue

Scale

Wall Hue

Scale + Floor Hue

Wall Hue + Object Hue

Wall Hue + Floor Hue

Scale + Object Hue

Lightning Intensity + Object Size

Lightning Intensity + Object Color

Wall Color + Object Color

Wall Color + Robot Y-Move

Camera Y-Pos

Lightning X-Dir

Camera Y-Pos

Camera Z-Pos

REP

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a new generative model called Flow Factorized Representation Learning that learns a distribution over sequences of observed variables using which one can learn both disentangled and equivariant representations. The model is based on the idea of representing transformations in the input space as latent flow paths generated by the gradient fields of learned potentials. In this way, disentanglement can be read from the set of tangent directions, while equivariance can be read by analyzing whether the probabilistic path in the input space and the corresponding latent path result in the same latent representation. The paper includes detailed formal derivation of the model which is realized as a Variational Autoencoder. Experiments are performed on MNIST and 3DShapes datasets.\n\nSTRENGTHS:\n- Solid model derivation\n- Strong experimental results showing empirical equivariance and disentanglement, though on basic datasets.\n\nWEAKNESSES:\nW1) It is not easy to follow the derivations, and it is not clear how the training is performed at the end. It would be helpful to at least add a pseudocode of the algorithm.\n\nW2) It is not clear to me what exactly are the input sequences of observations in the experimental section. In other words, how big is T and how to you obtain the intermediate input datapoints between x0 and xT?\n\nW3) It is unclear what is the complexity of the model.\n\nW4) Experiments are performed on very simple datasets. It is even argued in line 33 that prior methods have not gone beyond toy settings but this work does not do that either.\n\nQUESTIONS:\n1. What is the sequence length and how does this affect the results? Could you please add more details on how you obtain the input sequences?\n2. How would this method generalize to more complex data? While the method demonstrates strong results on MNIST and 3DShapes, these datasets are rather basic.\n3. In the Extrapolation case, how do you aggregate representations from each of the flows?\n\nLIMITATIONS:\n- Limitations are discussed.\n- No code provided.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
This paper proposes a flow factorized representation learning approach for disentanglement and equivariant representation learning in a learned latent space. This work introduces a generative model that defines a set of latent probability paths to specify different input transformations and each latent flow is generated following dynamic optimal transport. Results show that the proposed framework is able to learn both more efficient and more usefully structured representations and thus achieve improved performance on representation learning benchmarks.

STRENGTHS:
1. This work provides a new understanding of the disentanglement and equivariant, which are important problems for representation learning research.
2. This paper proposes flow factorized representation learning which is a new and reasonable viewpoint for structured representation learning.
3. Visualization results demonstrate that the proposed framework is able to switch transformations during the latent sample evolution and achieves both disentanglement and equivariance.
4. Experimental results show that the learned latent transformations have flexibility and can generalize to new data.

WEAKNESSES:
1. The datasets used in this work are toy (e.g., MNIST, Shapes3D). Large-scale and real-world datasets are expected to evaluate the efficacy of the proposed approach, such as CIFAR, ImageNet.
2. Advanced baseline methods are expected to be compared including StyleGAN, StyleGAN2, and Diffusion models.
3. Can the learned disentangled representations be transferred across datasets, such as from Shapes3D to MNIST?

QUESTIONS:
Please refer to detailed comments in weaknesses.

LIMITATIONS:
This works shows no potential negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes the Flow Factorized Representation to learn structured representations. It introduces a generative model that defines different input transformations with a set of latent flow paths. The latent evolution is generated by the gradient flow of learned potentials, which are parameterized by a set of MLPs. A PINN loss is introduced to ensure that the latent posterior flow follows the OT path. The proposed method is tested on the MNIST and Shapes3D datasets.\n\nSTRENGTHS:\n- The paper is well-written, and the proposed method is technically sound.\n- The learned transformation are well factorized and demonstrate flexible composability.\n\nWEAKNESSES:\n- The paper lacks a comparison in terms of disentanglement. Since the authors claim that the proposed method brings new understandings to both disentanglement and equivariance, it would be beneficial to compare it with state-of-the-art disentanglement methods.  As traditional disentanglement metrics in VAE are unsuitable for this method, the VP metric [a] and the correlation metric [b] could be helpful. Additionally, a qualitative comparison would also be beneficial.\n- More state-of-the-art disentanglement methods should be included in the comparison, such as [c, d].\n\n\n[a] Zhu, Xinqi, et al. \"Learning disentangled representations with latent variation predictability.\" ECCV 2020.\n[b] Shen, Yujun, et al. \"Interfacegan: Interpreting the disentangled face representation learned by gans.\" TPAMI 2020.\n[c] Estermann, Benjamin, et al. \"DAVA: Disentangling Adversarial Variational Autoencoder.\" ICLR 2023.  \n[d] Shao, Huajie, et al. \"Controlvae: Controllable variational autoencoder.\" ICML, 2020.\n\nQUESTIONS:\n- What value is $k$ set to in the experiments? If $k$ is larger than the actual factors of the datasets, how will the model perform?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors study disentangled and equivariant representation learning as the authors claim existing approaches are either ill-specified or can not effectively separate realistic factors of interest in the learned latent space. The key idea is to define a distinct set of tangent directions that associate different transformations. This is achieved by using flow factorized variational autoencoders, and encouraging each latent posterior flow to follow the optimal transport path, which is accomplished by using the PINN loss (equation 14) based on the Hamilton-Jacobi equation with an external driving force.\n\nSTRENGTHS:\nThe authors provide an alternative viewpoint at the intersection of disentangled and equivariant-based representation learning, and propose the flow factored representation learning which targets to define distinct sets of latent probability paths that explains different input transformations. The authors designed a sequential variational autoencoder with latent posterior flow following the OT path. This is achieved by making the gradient satisfy some PEDEs. \n\nThe graphical model, illustration and derivation overall look sound to me (I did not check Theorem 1). The authors also examine their proposed methods in two established datasets (MNIST and Shapes3D). In the two datasets, the proposed methods achieve clear benefits in terms of log-likelihood and equivariance error.\n\nAs the proposed method is targeting for learning equivariant representations, it is expected that the method does achieve lowest equivariance error; Meanwhile the latent flow does perform the target transformation precisely during evaluation while leaving other trains unaffected according to the qualitative study.\n\nWEAKNESSES:\nThe methods study the intersection of disentanglement and equivariance, but quantitative results associated with disentanglement are missing. It would still be interesting to evaluate the proposed methods in terms of some major disentanglement metrics (e.g., MIG,DCI,Modularity) when compared with other existing approaches. Also, it would be a good test to benchmark the proposed methods in some other well-established datasets like dSprites, Natural, Falor3D and Issac3D, to examine equivariance error, some disentanglement metrics, reconstruction, or some other downstream tasks.\n\nQUESTIONS:\nIn equation 6,  will t=0 be problematic?\n\nIn page 4 (right before equation 9), it’s unclear from the notation if you are defining one Neural potential for each value in the category, or you are using only one MLP. \n\nThe external force appeared abruptly in equation 13 without any introduction, and any details on neural parameterization (e.g., what does MLP looks like) are missing. Looking at equation 14, it’s unclear why an independent MLP (the f(z, t)) would affect the potentials (those u^k(z, t) which are also defined by MLP).\n\nAccording to table 1 and 2, the weakly supervised approach still outperforms other baseline methods (which is good), but are consistently worse than supervised ones (expected). It would be interesting to see if semi-supervised learning could improve over supervised one.\n\nLIMITATIONS:\nDo not identify any negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel approach to unsupervised representation learning; More precisely, the authors focus on the learning of latent dynamics with the aim to design an equivariant mapping from the observational to the representational space; Authors leverage variational posterior approximation and dynamic optimal transport to do so and evaluate their method on MNIST (augmented with standard transformation - rotation, scaling, ...) and dSprites showing a significant log-likelihood and equivariance error improvement in comparison to SlowVAE and TVAE; In addition, authors highlight the flexibility of the approach by showcasing the ability of the trained VAE to generate coherent dynamics in the pixel space with superposed transformations and transformations shifts; \n\nWhilst the motivation towards learning equivariant models as well as the idea of learning disentangled latent dynamics through variational posterior approximation is not new, here authors propose a novel set of assumptions over the posterior ($q(z_t|z_{t-1},k)$) and use optimal transport to efficiently train the VAE;\n\nSTRENGTHS:\nMy expertise in optimal transport being limited, I will mostly focus my evaluation on the novelty of the proposed contributions, the soundness of the ELBO objective proposed as well as the assumptions considered; \n\n- I genuinely enjoyed reading this paper and found the contribution exciting. The manuscript is well-written, mathematical derivations are sound, and the experimental setting is clear; The fact that authors show a competitive performance of the model whilst considering only the information from element $x_0$, the first observations of a sequence, in their posterior factorization is exciting. This structural assumption sounds to me more intuitive and sound than the ones I encountered in previous works. \n- Whilst the datasets considered in this paper remain quite simple, as generally observed in this sub-field, the improvements in terms of log-likelihood and equivariance error (standard evaluation metrics) are significant and the model is general enough to allow for nice properties to arise (i.e., superposition of transformation, transformation shift);\n\nWEAKNESSES:\n- The weakness of this work lies for me in the lack of diversity in the experimental setup considered as further detailed in the limitations sections.\n\nQUESTIONS:\n- Could authors clarify how the number of transformations, $k$ was chosen in the experimental work and whether an ablation over this hyperparameter was performed? \n- Could authors discuss the computational complexity of the proposed approach and provide some insights on how this compares to the baselines considered?\n\nLIMITATIONS:\n- As for most work in this sub-field (c.f., SlowVAE, TVAE), the experimental work is limited to simple datasets (i.e., MNIST, dSprites) and does not investigate how these methods perform on real-world images/videos;", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}