SE(n)-Invariant Flow Matching: A General Framework with Application to Object Reassembly

Reassembling $N$ fragments in $n$-dimensional space is a shape reconstruction task that is invariant to global rigid motions. Training directly on $\mathcal{M}=\mathrm{SE}(n)^N$ can be ill-posed: standard losses penalize solutions that differ only by a global transform. Existing methods often address this with ad-hoc anchoring which breaks permutation invariance across fragments and can introduce biases that must be mitigated with extensive and costly data augmentation. We propose a geometric framework that enforces invariance by construction. First, a **Global Gauge Fixing** (GGF) strategy deterministically aligns configurations using an intrinsic generalized-inertia rule. Second, we introduce a **quotient-invariant Flow Matching objective** that operates via orthogonal projection onto the horizontal tangent bundle. This construction factors out global pose at each timestep, enabling the model to learn only shape-changing dynamics on the quotient space $\mathcal{M}/\mathrm{SE}(n)$. Our unified $\mathrm{SE}(n)$-invariant framework admits efficient closed-form 2D/3D instantiations and improves accuracy on polygonal jigsaw puzzles and 3D fracture reassembly benchmarks.

Reassembling an object from fragments is challenging because the same correct assembly can be shifted or rotated without changing how the pieces fit together. This creates an ambiguity for machine learning methods: they may spend effort learning arbitrary global positions instead of learning the relative arrangement of the fragments. We address this by removing the irrelevant global motion from the learning problem. Our method first places each candidate assembly in a common reference frame, then trains the model to ignore motions that move all fragments together. This makes the model focus on the motions that actually change the shape of the assembled object. The approach applies to both 2D and 3D rigid-fragment reassembly. We evaluate it on polygonal jigsaw puzzles, floorplan reconstruction, archaeological reassembly, and 3D fracture benchmarks, where it improves reconstruction accuracy over the compared baselines.