ICML2026

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

Gaël Heck, Sylvie Le Hégarat-Mascle, Nicolas Lermé

摘要

Reassembling NN fragments in nn-dimensional space is a shape reconstruction task that is invariant to global rigid motions. Training directly on M=SE(n)N\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 M/SE(n)\mathcal{M}/\mathrm{SE}(n). Our unified SE(n)\mathrm{SE}(n)-invariant framework admits efficient closed-form 2D/3D instantiations and improves accuracy on polygonal jigsaw puzzles and 3D fracture reassembly benchmarks.