Set-Based Flow Matching for Point-Cloud Generation

Generative models over point clouds and graphs have to deal with two properties that generative models over images or text don't: the data has no canonical ordering (a point cloud is a set, invariant under permutation of its points), and it can have variable size (a shape can be represented with 16 points or 300). This project builds a minimal flow-matching model that respects both properties by construction, then empirically tests a theoretical prediction about the architecture's capacity.

🧩 Method

The velocity field follows the DeepSets template: a shared per-point MLP embeds every point (conditioned on flow-time and shape identity), a permutation-invariant mean-pool forms a global context vector, and a second shared MLP combines each point's embedding with that context to predict its velocity. Because the network is applied point-wise with a symmetric pooling step, it is permutation-equivariant and has no dependence on the number of points N in its parameter count — the same weights generate a set of 8 points or 800.

Training uses conditional flow matching (Lipman et al., 2022) with the linear interpolation path between Gaussian noise and data; sampling integrates the learned ODE with simple Euler steps.

📊 Results

Trained for 7000 steps (~4 minutes, CPU-only) on four synthetic 2D point-set shapes (circle, square, star, two-moons) with N uniformly random in [16, 64] per batch.

Real vs. generated, conditioned on shape identity

Size transferability

Generating at sizes never seen during training (trained on N in [16, 64]; here N = 8, 16, 64, 150, 300) with the exact same weights, no retraining:

Training loss

The loss plateaus above zero, as expected: with independently sampled (non-optimal-transport) noise/data pairing, the conditional flow-matching loss floor is the conditional variance of the pairing itself, not the model's approximation error.

Sample quality vs. pooling width

The architecture's pooled context vector is a bottleneck through which all cross-point information must flow. A companion write-up derives why this width matters for universal approximation of permutation-invariant functions. Retraining the same model with context widths 20 through 26 and measuring Chamfer distance to real shapes:

Quality degrades sharply below a context width of about 4, then plateaus — consistent with the theoretical prediction that a sufficiently wide pooling bottleneck is necessary (and, past a point, no longer the limiting factor) for representing the relevant permutation-invariant summary of the point set.

💻 Code

Full implementation: PointCloudFlowMatching on GitHub.