Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Equivariant Flow Matching for Point Cloud Assembly

Ziming Wang, Nan Xue, Rebecka Jörnsten

TL;DR

This work tackles multi-piece point cloud assembly under the symmetry of the group $SE(3)^N$ by developing an equivariant diffusion framework, Eda. A key theoretical result shows that learning an equivariant distribution via flow matching reduces to constructing related vector fields with invariant noise, avoiding the combinatorial cost of enforcing full $SE(3)^N$-equivariance. Eda parametrizes related vector fields with an $SO(3)^N$-equivariant network, uses a short, equivariant training path, and samples solutions through Runge-Kutta integration on $SE(3)^N$, achieving high data efficiency. Empirical results on 3DMatch and BB demonstrate strong performance, including non-overlapped inputs, outperforming correspondence-based and diffusion-based baselines while preserving efficiency and geometric priors.

Abstract

The goal of point cloud assembly is to reconstruct a complete 3D shape by aligning multiple point cloud pieces. This work presents a novel equivariant solver for assembly tasks based on flow matching models. We first theoretically show that the key to learning equivariant distributions via flow matching is to learn related vector fields. Based on this result, we propose an assembly model, called equivariant diffusion assembly (Eda), which learns related vector fields conditioned on the input pieces. We further construct an equivariant path for Eda, which guarantees high data efficiency of the training process. Our numerical results show that Eda is highly competitive on practical datasets, and it can even handle the challenging situation where the input pieces are non-overlapped.

Equivariant Flow Matching for Point Cloud Assembly

TL;DR

This work tackles multi-piece point cloud assembly under the symmetry of the group by developing an equivariant diffusion framework, Eda. A key theoretical result shows that learning an equivariant distribution via flow matching reduces to constructing related vector fields with invariant noise, avoiding the combinatorial cost of enforcing full -equivariance. Eda parametrizes related vector fields with an -equivariant network, uses a short, equivariant training path, and samples solutions through Runge-Kutta integration on , achieving high data efficiency. Empirical results on 3DMatch and BB demonstrate strong performance, including non-overlapped inputs, outperforming correspondence-based and diffusion-based baselines while preserving efficiency and geometric priors.

Abstract

The goal of point cloud assembly is to reconstruct a complete 3D shape by aligning multiple point cloud pieces. This work presents a novel equivariant solver for assembly tasks based on flow matching models. We first theoretically show that the key to learning equivariant distributions via flow matching is to learn related vector fields. Based on this result, we propose an assembly model, called equivariant diffusion assembly (Eda), which learns related vector fields conditioned on the input pieces. We further construct an equivariant path for Eda, which guarantees high data efficiency of the training process. Our numerical results show that Eda is highly competitive on practical datasets, and it can even handle the challenging situation where the input pieces are non-overlapped.

Paper Structure

This paper contains 33 sections, 11 theorems, 33 equations, 9 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Preliminaries
  4. Equivariances of PC assembly
  5. Vector fields and flow matching
  6. Method
  7. Problem formulation
  8. Equivariant flow
  9. Training
  10. Sampling via the Runge-Kutta method
  11. Implementation
  12. Equivariant attention layers
  13. Adaptive normalization and nonlinear layers
  14. Experiment
  15. Experiment settings
  16. ...and 18 more sections

Key Result

Theorem 4.2

Let $G$ be a smooth manifold, $F:G \rightarrow G$ be a diffeomorphism, and $P \in \mu(G)$. If vector field $v_X \in TG$ is $F$-related to vector field $v_{Y} \in TG$, then where $P_X = (\phi_X)_\# P_0$, $P_{Y}=(\phi_{Y})_\# (F_\#P_0)$. Here $\phi_X, \phi_{Y}: G \rightarrow G$ are generated by $v_X$ and $v_{Y}$ respectively.

Figures (9)

  • Figure 1: An overview of our model. The shapes of variables are shown in the brackets.
  • Figure 2: More details of Eda on 3DZ. \ref{['fig-3DZ-1']}: A result of Eda. Cameras are set to look at the room from above. Two PC pieces are marked by different colors. \ref{['fig-3DZ-3']}: the distribution of $\Delta r$ on the test set.
  • Figure 3: From left to right: the assembly process of a $8$-piece bottle by Eda.
  • Figure 4: The results of Eda on different number of pieces.
  • Figure 5: More details of Eda on 3DL. Left: the training curve. Middle and right: the influence of RK4/RK1 and the number of time steps on $\Delta r$ and $\Delta t$.
  • ...and 4 more figures

Theorems & Definitions (26)

  • Definition 3.1
  • Example 3.2
  • Remark 3.3
  • Theorem 4.2
  • Corollary 4.4
  • Proposition 4.5
  • Proposition 4.6
  • Proposition 4.7: Data efficiency
  • Definition D.1: Eqn. (10) in ryu2024diffusion and Def. (1) in ryu2022equivariant
  • Definition D.2: Restate $SO(3)^2$-equivariance and $SO(3)$-invariance in Def. \ref{['def:equivariance']} for pair-wise problems
  • ...and 16 more