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.
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.
The 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.
As 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.
Questions
In equation 6, will t=0 be problematic?
In 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.
The 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).
According 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.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Limitations
Do not identify any negative societal impact.