Shape Assembly via Equivariant Diffusion

Sangjin Lee, Woohyeon Joseph Shim, Jungtaek Kim, Minsu Cho

Keywords: Diffusion, Equivariant diffusion, Shape assembly

Abstract: We tackle the problem of solving shape puzzles, that is, reassembling randomly-partitioned and scattered pieces of 2D or 3D shapes into an original shape. This task is challenging since it only relies on geometric features without rich visual information. Specifically, we are supposed that target shapes and their randomly-partitioned pieces are pattern-free and irregular. Existing methods tend to rely on specific constraints regarding piece shapes and neglect the consideration of invariance and equivariance. We propose learning a robust puzzle solver through a generative diffusion process in which the roto-translational equivariance holds. Experiments on 2D and 3D puzzle benchmarks including the Breaking Bad dataset demonstrate that our method successfully assembles given geometric pieces into a target shape. We also provide in-depth ablation studies showing the effects of our equivariant design and the components in our proposed framework.

Primary Area: applications to computer vision, audio, language, and other modalities

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 5468