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

Abstract

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.