Flow Factorized Representation Learning
Authors: Yue Song
T. Anderson Keller
Nicu Sebe
Max Welling
Venue: NeurIPS 2023 poster (2023)
Keywords: Generative modelling
latent disentanglement
variational autoencoders
TL;DR: We introduce Flow Factorized Representation Learning which define a set of probabilitic latent path following dynamic optimal transport for modeling different input transformations.

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
Yue Song!2, T. Anderson Keller Nicu Sebe', and Max Welling
Department of Information Engineering and Computer Science, University of Trento, Italy Amsterdam Machine Learning Lab, University of Amsterdam, the Netherlands UvA-Bosch Delta Lab, University of Amsterdam, the Netherlands yue songounitn.
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, propose an alternative viewpoint 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 generative model which specifies distinct 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 both disentanglement and equivariance. We show that our model achieves higher likelihoods on standard representation learning benchmarks while simultancously 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 degree of robustness and generalizability approaching the ultimate goal of usefully factorized representation learning.
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 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 [3] In general, the definition of "disentanglemenn' often refers to learning and controlling statistically independent factors of variation [5] 36). Over the years, many disentanglement methods have been proposed, including axis-al akis-aligned single-dimensions manipulation [35] [3), linear multi limensiaman traversals [78] [77, [90] [66, and, more recently, dynamic non- 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 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 (3). Specifically, equivariant representations are those for which the output transforms in known predictable way for given input transformation.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 [6).. 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 [T6] 50) 73 (9). 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
In this work. we provide an alternative view- point the intersection of these two fields of work which we call Flow Factorized Represen- tation Learning. Fig. [ depicts the high-level illustration of our method. Given k different transformations p& ( in the input space, we have the corresponding latent probabilistic path 22,.=, (00 æo g (0 (2))) for each of the transformations. Each latent flow path (x<=z(=0) is generated by the gradient field of some learned potentials Vu following fluid me- chanical dynamic Optimal Transport (OT) [4]. Our framework allows for novel understandings of both disentanglement and equivariance. The definition of disentanglement refers to the dis- tinct set of tangent directions Vu that follow the OT paths to generate latent flows for mod- eling different factors of variation. The concept of equivariance in our case means that the two probabilistic paths, i.e., Pk (xe (zp) ) in the image space and J= (z (F (Ze (2))pp (x))) (z)) in the latent space, would eventually result in the same distribution of transformed data.
[ze Vu4
Vu
Pu
Vu
Space
p(x(
g(20) xn)
Image Space
Figure 1: Illustration of our flow factorized repre- sentation learning: at each point in the latent space we have a distinct set of tangent directions Vu which define different transformations we would like to model in the image space. For each path, the latent sample evolves to the target on the potential landscape following dynamic optimal transport.
We build formal generative model of quences and integrate the above latent proba- bility evolution condition updates of the fac- 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 approx- imate the true posterior of latent variables and train the parameters as Variational Autoencoder (VAE) 477. When the transformation type is not observed (i.e., available 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 effective- ness 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 com- bine to yield the highest likelihood on the test set in each setting. Code is publicly available at /ttps://g/ thub. com/ SaeninnaameSSon/latteeeee
label), we treat
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 distinct components by assuming that each observed sequence generated by one of the l separate flows of probability mass in latent space. Since in this work

---

80
X0
Figure 2: Depiction of our model in plate notation. (Left) Supervised, (Right) Weakly-supervised. White nodes denote latent variables, shaded nodes denote observed variables, solid lines denote the generative model, and dashed lines denote the approximate posterior. We see, in standard VAE framework, our model approximates the initial one-step posterior p(z0 (zo). but additionally approximates the conditional transition distribution p( k) through dynamic optimal transport over a potential landscape.
model discrete sequences of observations ã (mo. æ ........ Fr J. we aim to define joint distribution with similarly discrete sequence latent variables (zu, 21 .... 2r and categorical random variable k describing the sequence type (observed or unobserved). Explicitly, we assert the following factorization of the joint distribution over timesteps:
T p 2, = p()))((z))p(0oplppp Il
(I)
Here pkk) is a categorical distribution defining the transformation type, p(a:, (2,) asserts mapping from latents to observations with Gaussian noise, and p(=o) (0, 1). A plate diagram of this model depicted through the solid lines Fig.
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( k) from the continuous form of the continuity equation: 0.p(z) -V (p(=)Vu<())) where vt (z) is the k'th potential function which advects the density p(z) through the induced velocity field V V (z). Considering the discrete particle evolution corresponding to this density evolution, derive the conditional update from the continuous change of variables formula 669
f(= k) - V=1/"(2,11 we see that we can
(f(=111))
p(2/|z(-1, p(zz-1)
dz,-1
In this setting, we see that the choice of ib ultimately determines the prior on the transition probability our model. As minimally informative prior for random trajectories, we use diffusion equation achieved by simply taking v - D log p(zz). Then according to the continuity equation, the prior evolves as:
0,((2)) = - (p( D2 p(z.))
where D3 is constant coefficient that does not change over time. The density evolution the prior distribution thus follows constant diffusion process. We set Dj 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 variational approximation to the true posterior, and train the parameters of the model VAE. To this, parameterize approximate posterior for p(zol æo), and additionally parameterize set of K

---

functions (z) to approximate the true latent potentials 30", First, we will describe how we do this in the setting where the categorical random variable & is observed (which we call the supervised setting), then we will describe the model when & is also latent and thus additionally inferred (which we call the weakly supervised setting).
3.1 Inference with observed (supervised)
When k is observed, we define our approximate posterior to factorize as follows:
T
k=) ((20 mn) II g(zz|zz-1,
(4)
We see that, in effect, our approximate posterior only considers information from element vo: however, combined with supervision in the form of find this sufficient for the posterior 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+30 in future work.
Combing Eq. (4) with Eq.
we can derive the following lower bound to model evidence (ELBO):
p # kkk) (z)
logp(#)k)
Egg (#)2,A)
log << a., k) p()
log p 2, k))( k)
E go (E((,A)
(5)
( c k)
p(=
E (E)2,A)
[log E, k)] + (g log (Z| k.)
Substituting and simplifying, Eq. (5) can be re-written
logp(a/k) N
E (o( pax, k)] (x
(6)
T
E (2)R) DKL ())
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 time evolution defined by the conditional distributions.
3.2 Inference with latent & (weakly supervised)
When ie is not observed, we can treat it another latent variable, and simultaneously perform inference over it in addition to the sequential latent ž. To achieve this, we define our approximate posterior and instead factorize it as
q(z, k|z) = q(<[=))((z)) 2o II q(<z|
Following a similar procedure as in the supervised setting, we derive the new ELBO as
log p( 2, k)) g(2, k|
log pl (+) (x,A(a
(( k| p(z, k a))
log p (z) p(k)
(8)
q(2 L k)
(k(
p [+) ( x. k:)
log p(k)
(x,,aa)) [log p(x k), (=,
log
+ Ey, (Maa
q & ax)
We see that, compared with Eq. 5 only one additional KL divergence term DEL [q, (A| [A|))|pp(()]]]]]] is added. The prior pkk) set to follow categorical distribution, and we apply the Gumbel Gumbel-Softhe Sof tMax trick [43] to allow for categorical paramtterization and sampling of q, (k)æ).

---

3.3 Posterior time evolution
As noted, to approximate the true generative model which has some unknown latent potentials e we propose to parameterize a set of potentials as u² (z, t) = MLP([z; t)) and train them through the ELBOs above. Again, we use the continuity equation define the time evolution of the posterior, and thus we can derive the conditional time update g( 1, through the change of variables
formula. Given the function of the sample evolution 2t- = (( = - + V Vxu we have:
q (211/21-1, k:) = q(2 dg( 2
dz--1
Converting the above continuous equation to the discrete setting and taking the logarithm of both sides gives the normalizing- caapliiiingtttw--le flow-like density evolution of our posterior:
log (z 1 kb log ((x log V2u
(10)
The above relation can be equivalently derived from the continuity equation (i.e. 0zg(2) =
-V. (u²=) Notice that we only assume the initial posterior g(zu 22)) follows a Gaussian
distribution. For future timesteps, 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 Vu satisfies certain PDEs, the evolution of the probability density can be seen to minimize the L2 Wasserstein distance between the source distribution and the distribution the target transformation. Specifically, we have:
Theorem 1 (Benamou- Brenier Formula 44. For probability measures μo and μn the L2 Wasserstein distance can be defined as
{ ()((,)) dr
W W2(No: = min
(11)
p.v.
where the density and the velocity satisfy:
d d V (p(x, 1)), v(x,t) Vu(x,
(12)
dt
The optimality condition of the velocity is given by the generalized Hamilton-Jacobi (HJ) equation (i.e., deu + 1/2l/Vul|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
1
nd (z, t) +
f subject to f(z,t) 0
(13)
2
Here we another MLP parameterize the external force f(2, () and realize the negativity constraint by setting f(z, t) = - -MLP([>: { 2)) Notice that here we take the external force as learnable MLPs simply because we would like to obtain flexible negativity constraint. The MLP architecture is set the same for both u(z,8 t) and f(z, t). To achieve the PDE constraint, we impose Physics-Informed
Neural Network (PINN) [67 loss as
T CHJ Σ
(z.tt)
2 at
(z,t) + I1 (z, -(x,)) ( 0)
(14)
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 (=0, therefore matching the initial condition.
Experiments
This section starts with the experimental setup, followed by the main qualitative and quantitative results, then proceeds to discussions about the generalization ability different composability and unseen data, and ends with the results complex real-world datasets.
5

---

2
3
0
2
C
C
3
3
3 3
3 3 3
5
5
5
5
5
3
3
3
3
3
m
4
5
6
6
6
3
3
3
3
3
3
3
4
4
4
4
4
4 4
5
5
5
5
5
6
6
6
6
6
6
66
7 7
7
7 7
8
8
8
8
8
8
8
6
6
9
7 1
8
8
8
6 6
6
6
7 7
7
7
7
7
8
8
8
8
8
8
8
Figure 3: Exemplary latent evolution results of Scaling, Rotation, and Coloring on MNIST [54). The lop two rows are based on the supervised experiment, while the images of the bottom row are taken from the weakly-supervised setting of our experiment.
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 transforma- tions including Scaling, Rotation, and Coloring. For Shapes3D [10], we use the self celf-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 different lighting conditions and viewpoints, while Isaac3D is dataset of various robot movements in dynamic environments.
Figure 4: Exemplary latent flow results on Shapes3D [10] The transformations from top to bottom are Floor Hue, Wall Hue, Object Hue, and Scale, respectively. The images of the top row are from the supervised experiment, while the bottom row is based on the
experiment.
Baselines. We mainly compare our method with SlowVAE 49 and Topographic VAE (TVAE) 45
These two baselines could both achieve approximate equivariance. Specifically, TVAE intro- duces some learned latent operators, while Slow enforces the Laplacian prior p(z, x;-1) =
II a>/2r(1//)) exp (- (-X(zz,; =t=1,; ( to sequential pairs. Within the disentanglement literature, our method is compared with the supervised PoFlow 79 which adopts wave-like potential flow for sam- ple evolution, and the unsupervised (A-AE 35 and FactorVAE 46 which encourage independence between single latent dimensions. Finally, the vanilla VAE is used controlled bascline.

---

Equivariance Error (U)
Methods
Supervision?
Log-likelibood (1)
Scaling
Rotation
Coloring
VAE +7
No (X)
1275.3111.88 1.89
1310.72:12.19
1368.99 1368.92.12.33
120...7::1..83
3-VAE 35
(X)
741.58:1.557
751.32.15.22
10..16.15.03
-2224.67:12.355
FactorVAE [46
(7)
659.7144.89
632.4415.76
662.18:5:266
-2209.33+2.447
Slow 49
Weak (A)
461.39==5377
447.46:+3.46
398. 398.12:444 12+-4.83
-2197.68+-2.39
TVAE 45
Yes (V)
505.19
493.28-+3.37
451.25-+2.76
-2181.13+
PoFlow
79
Yes (V)
234.78+2.91
231.42--2.98
240.57+22588
-2145.03+2.00
Ours
Weak (/)
185.42
3.3.5413.100
158.57:12.95
.2112.45
Ours
193.8412.47
157.16.1.224
165.1912.78
12119.94.1..66
Table Equivariance error Eg and likeliiodoood log p((,) on MNIST [54].
Metrics. We use the approximate equivariance error Eg and the log-likelihood of transformed data log p(z,) as the evaluation protocols. The equivariance error is defined as & => |x,
Decode(z:)) where 2, 2g + Ei= V=u For TVAE, the latent operator is changed to Roll(zo, t)- For unsupervised disentanglement baselines [35 46 and Slow [4), 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 VAE does not have the corresponding learned transformation in the latent space, we simply set Vau* k 0 and take it 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 & -En Encode((x;) - zr). 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. and 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 supeevies 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 weakly- supervised 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.
Equivariance Error (1)
Methods
Supervision?
Log-Iikelihood (f)
Floor Hue
Wall Hue
Object Hue
Scale
VAE 47
No (X)
6924.63:1899
7746.37=8.77
4383.54:1.26
2609.59:17.11
-11784.69:1..8
3-VAE
(X)
1243.95.112.48
2279.23113..7
2186.73-12.61
2037.94111.722
09924.8.115.6
V.c.orVA..
46
(X)
1985.75+13.26
1876.41+11199
1902.83-12.27
1657.32++1.m0 11.05
-1 11802 17+55.9
Slow VAE VAE
40
Weak (7)
1247.36+:12.99
1314.86++11441
1102.28-12.17
1058 10.96
-11674.89-5.77
TVAE 45
Yes (V)
1225.47
1246.32-9.54
1261.79:19.86
1142.01::9.37
111455.88:5.18
PoFlow 79
Yes (V)
885. 10.337
10.499
912.48-9.86
924. 10.05
-11335 84:
Ours
Yes (V)
613.29=88.93
653445:9.48
605.79_8.63
599.7119.34
-11215.42:5.71
Ours
Weak )
690.84:19.57
10.65
681.59-9.02
653.58119.57
-11279.61:5.89
Table 2: Equivariance error E and log-likelihood log p(x.) on Shapes3D [10]
Quantitative results. Tables and 2 compare the equivariance error and the log-likelihood on MNIST [5 and Shapes3D [TO] 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- likeelildhood on MNIST. The performance gain is also consistent on Shapes3I:: 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-supervisee 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. 67.88 on Shapes3D.
setting, our method
4.3 Discussion
Extrapolation: switching transformations. Fig. we demonstrate that, empowered by our method, it is possible to switch latent transformation categories mid-way through the latent evolution
7

---

Aiteviion
Calrrigg
Wal Hum
Ohjwcr
Coloring
Seoling
Otywer Sioe
Otjwcr HUW
2
2
2
Figure 5: Exemplary visualization of switching transformations during the latent sample evolution.
and maintain coherence. That is, perform 2,- V v for / S 7/2 and then change to 2, - -1-1 + Vzw where j j for 7/2. As can seen, factor variation immediately changes after the transformation type is switched. Moreover, the transition phase is smooth and no other attributes of the image are influenced.
Tansformations
Scala
Object Hue
Scale Object Hho
Wall Hue Floor
Wall Object Hue
Figure 6: Examples of combining different transformations simultaneously during the latent evolution.
superposing transformations. Besides switching transformations, method also supports applying different transformations simultaneously. i.e., consistently performing 2; = + vzu during the latent flow process. Fig. 6 presents such exemplary visualizations of superposing two and transformations simultancously. In each case. the latent evolution corresponds to simultaneous smooth variations of multiple image attributes. This indicates that our method also generalizes well superposing different transformations.
Extrapolation:
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 measure the generalization ability of the learned equivariance to novel compositions.
extrapolation to
Equivariance generalization to new data. We also test whether the learned equivariance holds for aut-of-Distributian (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. displays the exemplary visualization results of the VAE trained on MNIST (54) but evaluated on dSprites (59). Although the reconstruction quality is poor, the learned equivariance still clearly effective 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 compare the equivariance error of our methods and the representative baselines on F and Isaac3D, respectively. Notice that the values are much larger than previous datasets due to the increased image resolution. Our method still outperforms other basclines by a large margin

---

Thomrra Phaats
Tovaraal Roznts
Theersal Hozds
GoD
Prowasal
Tavensa Rooels
Vowerooe Rocues
CoD Ieeu
Coteat
Figure Equivariance generalization unseen OoD input data. Here the model is trained on but the latent flow is tested on dSprites [59]]
MNIST
54
Methods
Lighting Intensity
Lighting X-dir
Lighting Y-dir
Lighting Z-dir
Camera X-pos
Camera Y-pos
Camera Y-pos
TYAE 45
11477.81 8312.97
12568.
11807,3 34
11829
11539
11736.78
11951.45 8994.66
PoFlow 77
79%6. 18
85 19.
8871 L22
81 16.882
15334,9 91
Ours
5798.42
6145.09
6334.87
6782.84
6312,95
6513.68
6614.27
Table 3: Equivariance error (4) on Falcol3D [61].
and achieves reasonable equivariance error. Fig. 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 method 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 3-VAE [35.. InfoGAN [3] achieves disentanglement by maximizing the mutual information between subset of latent dimensions and observations, while 8 3-VAE 335 induces the factorized posterior ((z) by penalizing the Total Correlation (TC) through an extra Syper-parameter B>> 1 controlling the
strength of the KI. divergence. Following infoGAN, many attempts have been made facilitate the
discovery of semantically meaningful traversal directions through regularization [33 42
89
34
100
66
90 98 84 99 78 62]. The follow-up research of 3-VAE mainly explored different methods
to factorize the aggregated posterior 22
25
52
46
12
96
23 76
588
80 28). More recently,
some works proposed 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 distinct tangent bundle V that moves every latent sample via dynamic OT.
We see the most similar method 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 really correspond the OT theory. Also, they do not consider the posterior evolution but instead use the loss || Encode (x;)) (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. A function is said to be equivariant map if it commutes with given transformation, i.e., [=) ]= f (T[xx) where T and T' represent operators in different domains. Equivariance has been considered desired inductive bias for deep neural networks as this property
can preserve geometric symmetries of the input space 38 75 56 57 Analytically equivariant
Methods
Robot Y-mooe Camera Height
Objject
Lighting Intensity
Lighting Y-die
Wall Object
TVAE
8251
8291.70
PoFlow
6572.19
6319
6188
6517
6712
7056.
Ours
3689.72
3993.33
4170.27
4359.78
4225.34
4019.54
$514.97
3876.01
Table 4: Equivariance error (L) Isaac3D [61].

---

Color
Lyghing claertty
PvRDA
Poler Vencon
Wyving Uluarty
Cocoon Toine
Figure 8: Qualitative comparison of our method against TVAE and PoFlow on Falcol3D and Isaac3D.
networks typically enforce explicit symmetry to group transformations in neural networks [6
68
Another line of research proposed to directly learn approximate equivariance Our framework re-defines approximate equivariance by matching the latent probabilistic flow to the actual path of the given transformation in the image space.
93
92
85
31
39
from data 21
49
20
45
Optimal transport in deep learning. There is 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 light-speed implementation of the Sinkhorn algorithm for fast computation of entropy regulaa-ree Wasserstein distances, which opened the way for many differentiable Sinkhorn algorithm- based applications [32] 22 [4 27 [5]] In generative modeling. the Wasserstein distance is often used to minimize the discrepancy between the data distribution and Inspired by the fluid mechanical interpretation of OT 4 some considered regularizing the velocity fields to satisfy the HJ 60 Our method applies PIÑNs [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.
the model distribution 2 81
72
normalizing flow methods 669
24
48
equation, thus matching the dynamic OT plan 94 30
Conclusion
In this paper, we introduce Flow Factorized Representation Learning which defines 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 re- interprets 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 ( k) to q( 20, k) because we assume & contains sufficient information of the whole sequence. To keep the posterior definition of q( (E, k)., we need make q(z,) also function of æ;- This can be achieved either by changing the potential to ul 2 Zg. 1-1) or modifying the external driving force f 1, Ft, 1-1). Nonetheless, we see these modifications would make the model less flexible than our current formulations as the element 2; might be needed during inference.
Acknowledgments and Disclosure of Funding
This work was supported by the MUR PNRR project FAIR (PE00000013) funded by the NexiGener- ationEU, by the PRIN project CREATIVE (Prot. 20207S1.999, by the EU H2020 project AI4Media (No. 951911), and by the Bosch Center for Artificial Intelligence.
10

---

References
[1] Pulkit Agrawal, Joao Carreira, and Jitendra Malik. Learning to by moving, In ICCV, 2015.
[2] Martin Arjovsky, Soumith Chintala, and Lčon Bottou. Wasserstein generative adversarial networks. In ICML, 2017.
[3] Chandrajit Bajaj. Luke McLennan, Timothy Andeen, and Avik Roy. Recipes for when physics fails: Recovering robust learning of physies informed neural networks. Muchine Learning: Science and Technology, 2023.
[4] Jean- Lean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the monge- kantorovich mass transfer problem. Numerische Mathematik, 2000.
[5] Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE T-PAMI, 2013.
[6] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint 1308.3330...43. 2013.
[7] Diane Bouchacourt, Mark Ibrahim, and Stéphane Deny. Addressing the topological defects of disentan- glement via distributed operators. arXiv preprint arXiv:2102.05623, 2021.
[8] Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. ICLR, 2022.
[9] 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.
[10] Chris Burgess and Hyunjik Kim. 3d shapes dataset. http://ggthh...oommmemmmmmm/aaae 2018.
[11] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations, 2019.
[12] Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. Isolating sources of disentangle- ment in variational autoencoders. NeurlPS, 2018.
[13] 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.
[14] Lenaic Chizat, Pierre Roussillon, Flavien Léger, François-Xavier Vialard, and Gabriel Peyré. Faster wasserstein distance estimation with the sinkhorn divergenee. NeurIPS, 2020.
[15] Taco Cohen and Max Welling. Transformation properties learned visual representations. ICLR, 2015.
[16] Taco Cohen and Max Welling, Group equivariant convolutional networks. In ICML. PMLR. 2016.
[17] Taco Cohen and Max Welling. Steerable enns. ICLR, 2017.
[18] Marissa Connor, Gregory Canal, and Christopher Rozell. Variational autoencoder with learned latent structure. AISTATS. PMLR, 2021.
[19] Marco Cuturi. Sinkhorn distances: Lightspeed computation optimal transport. NeurIPS, 2013.
[20] Neel Dey, Antong Chen, and Soheil Ghafurian. Group equivariant generative adversarial networks. JCLR, 2021.
[21] Nichita Diaconu and Daniel Worrall. Learning to convolve: A generalized eeight-tying approach. In ICML. PMLR, 2019.
[22] Nat Dilokthanakul, Pedro AM Mediano, Marta Garnelo, Matthew CH Lee, Hugh Salimbeni, Kai Arulku- maran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture variational autoen- coders. ICLR, 2016.
[23] Zheng Ding. Yifan Xu. Weijian Xu, Gaurav Parmar, Yang Yang. Max Welling, and Zhuowen Tu. Guided variational autoencoder for disentanglement learning. In CVPR, 2020.
11

---

[24] Laurent Dinh, Jascha Sohl- Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. ICLR, 2017.
[25] Emilien Dupont. Learning disentangled joint continuous and discrete representations. NewrIPS, 2018.
[26] Cian Eastwood and Christopher KI Williams. A framework for the quantitative evaluation of disentangled representations. In ICLR. 2018.
[27] Marvin Eisenberger, Aysim Toker, Laura Leal- Taixé, Florian Bernard, and Daniel Cremers. A unified framework for implicit sinkhorn diffcrentiation. In CVPR, 2022.
[28] Benjamin Estermann and Roger Wattenhofer, Dava: Disentangling adversarial variational autoencoder. ICLR, 2023.
[29] 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.
[30] 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 JCML. PMLR, 2020.
[31] Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to groups on arbitrary continuous data. In ICML. PMLR, 2020.
[32] Charlic Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio. Learning with a wasserstein loss. NeurIPS, 2015.
[33] Lore Goetschalckx, Alex Andonian, Aude Oliva, and Phillip Isola. Ganalyze: Toward visual definitions of cognitive image properties. In ICCV, 2019.
[34] Erik Härkönen, Aaron Hertzmann, Jaakko Lehtinen, and Sylvain Paris. Ganspace: Discovering inter- pretable gan controls. NeurlPS, 2020.
[35] Irina Higgins. Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vac: Learning basic visual concepts with a constrained variational framework. ICLR, 2016.
[36] Trina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexan- der Lerchner. Towards definition of disentangled representations. arXiv preprint arXiv:1812.02233.. 2018.
[37] Irina Higgins, Le Chang, Victoria Langston, Demis Hassabis, Christopher Summerficld, Doris Tsao, Matthew Botvinick. Unsupervised deep learning identifies semantic disentanglement in single inferotemporal face patch neurons. Nature communications, 2021.
[38] Geoffrey E Hinton, Alex Krizhevsky, and Sida D Wang. Transforming auto-encoders. In ICANN. Springer. 2011.
[39] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In ICML. PMLR, 2022.
[40] Jun-Ting Hsich, Shengjia Zhao, Stephan Eismann, Lucia Mirabella, and Stefano Ermon. Learning neural pde solvers with convergence guarantees. ICLR, 2019.
[41] 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.
[42] Ali Jahanian, Lucy Chai, and Phillip Isola. On the" steerability" of generative adversarial networks. ICLR, 2020.
[43] Eric Jang. Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. ICLR, 2017.
[44] Yeonwoo Jeong and Hyun Oh Song. Learning discrete and continuous factors of data via alternating disentanglement. /CML, 2019.
[45] T Anderson Keller and Max Welling. Topographic vaes leam equivariant capsules. NeurlPS, 2021.
[46] Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In ICML, 2018.
[47] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. ICLR, 2014.
12

---

[48] Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible Ixl convolutions. NeurIPS, 2018.
[49] 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.
[50] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In ICML. PMLR, 2020.
[51] Soheil Kolouri, Navid Naderializadeh, Gustavo Rohde, and Heiko Hoffmann. Wasserstein embedding for graph learning. ICLR, 2021.
[52] Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan, Variational inference of disentangled latent concepts from unlabeled observations. ICLR, 2018.
[53] Mingi Kwon, Jacscok Jeong, and Youngjung Uh. Diffusion models already have semantic latent space. ICLR, 2023.
[54] Yann LeCun. The mnist database of handwritten digits. 1998. URL http::/://y/a www.coo///eoo/oo /xtb/ mnist/
[55] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 2015.
[56] Chen Yu Lee. Saining Xie, Patrick Gallagher, Zhengyou Zhang, and Zhuowen Tu. nets. In AISTATS. PMLR, 2015.
[57] Karel Lenc and Andrea Vedaldi. Understanding image representations by measuring their equivariance and equivalence. In CVPR, 2015.
[58] Francesco Locatello, Ben Poole, Gunnar Rätsch, Berhhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-supervised disentanglement without compromises. In ICML. PMLR, 2020.
[59] Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerehner. dsprites: Disentanglement testing sprites dataset, 2017. httptttt:://
com http///eepmmnndddddde wwwpp///dddpprtessssaassttstt
[60] Kirill Neklyudov, Daniel Severo, and Alireza Makhzani. Action matching: A variational method for learning stochastic dynamics from samples. ICML, 2023.
[61] 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.
[62] James Oldfield, Christos Tzelepis, Yannis Panagakis, Mihalis A Nicolaou, and Joannis Patras. Panda: Unsupervised learning of parts and appearances the feature maps of gans. ICLR, 2023.
[63] Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. Ot-flow: Fast and accurate continuous normalizing flows via optimal transport. In AAAI, 2021.
[64] Yong-Ily Park, Mingi Kwon, Junghyo Jo. and Youngjung Uh. Unsupervised discovery of semantic latent directions in diffusson models. arXiv preprint arXiv:2302. 12469, 2023.
[65] Giorgio Patrini, Rianne den Berg, Patrick Forre, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. UAI, 2020.
[66] William Peebles, John Peebles, Jun-Yan Zhu, Alexei Efros, and Antonio Torralba. The hessian penalty: weak prior for unsupervised disentanglement. In ECCV, 2020.
[67] M. Raissi, 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.
[68] Siamak Ravanbakhsh, Jeff Schneider, and Barnabas Poczos. Equivariance through parameter-sharing. In ICML. PMLR, 2017.
[69] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In ICML PMLR, 2015.
[70] Jack Riwhter-Powell, Yaron Lipman, and Ricky TQ Chen. Neural conservation laws: A divergence free perspective. NeurIPS, 2022.
13

---

[71] Karl Ridgeway and Michael C Mozer. Learning deep disentangled embeddings with the f-statistic loss. NeurIPS, 2018.
[72] Tim Salimans, Han Zhang, Alec Radford, and Dimitris Metaxas. Improving gans using optimal transport. ICIR, 2018.
[73] Victor Garcia Satorras, Emiel Hoogeboom, Fabian B Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows. NeurlPS, 2021.
[74] Jürgen Schmidhuber. Learning factorial codes by prodictability minimization. Neural computation, 1992.
[75] Uwe Schmidt and Stefan Roth. Learning rotation-aware features: From invariant priors equivariant descriptors. In CVPR, 2012.
[76] Huajie Shao, Shuochao Yao, Dachun Sun, Aston Zhang, Shengzhong Liu, Dongxin Liu, Jun Wang, and Tarek Abdelzaher. Controlvae: Controllable variational autoencoder. In JCML. PMLR, 2020.
[77] Yujun Shen and Bolci Zhou. Closed-form factorization of latent semantics in gans. In CVPR, 2021.
[78) Yue Song, Nicu Sebe, and Wei Wang. Orthogonal svd covariance conditioning and latent disentanglement. IEEE T-PAMI, 2022.
[79) Yue Song, Andy Keller, Nicu Sebe, and Max Welling. Latent traversals in generative models as potential flows. In ICML. PMLR, 2023.
[80] Chang- Tai, Ming- Yao Li, and Lun-Wei Ku. Hyperbolic disentangled representation for fine-grained aspect extraction. In AAAI, 2022.
[81] Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto-encoders. ICLR, 2018.
[82] Jakub Tomezak and Max Welling. Vae with vampprior. AISTATS. PMLR, 2018
[83] 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.
[84] Christos Tzelepis. Georgios Tzimiropoulos, and Ioannis Patras. WarpedGANSppcee Finding non-linear rhf paths in GAN latent space. In ICCV, 2021.
[85] Elise Van der Pol, Daniel Worrall, Herke van Hoof, Frans Oliehoek, and Max Welling. Mdp homomorphic networks: Group symmetries in reinforcement learning. NeurIPS, 2020.
[86] Ashish Vaswani. Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NeurlPS, 2017.
[87] Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
[88]
Cédric Villani. Topics in optimal transportation, volume 58. American Mathematical Soc., 2021.
[89] Andrey Voynov and Artem Babenko. Unsupervised discovery of interpretable directions in the gan latent space. In ICML, 2020.
[90] Yuxiang Wei, Yupeng Shi, Xiao Liu, Zhilong Ji, Yuan Gao, Zhongqin Wu, and Wangmeng Zuo. Orthoga- nal jacobian regularization for unsupervised disentanglement in image generation. In ICCV, 2021.
[91] James CR Whittington, Will Dorrell, Surya Ganguli, and Timothy Behrens. Disentanglement with biological constraints: A theory of functional cell types. In ICLR, 2023.
[92] Daniel Worrall and Max Welling. Deep scale-spaces: Equivariance over scale. NeurlPS, 2019.
[93] Daniel E Worrall, Stephan J Garbin, Daniyar Turmukhambetov, and Gabriel Brostow. Harmonic networks: Deep translation and rotation equivariance. CVPR, 2017.
[94] Liu Yang and George Em Kamiadakis. Potential flow generator with optimal transport regularity for generative models. IEEE TNNLS, 2020.
[95] Tao Yang, Yuwang Wang, Yan Lv, and Nanning Zh. Disdiff: Unsupervised disentanglement of diffusion probabilistic models. arXiv preprint arXiv:2301.1372.1 2023.
[96] Cagatay Yildiz, Markus Heinonen, and Harri Labdesmaki. Ode2vae: Deep generative second order odes with bayesian neural networks. NewrIPS, 2019.
14

---

[97] Qi Zeng. Spencer H Bryngelson, and Florian Schäfer. Competitive physics informed networks. ICLR. 2023.
[98] Jiapeng Zhu, Ruili Feng, Yujun Shen, Deli Zhao, Zheng-Jun Zha, Jingren Zhou, and Qifeng Chen. Low-rank subspaces in gans, NeurIPS, 2021.
[99] Jiapeng Zhu, Yujun Shen, Yinghao Xu, Deli Zhao, and Qifeng Chen. Region-based semantic factorization in gans. ICML, 2022.
[100] Xinqi Zhu, Chang Xu, and Dacheng Tao. Learning disentangled representations with latent variation predictability. In ECCV, 2020.
15

---

A Supplementary Material
A.1 Pseudo codes
import torch
#Randomly sample transformation each iteration
index forch.randing (0, potential.number) x_bar sequence_generart (index)
#Generating index according to the supervision setting if training_mode "supervised";
index_potential index
10 elif training_mode mode = -uekkly-spprrvise""
index_potential = q_k(x_bar)
13 #initial element of the sequence
14 2, rho_2 flev_vae (x_bar (0])
15
16 #Future elements of the sequence obtained by latent flow for range T)
18
PDE_loss, delta_z, delat_rho_z HJ_PDE (index_
19
#Updates in the sample and probability space z delta,2
22
rho_2 rho_2 delat_rho_2
23
24
#Inference at every intermediate step
hat_xt = vae inference
27
#Loss: loss +- PDE_loss + CB x_rar (t]) KL (L prior_rho_z)
reconstrutction
div
28
KKL div for index prediction (uparly- setting)
if training_mode "weakly- supervised":
12
loss (u(iinex_potential index)
33
loss backward 34
optimizer step ()
Figure 9: Pytorch-like pseudo codes for training our flow-factorized VAE,
Fig. 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 each iteration. The batch size is set 128 for both datasets. We use Adam optimizer and the learning rate is set as 1e-4 for all the parameters. The encoder consists of four stacked convolution layers with the activation function RcLU, while the decoder comprised four stacked transposed convolution layers. For the prior evolution, the diffusion coefficient D is initialized with 0 and we set it learnable parameter for distinct k. For MLPs that parameterize the potential u 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
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.
16

---

Shapes3D. The input images are resized to 64x64. 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 128x 128. We use the self- contained 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-Sof tmax trick, we re-parameterize (x(k)) by
1/ Σ,e tet
(15)
where is the category prediction, 9i is the sample drawn from Gumbel distributions, and is the small temperature to make tmax 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 initialized with and gradually reduced to 0.05 with the annealing rate 3e-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 3-VAE and FactorVAB to learn the given transformations in an unsupervised manner. For tuning the interpolation range, we start from the initial value 2; and traverse till the appropriate bound which selected from the range [-5,5] 5| with the interval of 0.1.
A.3 Disentanglement metries
There are many traditional disentanglement metrics [7]] 26 12
but they are designed for single- dimension traversal methods. These metries 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 perturbed by vectors for meaningful output variations. When it comes these vector-based disentanglement methods, their scores of disentanglement metrics would drop considerably and cannot be compared with those migle-dimension baselines.
Table 5: VP Scores (%) on MNIST.
Table VP Scores (%) on Shapes3D.
Training Set Split
Ours
PoFlow
IYAE
FactorVAR
J-VAR
Training Set Split
Pollow
FactorVAK
10%
95.69
93.05
10% 1%
95.92 77.03
91.48
83.27
84.49
1%
92.71
91.27
$8.15
84.46
85.25
72.32
68.39
63.83
65.78
Nonetheless, certain disentanglement metrics such as VP scores [100] can leveraged as they do not pose any assumptions on the latent space but only require image pairs [T]] of different transformations for evaluation. The VP metric adopts the few-shot learning setting (using 1 1% or 10% of the dataset the training set) and takes lightweight neural network for learning classify image pairs [xo, of different attributes. The gencralization ability ( (i.e., validation accuracy) can be thus regarded
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 3-VAE 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 fow-frctorized VAE has better disentanglement ability.
A.4 Ablation studies
Impact different priors. We use diffusion equations to model the prior evolution as random parti- cle movement. It would also be intriguing to choose other priors commonly used in the VAE literature, such as Standard Gaussian (SG) priors \' 1), mixture of Gaussian (MoG) priors Σ Sw,N(1, of), and VAMP priors [22 which average aggregated posterior of N pseudo-inputs as 1/N V/ g(2,) Ta- ble 7 presents the equivariance error of different priors on MNIST. Among these priors, our diffusion equations achieve the best performance. This meets assumption that modeling the prior evolution
17

---

Table 7: Equivariance error of different priors.
Table 8: Equivariance error of different PDEs.
Prior
Scaling
Rotation
Coloring
Prior
Scaling
Rotation
Coloring
SG
190.24++2.18
158.93++3.25
164.18+22.7
Heat
223.95+3.38
212.47++3.85
207.666+2.9
MoG
188.23:2.45
157.7942.86
161.4942.22
FP
211.54=3.17
188,59:3.92
194.73-33.00
VAMP
192.81+3.67
161.47+4.12
162.97+3.89
OHJ
190.43+2448
163.87433.05
162.38=2.86
Diffusion
185.42:12.35
153.5413.10
158.5712.95
GHJ
185.42=2.35
153.5413.10
158.57=2.95
as 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 g 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 slight margin. We attribute this advantage the external driving force f(z, t) which gives us more flexibility and dynamics in modeling the velocity fields Vu
Table 9: Equivariance error on MNIST of a differ- Table 10: Equivariance error MNIST of differ- ent number of transformations (K).
ent sequence lengths (T).
K
Scaling
Rotation
Coloring
Sequnce Length (7)
Scaling
Rotation
1
285.27-2..99
2
185.78-2.21 154.2912.87
9
185,42+2.35
153.54+3.10 198.7242.89
3
185.42=2.45
153.5443.10
158.5742.95
12
214.47=2.59
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 As indicated above, in general, the performance 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 Vu on the potential landscape huving 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 worse. To better illustrate this point, we perform an ablation study on MNIST and present the results in Table To Specifically, we change the sequence length from 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 L2 Wasserstein distance can be re-formulated in the fluid mechanical interpretation as
W inf J, S. ",,,, t)o(t, dx dt
(16)
where the density satisfies the continuity equation (O.p = -V (p(x, t)u t)). If we introduce the momentum m(z,t) t) = p(x,,) t)v t) and two Lagrange multipliers u and X, the Lagrangian function of the Wasserstein distance would be:
L(p m, Φ) = [
|mml
+ uu(0pp V. - A(p - 22)
(17)
2p
18

---

Scaling
Rotation
Coloring
V
44
V
4
4
4
1
4
4 4 4 4444
4
4
9
4
4
4
4
155
1
11 5 3 6 9 5 8
4
55
5 35555565
5
5
6
5
5
5
5
5
5
3 3 3 3 33
3
3
3
3
3
3
3
3 6 9 5 8
6
6
6 9 5 8 0 9 T 4 5 9999
5
8 O
8 0 9 T 9 5
8
8
O
0
99
9
9
4
9
9
4
4 7
T
/
7
T
T 4 5 9
7
4|44/4444
1944
1
155
14
9 9
4 5 9
9 5
55
5 55555555
5
J
5
5
5 9
9 P999999p
9 9
9
6
9
9
9
Figure 10: More visualizations of the learned latent flows on MNIST (54]
where the second term is the equality constraint, and the third term is an equality constraint with slack variable s. Using integration by parts formula, the above equation can be re-written as
L(p. m. σ) -/
|m 2p
vylb-/oo S. (0.up Vu m) A(p -0- s²)
uplp
(18)
Based on the set of Karush-Kunn Kuhn-Kuuhke (KKT) conditions (0, = 0, 0.L = 0. 0,L = 0, and
A 0), would have:
0m L m uo=v-Vu=D
0,L = 0ep + V m= =0
(19)
0,L = |m]l 2p d.u - -= - 8,u x
where the first condition indicates that the gradient Vu acts as the velocity field, and the third condition implies the optimal solution given by the generalized HJ equation:
d.u -
(20)
We thus apply the generalized equation (i.e., 0,u +
| S 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 and 12 display more visualization results the latent evolution on MNIST, Shapes3D, Falcol3D and Isaac3D, respectively. Across all the datasets, our method presents precise control of the given transformations. Fig. and [4]show more latent evolution results of switching transformations (top) and combining transformations (bottom) MNIST and Shapes3D, respectively. Fig. [5 also visualizes a few examples of superposing and switching transformation on Falcol3D and Isa Our latent flows learn to compose or switch different transformations precisely and flexibly.
19

---

Floor Hue
Wal Hue
Object Hue
Scale
Figure 11: More visualizations of the learned latent flows on Shapes3D 10
Lighning erwessly
Chrmona Kpaae
Camrra Z-Pox
DIN
KCNG
ligtting X-DY
Camera X-fooe
Lightiing 2-De
BSN
RUVA
WVF
PPP
RONG
NANA
VDSSA
DITSA
Robot X-Movv
Robot V-ldoue
Lighting Inteasty
Well Coler
Cbject Color
Lightring Vrr
Figure 12: Exemplary visualizations of learned latent flows Falcol3D (top) and Isaac3D (bottom).
20

---

Rotation
Scaling
Scaling
Coloring
Scaling
Rotation
222
Scaling
Coloring
Scaling
Rotation
Scaling
bym
3
5
5
S
SSS
Coloring
Rotation
Rotation
Coloring
Scaling
Coloring
3
3
3
3
3
5
Coloring
Rotation
Coloring
Rotation
Coloring
Rotation
2
3
3
3
M
Scaling
Rotation
Scaling
Coloring
Scaling
Rotation
5
5
S
5
155
6
6 6 6
6
BAAA
Scaling Rotation
Rotation Coloring
Scaling Caloring
1
S S
5
8
5
5
5
5
5
5
5
5
5
4
4
4
4
4
4
7
7 0
7
7
1
7
2 2
2
2
b
h
h
5
5
5
5
0
D
D
0 6
0
8
6
6
6 5
6
6
19
9
9
9
5
5
5
5
5
5
5 5
5
O
O
9
9
2
2222222
Figure 13: More visualizations of switching and superposing transformations on MNIST [44]
21

---

Object Hoee
Wall
Floor Hue
Scale
Wail Hue
Foor Hue
Wall Huee
Object Hue
Floor Hue
Wall Hue
Object Hue
Scale
Objact
Rooo Hup
Wall Hue
Floor Hue
Scate
Fioor Hue
Well Hue
Scale
Object Hue
Woll Huve
Object Hue
Scale
Woll Huee
Scale Floor Hue
Object
Scale Wail Hue
Scale Object Hue
Floor Hue Otjjct Hue
Figure 14: More visualizations of switching and superposing transformations on Shapes3D 10
Lightning Intensll Object Color
Objectt Color
Lightning
REP
BNN
Lightning Intensity Object Soe
Lightning KDD Object Cater
Lightsing Kalr Rubot KoMove
Camera Yipos
Camera Z.P
Figure 15: Exemplary visualization results of superposing transformations on Isaac3D (left) and
switching transformations on Falcol3D (right).
22


Reviews:
[{"review_id": "ZXMbcIHzbg", "reviewer": "Reviewer_tD2D", "summary": "This 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; ", "questions": "- 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", "limitations": "- 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; ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "My 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", "weaknesses": "- 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank $\\textcolor{red}{Reviewer\\ tD2D}$ for the constructive suggestions and positive feedback. Below we respond to the questions point by point.\n\n---\n\n**Q1. Choosing the Value of $k$**\n\nFor MNIST, we choose the value of $k$ in line with TVAE [a] and PoFlow [b]. For Shapes3D, we use the self-contained transformations of the dataset but omit the transformation of \"Object Shape\". This is because the sequence is very short (4 elements) and its transition is very abrupt, which does not define a smooth and natural transformation like the other sequences. \n\nEquivariance error of different k on MNIST\nk|Scaling|Rotation|Coloring|\n|:---:|:---:|:---:|:---:| \n|1|185.27±2.59|-|-|\n|2|185.78±2.21|154.29±2.87|-|\n|3|185.42±2.45|153.54±3.10|158.57±2.95|\n\nWe did an ablation study on the impact of $k$ on MNIST and present the evaluation results in the above Table. In general, the performance is not affected by the number of transformations being applied. 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 $\\nabla u$ on the potential landscape $u$; _**having multiple latent flows defined on different potential landscapes therefore does not interfere with each other**_.\n\n>[a] Topographic vaes learn equivariant capsules. NeurIPS21.\n>\n>[b] Latent traversals in generative models as potential flows. ICML23.\n\n**Q2. Computational Complexity**\n\nWe present the training time cost of our method on MNIST in the Table below, as well as the results of TVAE and PoFlow. In the supervised setting, our method is slightly more time-consuming than PoFlow (around 10%) and than TVAE (around 20%). When it comes to the weakly-supervised fashion, our flow-factorized VAE costs more time, as we need to use the Gumbel-Softmax trick for predicting the transformation index given the whole image sequence ($q(k,\\bar{\\mathbf{x}})$) and aggregate flows according to the index prediction ($\\mathbf{z} _{t+1} = \\mathbf{z} _t + \\sum _i^K q(k, \\bar{\\mathbf{x}} ) _i  \\nabla u^i$). Nonetheless, we think the overall time cost is still acceptable considering the performance gain over the baselines (especially for the supervised setting).\n\n\n Training time cost on MNIST each iteration\n||Ours (weakly-supervised)|Ours (supervised)|PoFlow|TVAE|\n|:---:|:---:|:---:|:---:| :---:| \n|Time (s)|1.076|0.723|0.638|0.591|\n\n**Q3. Application on Real-world Datasets**\n\nThanks for the great advice! As suggested by $\\textcolor{magenta}{Reviewer\\ ww47}$, we take a step further to apply our method on Falcol3D and Isaac3D, 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.\n\nEquivariance errors on Falcor3D\n|Methods|Lighting Intensity|Lighting X-dir|Lighting Y-dir|Lighting Z-dir|Camera X-pos|Camera Y-pos|Camera Z-pos|\n|:---:|:---:|:---:|:---:| :---:| :---:| :---:|  :---:| \n|TVAE|11477.81|12568.32|11807.34|11829.33|11539.69|11736.78|11951.45|\n|PoFlow|8312.97|7956.18|8519.39|8871.62|8116.82|8534.91|8994.63|\n|Ours|**5798.42**|**6145.09**|**6334.87**|**6782.84**|**6312.95**|**6513.68**|**6614.27**|\n\n\nEquivariance errors on Issac3D\n|Methods|Robot X-move|Robot Y-move|Camera Height|Object Scale|Lighting Intensity|Lighting Y-dir|Object Color|Wall Color|\n|:---:|:---:|:---:|:---:| :---:| :---:| :---:|  :---:| :---:| \n|TVAE|8441.65|8348.23|8495.31|8251.34|8291.70|8741.07|8456.78|8512.09|\n|PoFlow|6572.19|6489.35|6319.82|6188.59|6517.40|6712.06|7056.98|6343.76|\n|Ours|**3659.72**|**3993.33**|**4170.27**|**4359.78**|**4225.34**|**4019.84**|**5514.97**|**3876.01**|\n\n\n\nThe two Tables above present the equivariance errors of our method and TVAE/PoFlow. Notice that the values are much larger due to the increased image resolution. Our method still outperforms other baselines by a large margin and achieves reasonable equivariance error. To visually display the learned latent transformations, we also conducted some qualitative evaluations and put them in the attached one-page pdf (which comprises exemplary latent evolutions of different transformations and exemplary comparisons against baselines). Our method can precisely control the image transformations through our latent flows.\n\n---\n\nThanks again and it is more than welcome to ask further questions.\n"}, {"review_id": "hHyhGZ9H9v", "reviewer": "Reviewer_ww47", "summary": "The 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. ", "questions": "In 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. ", "limitations": "Do not identify any negative societal impact. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank $\\textcolor{magenta}{Reviewer\\ ww47}$ for the encouraging feedback and valuable suggestions. Below we respond to the concerns in detail.\n\n---\n\n**Q1. Disentanglement Metrics**\n\nAs stated in line29 of the supplementary (and also indicated by $\\textcolor{blue}{Reviewer\\ XnNs}$), most traditional disentanglement metrics like MIG [a] and DCI [b] do not suit our method because 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 [c,d,e], there emerges a more realistic disengagement setting: all the latent dimensions are perturbed by vectors for meaningful output variations. \n\nNonetheless, as pointed out by $\\textcolor{blue}{Reviewer\\ XnNs}$, certain disentanglement metrics such as VP scores [f] can be leveraged as they do not pose any assumptions on the latent space but only require image pairs $\\mathbf{x} _0,\\mathbf{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 light-weighted neural network for learning to classify image pairs $\\mathbf{x} _0,\\mathbf{x} _T$ of different attributes. The generalization ability (validation accuracy) can be thus regarded as a reasonable surrogate for the disentanglement ability.\n\nDue to the limited space, we put the evaluation results in the response to $\\textcolor{blue}{Reviewer\\ xnNs}$. Our method still achieves superior performance on the VP scores and outperforms the previous disentanglement baselines. \n\n> [a] Isolating Sources of Disentanglement in Variational Autoencoders. NeurIPS18.\n> \n> [b] A Framework for the Quatitative Evaluation of Disentangled Representations . ICLR18.\n>\n> [c] Topographic vaes learn equivariant capsules. NeurIPS21.\n> \n> [d] WarpedGANSpace: Finding non-linear rbf paths in GAN latent space. ICCV21.\n>\n> [e] Closed-form factorization of latent semantics in gans. CVPR21.\n>\n> [f] Variation Predictability for Disentanglement. ECCV20.\n\n**Q2. Application on Complex Real-world Datasets**\n\nThanks for the constructive advice! We follow your suggestion and further evaluate our method on Falcor3D and Isaac3D. Due to the limited space, please refer to the response to $\\textcolor{red}{Reviewer\\ tD2D}$ and the attached one-page pdf for the relevant results and discussions. \n\n**Q3. Eq. (6) with t=0**\n\nThanks for the careful check! When $t=0$ the KL div of Eq. (6) is indeed a bit problematic as the conditioned posterior $q(z_t|z_{t-1},k)$ starts with $t=1$. The more rigorous form could be obtained by injecting Eq. (4) into the KL div part of Eq. (6):\n \n$$\\log p(\\bar{\\mathbf{x}} | k)\\geq \\sum _{t=0}^{T} E _{q _{\\theta}(\\bar{\\mathbf{z}}| k)} \\big[ \\log p(\\mathbf{x} _{t}|\\mathbf{z} _{t},k)\\big] -E _{q _{\\theta}(\\bar{\\mathbf{z}}| k)}\\big[\\mathrm{D} _{\\text{KL}}\\left[q _\\theta(\\mathbf{z} _0|\\mathbf{x} _0)||p(\\mathbf{z} _0)\\right] \\big] - \\sum _{t=1}^{T} E _{q _{\\theta}(\\bar{\\mathbf{z}}| k)} \\big[\\mathrm{D} _{\\text{KL}} \\big[\\left[q _\\theta(\\mathbf{z} _{t}|\\mathbf{z} _{t-1}, k) || p(\\mathbf{z} _{t} | \\mathbf{z} _{t-1}, k)\\right]\\big]$$\n\nThe modification decomposes the second r.h.s.term of Eq. (6) into the initial KL div at $t=0$ and the Kl div of the subsequent sequence $t=1:T$. This is also more coherent to our sequential assumption as the first element is independent of the specific transformation $k$. We will revise Eq. (6) for better clarity.\n\n**Q4. One MLP for One Transformation**\n\nThanks for the detailed comment. We define a single MLP for each transformation. The MLP $u(\\mathbf{z},t)\\in R^1$ is a scalar and the gradient $\\nabla _\\mathbf{z} u(\\mathbf{z},t)\\in R^d$ takes the same dimensionality as the latent variable. Therefore, the gradient of one MLP can drive the evolution of the entire latent variable. In the revised paper, we would introduce the dimensionality of $u$ and $\\nabla u$ in Sec 3.3 to avoid confusion.\n\n**Q5. External Force**\n\nWe agree that it might be abrupt to jump to the external force immediately after introducing the generalized HJ equation in Eq. (13). Actually, the external force is set as learnable MLPs simply because we would like to achieve a flexible negativity constraint (in order to satisfy the optimality condition). This can be also achieved by setting $f(z,t)$ as learnable parameters at different timesteps, but we empirically observed that parameterizing $f(z,t)$ as an MLP is more stable throughout the training process. We use the same MLP architectures for both $u(z,t)$ and $f(z,t)$.\n\nWe would re-organize this part and explain the external force more clearly in detail in the revised paper.\n\n**Q6. Semi-supervised Learning**\n\nThanks for the insightful advice. Of course, our method can be extended to the semi-supervised setting. One straightforward implementation could be switching between inferring the transformation index from sequences and directly enforcing the index during the training process. We test such a possibility on MNIST and present the results below:\n\nEquvariance error on MNIST under different supervision settings.\nSupervision|Scaling|Rotation|Coloring|\n|:---:|:---:|:---:|:---:| \n|Weakly|193.84±2.47|157.16±3.24|165.19±2.78|\n|Semi|186.07±2.58|152.87±3.36|160.13±2.82|\n|Full|185.42±2.35|153.54±3.10|158.57±2.95|\n\nAs can be seen, this semi-supervised setting has very close and even competitive performance against the supervised setting. We also observed accelerated convergence of this setting compared with the weakly-supervised fashion. \n\nAnother interesting semi-supervised learning setting is directly enforcing the index for some of the transformations while inferring the index for others. This setting might be useful when the number of transformations to model is very large. We would add a paragraph in the revised paper to discuss the many intriguing possibilities! \n\n---"}, {"review_id": "hbd8tgwdOB", "reviewer": "Reviewer_xnNs", "summary": "This 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.", "questions": "- 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?", "limitations": "yes", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written, and the proposed method is technically sound.\n- The learned transformation are well factorized and demonstrate flexible composability.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank $\\textcolor{blue}{Reviewer\\ xnNs}$ for the careful review and valuable advice. Below we address the concerns in detail.\n\n---\n\n**Q1. Disentanglement Metrics**\n\nThanks for the constructive suggestions! We strongly agree that the VP metric [a] indeed suits our method as it does not pose any assumptions on the latent space but only requires image pairs $\\mathbf{x} _0,\\mathbf{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 light-weighted neural network for learning to classify image pairs $\\mathbf{x} _0,\\mathbf{x} _T$ of different attributes. The generalization ability (validation accuracy) can be thus regarded as a reasonable surrogate for the disentanglement ability. The attribute correlation used in InterfaceGAN is also a good metric but it requires pre-trained attribute estimators to get scores, which is usually used for face image datasets. \n\n\nVP Scores (%) on MNIST\n|Training Set Split|Ours|PoFlow|TVAE|FactorVAE|$\\beta$-VAE|\n|:---:|:---:|:---:|:---:| :---:| :---:| \n|10%|**95.69**|93.05|89.91|85.92|87.31|\n|1%|**92.71**|91.27|88.15|84.86|85.25|\n\n\nVP Scores (%) on Shapes3D\n|Training Set Split|Ours|PoFlow|TVAE|FactorVAE|$\\beta$-VAE|\n|:---:|:---:|:---:|:---:| :---:| :---:|\n|10%|**95.92**|91.48|88.27|84.49|85.91|\n|1%|**77.03**|72.32|68.39|63.83|65.78|\n\n\nWe adopt the VP metric and evaluate the performance of the baseline methods on MNIST and Shapes3D. To ensure a fair comparison, for FactorVAE and $\\beta$-VAE, we choose the dimensions with the lowest equivariance errors to generate image pairs of different transformations. As indicated in the above two Tables, our method still achieves superior performance on the VP scores and outperforms the previous disentanglement baselines. This indicates that our flow-factorized VAE has better disentanglement performance. \n\n\n>[a] Learning disentangled representations with latent variation predictability. ECCV20.\n\n\n**Q2. Qualitative Comparison**\n\nSince the datasets used in the paper (MNIST and Shapes3D) are relatively simple, it is hard to visually inspect the difference unless cherry-picking extreme samples. However, as we have further applied our method to Falcol3D and Isaac3D (the two more complex real-world datasets), we did some qualitative comparisons on these two datasets. Please refer to the attached one-page pdf for the exemplary visual comparisons. Our method has obviously better equivariance properties and improves the image attribute control compared with other baselines.\n\n\n\n**Q3. Added Baselines**\n\nThanks for sharing the interesting works! DAVA [b] proposes an adversarial training procedure of VAEs to avoid manually tuning the dataset-specific regularization strength, while ControlVAE [c] defines a controller to automatically tune the regularization strength based on the proportional-integral-derivative (PID) control. \n\nSince both DAVA and ControlVAE can be regarded as improved extensions of $\\beta$-VAE, we use the same strategy of processing the results of $\\beta$-VAE to search for the best latent dimension that corresponds to the given transformation. The Table below displays the equivariance errors and VP scores of these two baselines against our methods. We would add these papers to the revised references and meanwhile do a more thorough literature review.\n\n\nEquivariance error ($\\downarrow$) and VP scores ($\\uparrow$) of DAVA and ControlVAE on MNIST\nMethods|Scaling|Rotation|Coloring|VP(10%)|VP(1%)|\n|:---:|:---:|:---:|:---:|:---:|:---:| \n|DAVA| 547.93±4.32| 523.98±5.47| 493.49±4.91|89.32|87.94 |\n|ControlVAE| 586.29±4.16|546.72±5.93 | 501.92±5.12|88.45|86.82 |\n|Ours|**185.42±2.35**|**153.54±3.10**|**158.57±2.95**|**95.69**|**92.71**|\n  \n \n>[b] DAVA: Disentangling Adversarial Variational Autoencoder. ICLR23.\n>\n>[c] Controlvae: Controllable variational autoencoder. ICML20.\n\n\n**Q4. Choosing the Value of $k$**\n\nFor MNIST, we choose the value of $k$ in line with TVAE and PoFlow. For Shapes3D, we use the self-contained transformations of the dataset but omit the transformation of \"Object Shape\". This is because the sequence is very short (4 elements) and its transition is very abrupt, which does not define a smooth and natural transformation like the other sequences. \n\nEquivariance error of different k on MNIST\nk|Scaling|Rotation|Coloring|\n|:---:|:---:|:---:|:---:| \n|1|185.27±2.59|-|-|\n|2|185.78±2.21|154.29±2.87|-|\n|3|185.42±2.45|153.54±3.10|158.57±2.95|\n\nWe did an ablation study on the impact of $k$ on MNIST and present the evaluation results in the above Table. 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 $\\nabla u$ on the potential landscape $u$; _**having multiple latent flows defined on different potential landscapes therefore does not interfere with each other**_.\n\nSince the value of $k$ is defined exactly the same as the number of transformations being applied, it is unlikely to set a $k$ larger than the actual number of given sequences. However, it might be possible to consider a larger $k$ in some unsupervised methods such as [d,e]. Then the model is likely to learn various combinations of the actual variation factors of the dataset. \n\n> [d] Unsupervised discovery of interpretable directions in the gan latent space. ICML20.\n>\n> [e] Latent Traversals in Generative Models as Potential Flows. ICML23.\n\n---\n\nThanks again and it is more than welcome to post further comments.\n"}, {"review_id": "djF4dn3W6c", "reviewer": "Reviewer_znKv", "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.", "questions": "Please refer to detailed comments in weaknesses.", "limitations": "This works shows no potential negative societal impact.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This work provides a new understanding of the disentanglement and equivariant, which are important problems for representation learning research.\n2. This paper proposes flow factorized representation learning which is a new and reasonable viewpoint for structured representation learning.\n3. Visualization results demonstrate that the proposed framework is able to switch transformations during the latent sample evolution and achieves both disentanglement and equivariance.\n4. 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.\n2. Advanced baseline methods are expected to be compared including StyleGAN, StyleGAN2, and Diffusion models.\n3. Can the learned disentangled representations be transferred across datasets, such as from Shapes3D to MNIST?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank $\\textcolor{green}{Reviewer\\ znKv}$ for the constructive suggestions. Below we respond to the concerns point by point.\n\n---\n\n**Q1. Application on Real-world Datasets**\n\nThanks for the great advice! As suggested by $\\textcolor{magenta}{Reviewer\\ ww47}$, we take a step further to apply our method on Falcol3D and Isaac3D, 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.\n\nEquivariance errors on Falcor3D\n|Methods|Lighting Intensity|Lighting X-dir|Lighting Y-dir|Lighting Z-dir|Camera X-pos|Camera Y-pos|Camera Z-pos|\n|:---:|:---:|:---:|:---:| :---:| :---:| :---:|  :---:| \n|TVAE|11477.81|12568.32|11807.34|11829.33|11539.69|11736.78|11951.45|\n|PoFlow|8312.97|7956.18|8519.39|8871.62|8116.82|8534.91|8994.63|\n|Ours|**5798.42**|**6145.09**|**6334.87**|**6782.84**|**6312.95**|**6513.68**|**6614.27**|\n\nEquivariance errors on Issac3D\n|Methods|Robot X-move|Robot Y-move|Camera Height|Object Scale|Lighting Intensity|Lighting Y-dir|Object Color|Wall Color|\n|:---:|:---:|:---:|:---:| :---:| :---:| :---:|  :---:| :---:| \n|TVAE|8441.65|8348.23|8495.31|8251.34|8291.70|8741.07|8456.78|8512.09|\n|PoFlow|6572.19|6489.35|6319.82|6188.59|6517.40|6712.06|7056.98|6343.76|\n|Ours|**3659.72**|**3993.33**|**4170.27**|**4359.78**|**4225.34**|**4019.84**|**5514.97**|**3876.01**|\n\n\n\nThe two Tables above present the equivariance errors of our method and TVAE/PoFlow. Notice that the values are much larger due to the increased resolution. Our method still outperforms other baselines by a large margin and achieves reasonable equivariance error. To visually display the learned latent transformations, we also conducted some qualitative evaluations and put them in the attached one-page pdf (which comprises exemplary latent evolutions of different transformations and exemplary comparisons against baselines). Our method can precisely control the image transformations through our latent flows.\n\n\n\n\n**Q2. Comparison against StyleGANs and Diffusion Models**\n\nIt is unlikely to compare our method with StyleGANs or diffusion models because their generative models are fundamentally different from ours. Our flow factorized VAE is defined on **sequences of data transformations**, while StyleGANs and diffusion models, despite being super powerful in generative modeling, are in essence generative models defined on **single images**. _Moreover, StyleGANs and diffusion models possess inherent architectural limitations of being simultaneously equivariant and disentangled._\n\nStyleGAN models and GANs in general do not have any encoders which could define the input space; the input to GANs is basically random noises. This architectural limitation makes it almost impossible to have equivariant image transformations. There exists a research line called GAN inversion which could identify the latent codes which closely reconstruct given images [a]. With the technique of GAN inversion, it might be possible to have an approximate equivariance mapping. However, the GAN inversion techniques usually require optimization for every single image, which makes it very complex and challenging to incorporate GANs into our framework. \n\nFor diffusion models, there is no concept of a well-defined latent space as they generate images by recursive denoising from a Gaussian. Consequently, there does not exist such a definition or concept of disentanglement for diffusion models. Some recent works hypothesize that the bottleneck of the diffusion network might contain some semantics [b]. However, the semantics discovery in the bottleneck requires complex optimization processes and even needs the help of the CLIP text encoder, which again makes it really challenging to design special diffuson models for our framework. \n\nIn a nutshell, the property of disentanglement and equivariance restricts us to implement our framework as encoder-decoder-like deep architectures for now. Nonetheless, we see that the flow-factorized VAE is a first step and we would keep investigating extensions to other generative models.\n\n\n> [a] PTI: Pivotal Tuning for Latent-based editing of Real Images. ACM TOG20.\n>\n> [b] Diffusion models already have a semantic latent space. ICLR23.\n\n\n**Q3. Transferability of Learned Transformations**\n\nThe transferability depends on the specific experimental setting. If we would like to validate an already-trained model on a new dataset, the reconstruction quality might be poor but the latent flow operator should be transferred (as indicated in Fig. 7 of the paper). \n\nOn the other hand, if we consider training an already-trained model on a new dataset, the learned transformations are not likely to be transferred because the model and the latent space would be completely different (as a consequence of training on new data). The corresponding latent flows would not lead to the prescribed transformations. \n\n---\n\nThanks again and it is more than welcome to ask further questions."}, {"review_id": "CCm7nzUv8x", "reviewer": "Reviewer_wDok", "summary": "The 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", "questions": "1. 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", "limitations": "- Limitations are discussed.\n- No code provided.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Solid model derivation\n- Strong experimental results showing empirical equivariance and disentanglement, though on basic datasets.\n", "weaknesses": "W1) 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank $\\textcolor{pink}{Reviewer\\ wDok}$ for the positive feedback and constructive suggestions. Below we respond to the concerns point by point.\n\n---\n\n**Q1. Pseudocode of the Algorithm**\n\nThanks for the constructive suggestion. We provide the following pseudo-codes for training our flow-factorized VAE:\n\n```python\nimport torch\n\n#Randomly sample a transformation at each iteration\nindex = torch.randint(0, potential_number) \nx_bar = sequence_generation(index)\n\n#Generating index according to the supervision setting\nif training_mode = \"supervised\":\n  index_potential = index\nelif training_mode = \"weakly-supervised\":\n  index_potential = q_k(x_bar)\n\n#initial element of the sequence\nz, rho_z = flow_vae(x_bar[0])\n\n#Future elements of the sequence obtained by latent evolution\nfor t in range(0,T)\n   PDE_loss, delta_z, delat_rho_z = HJ_PDE(index_potential,z,t) \n\n   #Updates in the sample and probability space\n   z = z + delta_z\n   rho_z = rho_z + delat_rho_z\n\n   #Inference at every intermediate steps\n   hat_xt = flow_vae.inference(z)\n\n   #Loss: PDE loss + reconstrutction loss + KL div \n   loss += PDE_loss + CE(hat_xt,x_bar[t]) + KL(rho_z, prior_rho_z)\n\n#KL div for index prediction (weakly-supervised setting)\nif training_mode = \"weakly-supervised\":\n   loss += KL(index_potential,index)\n\n\nloss.backward()\noptimizer.step()\n```\n\nHere we omit the computation of HJ PDEs for concisity. We will add the above pseudo-codes and an algorithm to the revised supplementary material. \n\n**Q2. Sequence Length and Intermediate Transformations**\n\nAs stated in line13 and line16 of the supplementary, the sequence length (T) for MNIST and Shapes3D is 9 and 8, respectively. For MNIST, we transform the data using the image geometric transformation functions imported from the $\\texttt{Kornia}$ library. For Shapes3D, this dataset itself has sequences of different transformations so we just use its self-contained sequences.  \n\nAs for the impact of sequence length, if the sequence is longer and has more variations, generally the equivariance error would be worse. To better illustrate this point, we did an ablation study on MNIST. 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. The Table below presents the results.  \n\nEquvariance error on MNIST for different sequence lengths\nSequence Length|Scaling|Rotation|\n|:---:|:---:|:---:| \n|9|185.42±2.35|153.54±3.10|\n|12|214.47±2.59|198.72±2.89|\n\nAs can be seen above, the equivariance error is larger as the sequence is longer and the variations are larger. Notice that even for the longer sequence, our method still outperforms other baselines with shorter sequences.\n\n**Q3. Computational Complexity**\n\nWe present the training time cost of our method on MNIST in the Table below, as well as the results of TVAE and PoFlow. In the supervised setting, our method is slightly more time-consuming than PoFlow (around 10%) and than TVAE (around 20%). When it comes to the weakly-supervised fashion, our flow-factorized VAE costs more time, as we need to use the Gumbel-Softmax trick for predicting the transformation index given the whole image sequence ($q(k,\\bar{\\mathbf{x}})$) and aggregate flows according to the index prediction ($\\mathbf{z} _{t+1} = \\mathbf{z} _t + \\sum _i^K q(k,\\bar{\\mathbf{x}}) _i  \\nabla u^i$). Nonetheless, we think the overall time cost is still acceptable considering the performance gain over the baselines (especially for the supervised setting).\n\n   Training time cost on MNIST every iteration\n||Ours (weakly-supervised)|Ours (supervised)|PoFlow|TVAE|\n|:---:|:---:|:---:|:---:| :---:| \n|Time (s)|1.076|0.723|0.638|0.591|\n\n\n**Q4. Application on Complex Real-world Datasets**\n\nThanks for the constructive advice! We follow your suggestion and take a step further to evaluate our method on Falcor3D and Isaac3D, 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. \n\nDue to the limited space, please refer to the response to $\\textcolor{red}{Reviewer\\ tD2D}$ and the attached one-page pdf for the relevant results and discussions.\n\n**Q5. Flow aggregation in extrapolation**\n\nAs indicated in line217 of the paper, the aggregation is simply done by linearly superposing different gradient fields. For example, if we would like to superpose scaling ($u^1$) and rotation ($u^2$), the new latent evolution would follow:\n\n$\\mathbf{z} _t = \\mathbf{z} _{t-1} + (\\nabla u^1 + \\nabla u^2)$\n\n\n**Q6. Code would be released upon acceptance**\n\nThanks for your interest! Our code is implemented in Pytorch and will be publicly released upon acceptance!\n\n---\n\nThanks again and it is more than welcome to post further comments."}]

Author rebuttal:
We sincerely thank all the reviewers for their constructive suggestions and valuable feedback. We appreciate the reviewers' common sentiment that our work is **intuitive and novel** ($\textcolor{red}{R1}$, $\textcolor{magenta}{R2}$, $\textcolor{green}{R4}$), **mathematically sound and technically solid** ($\textcolor{red}{R1}$, $\textcolor{magenta}{R2}$, $\textcolor{blue}{R3}$, $\textcolor{pink}{R5}$), **written and organized well** ($\textcolor{red}{R1}$, $\textcolor{magenta}{R2}$, $\textcolor{blue}{R3}$), and has **significant performance improvements** and **good generalization ability/flexible composability** ($\textcolor{red}{R1}$, $\textcolor{magenta}{R2}$, $\textcolor{blue}{R3}$, $\textcolor{green}{R4}$, $\textcolor{pink}{R5}$). We are also glad that $\textcolor{magenta}{R2}$ and $\textcolor{green}{R4}$ clearly agree that our work provides an **alternatively new viewpoint at the intersection of disentangled and equivariant representation learning**. 

We have replied to each reviewer individually to address the concerns. Below we answer the common questions and summarize the main changes in the response.

---

**Q1. Application on Complex Real-world Datasets ($\textcolor{red}{R1}$,$\textcolor{magenta}{R2}$,$\textcolor{green}{R4}$,$\textcolor{pink}{R5}$)**

One main concern of several reviewers is the applicability of our method on real-world datasets, as most works in the field benchmark their methods only on toy datasets. As suggested by $\textcolor{magenta}{R2}$, we take a step further to evaluate our approach on Falcol3D and Isaac3D, two complex real-world datasets that contain sequences of different transformations. Specifically, 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.

We have presented the quantitative evaluation in the individual responses, and the results demonstrate that our method still outperforms other approaches and achieves reasonable equivariance error. Moreover, to visually display the learned latent transformations, we also conducted some qualitative evaluations and attached them in the attached one-page pdf. Our method can precisely control the image transformations through our latent flows.

**Q2. Disentanglement Metrics ($\textcolor{magenta}{R2}$,$\textcolor{blue}{R3}$)** 

As stated in line29 of the supplementary (and also confirmed by $\textcolor{blue}{R3}$), most traditional disentanglement metrics like MIG and DCI do not suit our method because these metrics require that each latent dimension is responsible for one semantic and manipulating single dimensions of the latent variable would involve distinct transformations. However, for the recent disentanglement methods including ours, there emerges a more realistic disengagement setting: all the latent dimensions are perturbed by vectors for semantically meaningful output variations.

Nonetheless, also as pointed out by $\textcolor{blue}{R3}$, certain disentanglement metrics such as VP scores can be leveraged as they do not pose any assumptions of the latent space but only requires image pairs $\mathbf{x} _0,\mathbf{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 light-weighted neural network for learning to classify image pairs $\mathbf{x} _0,\mathbf{x} _T$. The generalization ability (validation accuracy) can be thus regarded as a reasonable surrogate for the disentanglement ability. 

We adopted the VP metric and evaluated the performance of the baseline methods on MNIST and Shapes3D. Please find the results in the individual responses. The evaluation results demonstrate that our method achieves consistently superior performance on the VP scores under different training set split ratios. This indicates that our flow-factorized VAE has better disentanglement performance. 


**Q3. Model Complexity ($\textcolor{red}{R1}$,$\textcolor{pink}{R5}$)** 

In the individual responses, we present the training time cost of our method. In the supervised setting, our method is only slightly more time-consuming than PoFlow (around 10%) and TVAE (around 20%). We hope the overall time cost of our method is acceptable considering the performance gain over the baselines.

 
**Q4. Ablation Study on $k$ ($\textcolor{red}{R1}$,$\textcolor{blue}{R3}$)** 

We choose the value of $k$ mainly following the literature and the specific dataset specifications. We conducted an ablation to investigate its impact and find out that the results are not affected by the number of transformations being applied. 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 $\nabla u$ on the potential landscape $u$; _**having multiple latent flows defined on different potential landscapes therefore does not interfere with each other**_.



---

*For brevity, we refer to reviewers $\textcolor{red}{tD2D}$ as $\textcolor{red}{R1}$, $\textcolor{magenta}{ww47}$ as $\textcolor{magenta}{R2}$, $\textcolor{blue}{xnNs}$ as $\textcolor{blue}{R3}$, $\textcolor{green}{znKv}$ as $\textcolor{green}{R4}$, and $\textcolor{pink}{wDok}$ as $\textcolor{pink}{R5}$, respectively.

Meta-review: The paper received mixed initial scores. The common concerns shared by reviewers include the absence of comparisons on complex datasets, the use of disentanglement metrics and the incorporation of modern generative baselines (such as styleGAN and diffusion models). 

The authors provide empirical results on two additional benchmarks, and provide disentanglement comparisons in their rebuttal. The authors’ responses steered reviewers towards a consensus to accept this paper. The final version should include the additional benchmark results and disentanglement comparisons discussed in the rebuttal, as well as more in-depth discussions concerning modern generative models.
Decision: Accept (poster)