SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Metrics. We use the approximate equivariance error  $\mathcal{E}_k$  and the log-likelihood of transformed data  $\log p(\pmb{x}_t)$  as the evaluation protocols. The equivariance error is defined as  $\mathcal{E}_k = \sum_{t=1}^T |\pmb{x}_t - \text{Decode}(\pmb{x}_t)|$  where  $\pmb{x}_t = \pmb{z}_0 + \sum_{t=1}^T \nabla_z u^k$ . For TVAE, the latent operator is changed to  $\text{Roi1}(\pmb{z}_0, t)$ . For unsupervised disentanglement baselines [35, 46] and SlowVAE [49], we carefully select the latent dimension and tune the interpolation range to attain the traversal direction and range that correspond to the smallest equivariance error. Since the vanilla VAE does not have the corresponding learned transformation in the latent space, we simply set  $\nabla_z u^k = 0$  and take it as a lower-bound baseline. For all the methods, the results are reported based on 5 runs.
Notice that the above equivariance error is defined in the output space. Another reasonable evaluation metric is instead measuring error in the latent space as  $\mathcal{E}_k = \sum_{t=1}^{T} |\operatorname{Encode}(\boldsymbol{x}_t) - \boldsymbol{z}_t|$ . We see the first evaluation method is more comprehensive as it further involves the decoder in the evaluation.
Qualitative results. Fig. 3 and 4 display decoded images of the latent evolution on MNIST [54] and Shapes3D [10], respectively. On both datasets, our latent flow can perform the target transformation precisely during evolution while leaving other traits of the image unaffected. In particular, for the weakly-supervised setting, the decoded images (i.e., the bottom rows of Fig. 3 and 4) can still reproduce the given transformations well and it is even hard to visually tell them apart from the generated images under the supervised setting. This demonstrates the effectiveness of the 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.
Methods	Supervision?	Equivariance Error (↓)	Log-likelihood (↑)
Floor Hue	Wall Hue	Object Hue	Scale
VAE [47]	No (X)	6924.63±8.92	7746.37±8.77	4383.54±9.26	2609.59±7.41	-11784.69±4.87
β-VAE [35]	No (X)	2243.95±12.48	2279.23±13.97	2188.73±12.61	2037.94±11.72	-11924.83±5.64
FactorVAE [46]	No (X)	1985.75±13.26	1876.41±11.93	1902.83±12.27	1657.32±11.05	-11802.17±5.69
SlowVAE [49]	Weak (✓)	1247.36±12.49	1314.86±11.41	1102.28±12.17	1058.74±10.96	-11674.89±5.74
TVAE [45]	Yes (✓)	1225.47±9.82	1246.32±9.54	1261.79±9.86	1142.01±9.37	-11475.48±5.18
PoFlow [79]	Yes (✓)	885.46±10.37	916.71±10.49	912.48±9.86	924.39±10.05	-11335.84±4.95
Ours	Yes (✓)	613.29±8.93	653.45±9.48	605.79±8.63	599.71±9.34	-11215.42±5.71
Ours	Weak (✓)	690.84±9.57	717.74±10.65	681.59±9.02	653.58±9.57	-11279.61±5.89
Table 2: Equivariance error  $\mathcal{E}_k$  and log-likelihood  $\log p(\pmb{x}_t)$  on Shapes3D [10].
Quantitative results. Tables 1 and 2 compare the equivariance error and the log-likelihood on MNIST [54] and Shapes3D [10], respectively. Our method learns the latent flows which model the transformations precisely, achieving the best performance across datasets under different supervision settings. Specifically, our method outperforms the previous best baseline by 69.74 on average in the equivariance error and by 32.58 in the log-likelihood on MNIST. The performance gain is also consistent on Shapes3D: our method surpasses the second-best baseline by 291.70 in the average equivariance error and by 120.42 in the log-likelihood. In the weakly-supervised setting, our method also achieves very competitive performance, falling behind that of the supervised setting in the average equivariance error slightly by 6.22 on MNIST and by 67.88 on Shapes3D.
4.3 Discussion
Extrapolation: switching transformations. In Fig. 5 we demonstrate that, empowered by our method, it is possible to switch latent transformation categories mid-way through the latent evolution
Figure 5: Exemplary visualization of switching transformations during the latent sample evolution.
and maintain coherence. That is, we perform  $\boldsymbol{z}_t = \boldsymbol{z}_{t-1} + \nabla_z u^k$  for  $t \leq T/2$  and then change to  $\boldsymbol{z}_t = \boldsymbol{z}_{t-1} + \nabla_z u^j$  where  $j \neq k$  for  $t > T/2$ . As can be seen, the factor of variation immediately changes after the transformation type is switched. Moreover, the transition phase is smooth and no other attributes of the image are influenced.
Figure 6: Examples of combining different transformations simultaneously during the latent evolution.
Extrapolation: superposing transformations. Besides switching transformations, our method also supports applying different transformations simultaneously, i.e., consistently performing  $z_{t} = z_{t-1} + \sum_{k}^{K} \nabla_{z} u^{k}$  during the latent flow process. Fig. 6 presents such exemplary visualizations of superposing two and all transformations simultaneously. In each case, the latent evolution corresponds to simultaneous smooth variations of multiple image attributes. This indicates that our method also generalizes well to superposing different transformations.
Notice that we only apply single and separate transformations in the training stage. Switching or superposing transformations in the test phase can be thus understood as an extrapolation test to measure the generalization ability of the learned equivariance to novel compositions.
Equivariance generalization to new data. We also test whether the learned equivariance holds for Out-of-Distribution (OoD) data. To verify this, we validate our method on a test dataset that is different from the training set and therefore unseen to the model. Fig. 7 displays the exemplary visualization results of the VAE trained on MNIST [54] but evaluated on dSprites [59]. Although the reconstruction quality is poor, the learned equivariance is still clearly effective as each transformation still operates as expected: scaling, rotation, and coloring transformations from top to bottom respectively.
4.4 Results on Complex Real-world and Large-scale Datasets
Table 3 and 4 compare the equivariance error of our methods and the representative baselines on Falcol3D and Isaac3D, respectively. Notice that the values are much larger than previous datasets due to the increased image resolution. Our method still outperforms other baselines by a large margin
Figure 7: Equivariance generalization to unseen OoD input data. Here the model is trained on MNIST [54] but the latent flow is tested on dSprites [59].
Methods	Lighting Intensity	Lighting X-dir	Lighting Y-dir	Lighting Z-dir	Camera X-pos	Camera Y-pos	Camera Y-pos
TVAE [45]	11477.81	12568.32	11807.34	11829.33	11539.69	11736.78	11951.45
PoFlow [79]	8312.97	7956.18	8519.39	8871.62	8116.82	8534.91	8994.63
Ours	5798.42	6145.09	6334.87	6782.84	6312.95	6513.68	6614.27
and achieves reasonable equivariance error. Fig. 8 displays the qualitative comparisons of our method against other baselines. Our method precisely can control the image transformations through our latent flows. Overall, the above results demonstrate that our method can go beyond the toy setting and can be further applied to more complex real-world scenarios.
More visualization results of exemplary latent flows are kindly referred to in the supplementary.
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  $\beta$ -VAE [35]. InfoGAN [13] achieves disentanglement by maximizing the mutual information between a subset of latent dimensions and observations, while  $\beta$ -VAE [35] induces the factorized posterior  $q(\mathbf{z})$  by penalizing the Total Correlation (TC) through an extra hyper-parameter  $\beta > 1$  controlling the strength of the KL divergence. Following infoGAN, many attempts have been made to facilitate the discovery of semantically meaningful traversal directions through regularization [33, 42, 89, 34, 100, 66, 77, 90, 98, 84, 99, 78, 62]. The follow-up research of  $\beta$ -VAE mainly explored different methods to factorize the aggregated posterior [22, 25, 52, 46, 12, 44, 96, 23, 76, 58, 80, 28]. More recently, some works proposed to discover meaningful directions of diffusion models in the bottleneck of denoising networks [53, 64, 95, 41]. The previous literature mainly considers disentanglement as learning different transformations per dimension or per linear direction. Our method generalizes this concept to learning a distinct tangent bundle  $\nabla u^k$  that moves every latent sample via dynamic OT.
We see the most similar method to ours is the work of [79]. In [79], the authors also apply the gradient of a potential function to move the latent code. However, their potentials are restricted to obey the wave equations, which do not really correspond to the OT theory. Also, they do not consider the posterior evolution but instead use the loss  $||\pmb{z}_t - \mathsf{Encode}(\pmb{x}_t)||^2$  to match the latent codes. By contrast, we propose a unified probabilistic generative model that encompasses the posterior flow that follows dynamic OT, the flow-like time evolution, and different supervision settings.
Equivariant neural networks. A function is said to be an equivariant map if it commutes with a given transformation, i.e.,  $T'[f(x)] = f(T[x])$  where  $T$  and  $T'$  represent operators in different domains. Equivariance has been considered a desired inductive bias for deep neural networks as this property can preserve geometric symmetries of the input space [38, 75, 56, 57, 1]. Analytically equivariant
Table 3: Equivariance error (↓) on Falcol3D [61].
Methods	Robot X-move	Robot Y-move	Camera Height	Object Scale	Lighting Intensity	Lighting Y-dir	Object Color	Wall Color
TVAE [45]	8441.65	8348.23	8495.31	8251.34	8291.70	8741.07	8456.78	8512.09
PoFlow [79]	6572.19	6489.35	6319.82	6188.59	6517.40	6712.06	7056.98	6343.76
Ours	3659.72	3993.33	4170.27	4359.78	4225.34	4019.84	5514.97	3876.01
Table 4: Equivariance error (↓) on Isaac3D [61].
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 [16, 17, 68, 93, 92, 85, 31, 39]. Another line of research proposed to directly learn approximate equivariance from data [21, 18, 49, 20, 45]. Our framework re-defines approximate equivariance by matching the latent probabilistic flow to the actual path of the given transformation in the image space.
Optimal transport in deep learning. There is a vast literature on OT theory and applications in various fields [87, 88]. Here we mainly highlight the relevant applications in deep learning. The pioneering work of [19] proposed a light-speed implementation of the Sinkhorn algorithm for fast computation of entropy-regularized Wasserstein distances, which opened the way for many differentiable Sinkhorn algorithm-based applications [32, 29, 14, 27, 51]. In generative modeling, the Wasserstein distance is often used to minimize the discrepancy between the data distribution and the model distribution [2, 81, 72, 65]. Inspired by the fluid mechanical interpretation of OT [4], some normalizing flow methods [69, 24, 48] considered regularizing the velocity fields to satisfy the HJ equation, thus matching the dynamic OT plan [94, 30, 83, 63, 60]. Our method applies PINNs [67] to directly model generalized HJ equations in the latent space and uses the gradient fields of learned potentials to generate latent flows, which also aligns to the theory of dynamic fluid mechanical OT.
In this paper, we introduce Flow Factorized Representation Learning which defines a set of latent flow paths that correspond to sequences of different input transformations. The latent evolution is generated by the gradient flow of learned potentials following dynamic optimal transport. Our setup reinterprets the concepts of both disentanglement and equivariance. Extensive experiments demonstrate that our model achieves higher likelihoods on standard representation learning benchmarks while simultaneously achieving smaller equivariance error. Furthermore, we show that the learned latent transformations generalize well, allowing for flexible composition and extrapolation to new data.
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 inferring with observed  $k$ , we change the posterior from  $q(\bar{z}|\bar{x},k)$  to  $q(\bar{z}|\pmb{x}_0,k)$  because we assume  $k$  contains sufficient information of the whole sequence. To keep the posterior definition of  $q(\bar{z}|\bar{x},k)$ , we need to make  $q(\pmb{z}_t)$  also a function of  $\pmb{x}_t$ . This can be achieved either by changing the potential to  $u(\pmb{z}_{t-1},\pmb{x}_t,t-1)$  or modifying the external driving force to  $f(\pmb{z}_{t-1},\pmb{x}_t,t-1)$ . Nonetheless, we see these modifications would make the model less flexible than our current formulations as the element  $\pmb{x}_t$  might be needed during inference.
This work was supported by the MUR PNRR project FAIR (PE00000013) funded by the NextGenerationEU, by the PRIN project CREATIVE (Prot. 2020ZSL9F9), by the EU H2020 project AI4Media (No. 951911), and by the Bosch Center for Artificial Intelligence.
[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.
A.1 Pseudo codes
import torch   
#Randomly sample a transformation at each iteration   
index  $=$  torch.randint(0,potential_number)   
x_bar  $\equiv$  sequence_generation(index)   
#Generating index according to the supervision setting   
if training_mode  $=$  "supervised": index_potential  $=$  index   
elif training_mode  $=$  "weakly-supervised": index_potential  $=$  q_k(x_bar)   
#initial element of the sequence   
z, rho_z  $=$  flow_vae(x_bar[0])   
#Future elements of the sequence obtained by latent flow   
for t in range(O,T) PDE_loss, delta_z, delat_rho_z  $=$  HJ_PDE(indexpotential,z,t) #Updates in the sample and probability space z  $=$  z + delta_z rho_z  $=$  rho_z + delat_rho_z #Inference at every intermediate step hat xt  $=$  flow_vae.inference(z) #Loss: PDE loss  $^+$  reconstruction loss  $^+$  KL div loss  $+ =$  PDE_loss  $^+$  CE(hat xt,x_bar[t]）  $^+$  KL(rho_z,prior_rho_z)   
#KL div for index prediction (weakly-supervised setting)   
if training_mode  $=$  "weakly-supervised": loss  $+ =$  KL(indexPotential,index)   
lossbackward()   
optimizer-step()
Figure 9: Pytorch-like pseudo codes for training our flow-factorized VAE.
Fig. 9 displays the Pytorch implementation for training our flow-factorized VAE under different supervision settings. Here we omit the computation of HJ PDEs for concisenity.
A.2 Implementation details
Common settings. During the training stage, we randomly sample one single transformation at each iteration. The batch size is set to 128 for both datasets. We use Adam optimizer and the learning rate is set as  $1e - 4$  for all the parameters. The encoder consists of four stacked convolution layers with the activation function ReLU, while the decoder is comprised of four stacked transposed convolution layers. For the prior evolution, the diffusion coefficient  $D_{k}$  is initialized with 0 and we set it as a learnable parameter for distinct  $k$ . For MLPs that parameterize the potential  $u(z,t)$  and the force  $f(z,t)$ , we use the sinusoidal positional embeddings [86] to embed the timestep  $t$ , and use linear layers for embedding the latent code  $z$ . Tanh gates are applied as the activation functions of the MLPs. All the experiments are run on a single NVIDIA Quadro RTX 6000 GPU.
MNIST. The input images are of the size  $28 \times 28$ . The sequence of each transformation contains 9 states of variations. The scaling transformation scales the image from 1.0 up to 1.8 times. The rotation transformation rotates the object by maximally 80 degrees, and the coloring transformation adjusts the image hue from 0 to 340 degrees. The model is trained for 90,000 iterations.
Shapes3D. The input images are resized to  $64 \times 64$ . Each transformation sequence consists of 8 images. The model is also trained for 90,000 iterations.
Falcon3D and Isaac3D. The input images are in the resolution of  $128 \times 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-Softmax trick, we re-parameterize  $q_{\gamma}(k|\bar{x})$  by
$$
y _ {i} = \frac {e ^ {\frac {x _ {i} + g _ {i}}{\tau}}}{\sum_ {i} e ^ {\frac {x _ {i} + g _ {i}}{\tau}}} \tag {15}
$$
where  $x_{i}$  is the category prediction,  $g_{i}$  is the sample drawn from Gumbel distributions, and  $\tau$  is the small temperature to make softmax behave like argmax. We take the 'hard' binary prediction in the forward pass and use the straight-through gradient estimator [6] during backpropagation. The temperature  $\tau$  is initialized with 1 and is 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  $\beta$ -VAE and FactorVAE to learn the given transformations in an unsupervised manner. For tuning the interpolation range, we start from the initial value  $z_{i}$  and traverse till the appropriate bound which is selected from the range  $[-5, 5]$  with the interval of 0.1.
A.3 Disentanglement metrics
There are many traditional disentanglement metrics [71, 26, 12], but they are designed for single-dimension traversal methods. These metrics assume and require that each latent dimension is responsible for one semantic and manipulating single dimensions of the latent variable would involve distinct output transformations. However, for the recent disentanglement methods including ours [77, 84, 79], there emerges a more realistic disengagement setting: all the latent dimensions are perturbed by vectors for meaningful output variations. When it comes to these vector-based disentanglement methods, their scores of disentanglement metrics would drop considerably and cannot be compared with those single-dimension baselines.
Table 5: VP Scores (%) on MNIST.
Training Set Split	Ours	PoFlow	TVAE	FactorVAE	β-VAE
10%	95.69	93.05	89.91	85.92	87.31
1%	92.71	91.27	88.15	84.46	85.25
Table 6: VP Scores  $(\%)$  on Shapes3D.
Training Set Split	Ours	PoFlow	TVAE	FactorVAE	β-VAE
10%	95.92	91.48	88.27	84.49	85.91
1%	77.03	72.32	68.39	63.83	65.78
Nonetheless, certain disentanglement metrics such as VP scores [100] can be leveraged as they do not pose any assumptions on the latent space but only require image pairs  $\left[\pmb{x}_0,\pmb{x}_T\right]$  of different transformations for evaluation. The VP metric adopts the few-shot learning setting (using  $1\%$  or  $10\%$  of the dataset as the training set) and takes a lightweight neural network for learning to classify image pairs  $\left[\pmb{x}_0,\pmb{x}_T\right]$  of different attributes. The generalization ability (i.e., validation accuracy) can be thus regarded as a reasonable surrogate for the disentanglement ability. Table 5 and 6 present the VP scores of all the baseline methods on MNIST and Shapes3D. To ensure a fair comparison, for FactorVAE and  $\beta$ -VAE, we choose the dimensions with the lowest equivariance errors to generate image pairs of different transformations. Our method outperforms the previous disentanglement baselines and achieves superior performance on the VP scores. This indicates that our flow-factorized VAE has better disentanglement ability.
A.4 Ablation studies
Impact of different priors. We use diffusion equations to model the prior evolution as random particle movement. It would also be intriguing to choose other priors commonly used in the VAE literature, such as Standard Gaussian (SG) priors  $\mathcal{N}(0,1)$ , mixture of Gaussian (MoG) priors  $\sum w_{i}\mathcal{N}(\mu_{i},\sigma_{i}^{2})$  and VAMP priors [82] which average aggregated posterior of  $N$  pseudo-inputs as  $1 / N\sum_{n}q(z_{n})$ . Table 7 presents the equivariance error of different priors on MNIST. Among these priors, our diffusion equations achieve the best performance. This meets our assumption that modeling the prior evolution
Table 7: Equivariance error of different priors.
Prior	Scaling	Rotation	Coloring
SG	190.24±2.18	158.93±3.25	164.18±2.77
MoG	188.23±2.45	157.79±2.86	161.49±2.62
VAMP	192.81±3.67	161.47±4.12	162.97±3.89
Diffusion	185.42±2.35	153.54±3.10	158.57±2.95
Table 8: Equivariance error of different PDEs.
Prior	Scaling	Rotation	Coloring
Heat	223.95±3.38	212.47±3.85	207.66±2.91
FP	211.54±3.17	188.59±3.92	194.73±3.09
OHJ	190.43±2.48	163.87±3.03	162.38±2.86
GHJ	185.42±2.35	153.54±3.10	158.57±2.95
as a diffusion process suits more the random motion. Nonetheless, we see that the performance gap between each baseline is narrow, which somehow implies that the impact of different priors is limited.
Impact of different PDEs. We apply the generalized HJ (GHJ) equation as the PINN constraint in order to achieve dynamic OT. It would be also interesting to try other commonly used PDEs. We compare our GHJ with the ordinary HJ (OHJ) equation, the Fokker Planck (FP) equation, and the heat equation. Table 8 compares the equivariance error of PDEs on MNIST. Our GHJ and OHJ equations achieve the best performance as they both satisfy the condition of dynamic OT. This empirical evidence indicates that the OT theory can indeed model better latent flow paths. Moreover, our GHJ outperforms the OHJ by a slight margin. We attribute this advantage to the external driving force  $f(z,t)$  which gives us more flexibility and dynamics in modeling the velocity fields  $\nabla u^k$ .
Table 9: Equivariance error on MNIST of a different number of transformations  $(K)$ .
K	Scaling	Rotation	Coloring
1	185.27±2.59	-	-
2	185.78±2.21	154.29±2.87	-
3	185.42±2.45	153.54±3.10	158.57±2.95
Table 10: Equivariance error on MNIST of different sequence lengths (T).
Sequence Length (T)	Scaling	Rotation
9	185.42±2.35	153.54±3.10
12	214.47±2.59	198.72±2.89
Impact of different  $K$ . We conduct an ablation study on the impact of the number of transformations on MNIST and present the evaluation results in Table 9. As indicated above, in general, the performance is not affected by the number of transformations being applied. The fluctuation of the results when  $K$  varies can be sufficiently negligible. We expect that this is because the transformations are learned by distinct potentials (which are implemented as  $K$  different MLPs). Each flow evolves along with the gradient field  $\nabla u$  on the potential landscape  $u$ ; having multiple latent flows defined on different potential landscapes therefore does not interfere with each other.