Not-So-Optimal Transport Flows for 3D Point Cloud Generation

Learning generative models of 3D point clouds is one of the fundamental problems in 3D generative learning. One of the key properties of point clouds is their permutation invariance, i.e., changing the order of points in a point cloud does not change the shape they represent. In this paper, we analyze the recently proposed equivariant OT flows that learn permutation invariant generative models for pointbased molecular data and we show that these models scale poorly on large point clouds. Also, we observe learning (equivariant) OT flows is generally challenging since straightening flow trajectories makes the learned flow model complex at the beginning of the trajectory. To remedy these, we propose not-so-optimal transport flow models that obtain an approximate OT by an offline OT precomputation, enabling an efficient construction of OT pairs for training. During training, we can additionally construct a hybrid coupling by combining our approximate OT and independent coupling to make the target flow models easier to learn. In an extensive empirical study, we show that our proposed model outperforms prior diffusion-and flow-based approaches on a wide range of unconditional generation and shape completion on the ShapeNet benchmark. Generating 3D point clouds is one of the fundamental problems in 3D modeling with applications in shape generation, 3D reconstruction, 3D design, and perception for robotics and autonomous systems. Recently, diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) and flow matching (Lipman et al., 2022) have become the de facto frameworks for learning generative models for 3D point clouds. These frameworks often overlook 3D point cloud permutation invariance, implying the rearrangement of points does not change the shape that they represent. In closely related areas, equivariant optimal transport (OT) flows (Klein et al., 2024; Song et al., 2024) have been recently developed for 3D molecules that can be considered as sets of 3D atom coordinates.